feat(tvix/store): implement iteration over Directory nodes r/5790

This provides a `Directory.nodes()` function, returning an iterator over
all three node types in an ordered fashion.

Change-Id: Ib98696c03a9db8b6c613d6e2bf5587c1ae35133f
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7955
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI

3 files changed, 163 insertions, 1 deletions

diff --git a/tvix/store/src/proto.rs b/tvix/store/src/proto.rs
index 82a2d09ee5b7..fc7fa1eff9a1 100644
--- a/tvix/store/src/proto.rs
+++ b/tvix/store/src/proto.rs
@@ -1,6 +1,6 @@
 #![allow(clippy::derive_partial_eq_without_eq)]
 // https://github.com/hyperium/tonic/issues/1056
-use std::collections::HashSet;
+use std::{collections::HashSet, iter::Peekable};
 use thiserror::Error;
 
 use prost::Message;
@@ -274,4 +274,83 @@ impl Directory {
 
         Ok(())
     }
+
+    /// Allows iterating over all three nodes ([DirectoryNode], [FileNode],
+    /// [SymlinkNode]) in an ordered fashion, as long as the individual lists
+    /// are sorted (which can be checked by the [Directory::validate]).
+    pub fn nodes(&self) -> DirectoryNodesIterator {
+        return DirectoryNodesIterator {
+            i_directories: self.directories.iter().peekable(),
+            i_files: self.files.iter().peekable(),
+            i_symlinks: self.symlinks.iter().peekable(),
+        };
+    }
+}
+
+/// Struct to hold the state of an iterator over all nodes of a Directory.
+///
+/// Internally, this keeps peekable Iterators over all three lists of a
+/// Directory message.
+pub struct DirectoryNodesIterator<'a> {
+    // directory: &Directory,
+    i_directories: Peekable<std::slice::Iter<'a, DirectoryNode>>,
+    i_files: Peekable<std::slice::Iter<'a, FileNode>>,
+    i_symlinks: Peekable<std::slice::Iter<'a, SymlinkNode>>,
+}
+
+/// looks at two elements implementing NamedNode, and returns true if "left
+/// is smaller / comes first".
+///
+/// Some(_) is preferred over None.
+fn left_name_lt_right<A: NamedNode, B: NamedNode>(left: Option<&A>, right: Option<&B>) -> bool {
+    match left {
+        // if left is None, right always wins
+        None => false,
+        Some(left_inner) => {
+            // left is Some.
+            match right {
+                // left is Some, right is None - left wins.
+                None => true,
+                Some(right_inner) => {
+                    // both are Some - compare the name.
+                    return left_inner.get_name() < right_inner.get_name();
+                }
+            }
+        }
+    }
+}
+
+impl Iterator for DirectoryNodesIterator<'_> {
+    type Item = node::Node;
+
+    // next returns the next node in the Directory.
+    // we peek at all three internal iterators, and pick the one with the
+    // smallest name, to ensure lexicographical ordering.
+    // The individual lists are already known to be sorted.
+    fn next(&mut self) -> Option<Self::Item> {
+        if left_name_lt_right(self.i_directories.peek(), self.i_files.peek()) {
+            // i_directories is still in the game, compare with symlinks
+            if left_name_lt_right(self.i_directories.peek(), self.i_symlinks.peek()) {
+                self.i_directories
+                    .next()
+                    .map(|x| x.clone())
+                    .map(node::Node::Directory)
+            } else {
+                self.i_symlinks
+                    .next()
+                    .map(|x| x.clone())
+                    .map(node::Node::Symlink)
+            }
+        } else {
+            // i_files is still in the game, compare with symlinks
+            if left_name_lt_right(self.i_files.peek(), self.i_symlinks.peek()) {
+                self.i_files.next().map(|x| x.clone()).map(node::Node::File)
+            } else {
+                self.i_symlinks
+                    .next()
+                    .map(|x| x.clone())
+                    .map(node::Node::Symlink)
+            }
+        }
+    }
 }
diff --git a/tvix/store/src/tests/directory_nodes_iterator.rs b/tvix/store/src/tests/directory_nodes_iterator.rs
new file mode 100644
index 000000000000..8f591f0560ce
--- /dev/null
+++ b/tvix/store/src/tests/directory_nodes_iterator.rs
@@ -0,0 +1,82 @@
+use crate::proto::node::Node;
+use crate::proto::Directory;
+use crate::proto::DirectoryNode;
+use crate::proto::FileNode;
+use crate::proto::SymlinkNode;
+
+#[test]
+fn iterator() -> anyhow::Result<()> {
+    let d = Directory {
+        directories: vec![
+            DirectoryNode {
+                name: "c".to_string(),
+                ..DirectoryNode::default()
+            },
+            DirectoryNode {
+                name: "d".to_string(),
+                ..DirectoryNode::default()
+            },
+            DirectoryNode {
+                name: "h".to_string(),
+                ..DirectoryNode::default()
+            },
+            DirectoryNode {
+                name: "l".to_string(),
+                ..DirectoryNode::default()
+            },
+        ],
+        files: vec![
+            FileNode {
+                name: "b".to_string(),
+                ..FileNode::default()
+            },
+            FileNode {
+                name: "e".to_string(),
+                ..FileNode::default()
+            },
+            FileNode {
+                name: "g".to_string(),
+                ..FileNode::default()
+            },
+            FileNode {
+                name: "j".to_string(),
+                ..FileNode::default()
+            },
+        ],
+        symlinks: vec![
+            SymlinkNode {
+                name: "a".to_string(),
+                ..SymlinkNode::default()
+            },
+            SymlinkNode {
+                name: "f".to_string(),
+                ..SymlinkNode::default()
+            },
+            SymlinkNode {
+                name: "i".to_string(),
+                ..SymlinkNode::default()
+            },
+            SymlinkNode {
+                name: "k".to_string(),
+                ..SymlinkNode::default()
+            },
+        ],
+    };
+
+    let mut node_names: Vec<String> = vec![];
+
+    for node in d.nodes() {
+        match node {
+            Node::Directory(n) => node_names.push(n.name),
+            Node::File(n) => node_names.push(n.name),
+            Node::Symlink(n) => node_names.push(n.name),
+        };
+    }
+
+    assert_eq!(
+        vec!["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"],
+        node_names
+    );
+
+    Ok(())
+}
diff --git a/tvix/store/src/tests/mod.rs b/tvix/store/src/tests/mod.rs
index b340fcff5d32..4022c329798b 100644
--- a/tvix/store/src/tests/mod.rs
+++ b/tvix/store/src/tests/mod.rs
@@ -1,4 +1,5 @@
 mod directory;
+mod directory_nodes_iterator;
 mod directory_service;
 mod path_info_service;
 mod pathinfo;