feat(tvix/glue): use TvixStoreIO as derivation builtin state r/7410

We propagate a `TvixStoreIO` as the `state` of our derivation-specific
builtins in the glue crate.

The evaluators `io_handle` itself is using a Rc<dyn EvalIO>.

An earlier version of TvixStoreIO was also introducing generics over the
different internal services themselves, but we opted for instead
hardcoding this to Arc<dyn …> for the sake of less macro voodoo.

Change-Id: I535c476f06b840858fa3070c4a237ece47f7a15b
Reviewed-on: https://cl.tvl.fyi/c/depot/+/10636
Reviewed-by: raitobezarius <tvl@lahfa.xyz>
Autosubmit: raitobezarius <tvl@lahfa.xyz>
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

6 files changed, 106 insertions, 86 deletions

diff --git a/tvix/glue/benches/eval.rs b/tvix/glue/benches/eval.rs
index f5c9813c9063..12a0f958963e 100644
--- a/tvix/glue/benches/eval.rs
+++ b/tvix/glue/benches/eval.rs
@@ -1,12 +1,13 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use lazy_static::lazy_static;
-use std::{cell::RefCell, env, rc::Rc, sync::Arc, time::Duration};
+use std::{env, rc::Rc, sync::Arc, time::Duration};
 use tvix_castore::{
     blobservice::{BlobService, MemoryBlobService},
     directoryservice::{DirectoryService, MemoryDirectoryService},
 };
+use tvix_eval::EvalIO;
 use tvix_glue::{
-    builtins::add_derivation_builtins, configure_nix_path, known_paths::KnownPaths,
+    builtins::add_derivation_builtins, configure_nix_path, tvix_io::TvixIO,
     tvix_store_io::TvixStoreIO,
 };
 use tvix_store::pathinfoservice::{MemoryPathInfoService, PathInfoService};
