path: root/build-aux/kicad-rp2040-check

#!/usr/bin/env python3
# build-aux/kicad-rp2040-check - Additional ERC-type checks for the RP2040
#
# Copyright (C) 2025  Luke T. Shumaker <lukeshu@lukeshu.com>
# SPDX-License-Identifier: AGPL-3.0-or-later

import argparse
import os.path
import re
import subprocess
import sys
import typing

common = os.path.abspath(os.path.dirname(__file__))
if common not in sys.path:
    sys.path.insert(0, common)

import kicad_sexpr  # pylint: disable=wrong-import-position

pin2name: dict[int, str] = {
    1: "IOVDD",
    2: "GPIO0",
    3: "GPIO1",
    4: "GPIO2",
    5: "GPIO3",
    6: "GPIO4",
    7: "GPIO5",
    8: "GPIO6",
    9: "GPIO7",
    10: "IOVDD",
    11: "GPIO8",
    12: "GPIO9",
    13: "GPIO10",
    14: "GPIO11",
    15: "GPIO12",
    16: "GPIO13",
    17: "GPIO14",
    18: "GPIO15",
    19: "TESTEN",
    20: "XIN",
    21: "XOUT",
    22: "IOVDD",
    23: "DVDD",
    24: "SWCLK",
    26: "SWDIO",
    27: "GPIO16",
    28: "GPIO17",
    29: "GPIO18",
    30: "GPIO19",
    31: "GPIO20",
    32: "GPIO21",
    33: "IOVDD",
    34: "GPIO22",
    35: "GPIO23",
    36: "GPIO24",
    37: "GPIO25",
    38: "GPIO26/ADC0",
    39: "GPIO27/ADC1",
    40: "GPIO28/ADC2",
    41: "GPIO29/ADC3",
    42: "IOVDD",
    43: "ADC_AVDD",
    44: "VREG_VIN",
    45: "VREG_VOUT",
    46: "USB_DM",
    47: "USB_DP",
    48: "USB_VDD",
    49: "IOVDD",
    50: "DVDD",
    51: "QSPI_SD3",
    52: "QSPI_SCLK",
    53: "QSPI_SD0",
    54: "QSPI_SD2",
    55: "QSPI_SD1",
    56: "QSPI_SS_N",
}

pin2gpio: dict[int, int]
gpio2pin: dict[int, int]
RED: str
RESET: str


def init() -> None:
    global pin2gpio
    global gpio2pin
    pin2gpio = {}
    gpio2pin = {}
    for pin, name in pin2name.items():
        if name.startswith("GPIO"):
            gpio = int(name[4:].split("/")[0], 10)
            pin2gpio[pin] = gpio
            gpio2pin[gpio] = pin
    global RED
    RED = subprocess.run(
        ["tput", "setaf", "1"], capture_output=True, encoding="utf-8", check=True
    ).stdout
    global RESET
    RESET = subprocess.run(
        ["tput", "sgr0"], capture_output=True, encoding="utf-8", check=True
    ).stdout


class GPIOFuncs(typing.NamedTuple):
    F1: str
    F2: str
    F3: str
    F4: str
    F5: str
    F6: str
    F7: str
    F8: str | None
    F9: str

    @property
    def SPI(self) -> str:
        return self.F1

    @property
    def UART(self) -> str:
        return self.F2

    @property
    def I2C(self) -> str:
        return self.F3

    @property
    def PWM(self) -> str:
        return self.F4

    @property
    def SIO(self) -> str:
        return self.F5

    @property
    def PIO0(self) -> str:
        return self.F6

    @property
    def PIO1(self) -> str:
        return self.F7

    @property
    def CLOCK(self) -> str | None:
        return self.F8

    @property
    def USB(self) -> str:
        return self.F9


