Merge pull request #31 from Balint-R/macro-default-arg

Add CallMacro default argument testcase

Author: Balint-R

src/pydsl/compiler.py

@@ -358,6 +358,10 @@ def visit_Expr(self, node: ast.Expr) -> SubtreeOut:
             case _:
                 self.visit(expr_val)
 
+    def visit_Expression(self, node: ast.Expression) -> SubtreeOut:
+        # ast.Expression is output by ast.parse(mode="eval")
+        return self.visit(node.body)
+
     def visit_Constant(self, node: ast.Constant) -> SubtreeOut:
         # TODO: Constant may not always be Number. It may also be e.g. string.
         # When other forms of constants are supported this needs to be updated.

tests/e2e/test_macro.py

@@ -182,6 +182,23 @@ def g(x: UInt32) -> UInt32:
     assert g(0) == 0 + 0
 
 
+def test_default_args():
+    ast_123 = ast.parse("UInt32(123)", mode="eval")
+
+    @CallMacro.generate()
+    def add_macro(
+        visitor: ToMLIRBase, x: Uncompiled = ast_123, y: Evaluated = 456
+    ) -> UInt32:
+        x = visitor.visit(x)
+        return x.op_add(y)
+
+    @compile()
+    def f() -> UInt32:
+        return add_macro()
+
+    assert f() == 123 + 456
+
+
 if __name__ == "__main__":
     run(test_Compiled_ArgCompiler)
     run(test_Evaluated_ArgCompiler)
@@ -193,3 +210,4 @@ def g(x: UInt32) -> UInt32:
     run(test_class_method)
     run(test_class_only_method)
     run(test_static_method)
+    run(test_default_args)