RFC: support overwriting existing method (#421)

this PR add support of overwriting existing method by directly updating
the existing method object in (current Python not Kirin) scope if the
name matches. We limit the compiler not to update existing method
because this may cause unexpected behaviour because we won't actually
update existing compiled code. However, after thinking a bit more, I
think this is still ok because we can keep the same method object but
update its body.

The caveat is we need to retrigger the whole compilation of call the
caller of this method, which is probably still ok? I didn't implement
the [entire world age mechanism in
Julia](https://d1wqtxts1xzle7.cloudfront.net/73563045/3428275-libre.pdf?1635141917=&response-content-disposition=inline%3B+filename%3DWorld_age_in_Julia_optimizing_method_dis.pdf&Expires=1749052521&Signature=OpvF9l3YO6ZRxHylvBFSMIH0O1Fl6AO2nyg1EZuD-5SWpZZw8E50yy7mgLcGMKhCwTWm6414MhLU9BwolYYnId-ichrBKTkxiDH3MPp050NFtXtmuGQ0Hxnc2wgeDgvKJPiNfVLydlPQLu7daEWw9uB4J2rAzb42m37YBOkfVvdjwvBHg8ysLncdDs-rc15mXKCmvsAYTFl7zRdaUxxhgphNYO~SOdNLB2QWxPpA8SVu3p-HGGitEzyjhtBJi9aQzH0xgBStX6~l2DIyMLEwlhYwjByrw0~moMhKWmOpRELyTfQdYeOxXD39YcztERPOFCE9dTHgzFv0fX-tMaDfkA__&Key-Pair-Id=APKAJLOHF5GGSLRBV4ZA)
but technically every time you define a new method it should trigger an
eval -ish process that recompiles the related methods and thus creates a
new world age. So I'd like to see if there are volunteers willing to
take a second look and play with this PR a bit before moving forward.

## Motivation

currently when supporting interactive environments, such as a notebook,
it is quite cumbersome to use notebook as a dev environment like how one
codes Python natively because we cannot update the code of an existing
`Method`.

Author: Roger-luo

src/kirin/analysis/callgraph.py

@@ -18,42 +18,49 @@ class CallGraph(Printable):
         [`.print()`][kirin.print.printable.Printable.print] method.
     """
 
-    """Mapping from symbol names to methods."""
+    edges: dict[ir.Method, set[ir.Method]] = field(default_factory=dict)
+    """Mapping from symbol names to edges (caller -> callee)."""
     backedges: dict[ir.Method, set[ir.Method]] = field(default_factory=dict)
-    """Mapping from symbol names to backedges."""
+    """Mapping from symbol names to backedges (callee -> caller)."""
 
     def __init__(self, mt: ir.Method):
         self.defs = {}
+        self.edges = {}
         self.backedges = {}
         self.__build(mt)
 
     def __build(self, mt: ir.Method):
         for stmt in mt.callable_region.walk():
             if isinstance(stmt, func.Invoke):
-                backedges = self.backedges.setdefault(stmt.callee, set())
-                backedges.add(mt)
+                edges = self.edges.setdefault(stmt.callee, set())
+                edges.add(mt)
                 self.__build(stmt.callee)
 
+        for caller in self.edges:
+            for callee in self.edges[caller]:
+                backedges = self.backedges.setdefault(callee, set())
+                backedges.add(caller)
+
     def get_neighbors(self, node: ir.Method) -> Iterable[ir.Method]:
         """Get the neighbors of a node in the call graph."""
-        return self.backedges.get(node, ())
+        return self.edges.get(node, ())
 
     def get_edges(self) -> Iterable[tuple[ir.Method, ir.Method]]:
         """Get the edges of the call graph."""
-        for node, neighbors in self.backedges.items():
+        for node, neighbors in self.edges.items():
             for neighbor in neighbors:
                 yield node, neighbor
 
     def get_nodes(self) -> Iterable[ir.Method]:
         """Get the nodes of the call graph."""
-        return self.backedges.keys()
+        return self.edges.keys()
 
     def print_impl(self, printer: Printer) -> None:
-        for idx, (caller, callee) in enumerate(self.backedges.items()):
+        for idx, (caller, callee) in enumerate(self.edges.items()):
             printer.plain_print(caller)
             printer.plain_print(" -> ")
             printer.print_seq(
                 callee, delim=", ", prefix="[", suffix="]", emit=printer.plain_print
             )
-            if idx < len(self.backedges) - 1:
+            if idx < len(self.edges) - 1:
                 printer.print_newline()

src/kirin/dialects/func/stmts.py

@@ -30,6 +30,13 @@ def align_input_args(
         return tuple(inputs)
 
 
+class InvokeCall(ir.StaticCall["Invoke"]):
+
+    @classmethod
+    def get_callee(cls, stmt: Invoke) -> ir.Method:
+        return stmt.callee
+
+
 @statement(dialect=dialect)
 class Function(ir.Statement):
     name = "func"
@@ -270,7 +277,7 @@ def check_type(self) -> None:
 @statement(dialect=dialect)
 class Invoke(ir.Statement):
     name = "invoke"
-    traits = frozenset({ir.MaybePure()})
+    traits = frozenset({ir.MaybePure(), InvokeCall()})
     callee: ir.Method = info.attribute()
     inputs: tuple[ir.SSAValue, ...] = info.argument()
     result: ir.ResultValue = info.result()

src/kirin/ir/__init__.py

@@ -30,6 +30,7 @@
     HasParent as HasParent,
     MaybePure as MaybePure,
     StmtTrait as StmtTrait,
+    StaticCall as StaticCall,
     RegionGraph as RegionGraph,
     SymbolTable as SymbolTable,
     ConstantLike as ConstantLike,

src/kirin/ir/group.py

@@ -200,46 +200,76 @@ def wrapper(py_func: Callable) -> Method:
                 raise ValueError("Cannot compile lambda functions")
 
             lineno_offset, file = 0, ""
+            mt = None
             if frame and frame.f_back is not None:
                 call_site_frame = frame.f_back
                 if py_func.__name__ in call_site_frame.f_locals:
-                    raise CompilerError(
-                        f"overwriting function definition of `{py_func.__name__}`"
-                    )
+                    mt = call_site_frame.f_locals[py_func.__name__]
+                    if not isinstance(mt, Method):
+                        raise CompilerError(
+                            f"`{py_func.__name__}` is already defined in the current scope and is not a Method."
+                        )
 
                 lineno_offset = call_site_frame.f_lineno - 1
                 file = call_site_frame.f_code.co_filename
 
             code = self.lowering.python_function(py_func, lineno_offset=lineno_offset)
             arg_names = ["#self#"] + inspect.getfullargspec(py_func).args
-            mt = Method(
-                dialects=self,
-                code=code,
-                nargs=len(arg_names),
-                mod=inspect.getmodule(py_func),
-                py_func=py_func,
-                sym_name=py_func.__name__,
-                arg_names=arg_names,
-                file=file,
-                lineno_begin=lineno_offset,
-            )
+
+            if mt:
+                mt.mod = inspect.getmodule(py_func)
+                mt.dialects = self
+                mt.code = code
+                mt.py_func = py_func
+                mt.nargs = len(arg_names)
+                mt.arg_names = arg_names
+                mt.sym_name = py_func.__name__
+                mt.file = file
+                mt.lineno_begin = lineno_offset
+                mt.run_passes = self.run_pass
+                mt.update_backedges()  # update the callee
+                self.recompile_callers(mt)
+            else:
+                mt = Method(
+                    dialects=self,
+                    code=code,
+                    nargs=len(arg_names),
+                    mod=inspect.getmodule(py_func),
+                    py_func=py_func,
+                    sym_name=py_func.__name__,
+                    arg_names=arg_names,
+                    file=file,
+                    lineno_begin=lineno_offset,
+                )
+
             if doc := inspect.getdoc(py_func):
                 mt.__doc__ = doc
 
-            if self.run_pass is not None:
-                try:
-                    self.run_pass(mt, *args, **options)
-                except ValidationError as e:
-                    e.attach(mt)
-                    raise e
-
+            def run_pass(mt: Method) -> None:
+                if self.run_pass is not None:
+                    try:
+                        self.run_pass(mt, *args, **options)
+                    except ValidationError as e:
+                        e.attach(mt)
+                        raise e
+
+            mt.run_passes = run_pass
+            run_pass(mt)
             self.update_symbol_table(mt)
             return mt
 
         if py_func is not None:
             return wrapper(py_func)
         return wrapper
 
+    def recompile_callers(self, method: Method) -> None:
+        for caller in method.backedges:
+            if caller.run_passes:
+                caller.run_passes(caller)
+            # propagate the changes to all callers
+            caller.dialects.recompile_callers(caller)
+        return
+
     def update_symbol_table(self, method: Method) -> None:
         trait = method.code.get_trait(SymbolTable)
         if trait is None:

src/kirin/ir/method.py

@@ -10,6 +10,7 @@
 from kirin.print.printable import Printable
 
 from .traits import (
+    StaticCall,
     HasSignature,
     SymbolOpInterface,
     EntryPointInterface,
@@ -52,6 +53,9 @@ class Method(Printable, typing.Generic[Param, RetType]):
     inferred: bool = False
     """if typeinfer has been run on this method
     """
+    backedges: set[Method] = field(init=False, repr=False)
+    """Cache for the backedges. (who calls this method)"""
+    run_passes: typing.Callable[[Method], None] | None = field(init=False, repr=False)
 
     def __init__(
         self,
@@ -96,6 +100,9 @@ def __init__(
         self.file = file
         self.lineno_begin = lineno_begin
         self.inferred = inferred
+        self.backedges = set()
+        self.update_backedges()
+        self.run_passes = None
 
     def __hash__(self) -> int:
         return id(self)
@@ -185,3 +192,13 @@ def verify_type(self) -> None:
         except ValidationError as e:
             e.attach(self)
             raise e
+
+    def update_backedges(self):
+        """Update the backedges of callee methods. (if they are static calls)"""
+        for stmt in self.code.walk():
+            trait = stmt.get_trait(StaticCall)
+            if not trait:
+                continue
+
+            callee = trait.get_callee(stmt)
+            callee.backedges.add(self)

src/kirin/ir/traits/__init__.py

@@ -31,6 +31,7 @@
     EntryPointInterface as EntryPointInterface,
 )
 from .callable import (
+    StaticCall as StaticCall,
     HasSignature as HasSignature,
     CallableStmtInterface as CallableStmtInterface,
 )

src/kirin/ir/traits/callable.py

@@ -7,7 +7,7 @@
 from kirin.ir.traits.abc import StmtTrait
 
 if TYPE_CHECKING:
-    from kirin.ir import Region, Statement
+    from kirin.ir import Method, Region, Statement
     from kirin.dialects.func.attrs import Signature
 
 StmtType = TypeVar("StmtType", bound="Statement")
@@ -63,3 +63,12 @@ def verify(self, node: "Statement"):
         signature = self.get_signature(node)
         if not isinstance(signature, Signature):
             raise ValueError(f"{signature} is not a Signature attribute")
+
+
+class StaticCall(StmtTrait, ABC, Generic[StmtType]):
+
+    @classmethod
+    @abstractmethod
+    def get_callee(cls, stmt: StmtType) -> "Method":
+        """Returns the callee of the static call statement."""
+        ...

test/ir/test_duplicated.py

@@ -1,4 +1,6 @@
 from kirin.ir import DialectGroup
+from kirin.prelude import basic
+from kirin.analysis import const
 from kirin.dialects import cf, func
 from kirin.dialects.py import base
 
@@ -45,3 +47,40 @@ def test_discard():
     assert "DialectGroup(" in target_b_repr
     assert base.dialect.name in target_b_repr
     assert cf.dialect.name in target_b_repr
+
+
+def test_overwrite():
+    @basic
+    def foo(x):  # type: ignore
+        return x * 2
+
+    @basic
+    def main(x):
+        return x + foo(x)
+
+    assert main(2) == 6
+
+    @basic
+    def foo(x):  # noqa: F811
+        return x * 3
+
+    assert main(2) == 8
+
+
+def test_recompile():
+    @basic
+    def foo(x):  # type: ignore
+        return x * 2
+
+    @basic(fold=True)
+    def main(x):
+        return x + foo(x)
+
+    ret = main.callable_region.blocks[0].stmts.at(0).results[0]
+    assert isinstance(ret.hints.get("const"), const.Unknown)
+
+    @basic
+    def foo(x):  # noqa: F811
+        return 3
+
+    assert isinstance(ret.hints.get("const"), const.Value)