4 files changed, 128 insertions, 0 deletions
diff --git a/tvix/eval/src/value/arbitrary.rs b/tvix/eval/src/value/arbitrary.rs
new file mode 100644
index 000000000000..56ccd3b28ab6
--- /dev/null
+++ b/tvix/eval/src/value/arbitrary.rs
@@ -0,0 +1,78 @@
+//! Support for configurable generation of arbitrary nix values
+
+use proptest::{prelude::*, strategy::BoxedStrategy};
+use std::{ffi::OsString, rc::Rc};
+
+use super::{NixAttrs, NixList, NixString, Value};
+
+#[derive(Clone)]
+pub enum Parameters {
+    Strategy(BoxedStrategy<Value>),
+    Parameters {
+        generate_internal_values: bool,
+        generate_functions: bool,
+        generate_nested: bool,
+    },
+}
+
+impl Default for Parameters {
+    fn default() -> Self {
+        Self::Parameters {
+            generate_internal_values: false,
+            generate_functions: false,
+            generate_nested: true,
+        }
+    }
+}
+
+impl Arbitrary for Value {
+    type Parameters = Parameters;
+    type Strategy = BoxedStrategy<Self>;
+
+    fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
+        match args {
+            Parameters::Strategy(s) => s,
+            Parameters::Parameters {
+                generate_internal_values,
+                generate_functions,
+                generate_nested,
+            } => {
+                if generate_internal_values || generate_functions {
+                    todo!("Generating internal values and functions not implemented yet")
+                } else if generate_nested {
+                    non_internal_value().boxed()
+                } else {
+                    leaf_value().boxed()
+                }
+            }
+        }
+    }
+}
+
+fn leaf_value() -> impl Strategy<Value = Value> {
+    use Value::*;
+
+    prop_oneof![
+        Just(Null),
+        any::<bool>().prop_map(Bool),
+        any::<i64>().prop_map(Integer),
+        any::<f64>().prop_map(Float),
+        any::<NixString>().prop_map(String),
+        any::<OsString>().prop_map(|s| Path(s.into())),
+    ]
+}
+
+fn non_internal_value() -> impl Strategy<Value = Value> {
+    leaf_value().prop_recursive(10, 256, 10, |inner| {
+        prop_oneof![
+            any_with::<NixAttrs>((
+                Default::default(),
+                Default::default(),
+                Parameters::Strategy(inner.clone())
+            ))
+            .prop_map(|a| Value::Attrs(Rc::new(a))),
+            any_with::<NixList>((Default::default(), Parameters::Strategy(inner)))
+                .prop_map(Value::List)
+        ]
+    })
+}
diff --git a/tvix/eval/src/value/attrs.rs b/tvix/eval/src/value/attrs.rs
index 86e52206b846..0551fe20d905 100644
--- a/tvix/eval/src/value/attrs.rs
+++ b/tvix/eval/src/value/attrs.rs
@@ -149,6 +149,35 @@ impl PartialEq for NixAttrs {
     }
 }
 
+#[cfg(feature = "arbitrary")]
+mod arbitrary {
+    use super::*;
+
+    use proptest::prelude::*;
+    use proptest::prop_oneof;
+    use proptest::strategy::{BoxedStrategy, Just, Strategy};
+
+    impl Arbitrary for NixAttrs {
+        type Parameters = <BTreeMap<NixString, Value> as Arbitrary>::Parameters;
+
+        type Strategy = BoxedStrategy<Self>;
+
+        fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
+            prop_oneof![
+                Just(Self(AttrsRep::Empty)),
+                (
+                    any_with::<Value>(args.2.clone()),
+                    any_with::<Value>(args.2.clone())
+                )
+                    .prop_map(|(name, value)| Self(AttrsRep::KV { name, value })),
+                any_with::<BTreeMap<NixString, Value>>(args)
+                    .prop_map(|map| Self(AttrsRep::Map(map)))
+            ]
+            .boxed()
+        }
+    }
+}
+
 impl NixAttrs {
     /// Return an attribute set containing the merge of the two
     /// provided sets. Keys from the `other` set have precedence.
diff --git a/tvix/eval/src/value/list.rs b/tvix/eval/src/value/list.rs
index c00ddd4191ea..19415c0d557f 100644
--- a/tvix/eval/src/value/list.rs
+++ b/tvix/eval/src/value/list.rs
@@ -20,6 +20,25 @@ impl Display for NixList {
     }
 }
 
+#[cfg(feature = "arbitrary")]
+mod arbitrary {
+    use proptest::{
+        prelude::{any_with, Arbitrary},
+        strategy::{BoxedStrategy, Strategy},
+    };
+
+    use super::*;
+
+    impl Arbitrary for NixList {
+        type Parameters = <Vec<Value> as Arbitrary>::Parameters;
+        type Strategy = BoxedStrategy<Self>;
+
+        fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
+            any_with::<Vec<Value>>(args).prop_map(Self).boxed()
+        }
+    }
+}
+
 impl NixList {
     pub fn concat(&self, other: &Self) -> Self {
         let mut lhs = self.clone();
diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs
index c4e8221e158d..9805af7f4bb2 100644
--- a/tvix/eval/src/value/mod.rs
+++ b/tvix/eval/src/value/mod.rs
@@ -3,6 +3,8 @@
 use std::rc::Rc;
 use std::{fmt::Display, path::PathBuf};
 
+#[cfg(feature = "arbitrary")]
+mod arbitrary;
 mod attrs;
 mod builtin;
 mod function;