codecraft 0.1.2

//! Forces that pull heavily linked containers together, laid over the tree layout.
use std::collections::HashMap;

use glam::Vec3;

use super::graph::{Graph, Id, Kind, Relation};
use super::layout::Layout;

/// Pull back towards where the tree put a container.
const ANCHOR: f32 = 1.2;
/// Pull between two containers per unit of link weight.
const LINK: f32 = 0.35;
/// Push between siblings that overlap.
const REPEL: f32 = 8.0;
const GAP: f32 = 0.1;
const DAMPING: f32 = 0.82;
const DT: f32 = 1.0 / 30.0;
/// Links heavier than this pull no harder, so one busy pair does not fold the map.
const MAX_WEIGHT: f32 = 8.0;
const SETTLED: f32 = 0.01;

pub struct Forces {
    displacement: HashMap<Id, Vec3>,
    velocity: HashMap<Id, Vec3>,
    /// Container pairs and how strongly their contents are tied.
    links: Vec<(Id, Id, f32)>,
    /// Containers sharing a parent, which keep out of each other's way.
    siblings: Vec<Vec<Id>>,
    settled: bool,
}

impl Forces {
    pub fn new(graph: &Graph) -> Self {
        let container_of = |id: Id| match graph.node(id).kind.is_container() {
            true => id,
            false => graph.parent(id).unwrap_or(id),
        };
        let mut weights: HashMap<(Id, Id), f32> = HashMap::new();
        for id in graph.ids() {
            for (other, relation, outgoing) in graph.relations(id) {
                if !outgoing {
                    continue;
                }
                let (a, b) = (container_of(id), container_of(other));
                if a == b || graph.ancestors(a).contains(&b) || graph.ancestors(b).contains(&a) {
                    continue;
                }
                let weight = match relation {
                    Relation::DependsOn => 3.0,
                    Relation::Implements => 1.5,
                    Relation::Uses => 1.0,
                    Relation::Contains => continue,
                };
                let key = if a < b { (a, b) } else { (b, a) };
                *weights.entry(key).or_default() += weight;
            }
        }
        let mut links: Vec<(Id, Id, f32)> = weights
            .into_iter()
            .map(|((a, b), w)| (a, b, w.min(MAX_WEIGHT)))
            .collect();
        links.sort_by_key(|&(a, b, _)| (a, b));

        let siblings = graph
            .ids()
            .filter(|&id| graph.node(id).kind.is_container())
            .map(|id| {
                graph
                    .children(id)
                    .into_iter()
                    .filter(|&child| graph.node(child).kind.is_container())
                    .collect::<Vec<Id>>()
            })
            .filter(|group| group.len() > 1)
            .collect();

        Self {
            displacement: HashMap::new(),
            velocity: HashMap::new(),
            links,
            siblings,
            settled: true,
        }
    }

    /// How many container pairs are tied.
    pub fn links(&self) -> usize {
        self.links.len()
    }

    /// Whether the last step hardly moved anything.
    pub fn is_settled(&self) -> bool {
        self.settled
    }

    /// One step at `strength`, moving the layout's positions.
    pub fn step(&mut self, graph: &Graph, layout: &mut Layout, strength: f32) {
        let mut force: HashMap<Id, Vec3> = HashMap::new();
        for (&id, displaced) in &self.displacement {
            *force.entry(id).or_default() -= *displaced * ANCHOR;
        }
        for &(a, b, weight) in &self.links {
            let apart = layout.position(b) - layout.position(a);
            let pull = apart * (LINK * strength * weight);
            *force.entry(a).or_default() += pull;
            *force.entry(b).or_default() -= pull;
        }
        for group in &self.siblings {
            for (i, &a) in group.iter().enumerate() {
                for &b in &group[i + 1..] {
                    let apart = layout.position(b) - layout.position(a);
                    let distance = apart.length();
                    let wanted = (layout.side_of(a) + layout.side_of(b)) * 0.5 + GAP;
                    if distance >= wanted {
                        continue;
                    }
                    let direction = match distance > 1e-3 {
                        true => apart / distance,
                        false => Vec3::X,
                    };
                    let push = direction * ((wanted - distance) * REPEL);
                    *force.entry(a).or_default() -= push;
                    *force.entry(b).or_default() += push;
                }
            }
        }

        let mut fastest = 0.0f32;
        for (id, force) in force {
            let node = graph.node(id);
            if node.kind == Kind::Workspace {
                continue;
            }
            let velocity = self.velocity.entry(id).or_default();
            *velocity = (*velocity + force * DT) * DAMPING;
            let displaced = self.displacement.entry(id).or_default();
            *displaced += *velocity * DT;
            // On a leash the length of the parent's footprint, so the tree still reads.
            let leash = match node.kind {
                Kind::Crate => 20.0,
                _ => graph
                    .parent(id)
                    .map(|parent| layout.radius_of(parent))
                    .unwrap_or(5.0)
                    .max(2.0),
            };
            if displaced.length() > leash {
                *displaced = displaced.normalize() * leash;
            }
            fastest = fastest.max(velocity.length());
        }
        self.settled = fastest < SETTLED;
        layout.displace(graph, &self.displacement);
    }

