refactor: with_thunk_scope

nix-js/src/context.rs

@@ -518,7 +518,7 @@ impl DowngradeContext for DowngradeCtx<'_> {
     fn downgrade(mut self, root: rnix::ast::Expr) -> Result<ExprId> {
         use crate::ir::TopLevel;
         let body = root.downgrade(&mut self)?;
-        let thunks = self.pop_thunk_scope();
+        let thunks = self.thunk_scopes.pop().expect("no thunk scope left???");
         let span = self.get_ir(body).span();
         let top_level = self.new_expr(TopLevel { body, thunks, span }.to_ir());
         self.ctx.irs.extend(self.irs);
@@ -553,14 +553,16 @@ impl DowngradeContext for DowngradeCtx<'_> {
         f(guard.as_ctx())
     }
 
-    fn push_thunk_scope(&mut self) {
+    fn with_thunk_scope<F, R>(&mut self, f: F) -> (R, Vec<(ExprId, ExprId)>)
+    where
+        F: FnOnce(&mut Self) -> R,
+    {
         self.thunk_scopes.push(Vec::new());
-    }
-
-    fn pop_thunk_scope(&mut self) -> Vec<(ExprId, ExprId)> {
-        self.thunk_scopes
-            .pop()
-            .expect("pop_thunk_scope without active scope")
+        let ret = f(self);
+        (
+            ret,
+            self.thunk_scopes.pop().expect("no thunk scope left???"),
+        )
     }
 
     fn register_thunk(&mut self, slot: ExprId, inner: ExprId) {

nix-js/src/downgrade.rs

@@ -18,6 +18,7 @@ pub trait DowngradeContext {
     fn new_expr(&mut self, expr: Ir) -> ExprId;
     fn new_arg(&mut self, span: TextRange) -> ExprId;
     fn maybe_thunk(&mut self, id: ExprId) -> ExprId;
+    fn register_thunk(&mut self, slot: ExprId, inner: ExprId);
 
     fn new_sym(&mut self, sym: String) -> SymId;
     fn get_sym(&self, id: SymId) -> Symbol<'_>;
@@ -37,10 +38,9 @@ pub trait DowngradeContext {
     fn with_with_scope<F, R>(&mut self, namespace: ExprId, f: F) -> R
     where
         F: FnOnce(&mut Self) -> R;
-
-    fn push_thunk_scope(&mut self);
-    fn pop_thunk_scope(&mut self) -> Vec<(ExprId, ExprId)>;
-    fn register_thunk(&mut self, slot: ExprId, inner: ExprId);
+    fn with_thunk_scope<F, R>(&mut self, f: F) -> (R, Vec<(ExprId, ExprId)>)
+    where
+        F: FnOnce(&mut Self) -> R;
 }
 
 pub trait Downgrade<Ctx: DowngradeContext> {
@@ -419,8 +419,12 @@ impl<Ctx: DowngradeContext> Downgrade<Ctx> for ast::Lambda {
         let raw_param = self.param().unwrap();
         let arg = ctx.new_arg(raw_param.syntax().text_range());
 
-        ctx.push_thunk_scope();
+        struct Ret {
+            param: Option<Param>,
+            body: ExprId,
+        }
 
+        let (ret, thunks) = ctx.with_thunk_scope(|ctx| {
             let param;
             let body;
 
@@ -431,8 +435,9 @@ impl<Ctx: DowngradeContext> Downgrade<Ctx> for ast::Lambda {
                     param = None;
 
                     // Downgrade body in Param scope
-                body = ctx
-                    .with_param_scope(param_sym, arg, |ctx| self.body().unwrap().downgrade(ctx))?;
+                    body = ctx.with_param_scope(param_sym, arg, |ctx| {
+                        self.body().unwrap().downgrade(ctx)
+                    })?;
                 }
                 ast::Param::Pattern(pattern) => {
                     let alias = pattern
@@ -460,7 +465,10 @@ impl<Ctx: DowngradeContext> Downgrade<Ctx> for ast::Lambda {
                 }
             }
 
-        let thunks = ctx.pop_thunk_scope();
+            Result::Ok(Ret { param, body })
+        });
+        let Ret { param, body } = ret?;
+
         let span = self.syntax().text_range();
         Ok(ctx.new_expr(
             Func {