Simplify graph viewer rendering logic

crates/viewer/re_space_view_graph/src/graph/index.rs

@@ -3,25 +3,25 @@ use re_types::components;
 
 use super::GraphNodeHash;
 
-#[derive(Clone, Copy, PartialEq, Eq)]
-pub struct NodeIndex {
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub struct NodeId {
     pub entity_hash: EntityPathHash,
     pub node_hash: GraphNodeHash,
 }
 
-impl nohash_hasher::IsEnabled for NodeIndex {}
+impl nohash_hasher::IsEnabled for NodeId {}
 
 // We implement `Hash` manually, because `nohash_hasher` requires a single call to `state.write_*`.
 // More info: https://crates.io/crates/nohash-hasher
-impl std::hash::Hash for NodeIndex {
+impl std::hash::Hash for NodeId {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
         // TODO(grtlr): Consider using `write_usize` here, to further decrease the risk of collision.
         let combined = self.entity_hash.hash64() << 32 | self.node_hash.hash64();
         state.write_u64(combined);
     }
 }
 
-impl NodeIndex {
+impl NodeId {
     pub fn from_entity_node(entity_path: &EntityPath, node: &components::GraphNode) -> Self {
         Self {
             entity_hash: entity_path.hash(),
@@ -30,7 +30,7 @@ impl NodeIndex {
     }
 }
 
-impl std::fmt::Debug for NodeIndex {
+impl std::fmt::Debug for NodeId {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "NodeIndex({:?}@{:?})", self.node_hash, self.entity_hash)
     }

crates/viewer/re_space_view_graph/src/graph/mod.rs

@@ -1,87 +1,185 @@
 mod hash;
+use std::hash::{Hash as _, Hasher as _};
+
+use egui::{Pos2, Vec2};
 pub(crate) use hash::GraphNodeHash;
 mod index;
-pub(crate) use index::NodeIndex;
+pub(crate) use index::NodeId;
 
-use re_types::components::{GraphNode, GraphType};
+use re_chunk::EntityPath;
+use re_log_types::Instance;
+use re_types::{components::GraphType, ArrowString};
 
-use crate::visualizers::{EdgeData, EdgeInstance, NodeData, NodeInstance};
+use crate::{
+    ui::DrawableLabel,
+    visualizers::{EdgeData, NodeData},
+};
 
-pub struct NodeInstanceImplicit {
-    pub node: GraphNode,
-    pub index: NodeIndex,
+pub enum Node {
+    Explicit {
+        id: NodeId,
+        instance: Instance,
+        node: ArrowString,
+        position: Option<Pos2>,
+        label: DrawableLabel,
+    },
+    Implicit {
+        id: NodeId,
+        node: ArrowString,
+        label: DrawableLabel,
+    },
 }
 
-impl std::hash::Hash for NodeInstanceImplicit {
-    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        self.index.hash(state);
+impl Node {
+    pub fn id(&self) -> NodeId {
+        match self {
+            Self::Explicit { id, .. } | Self::Implicit { id, .. } => *id,
+        }
     }
+
+    pub fn label(&self) -> &DrawableLabel {
+        match self {
+            Self::Explicit { label, .. } | Self::Implicit { label, .. } => label,
+        }
+    }
+
+    pub fn size(&self) -> Vec2 {
+        match self {
+            Self::Explicit { label, .. } | Self::Implicit { label, .. } => label.size(),
+        }
+    }
+
+    pub fn position(&self) -> Option<Pos2> {
+        match self {
+            Self::Explicit { position, .. } => *position,
+            Self::Implicit { .. } => None,
+        }
+    }
+}
+
+pub struct Edge {
+    pub from: NodeId,
+    pub to: NodeId,
+    pub arrow: bool,
 }
 
-#[derive(Hash)]
-pub struct Graph<'a> {
-    explicit: &'a [NodeInstance],
-    implicit: Vec<NodeInstanceImplicit>,
-    edges: &'a [EdgeInstance],
+pub struct Graph {
+    entity: EntityPath,
+    nodes: Vec<Node>,
+    edges: Vec<Edge>,
     kind: GraphType,
 }
 
-impl<'a> Graph<'a> {
-    pub fn new(node_data: Option<&'a NodeData>, edge_data: Option<&'a EdgeData>) -> Self {
+impl Graph {
+    pub fn new<'a>(
+        ui: &egui::Ui,
+        entity: EntityPath,
+        node_data: Option<&'a NodeData>,
+        edge_data: Option<&'a EdgeData>,
+    ) -> Self {
         // We keep track of the nodes to find implicit nodes.
         let mut seen = ahash::HashSet::default();
 
-        let explicit = if let Some(data) = node_data {
+        let mut nodes: Vec<Node> = if let Some(data) = node_data {
             seen.extend(data.nodes.iter().map(|n| n.index));
-            data.nodes.as_slice()
+            data.nodes
+                .iter()
+                .map(|n| Node::Explicit {
+                    id: n.index,
+                    instance: n.instance,
+                    node: n.node.0 .0.clone(),
+                    position: n.position,
+                    label: DrawableLabel::from_label(ui, &n.label),
+                })
+                .collect()
         } else {
-            &[][..]
+            Vec::new()
         };
 
-        let (edges, implicit, kind) = if let Some(data) = edge_data {
-            let mut implicit = Vec::new();
+        let (edges, kind) = if let Some(data) = edge_data {
             for edge in &data.edges {
                 if !seen.contains(&edge.source_index) {
-                    implicit.push(NodeInstanceImplicit {
-                        node: edge.source.clone(),
-                        index: edge.source_index,
+                    nodes.push(Node::Implicit {
+                        id: edge.source_index,
+                        node: edge.source.0 .0.clone(),
+                        label: DrawableLabel::implicit_circle(),
                     });
                     seen.insert(edge.source_index);
                 }
                 if !seen.contains(&edge.target_index) {
-                    implicit.push(NodeInstanceImplicit {
-                        node: edge.target.clone(),
-                        index: edge.target_index,
+                    nodes.push(Node::Implicit {
+                        id: edge.target_index,
+                        node: edge.target.0 .0.clone(),
+                        label: DrawableLabel::implicit_circle(),
                     });
                     seen.insert(edge.target_index);
                 }
             }
-            (data.edges.as_slice(), implicit, Some(data.graph_type))
+
+            let es = data.edges.iter().map(|e| Edge {
+                from: e.source_index,
+                to: e.target_index,
+                arrow: data.graph_type == GraphType::Directed,
+            });
+
+            (es.collect(), data.graph_type)
         } else {
-            (&[][..], Vec::new(), None)
+            (Vec::new(), GraphType::default())
         };
 
         Self {
-            explicit,
-            implicit,
+            entity,
+            nodes,
             edges,
-            kind: kind.unwrap_or_default(),
+            kind,
         }
     }
 
-    pub fn nodes_explicit(&self) -> impl Iterator<Item = &NodeInstance> {
-        self.explicit.iter()
-    }
-
-    pub fn nodes_implicit(&self) -> impl Iterator<Item = &NodeInstanceImplicit> + '_ {
-        self.implicit.iter()
+    pub fn nodes(&self) -> &[Node] {
+        &self.nodes
     }
 
-    pub fn edges(&self) -> impl Iterator<Item = &EdgeInstance> {
-        self.edges.iter()
+    pub fn edges(&self) -> &[Edge] {
+        &self.edges
     }
 
     pub fn kind(&self) -> GraphType {
         self.kind
     }
+
+    pub fn entity(&self) -> &EntityPath {
+        &self.entity
+    }
+
+    #[deprecated]
+    pub fn size_hash(&self) -> u64 {
+        let mut hasher = ahash::AHasher::default();
+        for node in &self.nodes {
+            match node {
+                Node::Explicit {
+                    id,
+                    position,
+                    label,
+                    // The following fields can be ignored:
+                    instance: _instance,
+                    node: _node,
+                } => {
+                    id.hash(&mut hasher);
+
+                    position.as_ref().map(bytemuck::bytes_of).hash(&mut hasher);
+                    bytemuck::bytes_of(&label.size()).hash(&mut hasher);
+                }
+                Node::Implicit {
+                    id,
+                    label,
+                    // The following fields can be ignored:
+                    node: _node,
+                } => {
+                    id.hash(&mut hasher);
+                    bytemuck::bytes_of(&label.size()).hash(&mut hasher);
+                }
+            }
+        }
+        hasher.finish()
+    }
 }