Merge pull request #3 from grtlr/feat/graph-auto-layout

Feat/graph auto layout

Cargo.lock

@@ -2418,6 +2418,11 @@ version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
 
+[[package]]
+name = "fjadra"
+version = "0.1.0"
+source = "git+https://github.com/grtlr/fjadra.git#79c2e6dccbc84c2823325cf7bcbcfaa5730064d0"
+
 [[package]]
 name = "flatbuffers"
 version = "23.5.26"
@@ -6159,9 +6164,11 @@ version = "0.20.0-alpha.4+dev"
 dependencies = [
  "ahash",
  "egui",
+ "fjadra",
  "nohash-hasher",
  "re_chunk",
  "re_format",
+ "re_log",
  "re_log_types",
  "re_query",
  "re_renderer",

Cargo.toml

@@ -181,6 +181,7 @@ enumset = "1.0.12"
 env_logger = { version = "0.10", default-features = false }
 ffmpeg-sidecar = { version = "2.0.2", default-features = false }
 fixed = { version = "1.28", default-features = false }
+fjadra = { version = "0.1", default-features = false, git = "https://github.com/grtlr/fjadra.git" }
 flatbuffers = "23.0"
 futures-channel = "0.3"
 futures-util = { version = "0.3", default-features = false }

crates/viewer/re_space_view_graph/Cargo.toml

@@ -21,6 +21,7 @@ all-features = true
 [dependencies]
 re_chunk.workspace = true
 re_format.workspace = true
+re_log.workspace = true
 re_log_types.workspace = true
 re_query.workspace = true
 re_renderer.workspace = true
@@ -33,4 +34,5 @@ re_viewport_blueprint.workspace = true
 
 ahash.workspace = true
 egui.workspace = true
+fjadra.workspace = true
 nohash-hasher.workspace = true

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

@@ -3,7 +3,7 @@ use re_types::components;
 
 /// A 64 bit hash of [`components::GraphNode`] with very small risk of collision.
 #[derive(Copy, Clone, Eq, PartialOrd, Ord)]
-pub(crate) struct GraphNodeHash(Hash64);
+pub struct GraphNodeHash(Hash64);
 
 impl nohash_hasher::IsEnabled for GraphNodeHash {}
 

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

@@ -4,7 +4,7 @@ use re_types::components;
 use super::GraphNodeHash;
 
 #[derive(Clone, Copy, PartialEq, Eq)]
-pub(crate) struct NodeIndex {
+pub struct NodeIndex {
     pub entity_hash: EntityPathHash,
     pub node_hash: GraphNodeHash,
 }

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

@@ -2,3 +2,79 @@ mod hash;
 pub(crate) use hash::GraphNodeHash;
 mod index;
 pub(crate) use index::NodeIndex;
+
+use re_types::components::{GraphNode, GraphType};
+
+use crate::visualizers::{EdgeData, EdgeInstance, NodeData, NodeInstance};
+
+pub struct NodeInstanceImplicit {
+    pub node: GraphNode,
+    pub index: NodeIndex,
+}
+
+pub struct Graph<'a> {
+    explicit: &'a [NodeInstance],
+    implicit: Vec<NodeInstanceImplicit>,
+    edges: &'a [EdgeInstance],
+    kind: GraphType,
+}
+
+impl<'a> Graph<'a> {
+    pub fn new(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 {
+            seen.extend(data.nodes.iter().map(|n| n.index));
+            data.nodes.as_slice()
+        } else {
+            &[][..]
+        };
+
+        let (edges, implicit, kind) = if let Some(data) = edge_data {
+            let mut implicit = Vec::new();
+            for edge in &data.edges {
+                if !seen.contains(&edge.source_index) {
+                    implicit.push(NodeInstanceImplicit {
+                        node: edge.source.clone(),
+                        index: edge.source_index,
+                    });
+                    seen.insert(edge.source_index);
+                }
+                if !seen.contains(&edge.target_index) {
+                    implicit.push(NodeInstanceImplicit {
+                        node: edge.target.clone(),
+                        index: edge.target_index,
+                    });
+                    seen.insert(edge.target_index);
+                }
+            }
+            (data.edges.as_slice(), implicit, Some(data.graph_type))
+        } else {
+            (&[][..], Vec::new(), None)
+        };
+
+        Self {
+            explicit,
+            implicit,
+            edges,
+            kind: kind.unwrap_or_default(),
+        }
+    }
+
+    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 edges(&self) -> impl Iterator<Item = &EdgeInstance> {
+        self.edges.iter()
+    }
+
+    pub fn kind(&self) -> GraphType {
+        self.kind
+    }
+}

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

@@ -0,0 +1,117 @@
+use egui::{Pos2, Rect, Vec2};
+use fjadra as fj;
+
+use crate::{
+    graph::{Graph, NodeIndex},
+    ui::bounding_rect_from_iter,
+    visualizers::NodeInstance,
+};
+
+#[derive(Debug, PartialEq, Eq)]
+pub struct Layout {
+    extents: ahash::HashMap<NodeIndex, Rect>,
+}
+
+impl Layout {
+    pub fn bounding_rect(&self) -> Rect {
+        bounding_rect_from_iter(self.extents.values().copied())
+    }
+
+    pub fn get(&self, node: &NodeIndex) -> Option<Rect> {
+        self.extents.get(node).copied()
+    }
+
+    pub fn update(&mut self, node: &NodeIndex, rect: Rect) {
+        *self
+            .extents
+            .get_mut(node)
+            .expect("node should exist in layout") = rect;
+    }
+}
+
+impl<'a> From<&'a NodeInstance> for fj::Node {
+    fn from(instance: &'a NodeInstance) -> Self {
+        let mut node = Self::default();
+        if let Some(pos) = instance.position {
+            node = node.fixed_position(pos.x as f64, pos.y as f64);
+        }
+        node
+    }
+}
+
+pub struct ForceLayout {
+    simulation: fj::Simulation,
+    node_index: ahash::HashMap<NodeIndex, usize>,
+}
+
+impl ForceLayout {
+    pub fn new<'a>(graphs: impl Iterator<Item = &'a Graph<'a>> + Clone) -> Self {
+        let explicit = graphs
+            .clone()
+            .flat_map(|g| g.nodes_explicit().map(|n| (n.index, fj::Node::from(n))));
+        let implicit = graphs
+            .clone()
+            .flat_map(|g| g.nodes_implicit().map(|n| (n.index, fj::Node::default())));
+
+        let mut node_index = ahash::HashMap::default();
+        let all_nodes: Vec<fj::Node> = explicit
+            .chain(implicit)
+            .enumerate()
+            .map(|(i, n)| {
+                node_index.insert(n.0, i);
+                n.1
+            })
+            .collect();
+
+        let all_edges = graphs.flat_map(|g| {
+            g.edges()
+                .map(|e| (node_index[&e.source_index], node_index[&e.target_index]))
+        });
+
+        // TODO(grtlr): Currently we guesstimate good forces. Eventually these should be exposed as blueprints.
+        let simulation = fj::SimulationBuilder::default()
+            .with_alpha_decay(0.01) // TODO(grtlr): slows down the simulation for demo
+            .build(all_nodes)
+            .add_force("link", fj::Link::new(all_edges))
+            .add_force("charge", fj::ManyBody::new())
+            // TODO(grtlr): detect if we need a center force or a position force, depending on if the graph is disjoint
+            // or not. More generally: how much heuristics do we want to bake in here?
+            .add_force("center", fj::Center::new());
+
+        Self {
+            simulation,
+            node_index,
+        }
+    }
+
+    pub fn init_layout(&self) -> Layout {
+        let positions = self.simulation.positions().collect::<Vec<_>>();
+        let mut extents = ahash::HashMap::default();
+
+        for (node, i) in &self.node_index {
+            let [x, y] = positions[*i];
+            let pos = Pos2::new(x as f32, y as f32);
+            let size = Vec2::ZERO;
+            let rect = Rect::from_min_size(pos, size);
+            extents.insert(*node, rect);
+        }
+
+        Layout { extents }
+    }
+
+    /// Returns `true` if finished.
+    pub fn tick(&mut self, layout: &mut Layout) -> bool {
+        self.simulation.tick(1);
+
+        let positions = self.simulation.positions().collect::<Vec<_>>();
+
+        for (node, extent) in layout.extents.iter_mut() {
+            let i = self.node_index[node];
+            let [x, y] = positions[i];
+            let pos = Pos2::new(x as f32, y as f32);
+            extent.set_center(pos);
+        }
+
+        self.simulation.finished()
+    }
+}

