fix: fix infinite loops

Author: DiogoRibeiro7

src/ssspx/frontier.py

@@ -148,9 +148,15 @@ def _consume_block_prefix(
         """Greedily take up to ``want`` keys from the head of ``blocks``."""
         got = 0
         idx_block = 0
-        while got < want and idx_block < len(blocks):
+        iterations = 0
+        max_iterations = len(blocks) * 100  # Safety limit to prevent infinite loops
+        
+        while got < want and idx_block < len(blocks) and iterations < max_iterations:
+            iterations += 1
             block = blocks[idx_block]
             keep: List[Tuple[Vertex, Float]] = []
+            processed_any = False
+            
             for k, v in block:
                 bestv = self._best.get(k)
                 if bestv is None or v != bestv:
@@ -162,13 +168,21 @@ def _consume_block_prefix(
                     chosen[k] = v
                     pulled_keys.add(k)
                     got += 1
+                    processed_any = True
                 else:
                     keep.append((k, v))
+            
             blocks[idx_block] = keep
             if not blocks[idx_block]:
                 blocks.pop(idx_block)
+                # Don't increment idx_block since we removed an element
             else:
                 idx_block += 1
+                
+            # If we didn't process anything useful, break to avoid infinite loop
+            if not processed_any and not keep:
+                break
+                
         return got
 
     def pull(self) -> Tuple[Set[Vertex], Float]:
@@ -258,11 +272,16 @@ def batch_prepend(self, pairs: Iterable[Tuple[Vertex, Float]]) -> None:
     def pull(self) -> Tuple[Set[Vertex], Float]:
         """Return up to ``M`` keys with the smallest values."""
         s: Set[Vertex] = set()
-        while self._heap and len(s) < self.M:
+        iterations = 0
+        max_iterations = len(self._heap) * 2  # Safety limit
+        
+        while self._heap and len(s) < self.M and iterations < max_iterations:
+            iterations += 1
             val, key = heapq.heappop(self._heap)
             if self._best.get(key) != val:
                 continue  # stale
             s.add(key)
+            
         if not s:
             return set(), self.B
         x = self._heap[0][0] if self._heap else self.B

src/ssspx/solver.py

@@ -99,6 +99,7 @@ def __init__(
             "pulls": 0,
             "findpivots_rounds": 0,
             "basecase_pops": 0,
+            "iterations_protected": 0,  # Track safety limit hits
         }
 
         # Optional transform for outdegree
@@ -133,19 +134,22 @@ def __init__(
             self.complete[s] = True
             self.root[s] = s
 
-        # Parameters (k, t, levels)
+        # Parameters (k, t, levels) with safety bounds
         n = max(2, self.G.n)
         if self.cfg.k_t_auto:
             log2n = math.log2(n)
             k = max(1, int(round(log2n ** (1.0 / 3.0))))
             t = max(1, int(round(log2n ** (2.0 / 3.0))))
             k = max(1, min(k, t))
+            # Cap to reasonable values to prevent runaway algorithms
+            k = min(k, 100)
+            t = min(t, 20)
         else:
-            k = max(1, self.cfg.k)
-            t = max(1, self.cfg.t)
+            k = max(1, min(self.cfg.k, 100))  # Safety cap
+            t = max(1, min(self.cfg.t, 20))   # Safety cap
         self.k: int = k
         self.t: int = t
-        self.L: int = max(1, math.ceil(math.log2(n) / t))
+        self.L: int = max(1, min(math.ceil(math.log2(n) / t), 10))  # Cap levels
 
         # Best-clone cache for each original vertex after solve()
         self._best_clone_for_orig: Optional[List[int]] = None
@@ -211,22 +215,32 @@ def _base_case(self, B: Float, S: Set[Vertex]) -> Tuple[Float, Set[Vertex]]:
         seen: Set[Vertex] = set()
         heap: List[Tuple[Float, Vertex]] = [(self.dhat[x], x)]
         in_heap: Set[Vertex] = {x}
+        
+        # Safety limits to prevent infinite loops
+        iterations = 0
+        max_iterations = min(self.k * 1000, self.G.n * 10)
 
-        while heap and len(U0) < self.k + 1:
+        while heap and len(U0) < self.k + 1 and iterations < max_iterations:
             self.counters["basecase_pops"] += 1
+            iterations += 1
+            
             du, u = heapq.heappop(heap)
             in_heap.discard(u)
             if du != self.dhat[u] or u in seen:
                 continue
             seen.add(u)
             self.complete[u] = True
             U0.append(u)
+            
             for v, w in self.G.adj[u]:
                 if self._relax(u, v, w) and self.dhat[u] + w < B:
                     if v not in in_heap:
                         heapq.heappush(heap, (self.dhat[v], v))
                         in_heap.add(v)
 
+        if iterations >= max_iterations:
+            self.counters["iterations_protected"] += 1
+
         if len(U0) <= self.k:
             return (B, set(U0))
         Bprime = max(self.dhat[v] for v in U0)
@@ -248,20 +262,38 @@ def _find_pivots(self, B: Float, S: Set[Vertex]) -> Tuple[Set[Vertex], Set[Verte
         """
         W: Set[Vertex] = set(S)
         current: Set[Vertex] = set(S)
-        for _ in range(1, self.k + 1):
+        
+        # Safety limits for findpivots
+        iterations = 0
+        max_iterations = min(self.k * len(S) * 100, self.G.n * 10)
+        
+        for round_num in range(1, self.k + 1):
+            if iterations >= max_iterations:
+                self.counters["iterations_protected"] += 1
+                break
+                
             self.counters["findpivots_rounds"] += 1
             nxt: Set[Vertex] = set()
+            
             for u in current:
+                iterations += 1
+                if iterations >= max_iterations:
+                    break
+                    
                 for v, w in self.G.adj[u]:
                     if self._relax(u, v, w) and (self.dhat[u] + w < B):
                         nxt.add(v)
+                        
             if not nxt:
                 break
             W |= nxt
-            if len(W) > self.k * max(1, len(S)):
+            
+            # Early termination if W gets too large
+            if len(W) > self.k * max(1, len(S)) * 5:  # More generous limit
                 return set(S), W
             current = nxt
 
+        # Build pivot tree with safety limits
         children: Dict[Vertex, List[Vertex]] = {u: [] for u in W}
         for v in W:
             p = self.pred[v]
@@ -273,7 +305,11 @@ def _find_pivots(self, B: Float, S: Set[Vertex]) -> Tuple[Set[Vertex], Set[Verte
             size = 0
             stack = [u]
             seen: Set[Vertex] = set()
-            while stack:
+            iterations = 0
+            max_tree_iterations = min(self.k * 10, len(W))
+            
+            while stack and iterations < max_tree_iterations:
+                iterations += 1
                 a = stack.pop()
                 if a in seen:
                     continue
@@ -288,14 +324,19 @@ def _find_pivots(self, B: Float, S: Set[Vertex]) -> Tuple[Set[Vertex], Set[Verte
     # ---------- BMSSP -----------------------------------------------------
 
     def _make_frontier(self, level: int, B: Float) -> FrontierProtocol:
-        M = max(1, 2 ** ((level - 1) * self.t))
+        # Cap the frontier size to prevent excessive memory usage
+        M = max(1, min(2 ** ((level - 1) * self.t), 10000))
         if self.cfg.frontier == "heap":
             return HeapFrontier(M=M, B=B)
         if self.cfg.frontier == "block":
             return BlockFrontier(M=M, B=B)
         raise ConfigError(f"unknown frontier '{self.cfg.frontier}'")
 
-    def _bmssp(self, level: int, B: Float, S: Set[Vertex]) -> Tuple[Float, Set[Vertex]]:
+    def _bmssp(self, level: int, B: Float, S: Set[Vertex], depth: int = 0) -> Tuple[Float, Set[Vertex]]:
+        # Prevent excessive recursion
+        if depth > 50 or level > 20:
+            return self._base_case(B, S)
+            
         if level == 0:
             return self._base_case(B, S)
 
@@ -305,10 +346,14 @@ def _bmssp(self, level: int, B: Float, S: Set[Vertex]) -> Tuple[Float, Set[Verte
             D.insert(x, self.dhat[x])
 
         U_accum: Set[Vertex] = set()
-        cap = self.k * max(1, 2 ** (level * self.t))
+        cap = min(self.k * max(1, 2 ** (level * self.t)), self.G.n)  # Cap to graph size
+        pull_iterations = 0
+        max_pull_iterations = min(cap * 10, 1000)  # Safety limit on pulls
 
-        while len(U_accum) < cap:
+        while len(U_accum) < cap and pull_iterations < max_pull_iterations:
             self.counters["pulls"] += 1
+            pull_iterations += 1
+            
             S_i, B_i = D.pull()
             if not S_i:
                 Bprime = B
@@ -317,7 +362,7 @@ def _bmssp(self, level: int, B: Float, S: Set[Vertex]) -> Tuple[Float, Set[Verte
                     self.complete[u] = True
                 return Bprime, U_accum
 
-            B_i_prime, U_i = self._bmssp(level - 1, B_i, S_i)
+            B_i_prime, U_i = self._bmssp(level - 1, B_i, S_i, depth + 1)
             for u in U_i:
                 self.complete[u] = True
             U_accum |= U_i
@@ -344,6 +389,9 @@ def _bmssp(self, level: int, B: Float, S: Set[Vertex]) -> Tuple[Float, Set[Verte
                     self.complete[u] = True
                 return Bprime, U_accum
 
+        if pull_iterations >= max_pull_iterations:
+            self.counters["iterations_protected"] += 1
+
         return B, U_accum
 
     # ---------- public API ------------------------------------------------
@@ -413,11 +461,15 @@ def path(self, target_original: Vertex) -> List[Vertex]:
             return []  # unreachable
         src_clone = self.root[start_clone]
 
-        # Walk predecessors in clone-space
+        # Walk predecessors in clone-space with safety limits
         chain: List[int] = []
         cur: Optional[int] = start_clone
         seen = set()
-        while cur is not None:
+        iterations = 0
+        max_iterations = self.G.n * 2  # Safety limit
+        
+        while cur is not None and iterations < max_iterations:
+            iterations += 1
             chain.append(cur)
             if cur == src_clone:
                 chain.reverse()