refactor(tvix): Centralize arbitrary definitions r/1948

To make it easier for future tests to use the arbitrary specialisations
we've defined for some nix types, centralize them all in a single
arbitrary.hh header file.

Change-Id: I767f27949cfe7ec55c79901f7d7aa538d2f98c6b
Reviewed-on: https://cl.tvl.fyi/c/depot/+/2182
Tested-by: BuildkiteCI
Reviewed-by: kanepyork <rikingcoding@gmail.com>

4 files changed, 180 insertions, 175 deletions

diff --git a/third_party/nix/src/tests/arbitrary.hh b/third_party/nix/src/tests/arbitrary.hh
new file mode 100644
index 0000000000..026f8522cf
--- /dev/null
+++ b/third_party/nix/src/tests/arbitrary.hh
@@ -0,0 +1,176 @@
+#pragma once
+
+#include <rapidcheck.h>
+#include <rapidcheck/Gen.h>
+#include <rapidcheck/gen/Arbitrary.h>
+
+#include "libexpr/attr-set.hh"
+#include "libexpr/nixexpr.hh"
+#include "libstore/derivations.hh"
+#include "libutil/hash.hh"
+
+namespace nix::tests {
+static nix::SymbolTable* symbol_table;
+}
+
+namespace rc {
+
+using nix::Derivation;
+using nix::DerivationOutput;
+using nix::Pos;
+using nix::Value;
+
+template <>
+struct Arbitrary<nix::Symbol> {
+  static Gen<nix::Symbol> arbitrary() {
+    return gen::map(gen::arbitrary<std::string>(), [](std::string s) {
+      return nix::tests::symbol_table->Create(s);
+    });
+  }
+};
+
+template <>
+struct Arbitrary<Value> {
+  static Gen<nix::Value> arbitrary() {
+    return gen::build(gen::construct<Value>(),
+                      // TODO(grfn) generalize to more types
+                      gen::set(&Value::type, gen::just(nix::ValueType::tInt)),
+                      gen::set(&Value::integer, gen::arbitrary<int64_t>()));
+  }
+};
+
+template <>
+struct Arbitrary<Value*> {
+  static Gen<nix::Value*> arbitrary() {
+    return gen::apply(
+        [](nix::ValueType typ, int i) {
+          auto ret = new Value();
+          ret->type = typ;
+          ret->integer = i;
+          return ret;
+        },
+        gen::just(nix::ValueType::tInt), gen::arbitrary<int64_t>());
+  }
+};
+
+template <>
+struct Arbitrary<nix::Pos> {
+  static Gen<nix::Pos> arbitrary() {
+    return gen::construct<nix::Pos>(gen::arbitrary<nix::Symbol>(),
+                                    gen::arbitrary<unsigned int>(),
+                                    gen::arbitrary<unsigned int>());
+  }
+};
+
+template <>
+struct Arbitrary<nix::Pos*> {
+  static Gen<nix::Pos*> arbitrary() {
+    return gen::apply(
+        [](unsigned int line, unsigned int column) {
+          return new Pos({}, line, column);
+        },
+        gen::arbitrary<unsigned int>(), gen::arbitrary<unsigned int>());
+  }
+};
+
+template <>
+struct Arbitrary<nix::Attr> {
+  static Gen<nix::Attr> arbitrary() {
+    return gen::construct<nix::Attr>(gen::arbitrary<nix::Symbol>(),
+                                     gen::arbitrary<Value*>(),
+                                     gen::arbitrary<nix::Pos*>());
+  }
+};
+
+template <>
+struct Arbitrary<nix::Bindings> {
+  static Gen<nix::Bindings> arbitrary() {
+    return gen::map(gen::arbitrary<std::vector<nix::Attr>>(), [](auto attrs) {
+      nix::Bindings res;
+      for (const auto& attr : attrs) {
+        res.push_back(attr);
+      }
+      return res;
+    });
+  }
+};
+
+template <class K, class V>
+struct Arbitrary<absl::btree_map<K, V>> {
+  static Gen<absl::btree_map<K, V>> arbitrary() {
+    return gen::map(gen::arbitrary<std::map<K, V>>(), [](std::map<K, V> map) {
+      absl::btree_map<K, V> out_map;
+      out_map.insert(map.begin(), map.end());
+      return out_map;
+    });
+  }
+};
+
+template <>
+struct Arbitrary<nix::Base> {
+  static Gen<nix::Base> arbitrary() {
+    return gen::element(nix::Base16, nix::Base32, nix::Base64);
+  }
+};
+
+template <>
+struct Arbitrary<DerivationOutput> {
+  static Gen<DerivationOutput> arbitrary() {
+    return gen::apply(
+        [](std::string content, std::string path, std::string hash_algo,
+           bool recursive, bool include_algo_in_hash, nix::Base base) {
+          auto hash_type = nix::parseHashType(hash_algo);
+          auto hash = nix::hashString(hash_type, content);
+          return DerivationOutput(
+              path, recursive ? absl::StrCat("r:", hash_algo) : hash_algo,
+              hash.to_string(base, include_algo_in_hash));
+        },
+        gen::arbitrary<std::string>(),
+        gen::map(gen::arbitrary<std::string>(),
+                 [](std::string s) { return absl::StrCat("/", s); }),
+        gen::element<std::string>("md5", "sha1", "sha256", "sha512"),
+        gen::arbitrary<bool>(), gen::arbitrary<bool>(),
+        gen::arbitrary<nix::Base>());
+  }
+};
+
+template <>
+struct Arbitrary<Derivation> {
+  static Gen<Derivation> arbitrary() {
+    auto gen_path = gen::map(gen::arbitrary<std::string>(), [](std::string s) {
+      return absl::StrCat("/", s);
+    });
+
+    return gen::build<Derivation>(
+        gen::set(&nix::BasicDerivation::outputs),
+        gen::set(&nix::BasicDerivation::inputSrcs,
+                 gen::container<nix::PathSet>(gen_path)),
+        gen::set(&nix::BasicDerivation::platform),
+        gen::set(&nix::BasicDerivation::builder, gen_path),
+        gen::set(&nix::BasicDerivation::args),
+        gen::set(&nix::BasicDerivation::env),
+        gen::set(&Derivation::inputDrvs,
+                 gen::container<nix::DerivationInputs>(
+                     gen_path, gen::arbitrary<nix::StringSet>())));
+  }
+};
+
+template <>
+struct Arbitrary<nix::BuildResult::Status> {
+  static Gen<nix::BuildResult::Status> arbitrary() {
+    return gen::element(nix::BuildResult::Status::Built,
+                        nix::BuildResult::Status::Substituted,
+                        nix::BuildResult::Status::AlreadyValid,
+                        nix::BuildResult::Status::PermanentFailure,
+                        nix::BuildResult::Status::InputRejected,
+                        nix::BuildResult::Status::OutputRejected,
+                        nix::BuildResult::Status::TransientFailure,
+                        nix::BuildResult::Status::CachedFailure,
+                        nix::BuildResult::Status::TimedOut,
+                        nix::BuildResult::Status::MiscFailure,
+                        nix::BuildResult::Status::DependencyFailed,
+                        nix::BuildResult::Status::LogLimitExceeded,
+                        nix::BuildResult::Status::NotDeterministic);
+  }
+};
+}  // namespace rc
diff --git a/third_party/nix/src/tests/attr-set.cc b/third_party/nix/src/tests/attr-set.cc
index 84756c6094..35932bbeff 100644
--- a/third_party/nix/src/tests/attr-set.cc
+++ b/third_party/nix/src/tests/attr-set.cc
@@ -20,93 +20,9 @@
 #include "libexpr/nixexpr.hh"
 #include "libexpr/symbol-table.hh"
 #include "libexpr/value.hh"
