pass in allocator to post-eval callback

Chia-Network · Apr 29, 2024 · 9639ed7 · 9639ed7
1 parent d6db84d
commit 9639ed7
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 14 deletions.
diff --git a/src/run_program.rs b/src/run_program.rs
@@ -24,7 +24,7 @@ pub type PreEval =
     Box<dyn Fn(&mut Allocator, NodePtr, NodePtr) -> Result<Option<Box<PostEval>>, EvalErr>>;
 
 #[cfg(feature = "pre-eval")]
-pub type PostEval = dyn Fn(Option<NodePtr>);
+pub type PostEval = dyn Fn(&mut Allocator, Option<NodePtr>);
 
 #[repr(u8)]
 enum Operation {
@@ -495,7 +495,7 @@ impl<'a, D: Dialect> RunProgramContext<'a, D> {
                 Operation::PostEval => {
                     let f = self.posteval_stack.pop().unwrap();
                     let peek: Option<NodePtr> = self.val_stack.last().copied();
-                    f(peek);
+                    f(&mut self.allocator, peek);
                     0
                 }
             };

diff --git a/src/test_ops.rs b/src/test_ops.rs
@@ -413,7 +413,7 @@ fn test_pre_eval_and_post_eval() {
             &mut Allocator,
             NodePtr,
             NodePtr,
-        ) -> Result<Option<Box<(dyn Fn(Option<NodePtr>))>>, EvalErr>,
+        ) -> Result<Option<Box<(dyn Fn(&mut Allocator, Option<NodePtr>))>>, EvalErr>,
     > = Box::new(move |_allocator, prog, args| {
         let tracking_key = pre_eval_tracking.borrow().len();
         // Ensure lifetime of mutable borrow is contained.
@@ -430,17 +430,18 @@ fn test_pre_eval_and_post_eval() {
             );
         }
         let post_eval_tracking = pre_eval_tracking.clone();
-        let post_eval_f: Box<dyn Fn(Option<NodePtr>)> = Box::new(move |outcome| {
-            let mut tracking_mutable = post_eval_tracking.borrow_mut();
-            tracking_mutable.insert(
-                tracking_key,
-                EvalFTracker {
-                    prog,
-                    args,
-                    outcome,
-                },
-            );
-        });
+        let post_eval_f: Box<dyn Fn(&mut Allocator, Option<NodePtr>)> =
+            Box::new(move |_a, outcome| {
+                let mut tracking_mutable = post_eval_tracking.borrow_mut();
+                tracking_mutable.insert(
+                    tracking_key,
+                    EvalFTracker {
+                        prog,
+                        args,
+                        outcome,
+                    },
+                );
+            });
         Ok(Some(post_eval_f))
     });