diff --git a/src/pystencils/backend/transformations/canonicalize_symbols.py b/src/pystencils/backend/transformations/canonicalize_symbols.py index 17875caa17d470484d6d509fdd023a1d7867ba1f..eb1aa1851ebedf801cd51bb24c054f7acf1c2895 100644 --- a/src/pystencils/backend/transformations/canonicalize_symbols.py +++ b/src/pystencils/backend/transformations/canonicalize_symbols.py @@ -3,7 +3,15 @@ from ..memory import PsSymbol from ..exceptions import PsInternalCompilerError from ..ast import PsAstNode -from ..ast.structural import PsDeclaration, PsAssignment, PsLoop, PsConditional, PsBlock +from ..ast.structural import ( + PsDeclaration, + PsAssignment, + PsLoop, + PsConditional, + PsBlock, + PsStatement, + PsEmptyLeafMixIn, +) from ..ast.expressions import PsSymbolExpr, PsExpression from ...types import constify @@ -117,3 +125,14 @@ class CanonicalizeSymbols: case PsBlock(statements): for stmt in statements[::-1]: self.visit(stmt, cc) + + case PsStatement(expr): + self.visit(expr, cc) + + case PsEmptyLeafMixIn(): + ... + + case unknown: + raise PsInternalCompilerError( + f"Can't canonicalize symbols in {unknown} ({repr(unknown)})." + )