zynq devicetree overlay Generator

#!/usr/bin/env python3
"""
Device Tree Overlay Generator - Generates DTS overlays with fragments.
Creates overlay files that can be dynamically loaded to modify existing device trees
"""

import argparse
import sys
import os
import json
import subprocess

class DTSOverlayGenerator:
    def __init__(self, overlay_name="soc.dts", compatible=""):
        """Initialize the DTS overlay generator"""
        self.overlay_name = overlay_name
        self.compatible   = compatible # overlay compatible string
        self.fragments    = []         # Store fragments
        self.labels       = {}         # Store labels for overlay nodes
    
    """
    Create a new fragment node. target_path xor target_label must be
        provided.
    Parameters:
    ===========
    target_path: str
        Full path to the node to be overriden or path where the overlay must
        be applied
    target_label: str
        phandle to the node to be overriden
    label: str
        fragment label (optional)
    properties: dict of str
        dictionary of properties (key/value)
    """
    def add_fragment(self, target_path=None, target_label=None, label=None, properties=None):
        """Add a fragment - can target by path or label"""
        if not target_path and not target_label:
            raise ValueError("Fragment must have either target_path or target_label")
        
        fragment = {
            'target_path'  : target_path,
            'target_label' : target_label,
            'properties'   : {True: properties, False: {}}[properties is not None],
            'child_nodes'  : {}
        }

        #if properties:
        #    fragment["properties"].update(properties)

        if label:
            fragment["label"] = label
            self.labels[f"fragment_{len(self.fragments)}"] = label
            
        self.fragments.append(fragment)

        return len(self.fragments) - 1  # Return fragment index
    
    """
    Add a property to a fragment.
    Parameters:
    ===========
    fragment_idx: int
        fragment index.
    prop_name: str
        property name to be added/updated.
    prop_value: str
        property value.
    """
    def add_fragment_property(self, fragment_idx, prop_name, prop_value):
        assert fragment_idx < len(self.fragments)

        self.fragments[fragment_idx]['properties'][prop_name] = prop_value
    
    def _add_node(self, node_name, label=None, properties=None):
        node_data = {
            'label'      : label,
            'properties' : properties or {},
            'children'   : {}
        }
        
        if label:
            self.labels[node_name] = label
        return node_data

    """
    Add a child node to a fragment
    """
    def add_fragment_node(self, fragment_idx, node_name, label=None, properties=None):
        assert fragment_idx < len(self.fragments)

        node_data = self._add_node(node_name, label, properties)
        self.fragments[fragment_idx]['child_nodes'][node_name] = node_data
        return node_data
    
    """
    Add a child node to an existing node
    """
    def add_child_node(self, parent_node, node_name, label=None, properties=None):
        child_data = self._add_node(node_name, label, properties)
        parent_node['children'][node_name] = child_data
        return child_data
    
    def format_int(self, value, is_hex=False):
        if is_hex:
            if value < 0:
                return f"0x{value & 0xFFFFFFFF:x}"  # Handle negative as unsigned
            else:
                return f"0x{value:x}"
        else:
            if value < 0:
                return f"{value & 0xFFFFFFFF}"  # Handle negative as unsigned
            else:
                return f"{value}"

    """
    Format property value for DTS output
    """
    def format_property_value(self, value, is_hex=False):
        if isinstance(value, str):
            if value.startswith("["):
                return f'{value}'
            elif value.startswith("&"):
                return f'<{value}>'
            else:
                return f'"{value}"'
        elif isinstance(value, int):
            return "<" + self.format_int(value, is_hex) + ">"
        elif isinstance(value, list):
            if all(isinstance(x, int) for x in value):
                if len(value) == 1:
                    return "<" + self.format_int(value[0], is_hex) + ">"
                else:
                    return f'<{" ".join(self.format_int(v, is_hex) for v in value)}>'
            else:
                # String array or mixed
                return f'"{", ".join(map(str, value))}"'
        elif value == "" or value is None:
            # Empty property
            return None
        elif isinstance(value, bytes):
            # Hex data
            hex_str = " ".join(f"{b:02x}" for b in value)
            return f'[{hex_str}]'
        else:
            return f'"{str(value)}"'
    
    """
    Generate complete overlay DTS
    """
    def generate_overlay_dts(self):
        lines = ["/dts-v1/;", "/plugin/;", ""]
        
        # Add overlay metadata
        lines.append("/ {")
        lines.append(f'\tcompatible = {self.compatible};')
        lines.append("")
        
        # Generate fragments
        for i, fragment in enumerate(self.fragments):
            self._generate_fragment(i, fragment, lines)
        
        lines.append("};")
        return "\n".join(lines)
    
    def _generate_fragment(self, fragment_idx, fragment, lines):
        """Generate a single fragment"""
        INDENT         = "\t"
        label          = fragment.get("label", None)
        fragment_label = {
            True  : f"{label}: ",
            False : ""
        }[label is not None]
        
        lines.append(f"{INDENT}{fragment_label}fragment@{fragment_idx} {{")
        
        # Target specification
        INDENT += '\t'
        if fragment['target_label']:
            lines.append(f'{INDENT}target = <&{fragment["target_label"]}>;')
        elif fragment['target_path']:
            lines.append(f'{INDENT}target-path = "{fragment["target_path"]}";')
        
        lines.append("")
        lines.append(f"{INDENT}__overlay__ {{")
        
        # Fragment properties
        INDENT += "\t"
        for prop_name, prop_value in fragment['properties'].items():
            formatted_value = self.format_property_value(prop_value)
            if formatted_value is None:
                lines.append(f"{INDENT}{prop_name};")
            else:
                lines.append(f"{INDENT}{prop_name} = {formatted_value};")
        
        # Child nodes
        if fragment['child_nodes']:
            if fragment['properties']:
                lines.append("")
            
            for node_name, node_data in fragment['child_nodes'].items():
                self._generate_node(node_name, node_data, lines, 3)
        
        INDENT = INDENT[:-1]
        lines.append(f'{INDENT}}};')
        INDENT = INDENT[:-1]
        lines.append(f"{INDENT}}};")
        lines.append("")
    
    def _generate_node(self, node_name, node_data, lines, indent_level):
        """Generate a node recursively"""
        indent = "\t\t\t" + "\t" * (indent_level - 3)
        
        # Node with label
        if node_data['label']:
            lines.append(f"{indent}{node_data['label']}: {node_name} {{")
        else:
            lines.append(f"{indent}{node_name} {{")
        
        # Properties
        for prop_name, prop_value in node_data['properties'].items():
            formatted_value = self.format_property_value(prop_value, prop_name=="reg")
            if formatted_value is None:
                lines.append(f"{indent}    {prop_name};")
            else:
                lines.append(f"{indent}    {prop_name} = {formatted_value};")
        
        # Child nodes
        if node_data['children']:
            if node_data['properties']:
                lines.append("")
            
            for child_name, child_data in node_data['children'].items():
                self._generate_node(child_name, child_data, lines, indent_level + 1)
        
        lines.append(f"{indent}}};")
        
        if indent_level == 3:  # Top level nodes get extra spacing
            lines.append("")