gpio_functions: dict[int, GPIOFuncs] = {
    # fmt: off
    0: GPIOFuncs("SPI0 RX", "UART0 TX", "I2C0 SDA", "PWM0 A", "SIO", "PIO0", "PIO1", None, "USB OVCUR DET"),
    1: GPIOFuncs("SPI0 CSn", "UART0 RX", "I2C0 SCL", "PWM0 B", "SIO", "PIO0", "PIO1", None, "USB VBUS DET"),
    2: GPIOFuncs("SPI0 SCK", "UART0 CTS", "I2C1 SDA", "PWM1 A", "SIO", "PIO0", "PIO1", None, "USB VBUS EN"),
    3: GPIOFuncs("SPI0 TX", "UART0 RTS", "I2C1 SCL", "PWM1 B", "SIO", "PIO0", "PIO1", None, "USB OVCUR DET"),
    4: GPIOFuncs("SPI0 RX", "UART1 TX", "I2C0 SDA", "PWM2 A", "SIO", "PIO0", "PIO1", None, "USB VBUS DET"),
    5: GPIOFuncs("SPI0 CSn", "UART1 RX", "I2C0 SCL", "PWM2 B", "SIO", "PIO0", "PIO1", None, "USB VBUS EN"),
    6: GPIOFuncs("SPI0 SCK", "UART1 CTS", "I2C1 SDA", "PWM3 A", "SIO", "PIO0", "PIO1", None, "USB OVCUR DET"),
    7: GPIOFuncs("SPI0 TX", "UART1 RTS", "I2C1 SCL", "PWM3 B", "SIO", "PIO0", "PIO1", None, "USB VBUS DET"),
    8: GPIOFuncs("SPI1 RX", "UART1 TX", "I2C0 SDA", "PWM4 A", "SIO", "PIO0", "PIO1", None, "USB VBUS EN"),
    9: GPIOFuncs("SPI1 CSn", "UART1 RX", "I2C0 SCL", "PWM4 B", "SIO", "PIO0", "PIO1", None, "USB OVCUR DET"),
    10: GPIOFuncs("SPI1 SCK", "UART1 CTS", "I2C1 SDA", "PWM5 A", "SIO", "PIO0", "PIO1", None, "USB VBUS DET"),
    11: GPIOFuncs("SPI1 TX", "UART1 RTS", "I2C1 SCL", "PWM5 B", "SIO", "PIO0", "PIO1", None,  "USB VBUS EN"),
    12: GPIOFuncs("SPI1 RX", "UART0 TX", "I2C0 SDA", "PWM6 A", "SIO", "PIO0", "PIO1", None, "USB OVCUR DET"),
    13: GPIOFuncs("SPI1 CSn", "UART0 RX", "I2C0 SCL", "PWM6 B", "SIO", "PIO0", "PIO1", None, "USB VBUS DET"),
    14: GPIOFuncs("SPI1 SCK", "UART0 CTS", "I2C1 SDA", "PWM7 A", "SIO", "PIO0", "PIO1", None, "USB VBUS EN"),
    15: GPIOFuncs("SPI1 TX", "UART0 RTS", "I2C1 SCL", "PWM7 B", "SIO", "PIO0", "PIO1", None, "USB OVCUR DET"),
    16: GPIOFuncs("SPI0 RX", "UART0 TX", "I2C0 SDA", "PWM0 A", "SIO", "PIO0", "PIO1", None, "USB VBUS DET"),
    17: GPIOFuncs("SPI0 CSn", "UART0 RX", "I2C0 SCL", "PWM0 B", "SIO", "PIO0", "PIO1", None, "USB VBUS EN"),
    18: GPIOFuncs("SPI0 SCK", "UART0 CTS", "I2C1 SDA", "PWM1 A", "SIO", "PIO0", "PIO1", None, "USB OVCUR DET"),
    19: GPIOFuncs("SPI0 TX", "UART0 RTS", "I2C1 SCL", "PWM1 B", "SIO", "PIO0", "PIO1", None, "USB VBUS DET"),
    20: GPIOFuncs("SPI0 RX", "UART1 TX", "I2C0 SDA", "PWM2 A", "SIO", "PIO0", "PIO1", "CLOCK GPIN0", "USB VBUS EN"),
    21: GPIOFuncs("SPI0 CSn", "UART1 RX", "I2C0 SCL", "PWM2 B", "SIO", "PIO0", "PIO1", "CLOCK GPOUT0", "USB OVCUR DET"),
    22: GPIOFuncs("SPI0 SCK", "UART1 CTS", "I2C1 SDA", "PWM3 A", "SIO", "PIO0", "PIO1", "CLOCK GPIN1", "USB VBUS DET"),
    23: GPIOFuncs("SPI0 TX", "UART1 RTS", "I2C1 SCL", "PWM3 B", "SIO", "PIO0", "PIO1", "CLOCK GPOUT1", "USB VBUS EN"),
    24: GPIOFuncs("SPI1 RX", "UART1 TX", "I2C0 SDA", "PWM4 A", "SIO", "PIO0", "PIO1", "CLOCK GPOUT2", "USB OVCUR DET"),
    25: GPIOFuncs("SPI1 CSn", "UART1 RX", "I2C0 SCL", "PWM4 B", "SIO", "PIO0", "PIO1", "CLOCK GPOUT3", "USB VBUS DET"),
    26: GPIOFuncs("SPI1 SCK", "UART1 CTS", "I2C1 SDA", "PWM5 A", "SIO", "PIO0", "PIO1", None, "USB VBUS EN"),
    27: GPIOFuncs("SPI1 TX", "UART1 RTS", "I2C1 SCL", "PWM5 B", "SIO", "PIO0", "PIO1", None, "USB OVCUR DET"),
    28: GPIOFuncs("SPI1 RX", "UART0 TX", "I2C0 SDA", "PWM6 A", "SIO", "PIO0", "PIO1", None, "USB VBUS DET"),
    29: GPIOFuncs("SPI1 CSn", "UART0 RX", "I2C0 SCL", "PWM6 B", "SIO", "PIO0", "PIO1", None, "USB VBUS EN"),
    # fmt: on
}


