1# Copyright (c) 2025 Espressif Systems (Shanghai) Co., Ltd.
2# SPDX-License-Identifier: Apache-2.0
3
4import sys
5import os
6from datetime import datetime
7
8output_path = ""
9start_year = 2024
10
11scheme_default = [
12    ("boot",         "mcuboot",        128),
13    ("slot0",        "image-0",        -1),
14    ("slot1",        "image-1",        -1),
15    ("slot0-lpcore", "image-0-lpcore", 32),
16    ("slot1-lpcore", "image-1-lpcore", 32),
17    ("storage",      "storage",        192),
18    ("scratch",      "image-scratch",  124),
19    ("coredump",     "coredump",       4)
20]
21
22scheme_amp = [
23    ("boot",         "mcuboot",        128),
24    ("slot0",        "image-0",        -3),
25    ("slot1",        "image-1",        -3),
26    ("slot0-appcpu", "image-0-appcpu", -1),
27    ("slot1-appcpu", "image-1-appcpu", -1),
28    ("slot0-lpcore", "image-0-lpcore", 32),
29    ("slot1-lpcore", "image-1-lpcore", 32),
30    ("storage",      "storage",        192),
31    ("scratch",      "image-scratch",  124),
32    ("coredump",     "coredump",       4)
33]
34
35flash_sizes_mb = [2, 4, 8, 16, 32]
36flash_offsets_kb = [0, 4]
37default_flash_mb = 4
38
39def make_layout(flash_size_mb, flash_offset_kb, scheme):
40    """
41    Create a valid partitions layout based on provided scheme, flash size and flash offset.
42
43    :param flash_size_mb: Total size of flash in MBytes
44    :param flash_offset_kb: Flash start offset in kBytes
45    :param scheme: Partition scheme used.
46    """
47
48    fixed_size = 0
49    sum_weights = 0
50    flash_size = flash_size_mb * 1024 * 1024
51    offset = flash_offset_kb * 1024
52
53    for i, (_, _, size) in enumerate(scheme):
54        if size < 0:
55            sum_weights += abs(size)
56        else:
57            fixed_size += size * 1024
58
59    variable_size = flash_size - fixed_size
60    if (variable_size % 0x10000) != 0:
61        print("Variable size area must be in orders of 64kB due to the cache mappings. Correct the size of fixed allocations and try again.")
62        sys.exit(1)
63
64    variable_blocks = variable_size / 0x10000
65
66    variable_alloc = []
67    allocated_blocks = 0
68    for i, (n, l, w) in enumerate(scheme):
69        if w < 0:
70            alloc = abs(w) * variable_blocks // sum_weights
71            allocated_blocks += alloc
72            variable_alloc.append([alloc, i])
73
74    leftover_blocks = variable_blocks - allocated_blocks
75
76    if leftover_blocks % 2 != 0:
77        print("Odd number of leftower blocks cannot be distributed. Increase or decrease fixed partitions size by 64kB and try again.")
78        sys.exit(1)
79
80    if leftover_blocks:
81        variable_alloc[0][0] += 1
82        variable_alloc[1][0] += 1
83
84    sum_kb = 0
85    partition_layout = []
86
87    for i, (n, l, s) in enumerate(scheme):
88        size_kb = 0
89        if (s < 0):
90            for alloc, index in variable_alloc:
91                if i == index:
92                    size_kb = alloc * 64
93        else:
94            size_kb = s
95        if i == 0:
96            size_kb -= flash_offset_kb
97        sum_kb += size_kb
98        partition_layout.append((n, l, int(size_kb)))
99    if sum_kb + flash_offset_kb != flash_size_mb * 1024:
100        print("Something went wrong. Check the scheme and try again!")
101        sys.exit(1)
102
103    return partition_layout
104
105def generate_partition_table_dtsi(flash_size_mb, flash_offset_kb, scheme, scheme_name):
106    """
107    Generate a DTSI file with a static flash partition layout.
108
109    :param flash_size_mb: Total size of flash in MBytes
110    :param flash_offset_kb: Flash start offset in kBytes
111    :param scheme_name: Scheme name used in file name
112    :param scheme: Partition scheme used.
113    """
114
115    print(f"Generate partition for {flash_size_mb}MB flash with offset 0x{flash_offset_kb*1024:x}kB")
116
117    # Generate layout
118    partitions = make_layout(flash_size_mb, flash_offset_kb, scheme)
119
120    flash_size = flash_size_mb * 1024 * 1024
121    current_offset = flash_offset_kb * 1024
122    output = []
123
124    output.append("/*")
125    if start_year < datetime.now().year:
126        output.append(f" * Copyright (c) {start_year}-{datetime.now().year} Espressif Systems (Shanghai) Co., Ltd.")
127    else:
128        output.append(f" * Copyright (c) {datetime.now().year} Espressif Systems (Shanghai) Co., Ltd.")
129    output.append(" *")
130    output.append(" * SPDX-License-Identifier: Apache-2.0")
131    output.append(" */")
132    output.append("")
133    output.append("&flash0 {")
134    output.append("\tpartitions {")
135    output.append("\t\tcompatible = \"fixed-partitions\";")
136    output.append("\t\t#address-cells = <1>;")
137    output.append("\t\t#size-cells = <1>;")
138
139    for i, (name, label, size_kb) in enumerate(partitions):
140        size = size_kb * 1024
141        if current_offset + size > flash_size:
142            print(f"Error: Partition {name} exceeds flash size. Aborting.")
143            sys.exit(1)
144
145        output.append("")
146        output.append(f"\t\t{name.replace('-', '_')}_partition: partition@{current_offset:x} {{")
147        output.append(f"\t\t\tlabel = \"{label}\";")
148        output.append(f"\t\t\treg = <0x{current_offset:X} DT_SIZE_K({int(size_kb)})>;")
149        output.append("\t\t};")
150        print(f"{i}, {name:>16},{label:>16}, {current_offset:>8X}, DT_SIZE_K({size_kb})")
151        current_offset += size
152
153    output.append("\t};")
154    output.append("};")
155    output.append("")
156    output.append(f"/* Remaining flash size is {int((flash_size - current_offset) / 1024)}kB")
157    output.append(f" * Last used address is 0x{current_offset-1:X}")
158    output.append(f" */")
159    output.append("")
160
161    output_file = f"{output_path}/partitions_{hex(flash_offset_kb * 1024)}_{scheme_name}_{int(flash_size_mb)}M.dtsi"
162
163    with open(output_file, "w") as f:
164        f.write("\n".join(output))
165
166    print(f"Writing file {output_file}")
167
168def generate_default_partition_table_dtsi(flash_size_mb, flash_offset_kb, scheme_name):
169
170    output = []
171
172    output.append("/*")
173    if start_year < datetime.now().year:
174        output.append(f" * Copyright (c) {start_year}-{datetime.now().year} Espressif Systems (Shanghai) Co., Ltd.")
175    else:
176        output.append(f" * Copyright (c) {datetime.now().year} Espressif Systems (Shanghai) Co., Ltd.")
177    output.append(" *")
178    output.append(" * SPDX-License-Identifier: Apache-2.0")
179    output.append(" */")
180    output.append("")
181    output.append(f"#include <espressif/partitions_0x{(flash_offset_kb*1024):X}_{str(scheme_name)}_{int(flash_size_mb)}M.dtsi>")
182    output.append("")
183
184    output_file = f"{output_path}/partitions_{hex(offset_kb * 1024)}_{scheme_name}.dtsi"
185
186    with open(output_file, "w") as f:
187        f.write("\n".join(output))
188
189    print(f"Writing file {output_file}")
190
191if __name__ == "__main__":
192
193    output_path = os.getcwd()
194
195    if len(sys.argv) > 1:
196        output_path = sys.argv[1]
197
198    for offset_kb in flash_offsets_kb:
199        for flash_mb in flash_sizes_mb:
200            generate_partition_table_dtsi(flash_mb, offset_kb, scheme_default, "default")
201        generate_default_partition_table_dtsi(default_flash_mb, offset_kb, "default")
202
203    for offset_kb in flash_offsets_kb:
204        for flash_mb in flash_sizes_mb:
205            generate_partition_table_dtsi(flash_mb, offset_kb, scheme_amp, "amp")
206        generate_default_partition_table_dtsi(default_flash_mb, offset_kb, "amp")
207
208