From 9f9ff0bf76faba00cbc357589e2d5b5c33d81819 Mon Sep 17 00:00:00 2001
From: timofey <tim@ongoteam.yaconnect.com>
Date: Sun, 14 May 2023 05:18:32 +0000
Subject: [PATCH] more flexible `UnbalancedConstructor`

---
 src/flow/traversible/unbalanced.rs | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/flow/traversible/unbalanced.rs b/src/flow/traversible/unbalanced.rs
index 72b45a4..b0f67fe 100644
--- a/src/flow/traversible/unbalanced.rs
+++ b/src/flow/traversible/unbalanced.rs
@@ -112,8 +112,9 @@ impl<'a, T: 'a + Monad, A: 'a> TraversibleBinaryTree<'a, T, A, UnbalancedData>
     }
 }
 
-type WrapType<'a, T, A> =
-    Box<dyn 'a + Fn(Rc<UnbalancedNode<'a, T, A>>) -> UnbalancedResolution<'a, T, A>>;
+type WrapType<'a, T, A> = Box<
+    dyn 'a + Fn(Rc<UnbalancedNode<'a, T, A>>) -> Box<dyn Fn() -> UnbalancedResolution<'a, T, A>>,
+>;
 
 pub struct UnbalancedConstructor<'a, T: 'a + Monad, A: 'a> {
     wrap: WrapType<'a, T, A>,
@@ -140,10 +141,7 @@ impl<'a, T: 'a + Monad, A: 'a> UnbalancedConstructor<'a, T, A> {
             key: key.into(),
         });
         let ctr = self.clone();
-        UnbalancedTree::Node(Rc::new(UnbalancedReference(Box::new(move || {
-            (ctr.wrap)(node.clone())
-        }))))
-        .into()
+        UnbalancedTree::Node(Rc::new(UnbalancedReference((ctr.wrap)(node)))).into()
     }
 
     #[cfg(test)]
@@ -182,7 +180,7 @@ mod tests {
     #[test]
     fn test_simple_slices() {
         let ctr: Rc<UnbalancedConstructor<classes::result::ResultClass<()>, _>> =
-            UnbalancedConstructor::rc(Box::new(|node| Ok(node)));
+            UnbalancedConstructor::rc(Box::new(|node| Box::new(move || Ok(node.clone()))));
         let mut rng = rand::thread_rng();
         let t_set = ctr.from_slice(&mut rng, &[0]);
         assert!(t_contains(&DefaultComparator, t_set.clone(), 0.into()).is_ok());
@@ -215,7 +213,7 @@ mod tests {
     #[test]
     fn test_random_slices() {
         let ctr: Rc<UnbalancedConstructor<classes::result::ResultClass<()>, _>> =
-            UnbalancedConstructor::rc(Box::new(|node| Ok(node)));
+            UnbalancedConstructor::rc(Box::new(|node| Box::new(move || Ok(node.clone()))));
         let mut rng = rand::thread_rng();
         for _ in 0..1000 {
             let big: Vec<i32> = (0..(rng.gen_range(2..10))).collect();
@@ -263,7 +261,9 @@ mod tests {
     #[test]
     fn trace_one_slice() {
         let ctr: Rc<UnbalancedConstructor<TracedMonad, _>> =
-            UnbalancedConstructor::rc(Box::new(|node| TracedMonad::pure(node).after_resolution()));
+            UnbalancedConstructor::rc(Box::new(|node| {
+                Box::new(move || TracedMonad::pure(node.clone()).after_resolution())
+            }));
         // let mut rng = rand::thread_rng();
         let mut rng = rand::rngs::StdRng::seed_from_u64(426);
         let big: Vec<i32> = (0..(rng.gen_range(1000..2000))).collect();