refactor(tvix/castore): use Directory struct separate from proto one r/8484

This uses our own data type to deal with Directories in the castore model.

It makes some undesired states unrepresentable, removing the need for conversions and checking in various places:

 - In the protobuf, blake3 digests could have a wrong length, as proto doesn't know fixed-size fields. We now use `B3Digest`, which makes cloning cheaper, and removes the need to do size-checking everywhere.
 - In the protobuf, we had three different lists for `files`, `symlinks` and `directories`. This was mostly a protobuf size optimization, but made interacting with them a bit awkward. This has now been replaced with a list of enums, and convenience iterators to get various nodes, and add new ones.

Change-Id: I7b92691bb06d77ff3f58a5ccea94a22c16f84f04
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12057
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

5 files changed, 79 insertions, 97 deletions

diff --git a/tvix/glue/src/builtins/derivation.rs b/tvix/glue/src/builtins/derivation.rs
index b17b90466965..9074d418422e 100644
--- a/tvix/glue/src/builtins/derivation.rs
+++ b/tvix/glue/src/builtins/derivation.rs
@@ -179,9 +179,7 @@ pub(crate) mod derivation_builtins {
     use nix_compat::nixhash::CAHash;
     use nix_compat::store_path::{build_ca_path, hash_placeholder};
     use sha2::Sha256;
-    use tvix_castore::proto as castorepb;
-    use tvix_castore::proto::node::Node;
-    use tvix_castore::proto::FileNode;
+    use tvix_castore::directoryservice::{FileNode, Node};
     use tvix_eval::generators::Gen;
     use tvix_eval::{NixContext, NixContextElement, NixString};
     use tvix_store::proto::{NarInfo, PathInfo};
@@ -579,12 +577,10 @@ pub(crate) mod derivation_builtins {
                     })
                     .map_err(DerivationError::InvalidDerivation)?;
 
-            let root_node = Node::File(FileNode {
-                name: store_path.to_string().into(),
-                digest: blob_digest.into(),
-                size: blob_size,
-                executable: false,
-            });
+            let root_node = Node::File(
+                FileNode::new(store_path.to_string().into(), blob_digest, blob_size, false)
+                    .map_err(|e| ErrorKind::TvixError(Rc::new(e)))?,
+            );
 
             // calculate the nar hash
             let (nar_size, nar_sha256) = state