def create_ethernet_node(overlay, key, constants):
    eth_node = None

    # key is config_ps7_gemX_enable
    t        = key.split("_")
    gem_name = t[2]
    gem_base = "_".join(t[0:3])

    gem_dict = {}
    for k, v in constants.items():
        if k.startswith(gem_base):
            nk = k.replace(gem_base + "_", "")
            gem_dict[nk] = v

    if gem_dict["io"] != "emio":
        return None


    print("\n\n")
    print(gem_base)

    # Prepare
    phy_type    = gem_dict["type"]
    ext_phyaddr = gem_dict["ext_phyaddr"]
    ext_phyname = f"phy{ext_phyaddr}"
    ext_phy     = f"phy@{ext_phyaddr}"
    int_phyaddr = gem_dict["int_phyaddr"]
    int_phyname = f"gmiitorgmii"
    int_phy     = f"gmiitorgmii@{int_phyaddr}"

    # GEMx node
    properties = {
        "status"            : "okay",
        "local_mac_address" : "[00 0a 35 00 01 23]", # FIXME: hardcoded
        "xlnx,has_mdio"     : "0x1",
    }

    if phy_type == "rgmii":
        properties["gmii2rgmii-phy-handle"] = "&gmiitorgmii"
    else:
        properties["phy-handle"] = f"&{ext_phyname}"

    eth_node = overlay.add_fragment(target_label=gem_name)
    for fk, fv in properties.items():
        overlay.add_fragment_property(eth_node, fk, fv)

    # mdio sub-node
    properties = {
        "#address-cells" : 1,
        "#size-cells"    : 0,
    }
    mdio_node = overlay.add_fragment_node(eth_node, "mdio", "my_mdio", properties)

    # external phy sub-sub-node
    properties = {
        "device_type" : "ethernet-phy",
        "reg"         : ext_phyaddr,
    }
    ext_phy_node = overlay.add_child_node(mdio_node, ext_phy, ext_phyname, properties)

    if phy_type == "rgmii":
        # gmii_to_rgmii phy sub-sub-node
        properties = {
            "compatible"  : "xlnx,gmii-to-rgmii-1.0",
            #"device_type" : "ethernet-phy",
            #phy-mode = "rgmii-rxid";
            "reg"         : int_phyaddr,
            "phy-handle"  : f"&{ext_phyname}",
        }
        int_phy_node = overlay.add_child_node(mdio_node, int_phy, int_phyname, properties)

    return eth_node

