change return type of map to tuple instead of list

Author: kaihsin

src/kirin/dialects/fcf/interp.py

@@ -36,7 +36,7 @@ def foldr(self, interp: Interpreter, frame: Frame, stmt: Foldr):
         return (acc,)
 
     @impl(Map)
-    def map_list(self, interp: Interpreter, frame: Frame, stmt: Map):
+    def map_tuple(self, interp: Interpreter, frame: Frame, stmt: Map):
         fn: ir.Method = frame.get(stmt.fn)
         coll = frame.get(stmt.coll)
         ret = []
@@ -47,7 +47,7 @@ def map_list(self, interp: Interpreter, frame: Frame, stmt: Map):
                 return _ret
             else:
                 ret.append(_ret)
-        return (ret,)
+        return (tuple(ret),)
 
     @impl(Scan)
     def scan(self, interp: Interpreter, frame: Frame, stmt: Scan):

src/kirin/dialects/fcf/rewrite/fcfmap_inline.py

@@ -4,7 +4,6 @@
 from kirin import ir
 from kirin.analysis import const
 from kirin.dialects import fcf, func
-from kirin.exceptions import DialectInterpretationError
 from kirin.dialects.py import stmts as py
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 from kirin.ir.nodes.stmt import Statement
@@ -14,13 +13,6 @@
 class InlineFcfMap(RewriteRule):
     cp_results: Dict[ir.SSAValue, const.JointResult]
 
-    def get_const_value(self, p: ir.SSAValue):
-        tmp = self.cp_results.get(p, None)
-        if (tmp is None) or (not isinstance(tmp.const, const.Value)):
-            raise DialectInterpretationError(f"not a const value: {p}")
-
-        return tmp.const.data
-
     def rewrite_Statement(self, node: Statement) -> RewriteResult:
         match node:
             case fcf.Map():
@@ -30,7 +22,12 @@ def rewrite_Statement(self, node: Statement) -> RewriteResult:
 
     def rewrite_fcf_map(self, node: fcf.Map) -> RewriteResult:
         # TODO make this more generic without the need for the constprop results
-        coll = self.get_const_value(node.coll)
+        tmp = self.cp_results.get(node.coll, None)
+
+        if (tmp is None) or (not isinstance(tmp.const, const.Value)):
+            return RewriteResult()
+
+        coll = tmp.const.data
 
         # rewrite to directly inline:
         # get the method:

src/kirin/dialects/fcf/stmts.py

@@ -23,7 +23,7 @@ class Foldr(ir.Statement):
 class Map(ir.Statement):
     fn: ir.SSAValue = info.argument(ir.types.PyClass(ir.Method))
     coll: ir.SSAValue = info.argument(ir.types.Any)
-    result: ir.ResultValue = info.result(ir.types.List)
+    result: ir.ResultValue = info.result(ir.types.Tuple)
 
 
 @statement(dialect=dialect)

test/dialects/fcf/test_fold.py

@@ -43,6 +43,6 @@ def test_fold():
     assert foldl(xs) == sum(xs)
     assert foldr(xs) == sum(xs)
     assert map_list.return_type.is_subseteq(types.List[types.Float])
-    assert map_list([1, 2, 3]) == [2.0, 3.0, 4.0]
+    assert map_list([1, 2, 3]) == (2.0, 3.0, 4.0)
     assert scan([1, 2, 3, 4, 5]) == (15, [1, 2, 3, 4, 5])
     assert scan.return_type.is_subseteq(types.Tuple[types.Int, types.List[types.Int]])

test/dialects/fcf/test_map.py

@@ -11,4 +11,4 @@ def map_func(i):
 
 
 def test_enumerate_kirin():
-    assert enumerate_kirin([1, 2, 3, 4, 5]) == [(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)]
+    assert enumerate_kirin([1, 2, 3, 4, 5]) == ((0, 1), (1, 2), (2, 3), (3, 4), (4, 5))

test/dialects/pystmts/test_range.py

@@ -44,4 +44,4 @@ def map_range(x: range):
 
 def test_map_range():
     assert map_range.return_type.is_subseteq(types.List[types.Float])
-    assert map_range(range(10)) == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
+    assert map_range(range(10)) == (1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0)

test/rules/test_fcfmap_inline.py

@@ -1,5 +1,6 @@
 from kirin.prelude import basic
 from kirin.rewrite import Walk
+from kirin.analysis import const
 from kirin.dialects import fcf
 from kirin.analysis.const.prop import Propagate
 from kirin.dialects.fcf.rewrite.fcfmap_inline import InlineFcfMap
@@ -25,3 +26,28 @@ def _simple(i: int):
     val = fcf_map_rewrite()
 
     assert val == (0, 1, 2, 3, 4)
+
+
+def test_fcfmap_rewrite_with_arg():
+
+    @basic(fold=False)
+    def fcf_map_rewrite_with_arg(x: int):
+
+        def _simple(i: int):
+            return i
+
+        tmp = fcf.Map(_simple, (x, x + 1, x + 2))
+        return tmp
+
+    fcf_map_rewrite_with_arg.code.print()
+    cp = Propagate(dialects=fcf_map_rewrite_with_arg.dialects)
+    cp.eval(fcf_map_rewrite_with_arg, (const.JointResult.top(),))
+    Walk(InlineFcfMap(cp.results)).rewrite(fcf_map_rewrite_with_arg.code)
+    fcf_map_rewrite_with_arg.code.print()
+
+    val = fcf_map_rewrite_with_arg(4)
+    print(val)
+    assert val == (4, 5, 6)
+
+
+test_fcfmap_rewrite_with_arg()