@@ -26,10 +27,21 @@ fn interpret(code: &str) {
     // TODO: this is a bit annoying.
     // It'd be nice if we could set this up once and then run evaluate() with a
     // piece of code. b/262
-    let mut eval = tvix_eval::Evaluation::new_impure();
 
-    let known_paths: Rc<RefCell<KnownPaths>> = Default::default();
-    add_derivation_builtins(&mut eval, known_paths);
+    // We assemble a complete store in memory.
+    let tvix_store_io = Rc::new(TvixStoreIO::new(
+        BLOB_SERVICE.clone(),
+        DIRECTORY_SERVICE.clone(),
+        PATH_INFO_SERVICE.clone(),
+        TOKIO_RUNTIME.handle().clone(),
+    ));
+
+    let mut eval = tvix_eval::Evaluation::new(
+        Box::new(TvixIO::new(tvix_store_io.clone() as Rc<dyn EvalIO>)) as Box<dyn EvalIO>,
+        true,
+    );
+
+    add_derivation_builtins(&mut eval, tvix_store_io);
     configure_nix_path(
         &mut eval,
         // The benchmark requires TVIX_BENCH_NIX_PATH to be set, so barf out
@@ -37,13 +49,6 @@ fn interpret(code: &str) {
         &Some(env::var("TVIX_BENCH_NIX_PATH").expect("TVIX_BENCH_NIX_PATH must be set")),
     );
 
-    eval.io_handle = Box::new(tvix_glue::tvix_io::TvixIO::new(TvixStoreIO::new(
-        BLOB_SERVICE.clone(),
-        DIRECTORY_SERVICE.clone(),
-        PATH_INFO_SERVICE.clone(),
-        TOKIO_RUNTIME.handle().clone(),
-    )));
-
     let result = eval.evaluate(code, None);
 
     assert!(result.errors.is_empty());
diff --git a/tvix/glue/src/builtins/derivation.rs b/tvix/glue/src/builtins/derivation.rs
index c5ed8ec2fffe..180567b8c4a8 100644
--- a/tvix/glue/src/builtins/derivation.rs
+++ b/tvix/glue/src/builtins/derivation.rs
@@ -1,10 +1,9 @@
 //! Implements `builtins.derivation`, the core of what makes Nix build packages.
 use crate::builtins::DerivationError;
-use crate::known_paths::KnownPaths;
+use crate::tvix_store_io::TvixStoreIO;
 use bstr::BString;
 use nix_compat::derivation::{Derivation, Output};
 use nix_compat::nixhash;
-use std::cell::RefCell;
 use std::collections::{btree_map, BTreeSet};
 use std::rc::Rc;
 use tvix_eval::builtin_macros::builtins;
@@ -121,7 +120,7 @@ fn handle_fixed_output(
     Ok(None)
 }
 
-#[builtins(state = "Rc<RefCell<KnownPaths>>")]
+#[builtins(state = "Rc<TvixStoreIO>")]
 pub(crate) mod derivation_builtins {
     use std::collections::BTreeMap;
 
@@ -152,7 +151,7 @@ pub(crate) mod derivation_builtins {
     /// use the higher-level `builtins.derivation` instead.
     #[builtin("derivationStrict")]
     async fn builtin_derivation_strict(
-        state: Rc<RefCell<KnownPaths>>,
+        state: Rc<TvixStoreIO>,
         co: GenCo,
         input: Value,
     ) -> Result<Value, ErrorKind> {
@@ -425,7 +424,7 @@ pub(crate) mod derivation_builtins {
         }
 
         populate_inputs(&mut drv, input_context);
-        let mut known_paths = state.borrow_mut();
+        let mut known_paths = state.as_ref().known_paths.borrow_mut();
 
         // At this point, derivation fields are fully populated from
         // eval data structures.
diff --git a/tvix/glue/src/builtins/mod.rs b/tvix/glue/src/builtins/mod.rs
index 22166055d067..d2378779319b 100644
--- a/tvix/glue/src/builtins/mod.rs
+++ b/tvix/glue/src/builtins/mod.rs
@@ -1,12 +1,14 @@
 //! Contains builtins that deal with the store or builder.
-use std::{cell::RefCell, rc::Rc};
 
-use crate::known_paths::KnownPaths;
+use std::rc::Rc;
+
+use crate::tvix_store_io::TvixStoreIO;
 
 mod derivation;
 mod derivation_error;
 
 pub use derivation_error::Error as DerivationError;
+use tvix_eval::EvalIO;
 
 /// Adds derivation-related builtins to the passed [tvix_eval::Evaluation].
 ///
@@ -14,12 +16,12 @@ pub use derivation_error::Error as DerivationError;
 ///
 /// As they need to interact with `known_paths`, we also need to pass in
 /// `known_paths`.
-pub fn add_derivation_builtins<IO>(
-    eval: &mut tvix_eval::Evaluation<IO>,
-    known_paths: Rc<RefCell<KnownPaths>>,
-) {
+pub fn add_derivation_builtins<IO>(eval: &mut tvix_eval::Evaluation<IO>, io: Rc<TvixStoreIO>)
+where
+    IO: AsRef<dyn EvalIO>,
+{
     eval.builtins
-        .extend(derivation::derivation_builtins::builtins(known_paths));
+        .extend(derivation::derivation_builtins::builtins(io));
 
     // Add the actual `builtins.derivation` from compiled Nix code
     eval.src_builtins
@@ -28,22 +30,36 @@ pub fn add_derivation_builtins<IO>(
 
 #[cfg(test)]
 mod tests {
+    use std::rc::Rc;
+
+    use crate::tvix_store_io::TvixStoreIO;
+
     use super::add_derivation_builtins;
-    use crate::known_paths::KnownPaths;
     use nix_compat::store_path::hash_placeholder;
-    use std::{cell::RefCell, rc::Rc};
     use test_case::test_case;
-    use tvix_eval::EvaluationResult;
+    use tvix_eval::{EvalIO, EvaluationResult};
+    use tvix_store::utils::construct_services;
 
     /// evaluates a given nix expression and returns the result.
     /// Takes care of setting up the evaluator so it knows about the
     // `derivation` builtin.
     fn eval(str: &str) -> EvaluationResult {
-        let mut eval = tvix_eval::Evaluation::new_impure();
-
-        let known_paths: Rc<RefCell<KnownPaths>> = Default::default();
-
-        add_derivation_builtins(&mut eval, known_paths.clone());
+        // We assemble a complete store in memory.
+        let runtime = tokio::runtime::Runtime::new().expect("Failed to build a Tokio runtime");
+        let (blob_service, directory_service, path_info_service) = runtime
+            .block_on(async { construct_services("memory://", "memory://", "memory://").await })
+            .expect("Failed to construct store services in memory");
+
+        let io = Rc::new(TvixStoreIO::new(
+            blob_service,
+            directory_service,
+            path_info_service.into(),
+            runtime.handle().clone(),
+        ));
+
+        let mut eval = tvix_eval::Evaluation::new(io.clone() as Rc<dyn EvalIO>, false);
+
+        add_derivation_builtins(&mut eval, io);
 
         // run the evaluation itself.
         eval.evaluate(str, None)
diff --git a/tvix/glue/src/tests/mod.rs b/tvix/glue/src/tests/mod.rs
index 9ae17c219ac9..b5b62c5d7e4c 100644
--- a/tvix/glue/src/tests/mod.rs
+++ b/tvix/glue/src/tests/mod.rs
@@ -1,4 +1,4 @@
-use std::{cell::RefCell, rc::Rc, sync::Arc};
+use std::{rc::Rc, sync::Arc};
 
 use pretty_assertions::assert_eq;
 use std::path::PathBuf;
@@ -6,14 +6,12 @@ use tvix_castore::{
     blobservice::{BlobService, MemoryBlobService},
     directoryservice::{DirectoryService, MemoryDirectoryService},
 };
-use tvix_eval::Value;
+use tvix_eval::{EvalIO, Value};
 use tvix_store::pathinfoservice::{MemoryPathInfoService, PathInfoService};
 
 use rstest::rstest;
 
-use crate::{
-    builtins::add_derivation_builtins, known_paths::KnownPaths, tvix_store_io::TvixStoreIO,
-};
+use crate::{builtins::add_derivation_builtins, tvix_store_io::TvixStoreIO};
 
 fn eval_test(code_path: PathBuf, expect_success: bool) {
     assert_eq!(
@@ -32,8 +30,6 @@ fn eval_test(code_path: PathBuf, expect_success: bool) {
         return;
     }
 
-    let mut eval = tvix_eval::Evaluation::new_impure();
-
     let blob_service = Arc::new(MemoryBlobService::default()) as Arc<dyn BlobService>;
     let directory_service =
         Arc::new(MemoryDirectoryService::default()) as Arc<dyn DirectoryService>;
@@ -41,19 +37,18 @@ fn eval_test(code_path: PathBuf, expect_success: bool) {
         blob_service.clone(),
         directory_service.clone(),
     )) as Box<dyn PathInfoService>;
-    let runtime = tokio::runtime::Runtime::new().unwrap();
+    let tokio_runtime = tokio::runtime::Runtime::new().unwrap();
 
-    eval.io_handle = Box::new(TvixStoreIO::new(
+    let tvix_store_io = Rc::new(TvixStoreIO::new(
         blob_service,
         directory_service,
-        path_info_service,
-        runtime.handle().clone(),
+        path_info_service.into(),
+        tokio_runtime.handle().clone(),
     ));
-
-    let known_paths: Rc<RefCell<KnownPaths>> = Default::default();
+    let mut eval = tvix_eval::Evaluation::new(tvix_store_io.clone() as Rc<dyn EvalIO>, true);
 
     eval.strict = true;
-    add_derivation_builtins(&mut eval, known_paths.clone());
+    add_derivation_builtins(&mut eval, tvix_store_io.clone());
 
     let result = eval.evaluate(code, Some(code_path.clone()));
     let failed = match result.value {
diff --git a/tvix/glue/src/tvix_io.rs b/tvix/glue/src/tvix_io.rs
index 77dcb9291032..19e5dd0b41df 100644
--- a/tvix/glue/src/tvix_io.rs
+++ b/tvix/glue/src/tvix_io.rs
@@ -13,24 +13,27 @@ use std::path::{Path, PathBuf};
 use tvix_eval::{EvalIO, FileType};
 
 // TODO: Merge this together with TvixStoreIO?
-pub struct TvixIO<T: EvalIO> {
+pub struct TvixIO<T> {
     // Actual underlying [EvalIO] implementation.
     actual: T,
 }
 
-impl<T: EvalIO> TvixIO<T> {
+impl<T> TvixIO<T> {
     pub fn new(actual: T) -> Self {
         Self { actual }
     }
 }
 
-impl<T: EvalIO> EvalIO for TvixIO<T> {
+impl<T> EvalIO for TvixIO<T>
+where
+    T: AsRef<dyn EvalIO>,
+{
     fn store_dir(&self) -> Option<String> {
-        self.actual.store_dir()
+        self.actual.as_ref().store_dir()
     }
 
     fn import_path(&self, path: &Path) -> io::Result<PathBuf> {
-        let imported_path = self.actual.import_path(path)?;
+        let imported_path = self.actual.as_ref().import_path(path)?;
         Ok(imported_path)
     }
 
@@ -39,7 +42,7 @@ impl<T: EvalIO> EvalIO for TvixIO<T> {
             return Ok(true);
         }
 
-        self.actual.path_exists(path)
+        self.actual.as_ref().path_exists(path)
     }
 
     fn read_to_string(&self, path: &Path) -> io::Result<String> {
@@ -56,10 +59,10 @@ impl<T: EvalIO> EvalIO for TvixIO<T> {
             return Ok(include_str!("fetchurl.nix").to_string());
         }
 
-        self.actual.read_to_string(path)
+        self.actual.as_ref().read_to_string(path)
     }
 
     fn read_dir(&self, path: &Path) -> io::Result<Vec<(bytes::Bytes, FileType)>> {
-        self.actual.read_dir(path)
+        self.actual.as_ref().read_dir(path)
     }
 }
diff --git a/tvix/glue/src/tvix_store_io.rs b/tvix/glue/src/tvix_store_io.rs
index 251023ba3eab..bac0d5e3d3a5 100644
--- a/tvix/glue/src/tvix_store_io.rs
+++ b/tvix/glue/src/tvix_store_io.rs
@@ -2,8 +2,10 @@
 
 use nix_compat::store_path::StorePath;
 use std::{
+    cell::RefCell,
     io,
     path::{Path, PathBuf},
+    sync::Arc,
 };
 use tokio::io::AsyncReadExt;
 use tracing::{error, instrument, warn};
@@ -17,6 +19,8 @@ use tvix_castore::{
 };
 use tvix_store::pathinfoservice::PathInfoService;
 
+use crate::known_paths::KnownPaths;
+
 /// Implements [EvalIO], asking given [PathInfoService], [DirectoryService]
 /// and [BlobService].
 ///
@@ -24,23 +28,28 @@ use tvix_store::pathinfoservice::PathInfoService;
 /// This is to both cover cases of syntactically valid store paths, that exist
 /// on the filesystem (still managed by Nix), as well as being able to read
 /// files outside store paths.
-pub struct TvixStoreIO<BS, DS, PS> {
-    blob_service: BS,
-    directory_service: DS,
-    path_info_service: PS,
+///
+/// This structure is also directly used by the derivation builtins
+/// and tightly coupled to it.
+///
+/// In the future, we may revisit that coupling and figure out how to generalize this interface and
+/// hide this implementation detail of the glue itself so that glue can be used with more than one
+/// implementation of "Tvix Store IO" which does not necessarily bring the concept of blob service,
+/// directory service or path info service.
+pub struct TvixStoreIO {
+    blob_service: Arc<dyn BlobService>,
+    directory_service: Arc<dyn DirectoryService>,
+    path_info_service: Arc<dyn PathInfoService>,
     std_io: StdIO,
     tokio_handle: tokio::runtime::Handle,
+    pub(crate) known_paths: RefCell<KnownPaths>,
 }
 
-impl<BS, DS, PS> TvixStoreIO<BS, DS, PS>
-where
-    DS: AsRef<dyn DirectoryService>,
-    PS: AsRef<dyn PathInfoService>,
-{
+impl TvixStoreIO {
     pub fn new(
-        blob_service: BS,
-        directory_service: DS,
-        path_info_service: PS,
+        blob_service: Arc<dyn BlobService>,
+        directory_service: Arc<dyn DirectoryService>,
+        path_info_service: Arc<dyn PathInfoService>,
         tokio_handle: tokio::runtime::Handle,
     ) -> Self {
         Self {
@@ -49,6 +58,7 @@ where
             path_info_service,
             std_io: StdIO {},
             tokio_handle,
+            known_paths: Default::default(),
         }
     }
 
@@ -101,12 +111,7 @@ where
     }
 }
 
-impl<BS, DS, PS> EvalIO for TvixStoreIO<BS, DS, PS>
-where
-    BS: AsRef<dyn BlobService> + Clone,
-    DS: AsRef<dyn DirectoryService>,
-    PS: AsRef<dyn PathInfoService>,
-{
+impl EvalIO for TvixStoreIO {
     #[instrument(skip(self), ret, err)]
     fn path_exists(&self, path: &Path) -> io::Result<bool> {
         if let Ok((store_path, sub_path)) =
@@ -284,17 +289,17 @@ where
 
 #[cfg(test)]
 mod tests {
-    use std::{cell::RefCell, path::Path, rc::Rc, sync::Arc};
+    use std::{path::Path, rc::Rc, sync::Arc};
 
     use tempfile::TempDir;
     use tvix_castore::{
         blobservice::{BlobService, MemoryBlobService},
         directoryservice::{DirectoryService, MemoryDirectoryService},
     };
-    use tvix_eval::EvaluationResult;
-    use tvix_store::pathinfoservice::{MemoryPathInfoService, PathInfoService};
+    use tvix_eval::{EvalIO, EvaluationResult};
+    use tvix_store::pathinfoservice::MemoryPathInfoService;
 
-    use crate::{builtins::add_derivation_builtins, known_paths::KnownPaths};
+    use crate::builtins::add_derivation_builtins;
 
     use super::TvixStoreIO;
 
@@ -302,27 +307,24 @@ mod tests {
     /// Takes care of setting up the evaluator so it knows about the
     // `derivation` builtin.
     fn eval(str: &str) -> EvaluationResult {
-        let mut eval = tvix_eval::Evaluation::new_impure();
-
         let blob_service = Arc::new(MemoryBlobService::default()) as Arc<dyn BlobService>;
         let directory_service =
             Arc::new(MemoryDirectoryService::default()) as Arc<dyn DirectoryService>;
-        let path_info_service = Box::new(MemoryPathInfoService::new(
+        let path_info_service = Arc::new(MemoryPathInfoService::new(
             blob_service.clone(),
             directory_service.clone(),
-        )) as Box<dyn PathInfoService>;
+        ));
         let runtime = tokio::runtime::Runtime::new().unwrap();
 
-        eval.io_handle = Box::new(TvixStoreIO::new(
-            blob_service,
-            directory_service,
+        let io = Rc::new(TvixStoreIO::new(
+            blob_service.clone(),
+            directory_service.clone(),
             path_info_service,
             runtime.handle().clone(),
         ));
+        let mut eval = tvix_eval::Evaluation::new(io.clone() as Rc<dyn EvalIO>, true);
 
-        let known_paths: Rc<RefCell<KnownPaths>> = Default::default();
-
-        add_derivation_builtins(&mut eval, known_paths.clone());
+        add_derivation_builtins(&mut eval, io.clone());
 
         // run the evaluation itself.
         eval.evaluate(str, None)