Enhance TraceCleanupSolver with merging logic

lib/solvers/TraceCleanupSolver/TraceCleanupSolver.ts

@@ -6,6 +6,8 @@ import { BaseSolver } from "lib/solvers/BaseSolver/BaseSolver"
 import type { SolvedTracePath } from "lib/solvers/SchematicTraceLinesSolver/SchematicTraceLinesSolver"
 import { visualizeInputProblem } from "lib/solvers/SchematicTracePipelineSolver/visualizeInputProblem"
 import type { NetLabelPlacement } from "../NetLabelPlacementSolver/NetLabelPlacementSolver"
+import { UntangleTraceSubsolver } from "./sub-solver/UntangleTraceSubsolver"
+import { is4PointRectangle } from "./is4PointRectangle"
 
 /**
  * Defines the input structure for the TraceCleanupSolver.
@@ -18,9 +20,6 @@ interface TraceCleanupSolverInput {
   paddingBuffer: number
 }
 
-import { UntangleTraceSubsolver } from "./sub-solver/UntangleTraceSubsolver"
-import { is4PointRectangle } from "./is4PointRectangle"
-
 /**
  * Represents the different stages or steps within the trace cleanup pipeline.
  */
@@ -31,38 +30,30 @@ type PipelineStep =
 
 /**
  * The TraceCleanupSolver is responsible for improving the aesthetics and readability of schematic traces.
- * It operates in a multi-step pipeline:
- * 1. **Untangling Traces**: It first attempts to untangle any overlapping or highly convoluted traces using a sub-solver.
- * 2. **Minimizing Turns**: After untangling, it iterates through each trace to minimize the number of turns, simplifying their paths.
- * 3. **Balancing L-Shapes**: Finally, it balances L-shaped trace segments to create more visually appealing and consistent layouts.
- * The solver processes traces one by one, applying these cleanup steps sequentially to refine the overall trace layout.
  */
 export class TraceCleanupSolver extends BaseSolver {
   private input: TraceCleanupSolverInput
   private outputTraces: SolvedTracePath[]
   private traceIdQueue: string[]
   private tracesMap: Map<string, SolvedTracePath>
   private pipelineStep: PipelineStep = "untangling_traces"
-  private activeTraceId: string | null = null // New property
+  private activeTraceId: string | null = null
   override activeSubSolver: BaseSolver | null = null
 
   constructor(solverInput: TraceCleanupSolverInput) {
     super()
     this.input = solverInput
     this.outputTraces = [...solverInput.allTraces]
     this.tracesMap = new Map(this.outputTraces.map((t) => [t.mspPairId, t]))
-    this.traceIdQueue = Array.from(
-      solverInput.allTraces.map((e) => e.mspPairId),
-    )
+    this.traceIdQueue = solverInput.allTraces.map((e) => e.mspPairId)
   }
 
   override _step() {
     if (this.activeSubSolver) {
       this.activeSubSolver.step()
+
       if (this.activeSubSolver.solved) {
-        const output = (
-          this.activeSubSolver as UntangleTraceSubsolver
-        ).getOutput()
+        const output = (this.activeSubSolver as UntangleTraceSubsolver).getOutput()
         this.outputTraces = output.traces
         this.tracesMap = new Map(this.outputTraces.map((t) => [t.mspPairId, t]))
         this.activeSubSolver = null
@@ -97,21 +88,18 @@ export class TraceCleanupSolver extends BaseSolver {
   private _runMinimizeTurnsStep() {
     if (this.traceIdQueue.length === 0) {
       this.pipelineStep = "balancing_l_shapes"
-      this.traceIdQueue = Array.from(
-        this.input.allTraces.map((e) => e.mspPairId),
-      )
+      this.traceIdQueue = this.input.allTraces.map((e) => e.mspPairId)
       return
     }
-
     this._processTrace("minimizing_turns")
   }
 
   private _runBalanceLShapesStep() {
     if (this.traceIdQueue.length === 0) {
+      this._mergeNearbySameNetTraces()
       this.solved = true
       return
     }
-
     this._processTrace("balancing_l_shapes")
   }
 
@@ -120,12 +108,9 @@ export class TraceCleanupSolver extends BaseSolver {
     this.activeTraceId = targetMspConnectionPairId
     const originalTrace = this.tracesMap.get(targetMspConnectionPairId)!
 
-    if (is4PointRectangle(originalTrace.tracePath)) {
-      return
-    }
+    if (is4PointRectangle(originalTrace.tracePath)) return
 
     const allTraces = Array.from(this.tracesMap.values())
-
     let updatedTrace: SolvedTracePath
 
     if (step === "minimizing_turns") {
@@ -146,35 +131,82 @@ export class TraceCleanupSolver extends BaseSolver {
     this.outputTraces = Array.from(this.tracesMap.values())
   }
 
-  getOutput() {
-    return {
-      traces: this.outputTraces,
+  // ✅ NEW LOGIC (BUG FIX)
+  private _mergeNearbySameNetTraces() {
+    const threshold = 5
+    const traces = [...this.outputTraces]
+    const merged: SolvedTracePath[] = []
+    const used = new Set<number>()
+
+    const dist = (a: { x: number; y: number }, b: { x: number; y: number }) =>
+      Math.hypot(a.x - b.x, a.y - b.y)
+
+    for (let i = 0; i < traces.length; i++) {
+      if (used.has(i)) continue
+      let base = traces[i]
+
+      for (let j = i + 1; j < traces.length; j++) {
+        if (used.has(j)) continue
+        const other = traces[j]
+
+        const sameNet = Object.values(this.input.mergedLabelNetIdMap).some(
+          (netSet) =>
+            base.mspConnectionPairIds.some((id) => netSet.has(id)) &&
+            other.mspConnectionPairIds.some((id) => netSet.has(id)),
+        )
+
+        if (!sameNet) continue
+
+        const aEnd = base.tracePath.at(-1)
+        const bStart = other.tracePath[0]
+        if (!aEnd || !bStart) continue
+
+        if (dist(aEnd, bStart) < threshold) {
+          base = {
+            ...base,
+            tracePath: [...base.tracePath, ...other.tracePath],
+            mspConnectionPairIds: Array.from(
+              new Set([
+                ...base.mspConnectionPairIds,
+                ...other.mspConnectionPairIds,
+              ]),
+            ),
+            pinIds: Array.from(new Set([...base.pinIds, ...other.pinIds])),
+          }
+          used.add(j)
+        }
+      }
+
+      used.add(i)
+      merged.push(base)
     }
+
+    this.outputTraces = merged
+    this.tracesMap = new Map(merged.map((t) => [t.mspPairId, t]))
+  }
+
+  getOutput() {
+    return { traces: this.outputTraces }
   }
 
   override visualize(): GraphicsObject {
-    if (this.activeSubSolver) {
-      return this.activeSubSolver.visualize()
-    }
+    if (this.activeSubSolver) return this.activeSubSolver.visualize()
 
     const graphics = visualizeInputProblem(this.input.inputProblem, {
       chipAlpha: 0.1,
       connectionAlpha: 0.1,
     })
 
     if (!graphics.lines) graphics.lines = []
-    if (!graphics.points) graphics.points = []
-    if (!graphics.rects) graphics.rects = []
-    if (!graphics.circles) graphics.circles = []
-    if (!graphics.texts) graphics.texts = []
 
     for (const trace of this.outputTraces) {
       const line: Line = {
         points: trace.tracePath.map((p) => ({ x: p.x, y: p.y })),
-        strokeColor: trace.mspPairId === this.activeTraceId ? "red" : "blue", // Highlight active trace
+        strokeColor: trace.mspPairId === this.activeTraceId ? "red" : "blue",
       }
-      graphics.lines!.push(line)
+      graphics.lines.push(line)
     }
     return graphics
   }
 }
+