crates/viewer/re_space_view_graph/src/lib.rs

@@ -3,8 +3,8 @@
 //! A Space View that shows a graph (node-link diagram).
 
 mod graph;
+mod layout;
 mod properties;
-mod types;
 mod ui;
 mod view;
 mod visualizers;

crates/viewer/re_space_view_graph/src/properties.rs

@@ -1,10 +1,7 @@
 use re_types::blueprint::components::VisualBounds2D;
 use re_viewer_context::{SpaceViewStateExt as _, TypedComponentFallbackProvider};
 
-use crate::{
-    ui::{bounding_rect_from_iter, GraphSpaceViewState},
-    GraphSpaceView,
-};
+use crate::{ui::GraphSpaceViewState, GraphSpaceView};
 
 fn valid_bound(rect: &egui::Rect) -> bool {
     rect.is_finite() && rect.is_positive()
@@ -16,17 +13,9 @@ impl TypedComponentFallbackProvider<VisualBounds2D> for GraphSpaceView {
             return VisualBounds2D::default();
         };
 
-        let Some(layout) = &state.layout else {
-            return VisualBounds2D::default();
-        };
-
-        let default_scene_rect = bounding_rect_from_iter(layout.values());
-
-        if valid_bound(&default_scene_rect) {
-            default_scene_rect.into()
-        } else {
-            // Nothing in scene, probably.
-            VisualBounds2D::default()
+        match state.layout.bounding_rect() {
+            Some(rect) if valid_bound(&rect) => rect.into(),
+            _ => VisualBounds2D::default(),
         }
     }
 }

