region rebalancing

.gitignore

@@ -4,6 +4,7 @@ node_modules
 # output
 out
 dist
+cosmos-export
 *.tgz
 
 # code coverage

lib/core.ts

@@ -210,6 +210,21 @@ export const getTinyHyperGraphSolverOptions = (
 const compareCandidatesByF = (left: Candidate, right: Candidate) =>
   left.f - right.f
 
+const compareRegionCostSummaries = (
+  left: RegionCostSummary,
+  right: RegionCostSummary,
+) => {
+  if (left.totalRegionCost !== right.totalRegionCost) {
+    return left.totalRegionCost - right.totalRegionCost
+  }
+
+  if (left.maxRegionCost !== right.maxRegionCost) {
+    return left.maxRegionCost - right.maxRegionCost
+  }
+
+  return 0
+}
+
 interface SegmentGeometryScratch {
   lesserAngle: number
   greaterAngle: number
@@ -488,6 +503,36 @@ export class TinyHyperGraphSolver extends BaseSolver {
     )
   }
 
+  getSharedRegionIdsForPorts(port1Id: PortId, port2Id: PortId): RegionId[] {
+    const port1IncidentRegions = this.topology.incidentPortRegion[port1Id] ?? []
+    const port2IncidentRegions = this.topology.incidentPortRegion[port2Id] ?? []
+
+    return port1IncidentRegions.filter((regionId) =>
+      port2IncidentRegions.includes(regionId),
+    )
+  }
+
+  getAlternateRegionIdForSegment(
+    regionId: RegionId,
+    port1Id: PortId,
+    port2Id: PortId,
+  ): RegionId | undefined {
+    const sharedRegionIds = this.getSharedRegionIdsForPorts(port1Id, port2Id)
+
+    if (sharedRegionIds.length < 2 || !sharedRegionIds.includes(regionId)) {
+      return undefined
+    }
+
+    return sharedRegionIds.find((sharedRegionId) => sharedRegionId !== regionId)
+  }
+
+  canRebalanceRegionAssignment(
+    currentRegionId: RegionId,
+    nextRegionId: RegionId,
+  ): boolean {
+    return true
+  }
+
   isPortReservedForDifferentNet(portId: PortId): boolean {
     const reservedNetIds = this.problemSetup.portEndpointNetIds[portId]
     if (!reservedNetIds) {
@@ -511,7 +556,8 @@ export class TinyHyperGraphSolver extends BaseSolver {
   }
 
   isKnownSingleLayerRegion(regionId: RegionId): boolean {
-    const regionAvailableZMask = this.topology.regionAvailableZMask?.[regionId] ?? 0
+    const regionAvailableZMask =
+      this.topology.regionAvailableZMask?.[regionId] ?? 0
     return isKnownSingleLayerMask(regionAvailableZMask)
   }
 
@@ -525,15 +571,17 @@ export class TinyHyperGraphSolver extends BaseSolver {
     const port1IncidentRegions = topology.incidentPortRegion[port1Id]
     const port2IncidentRegions = topology.incidentPortRegion[port2Id]
     const angle1 =
-      port1IncidentRegions[0] === regionId || port1IncidentRegions[1] !== regionId
+      port1IncidentRegions[0] === regionId ||
+      port1IncidentRegions[1] !== regionId
         ? topology.portAngleForRegion1[port1Id]
-        : topology.portAngleForRegion2?.[port1Id] ??
-          topology.portAngleForRegion1[port1Id]
+        : (topology.portAngleForRegion2?.[port1Id] ??
+          topology.portAngleForRegion1[port1Id])
     const angle2 =
-      port2IncidentRegions[0] === regionId || port2IncidentRegions[1] !== regionId
+      port2IncidentRegions[0] === regionId ||
+      port2IncidentRegions[1] !== regionId
         ? topology.portAngleForRegion1[port2Id]
-        : topology.portAngleForRegion2?.[port2Id] ??
-          topology.portAngleForRegion1[port2Id]
+        : (topology.portAngleForRegion2?.[port2Id] ??
+          topology.portAngleForRegion1[port2Id])
     const z1 = topology.portZ[port1Id]
     const z2 = topology.portZ[port2Id]
     scratch.lesserAngle = angle1 < angle2 ? angle1 : angle2
@@ -548,9 +596,26 @@ export class TinyHyperGraphSolver extends BaseSolver {
     regionId: RegionId,
     port1Id: PortId,
     port2Id: PortId,
-  ) {
-    const { topology, state } = this
-    const regionCache = state.regionIntersectionCaches[regionId]
+  ): RegionIntersectionCache {
+    const nextRegionIntersectionCache = this.createNextRegionIntersectionCache(
+      regionId,
+      this.state.regionIntersectionCaches[regionId]!,
+      this.state.currentRouteNetId!,
+      port1Id,
+      port2Id,
+    )
+    this.state.regionIntersectionCaches[regionId] = nextRegionIntersectionCache
+    return nextRegionIntersectionCache
+  }
+
+  createNextRegionIntersectionCache(
+    regionId: RegionId,
+    regionCache: RegionIntersectionCache,
+    routeNetId: NetId,
+    port1Id: PortId,
+    port2Id: PortId,
+  ): RegionIntersectionCache {
+    const { topology } = this
     const segmentGeometry = this.populateSegmentGeometryScratch(
       regionId,
       port1Id,
@@ -562,7 +627,7 @@ export class TinyHyperGraphSolver extends BaseSolver {
       newEntryExitLayerChanges,
     ] = countNewIntersectionsWithValues(
       regionCache,
-      state.currentRouteNetId!,
+      routeNetId,
       segmentGeometry.lesserAngle,
       segmentGeometry.greaterAngle,
       segmentGeometry.layerMask,
@@ -572,7 +637,7 @@ export class TinyHyperGraphSolver extends BaseSolver {
 
     const netIds = new Int32Array(nextLength)
     netIds.set(regionCache.netIds)
-    netIds[nextLength - 1] = state.currentRouteNetId!
+    netIds[nextLength - 1] = routeNetId
 
     const lesserAngles = new Int32Array(nextLength)
     lesserAngles.set(regionCache.lesserAngles)
@@ -595,7 +660,7 @@ export class TinyHyperGraphSolver extends BaseSolver {
       regionCache.existingEntryExitLayerChanges + newEntryExitLayerChanges
     const existingSegmentCount = lesserAngles.length
 
-    state.regionIntersectionCaches[regionId] = {
+    return {
       netIds,
       lesserAngles,
       greaterAngles,
@@ -616,6 +681,114 @@ export class TinyHyperGraphSolver extends BaseSolver {
     }
   }
 
+  createRegionIntersectionCacheForSegments(
+    regionId: RegionId,
+    regionSegments: Array<[RouteId, PortId, PortId]>,
+  ): RegionIntersectionCache {
+    let regionIntersectionCache = createEmptyRegionIntersectionCache()
+
+    for (const [routeId, port1Id, port2Id] of regionSegments) {
+      regionIntersectionCache = this.createNextRegionIntersectionCache(
+        regionId,
+        regionIntersectionCache,
+        this.problem.routeNet[routeId]!,
+        port1Id,
+        port2Id,
+      )
+    }
+
+    return regionIntersectionCache
+  }
+
+  rebalanceRegionAssignments() {
+    let didChange = true
+
+    while (didChange) {
+      didChange = false
+
+      for (let regionId = 0; regionId < this.topology.regionCount; regionId++) {
+        const regionSegments = this.state.regionSegments[regionId]!
+        let segmentIndex = 0
+
+        while (segmentIndex < regionSegments.length) {
+          const [routeId, fromPortId, toPortId] = regionSegments[segmentIndex]!
+          const alternateRegionId = this.getAlternateRegionIdForSegment(
+            regionId,
+            fromPortId,
+            toPortId,
+          )
+
+          if (
+            alternateRegionId === undefined ||
+            !this.canRebalanceRegionAssignment(regionId, alternateRegionId)
+          ) {
+            segmentIndex += 1
+            continue
+          }
+
+          const nextSourceRegionSegments = regionSegments.filter(
+            (_, currentSegmentIndex) => currentSegmentIndex !== segmentIndex,
+          )
+          const movedSegment: [RouteId, PortId, PortId] = [
+            routeId,
+            fromPortId,
+            toPortId,
+          ]
+          const nextTargetRegionSegments = [
+            ...this.state.regionSegments[alternateRegionId]!,
+            movedSegment,
+          ]
+          const currentSummary: RegionCostSummary = {
+            maxRegionCost: Math.max(
+              this.state.regionIntersectionCaches[regionId]!.existingRegionCost,
+              this.state.regionIntersectionCaches[alternateRegionId]!
+                .existingRegionCost,
+            ),
+            totalRegionCost:
+              this.state.regionIntersectionCaches[regionId]!
+                .existingRegionCost +
+              this.state.regionIntersectionCaches[alternateRegionId]!
+                .existingRegionCost,
+          }
+          const nextSourceRegionCache =
+            this.createRegionIntersectionCacheForSegments(
+              regionId,
+              nextSourceRegionSegments,
+            )
+          const nextTargetRegionCache =
+            this.createRegionIntersectionCacheForSegments(
+              alternateRegionId,
+              nextTargetRegionSegments,
+            )
+          const nextSummary: RegionCostSummary = {
+            maxRegionCost: Math.max(
+              nextSourceRegionCache.existingRegionCost,
+              nextTargetRegionCache.existingRegionCost,
+            ),
+            totalRegionCost:
+              nextSourceRegionCache.existingRegionCost +
+              nextTargetRegionCache.existingRegionCost,
+          }
+
+          if (
+            compareRegionCostSummaries(nextSummary, currentSummary) >=
+            -Number.EPSILON
+          ) {
+            segmentIndex += 1
+            continue
+          }
+
+          regionSegments.splice(segmentIndex, 1)
+          this.state.regionSegments[alternateRegionId]!.push(movedSegment)
+          this.state.regionIntersectionCaches[regionId] = nextSourceRegionCache
+          this.state.regionIntersectionCaches[alternateRegionId] =
+            nextTargetRegionCache
+          didChange = true
+        }
+      }
+    }
+  }
+
   getSolvedPathSegments(finalCandidate: Candidate): Array<{
     regionId: RegionId
     fromPortId: PortId
@@ -769,6 +942,8 @@ export class TinyHyperGraphSolver extends BaseSolver {
       this.appendSegmentToRegionCache(regionId, fromPortId, toPortId)
     }
 
+    this.rebalanceRegionAssignments()
+
     state.candidateQueue.clear()
     state.currentRouteNetId = undefined
     state.currentRouteId = undefined

lib/section-solver/index.ts

@@ -52,9 +52,7 @@ export interface TinyHyperGraphSectionSolverOptions
 }
 
 const applyTinyHyperGraphSectionSolverOptions = (
-  solver:
-    | TinyHyperGraphSectionSearchSolver
-    | TinyHyperGraphSectionSolver,
+  solver: TinyHyperGraphSectionSearchSolver | TinyHyperGraphSectionSolver,
   options?: TinyHyperGraphSectionSolverOptions,
 ) => {
   applyTinyHyperGraphSolverOptions(solver, options)
@@ -131,7 +129,8 @@ const restoreSolvedStateSnapshot = (
   const clonedSnapshot = cloneSolvedStateSnapshot(snapshot)
   solver.state.portAssignment = clonedSnapshot.portAssignment
   solver.state.regionSegments = clonedSnapshot.regionSegments
-  solver.state.regionIntersectionCaches = clonedSnapshot.regionIntersectionCaches
+  solver.state.regionIntersectionCaches =
+    clonedSnapshot.regionIntersectionCaches
 }
 
 const summarizeRegionIntersectionCaches = (
@@ -245,7 +244,8 @@ const getOrderedRoutePath = (
   orderedRegionIds: RegionId[]
 } => {
   const routeSegments = solution.solvedRoutePathSegments[routeId] ?? []
-  const routeSegmentRegionIds = solution.solvedRoutePathRegionIds?.[routeId] ?? []
+  const routeSegmentRegionIds =
+    solution.solvedRoutePathRegionIds?.[routeId] ?? []
   const startPortId = problem.routeStartPort[routeId]
   const endPortId = problem.routeEndPort[routeId]
 
@@ -588,6 +588,7 @@ class TinyHyperGraphSectionSearchSolver extends TinyHyperGraphSolver {
   baselineBeatRipCount?: number
   previousBestMaxRegionCost = Number.POSITIVE_INFINITY
   ripsSinceBestMaxRegionCostImprovement = 0
+  mutableRegionIdSet: Set<RegionId>
 
   MAX_RIPS = Number.POSITIVE_INFINITY
   MAX_RIPS_WITHOUT_MAX_REGION_COST_IMPROVEMENT = Number.POSITIVE_INFINITY
@@ -605,6 +606,7 @@ class TinyHyperGraphSectionSearchSolver extends TinyHyperGraphSolver {
   ) {
     super(topology, problem, options)
     applyTinyHyperGraphSectionSolverOptions(this, options)
+    this.mutableRegionIdSet = new Set(mutableRegionIds)
     this.state.unroutedRoutes = [...activeRouteIds]
     this.applyFixedSegments()
     this.fixedSnapshot = cloneSolvedStateSnapshot({
@@ -616,7 +618,11 @@ class TinyHyperGraphSectionSearchSolver extends TinyHyperGraphSolver {
 
   applyFixedSegments() {
     for (const routePlan of this.routePlans) {
-      for (const { regionId, fromPortId, toPortId } of routePlan.fixedSegments) {
+      for (const {
+        regionId,
+        fromPortId,
+        toPortId,
+      } of routePlan.fixedSegments) {
         this.state.currentRouteNetId = this.problem.routeNet[routePlan.routeId]
         this.state.regionSegments[regionId]!.push([
           routePlan.routeId,
@@ -698,6 +704,16 @@ class TinyHyperGraphSectionSearchSolver extends TinyHyperGraphSolver {
     this.state.goalPortId = -1
   }
 
+  override canRebalanceRegionAssignment(
+    currentRegionId: RegionId,
+    nextRegionId: RegionId,
+  ): boolean {
+    return (
+      this.mutableRegionIdSet.has(currentRegionId) &&
+      this.mutableRegionIdSet.has(nextRegionId)
+    )
+  }
+
   override onAllRoutesRouted() {
     const { state } = this
     const maxRips = Math.min(this.MAX_RIPS, this.RIP_THRESHOLD_RAMP_ATTEMPTS)
@@ -944,7 +960,8 @@ export class TinyHyperGraphSectionSolver extends BaseSolver {
     this.stats = {
       ...this.stats,
       sectionBaselineMaxRegionCost: this.sectionBaselineSummary.maxRegionCost,
-      sectionBaselineTotalRegionCost: this.sectionBaselineSummary.totalRegionCost,
+      sectionBaselineTotalRegionCost:
+        this.sectionBaselineSummary.totalRegionCost,
       effectiveRipThresholdStart: this.RIP_THRESHOLD_START,
       effectiveRipThresholdEnd: this.RIP_THRESHOLD_END,
       effectiveMaxRips: this.MAX_RIPS,

package.json

@@ -3,6 +3,7 @@
   "module": "index.ts",
   "type": "module",
   "scripts": {
+    "build": "cosmos-export",
     "start": "cosmos",
     "typecheck": "tsc -p tsconfig.json --pretty false",
     "benchmark1": "bun run --cpu-prof-md scripts/profiling/hg07-first10.ts",
@@ -15,10 +16,12 @@
     "@tscircuit/hypergraph": "^0.0.71",
     "@tscircuit/solver-utils": "^0.0.17",
     "@types/bun": "latest",
-    "cosmos": "^0.1.2",
     "dataset-hg07": "https://github.com/tscircuit/dataset-hg07#6adb2ed998a16675b11cf59c25789ae95f0e1a6f",
     "graphics-debug": "^0.0.90",
+    "react": "^19.2.0",
+    "react-cosmos": "^7.2.0",
     "react-cosmos-plugin-vite": "^7.2.0",
+    "react-dom": "^19.2.0",
     "stack-svgs": "^0.0.1",
     "vite": "^8.0.1"
   },