feat(tvix/castore): add convenience `add` method to Directory r/7978

This adds `Directory::add` which is a convenience helper for adding
nodes into a `Directory` while preserving sorted order.

This implements `Ord` and `PartialOrd` for `FileNode`, `SymlinkNode`,
and `DirectoryNode` so `binary_search` can be used.

Change-Id: I94b86bdef5d0da55aa352e098988b9704cafca19
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11481
Autosubmit: Connor Brewster <cbrewster@hey.com>
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

1 files changed, 86 insertions, 0 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 {