#16 [asm] save context to stack and recusrive lambdas work (#22)

Author: vityaman

sleepy/asmik/emit.py

@@ -1,7 +1,7 @@
 from typing import override
 
 import sleepy.tafka.representation as taf
-from sleepy.tafka.walker import TafkaWalker
+from sleepy.tafka import Context, TafkaWalker, Usages
 
 from .argument import Immediate, Integer, Unassigned
 from .argument import PhysicalRegister as PhysReg
@@ -20,6 +20,7 @@
     Orb,
     Remi,
     Slti,
+    Stor,
     Xorb,
     mov,
     movi,
@@ -42,17 +43,27 @@ def temporary(self) -> VirtReg:
         return next(self.sequence)
 
 
-class AsmikEmitListener(TafkaWalker.Listener):
+class AsmikEmitListener(TafkaWalker.ContextedListener):
     def __init__(self) -> None:
+        super().__init__()
+
         self.memory = Memory()
         self.registers = VirtualRegisters()
 
         self.resolved: dict[str, int] = {}
 
+        self.procedure: taf.Procedure
+        self.usages: Usages
+
     @override
     def enter_procedure(self, procedure: taf.Procedure) -> None:
+        super().enter_procedure(procedure)
+
+        self.usages = Usages.analyzed(procedure)
+        self.procedure = procedure
+
         addr = self.memory.data_put(IntegerData(self.next_instr_addr))
-        self.resolved[repr(procedure.const)] = addr
+        self.resolved[f"${procedure.const.name}"] = addr
         for i, param in enumerate(procedure.parameters):
             register = self.registers.binded_to(param)
             self.emit(mov(register, PhysReg.arg(i + 1)))
@@ -63,6 +74,7 @@ def exit_procedure(self, procedure: taf.Procedure) -> None:
 
     @override
     def enter_block(self, block: taf.Block) -> None:
+        super().enter_block(block)
         self.resolved[repr(block.label)] = self.next_instr_addr
 
     @override
@@ -71,7 +83,7 @@ def exit_block(self, block: taf.Block) -> None:
 
     @override
     def enter_statement(self, statement: taf.Statement) -> None:
-        pass
+        super().enter_statement(statement)
 
     @override
     def exit_statement(self, statement: taf.Statement) -> None:
@@ -98,27 +110,48 @@ def on_conditional(self, conditional: taf.Conditional) -> None:
         self.emit(movi(else_address, Unassigned(else_label)))
         self.emit(Brn(condition, else_address))
 
+    def push_context(self, variables: list[taf.Var]) -> None:
+        def push(register: Reg) -> None:
+            self.emit(Stor(Reg.sp(), register))
+            self.emit(Addim(Reg.sp(), Reg.sp(), Integer(8)))
+
+        for local in variables:
+            if self.is_alive(local):
+                push(self.registers.binded_to(local))
+        push(Reg.ra())
+
+    def pop_context(self, variables: list[taf.Var]) -> None:
+        def pop(register: Reg) -> None:
+            self.emit(Addim(Reg.sp(), Reg.sp(), Integer(-8)))
+            self.emit(Load(register, Reg.sp()))
+
+        pop(Reg.ra())
+        for local in variables[::-1]:
+            if self.is_alive(local):
+                pop(self.registers.binded_to(local))
+
     @override
     def on_invokation(
         self,
         target: taf.Var,
         source: taf.Invokation,
     ) -> None:
+        variables = list(self.procedure.locals)
+
+        self.push_context(variables)
+
         for i, arg in enumerate(source.args):
             arg_reg = self.registers.binded_to(arg)
             self.emit(mov(PhysReg.arg(i + 1), arg_reg))
 
-        prev_ra = self.registers.temporary()
-        self.emit(mov(prev_ra, Reg.ra()))
-
         proc_reg = self.registers.binded_to(source.closure)
         self.emit(Addim(Reg.ra(), Reg.ip(), Integer(4)))
         self.emit(Brn(Reg.ze(), proc_reg))
 
-        res_reg = self.registers.binded_to(target)
-        self.emit(mov(res_reg, Reg.a1()))
+        result = self.registers.binded_to(target)
+        self.emit(mov(result, Reg.a1()))
 
-        self.emit(mov(Reg.ra(), prev_ra))
+        self.pop_context(variables)
 
     @override
     def on_load(self, target: taf.Var, source: taf.Load) -> None:
