1#!/usr/bin/env python
2#
3# SPDX-FileCopyrightText: 2022 Espressif Systems (Shanghai) CO LTD
4# SPDX-License-Identifier: Apache-2.0
5#
6import argparse
7import json
8import os
9import re
10from typing import Dict
11
12IDF_PATH = os.environ['IDF_PATH']
13MAX_SIZE_DIFF = 50
14
15
16def mem_test(size_json: dict, esptool_output: list) -> None:
17    seg_len = {}  # type: Dict[str, int]
18    for i in esptool_output:
19        tmp = i.split(' ')
20        if tmp[0] == 'Segment':
21            # tmp look like ['Segment', '2:', 'len', '0x02780', 'load', '0x3fc90610', 'file_offs', '0x00007ab0', '[BYTE_ACCESSIBLE,MEM_INTERNAL,DRAM]']
22            # tmp[3] contains the size of the segment and tmp[8] contains the name of the memory segment
23            esptool_mem = {'mem_type':tmp[8], 'size':tmp[3]}
24            seg = re.sub(r'MEM_INTERNAL|,|BYTE_ACCESSIBLE|\n|\[|\]', '', esptool_mem['mem_type'])
25            # If there are two IRAMs in esptool output it will compute these two IRAM lengths in a seg_len['IRAM']
26            seg_len[seg] = int(esptool_mem['size'], 16) if seg not in seg_len else seg_len[seg] + int(esptool_mem['size'], 16)
27    # including flash_other to DROM because flash_other contain .flash.appdesc that includes in DROM that produced by esptool
28    size_from_map = [('IROM', size_json['flash_code']), ('IRAM', size_json['iram_text'] + size_json['iram_vectors'] + size_json['diram_text']
29                     + size_json['diram_vectors']), ('DROM', size_json['flash_rodata'] + size_json['flash_other']), ('DRAM', size_json
30                     ['dram_data'] + size_json['diram_data'])]
31    for mem_type, size in size_from_map:
32        if abs(size - seg_len[mem_type]) > MAX_SIZE_DIFF:
33            raise RuntimeError(mem_type + " segment in idf_size isn't correct regarding esptool")
34    print('Test complete without errors')
35
36
37def main() -> None:
38    parser = argparse.ArgumentParser(description='mem_test.py - a tool to test accuracy of the sizes of the memory segments regarding idf.py size by esptool')
39
40    parser.add_argument(
41        'size_json', help='JSON file with the output of the idf.py size',
42        type=argparse.FileType('r'))
43    parser.add_argument(
44        'esptool_output', help='File with the output of the esptool',
45        type=argparse.FileType('r'))
46
47    args = parser.parse_args()
48    mem_test(json.loads(args.size_json.read()), args.esptool_output.read().split('\n'))
49
50
51if __name__ == '__main__':
52    main()
53