feat: implement TraceAlignmentSolver to reduce trace zig-zags (#12)

lib/solvers/LayoutPipelineSolver/LayoutPipelineSolver.ts

@@ -12,12 +12,21 @@ import {
   type PackedPartition,
 } from "lib/solvers/PackInnerPartitionsSolver/PackInnerPartitionsSolver"
 import { PartitionPackingSolver } from "lib/solvers/PartitionPackingSolver/PartitionPackingSolver"
+import { TraceAlignmentSolver } from "lib/solvers/TraceAlignmentSolver/TraceAlignmentSolver"
 import type { ChipPin, InputProblem, PinId } from "lib/types/InputProblem"
 import type { OutputLayout } from "lib/types/OutputLayout"
 import { doBasicInputProblemLayout } from "./doBasicInputProblemLayout"
 import { visualizeInputProblem } from "./visualizeInputProblem"
 import { getPinIdToStronglyConnectedPinsObj } from "./getPinIdToStronglyConnectedPinsObj"
 
+export type LayoutPipelinePhase =
+  | "identifyDecouplingCapsSolver"
+  | "chipPartitionsSolver"
+  | "packInnerPartitionsSolver"
+  | "partitionPackingSolver"
+  | "traceAlignmentSolver"
+  | "none"
+
 type PipelineStep<T extends new (...args: any[]) => BaseSolver> = {
   solverName: string
   solverClass: T
@@ -53,6 +62,7 @@ export class LayoutPipelineSolver extends BaseSolver {
   chipPartitionsSolver?: ChipPartitionsSolver
   packInnerPartitionsSolver?: PackInnerPartitionsSolver
   partitionPackingSolver?: PartitionPackingSolver
+  traceAlignmentSolver?: TraceAlignmentSolver
 
   startTimeOfPhase: Record<string, number>
   endTimeOfPhase: Record<string, number>
@@ -124,6 +134,21 @@ export class LayoutPipelineSolver extends BaseSolver {
         },
       },
     ),
+    definePipelineStep(
+      "traceAlignmentSolver",
+      TraceAlignmentSolver,
+      () => [
+        {
+          inputProblem: this.inputProblem,
+          outputLayout: this.partitionPackingSolver!.finalLayout!,
+        },
+      ],
+      {
+        onSolved: (_solver) => {
+          // Alignment complete
+        },
+      },
+    ),
   ]
 
   constructor(inputProblem: InputProblem) {
@@ -188,6 +213,15 @@ export class LayoutPipelineSolver extends BaseSolver {
     if (!this.solved && this.activeSubSolver)
       return this.activeSubSolver.visualize()
 
+    // If the pipeline is complete and we have a trace alignment solver,
+    // show its final layout
+    if (this.solved && this.traceAlignmentSolver?.solved) {
+      return visualizeInputProblem(
+        this.inputProblem,
+        this.traceAlignmentSolver.outputLayout,
+      )
+    }
+
     // If the pipeline is complete and we have a partition packing solver,
     // show only the final chip placements
     if (this.solved && this.partitionPackingSolver?.solved) {
@@ -199,6 +233,12 @@ export class LayoutPipelineSolver extends BaseSolver {
     const chipPartitionsViz = this.chipPartitionsSolver?.visualize()
     const packInnerPartitionsViz = this.packInnerPartitionsSolver?.visualize()
     const partitionPackingViz = this.partitionPackingSolver?.visualize()
+    const traceAlignmentViz = this.traceAlignmentSolver?.solved
+      ? visualizeInputProblem(
+          this.inputProblem,
+          this.traceAlignmentSolver.outputLayout,
+        )
+      : null
 
     // Get basic layout positions to avoid overlapping at (0,0)
     const basicLayout = doBasicInputProblemLayout(this.inputProblem)
@@ -210,6 +250,7 @@ export class LayoutPipelineSolver extends BaseSolver {
       chipPartitionsViz,
       packInnerPartitionsViz,
       partitionPackingViz,
+      traceAlignmentViz,
     ]
       .filter(Boolean)
       .map((viz, stepIndex) => {
@@ -400,8 +441,10 @@ export class LayoutPipelineSolver extends BaseSolver {
 
     let finalLayout: OutputLayout
 
-    // Get the final layout from the partition packing solver
-    if (
+    // Get the final layout from the partition packing solver or trace alignment solver
+    if (this.traceAlignmentSolver?.solved) {
+      finalLayout = this.traceAlignmentSolver.outputLayout
+    } else if (
       this.partitionPackingSolver?.solved &&
       this.partitionPackingSolver.finalLayout
     ) {

lib/solvers/PackInnerPartitionsSolver/SingleInnerPartitionPackingSolver.ts

@@ -38,6 +38,13 @@ export class SingleInnerPartitionPackingSolver extends BaseSolver {
   }
 
   override _step() {
+    // Early return for decoupling caps: arrange them in a clean horizontal row
+    if (this.partitionInputProblem.partitionType === "decoupling_caps") {
+      this.layout = this.createLinearDecouplingCapLayout()
+      this.solved = true
+      return
+    }
+
     // Initialize PackSolver2 if not already created
     if (!this.activeSubSolver) {
       const packInput = this.createPackInput()
@@ -141,6 +148,45 @@ export class SingleInnerPartitionPackingSolver extends BaseSolver {
     }
   }
 
+  private createLinearDecouplingCapLayout(): OutputLayout {
+    const chipPlacements: Record<string, Placement> = {}
+
+    // Get all chips in this partition (which should all be decoupling caps)
+    // Sort by chipId for deterministic ordering
+    const chips = Object.entries(this.partitionInputProblem.chipMap).sort(
+      ([idA], [idB]) => idA.localeCompare(idB),
+    )
+
+    let minGap =
+      this.partitionInputProblem.decouplingCapsGap ??
+      this.partitionInputProblem.chipGap
+
+    // Calculate total width to center the row
+    const chipWidths = chips.map(([_, chip]) => chip.size.x)
+    const totalWidth =
+      chipWidths.reduce((sum, w) => sum + w, 0) +
+      minGap * Math.max(0, chips.length - 1)
+
+    // Start placing from the left side to center around x=0
+    let currentX = -totalWidth / 2
+
+    for (let i = 0; i < chips.length; i++) {
+      const chipEntry = chips[i]
+      const [chipId, chip] = chipEntry!
+      const halfWidth = chip.size.x / 2
+
+      chipPlacements[chipId] = {
+        x: currentX + halfWidth,
+        y: 0, // centered vertically
+        ccwRotationDegrees: 0, // Keep them uniformly rotated
+      }
+
+      currentX += chip.size.x + minGap
+    }
+
+    return { chipPlacements, groupPlacements: {} }
+  }
+
   private createLayoutFromPackingResult(
     packedComponents: PackSolver2["packedComponents"],
   ): OutputLayout {

lib/solvers/TraceAlignmentSolver/TraceAlignmentSolver.ts

@@ -0,0 +1,187 @@
+import { BaseSolver } from "../BaseSolver"
+import type { InputProblem, PinId, ChipId } from "lib/types/InputProblem"
+import type { OutputLayout, Placement } from "lib/types/OutputLayout"
+import type { Point } from "@tscircuit/math-utils"
+
+export class TraceAlignmentSolver extends BaseSolver {
+  inputProblem: InputProblem
+  outputLayout: OutputLayout
+  override MAX_ITERATIONS = 20
+
+  constructor(params: {
+    inputProblem: InputProblem
+    outputLayout: OutputLayout
+  }) {
+    super()
+    this.inputProblem = params.inputProblem
+    this.outputLayout = JSON.parse(JSON.stringify(params.outputLayout))
+  }
+
+  override _step() {
+    let nudgedAny = false
+
+    const strongConnections = Object.entries(this.inputProblem.pinStrongConnMap)
+      .filter(([_, connected]) => connected)
+      .map(([connKey]) => connKey.split("-") as [PinId, PinId])
+
+    for (const [pinIdA, pinIdB] of strongConnections) {
+      const chipIdA = this.getChipIdForPin(pinIdA)
+      const chipIdB = this.getChipIdForPin(pinIdB)
+
+      if (!chipIdA || !chipIdB || chipIdA === chipIdB) continue
+
+      const posA = this.getAbsolutePositionForPin(pinIdA)
+      const posB = this.getAbsolutePositionForPin(pinIdB)
+      if (!posA || !posB) continue
+
+      const dx = Math.abs(posA.x - posB.x)
+      const dy = Math.abs(posA.y - posB.y)
+
+      const threshold = 1.0
+
+      // Determine which axis to align based on proximity
+      const alignX = dx > 0 && dx < threshold && (dy === 0 || dx < dy)
+      const alignY = dy > 0 && dy < threshold && (dx === 0 || dy < dx)
+
+      if (!alignX && !alignY) continue
+
+      const pinsA = this.inputProblem.chipMap[chipIdA]?.pins.length || 0
+      const pinsB = this.inputProblem.chipMap[chipIdB]?.pins.length || 0
+
+      // Sort chips: nudge the one with fewer pins (or smaller ID)
+      const [toNudge, target] =
+        pinsA < pinsB || (pinsA === pinsB && chipIdA < chipIdB)
+          ? [
+              { id: chipIdA, pos: posA },
+              { id: chipIdB, pos: posB },
+            ]
+          : [
+              { id: chipIdB, pos: posB },
+              { id: chipIdA, pos: posA },
+            ]
+
+      const nudge = alignX
+        ? { x: target.pos.x - toNudge.pos.x, y: 0 }
+        : { x: 0, y: target.pos.y - toNudge.pos.y }
+
+      if (this.tryNudge(toNudge.id, nudge)) {
+        nudgedAny = true
+      }
+    }
+
+    if (!nudgedAny || this.iterations >= this.MAX_ITERATIONS) {
+      this.solved = true
+    }
+  }
+
+  private tryNudge(chipId: ChipId, nudge: { x: number; y: number }): boolean {
+    if (Math.abs(nudge.x) < 1e-6 && Math.abs(nudge.y) < 1e-6) return false
+
+    const placement = this.outputLayout.chipPlacements[chipId]
+    if (!placement) return false
+
+    const originalX = placement.x
+    const originalY = placement.y
+
+    placement.x += nudge.x
+    placement.y += nudge.y
+
+    if (this.hasOverlaps(chipId)) {
+      placement.x = originalX
+      placement.y = originalY
+      return false
+    }
+
+    return true
+  }
+
+  private hasOverlaps(chipId: ChipId): boolean {
+    const chipIds = Object.keys(this.outputLayout.chipPlacements)
+    const placement1 = this.outputLayout.chipPlacements[chipId]!
+    const chip1 = this.inputProblem.chipMap[chipId]!
+    const bounds1 = this.getRotatedBounds(placement1, chip1.size)
+
+    for (const otherId of chipIds) {
+      if (otherId === chipId) continue
+
+      const placement2 = this.outputLayout.chipPlacements[otherId]!
+      const chip2 = this.inputProblem.chipMap[otherId]!
+      const bounds2 = this.getRotatedBounds(placement2, chip2.size)
+
+      // Use a slightly larger epsilon for overlaps to avoid precision issues
+      if (this.calculateOverlapArea(bounds1, bounds2) > 0.0001) {
+        return true
+      }
+    }
+    return false
+  }
+
+  private getChipIdForPin(pinId: PinId): ChipId | null {
+    for (const [chipId, chip] of Object.entries(this.inputProblem.chipMap)) {
+      if (chip.pins.includes(pinId)) return chipId
+    }
+    return null
+  }
+
+  private getAbsolutePositionForPin(pinId: PinId): Point | null {
+    const chipPin = this.inputProblem.chipPinMap[pinId]
+    const chipId = this.getChipIdForPin(pinId)
+    if (!chipPin || !chipId) return null
+
+    const placement = this.outputLayout.chipPlacements[chipId]
+    if (!placement) return null
+
+    const rotatedOffset = this.rotatePoint(
+      chipPin.offset,
+      placement.ccwRotationDegrees,
+    )
+    return {
+      x: placement.x + rotatedOffset.x,
+      y: placement.y + rotatedOffset.y,
+    }
+  }
+
+  private rotatePoint(point: Point, angleDegrees: number): Point {
+    const angleRad = (angleDegrees * Math.PI) / 180
+    const cos = Math.cos(angleRad)
+    const sin = Math.sin(angleRad)
+    return {
+      x: point.x * cos - point.y * sin,
+      y: point.x * sin + point.y * cos,
+    }
+  }
+
+  private getRotatedBounds(
+    placement: Placement,
+    size: Point,
+  ): { minX: number; maxX: number; minY: number; maxY: number } {
+    const angleRad = (placement.ccwRotationDegrees * Math.PI) / 180
+    const cos = Math.abs(Math.cos(angleRad))
+    const sin = Math.abs(Math.sin(angleRad))
+    const rotatedWidth = size.x * cos + size.y * sin
+    const rotatedHeight = size.x * sin + size.y * cos
+    return {
+      minX: placement.x - rotatedWidth / 2,
+      maxX: placement.x + rotatedWidth / 2,
+      minY: placement.y - rotatedHeight / 2,
+      maxY: placement.y + rotatedHeight / 2,
+    }
+  }
+
+  private calculateOverlapArea(
+    b1: { minX: number; maxX: number; minY: number; maxY: number },
+    b2: { minX: number; maxX: number; minY: number; maxY: number },
+  ): number {
+    const overlapWidth = Math.min(b1.maxX, b2.maxX) - Math.max(b1.minX, b2.minX)
+    const overlapHeight =
+      Math.min(b1.maxY, b2.maxY) - Math.max(b1.minY, b2.minY)
+    if (overlapWidth <= 0 || overlapHeight <= 0) return 0
+    return overlapWidth * overlapHeight
+  }
+
+  override getConstructorParams(): [any] {
+    return [
+      { inputProblem: this.inputProblem, outputLayout: this.outputLayout },
+    ]
+  }
+}