@@ -215,10 +248,20 @@ def addr_of(self, cnst: taf.Const) -> Immediate:
                 addr = self.memory.data_put(data)
                 return Integer(addr)
             case taf.Signature():
-                return Unassigned(repr(cnst))
+                return Unassigned(f"${cnst.name}")
             case _:
                 raise NotImplementedError
 
+    def is_alive(self, var: taf.Var) -> bool:
+        read = self.usages.next_read(var, self.context)
+        write = self.usages.next_write(var, self.context)
+        init = self.usages.next_write(var, Context(-1, self.procedure.entry))
+
+        assert init is not None  # noqa: S101
+        if write is not None:
+            return read is not None and read <= write
+        return read is not None and init < self.position
+
     @property
     def next_instr_addr(self) -> int:
         return len(self.memory.instr) * 4

sleepy/asmik/unit.py

@@ -26,7 +26,7 @@ def resolve_addresses(asmik: AsmikEmitListener) -> None:
         asmik = AsmikEmitListener()
         walker = TafkaWalker(asmik)
 
-        walker.explore_block(tafka.main)
+        walker.explore_procedure(tafka.main)
         for proc in tafka.procedures:
             walker.explore_procedure(proc)
 

sleepy/interpreter/asmik.py

@@ -35,6 +35,7 @@ def __init__(self) -> None:
         self.registers["ze"] = 0
         self.registers["ip"] = 0
         self.registers["ra"] = self.STOP
+        self.registers["sp"] = 10000
 
         self.running = False
 
@@ -101,9 +102,7 @@ def read(self, reg: Register) -> int:
         return self.registers[repr(reg)]
 
     def write(self, reg: Register, value: int) -> None:
-        match reg:
-            case "ze":
-                message = "ze is readonly"
-                raise SleepyError(message)
-            case _:
-                self.registers[repr(reg)] = value
+        if reg == Register.ze():
+            message = "ze is readonly"
+            raise SleepyError(message)
+        self.registers[repr(reg)] = value

sleepy/main.py

@@ -6,13 +6,13 @@
 
 def main() -> None:
     source = """
-        (def id (lambda (n int) n))
-        (def a (id 1))
-        (def b (id 11))
-        (def c (id 111))
-        (if (and (eq a 1)
-            (and (eq b 11)
-                 (eq c 111))) 1 0)
+    (def fibb (lambda (n int)
+        (if (or (eq n 0) (eq n 1))
+            1
+            (sum
+                (self (sum n -1))
+                (self (sum n -2))))))
+    (fibb 13)
     """
 
     parser = LarkParser()

sleepy/program/representation.py

@@ -54,8 +54,11 @@ class Intrinsic(Atom):
 
 @dataclass
 class Closure(Expression):
+    from .namespace import Namespace
+
     parameters: list[Parameter]
     statements: list[Expression]
+    namespace: Namespace
 
 
 @dataclass

sleepy/syntax/s2p.py

@@ -78,7 +78,7 @@ def visit_lambda(self, tree: LambdaAST) -> Closure:
         for parameter in parameters:
             self.bindings.bind(parameter.name, parameter)
 
-        closure = Closure(parameters, statements=[])
+        closure = Closure(parameters, statements=[], namespace=self.namespace)
         self.bindings.bind(self.namespace.define(Symbol("self")), closure)
 
         closure.statements = [

sleepy/tafka/__init__.py

@@ -33,4 +33,4 @@
 )
 from .unit import TafkaUnit
 from .usage import Usages
-from .walker import TafkaWalker
+from .walker import Context, TafkaWalker

sleepy/tafka/emit.py

@@ -10,22 +10,29 @@ class TafkaEmitVisitor(ProgramVisitor[None]):
     def __init__(self, unit: ProgramUnit) -> None:
         self.unit = unit
 
-        self.main = taf.Block(taf.Label("main"), [])
+        self.main = taf.Procedure(
+            name="main",
+            entry=taf.Block(taf.Label("main"), []),
+            parameters=[],
+            value=taf.Unknown(),
+        )
         self.procedures: list[taf.Procedure] = []
 
         self.var_names = map(str, range(10000000))
         self.lbl_names = map(str, range(10000000))
 
         self.vars = MetaTable[taf.Var]()
 
-        self.current_block = self.main
+        self.current_procedure = self.main
+        self.current_block = self.current_procedure.entry
         self.last_result = taf.Var("0", taf.Int())
 
     @override
     def visit_program(self, tree: program.Program) -> None:
         for statement in tree.statements:
             self.visit_expression(statement)
         self.emit_statement(taf.Return(self.last_result))