@@ -604,9 +600,7 @@ pub(crate) mod derivation_builtins {
             state
                 .path_info_service
                 .put(PathInfo {
-                    node: Some(castorepb::Node {
-                        node: Some(root_node),
-                    }),
+                    node: Some((&root_node).into()),
                     references: reference_paths
                         .iter()
                         .map(|x| bytes::Bytes::copy_from_slice(x.digest()))
diff --git a/tvix/glue/src/builtins/import.rs b/tvix/glue/src/builtins/import.rs
index 273be08ef7b6..aa0d0d54fc8a 100644
--- a/tvix/glue/src/builtins/import.rs
+++ b/tvix/glue/src/builtins/import.rs
@@ -2,6 +2,8 @@
 
 use crate::builtins::errors::ImportError;
 use std::path::Path;
+use tvix_castore::directoryservice::FileNode;
+use tvix_castore::directoryservice::Node;
 use tvix_castore::import::ingest_entries;
 use tvix_eval::{
     builtin_macros::builtins,
@@ -16,7 +18,7 @@ async fn filtered_ingest(
     co: GenCo,
     path: &Path,
     filter: Option<&Value>,
-) -> Result<tvix_castore::proto::node::Node, ErrorKind> {
+) -> Result<Node, ErrorKind> {
     let mut entries: Vec<walkdir::DirEntry> = vec![];
     let mut it = walkdir::WalkDir::new(path)
         .follow_links(false)
@@ -114,8 +116,6 @@ mod import_builtins {
     use nix_compat::store_path::StorePath;
     use sha2::Digest;
     use tokio::io::AsyncWriteExt;
-    use tvix_castore::proto::node::Node;
-    use tvix_castore::proto::FileNode;
     use tvix_eval::builtins::coerce_value_to_path;
     use tvix_eval::generators::Gen;
     use tvix_eval::{generators::GenCo, ErrorKind, Value};
@@ -214,13 +214,16 @@ mod import_builtins {
                     .tokio_handle
                     .block_on(async { blob_writer.close().await })?;
 
-                let root_node = Node::File(FileNode {
-                    // The name gets set further down, while constructing the PathInfo.
-                    name: "".into(),
-                    digest: blob_digest.into(),
-                    size: blob_size,
-                    executable: false,
-                });
+                let root_node = Node::File(
+                    FileNode::new(
+                        // The name gets set further down, while constructing the PathInfo.
+                        "".into(),
+                        blob_digest,
+                        blob_size,
+                        false,
+                    )
+                    .map_err(|e| tvix_eval::ErrorKind::TvixError(Rc::new(e)))?,
+                );
 
                 let ca_hash = if recursive_ingestion {
                     let (_nar_size, nar_sha256) = state
diff --git a/tvix/glue/src/fetchers/mod.rs b/tvix/glue/src/fetchers/mod.rs
index eb035a5a905c..9cab1adc04a3 100644
--- a/tvix/glue/src/fetchers/mod.rs
+++ b/tvix/glue/src/fetchers/mod.rs
@@ -12,8 +12,8 @@ use tracing::{instrument, warn, Span};
 use tracing_indicatif::span_ext::IndicatifSpanExt;
 use tvix_castore::{
     blobservice::BlobService,
-    directoryservice::DirectoryService,
-    proto::{node::Node, FileNode},
+    directoryservice::{DirectoryService, FileNode, Node},
+    ValidateNodeError,
 };
 use tvix_store::{nar::NarCalculationService, pathinfoservice::PathInfoService, proto::PathInfo};
 use url::Url;
@@ -331,12 +331,10 @@ where
 
                 // Construct and return the FileNode describing the downloaded contents.
                 Ok((
-                    Node::File(FileNode {
-                        name: vec![].into(),
-                        digest: blob_writer.close().await?.into(),
-                        size: blob_size,
-                        executable: false,
-                    }),
+                    Node::File(
+                        FileNode::new(vec![].into(), blob_writer.close().await?, blob_size, false)
+                            .map_err(|e| FetcherError::Io(std::io::Error::other(e.to_string())))?,
+                    ),
                     CAHash::Flat(actual_hash),
                     blob_size,
                 ))
@@ -531,12 +529,13 @@ where
 
                 // Construct and return the FileNode describing the downloaded contents,
                 // make it executable.
-                let root_node = Node::File(FileNode {
-                    name: vec![].into(),
-                    digest: blob_digest.into(),
-                    size: file_size,
-                    executable: true,
-                });
+                let root_node = Node::File(
+                    FileNode::new(vec![].into(), blob_digest, file_size, true).map_err(
+                        |e: ValidateNodeError| {
+                            FetcherError::Io(std::io::Error::other(e.to_string()))
+                        },
+                    )?,
+                );
 
                 Ok((root_node, CAHash::Nar(actual_hash), file_size))
             }
@@ -580,7 +579,7 @@ where
 
         // Construct the PathInfo and persist it.
         let path_info = PathInfo {
-            node: Some(tvix_castore::proto::Node { node: Some(node) }),
+            node: Some((&node).into()),
             references: vec![],
             narinfo: Some(tvix_store::proto::NarInfo {
                 nar_size,
@@ -598,7 +597,14 @@ where
             .await
             .map_err(|e| FetcherError::Io(e.into()))?;
 
-        Ok((store_path, path_info.node.unwrap().node.unwrap()))
+        Ok((
+            store_path,
+            (&path_info.node.unwrap().node.unwrap())
+                .try_into()
+                .map_err(|e: ValidateNodeError| {
+                    FetcherError::Io(std::io::Error::other(e.to_string()))
+                })?,
+        ))
     }
 }
 
diff --git a/tvix/glue/src/tvix_build.rs b/tvix/glue/src/tvix_build.rs
index e9eb1725ef3e..0f930bfe5099 100644
--- a/tvix/glue/src/tvix_build.rs
+++ b/tvix/glue/src/tvix_build.rs
@@ -10,7 +10,7 @@ use tvix_build::proto::{
     build_request::{AdditionalFile, BuildConstraints, EnvVar},
     BuildRequest,
 };
-use tvix_castore::proto::{self, node::Node};
+use tvix_castore::directoryservice::Node;
 
 /// These are the environment variables that Nix sets in its sandbox for every
 /// build.
@@ -109,10 +109,7 @@ pub(crate) fn derivation_to_build_request(
             .into_iter()
             .map(|(key, value)| EnvVar { key, value })
             .collect(),
-        inputs: inputs
-            .into_iter()
-            .map(|n| proto::Node { node: Some(n) })
-            .collect(),
+        inputs: inputs.iter().map(Into::into).collect(),
         inputs_dir: nix_compat::store_path::STORE_DIR[1..].into(),
         constraints,
         working_dir: "build".into(),
@@ -200,10 +197,8 @@ mod test {
         build_request::{AdditionalFile, BuildConstraints, EnvVar},
         BuildRequest,
     };
-    use tvix_castore::{
-        fixtures::DUMMY_DIGEST,
-        proto::{self, node::Node, DirectoryNode},
-    };
+    use tvix_castore::directoryservice::{DirectoryNode, Node};
+    use tvix_castore::fixtures::DUMMY_DIGEST;
 
     use crate::tvix_build::NIX_ENVIRONMENT_VARS;
 
@@ -211,11 +206,14 @@ mod test {
     use lazy_static::lazy_static;
 
     lazy_static! {
-        static ref INPUT_NODE_FOO: Node = Node::Directory(DirectoryNode {
-            name: Bytes::from("mp57d33657rf34lzvlbpfa1gjfv5gmpg-bar"),
-            digest: DUMMY_DIGEST.clone().into(),
-            size: 42,
-        });
+        static ref INPUT_NODE_FOO: Node = Node::Directory(
+            DirectoryNode::new(
+                Bytes::from("mp57d33657rf34lzvlbpfa1gjfv5gmpg-bar"),
+                DUMMY_DIGEST.clone(),
+                42,
+            )
+            .unwrap()
+        );
     }
 
     #[test]
@@ -263,9 +261,7 @@ mod test {
                 command_args: vec![":".into()],
                 outputs: vec!["nix/store/fhaj6gmwns62s6ypkcldbaj2ybvkhx3p-foo".into()],
                 environment_vars: expected_environment_vars,
-                inputs: vec![proto::Node {
-                    node: Some(INPUT_NODE_FOO.clone())
-                }],
+                inputs: vec![(&*INPUT_NODE_FOO).into()],
                 inputs_dir: "nix/store".into(),
                 constraints: Some(BuildConstraints {
                     system: derivation.system.clone(),
diff --git a/tvix/glue/src/tvix_store_io.rs b/tvix/glue/src/tvix_store_io.rs
index 54e3550fc940..c92a9c57e0c7 100644
--- a/tvix/glue/src/tvix_store_io.rs
+++ b/tvix/glue/src/tvix_store_io.rs
@@ -15,15 +15,12 @@ use tokio_util::io::SyncIoBridge;
 use tracing::{error, instrument, warn, Level, Span};
 use tracing_indicatif::span_ext::IndicatifSpanExt;
 use tvix_build::buildservice::BuildService;
-use tvix_castore::proto::node::Node;
 use tvix_eval::{EvalIO, FileType, StdIO};
 use tvix_store::nar::NarCalculationService;
 
 use tvix_castore::{
     blobservice::BlobService,
-    directoryservice::{self, DirectoryService},
-    proto::NamedNode,
-    B3Digest,
+    directoryservice::{self, DirectoryService, NamedNode, Node},
 };
 use tvix_store::{pathinfoservice::PathInfoService, proto::PathInfo};
 
@@ -122,7 +119,12 @@ impl TvixStoreIO {
             .await?
         {
             // if we have a PathInfo, we know there will be a root_node (due to validation)
-            Some(path_info) => path_info.node.expect("no node").node.expect("no node"),
+            Some(path_info) => path_info
+                .node
+                .as_ref()
+                .expect("no node")
+                .try_into()
+                .expect("invalid node"),
             // If there's no PathInfo found, this normally means we have to
             // trigger the build (and insert into PathInfoService, after
             // reference scanning).
@@ -284,19 +286,17 @@ impl TvixStoreIO {
 
                         // For each output, insert a PathInfo.
                         for output in &build_result.outputs {
-                            let root_node = output.node.as_ref().expect("invalid root node");
+                            let root_node = output.try_into().expect("invalid root node");
 
                             // calculate the nar representation
                             let (nar_size, nar_sha256) = self
                                 .nar_calculation_service
-                                .calculate_nar(root_node)
+                                .calculate_nar(&root_node)
                                 .await?;
 
                             // assemble the PathInfo to persist
                             let path_info = PathInfo {
-                                node: Some(tvix_castore::proto::Node {
-                                    node: Some(root_node.clone()),
-                                }),
+                                node: Some((&root_node).into()),
                                 references: vec![], // TODO: refscan
                                 narinfo: Some(tvix_store::proto::NarInfo {
                                     nar_size,
@@ -332,13 +332,11 @@ impl TvixStoreIO {
                         build_result
                             .outputs
                             .into_iter()
+                            .map(|output_node| Node::try_from(&output_node).expect("invalid node"))
                             .find(|output_node| {
-                                output_node.node.as_ref().expect("invalid node").get_name()
-                                    == store_path.to_string().as_bytes()
+                                output_node.get_name() == store_path.to_string().as_bytes()
                             })
                             .expect("build didn't produce the store path")
-                            .node
-                            .expect("invalid node")
                     }
                 }
             }
@@ -460,20 +458,12 @@ impl EvalIO for TvixStoreIO {
                         ))
                     }
                     Node::File(file_node) => {
-                        let digest: B3Digest =
-                            file_node.digest.clone().try_into().map_err(|_e| {
-                                error!(
-                                    file_node = ?file_node,
-                                    "invalid digest"
-                                );
-                                io::Error::new(
-                                    io::ErrorKind::InvalidData,
-                                    format!("invalid digest length in file node: {:?}", file_node),
-                                )
-                            })?;
-
                         self.tokio_handle.block_on(async {
-                            let resp = self.blob_service.as_ref().open_read(&digest).await?;
+                            let resp = self
+                                .blob_service
+                                .as_ref()
+                                .open_read(file_node.digest())
+                                .await?;
                             match resp {
                                 Some(blob_reader) => {
                                     // The VM Response needs a sync [std::io::Reader].
@@ -482,12 +472,12 @@ impl EvalIO for TvixStoreIO {
                                 }
                                 None => {
                                     error!(
-                                        blob.digest = %digest,
+                                        blob.digest = %file_node.digest(),
                                         "blob not found",
                                     );
                                     Err(io::Error::new(
                                         io::ErrorKind::NotFound,
-                                        format!("blob {} not found", &digest),
+                                        format!("blob {} not found", &file_node.digest()),
                                     ))
                                 }
                             }
@@ -543,16 +533,7 @@ impl EvalIO for TvixStoreIO {
                 match node {
                     Node::Directory(directory_node) => {
                         // fetch the Directory itself.
-                        let digest: B3Digest =
-                            directory_node.digest.clone().try_into().map_err(|_e| {
-                                io::Error::new(
-                                    io::ErrorKind::InvalidData,
-                                    format!(
-                                        "invalid digest length in directory node: {:?}",
-                                        directory_node
-                                    ),
-                                )
-                            })?;
+                        let digest = directory_node.digest().clone();
 
                         if let Some(directory) = self.tokio_handle.block_on(async {
                             self.directory_service.as_ref().get(&digest).await
@@ -560,9 +541,11 @@ impl EvalIO for TvixStoreIO {
                             let mut children: Vec<(bytes::Bytes, FileType)> = Vec::new();
                             for node in directory.nodes() {
                                 children.push(match node {
-                                    Node::Directory(e) => (e.name, FileType::Directory),
-                                    Node::File(e) => (e.name, FileType::Regular),
-                                    Node::Symlink(e) => (e.name, FileType::Symlink),
+                                    Node::Directory(e) => {
+                                        (e.get_name().clone(), FileType::Directory)
+                                    }
+                                    Node::File(e) => (e.get_name().clone(), FileType::Regular),
+                                    Node::Symlink(e) => (e.get_name().clone(), FileType::Symlink),
                                 })
                             }
                             Ok(children)