pass in allocator to post-eval callback

Chia-Network · Apr 27, 2024 · ef3795c · ef3795c
1 parent d6db84d
commit ef3795c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 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,7 +430,7 @@ 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 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,