re_view_graph 0.29.0-alpha.6

A view that shows a graph (node-link diagram).
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

use re_chunk::LatestAtQuery;
use re_log_types::{EntityPath, Instance};
use re_sdk_types::archetypes::{self, GraphEdges};
use re_sdk_types::{self, components, datatypes};
use re_view::{DataResultQuery as _, RangeResultsExt as _};
use re_viewer_context::{
    self, IdentifiedViewSystem, ViewContext, ViewContextCollection, ViewQuery,
    ViewSystemExecutionError, ViewSystemIdentifier, VisualizerExecutionOutput, VisualizerQueryInfo,
    VisualizerSystem,
};

use crate::graph::NodeId;

#[derive(Default)]
pub struct EdgesVisualizer {
    pub data: ahash::HashMap<EntityPath, EdgeData>,
}

pub struct EdgeInstance {
    // We will need this in the future, when we want to select individual edges.
    pub instance: Instance,
    pub source: components::GraphNode,
    pub target: components::GraphNode,
    pub source_index: NodeId,
    pub target_index: NodeId,
}

pub struct EdgeData {
    pub graph_type: components::GraphType,
    pub edges: Vec<EdgeInstance>,
}

impl IdentifiedViewSystem for EdgesVisualizer {
    fn identifier() -> ViewSystemIdentifier {
        "GraphEdges".into()
    }
}

impl VisualizerSystem for EdgesVisualizer {
    fn visualizer_query_info(
        &self,
        _app_options: &re_viewer_context::AppOptions,
    ) -> VisualizerQueryInfo {
        VisualizerQueryInfo::from_archetype::<archetypes::GraphEdges>()
    }

    /// Populates the visualizer with data from the store.
    fn execute(
        &mut self,
        ctx: &ViewContext<'_>,
        query: &ViewQuery<'_>,
        _context_systems: &ViewContextCollection,
    ) -> Result<VisualizerExecutionOutput, ViewSystemExecutionError> {
        let timeline_query = LatestAtQuery::new(query.timeline, query.latest_at);

        // TODO(cmc): could we (improve and then) use reflection for this?
        re_sdk_types::static_assert_struct_has_fields!(
            datatypes::Utf8Pair,
            first: datatypes::Utf8,
            second: datatypes::Utf8,
        );
        const SOURCE: &str = "first";
        const TARGET: &str = "second";

        for (data_result, instruction) in query.iter_visualizer_instruction_for(Self::identifier())
        {
            let results = data_result
                .latest_at_with_blueprint_resolved_data::<archetypes::GraphEdges>(
                    ctx,
                    &timeline_query,
                    Some(instruction),
                );

            let all_edges =
                results.iter_as(query.timeline, GraphEdges::descriptor_edges().component);
            let graph_type = results
                .get_mono::<components::GraphType>(GraphEdges::descriptor_graph_type().component)
                .unwrap_or_default();

            let sources = all_edges
                .slice_from_struct_field::<String>(SOURCE)
                .map(|(_index, source)| source);
            let targets = all_edges
                .slice_from_struct_field::<String>(TARGET)
                .map(|(_index, target)| target);

            for (sources, targets) in itertools::izip!(sources, targets) {
                let edges = itertools::izip!(sources, targets)
                    .enumerate()
                    .map(|(i, (source, target))| {
                        let source = components::GraphNode(source.into());
                        let target = components::GraphNode(target.into());

                        let entity_path = &data_result.entity_path;
                        let source_index = NodeId::from_entity_node(entity_path, &source);
                        let target_index = NodeId::from_entity_node(entity_path, &target);

                        EdgeInstance {
                            instance: Instance::from(i as u64),
                            source,
                            target,
                            source_index,
                            target_index,
                        }
                    })
                    .collect();

                self.data.insert(
                    data_result.entity_path.clone(),
                    EdgeData { graph_type, edges },
                );
            }
        }

        // We're not using `re_renderer` here, so return an empty vector.
        // If you want to draw additional primitives here, you can emit re_renderer draw data here directly.
        Ok(VisualizerExecutionOutput::default())
    }
}