def test_gen_litex(filename="csr.json"):
    d = json.load(open(filename))
    overlay = DTSOverlayGenerator("digilent_arty_z7", '"litex,digilent_arty_z7", "xlnx,zynq-7000"')

    # csr_bases.
    csr_bases = d["csr_bases"]

    # Fragment 0: overrides fpga_full (ie with LiteX Cores).
    fpga_frag = overlay.add_fragment(target_label="fpga_full", properties={
        "status"         : "okay",
        "#address-cells" :  1,
        "#size-cells"    : 0,
    })

    # LiteX Cores
    for name, core_type in d["cores"].items():
        addr      = csr_bases[name]
        node_name = f"{name}@{addr:08x}"

        if core_type == "SPIMaster":
            properties = {
                "compatible"            : "litex,litespi",
                "reg"                   : [addr, 0x100],
                "litespi,max-bpw"       : 8,
                "litespi,sck-frequency" : 1000000,
                "#address-cells"        : 1,
                "#size-cells"           : 0,
            }
        elif core_type == "I2CMaster":
            properties = {
                "compatible"     : "litex,i2c",
                "reg"            : [addr, 0x5],
                "#address-cells" : 1,
                "#size-cells"    : 0,
            }

        elif core_type == "SoCController":
            node_name = f"soc_controler@{addr:08x}"
            properties = {
                "compatible"     : "litex,soc-controller",
                "reg"            : [addr, 0x5],
            }
        elif core_type == "UART":
            properties = {
                "compatible"    : "litex,liteuart",
                "reg"           : [addr, 0x100],
                #{uart_interrupt : FIXME
            }
        else:
            continue
        properties["status"] = "okay"
        node = overlay.add_fragment_node(fpga_frag, node_name, name, properties)

    # Fragment 1->n: overrides peripherals nodes to enable PS controlers.
    constants = d["constants"]
    for k, v in constants.items():
        # search for key config_ps7_PERIPHx_enable
        if not (k.startswith("config_ps7") and k.endswith("enable")):
            print("wrong key: ",k)
            continue
        periph_name = k.split("_")[2]
        args = {"status" : "okay"}
        if "i2c" in k:
            args["clock-frequency"] = 400000
        if "spi" in k:
            pass
        if "uart" in k:
            pass
        if "gem" in k:
            if "mdio" in k:
                continue
            print("GEM", k)
            #print(k)
            create_ethernet_node(overlay, k, constants)
            continue
        ps_frag = overlay.add_fragment(target_label=periph_name)
        for fk, fv in args.items():
            overlay.add_fragment_property(ps_frag, fk, fv)
    
    # TBD: Another fragment to overrides pincfg with pinout when
    # controler is connected to MMIO pins.

    return overlay

