Can You Reverse Engineer An ASIC?

"""SKY130 GDSII -> structural Verilog extractor for the Jane Street ASIC puzzle.

GDSII represents only ASIC geometry: polygons on metal layers, placements of
standard-cell masters, and a reuse hierarchy.  Recovering a netlist means:

  1. flatten every routing polygon into top-cell coordinates and extract
     connectivity (touching polygons on a layer, plus metal-cut-metal)
  2. transform each placed cell's master-local pin labels by its placement
     transform and probe the flat graph for that pin's net
  3. index the pin rows into power rails, drivers and loads
  4. find boundary nets from top-level /16 pin shapes and name the interface
  5. validate, then emit structural Verilog

Flattening in pass 1 is the point of the whole design.  A hierarchical probe
returns a name local to the master (``clkbuf_4:$2``) that every placement of
that master shares, so it cannot identify a chip net.

Specialised to SKY130: relies on the library's layer map, cell naming and
pin-label conventions.

    python3 gds_to_verilog.py puzzle.gds --top puzzle --outdir recovered

pin_nets.csv, top_ports.csv and extraction.txt are always written,
and recovered.v only when validation passes.
"""

from __future__ import annotations

import argparse
import csv
import os
import re
import sys
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Sequence, Set, Tuple

import klayout.db as db


# --- SKY130 layer map --------------------------------------------------------
# Conductors are datatype 20 and cuts datatype 44, both on layer 67+i.  The same
# layer numbers carry pin labels on /5 (inside cell masters) and pin-purpose
# shapes on /16 (at the top-level boundary).  Transistor layers are not needed:
# the master name already says which gate was placed.

CONDUCTORS = ["li1", "met1", "met2", "met3", "met4", "met5"]
CUTS = ["mcon", "via1", "via2", "via3", "via4"]

ROUTING_LAYERS: Dict[str, Tuple[int, int]] = {
    **{name: (67 + i, 20) for i, name in enumerate(CONDUCTORS)},
    **{name: (67 + i, 44) for i, name in enumerate(CUTS)},
}
# (lower conductor, cut, upper conductor), up the stack from li1 to met5.
STACK = [(CONDUCTORS[i], CUTS[i], CONDUCTORS[i + 1]) for i in range(len(CUTS))]
PIN_TEXT_LAYERS = [(67 + i, 5, name) for i, name in enumerate(CONDUCTORS)]
TOP_PIN_LAYERS = [(67 + i, 16, name) for i, name in enumerate(CONDUCTORS)]

PREFIX = "sky130_fd_sc_hd__"
CONB_CELL = PREFIX + "conb_1"           # constant generator; HI/LO are shared sources
OUTPUT_BANK_CELL = PREFIX + "and3_2"    # drives the eight out[] pins
PHYSICAL_ONLY = {PREFIX + n for n in ("tapvpwrvgnd_1", "decap_3", "diode_2")}

POWER_PINS = {"VPWR", "VGND", "VPB", "VNB"}
OUTPUT_PINS = {"X", "Y", "Q", "Q_N", "HI", "LO"}
EXPECTED_INPUTS = {"clk", "rst_n", "enable", "input"}
EXPECTED_OUTPUTS = {"success"} | {f"out[{i}]" for i in range(8)}

BIT_SELECT = re.compile(r"([A-Za-z_][A-Za-z0-9_$]*)\[(\d+)\]")
PLAIN_ID = re.compile(r"[A-Za-z_][A-Za-z0-9_$]*")
VERILOG_KEYWORDS = set("""
always and assign automatic begin buf bufif0 bufif1 case casex casez cell cmos
config deassign default defparam design disable edge else end endcase endconfig
endfunction endgenerate endmodule endprimitive endspecify endtable endtask event
for force forever fork function generate genvar highz0 highz1 if ifnone incdir
include initial inout input instance integer join large liblist library
localparam macromodule medium module nand negedge nmos nor noshowcancelled not
notif0 notif1 or output parameter pmos posedge primitive pull0 pull1 pulldown
pullup pulsestyle_onevent pulsestyle_ondetect rcmos real realtime reg release
repeat rnmos rpmos rtran rtranif0 rtranif1 scalared showcancelled signed small
specify specparam strong0 strong1 supply0 supply1 table task time tran tranif0
tranif1 tri tri0 tri1 triand trior trireg unsigned use vectored wait wand weak0
weak1 while wire wor xnor xor
""".split())


# --- KLayout helpers ---------------------------------------------------------

def attr_or_call(obj, name):
    """Read a KLayout property exposed as either a value or a method."""
    value = getattr(obj, name)
    return value() if callable(value) else value


def cell_index_of(obj) -> int:
    """Master ID, for a cell or for an instance's master."""
    return int(attr_or_call(obj, "cell_index"))


def net_name(net) -> Optional[str]:
    """Name one connected component of the flat graph.

    The spelling is irrelevant (the puzzle yields names like ``$27``); only
    uniqueness matters.
    """
    if net is None:
        return None
    for attr in ("expanded_name", "name"):
        try:
            if value := attr_or_call(net, attr):
                return str(value)
        except Exception:
            pass
    return str(net)


