From 9e809e21ccb1768567fc2516c5526ad0cdd56df0 Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Sun, 29 Jan 2023 18:41:12 +0100 Subject: feat(tvix/store): implement iteration over Directory nodes 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 Tested-by: BuildkiteCI --- tvix/store/src/proto.rs | 81 ++++++++++++++++++++++- tvix/store/src/tests/directory_nodes_iterator.rs | 82 ++++++++++++++++++++++++ tvix/store/src/tests/mod.rs | 1 + 3 files changed, 163 insertions(+), 1 deletion(-) create mode 100644 tvix/store/src/tests/directory_nodes_iterator.rs 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>, + i_files: Peekable>, + i_symlinks: Peekable>, +} + +/// 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(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 { + 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 = 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; -- cgit 1.4.1