From 15e4762d4ac8993d12fe2dd54e2b2d842c8a034b Mon Sep 17 00:00:00 2001
From: Jakub Doka <jakub.doka2@gmail.com>
Date: Sat, 19 Oct 2024 19:53:43 +0200
Subject: [PATCH] cleanup: 2

---
 lang/src/son.rs                    | 54 +++++++++++++++++++++---------
 lang/tests/son_tests_fb_driver.txt |  4 +--
 2 files changed, 41 insertions(+), 17 deletions(-)

diff --git a/lang/src/son.rs b/lang/src/son.rs
index f2deee9..8d329cb 100644
--- a/lang/src/son.rs
+++ b/lang/src/son.rs
@@ -990,6 +990,7 @@ impl ItemCtx {
 
     fn emit_body_code(&mut self, sig: Sig, tys: &Types) -> usize {
         let mut nodes = core::mem::take(&mut self.nodes);
+        nodes.visited.clear(nodes.values.len());
 
         let func = Function::new(&mut nodes, tys, sig);
         let mut ralloc = Regalloc::default(); // TODO: reuse
@@ -1165,7 +1166,7 @@ impl ItemCtx {
         saved_regs.len()
     }
 
-    fn emit_body(&mut self, id: ty::Func, tys: &mut Types, sig: Sig) {
+    fn emit_body(&mut self, tys: &mut Types, sig: Sig) {
         '_open_function: {
             self.emit(instrs::addi64(reg::STACK_PTR, reg::STACK_PTR, 0));
             self.emit(instrs::st(reg::RET_ADDR, reg::STACK_PTR, 0, 0));
@@ -1184,7 +1185,6 @@ impl ItemCtx {
             self.nodes[MEM].outputs = mems;
         }
 
-        self.nodes.visited.clear(self.nodes.values.len());
         let saved = self.emit_body_code(sig, tys);
 
         if let Some(last_ret) = self.ret_relocs.last()
@@ -1237,6 +1237,7 @@ impl ItemCtx {
         self.relocs.iter_mut().for_each(|r| r.reloc.offset -= stripped_prelude_size as u32);
         self.emit(instrs::jala(reg::ZERO, reg::RET_ADDR, 0));
 
+        let ty::Kind::Func(id) = self.id.expand() else { unreachable!() };
         tys.ins.funcs[id as usize].code.append(&mut self.code);
         tys.ins.funcs[id as usize].relocs.append(&mut self.relocs);
     }
@@ -1266,6 +1267,32 @@ impl Ctx {
 #[derive(Default)]
 struct Pool {
     cis: Vec<ItemCtx>,
+    used_cis: usize,
+}
+
+impl Pool {
+    pub fn push_ci(
+        &mut self,
+        target: &mut ItemCtx,
+        file: FileId,
+        id: ty::Id,
+        ret: Option<ty::Id>,
+        task_base: usize,
+    ) {
+        if let Some(slot) = self.cis.get_mut(self.used_cis) {
+            core::mem::swap(slot, target);
+        } else {
+            self.cis.push(ItemCtx::default());
+            core::mem::swap(self.cis.last_mut().unwrap(), target);
+        }
+        target.init(file, id, ret, task_base);
+        self.used_cis += 1;
+    }
+
+    pub fn pop_ci(&mut self, target: &mut ItemCtx) {
+        self.used_cis -= 1;
+        core::mem::swap(&mut self.cis[self.used_cis], target);
+    }
 }
 
 struct Regalloc {
@@ -2227,22 +2254,15 @@ impl<'a> Codegen<'a> {
 
     fn emit_func(&mut self, FTask { file, id }: FTask) {
         let func = &mut self.tys.ins.funcs[id as usize];
-        debug_assert!(func.file == file);
+        debug_assert_eq!(func.file, file);
         func.offset = u32::MAX - 1;
         let sig = func.sig.expect("to emmit only concrete functions");
         let ast = &self.files[file as usize];
         let expr = func.expr.get(ast).unwrap();
 
-        let mut prev_ci = self.pool.cis.pop().unwrap_or_default();
-        prev_ci.init(file, ty::Kind::Func(id).compress(), Some(sig.ret), 0);
-        core::mem::swap(&mut self.ci, &mut prev_ci);
+        self.pool.push_ci(&mut self.ci, file, ty::Kind::Func(id).compress(), Some(sig.ret), 0);
 
-        let Expr::BinOp {
-            left: Expr::Ident { .. },
-            op: TokenKind::Decl,
-            right: &Expr::Closure { body, args, .. },
-        } = expr
-        else {
+        let Expr::BinOp { right: &Expr::Closure { body, args, .. }, .. } = expr else {
             unreachable!("{}", self.ast_display(expr))
         };
 
@@ -2257,7 +2277,11 @@ impl<'a> Codegen<'a> {
         }
 
         if self.expr(body).is_some() && self.tys.size_of(sig.ret) != 0 {
-            self.report(body.pos(), "expected all paths in the fucntion to return");
+            self.report(
+                body.pos(),
+                "expected all paths in the fucntion to return \
+                or the return type to be zero sized",
+            );
         }
 
         self.ci.finalize();
@@ -2271,10 +2295,10 @@ impl<'a> Codegen<'a> {
 
             self.graphviz();
 
-            self.ci.emit_body(id, &mut self.tys, sig);
+            self.ci.emit_body(&mut self.tys, sig);
         }
 
-        self.pool.cis.push(core::mem::replace(&mut self.ci, prev_ci));
+        self.pool.pop_ci(&mut self.ci);
     }
 
     fn ty(&mut self, expr: &Expr) -> ty::Id {
diff --git a/lang/tests/son_tests_fb_driver.txt b/lang/tests/son_tests_fb_driver.txt
index 6d3b629..ae33add 100644
--- a/lang/tests/son_tests_fb_driver.txt
+++ b/lang/tests/son_tests_fb_driver.txt
@@ -42,8 +42,8 @@ main:
     ADDI64 r254, r254, 64d
     JALA r0, r31, 0a
 set_pixel:
-    MUL64 r8, r4, r3
-    ADD64 r1, r8, r2
+    MUL64 r7, r3, r4
+    ADD64 r1, r7, r2
     JALA r0, r31, 0a
 x86_fb_ptr:
     LI64 r1, 100d