Add unit tests

Author: rafaelha

test/ir/test_stmt.py

@@ -1,6 +1,7 @@
 import pytest
 
 from kirin.ir import Block
+from kirin.source import SourceInfo
 from kirin.dialects import py
 
 
@@ -50,3 +51,12 @@ def test_stmt_from_stmt():
     y = x.from_stmt(x)
 
     assert y.result.hints["const"] == py.constant.types.Int
+
+
+def test_stmt_from_stmt_preserves_source_info():
+    x = py.Constant(1)
+    x.source = SourceInfo(lineno=1, col_offset=0, end_lineno=None, end_col_offset=None)
+
+    y = x.from_stmt(x)
+    assert y.source == x.source
+    assert y.source is x.source

test/lowering/test_source_info.py

@@ -0,0 +1,36 @@
+import pytest
+
+from kirin.source import SourceInfo
+from kirin.prelude import basic_no_opt
+
+
+def get_line_of(target: str) -> int:
+    for i, line in enumerate(open(__file__), 1):
+        if target in line:
+            return i
+
+
+@pytest.mark.parametrize("similar", [True, False])
+def test_stmt_source_info(similar: bool):
+    @basic_no_opt
+    def test(x: int):
+        y = 2
+        a = 4**2
+        return y + 2 + a
+
+    if similar:
+        test = test.similar()
+
+    stmts = test.callable_region.blocks[0].stmts
+
+    def get_line_from_source_info(source: SourceInfo) -> int:
+        return source.lineno + source.lineno_begin
+
+    for stmt in stmts:
+        assert stmt.source.file == __file__
+
+    assert get_line_from_source_info(stmts.at(0).source) == get_line_of("y = 2")
+    assert get_line_from_source_info(stmts.at(2).source) == get_line_of("a = 4**2")
+    assert get_line_from_source_info(stmts.at(4).source) == get_line_of(
+        "return y + 2 + a"
+    )

test/passes/test_inline_pass.py

@@ -1,5 +1,6 @@
 from kirin.prelude import basic_no_opt
 from kirin.passes.inline import InlinePass
+from kirin.dialects.py.constant import Constant
 
 
 @basic_no_opt
@@ -40,3 +41,25 @@ def main_inline_pass2(x: int):
     assert a == b
 
     assert len(main_inline_pass2.callable_region.blocks[0].stmts) == 4
+
+
+def test_inline_preserves_source_info():
+    def get_line_of(target: str) -> int:
+        for i, line in enumerate(open(__file__), 1):
+            if target in line:
+                return i
+
+    @basic_no_opt
+    def main_inline_pass(x: int):
+        y = inline_func(x)
+        return y + 2
+
+    inline = InlinePass(main_inline_pass.dialects)
+    inline(main_inline_pass)
+
+    stmt = main_inline_pass.callable_region.blocks[0].stmts.at(0)
+    line = stmt.source.lineno + stmt.source.lineno_begin
+    assert stmt.value.data == 1
+    assert isinstance(stmt, Constant)
+
+    assert get_line_of("return x - 1") == line