def gather_rp2040_nets(
    netlist: kicad_sexpr.Expr,
) -> dict[str, dict[int, set[str]]]:  # ref=>pin=>netnames
    assert isinstance(netlist, list) and netlist[:2] == [
        kicad_sexpr.Symbol("export"),
        [kicad_sexpr.Symbol("version"), "E"],
    ]
    components: list[kicad_sexpr.Expr] = next(
        child[1:]
        for child in netlist[2:]
        if isinstance(child, list) and child[0] == kicad_sexpr.Symbol("components")
    )
    nets: list[kicad_sexpr.Expr] = next(
        child[1:]
        for child in netlist[2:]
        if isinstance(child, list) and child[0] == kicad_sexpr.Symbol("nets")
    )

    rp2040_nets: dict[str, dict[int, set[str]]] = {}  # ref=>pin=>netnames

    for comp in components:
        assert isinstance(comp, list) and comp[0] == kicad_sexpr.Symbol("comp")
        ref = ""
        value = ""
        for comp_kv in comp[1:]:
            assert isinstance(comp_kv, list) and len(comp_kv) >= 2
            [comp_k, *comp_vs] = comp_kv
            match comp_k:
                case kicad_sexpr.Symbol("ref"):
                    assert len(comp_vs) == 1 and isinstance(comp_vs[0], str)
                    ref = comp_vs[0]
                case kicad_sexpr.Symbol("value"):
                    assert len(comp_vs) == 1 and isinstance(comp_vs[0], str)
                    value = comp_vs[0]
        if value == "RP2040":
            rp2040_nets[ref] = {}

    for net in nets:
        assert isinstance(net, list) and net[0] == kicad_sexpr.Symbol("net")
        net_name = ""
        for net_kv in net[1:]:
            assert isinstance(net_kv, list) and len(net_kv) >= 2
            [net_k, *net_vs] = net_kv
            match net_k:
                case kicad_sexpr.Symbol("name"):
                    assert len(net_vs) == 1 and isinstance(net_vs[0], str)
                    net_name = net_vs[0]
                case kicad_sexpr.Symbol("node"):
                    node_ref = ""
                    node_pin = -1
                    for node_kv in net_vs:
                        assert isinstance(node_kv, list) and len(node_kv) == 2
                        [node_k, node_v] = node_kv
                        match node_k:
                            case kicad_sexpr.Symbol("ref"):
                                assert isinstance(node_v, str)
                                node_ref = node_v
                            case kicad_sexpr.Symbol("pin"):
                                assert isinstance(node_v, str)
                                try:
                                    node_pin = int(node_v, 10)
                                except ValueError:
                                    pass
                    if node_ref and (node_ref in rp2040_nets) and node_pin > 0:
                        if node_pin not in rp2040_nets[node_ref]:
                            rp2040_nets[node_ref][node_pin] = set()
                        rp2040_nets[node_ref][node_pin].add(net_name)

    return rp2040_nets