crates/viewer/re_space_view_graph/src/ui/canvas.rs

@@ -3,14 +3,14 @@ use egui::{
     Stroke, Ui, Vec2,
 };
 use re_chunk::EntityPath;
-use re_types::{
-    components::{GraphNode, Radius},
-    datatypes::Float32,
-};
+use re_types::{components::Radius, datatypes::Float32};
 use re_viewer_context::SpaceViewHighlights;
 use std::hash::Hash;
 
-use crate::types::{EdgeInstance, NodeInstance};
+use crate::{
+    graph::NodeInstanceImplicit,
+    visualizers::{EdgeInstance, NodeInstance},
+};
 
 use super::draw::{draw_edge, draw_entity, draw_explicit, draw_implicit};
 
@@ -178,26 +178,15 @@ pub struct Canvas<'a> {
 }
 
 impl<'a> Canvas<'a> {
-    /// Try to estimate a good starting `Rect` for a node that has not been layed out yet.
-    pub fn initial_rect(&self, node: &NodeInstance) -> Rect {
-        let size = node
-            .radius
-            .map(|r| self.context.radius_to_world(r))
-            .unwrap_or(0.0)
-            * 2.0;
-        let pos = node.position.unwrap_or(Pos2::ZERO);
-        Rect::from_center_size(pos, Vec2::splat(size))
-    }
-
     /// Draws a regular node, i.e. an explicit node instance.
     pub fn explicit_node(&mut self, pos: Pos2, node: &NodeInstance) -> Response {
         self.node_wrapper(node.index, pos, |ui, world_to_ui| {
             draw_explicit(ui, world_to_ui, node)
         })
     }
 
-    pub fn implicit_node(&mut self, pos: Pos2, node: &GraphNode) -> Response {
-        self.node_wrapper(node, pos, |ui, _| draw_implicit(ui, node))
+    pub fn implicit_node(&mut self, pos: Pos2, node: &NodeInstanceImplicit) -> Response {
+        self.node_wrapper(node.index, pos, |ui, _| draw_implicit(ui, node))
     }
 
     /// `pos` is the top-left position of the node in world coordinates.