Fix source propagation and offset

Author: rafaelha

src/kirin/ir/nodes/stmt.py

@@ -267,6 +267,9 @@ def insert_after(self, stmt: Statement) -> None:
         self.parent = stmt.parent
         stmt._next_stmt = self
 
+        if self.source is None and stmt.source is not None:
+            self.source = stmt.source
+
         if self.parent:
             self.parent._stmt_len += 1
 
@@ -302,6 +305,9 @@ def insert_before(self, stmt: Statement) -> None:
         self.parent = stmt.parent
         stmt._prev_stmt = self
 
+        if self.source is None and stmt.source is not None:
+            self.source = stmt.source
+
         if self.parent:
             self.parent._stmt_len += 1
 
@@ -506,6 +512,7 @@ def from_stmt(
             attributes=attributes or other.attributes,
             result_types=[result.type for result in other._results],
             args_slice=other._name_args_slice,
+            source=other.source,
         )
         # inherit the hint:
         for result, other_result in zip(obj._results, other._results):

src/kirin/lowering/frame.py

@@ -75,9 +75,9 @@ def _push_stmt(self, stmt: StmtType) -> StmtType:
             raise BuildError(
                 f"Unsupported dialect `{stmt.dialect.name}` from statement {stmt.name}"
             )
-        self.curr_block.stmts.append(stmt)
         if stmt.source is None:
             stmt.source = self.state.source
+        self.curr_block.stmts.append(stmt)
         return stmt
 
     def _push_block(self, block: Block):

src/kirin/lowering/python/lowering.py

@@ -138,6 +138,7 @@ def lower_global(self, state: State[ast.AST], node: ast.AST) -> LoweringABC.Resu
     def visit(self, state: State[ast.AST], node: ast.AST) -> Result:
         if hasattr(node, "lineno"):
             state.source = SourceInfo.from_ast(node, state.file)
+            state.source.offset(state.lineno_offset - 1, state.col_offset)
         name = node.__class__.__name__
         if name in self.registry.ast_table:
             return self.registry.ast_table[name].lower(state, node)
@@ -148,7 +149,8 @@ def generic_visit(self, state: State[ast.AST], node: ast.AST) -> Result:
 
     def visit_Call(self, state: State[ast.AST], node: ast.Call) -> Result:
         if hasattr(node.func, "lineno"):
-            state.source = SourceInfo.from_ast(node.func, state.file)
+            state.source = SourceInfo.from_ast(node, state.file)
+            state.source.offset(state.lineno_offset - 1, state.col_offset)
 
         global_callee_result = state.get_global(node.func, no_raise=True)
         if global_callee_result is None:

src/kirin/rewrite/inline.py

@@ -97,6 +97,16 @@ def inline_call_like(
 
         # NOTE: we cannot change region because it may be used elsewhere
         inline_region: ir.Region = region.clone()
+
+        # Preserve source information by attributing inlined code to the call site
+        if call_like.source is not None:
+            for block in inline_region.blocks:
+                if block.source is None:
+                    block.source = call_like.source
+                for stmt in block.stmts:
+                    if stmt.source is None:
+                        stmt.source = call_like.source
+
         parent_block: ir.Block = call_like.parent_block
         parent_region: ir.Region = call_like.parent_region