Merge pull request #40 from Centril/feature/recursion-improvements

Reduce heap allocation in .prop_recursive(..) + make it more ergonomic

CHANGELOG.md

@@ -2,6 +2,13 @@
 
 ### Potential Breaking Changes
 
+- There is a small chance of breakage if you've relied on the constraints put
+  on type inference by the closure in `leaf.prop_recursive(..)` having a fixed
+  output type. The output type is now any strategy that generates the same type
+  as `leaf`. This change is intended to make working with recursive types a bit
+  easier as you no longer have to use `.boxed()` inside the closure you pass to
+  `.prop_recursive(..)`.
+
 - There is a small chance of breakage wrt. type inference due to the
   introduction of `SizeRange`.
 
@@ -63,6 +70,8 @@
 
 - Relaxed the constraints of `btree_map` removing `'static`.
 
+- Reduced the heap allocation inside `Recursive` somewhat.
+
 ## 0.4.2
 
 ### Bug Fixes

src/arbitrary/_std/env.rs

@@ -71,7 +71,7 @@ fn make_utf16_invalid(buf: &mut Vec<u16>, p: usize) {
 #[cfg(target_os = "windows")]
 fn osstring_invalid_string() -> BoxedStrategy<OsString> {
     use std::os::windows::ffi::OsStringExt;
-    let size = 0..::std::u16::MAX as usize;
+    let size = 1..::std::u16::MAX as usize;
     let vec_gen = ::collection::vec(..::std::u16::MAX, size.clone());
     (size, vec_gen).prop_map(|(p, mut sbuf)| {
         // Not quite a uniform distribution due to clamping,

src/lib.rs

@@ -1277,7 +1277,7 @@
 //!               .prop_map(Json::Array),
 //!           prop::collection::hash_map(".*", inner, 0..10)
 //!               .prop_map(Json::Map),
-//!       ].boxed()).boxed()
+//!       ]).boxed()
 //! }
 //! # fn main() { }
 //! ```

src/strategy/recursive.rs

@@ -11,8 +11,32 @@ use std::fmt;
 use std::sync::Arc;
 
 use strategy::traits::*;
+use strategy::IndFlatten;
+use strategy::statics::{Map, MapFn};
 use test_runner::*;
 
+/// A branching `MapFn` that picks `yes` (true) or
+/// `no` (false) and clones the `Arc`.
+struct BranchFn<T> {
+    /// Result on true.
+    yes: Arc<T>,
+    /// Result on false.
+    no: Arc<T>,
+}
+
+impl<T> Clone for BranchFn<T> {
+    fn clone(&self) -> Self {
+        Self { yes: Arc::clone(&self.yes), no: Arc::clone(&self.no) }
+    }
+}
+
+impl<T: fmt::Debug> MapFn<bool> for BranchFn<T> {
+    type Output = Arc<T>;
+    fn apply(&self, branch: bool) -> Self::Output {
+        Arc::clone(if branch { &self.yes } else { &self.no })
+    }
+}
+
 /// Return type from `Strategy::prop_recursive()`.
 pub struct Recursive<B, F> {
     pub(super) base: Arc<B>,
@@ -47,8 +71,35 @@ impl<B, F> Clone for Recursive<B, F> {
 }
 
 impl<T : fmt::Debug + 'static,
-     F : Fn (Arc<BoxedStrategy<T>>) -> BoxedStrategy<T>>
-Strategy for Recursive<BoxedStrategy<T>, F> {
+     R : Strategy + 'static,
+     F : Fn(Arc<BoxedStrategy<T>>) -> R>
+Recursive<BoxedStrategy<T>, F>
+where
+    R::Value : ValueTree<Value = T>
+{
+    pub(super) fn new<S>
+        (base: S, depth: u32, desired_size: u32, expected_branch_size: u32,
+        recurse: F)
+        -> Self
+    where
+        S : Strategy + 'static,
+        S::Value : ValueTree<Value = T>
+    {
+        Self {
+            base: Arc::new(base.boxed()),
+            recurse: Arc::new(recurse),
+            depth, desired_size, expected_branch_size,
+        }
+    }
+}
+
+impl<T : fmt::Debug + 'static,
+     R : Strategy + 'static,
+     F : Fn(Arc<BoxedStrategy<T>>) -> R>
+Strategy for Recursive<BoxedStrategy<T>, F>
+where
+    R::Value : ValueTree<Value = T>
+{
     type Value = Box<ValueTree<Value = T>>;
 
     fn new_value(&self, runner: &mut TestRunner) -> NewTree<Self> {
@@ -90,15 +141,15 @@ Strategy for Recursive<BoxedStrategy<T>, F> {
 
         let mut strat = Arc::clone(&self.base);
         while let Some(branch_probability) = branch_probabilities.pop() {
-            let recursive_choice = Arc::new((self.recurse)(Arc::clone(&strat)));
+            let recursed = (self.recurse)(Arc::clone(&strat));
+            let recursive_choice = Arc::new(recursed.boxed());
             let non_recursive_choice = strat;
-            strat = Arc::new(
-                ::bool::weighted(branch_probability.min(0.9))
-                    .prop_ind_flat_map(move |branch| if branch {
-                        Arc::clone(&recursive_choice)
-                    } else {
-                        Arc::clone(&non_recursive_choice)
-                    }).boxed());
+            let branch_cond = ::bool::weighted(branch_probability.min(0.9));
+            let branch = IndFlatten(Map::new(branch_cond, BranchFn {
+                yes: recursive_choice,
+                no: non_recursive_choice,
+            }));
+            strat = Arc::new(branch.boxed());
         }
 
         strat.new_value(runner)
@@ -142,11 +193,8 @@ mod test {
         let mut max_depth = 0;
         let mut max_count = 0;
 
-        let strat = Just(Tree::Leaf).prop_recursive(
-            4, 64, 16,
-            |element| ::collection::vec(element, 8..16)
-                .prop_map(Tree::Branch).boxed());
-
+        let strat = Just(Tree::Leaf).prop_recursive(4, 64, 16,
+            |element| ::collection::vec(element, 8..16).prop_map(Tree::Branch));
 
         let mut runner = TestRunner::default();
         for _ in 0..65536 {

src/strategy/traits.rs

@@ -16,6 +16,10 @@ use rand::XorShiftRng;
 use strategy::*;
 use test_runner::*;
 
+//==============================================================================
+// Traits
+//==============================================================================
+
 /// A new [`ValueTree`] from a [`Strategy`] when [`Ok`] or otherwise [`Err`]
 /// when a new value-tree can not be produced for some reason such as
 /// in the case of filtering with a predicate which always returns false.
@@ -363,20 +367,19 @@ pub trait Strategy : fmt::Debug {
     ///       .prop_map(JsonNode::Array),
     ///     prop::collection::hash_map(".*", element, 0..16)
     ///       .prop_map(JsonNode::Map)
-    ///   ].boxed());
+    ///   ]);
     /// # }
     /// ```
-    fn prop_recursive<
-            F : Fn (Arc<BoxedStrategy<ValueFor<Self>>>)
-                    -> BoxedStrategy<ValueFor<Self>>>
+    fn prop_recursive<R, F>
         (self, depth: u32, desired_size: u32, expected_branch_size: u32, recurse: F)
         -> Recursive<BoxedStrategy<ValueFor<Self>>, F>
-    where Self : Sized + 'static {
-        Recursive {
-            base: Arc::new(self.boxed()),
-            recurse: Arc::new(recurse),
-            depth, desired_size, expected_branch_size,
-        }
+    where
+        Self : Sized + 'static,
+        F : Fn(Arc<BoxedStrategy<ValueFor<Self>>>) -> R,
+        R : Strategy + 'static,
+        R::Value : ValueTree<Value = ValueFor<Self>>
+    {
+        Recursive::new(self, depth, desired_size, expected_branch_size, recurse)
     }
 
     /// Shuffle the contents of the values produced by this strategy.
@@ -455,23 +458,6 @@ pub trait Strategy : fmt::Debug {
     }
 }
 
-macro_rules! proxy_strategy {
-    ($typ:ty $(, $lt:tt)*) => {
-        impl<$($lt,)* S : Strategy + ?Sized> Strategy for $typ {
-            type Value = S::Value;
-
-            fn new_value(&self, runner: &mut TestRunner) -> NewTree<Self> {
-                (**self).new_value(runner)
-            }
-        }
-    };
-}
-proxy_strategy!(Box<S>);
-proxy_strategy!(&'a S, 'a);
-proxy_strategy!(&'a mut S, 'a);
-proxy_strategy!(::std::rc::Rc<S>);
-proxy_strategy!(::std::sync::Arc<S>);
-
 /// A generated value and its associated shrinker.
 ///
 /// Conceptually, a `ValueTree` represents a spectrum between a "minimally
@@ -540,6 +526,57 @@ pub trait ValueTree {
     fn complicate(&mut self) -> bool;
 }
 
+//==============================================================================
+// NoShrink
+//==============================================================================
+
+/// Wraps a `Strategy` or `ValueTree` to suppress shrinking of generated
+/// values.
+///
+/// See `Strategy::no_shrink()` for more details.
+#[derive(Clone, Copy, Debug)]
+pub struct NoShrink<T>(T);
+
+impl<T : Strategy> Strategy for NoShrink<T> {
+    type Value = NoShrink<T::Value>;
+
+    fn new_value(&self, runner: &mut TestRunner) -> NewTree<Self> {
+        self.0.new_value(runner).map(NoShrink)
+    }
+}
+
+impl<T : ValueTree> ValueTree for NoShrink<T> {
+    type Value = T::Value;
+
+    fn current(&self) -> T::Value {
+        self.0.current()
+    }
+
+    fn simplify(&mut self) -> bool { false }
+    fn complicate(&mut self) -> bool { false }
+}
+
+//==============================================================================
+// Trait objects
+//==============================================================================
+
+macro_rules! proxy_strategy {
+    ($typ:ty $(, $lt:tt)*) => {
+        impl<$($lt,)* S : Strategy + ?Sized> Strategy for $typ {
+            type Value = S::Value;
+
+            fn new_value(&self, runner: &mut TestRunner) -> NewTree<Self> {
+                (**self).new_value(runner)
+            }
+        }
+    };
+}
+proxy_strategy!(Box<S>);
+proxy_strategy!(&'a S, 'a);
+proxy_strategy!(&'a mut S, 'a);
+proxy_strategy!(::std::rc::Rc<S>);
+proxy_strategy!(::std::sync::Arc<S>);
+
 impl<T : ValueTree + ?Sized> ValueTree for Box<T> {
     type Value = T::Value;
     fn current(&self) -> Self::Value { (**self).current() }
@@ -550,12 +587,11 @@ impl<T : ValueTree + ?Sized> ValueTree for Box<T> {
 /// A boxed `ValueTree`.
 type BoxedVT<T> = Box<ValueTree<Value = T>>;
 
-/// Shorthand for a boxed `Strategy` trait object as produced by
-/// `Strategy::boxed()`.
+/// A boxed `Strategy` trait object as produced by `Strategy::boxed()`.
 #[derive(Debug)]
 pub struct BoxedStrategy<T>(Box<Strategy<Value = BoxedVT<T>>>);
 
-/// Shorthand for a boxed `Strategy` trait object which is also `Sync` and
+/// A boxed `Strategy` trait object which is also `Sync` and
 /// `Send`, as produced by `Strategy::sboxed()`.
 #[derive(Debug)]
 pub struct SBoxedStrategy<T>(Box<Strategy<Value = BoxedVT<T>> + Sync + Send>);
@@ -606,31 +642,9 @@ where T::Value : 'static {
     }
 }
 
-/// Wraps a `Strategy` or `ValueTree` to suppress shrinking of generated
-/// values.
-///
-/// See `Strategy::no_shrink()` for more details.
-#[derive(Clone, Copy, Debug)]
-pub struct NoShrink<T>(T);
-
-impl<T : Strategy> Strategy for NoShrink<T> {
-    type Value = NoShrink<T::Value>;
-
-    fn new_value(&self, runner: &mut TestRunner) -> NewTree<Self> {
-        self.0.new_value(runner).map(NoShrink)
-    }
-}
-
-impl<T : ValueTree> ValueTree for NoShrink<T> {
-    type Value = T::Value;
-
-    fn current(&self) -> T::Value {
-        self.0.current()
-    }
-
-    fn simplify(&mut self) -> bool { false }
-    fn complicate(&mut self) -> bool { false }
-}
+//==============================================================================
+// Sanity checking
+//==============================================================================
 
 /// Options passed to `check_strategy_sanity()`.
 #[derive(Clone, Copy, Debug)]