+        self.main.value = self.last_result.kind
 
     @override
     def visit_conditional(self, tree: program.Conditional) -> None:
@@ -116,31 +123,46 @@ def visit_application_variable(
 
     @override
     def visit_lambda(self, tree: program.Closure) -> None:
+        current_procedure = self.current_procedure
         current_block = self.current_block
 
         label = self.next_lbl()
 
-        params = [
+        procedure = taf.Procedure(
+            name=label.name,
+            entry=taf.Block(label, statements=[]),
+            parameters=[],
+            value=taf.Unknown(),
+        )
+
+        self.current_procedure = procedure
+        self.current_block = procedure.entry
+
+        procedure.parameters = [
             self.next_var(taf.Kind.from_sleepy(param.kind))
             for param in tree.parameters
         ]
 
-        for param, var in zip(tree.parameters, params, strict=True):
+        for param, var in zip(
+            tree.parameters,
+            procedure.parameters,
+            strict=True,
+        ):
             self.vars[param.name] = var
 
-        body = taf.Block(label, statements=[])
+        self.emit_intermidiate(
+            taf.Load(taf.Const(label.name, procedure.signature)),
+        )
+        self.vars[tree.namespace.resolved("self")] = self.last_result
 
-        self.current_block = body
         for statement in tree.statements:
             self.visit_expression(statement)
-
         self.emit_statement(taf.Return(self.last_result))
-
-        value = self.last_result.kind
+        procedure.value = self.last_result.kind
 
         self.current_block = current_block
+        self.current_procedure = current_procedure
 
-        procedure = taf.Procedure(label.name, body, params, value)
         self.procedures.append(procedure)
 
         self.emit_intermidiate(
@@ -167,13 +189,17 @@ def visit_definition(self, tree: program.Definition) -> None:
     def emit_statement(self, statement: taf.Statement) -> None:
         if isinstance(statement, taf.Set):
             self.last_result = statement.target
+        if isinstance(statement, taf.Return):
+            self.current_procedure.value = self.last_result.kind
         self.current_block.statements.append(statement)
 
     def emit_intermidiate(self, rvalue: taf.RValue) -> None:
         self.emit_statement(taf.Set(self.next_var(rvalue.value), rvalue))
 
     def next_var(self, kind: taf.Kind) -> taf.Var:
-        return taf.Var(next(self.var_names), kind)
+        var = taf.Var(next(self.var_names), kind)
+        self.current_procedure.locals.append(var)
+        return var
 
     def next_lbl(self) -> taf.Label:
         return taf.Label(next(self.lbl_names))

sleepy/tafka/representation/__init__.py

@@ -9,7 +9,7 @@
     Set,
     Statement,
 )
-from .kind import Int, Kind, Signature
+from .kind import Int, Kind, Signature, Unknown
 from .node import Node
 from .rvalue import (
     And,

sleepy/tafka/representation/block.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import cast
 
 from .kind import Kind, Signature
@@ -60,6 +60,7 @@ class Procedure(Node):
     entry: Block
     parameters: list[Var]
     value: Kind
+    locals: list[Var] = field(default_factory=list, init=False)
 
     @property
     def signature(self) -> Signature:

sleepy/tafka/representation/kind.py

@@ -18,14 +18,21 @@ def from_sleepy(cls, kind: SleepyKind) -> "Kind":
                 raise NotImplementedError
 
 
-@dataclass(repr=False)
+@dataclass(repr=False, unsafe_hash=True)
+class Unknown(Kind):
+    @override
+    def __repr__(self) -> str:
+        return "?"
+
+
+@dataclass(repr=False, unsafe_hash=True)
 class Int(Kind):
     @override
     def __repr__(self) -> str:
         return "int"
 
 
-@dataclass(repr=False)
+@dataclass(repr=False, unsafe_hash=True)
 class Bool(Kind):
     @override
     def __repr__(self) -> str:
@@ -37,9 +44,10 @@ class Signature(Kind):
     params: list[Kind]
     value: Kind
 
+    @override
+    def __hash__(self) -> int:
+        return hash(str(self.params)) + hash(self.value)
+
     @override
     def __repr__(self) -> str:
-        return (
-            f"({', '.join(repr(_) for _ in self.params)}) "
-            f"-> {self.value!r}"
-        )
+        return f"({', '.join(repr(_) for _ in self.params)}) -> {self.value!r}"