From 1a69267fac9e5f4e6f01f76e6916d8b5fc6761e6 Mon Sep 17 00:00:00 2001
From: julia <jellemartijndijkstra@gmail.com>
Date: Thu, 30 Nov 2023 14:28:52 +0100
Subject: [PATCH] Moved the problem of uncover_get!. We still need to pass
 whether the lhs of let-bindings are actually mutable.

---
 compiler/src/passes/atomize/atomize.rs | 74 +++++++++++++++-----------
 1 file changed, 43 insertions(+), 31 deletions(-)

diff --git a/compiler/src/passes/atomize/atomize.rs b/compiler/src/passes/atomize/atomize.rs
index ad3d269..04a0cef 100644
--- a/compiler/src/passes/atomize/atomize.rs
+++ b/compiler/src/passes/atomize/atomize.rs
@@ -3,6 +3,7 @@ use crate::passes::parse::types::Type;
 use crate::passes::parse::{Meta, Typed};
 use crate::passes::reveal::{DefRevealed, PrgRevealed, RExpr};
 use crate::utils::gen_sym::{gen_sym, UniqueSym};
+use crate::utils::push_map::PushMap;
 
 impl<'p> PrgRevealed<'p> {
     #[must_use]
@@ -26,17 +27,25 @@ fn atomize_def(def: DefRevealed) -> DefAtomized {
             params,
             typ,
             bdy,
-        } => DefAtomized::Fn {
-            sym,
-            params,
-            typ,
-            bdy: atomize_expr(bdy),
-        },
+        } => {
+            let mut scope =
+                PushMap::from_iter(params.iter().map(|param| (param.sym, param.mutable)));
+
+            DefAtomized::Fn {
+                sym,
+                params,
+                typ,
+                bdy: atomize_expr(bdy, &mut scope),
+            }
+        }
         DefRevealed::TypeDef { sym, def } => DefAtomized::TypeDef { sym, def },
     }
 }
 
-fn atomize_expr<'p>(expr: Typed<'p, RExpr<'p>>) -> Typed<'p, AExpr<'p>> {
+fn atomize_expr<'p>(
+    expr: Typed<'p, RExpr<'p>>,
+    scope: &mut PushMap<UniqueSym<'p>, bool>,
+) -> Typed<'p, AExpr<'p>> {
     // Keep track of all the priors. These are bindings that should come before evaluating the expression.
     let mut priors = Vec::new();
 
@@ -54,23 +63,25 @@ fn atomize_expr<'p>(expr: Typed<'p, RExpr<'p>>) -> Typed<'p, AExpr<'p>> {
         } => AExpr::BinaryOp {
             op,
             exprs: [
-                atomize_atom(*lhs, &mut priors),
-                atomize_atom(*rhs, &mut priors),
+                atomize_atom(*lhs, &mut priors, scope),
+                atomize_atom(*rhs, &mut priors, scope),
             ],
         },
         RExpr::UnaryOp { op, expr: arg } => AExpr::UnaryOp {
             op,
-            expr: atomize_atom(*arg, &mut priors),
-        },
-        RExpr::Let { sym, bnd, bdy } => AExpr::Let {
-            sym,
-            bnd: Box::new(atomize_expr(*bnd)),
-            bdy: Box::new(atomize_expr(*bdy)),
+            expr: atomize_atom(*arg, &mut priors, scope),
         },
