simple example of rewrite with insert of multiple stmts + codegen (#193)

This PR address #145

TODO: modify/add doc

Author: kaihsin

example/beer/emit.py

@@ -0,0 +1,93 @@
+from io import StringIO
+from dataclasses import field, dataclass
+
+import stmts
+from group import beer
+from dialect import dialect
+
+from kirin import ir, interp
+from lattice import Item, ItemPints, AtLeastXItem, ConstIntItem
+from kirin.emit import EmitStr, EmitStrFrame
+from kirin.dialects import func
+from kirin.emit.exceptions import EmitError
+
+
+def default_menu_price():
+    return {
+        "budlight": 3.0,
+        "heineken": 4.0,
+        "tsingdao": 2.0,
+    }
+
+
+@dataclass
+class EmitReceptMain(EmitStr):
+    keys = ["emit.recept"]
+    dialects: ir.DialectGroup = field(default=beer)
+    file: StringIO = field(default_factory=StringIO)
+    menu_price: dict[str, float] = field(default_factory=default_menu_price)
+    recept_analysis_result: dict[ir.SSAValue, Item] = field(default_factory=dict)
+
+    def initialize(self):
+        super().initialize()
+        self.file.truncate(0)
+        self.file.seek(0)
+        return self
+
+    def eval_stmt_fallback(
+        self, frame: EmitStrFrame, stmt: ir.Statement
+    ) -> tuple[str, ...]:
+        return (stmt.name,)
+
+    def emit_block(self, frame: EmitStrFrame, block: ir.Block) -> str | None:
+        for stmt in block.stmts:
+            result = self.eval_stmt(frame, stmt)
+            if isinstance(result, tuple):
+                frame.set_values(stmt.results, result)
+        return None
+
+    def get_output(self) -> str:
+        self.file.seek(0)
+        return "\n".join(
+            [
+                "item    \tamount \t  price",
+                "-----------------------------------",
+                self.file.read(),
+            ]
+        )
+
+
+@func.dialect.register(key="emit.recept")
+class FuncEmit(interp.MethodTable):
+
+    @interp.impl(func.Function)
+    def emit_func(self, emit: EmitReceptMain, frame: EmitStrFrame, stmt: func.Function):
+        _ = emit.run_ssacfg_region(frame, stmt.body)
+        return ()
+
+
+@dialect.register(key="emit.recept")
+class BeerEmit(interp.MethodTable):
+
+    @interp.impl(stmts.Pour)
+    def emit_pour(self, emit: EmitReceptMain, frame: EmitStrFrame, stmt: stmts.Pour):
+        pints_item: ItemPints = emit.recept_analysis_result[stmt.result]
+
+        amount_str = ""
+        price_str = ""
+        if isinstance(pints_item.count, AtLeastXItem):
+            amount_str = f">={pints_item.count.data}"
+            price_str = (
+                f"  >=${emit.menu_price[pints_item.brand] * pints_item.count.data}"
+            )
+        elif isinstance(pints_item.count, ConstIntItem):
+            amount_str = f"  {pints_item.count.data}"
+            price_str = (
+                f"  ${emit.menu_price[pints_item.brand] * pints_item.count.data}"
+            )
+        else:
+            raise EmitError("invalid analysis result.")
+
+        emit.writeln(frame, f"{pints_item.brand}\t{amount_str}\t{price_str}")
+
+        return ()

example/beer/group.py

@@ -0,0 +1,23 @@
+from dialect import dialect
+
+from rewrite import RandomWalkBranch
+from kirin.ir import dialect_group
+from kirin.prelude import basic_no_opt
+from kirin.rewrite import Walk, Fixpoint
+from kirin.passes.fold import Fold
+
+
+# create our own beer dialect, it runs a random walk on the branches
+@dialect_group(basic_no_opt.add(dialect))
+def beer(self):
+
+    fold_pass = Fold(self)
+
+    def run_pass(mt, *, fold=True):
+        Fixpoint(Walk(RandomWalkBranch())).rewrite(mt.code)
+
+        # add const fold
+        if fold:
+            fold_pass(mt)
+
+    return run_pass

example/beer/recept.py

@@ -89,12 +89,6 @@ def add(
 @dialect.register(key="beer.fee")
 class BeerMethodTable(interp.MethodTable):
 
-    menu_price: dict[str, float] = {
-        "budlight": 1.0,
-        "heineken": 2.0,
-        "tsingdao": 3.0,
-    }
-
     @interp.impl(NewBeer)
     def new_beer(
         self,

example/beer/rewrite.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 
-from stmts import RandomBranch
+from stmts import Puke, Drink, NewBeer, RandomBranch
 
 from kirin.dialects import cf
 from kirin.rewrite.abc import RewriteRule, RewriteResult
@@ -23,3 +23,21 @@ def rewrite_Statement(self, node: Statement) -> RewriteResult:
             )
         )
         return RewriteResult(has_done_something=True)
+
+
+@dataclass
+class NewBeerAndPukeOnDrink(RewriteRule):
+    # sometimes someone get drunk, so they keep getting new beer and puke after they drink
+    def rewrite_Statement(self, node: Statement) -> RewriteResult:
+        if not isinstance(node, Drink):
+            return RewriteResult()
+
+        # 1. create new stmts:
+        new_beer_stmt = NewBeer(brand="saporo")
+        puke_stmt = Puke()
+
+        # 2. put them in the ir
+        new_beer_stmt.insert_after(node)
+        puke_stmt.insert_after(new_beer_stmt)
+
+        return RewriteResult(has_done_something=True)

example/beer/script.py

@@ -1,34 +1,12 @@
+from group import beer
 from stmts import Pour, Puke, Drink, NewBeer
 from recept import FeeAnalysis
-from dialect import dialect
 
+from emit import EmitReceptMain
 from interp import BeerMethods as BeerMethods
 from lattice import AtLeastXItem
-from rewrite import RandomWalkBranch
-from kirin.ir import dialect_group
-from kirin.prelude import basic_no_opt
-from kirin.rewrite import Walk, Fixpoint
-from kirin.passes.fold import Fold
-
-
-# create our own beer dialect, it runs a random walk on the branches
-@dialect_group(basic_no_opt.add(dialect))
-def beer(self):
-
-    fold_pass = Fold(self)
-
-    def run_pass(mt, *, fold=True):
-        Fixpoint(Walk(RandomWalkBranch())).rewrite(mt.code)
-
-        # add const fold
-        if fold:
-            fold_pass(mt)
-
-    return run_pass
-
-
-# we are going to get drunk!
-# add our beer dialect to the default dialect (builtin, cf, func, ...)
+from rewrite import NewBeerAndPukeOnDrink
+from kirin.rewrite import Walk
 
 
 # type: ignore
@@ -63,7 +41,26 @@ def some_closure(beer, amount):
 #     main(i)  # now drink a random beer!
 
 
-# simple analysis example:
+# 2. simple rewrite:
+@beer
+def main3():
+
+    bud = NewBeer(brand="budlight")
+    heineken = NewBeer(brand="heineken")
+
+    bud_pints = Pour(bud, 2)
+    heineken_pints = Pour(heineken, 10)
+
+    Drink(bud_pints)
+    Drink(heineken_pints)
+
+
+main3.print()
+Walk(NewBeerAndPukeOnDrink()).rewrite(main3.code)
+main3.print()
+
+
+# 3. simple analysis example:
 @beer
 def main2(x: int):
 
@@ -91,3 +88,10 @@ def main2(x: int):
 print(results)
 print(fee_analysis.puke_count)
 main2.print(analysis=results)
+
+
+emitter = EmitReceptMain()
+emitter.recept_analysis_result = results
+
+emitter.run(main2, ("",))
+print(emitter.get_output())

src/kirin/interp/base.py

@@ -85,7 +85,7 @@ def __post_init__(self) -> None:
         self.registry = self.dialects.registry.interpreter(keys=self.keys)
 
     def initialize(self) -> Self:
-        """Initialize the interpreter global states. This method is called before
+        """Initialize the interpreter global states. This method is called right upon
         calling [`run`][kirin.interp.base.BaseInterpreter.run] to initialize the
         interpreter global states.