def check_rp2040_nets(
    sch_filename: str, rp2040_nets: dict[str, dict[int, set[str]]]
) -> bool:
    erred = False

    for ref, chip in rp2040_nets.items():
        if check_rp2040(sch_filename, ref, chip):
            erred = True
    return erred


def check_rp2040(sch_filename: str, ref: str, chip: dict[int, set[str]]) -> bool:
    net2pin: dict[str, set[int]] = {}
    for pin in chip:
        for netname in chip[pin]:
            if netname not in net2pin:
                net2pin[netname] = set()
            net2pin[netname].add(pin)

    erred = False

    def err(msg: str) -> None:
        print(f"{RED}{sch_filename}{RESET}: {ref}: {msg}", file=sys.stderr)
        nonlocal erred
        erred = True

    func2pin: dict[str, set[int]] = {}

    def check_instance(
        pin: int, netname: str, attr: str, base_net: str, all_base_nets: list[str]
    ) -> None:
        assert base_net in netname
        func = getattr(gpio_functions[gpio], attr)
        if func not in func2pin:
            func2pin[func] = set()
        func2pin[func].add(pin)
        inst = func.split()[0]
        for rep in all_base_nets:
            if rep == base_net:
                continue
            sib_netname = netname.replace(base_net, rep)
            for sib_pin in net2pin.get(sib_netname, set()):
                sib_gpio = pin2gpio[sib_pin]
                sib_inst = getattr(gpio_functions[sib_gpio], attr).split()[0]
                if sib_inst != inst:
                    err(
                        f"pin={pin} (gpio={gpio2pin[pin]}): disagrees with pin={sib_pin} (gpio={sib_gpio}) about {attr} instance"
                    )

    for pin in chip:
        if pin not in pin2gpio:
            continue
        gpio = pin2gpio[pin]
        for netname in sorted(chip[pin]):
            # SPI ##############################################################
            if "SPI_CLK" in netname:
                if "SCK" not in gpio_functions[gpio].SPI:
                    err(
                        f"pin={pin}: net={netname!r} but GPIO{gpio}.F1={gpio_functions[gpio].SPI!r}"
                    )
                    continue
                check_instance(
                    pin, netname, "SPI", "CLK", ["CLK", "MISO", "MOSI", "CS"]
                )
            elif "SPI_MISO" in netname:
                if "RX" not in gpio_functions[gpio].SPI:
                    err(
                        f"pin={pin}: net={netname!r} but GPIO{gpio}.F1={gpio_functions[gpio].SPI!r}"
                    )
                    continue
                check_instance(
                    pin, netname, "SPI", "MISO", ["CLK", "MISO", "MOSI", "CS"]
                )
            elif "SPI_MOSI" in netname:
                if "TX" not in gpio_functions[gpio].SPI:
                    err(
                        f"pin={pin}: net={netname!r} but GPIO{gpio}.F1={gpio_functions[gpio].SPI!r}"
                    )
                    continue
                check_instance(
                    pin, netname, "SPI", "MOSI", ["CLK", "MISO", "MOSI", "CS"]
                )
            elif "SPI_CS" in netname:
                if "CS" not in gpio_functions[gpio].SPI:
                    err(
                        f"pin={pin}: net={netname!r} but GPIO{gpio}.F1={gpio_functions[gpio].SPI!r}"
                    )
                    continue
                check_instance(pin, netname, "SPI", "CS", ["CLK", "MISO", "MOSI", "CS"])
            # SD Card SPI ######################################################
            elif re.search(r"SD_.*CLK", netname):
                if "SCK" not in gpio_functions[gpio].SPI:
                    err(
                        f"pin={pin}: net={netname!r} (SPI SCK) but GPIO{gpio}.F1={gpio_functions[gpio].SPI!r}"
                    )
                    continue
                check_instance(
                    pin, netname, "SPI", "CLK", ["CLK", "CMD", "DAT0", "DAT3"]
                )
            elif re.search(r"SD_.*CMD", netname):
                if "TX" not in gpio_functions[gpio].SPI:
                    err(
                        f"pin={pin}: net={netname!r} (SPI TX) but GPIO{gpio}.F1={gpio_functions[gpio].SPI!r}"
                    )
                    continue
                check_instance(
                    pin, netname, "SPI", "CMD", ["CLK", "CMD", "DAT0", "DAT3"]
                )
            elif re.search(r"SD_.*DAT0", netname):
                if "RX" not in gpio_functions[gpio].SPI:
                    err(
                        f"pin={pin}: net={netname!r} (SPI RX) but GPIO{gpio}.F1={gpio_functions[gpio].SPI!r}"
                    )
                    continue
                check_instance(
                    pin, netname, "SPI", "DAT0", ["CLK", "CMD", "DAT0", "DAT3"]
                )
            elif re.search(r"SD_.*DAT3", netname):
                if "CS" not in gpio_functions[gpio].SPI:
                    err(
                        f"pin={pin}: net={netname!r} (SPI CS) but GPIO{gpio}.F1={gpio_functions[gpio].SPI!r}"
                    )
                    continue
                check_instance(
                    pin, netname, "SPI", "DAT3", ["CLK", "CMD", "DAT0", "DAT3"]
                )
            # UART #############################################################
            elif "UART_TX" in netname:
                if "TX" not in gpio_functions[gpio].UART:
                    err(
                        f"pin={pin}: net={netname!r} but GPIO{gpio}.F2={gpio_functions[gpio].UART!r}"
                    )
                    continue
                check_instance(pin, netname, "UART", "TX", ["RX", "TX"])
            elif "UART_RX" in netname:
                if "RX" not in gpio_functions[gpio].UART:
                    err(
                        f"pin={pin}: net={netname!r} but GPIO{gpio}.F2={gpio_functions[gpio].UART!r}"
                    )
                    continue
                check_instance(pin, netname, "UART", "RX", ["RX", "TX"])
            # I2C ##############################################################
            elif "SDA" in netname:
                if "SDA" not in gpio_functions[gpio].I2C:
                    err(
                        f"pin={pin}: net={netname!r} but GPIO{gpio}.F3={gpio_functions[gpio].I2C!r}"
                    )
                    continue
                check_instance(pin, netname, "I2C", "SDA", ["SDA", "SCL"])
            elif "SCL" in netname:
                if "SCL" not in gpio_functions[gpio].I2C:
                    err(
                        f"pin={pin}: net={netname!r} but GPIO{gpio}.F3={gpio_functions[gpio].I2C!r}"
                    )
                    continue
                check_instance(pin, netname, "I2C", "SCL", ["SDA", "SCL"])
            elif "CLK" in netname and not any(
                s in netname for s in ["SPI_CLK", "SWCLK", "SD_CLK"]
            ):
                if gpio_functions[gpio].CLOCK is None:
                    err(f"pin={pin}: net={netname!r} but no GPIO{gpio}.F8")

    for func, pins in func2pin.items():
        if len(pins) > 1:
            err(f"{func}: assigned to multiple pins: {sorted(pins)}")

    return erred


def main() -> int:
    parser = argparse.ArgumentParser(
        prog=sys.argv[0],
        description="Additional ERC-type checks for the RP2040",
    )
    parser.add_argument(
        "sch_filename",
        metavar="SCH_FILENAME.kicad_sch",
    )
    args = parser.parse_args(sys.argv[1:])
    if not args.sch_filename.endswith(".kicad_sch"):
        parser.error("argument SCH_FILENAME: must be a .kicad_sch file")

    netlist = kicad_sexpr.unmarshal(
        subprocess.run(
            [
                "kicad-cli",
                "sch",
                "export",
                "netlist",
                "--output=/dev/stdout",
                args.sch_filename,
            ],
            capture_output=True,
            encoding="utf-8",
            check=True,
        ).stdout
    )

    rp2040_nets: dict[str, dict[int, set[str]]] = gather_rp2040_nets(
        netlist
    )  # ref=>pin=>netnames

    if check_rp2040_nets(args.sch_filename, rp2040_nets):
        return 1
    return 0


if __name__ == "__main__":
    init()
    sys.exit(main())