snix_castore/directoryservice/traversal/
bfs.rs

1use super::Error;
2use crate::{B3Digest, Directory, Node};
3use futures::StreamExt;
4use std::collections::{HashSet, VecDeque};
5use tracing::instrument;
6use tracing::warn;
7
8/// Traverses a [Directory] from the root to the children.
9///
10/// This is mostly BFS, but directories are only returned once.
11#[instrument(skip(get_directory))]
12pub fn root_to_leaves<F, Fut>(
13    root_directory_digest: B3Digest,
14    get_directory: F,
15) -> impl futures::Stream<Item = Result<Directory, Error>> + use<F, Fut>
16where
17    F: Fn(B3Digest) -> Fut + Sync + Send + 'static,
18    Fut: Future<Output = Result<Option<Directory>, crate::directoryservice::Error>> + Send,
19{
20    // The list of all directories that still need to be traversed. The next
21    // element is picked from the front, new elements are enqueued at the
22    // back.
23    let mut worklist_directory_digests = VecDeque::from([root_directory_digest]);
24    // The list of directory digests already sent to the consumer.
25    // We omit sending the same directories multiple times.
26    let mut sent_directory_digests: HashSet<B3Digest> = HashSet::new();
27
28    async_stream::try_stream! {
29        while let Some(current_directory_digest) = worklist_directory_digests.pop_front() {
30            let current_directory = match get_directory(current_directory_digest).await.map_err(|e| {
31                Error::GetFailure(current_directory_digest, e)
32            })? {
33                // the root node of the requested closure was not found, return an empty list
34                None if current_directory_digest == root_directory_digest => break,
35                // if a child directory of the closure is not there, we have an inconsistent store!
36                None => {
37                    Err(Error::NotFound(current_directory_digest))?;
38                    break;
39                }
40                Some(dir) => dir,
41            };
42
43            // We're about to send this directory, so let's avoid sending it again if a
44            // descendant has it.
45            sent_directory_digests.insert(current_directory_digest);
46
47            // enqueue all child directory digests to the work queue, as
48            // long as they're not part of the worklist or already sent.
49            // This panics if the digest looks invalid, it's supposed to be checked first.
50            for (_, child_directory_node) in current_directory.nodes() {
51                if let Node::Directory{digest: child_digest, ..} = child_directory_node {
52                    if worklist_directory_digests.contains(child_digest)
53                        || sent_directory_digests.contains(child_digest)
54                    {
55                        continue;
56                    }
57                    worklist_directory_digests.push_back(*child_digest);
58                }
59            }
60
61            yield current_directory;
62        }
63    }.boxed()
64}
snix_castore/directoryservice/traversal/bfs.rs

snix_castore/directoryservice/traversal/
bfs.rs