+#include "tests/arbitrary.hh"
 #include "tests/dummy-store.hh"
 
-static nix::SymbolTable* symbol_table;
-
-namespace rc {
-using nix::Pos;
-using nix::Value;
-
-// TODO(grfn): These arbitrary implementations should be pulled out to a util
-// file sooner rather than later
-
-template <>
-struct Arbitrary<nix::Symbol> {
-  static Gen<nix::Symbol> arbitrary() {
-    return gen::map(gen::arbitrary<std::string>(),
-                    [](std::string s) { return symbol_table->Create(s); });
-  }
-};
-
-template <>
-struct Arbitrary<Value> {
-  static Gen<nix::Value> arbitrary() {
-    return gen::build(gen::construct<Value>(),
-                      // TODO(grfn) generalize to more types
-                      gen::set(&Value::type, gen::just(nix::ValueType::tInt)),
-                      gen::set(&Value::integer, gen::arbitrary<int64_t>()));
-  }
-};
-
-template <>
-struct Arbitrary<Value*> {
-  static Gen<nix::Value*> arbitrary() {
-    return gen::apply(
-        [](nix::ValueType typ, int i) {
-          auto ret = new Value();
-          ret->type = typ;
-          ret->integer = i;
-          return ret;
-        },
-        gen::just(nix::ValueType::tInt), gen::arbitrary<int64_t>());
-  }
-};
-
-template <>
-struct Arbitrary<nix::Pos> {
-  static Gen<nix::Pos> arbitrary() {
-    return gen::construct<nix::Pos>(gen::arbitrary<nix::Symbol>(),
-                                    gen::arbitrary<unsigned int>(),
-                                    gen::arbitrary<unsigned int>());
-  }
-};
-
-template <>
-struct Arbitrary<nix::Pos*> {
-  static Gen<nix::Pos*> arbitrary() {
-    return gen::apply(
-        [](unsigned int line, unsigned int column) {
-          return new Pos({}, line, column);
-        },
-        gen::arbitrary<unsigned int>(), gen::arbitrary<unsigned int>());
-  }
-};
-
-template <>
-struct Arbitrary<nix::Attr> {
-  static Gen<nix::Attr> arbitrary() {
-    return gen::construct<nix::Attr>(gen::arbitrary<nix::Symbol>(),
-                                     gen::arbitrary<Value*>(),
-                                     gen::arbitrary<nix::Pos*>());
-  }
-};
-
-template <>
-struct Arbitrary<nix::Bindings> {
-  static Gen<nix::Bindings> arbitrary() {
-    return gen::map(gen::arbitrary<std::vector<nix::Attr>>(), [](auto attrs) {
-      nix::Bindings res;
-      for (const auto& attr : attrs) {
-        res.push_back(attr);
-      }
-      return res;
-    });
-  }
-};
-
-}  // namespace rc
-
 namespace nix {
 
 using nix::tests::DummyStore;
@@ -118,7 +34,7 @@ class AttrSetTest : public ::testing::Test {
     nix::expr::InitGC();
     auto store = std::make_shared<DummyStore>();
     eval_state_ = new EvalState({"."}, ref<Store>(store));
-    symbol_table = &eval_state_->symbols;
+    tests::symbol_table = &eval_state_->symbols;
   }
 
   void assert_bindings_equal(nix::Bindings* lhs, nix::Bindings* rhs) {
diff --git a/third_party/nix/src/tests/derivations_test.cc b/third_party/nix/src/tests/derivations_test.cc
index 50a540e84a..6ebe17824c 100644
--- a/third_party/nix/src/tests/derivations_test.cc
+++ b/third_party/nix/src/tests/derivations_test.cc
@@ -16,73 +16,7 @@
 #include "libexpr/eval.hh"
 #include "libutil/hash.hh"
 #include "libutil/types.hh"
-
-namespace rc {
-
-using nix::Derivation;
-using nix::DerivationOutput;
-
-template <class K, class V>
-struct Arbitrary<absl::btree_map<K, V>> {
-  static Gen<absl::btree_map<K, V>> arbitrary() {
-    return gen::map(gen::arbitrary<std::map<K, V>>(), [](std::map<K, V> map) {
-      absl::btree_map<K, V> out_map;
-      out_map.insert(map.begin(), map.end());
-      return out_map;
-    });
-  }
-};
-
-template <>
-struct Arbitrary<nix::Base> {
-  static Gen<nix::Base> arbitrary() {
-    return gen::element(nix::Base16, nix::Base32, nix::Base64);
-  }
-};
-
-template <>
-struct Arbitrary<DerivationOutput> {
-  static Gen<DerivationOutput> arbitrary() {
-    return gen::apply(
-        [](std::string content, std::string path, std::string hash_algo,
-           bool recursive, bool include_algo_in_hash, nix::Base base) {
-          auto hash_type = nix::parseHashType(hash_algo);
-          auto hash = nix::hashString(hash_type, content);
-          return DerivationOutput(
-              path, recursive ? absl::StrCat("r:", hash_algo) : hash_algo,
-              hash.to_string(base, include_algo_in_hash));
-        },
-        gen::arbitrary<std::string>(),
-        gen::map(gen::arbitrary<std::string>(),
-                 [](std::string s) { return absl::StrCat("/", s); }),
-        gen::element<std::string>("md5", "sha1", "sha256", "sha512"),
-        gen::arbitrary<bool>(), gen::arbitrary<bool>(),
-        gen::arbitrary<nix::Base>());
-  }
-};
-
-template <>
-struct Arbitrary<Derivation> {
-  static Gen<Derivation> arbitrary() {
-    auto gen_path = gen::map(gen::arbitrary<std::string>(), [](std::string s) {
-      return absl::StrCat("/", s);
-    });
-
-    return gen::build<Derivation>(
-        gen::set(&nix::BasicDerivation::outputs),
-        gen::set(&nix::BasicDerivation::inputSrcs,
-                 gen::container<nix::PathSet>(gen_path)),
-        gen::set(&nix::BasicDerivation::platform),
-        gen::set(&nix::BasicDerivation::builder, gen_path),
-        gen::set(&nix::BasicDerivation::args),
-        gen::set(&nix::BasicDerivation::env),
-        gen::set(&Derivation::inputDrvs,
-                 gen::container<nix::DerivationInputs>(
-                     gen_path, gen::arbitrary<nix::StringSet>())));
-  }
-};
-
-}  // namespace rc
+#include "tests/arbitrary.hh"
 
 namespace nix {
 
diff --git a/third_party/nix/src/tests/store-api-test.cc b/third_party/nix/src/tests/store-api-test.cc
index 515613b799..259e4b991b 100644
--- a/third_party/nix/src/tests/store-api-test.cc
+++ b/third_party/nix/src/tests/store-api-test.cc
@@ -2,31 +2,10 @@
 
 #include <gtest/gtest.h>
 #include <rapidcheck/Assertions.h>
-#include <rapidcheck/Gen.h>
 #include <rapidcheck/gtest.h>
 
 #include "libproto/worker.pb.h"
-
-namespace rc {
-template <>
-struct Arbitrary<nix::BuildResult::Status> {
-  static Gen<nix::BuildResult::Status> arbitrary() {
-    return gen::element(nix::BuildResult::Status::Built,
-                        nix::BuildResult::Status::Substituted,
-                        nix::BuildResult::Status::AlreadyValid,
-                        nix::BuildResult::Status::PermanentFailure,
-                        nix::BuildResult::Status::InputRejected,
-                        nix::BuildResult::Status::OutputRejected,
-                        nix::BuildResult::Status::TransientFailure,
-                        nix::BuildResult::Status::CachedFailure,
-                        nix::BuildResult::Status::TimedOut,
-                        nix::BuildResult::Status::MiscFailure,
-                        nix::BuildResult::Status::DependencyFailed,
-                        nix::BuildResult::Status::LogLimitExceeded,
-                        nix::BuildResult::Status::NotDeterministic);
-  }
-};
-}  // namespace rc
+#include "tests/arbitrary.hh"
 
 namespace nix {