From d09b84676c4de3c7eb037d6cc545dc81853a1449 Mon Sep 17 00:00:00 2001
From: imxyy_soope_ <imxyy1soope1@gmail.com>
Date: Sun, 8 Feb 2026 12:41:34 +0800
Subject: [PATCH] refactor: with_thunk_scope

---
 nix-js/src/context.rs   | 18 ++++----
 nix-js/src/downgrade.rs | 96 ++++++++++++++++++++++-------------------
 2 files changed, 62 insertions(+), 52 deletions(-)

diff --git a/nix-js/src/context.rs b/nix-js/src/context.rs
index 415b2f4..2112afc 100644
--- a/nix-js/src/context.rs
+++ b/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) {
diff --git a/nix-js/src/downgrade.rs b/nix-js/src/downgrade.rs
index 42020d8..1e9bf01 100644
--- a/nix-js/src/downgrade.rs
+++ b/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,48 +419,56 @@ 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();
-
-        let param;
-        let body;
-
-        match raw_param {
-            ast::Param::IdentParam(id) => {
-                // Simple case: `x: body`
-                let param_sym = ctx.new_sym(id.to_string());
-                param = None;
-
-                // Downgrade body in Param scope
-                body = ctx
-                    .with_param_scope(param_sym, arg, |ctx| self.body().unwrap().downgrade(ctx))?;
-            }
-            ast::Param::Pattern(pattern) => {
-                let alias = pattern
-                    .pat_bind()
-                    .map(|alias| ctx.new_sym(alias.ident().unwrap().to_string()));
-
-                let ellipsis = pattern.ellipsis_token().is_some();
-                let pat_entries = pattern.pat_entries();
-
-                let PatternBindings {
-                    body: inner_body,
-                    required,
-                    optional,
-                } = downgrade_pattern_bindings(pat_entries, alias, arg, ctx, |ctx, _| {
-                    self.body().unwrap().downgrade(ctx)
-                })?;
-
-                param = Some(Param {
-                    required,
-                    optional,
-                    ellipsis,
-                });
-
-                body = inner_body;
-            }
+        struct Ret {
+            param: Option<Param>,
+            body: ExprId,
         }
 
-        let thunks = ctx.pop_thunk_scope();
+        let (ret, thunks) = ctx.with_thunk_scope(|ctx| {
+            let param;
+            let body;
+
+            match raw_param {
+                ast::Param::IdentParam(id) => {
+                    // Simple case: `x: body`
+                    let param_sym = ctx.new_sym(id.to_string());
+                    param = None;
+
+                    // Downgrade body in Param scope
+                    body = ctx.with_param_scope(param_sym, arg, |ctx| {
+                        self.body().unwrap().downgrade(ctx)
+                    })?;
+                }
+                ast::Param::Pattern(pattern) => {
+                    let alias = pattern
+                        .pat_bind()
+                        .map(|alias| ctx.new_sym(alias.ident().unwrap().to_string()));
+
+                    let ellipsis = pattern.ellipsis_token().is_some();
+                    let pat_entries = pattern.pat_entries();
+
+                    let PatternBindings {
+                        body: inner_body,
+                        required,
+                        optional,
+                    } = downgrade_pattern_bindings(pat_entries, alias, arg, ctx, |ctx, _| {
+                        self.body().unwrap().downgrade(ctx)
+                    })?;
+
+                    param = Some(Param {
+                        required,
+                        optional,
+                        ellipsis,
+                    });
+
+                    body = inner_body;
+                }
+            }
+
+            Result::Ok(Ret { param, body })
+        });
+        let Ret { param, body } = ret?;
+
         let span = self.syntax().text_range();
         Ok(ctx.new_expr(
             Func {