# Merging two dtbo is not really possible with fdtoverlay
# but it's possible to treat both dts
# two merge it into one dts and convert dts to dtb
def merge_dts(dts_names=[], dts_out=""):
    # We assume the first overlay is the reference
    dts_ref_in = []
    with open(dts_names[0], "r") as fd:
        dts_ref_in = fd.read().splitlines()
    # Removes the last line (close brace)
    dts_ref_in = dts_ref_in[:-1]
    # Extracts list of fragments lines.
    fragments_list = [ l for l in dts_ref_in if "fragment" in l]
    # only keep fragment's ID.
    ff             = [int(l.split("@")[1].split()[0]) for l in fragments_list]
    # Sort: last is the highest value
    ff.sort()

    # Save highest fragment ID + 1
    last_fragment_id = ff[-1] + 1

    # For remainings dts removes headers and only keeps fragments sections
    # ie we waits until first fragment
    dts_cnt = []
    for dts in dts_names[1:]:
        with open(dts, "r") as fd:
            c         = []
            in_header = True
            for v in fd.readlines():
                v = v.replace("\n", "")
                if in_header:
                    if "fragment" not in v:
                        continue
                    else:
                        in_header = False
                c.append(v)
            c = c[:-1]
            dts_cnt.append(c)


    output            = []
    # Directly integrates first dts (reference with headers).
    output           += dts_ref_in

    # Merge all dts in one file
    for dts in dts_cnt:
        for l in dts:
            if "fragment" in l:
                ll = l.split("@")
                ll2 = ll[1].split(" ")
                l = ll[0] + str(last_fragment_id) + " " + ll2[1]
                last_fragment_id += 1
            output.append(l)
    
    output.append("};")
    with open(dts_out, "w") as fd:
        fd.write("\n".join(output))

def gen_dtbo(dts_name, dtbo_name, symbols=True):
    subprocess.check_call(
        "dtc {} -O dtb -o {} {}".format("-@" if symbols else "", dtbo_name, dts_name), shell=True)

# FIXME: instead of reading the first dts (ie the one produces here
# it's more better /easy to keep information and buffer
def main():
    parser = argparse.ArgumentParser(description="LiteX's CSR JSON to Linux DTS overlay generator")
    parser.add_argument("csr_json", help="CSR JSON file")
    parser.add_argument("--dts",    help="optional devicetree overlay to merge in dtbo.")
    parser.add_argument("--out",    help="devicetree overlay (dtbo) name.")
    args = parser.parse_args()
    
    print("Generating Device Tree Overlays with fragments...")
    dtsfile = args.out.replace("dtbo", "dts")
    # Create overlay
    overlay = test_gen_litex(args.csr_json)
    
    # Generate and save overlay DTS
    overlay_content = overlay.generate_overlay_dts()
    
    # Write overlay
    with open(dtsfile, 'w') as f:
        f.write(overlay_content)

    # If another dts is provided by the user
    # merge it into dts writen in the previous step.
    if args.dts is not None:
        merge_dts([dtsfile, args.dts], dtsfile)

    # Gen final devicetree overlay from the intermediate dts
    gen_dtbo(dtsfile, args.out)

    print(f"Hardware overlay saved to: {dtsfile}")
    print(f"Size: {len(overlay_content)} characters")
    print(f"Fragments: {len(overlay.fragments)}")
    
    # Show all labels
    if overlay.labels:
        print(f"\nGenerated {len(overlay.labels)} labeled nodes:")
        for node, label in sorted(overlay.labels.items()):
            print(f"  {label}: {node}")

    return 0

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