def iter_transforms(inst):
    """Yield each placement in an instance record (one transform per instance)."""
    try:
        yield from attr_or_call(inst, "cell_inst").each_cplx_trans()
        return
    except Exception:
        pass
    yield attr_or_call(inst, "cplx_trans")


def each_logic_instance(ly: db.Layout, top: db.Cell):
    """Yield (instance, master) for placed cells that are real gates.

    Tap, decap and diode cells are fabrication support, not Boolean logic.
    """
    for inst in top.each_inst():
        child = ly.cell(cell_index_of(inst))
        if child and child.name.startswith(PREFIX) and child.name not in PHYSICAL_ONLY:
            yield inst, child


def sample_points(box: db.Box) -> List[db.Point]:
    """Centre plus quarter-offsets, so an ambiguous shape can be detected."""
    cx, cy = (box.left + box.right) // 2, (box.bottom + box.top) // 2
    if box.right <= box.left or box.top <= box.bottom:
        return [db.Point(cx, cy)]
    dx = max(1, (box.right - box.left) // 4)
    dy = max(1, (box.top - box.bottom) // 4)
    return [db.Point(cx, cy), db.Point(cx - dx, cy), db.Point(cx + dx, cy),
            db.Point(cx, cy - dy), db.Point(cx, cy + dy)]


# --- Pass 1: flat connectivity -----------------------------------------------

class Routing:
    """The chip's routing as one flat graph, probeable by point.

    RecursiveShapeIterator walks the top cell and all descendants applying
    instance transforms, so the regions land in top-cell coordinates.  After
    extract_netlist, probe() is a point-to-component lookup.
    """

    def __init__(self, ly: db.Layout, top: db.Cell):
        self.l2n = db.LayoutToNetlist(top.name, ly.dbu)
        self.regions: Dict[str, db.Region] = {}
        for name, (layer, datatype) in ROUTING_LAYERS.items():
            region = db.Region(db.RecursiveShapeIterator(ly, top, ly.layer(layer, datatype)))
            self.l2n.register(region, name)
            self.regions[name] = region
        for region in self.regions.values():
            self.l2n.connect(region)  # touching polygons on one layer are one wire
        for lower, cut, upper in STACK:
            self.l2n.connect(self.regions[lower], self.regions[cut])
            self.l2n.connect(self.regions[cut], self.regions[upper])
        self.l2n.extract_netlist()

    def probe(self, layer: str, point: db.Point) -> Optional[str]:
        return net_name(self.l2n.probe_net(self.regions[layer], point))


# --- Pass 2: standard-cell pins ----------------------------------------------

def build_pin_catalog(ly: db.Layout, masters):
    """master ID -> pin -> [(layer, master-local point), ...].

    A pin may be labelled on several layers; all are kept as a consistency
    check.  Coordinates stay master-local until a placement is known.
    """
    label_layers = [(ly.layer(layer, dt), route) for layer, dt, route in PIN_TEXT_LAYERS]
    catalog = {}
    for master in masters:
        pins = defaultdict(list)
        for idx, route in label_layers:
            for shape in master.shapes(idx).each():
                if shape.is_text() and (pin := str(shape.text.string).strip()):
                    disp = shape.text.trans.disp
                    pins[pin].append((route, db.Point(int(disp.x), int(disp.y))))
        catalog[cell_index_of(master)] = dict(pins)
    return catalog


def extract_pins(ly, top, routing: Routing, catalog):
    """Map every placed functional pin to a flat net.

    master-local pin -> instance transform -> top coordinate -> probe -> net.

    OK means every label location for that pin agreed on one net.  UNRESOLVED
    and CONFLICT rows are kept as evidence rather than dropped, so a bad
    extraction fails loudly instead of emitting a plausible wrong circuit.
    """
    rows, stats, count = [], Counter(), 0
    for inst, master in each_logic_instance(ly, top):
        pins = catalog.get(cell_index_of(master), {})
        if not pins:
            raise RuntimeError(f"No /5 pin labels found for functional cell {master.name}")

        for tr in iter_transforms(inst):
            count += 1
            # GDS carries no logical instance name, so number deterministically;
            # identity still rests on the cell type and the recovered nets.
            instance = f"U{count:04d}"
            for pin, locations in sorted(pins.items()):
                hits = {n for route, local in locations if (n := routing.probe(route, tr * local))}
                if len(hits) == 1:
                    net, status, key = next(iter(hits)), "OK", "mapped_pins"
                elif not hits:
                    net, status, key = "", "UNRESOLVED", "unresolved_pins"
                else:
                    net, status, key = "|".join(sorted(hits)), "CONFLICT", "conflicting_pins"
                stats[key] += 1
                rows.append({"instance": instance, "cell": master.name, "pin": pin,
                             "net": net, "status": status, "transform": str(tr)})

    stats["logic_instances"] = count
    return rows, stats


# --- Pass 3: driver/load graph -----------------------------------------------

@dataclass
class Graph:
    """The recovered netlist indexed by role."""
    power_hits: Dict[str, Counter]
    power_by_pin: Dict[str, str]
    power_nets: Set[str]
    drivers: Dict[str, List[Tuple[str, str, str]]]
    loads: Dict[str, List[Tuple[str, str, str]]]
    signal_nets: Set[str]


def analyze(rows: Sequence[dict]) -> Graph:
    """Index pin rows by net, separating rails from Boolean signals.

    Every cell should see the same rails, so the most-observed net per power
    pin is that rail's identity.
    """
    power_hits = defaultdict(Counter)
    drivers, loads = defaultdict(list), defaultdict(list)
    for r in rows:
        if r["status"] != "OK":
            continue
        if r["pin"] in POWER_PINS:
            power_hits[r["pin"]][r["net"]] += 1
        else:
            side = drivers if r["pin"] in OUTPUT_PINS else loads
            side[r["net"]].append((r["instance"], r["cell"], r["pin"]))

    power_by_pin = {pin: hits.most_common(1)[0][0] for pin, hits in power_hits.items() if hits}
    rails = set(power_by_pin.values())
    return Graph(power_hits, power_by_pin, rails, dict(drivers), dict(loads),
                 (set(drivers) | set(loads)) - rails)


# --- Pass 4: top-level ports -------------------------------------------------

@dataclass
class TopPort:
    """Every physical observation belonging to one boundary net.

    A port may own several pin-purpose rectangles and several labels, so the
    flat net ID is the identity.  Coordinates are integer database units.
    """
    net: str
    shape_boxes: List[Tuple[int, int, int, int]] = field(default_factory=list)
    labels: Set[str] = field(default_factory=set)
    direction: str = "unknown"
    name: str = ""
    side: str = ""
    x: int = 0
    y: int = 0
    edge_distance: int = 0
    drivers: int = 0
    loads: int = 0

    def finalize_geometry(self, bbox: db.Box):
        """Reduce the shapes to one point and a nearest die edge."""
        self.x = (min(b[0] for b in self.shape_boxes) + max(b[2] for b in self.shape_boxes)) // 2
        self.y = (min(b[1] for b in self.shape_boxes) + max(b[3] for b in self.shape_boxes)) // 2
        distances = {"LEFT": abs(self.x - bbox.left), "RIGHT": abs(bbox.right - self.x),
                     "BOTTOM": abs(self.y - bbox.bottom), "TOP": abs(bbox.top - self.y)}
        self.side = min(distances, key=distances.get)
        self.edge_distance = int(distances[self.side])


def discover_top_ports(ly, top, routing: Routing):
    """Find nets reaching the outside of the top cell.

    Only geometry drawn directly in top.shapes counts; a /16 shape inside a
    master is not a boundary port.  The net, not the shape, is the identity —
    one port may surface on several metal layers.
    """
    ports: Dict[str, TopPort] = {}
    stats = Counter()
    for layer, datatype, route in TOP_PIN_LAYERS:
        for shape in top.shapes(ly.layer(layer, datatype)).each():
            if shape.is_text():
                continue
            box = shape.bbox()
            hits = {n for p in sample_points(box) if (n := routing.probe(route, p))}
            if not hits:
                stats["unresolved_shapes"] += 1
            elif len(hits) > 1:
                stats["conflicting_shapes"] += 1
            else:
                key = next(iter(hits))
                ports.setdefault(key, TopPort(net=key)).shape_boxes.append(
                    (box.left, box.bottom, box.right, box.top))
                stats["accepted_shapes"] += 1

    bbox = top.bbox()
    for port in ports.values():
        port.finalize_geometry(bbox)
    return ports, stats


def attach_top_labels(ly, top, routing: Routing, ports):
    """Attach top-level text to ports through the routing graph, not proximity."""
    for layer, datatype, route in PIN_TEXT_LAYERS:
        for shape in top.shapes(ly.layer(layer, datatype)).each():
            if not shape.is_text():
                continue
            disp = shape.text.trans.disp
            key = routing.probe(route, db.Point(int(disp.x), int(disp.y)))
            if (label := str(shape.text.string).strip()) and key in ports:
                ports[key].labels.add(label)


def classify_top_ports(ports, g: Graph):
    """Label each boundary net power/input/output from circuit structure.

    Outputs may also feed internal logic (success does), so having loads does
    not disqualify a port from being an output.
    """
    rail_name = {net: pin for pin, net in g.power_by_pin.items()}
    for net, p in ports.items():
        p.drivers, p.loads = len(g.drivers.get(net, [])), len(g.loads.get(net, []))
        if net in rail_name:
            p.direction, p.name = "power", rail_name[net]
        elif p.drivers:
            p.direction = "output"
        elif p.loads:
            p.direction = "input"


def canonical_label(label: str) -> Optional[str]:
    """Normalise a physical label to its Verilog name, or None if unrecognised."""
    label = label.strip()
    direct = {"clk": "clk", "rst_n": "rst_n", "enable": "enable",
              "I": "input", "input": "input", "success": "success"}
    if label in direct:
        return direct[label]
    m = re.fullmatch(r"(?:O|out)\[(\d+)\]", label)
    return f"out[{int(m.group(1))}]" if m and int(m.group(1)) <= 7 else None


def clock_score(net, loads) -> int:
    """A clock tree ends in sequential CLK pins; a clkbuf A is one step upstream."""
    return sum(100 if pin == "CLK" else 50 if pin == "A" and "__clkbuf_" in cell else 0
               for _inst, cell, pin in loads.get(net, []))


def reset_score(net, loads) -> int:
    return sum(10 for _i, _c, pin in loads.get(net, []) if pin in {"RESET_B", "SET_B"})


def assign_puzzle_names(ports, g: Graph) -> Tuple[bool, str]:
    """Turn anonymous boundary nets into the puzzle's interface.

    Preference order: a recognised top-level label; then reset and clock from
    the pin types they drive; then the last two inputs by vertical position;
    then the eight external AND3 outputs numbered top-to-bottom, leaving
    success.  The expected counts and name sets are assertions about this
    puzzle, there to stop a plausible but misnamed module being emitted.
    """
    functional = [p for p in ports.values() if p.direction in ("input", "output")]
    for p in functional:
        recognized = {canonical_label(label) for label in p.labels} - {None}
        if len(recognized) > 1:
            return False, f"Port net {p.net} has contradictory labels {sorted(recognized)}"
        if recognized:
            p.name = recognized.pop()

    inputs = [p for p in functional if p.direction == "input"]
    outputs = [p for p in functional if p.direction == "output"]
    if (len(inputs), len(outputs)) != (4, 9):
        return False, f"Expected 4 inputs and 9 outputs; got {len(inputs)} and {len(outputs)}"

    # Structural fallback, only where a label was missing.
    for name, score in (("rst_n", reset_score), ("clk", clock_score)):
        if any(p.name == name for p in inputs):
            continue
        best = max((p for p in inputs if not p.name), default=None,
                   key=lambda p: score(p.net, g.loads))
        if best is None or score(best.net, g.loads) == 0:
            return False, f"Could not identify {name}"
        best.name = name

    # Only enable/input can remain; the input bank's top-to-bottom order is stable.
    remaining = [n for n in ("enable", "input") if n not in {p.name for p in inputs}]
    unnamed = sorted((p for p in inputs if not p.name), key=lambda p: -p.y)
    if len(unnamed) != len(remaining):
        return False, "Could not uniquely assign enable/input"
    for p, name in zip(unnamed, remaining):
        p.name = name

    if not EXPECTED_OUTPUTS <= {p.name for p in outputs}:
        bank = [p for p in outputs if not g.loads.get(p.net)
                and [d[1:] for d in g.drivers.get(p.net, [])] == [(OUTPUT_BANK_CELL, "X")]]
        if len(bank) != 8:
            return False, f"Expected 8 AND3/X external outputs; got {len(bank)}"
        for i, p in enumerate(sorted(bank, key=lambda p: -p.y)):
            p.name = p.name or f"out[{i}]"
        if len(rest := [p for p in outputs if not p.name]) == 1:
            rest[0].name = "success"

    for kind, got, want in (("Input", {p.name for p in inputs}, EXPECTED_INPUTS),
                            ("Output", {p.name for p in outputs}, EXPECTED_OUTPUTS)):
        if got != want:
            return False, f"{kind} name set mismatch: {sorted(got)}"
    return True, "Puzzle interface assigned successfully from physical top ports."


# --- Pass 5: validation ------------------------------------------------------

def validate(stats, g: Graph, ports, port_stats) -> List[str]:
    """Check the invariants a correct extraction must satisfy.

    A GDS parser can emit syntactically valid Verilog even when one layer or
    transform is wrong; these turn that into a visible failure.
    """
    problems = []
    for counts, suffix, what in ((stats, "pins", "logical pins"),
                                 (port_stats, "shapes", "top-level /16 pin shapes")):
        for label in ("unresolved", "conflicting"):
            if n := counts[f"{label}_{suffix}"]:
                problems.append(f"{n} {label} {what}")

    vpwr, vgnd = g.power_by_pin.get("VPWR"), g.power_by_pin.get("VGND")
    if not vpwr or not vgnd:
        problems.append("Could not recover both VPWR and VGND")
    elif vpwr == vgnd:
        problems.append("VPWR and VGND resolved to the same flat net")

    rails = {n for n in (vpwr, vgnd) if n}
    for net, ds in g.drivers.items():
        # conb_1 HI/LO are intentional shared sources, not contention.
        real = [d for d in ds if not (d[1] == CONB_CELL and d[2] in {"HI", "LO"})]
        if net in rails and real:
            problems.append(f"Supply net {net} has functional drivers: {real[:8]}")
        elif net not in rails and len(real) > 1:
            problems.append(f"Logic net {net} has multiple non-constant drivers: {real[:8]}")

    counts = Counter(p.direction for p in ports.values())
    if (counts["input"], counts["output"]) != (4, 9):
        problems.append(f"Physical interface is {counts['input']} inputs / "
                        f"{counts['output']} outputs, expected 4 / 9")
    return problems


# --- Pass 6: Verilog ---------------------------------------------------------

def verilog_id(name: str) -> str:
    """Make a recovered name safe as an identifier (escaped ones end at space)."""
    return name if PLAIN_ID.fullmatch(name) and name not in VERILOG_KEYWORDS else "\\" + name + " "


def port_expr(name: str) -> str:
    m = BIT_SELECT.fullmatch(name)
    return f"{verilog_id(m.group(1))}[{m.group(2)}]" if m else verilog_id(name)


def build_port_declarations(ports: Sequence[TopPort]):
    """Regroup individually discovered bits into scalar and bus declarations.

    Discovery sees eight separate boundary nets; Verilog can present them as
    one out[7:0], provided the bus is contiguous and single-direction.
    """
    scalars = []
    buses = defaultdict(lambda: {"bits": set(), "direction": None})
    for p in ports:
        if not (m := BIT_SELECT.fullmatch(p.name)):
            scalars.append((p.direction, p.name))
            continue
        bus = buses[m.group(1)]
        if bus["direction"] not in (None, p.direction):
            raise RuntimeError(f"Mixed directions on bus {m.group(1)}")
        bus["bits"].add(int(m.group(2)))
        bus["direction"] = p.direction

    order = {"clk": 0, "rst_n": 1, "enable": 2, "input": 3, "success": 4}
    scalars.sort(key=lambda s: order.get(s[1], 100))
    header = [verilog_id(name) for _d, name in scalars]
    decls = [f"{d} {verilog_id(name)};" for d, name in scalars]

    for base in sorted(buses):
        bits = sorted(buses[base]["bits"])
        if bits != list(range(bits[0], bits[-1] + 1)):
            raise RuntimeError(f"Sparse bus {base}: {bits}")
        header.append(verilog_id(base))
        decls.append(f"{buses[base]['direction']} [{bits[-1]}:{bits[0]}] {verilog_id(base)};")
    return header, decls


def generate_verilog(path, module_name, rows, ports, g: Graph):
    """Serialise the validated endpoint graph as structural Verilog.

    No Boolean behaviour is re-derived: the SKY130 cell names are kept so a
    library-aware simulator interprets the same gates.  Power and bulk pins are
    dropped and the rails appear as 1'b1/1'b0, which suffices for the logical
    model and avoids a rail port interface.
    """
    functional = [p for p in ports.values() if p.direction in ("input", "output")]
    if any(not p.name for p in functional):
        raise RuntimeError("Refusing to emit unnamed functional ports")

    boundary = {p.net: port_expr(p.name) for p in functional}
    instances, signal_nets = defaultdict(list), set()
    for r in rows:
        if r["status"] != "OK":
            continue
        instances[r["instance"]].append(r)
        if r["pin"] not in POWER_PINS:
            signal_nets.add(r["net"])

    # Sorting the opaque IDs keeps wire names stable across runs.
    internal = {net: f"n_{i:04d}"
                for i, net in enumerate(sorted(signal_nets - set(boundary) - g.power_nets))}
    vpwr, vgnd = g.power_by_pin.get("VPWR"), g.power_by_pin.get("VGND")

    def expr(net: str) -> str:
        if net in boundary:
            return boundary[net]
        if net in internal:
            return internal[net]
        if net in (vpwr, vgnd):
            return "1'b1" if net == vpwr else "1'b0"
        raise RuntimeError(f"No Verilog expression for net {net}")

    header, decls = build_port_declarations(functional)
    with open(path, "w") as f:
        f.write("// Auto-generated from GDSII.\n"
                "// Structural SKY130 HD gate-level netlist.\n"
                "// Power/bulk pins are omitted from cell instances.\n\n")
        f.write(f"module {verilog_id(module_name)} (\n")
        f.write(",\n".join(f"    {name}" for name in header) + "\n);\n\n")
        f.write("\n".join(decls) + "\n\n")
        f.write("".join(f"wire {name};\n" for name in internal.values()) + "\n")

        for instance in sorted(instances, key=lambda n: int(re.search(r"(\d+)$", n).group(1))):
            pin_rows = instances[instance]
            pin_to_net = {}
            for r in pin_rows:
                # A final guard: repeated labels for one pin must have agreed.
                if r["pin"] not in POWER_PINS and pin_to_net.setdefault(r["pin"], r["net"]) != r["net"]:
                    raise RuntimeError(f"{instance}/{r['pin']} maps to multiple nets")
            f.write(f"{verilog_id(pin_rows[0]['cell'])} {verilog_id(instance)} (\n")
            f.write(",\n".join(f"    .{verilog_id(pin)}({expr(net)})"
                               for pin, net in sorted(pin_to_net.items())) + "\n);\n\n")
        f.write("endmodule\n")


# --- Reporting ---------------------------------------------------------------

PIN_FIELDS = ["instance", "cell", "pin", "net", "status", "transform"]
PORT_FIELDS = ["net", "direction", "name", "side", "x", "y", "edge_distance",
               "shape_count", "drivers", "loads", "labels"]


def write_csv(path: str, fieldnames: Sequence[str], rows):
    with open(path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def by_position(p: TopPort):
    return (p.direction, -p.y, p.x, p.net)


def port_rows(ports):
    for p in sorted(ports.values(), key=by_position):
        yield {"net": p.net, "direction": p.direction, "name": p.name, "side": p.side,
               "x": p.x, "y": p.y, "edge_distance": p.edge_distance,
               "shape_count": len(p.shape_boxes), "drivers": p.drivers, "loads": p.loads,
               "labels": "|".join(sorted(p.labels))}


def make_summary(gds, top, dbu, stats, g: Graph, ports, port_stats, naming, problems):
    """Format the extraction audit report."""
    lines = [f"GDS: {gds}", f"TOP: {top.name}", f"DBU: {dbu}"]

    def section(title, body):
        lines.extend(["", title, "=" * 78, *body])

    section("FLAT STANDARD-CELL PIN EXTRACTION",
            [f"{k:24s}: {stats[k]}" for k in
             ("logic_instances", "mapped_pins", "unresolved_pins", "conflicting_pins")])
    section("POWER NETS", [
        f"{pin:8s}: {g.power_by_pin[pin]} ({g.power_hits[pin][g.power_by_pin[pin]]} hits)"
        if pin in g.power_by_pin else f"{pin:8s}: <not recovered / not needed>"
        for pin in ("VPWR", "VGND", "VPB", "VNB")])
    section("SIGNAL GRAPH", [f"signal nets: {len(g.signal_nets)}",
                             f"driver nets: {len(g.drivers)}",
                             f"load nets  : {len(g.loads)}"])
    section("TOP PORTS",
            [f"{kind} /16 shapes: {port_stats[key]}" for kind, key in
             (("accepted   ", "accepted_shapes"), ("unresolved ", "unresolved_shapes"),
              ("conflicting", "conflicting_shapes"))] +
            [f"{p.direction:8s} {p.name or '<unnamed>':12s} {p.net:10s} side={p.side:6s} "
             f"x={p.x:7d} y={p.y:7d} drivers={p.drivers:2d} loads={p.loads:3d} "
             f"labels={','.join(sorted(p.labels))}"
             for p in sorted(ports.values(), key=by_position)])
    section("PUZZLE PORT NAMING", [naming])
    section("VALIDATION", ["ERROR: " + p for p in problems] or ["OK"])
    return "\n".join(lines) + "\n"


# --- Main --------------------------------------------------------------------

def main() -> int:
    """Routing before pins, pins before rails and directions, directions before
    interface naming."""
    ap = argparse.ArgumentParser()
    ap.add_argument("gds")
    ap.add_argument("--top", default="puzzle")
    ap.add_argument("--outdir", default="recovered")
    ap.add_argument("--module", default="puzzle_recovered")
    args = ap.parse_args()

    os.makedirs(args.outdir, exist_ok=True)
    out = lambda name: os.path.join(args.outdir, name)

    ly = db.Layout()
    ly.read(args.gds)
    if (top := ly.cell(args.top)) is None:
        raise RuntimeError(f"Top cell {args.top!r} not found")
    print(f"DBU: {ly.dbu}\nTOP: {top.name}\nExtracting flat routing connectivity ...")

    routing = Routing(ly, top)
    masters = {cell_index_of(m): m for _inst, m in each_logic_instance(ly, top)}
    rows, stats = extract_pins(ly, top, routing, build_pin_catalog(ly, masters.values()))
    write_csv(out("pin_nets.csv"), PIN_FIELDS, rows)

    g = analyze(rows)
    ports, port_stats = discover_top_ports(ly, top, routing)
    attach_top_labels(ly, top, routing, ports)
    classify_top_ports(ports, g)
    naming_ok, naming = assign_puzzle_names(ports, g)
    write_csv(out("top_ports.csv"), PORT_FIELDS, port_rows(ports))

    # Fail closed, but only after the diagnostics are on disk, so a bad run can
    # be investigated without repeating the expensive GDS pass.
    problems = validate(stats, g, ports, port_stats)
    if not naming_ok:
        problems.append(naming)
    summary = make_summary(args.gds, top, ly.dbu, stats, g, ports, port_stats, naming, problems)
    with open(out("extraction.txt"), "w") as f:
        f.write(summary)
    print("\n" + summary)

    written = ["pin_nets.csv", "top_ports.csv", "extraction.txt"]
    if problems:
        print("VALIDATION FAILED")
        print("Diagnostic files:\n  " + "\n  ".join(out(n) for n in written))
        print("recovered.v was NOT emitted.")
        return 2

    generate_verilog(out("recovered.v"), args.module, rows, ports, g)
    print("WROTE\n  " + "\n  ".join(out(n) for n in written + ["recovered.v"]))
    return 0


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

open Hardcaml
module S = Signal

module H = Hardcaml_of_verilog
module Y = H.Expert.Yosys_netlist
module N = H.Netlist

type bit = Wire of int | Constant of int
type gate = And of bit * bit | Or of bit * bit | Not of bit
type reset_kind = Reset_zero | Reset_one | No_reset

type dff = { name : string; kind : reset_kind; d : bit; q : int }

type netlist =
  { clk_bits : int list
  ; rst_bits : int list
  ; enable_bits : int list
  ; input_bits : int list
  ; success_bits : int list
  ; out_bits : int list
  ; gates : (int, gate) Hashtbl.t
  ; dffs : dff list
  }

exception Parse_error of string

let fail message = raise (Parse_error message)

let read_file path =
  let channel = open_in_bin path in
  let length = in_channel_length channel in
  let contents = really_input_string channel length in
  close_in channel;
  contents

let unwrap = function
  | Ok value -> value
  | Error error -> fail (Base.Error.to_string_hum error)

let bit = function
  | Y.Bit.Index value -> Wire value
  | Y.Bit.Gnd -> Constant 0
  | Y.Bit.Vdd -> Constant 1
  | Y.Bit.X -> fail "unsupported X value in normalized netlist"

let one_bit = function
  | [ value ] -> bit value
  | _ -> fail "expected a one-bit connection"

let port_value ports name =
  match List.find_opt (fun port -> String.equal port.N.Port.name name) ports with
  | Some port -> port.N.Port.value
  | None -> fail ("missing cell port " ^ name)

let cell_input cell name = one_bit (port_value cell.N.Cell.inputs name)
let cell_output cell name = one_bit (port_value cell.N.Cell.outputs name)

let parse path =
  let yosys_netlist = Y.of_string (read_file path) |> unwrap in
  let netlist = N.of_yosys_netlist yosys_netlist |> unwrap in
  let module_ = N.find_module_by_name netlist "puzzle_recovered" |> unwrap in
  let port_bits name =
    let ports = module_.N.Module.inputs @ module_.N.Module.outputs in
    (* [out] is an eight-bit bus, so keep the bus instead of collapsing it. *)
    List.map
      (function
        | Y.Bit.Index value -> value
        | Y.Bit.Gnd | Y.Bit.Vdd | Y.Bit.X ->
            fail ("top-level port " ^ name ^ " is not a wire"))
      (port_value ports name)
  in
  let cells = module_.N.Module.cells in
  let gates = Hashtbl.create (List.length cells) in
  let dffs = ref [] in
  List.iter
    (fun cell ->
      let kind = cell.N.Cell.module_name in
      match kind with
      | "$_AND_" ->
          Hashtbl.replace gates
            (match cell_output cell "Y" with
            | Wire bit -> bit
            | Constant _ -> fail "gate output is constant")
            (And (cell_input cell "A", cell_input cell "B"))
      | "$_OR_" ->
          Hashtbl.replace gates
            (match cell_output cell "Y" with
            | Wire bit -> bit
            | Constant _ -> fail "gate output is constant")
            (Or (cell_input cell "A", cell_input cell "B"))
      | "$_NOT_" ->
          Hashtbl.replace gates
            (match cell_output cell "Y" with
            | Wire bit -> bit
            | Constant _ -> fail "gate output is constant")
            (Not (cell_input cell "A"))
      | "$_DFF_PN0_" -> (
          match cell_output cell "Q" with
          | Wire q ->
              dffs :=
                { name = cell.N.Cell.instance_name
                ; kind = Reset_zero
                ; d = cell_input cell "D"
                ; q
                }
                :: !dffs
          | Constant _ -> fail "DFF Q is constant")
      | "$_DFF_PN1_" -> (
          match cell_output cell "Q" with
          | Wire q ->
              dffs :=
                { name = cell.N.Cell.instance_name
                ; kind = Reset_one
                ; d = cell_input cell "D"
                ; q
                }
                :: !dffs
          | Constant _ -> fail "DFF Q is constant")
      | "$_DFF_P_" -> (
          match cell_output cell "Q" with
          | Wire q ->
              dffs :=
                { name = cell.N.Cell.instance_name
                ; kind = No_reset
                ; d = cell_input cell "D"
                ; q
                }
                :: !dffs
          | Constant _ -> fail "DFF Q is constant")
      | other -> fail ("unsupported optimized cell " ^ other))
    cells;
  { clk_bits = port_bits "clk"
  ; rst_bits = port_bits "rst_n"
  ; enable_bits = port_bits "enable"
  ; input_bits = port_bits "input"
  ; success_bits = port_bits "success"
  ; out_bits = port_bits "out"
  ; gates
  ; dffs = List.rev !dffs
  }

let circuit
    { clk_bits; rst_bits; enable_bits; input_bits; success_bits; out_bits; gates; dffs } =
  let one_port name = function
    | [ bit ] -> (S.input name 1, bit)
    | _ -> fail ("expected one-bit port " ^ name)
  in
  let open S in
  let clk, _ = one_port "clk" clk_bits in
  let rst_n, _ = one_port "rst_n" rst_bits in
  let enable = input "enable" 1 in
  let serial_input = input "input" 1 in
  let input_signals = Hashtbl.create 8 in
  let add_input bits signal name =
    match bits with
    | [ bit ] -> Hashtbl.add input_signals bit signal
    | _ -> fail ("expected one-bit input " ^ name)
  in
  add_input clk_bits clk "clk";
  add_input rst_bits rst_n "rst_n";
  add_input enable_bits enable "enable";
  add_input input_bits serial_input "input";
  let q_signals = Hashtbl.create (List.length dffs) in
  List.iter (fun dff -> Hashtbl.add q_signals dff.q (wire 1)) dffs;
  let memo = Hashtbl.create 4096 in
  let rec get = function
    | Constant value -> if value = 0 then gnd else vdd
    | Wire bit -> (
        match Hashtbl.find_opt memo bit with
        | Some value -> value
        | None ->
            let value =
              match
                ( Hashtbl.find_opt input_signals bit,
                  Hashtbl.find_opt q_signals bit )
              with
              | Some value, _ -> value
              | _, Some value -> value
              | None, None -> (
                  match Hashtbl.find_opt gates bit with
                  | Some (And (a, b)) -> get a &: get b
                  | Some (Or (a, b)) -> get a |: get b
                  | Some (Not a) -> ~:(get a)
                  | None -> fail ("net has no driver: " ^ string_of_int bit))
            in
            Hashtbl.add memo bit value;
            value)
  in
  let spec = Reg_spec.create ~clock:clk ~reset:rst_n ~reset_level:Level.Low () in
  let spec_no_reset = Reg_spec.create ~clock:clk () in
  List.iter
    (fun dff ->
      let q = Hashtbl.find q_signals dff.q in
      let d = get dff.d in
      let registered =
        match dff.kind with
        | Reset_zero -> reg spec ~reset_to:(Bits.of_int_trunc ~width:1 0) d
        | Reset_one -> reg spec ~reset_to:(Bits.of_int_trunc ~width:1 1) d
        (* Four flops in the extracted netlist have no reset. Their value does
           not matter until success, so starting them at zero gives Cyclesim
           the same deterministic start as the reference VCD. *)
        | No_reset ->
            reg spec_no_reset ~initialize_to:(Bits.of_int_trunc ~width:1 0) d
      in
      assign q registered)
    dffs;
  let port_value name bits =
    let bits = List.rev bits |> List.map (fun bit -> get (Wire bit)) in
    match bits with
    | [] -> fail ("empty output port " ^ name)
    | bits -> concat_msb bits
  in
  Circuit.create_exn ~name:"recovered_netlist_oracle"
    [
      output "success" (port_value "success" success_bits);
      output "out" (port_value "out" out_bits);
    ]

open Hardcaml

(*
   The serial protocol and simulator used to drive the recovered gate oracle.

   Hold rst_n low to reset, then present one bit per enabled clock.  A cycle
   with enable=0 is a stall and must leave every bit of puzzle state untouched.
   After the last cell, further clocks pump the output sequencer.
*)

type sample = { cycle : int; success : bool; out : int }

type t = {
  sim : Cyclesim.t_port_list;
  rst_n_port : Bits.t ref;
  enable_port : Bits.t ref;
  input_port : Bits.t ref;
  success_port : Bits.t ref;
  out_port : Bits.t ref;
}

let create circuit =
  let sim = Cyclesim.create circuit in
  {
    sim;
    rst_n_port = Cyclesim.in_port sim "rst_n";
    enable_port = Cyclesim.in_port sim "enable";
    input_port = Cyclesim.in_port sim "input";
    success_port = Cyclesim.out_port sim "success";
    out_port = Cyclesim.out_port sim "out";
  }

let set port value = port := if value then Bits.vdd else Bits.gnd

let reset t =
  Cyclesim.reset t.sim;
  set t.rst_n_port true;
  set t.enable_port false;
  set t.input_port false

let sample t cycle : sample =
  {
    cycle;
    success = Bits.to_int_trunc !(t.success_port) = 1;
    out = Bits.to_int_trunc !(t.out_port);
  }

(* [cycles] is a list of (enable, bit) pairs so callers can inject stalls. *)
let run ?(trailing_cycles = 16) circuit cycles =
  let t = create circuit in
  reset t;
  let samples = ref [] in
  let cycle = ref 0 in
  let clock enable bit =
    set t.enable_port enable;
    set t.input_port bit;
    Cyclesim.cycle t.sim;
    samples := sample t !cycle :: !samples;
    incr cycle
  in
  List.iter (fun (enable, bit) -> clock enable bit) cycles;
  for _ = 1 to trailing_cycles do
    clock false false
  done;
  List.rev !samples

(* Whitespace is ignored so board rows can be pasted in as eleven lines. *)
let bits_of_string text =
  let significant =
    String.to_seq text
    |> Seq.filter (fun c -> not (List.mem c [ ' '; '\n'; '\r'; '\t' ]))
    |> List.of_seq
  in
  if List.exists (fun c -> c <> '0' && c <> '1') significant then
    invalid_arg "input must contain only 0 and 1 (whitespace is ignored)";
  List.map (fun c -> c = '1') significant

let enabled_cycles text = bits_of_string text |> List.map (fun bit -> (true, bit))

let simulate ?trailing_cycles circuit text =
  run ?trailing_cycles circuit (enabled_cycles text)

let nonzero_outputs samples =
  List.filter_map (fun (s : sample) -> if s.out = 0 then None else Some s.out) samples

let output_string samples =
  nonzero_outputs samples |> List.map Char.chr |> List.to_seq |> String.of_seq

let first_success samples = List.find_opt (fun (s : sample) -> s.success) samples