+        RExpr::Let { sym, bnd, bdy } => {
+            let bnd = Box::new(atomize_expr(*bnd, scope));
+            // todo: this is too conservative, not every variable is mutable!
+            let bdy = Box::new(scope.push(sym, true, |scope| atomize_expr(*bdy, scope)));
+
+            AExpr::Let { sym, bnd, bdy }
+        }
         RExpr::If { cnd, thn, els } => AExpr::If {
-            cnd: Box::new(atomize_expr(*cnd)),
-            thn: Box::new(atomize_expr(*thn)),
-            els: Box::new(atomize_expr(*els)),
+            cnd: Box::new(atomize_expr(*cnd, scope)),
+            thn: Box::new(atomize_expr(*thn, scope)),
+            els: Box::new(atomize_expr(*els, scope)),
         },
         RExpr::Apply { fun, args } => {
             let Type::Fn { params, .. } = fun.meta.clone() else {
@@ -78,45 +89,45 @@ fn atomize_expr<'p>(expr: Typed<'p, RExpr<'p>>) -> Typed<'p, AExpr<'p>> {
             };
 
             AExpr::Apply {
-                fun: atomize_atom(*fun, &mut priors),
+                fun: atomize_atom(*fun, &mut priors, scope),
                 args: args
                     .into_iter()
-                    .map(|arg| atomize_atom(arg, &mut priors))
+                    .map(|arg| atomize_atom(arg, &mut priors, scope))
                     .zip(params)
                     .collect(),
             }
         }
         RExpr::FunRef { sym } => AExpr::FunRef { sym },
         RExpr::Loop { bdy } => AExpr::Loop {
-            bdy: Box::new(atomize_expr(*bdy)),
+            bdy: Box::new(atomize_expr(*bdy, scope)),
         },
         RExpr::Break { bdy } => AExpr::Break {
-            bdy: Box::new(atomize_expr(*bdy)),
+            bdy: Box::new(atomize_expr(*bdy, scope)),
         },
         RExpr::Seq { stmt, cnt } => AExpr::Seq {
-            stmt: Box::new(atomize_expr(*stmt)),
-            cnt: Box::new(atomize_expr(*cnt)),
+            stmt: Box::new(atomize_expr(*stmt, scope)),
+            cnt: Box::new(atomize_expr(*cnt, scope)),
         },
         RExpr::Assign { sym, bnd } => AExpr::Assign {
             sym,
-            bnd: Box::new(atomize_expr(*bnd)),
+            bnd: Box::new(atomize_expr(*bnd, scope)),
         },
         RExpr::Continue => AExpr::Continue,
         RExpr::Return { bdy } => AExpr::Return {
-            bdy: Box::new(atomize_expr(*bdy)),
+            bdy: Box::new(atomize_expr(*bdy, scope)),
         },
         RExpr::Struct { sym, fields } => AExpr::Struct {
             sym,
             fields: fields
                 .into_iter()
                 .map(|(sym, expr)| {
-                    let field = atomize_atom(expr, &mut priors);
+                    let field = atomize_atom(expr, &mut priors, scope);
                     (sym, field)
                 })
                 .collect(),
         },
         RExpr::AccessField { strct, field } => AExpr::AccessField {
-            strct: atomize_atom(*strct, &mut priors),
+            strct: atomize_atom(*strct, &mut priors, scope),
             field,
         },
         RExpr::Asm { instrs } => AExpr::Asm { instrs },
@@ -143,14 +154,15 @@ fn atomize_expr<'p>(expr: Typed<'p, RExpr<'p>>) -> Typed<'p, AExpr<'p>> {
 fn atomize_atom<'p>(
     expr: Typed<'p, RExpr<'p>>,
     priors: &mut Vec<(UniqueSym<'p>, Typed<'p, AExpr<'p>>)>,
+    scope: &mut PushMap<UniqueSym<'p>, bool>,
 ) -> Atom<'p> {
     match expr.inner {
         RExpr::Lit { val } => Atom::Val { val },
-        RExpr::Var { sym} => Atom::Var { sym },
+        RExpr::Var { sym } if !scope[&sym] => Atom::Var { sym },
         _ => {
             let tmp = gen_sym("tmp");
-            priors.push((tmp, atomize_expr(expr)));
+            priors.push((tmp, atomize_expr(expr, scope)));
             Atom::Var { sym: tmp }
         }
     }
-}
\ No newline at end of file
+}