2 files changed, 166 insertions, 1 deletions
diff --git a/tvix/castore/src/proto/mod.rs b/tvix/castore/src/proto/mod.rs
index 97ef1835881e..39c1bcc6faa0 100644
--- a/tvix/castore/src/proto/mod.rs
+++ b/tvix/castore/src/proto/mod.rs
@@ -179,6 +179,42 @@ impl Ord for node::Node {
     }
 }
 
+impl PartialOrd for FileNode {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for FileNode {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.get_name().cmp(other.get_name())
+    }
+}
+
+impl PartialOrd for SymlinkNode {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for SymlinkNode {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.get_name().cmp(other.get_name())
+    }
+}
+
+impl PartialOrd for DirectoryNode {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for DirectoryNode {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.get_name().cmp(other.get_name())
+    }
+}
+
 /// Accepts a name, and a mutable reference to the previous name.
 /// If the passed name is larger than the previous one, the reference is updated.
 /// If it's not, an error is returned.
@@ -303,6 +339,56 @@ impl Directory {
             i_symlinks: self.symlinks.iter().peekable(),
         };
     }
+
+    /// Adds the specified [node::Node] to the [Directory], preserving sorted entries.
+    /// This assumes the [Directory] to be sorted prior to adding the node.
+    ///
+    /// Inserting an element that already exists with the same name in the directory is not
+    /// supported.
+    pub fn add(&mut self, node: node::Node) {
+        debug_assert!(
+            !self.files.iter().any(|x| x.get_name() == node.get_name()),
+            "name already exists in files"
+        );
+        debug_assert!(
+            !self
+                .directories
+                .iter()
+                .any(|x| x.get_name() == node.get_name()),
+            "name already exists in directories"
+        );
+        debug_assert!(
+            !self
+                .symlinks
+                .iter()
+                .any(|x| x.get_name() == node.get_name()),
+            "name already exists in symlinks"
+        );
+
+        match node {
+            node::Node::File(node) => {
+                let pos = self
+                    .files
+                    .binary_search(&node)
+                    .expect_err("Tvix bug: dir entry with name already exists");
+                self.files.insert(pos, node);
+            }
+            node::Node::Directory(node) => {
+                let pos = self
+                    .directories
+                    .binary_search(&node)
+                    .expect_err("Tvix bug: dir entry with name already exists");
+                self.directories.insert(pos, node);
+            }
+            node::Node::Symlink(node) => {
+                let pos = self
+                    .symlinks
+                    .binary_search(&node)
+                    .expect_err("Tvix bug: dir entry with name already exists");
+                self.symlinks.insert(pos, node);
+            }
+        }
+    }
 }
 
 impl StatBlobResponse {
diff --git a/tvix/castore/src/proto/tests/directory.rs b/tvix/castore/src/proto/tests/directory.rs
index 5fda394775b5..81b73a048d52 100644
--- a/tvix/castore/src/proto/tests/directory.rs
+++ b/tvix/castore/src/proto/tests/directory.rs
@@ -1,5 +1,6 @@
 use crate::proto::{
-    Directory, DirectoryNode, FileNode, SymlinkNode, ValidateDirectoryError, ValidateNodeError,
+    node, Directory, DirectoryNode, FileNode, SymlinkNode, ValidateDirectoryError,
+    ValidateNodeError,
 };
 
 use hex_literal::hex;
@@ -371,3 +372,81 @@ fn validate_overflow() {
         _ => panic!("unexpected error"),
     }
 }
+
+#[test]
+fn add_nodes_to_directory() {
+    let mut d = Directory {
+        ..Default::default()
+    };
+
+    d.add(node::Node::Directory(DirectoryNode {
+        name: "b".into(),
+        digest: DUMMY_DIGEST.to_vec().into(),
+        size: 1,
+    }));
+    d.add(node::Node::Directory(DirectoryNode {
+        name: "a".into(),
+        digest: DUMMY_DIGEST.to_vec().into(),
+        size: 1,
+    }));
+    d.add(node::Node::Directory(DirectoryNode {
+        name: "z".into(),
+        digest: DUMMY_DIGEST.to_vec().into(),
+        size: 1,
+    }));
+
+    d.add(node::Node::File(FileNode {
+        name: "f".into(),
+        digest: DUMMY_DIGEST.to_vec().into(),
+        size: 1,
+        executable: true,
+    }));
+    d.add(node::Node::File(FileNode {
+        name: "c".into(),
+        digest: DUMMY_DIGEST.to_vec().into(),
+        size: 1,
+        executable: true,
+    }));
+    d.add(node::Node::File(FileNode {
+        name: "g".into(),
+        digest: DUMMY_DIGEST.to_vec().into(),
+        size: 1,
+        executable: true,
+    }));
+
+    d.add(node::Node::Symlink(SymlinkNode {
+        name: "t".into(),
+        target: "a".into(),
+    }));
+    d.add(node::Node::Symlink(SymlinkNode {
+        name: "o".into(),
+        target: "a".into(),
+    }));
+    d.add(node::Node::Symlink(SymlinkNode {
+        name: "e".into(),
+        target: "a".into(),
+    }));
+
+    d.validate().expect("directory should be valid");
+}
+
+#[test]
+#[cfg_attr(not(debug_assertions), ignore)]
+#[should_panic = "name already exists in directories"]
+fn add_duplicate_node_to_directory_panics() {
+    let mut d = Directory {
+        ..Default::default()
+    };
+
+    d.add(node::Node::Directory(DirectoryNode {
+        name: "a".into(),
+        digest: DUMMY_DIGEST.to_vec().into(),
+        size: 1,
+    }));
+    d.add(node::Node::File(FileNode {
+        name: "a".into(),
+        digest: DUMMY_DIGEST.to_vec().into(),
+        size: 1,
+        executable: true,
+    }));
+}