    /// Lets everything drift back to where the tree put it; false once it is there.
    pub fn relax(&mut self, graph: &Graph, layout: &mut Layout) -> bool {
        if self.displacement.is_empty() {
            return false;
        }
        self.velocity.clear();
        let mut biggest = 0.0f32;
        for displaced in self.displacement.values_mut() {
            *displaced *= 0.8;
            biggest = biggest.max(displaced.length());
        }
        if biggest < SETTLED {
            self.displacement.clear();
        }
        layout.displace(graph, &self.displacement);
        true
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::code::graph::Node;

    /// Two crates: `a` in the first is tied to `b` in the second; `c` beside `a` is tied to nothing.
    fn tied() -> (Graph, Id, Id, Id) {
        let mut graph = Graph::new("ws");
        let lower = graph.add(graph.root(), Node::new(Kind::Crate, "lower"));
        let upper = graph.add(graph.root(), Node::new(Kind::Crate, "upper"));
        graph.relate(upper, lower, Relation::DependsOn);
        let a = graph.add(lower, Node::new(Kind::Module, "a"));
        let c = graph.add(lower, Node::new(Kind::Module, "c"));
        let b = graph.add(upper, Node::new(Kind::Module, "b"));
        for i in 0..6 {
            let s = graph.add(a, Node::new(Kind::Struct, format!("S{i}")));
            let t = graph.add(b, Node::new(Kind::Trait, format!("T{i}")));
            graph.relate(s, t, Relation::Implements);
            graph.add(c, Node::new(Kind::Fn, format!("f{i}")));
        }
        (graph, a, b, c)
    }

    #[test]
    fn links_are_counted_between_containers_not_symbols() {
        let (graph, _, _, _) = tied();
        let forces = Forces::new(&graph);
        assert_eq!(
            forces.links(),
            2,
            "six impls fold into one module pair, plus the crate dependency"
        );
    }

    #[test]
    fn tied_modules_draw_together_and_the_outlier_stays_put() {
        let (graph, a, b, c) = tied();
        let mut layout = Layout::of(&graph);
        let mut forces = Forces::new(&graph);
        let before = layout.position(a).distance(layout.position(b));
        for _ in 0..300 {
            forces.step(&graph, &mut layout, 1.0);
        }
        let after = layout.position(a).distance(layout.position(b));
        assert!(after < before, "{after} < {before}");
        assert!(forces.is_settled());
        let moved = |id: Id| layout.position(id).distance(layout.rest(id));
        assert!(moved(a) > 0.1, "a was pulled");
        assert!(
            moved(c) - moved(graph.parent(c).unwrap()) < 1e-3,
            "c only went where its crate went"
        );
        let symbol = graph.children(a)[0];
        assert!(
            (layout.position(symbol) - layout.position(a) - (layout.rest(symbol) - layout.rest(a)))
                .length()
                < 1e-4,
            "a module's symbols move with it"
        );
    }

    #[test]
    fn relaxing_returns_everything_to_rest() {
        let (graph, a, _, _) = tied();
        let mut layout = Layout::of(&graph);
        let mut forces = Forces::new(&graph);
        for _ in 0..50 {
            forces.step(&graph, &mut layout, 1.0);
        }
        while forces.relax(&graph, &mut layout) {}
        assert!(layout.position(a).distance(layout.rest(a)) < 1e-3);
    }
}