bee_tangle/
traversal.rs

1// Copyright 2020-2021 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4//! Collection of Tangle traversal functions.
5
6// TODO: Refactor all of this into methods on `Tangle`.
7
8use crate::{metadata::MessageMetadata, storage::StorageBackend, tangle::Tangle, MessageRef};
9
10use bee_message::MessageId;
11
12use std::{collections::HashSet, future::Future};
13
14/// A Tangle walker that - given a starting vertex - visits all of its ancestors that are connected through
15/// either the *parent1* or the *parent2* edge. The walk continues as long as the visited vertices match a certain
16/// condition. For each visited vertex customized logic can be applied depending on the availability of the
17/// vertex. Each traversed vertex provides read access to its associated data and metadata.
18pub async fn visit_parents_depth_first<Fut, Match, Apply, ElseApply, MissingApply, B: StorageBackend>(
19    tangle: &Tangle<B>,
20    root: MessageId,
21    matches: Match,
22    mut apply: Apply,
23    mut else_apply: ElseApply,
24    mut missing_apply: MissingApply,
25) where
26    Fut: Future<Output = bool>,
27    Match: Fn(MessageId, MessageRef, MessageMetadata) -> Fut,
28    Apply: FnMut(&MessageId, &MessageRef, &MessageMetadata),
29    ElseApply: FnMut(&MessageId, &MessageRef, &MessageMetadata),
30    MissingApply: FnMut(&MessageId),
31{
32    let mut parents = vec![root];
33    let mut visited = HashSet::new();
34
35    while let Some(message_id) = parents.pop() {
36        if visited.insert(message_id) {
37            let msg_meta = tangle
38                .get_vertex(&message_id)
39                .await
40                .as_ref()
41                .and_then(|v| v.message_and_metadata().cloned());
42            match msg_meta {
43                Some((msg, meta)) => {
44                    if matches(message_id, msg.clone(), meta).await {
45                        apply(&message_id, &msg, &meta);
46
47                        parents.extend_from_slice(msg.parents());
48                    } else {
49                        else_apply(&message_id, &msg, &meta);
50                    }
51                }
52                None => {
53                    missing_apply(&message_id);
54                }
55            }
56        }
57    }
58}