Fix loop handling in CFG traversal (#1116)

* Fixes bug in one-iteration-checker

CHANGELOG.md

@@ -32,6 +32,8 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - Fixed issue where empty preconditions were preventing CFGs from being generated
 - Added strict numeric type checking to enforce type distinctions across the entire numeric hierarchy, including complex numbers.
 - Added strict type checking support for nested and union types (e.g., `list[int]`, `dict[float, int]`, `Union[int, float]`)
+- Fixed issue where CFG edges from loop body to loop condition block was ignored during augmenting edge z3 constraints
+- Fixed issue in `one-iteration-checker` where the message was not correctly reported for `while` loops when `z3` option is on
 
 ### 🔧 Internal changes
 

python_ta/cfg/graph.py

@@ -225,28 +225,23 @@ def get_paths(self) -> List[List[CFGEdge]]:
         """Get edges that represent paths from start to end node in depth-first order."""
         paths = []
 
-        def _visited(
-            edge: CFGEdge, visited_edges: Set[CFGEdge], visited_nodes: Set[CFGBlock]
-        ) -> bool:
-            return edge in visited_edges or edge.target in visited_nodes
-
         def _dfs(
             current_edge: CFGEdge,
             current_path: List[CFGEdge],
             visited_edges: Set[CFGEdge],
             visited_nodes: Set[CFGBlock],
         ):
-            # note: both visited edges and visited nodes need to be tracked to correctly handle cycles
-            if _visited(current_edge, visited_edges, visited_nodes):
+            if current_edge in visited_edges:
                 return
 
             visited_edges.add(current_edge)
-            visited_nodes.add(current_edge.source)
             current_path.append(current_edge)
+            visited_nodes.add(current_edge.source)
 
-            if current_edge.target == self.end or all(
-                _visited(edge, visited_edges, visited_nodes)
-                for edge in current_edge.target.successors
+            if (
+                current_edge.target == self.end
+                or current_edge.target in visited_nodes
+                or set(current_edge.target.successors).issubset(visited_edges)
             ):
                 paths.append(current_path.copy())
             else:

python_ta/checkers/one_iteration_checker.py

@@ -74,8 +74,8 @@ def _check_one_iteration(self, node: Union[nodes.For, nodes.While]) -> bool:
                 for succ in start.cfg_block.successors
             ):
                 return False
-            # filter out edges from unfeasible blocks
-            preds = [pred for pred in preds if pred.source.is_feasible]
+            # filter out unfeasible edges
+            preds = [pred for pred in preds if pred.is_feasible]
 
         for pred in preds:
             stmt = pred.source.statements[0]

tests/test_cfg/test_edge_feasibility.py

@@ -71,6 +71,31 @@ def func(x: str, y: int) -> None:
 
 
 def test_unfeasible_while_condition() -> None:
+    src = """
+    def func(x: int, y: int) -> None:
+        '''
+        Preconditions:
+            - x > 10
+            - y > 10
+        '''
+        while x + y < 15:
+            print("unreachable")
+        print("end")
+    """
+    cfg = _create_cfg(src, "func")
+    expected_while_path = [True, False, False]
+    expected_other_path = [True, True, True]
+
+    paths = cfg.get_paths()
+    assert all(
+        edge.is_feasible == expected for edge, expected in zip(paths[0], expected_while_path)
+    )
+    assert all(
+        edge.is_feasible == expected for edge, expected in zip(paths[1], expected_other_path)
+    )
+
+
+def test_unfeasible_while_condition_with_reassignment() -> None:
     src = """
     def func(x: int, y: int) -> None:
         '''
@@ -85,7 +110,7 @@ def func(x: int, y: int) -> None:
         print("end")
     """
     cfg = _create_cfg(src, "func")
-    expected_while_path = [True, False]
+    expected_while_path = [True, False, True]
     expected_other_path = [True, True, True]
 
     paths = cfg.get_paths()
@@ -226,8 +251,40 @@ def func(x: int, condition: bool) -> None:
 
     paths = cfg.get_paths()
     assert all(edge.is_feasible for edge in paths[0])
+    assert all(edge.is_feasible for edge in paths[1])
     assert all(
-        edge.is_feasible == expected for edge, expected in zip(paths[1], expected_other_path)
+        edge.is_feasible == expected for edge, expected in zip(paths[2], expected_other_path)
+    )
+
+
+def test_unfeasible_if_in_while() -> None:
+    src = """
+    def func(x: int) -> None:
+        '''
+        Preconditions:
+            - x in (1, 2, 3, 4, 5)
+        '''
+        while x > 2:
+            print(x)
+            if x > 10:
+                break
+        print("end")
+    """
+    cfg = _create_cfg(src, "func")
+
+    expected_break_path = [True, True, False, False, True]
+    expected_not_break_path = [True, True, True]
+    expected_other_path = [True, True, True]
+
+    paths = cfg.get_paths()
+    assert all(
+        edge.is_feasible == expected for edge, expected in zip(paths[0], expected_break_path)
+    )
+    assert all(
+        edge.is_feasible == expected for edge, expected in zip(paths[1], expected_not_break_path)
+    )
+    assert all(
+        edge.is_feasible == expected for edge, expected in zip(paths[2], expected_other_path)
     )
 
 

tests/test_cfg/test_z3_constraints.py

@@ -154,7 +154,59 @@ def func(x: int, y: int) -> None:
     cfg = _create_cfg(src, "func")
     x = z3.Int("x")
     y = z3.Int("y")
-    expected_while_true_path = [{x > 5, y > 10}, {x > 5, y > 10, x + y > 15}]
+    expected_while_true_path = [
+        {x > 5, y > 10},
+        {x > 5, y > 10, x + y > 15},
+        set(),
+    ]
+    expected_while_false_path = [
+        {x > 5, y > 10},
+        {x > 5, y > 10, z3.Not(x + y > 15)},
+        {x > 5, y > 10, z3.Not(x + y > 15)},
+    ]
+
+    actual_path_first = []
+    actual_path_second = []
+    for edge in cfg.get_edges():
+        actual1 = edge.z3_constraints.get(0)
+        actual2 = edge.z3_constraints.get(1)
+        if actual1 is not None:
+            actual_path_first.append(actual1)
+        if actual2 is not None:
+            actual_path_second.append(actual2)
+
+    assert len(actual_path_first) == len(expected_while_true_path)
+    assert len(actual_path_second) == len(expected_while_false_path)
+    assert (
+        set(actual) == expected
+        for actual, expected in zip(actual_path_first, expected_while_true_path)
+    )
+    assert (
+        set(actual) == expected
+        for actual, expected in zip(actual_path_second, expected_while_false_path)
+    )
+
+
+def test_while_loop_infinite() -> None:
+    src = """
+    def func(x: int, y: int) -> None:
+        '''
+        Preconditions:
+            - x > 5
+            - y > 10
+        '''
+        while x + y > 15:
+            print(x + y)
+        print(x + y)
+    """
+    cfg = _create_cfg(src, "func")
+    x = z3.Int("x")
+    y = z3.Int("y")
+    expected_while_true_path = [
+        {x > 5, y > 10},
+        {x > 5, y > 10, x + y > 15},
+        {x > 5, y > 10, x + y > 15},
+    ]
     expected_while_false_path = [
         {x > 5, y > 10},
         {x > 5, y > 10, z3.Not(x + y > 15)},
@@ -294,11 +346,13 @@ def func(x: int, y: int) -> None:
         {x > 10, y > 10},
         {x > 10, y > 10, x > 0},
         {x > 10, y > 10, x > 0, y > 5},
+        set(),
     ]
     expected_outer_while_path = [
         {x > 10, y > 10},
         {x > 10, y > 10, x > 0},
         {x > 10, y > 10, x > 0, z3.Not(y > 5)},
+        {x > 10, y > 10, x > 0, z3.Not(y > 5)},
     ]
     expected_other_path = [
         {x > 10, y > 10},
@@ -365,6 +419,7 @@ def func(x: int, y: int, condition: bool) -> None:
         {x > y, condition},
         {x > y, condition, condition},
         {x > y, condition, condition, z3.Not(x < y)},
+        {x > y, condition, condition, z3.Not(x < y)},
     ]
     expected_other_path = [
         {x > y, condition},
@@ -423,11 +478,13 @@ def func(x: int, y: int) -> None:
         {x > 10, y > 0},
         {x > 10, y > 0, x > 0},
         {x > 10, y > 0, x > 0, x < y},
+        {x > 10, y > 0, x > 0, x < y},
     ]
     expected_not_continue_path = [
         {x > 10, y > 0},
         {x > 10, y > 0, x > 0},
         {x > 10, y > 0, x > 0, z3.Not(x < y)},
+        set(),
     ]
     expected_other_path = [
         {x > 10, y > 0},
@@ -534,8 +591,8 @@ def func(x: float) -> None:
     expected_if_path = [
         {z3.Or(x == 1.0, x == 2.0, x == 3.0)},
         {z3.Or(x == 1.0, x == 2.0, x == 3.0), x < 5},
-        {},
-        {},
+        set(),
+        set(),
     ]
     expected_else_path = [
         {z3.Or(x == 1.0, x == 2.0, x == 3.0)},

tests/test_custom_checkers/test_one_iteration_checker.py

@@ -473,6 +473,32 @@ def func(x: int) -> int:
         ):
             self.checker.visit_while(while_node)
 
+    def test_z3_one_iteration_break_by_precondition_no_loop_body(self):
+        src = """
+        def func(x: int) -> int:
+            '''
+            Preconditions:
+                - x > 5
+            '''
+            while x > 0:
+                if x > 3:
+                    break
+            return x
+        """
+        z3v = Z3Visitor()
+        mod = z3v.visitor.visit(astroid.parse(src))
+        mod.accept(CFGVisitor())
+        while_node = next(mod.nodes_of_class(nodes.While))
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="one-iteration",
+                node=while_node,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_while(while_node)
+
     def test_z3_multiple_iterations(self):
         src = """
         def func(x: int, y: bool) -> int: