Add support for PsConditional to UndefinedSymbolsCollector and PsStatement to CanonicalClone

src/pystencils/backend/ast/analysis.py

@@ -6,6 +6,7 @@ from .structural import (
     PsAstNode,
     PsBlock,
     PsComment,
+    PsConditional,
     PsDeclaration,
     PsExpression,
     PsLoop,
@@ -56,6 +57,12 @@ class UndefinedSymbolsCollector:
                 undefined_vars.discard(ctr.symbol)
                 return undefined_vars
 
+            case PsConditional(cond, branch_true, branch_false):
+                undefined_vars = self(cond) | self(branch_true)
+                if branch_false is not None:
+                    undefined_vars |= self(branch_false)
+                return undefined_vars
+
             case PsComment():
                 return set()
 
@@ -86,6 +93,7 @@ class UndefinedSymbolsCollector:
                 PsAssignment()
                 | PsBlock()
                 | PsComment()
+                | PsConditional()
                 | PsExpression()
                 | PsLoop()
                 | PsStatement()

src/pystencils/backend/transformations/canonical_clone.py

@@ -12,6 +12,7 @@ from ..ast.structural import (
     PsDeclaration,
     PsAssignment,
     PsComment,
+    PsStatement,
 )
 from ..ast.expressions import PsExpression, PsSymbolExpr
 
@@ -99,6 +100,9 @@ class CanonicalClone:
                 self._replace_symbols(expr_clone, cc)
                 return cast(Node_T, expr_clone)
 
+            case PsStatement(expr):
+                return cast(Node_T, PsStatement(self.visit(expr, cc)))
+
             case _:
                 raise PsInternalCompilerError(
                     f"Don't know how to canonically clone {type(node)}"