1# Copyright (c) 2024 Chen Xingyu <hi@xingrz.me>
2# SPDX-License-Identifier: Apache-2.0
3
4'''Runner for probe-rs.'''
5
6from runners.core import RunnerCaps, ZephyrBinaryRunner
7
8DEFAULT_PROBE_RS_GDB_HOST = 'localhost'
9DEFAULT_PROBE_RS_GDB_PORT = 1337
10
11class ProbeRsBinaryRunner(ZephyrBinaryRunner):
12    '''Runner front-end for probe-rs.'''
13
14    def __init__(self, cfg, chip,
15                 probe_rs='probe-rs',
16                 gdb_host=DEFAULT_PROBE_RS_GDB_HOST,
17                 gdb_port=DEFAULT_PROBE_RS_GDB_PORT,
18                 dev_id=None,
19                 erase=False,
20                 tool_opt=None):
21        super().__init__(cfg)
22
23        self.probe_rs = probe_rs
24        self.erase = erase
25
26        self.args = ['--chip', chip]
27
28        if dev_id is not None:
29            self.args += ['--probe', dev_id]
30
31        if tool_opt is not None:
32            self.args += tool_opt
33
34        self.elf_name = cfg.elf_file
35
36        self.gdb_cmd = cfg.gdb
37        self.gdb_host = gdb_host
38        self.gdb_port = gdb_port
39
40    @classmethod
41    def name(cls):
42        return 'probe-rs'
43
44    @classmethod
45    def capabilities(cls):
46        return RunnerCaps(commands={'flash', 'debug', 'debugserver'},
47                          dev_id=True,
48                          erase=True,
49                          tool_opt=True)
50
51    @classmethod
52    def do_add_parser(cls, parser):
53        parser.add_argument('--chip', required=True,
54                            help='chip name')
55        parser.add_argument('--probe-rs', default='probe-rs',
56                            help='path to probe-rs tool, default is probe-rs')
57        parser.add_argument('--gdb-host', default=DEFAULT_PROBE_RS_GDB_HOST,
58                            help=f'probe-rs gdb host, defaults to {DEFAULT_PROBE_RS_GDB_HOST}')
59        parser.add_argument('--gdb-port', default=DEFAULT_PROBE_RS_GDB_PORT,
60                            help=f'probe-rs gdb port, defaults to {DEFAULT_PROBE_RS_GDB_PORT}')
61
62
63    @classmethod
64    def dev_id_help(cls) -> str:
65        return '''select a specific probe, in the form `VID:PID:<Serial>`'''
66
67    @classmethod
68    def tool_opt_help(cls) -> str:
69        return '''additional options for probe-rs,
70                  e.g. --chip-description-path=/path/to/chip.yml'''
71
72    @classmethod
73    def do_create(cls, cfg, args):
74        return ProbeRsBinaryRunner(cfg, args.chip,
75                                   probe_rs=args.probe_rs,
76                                   dev_id=args.dev_id,
77                                   erase=args.erase,
78                                   tool_opt=args.tool_opt)
79
80    def do_run(self, command, **kwargs):
81        self.require(self.probe_rs)
82        if command == 'flash':
83            self.do_flash(**kwargs)
84        elif command in ('debug', 'debugserver'):
85            self.do_debug_debugserver(command, **kwargs)
86
87    def do_flash(self, **kwargs):
88        download_args = []
89        if self.erase:
90            download_args += ['--chip-erase']
91        download_args += [self.elf_name]
92
93        self.check_call([self.probe_rs, 'download']
94                        + self.args + download_args)
95
96        self.check_call([self.probe_rs, 'reset']
97                        + self.args)
98
99    def do_debug_debugserver(self, command, **kwargs):
100        debug_args = ['--gdb-connection-string', f"{self.gdb_host}:{self.gdb_port}"]
101        if command == 'debug':
102            debug_args += [self.elf_name]
103            debug_args += ['--gdb', self.gdb_cmd]
104
105        self.check_call([self.probe_rs, 'gdb']
106                        + self.args + debug_args)
107