steel_registry/advancement/
positioner.rs

1use crate::advancement::registry::AdvancementRegistry;
2
3pub enum PositionError {
4    InvalidRootIndex(usize),
5    RootMustHaveDisplay(String),
6}
7
8impl PositionError {
9    #[must_use]
10    pub fn get_message(&self) -> String {
11        match self {
12            PositionError::InvalidRootIndex(index) => format!("Invalid root index at {index}"),
13            PositionError::RootMustHaveDisplay(key) => format!("{key} must have display data"),
14        }
15    }
16}
17
18pub type NodePositionIdx = usize;
19
20/// Calculates advancement positions using the Reingold tilford algorithm X represents display depth and Y is determined by the tree layout
21pub fn run(tree: &mut AdvancementRegistry, root_index: usize) -> Result<(), PositionError> {
22    let Some(root_node) = tree.adv_nodes.get(root_index) else {
23        return Err(PositionError::InvalidRootIndex(root_index));
24    };
25    if !root_node.has_display() {
26        return Err(PositionError::RootMustHaveDisplay(
27            root_node.value().key.to_string(),
28        ));
29    }
30    //store everything inside a vector to not have to deal with pointer
31    let mut nodes: Vec<TreeNodePosition> = Vec::with_capacity(32);
32    let root_idx = nodes.len();
33    nodes.push(TreeNodePosition {
34        node: root_index,
35        parent: None,
36        previous_sibling: None,
37        child_index: 1,
38        children: Vec::new(),
39        ancestor: root_idx,
40        thread: None,
41        // We can put 0 here because it is recalculated later inside the second_walk.
42        x: 0,
43        y: -1.0,
44        r#mod: 0.0,
45        change: 0.0,
46        shift: 0.0,
47    });
48
49    let mut previous_idx = None;
50    for child in root_node.children() {
51        previous_idx = TreeNodePosition::add_child(&mut nodes, tree, root_idx, child, previous_idx);
52    }
53
54    TreeNodePosition::first_walk(&mut nodes, root_idx);
55
56    let root_y = nodes[root_idx].y;
57    let min = TreeNodePosition::second_walk(&mut nodes, root_idx, 0.0, 0, root_y);
58
59    if min < 0.0 {
60        TreeNodePosition::normalize_y(&mut nodes, -min);
61    }
62
63    TreeNodePosition::set_position(tree, &nodes, root_idx);
64    Ok(())
65}
66
67/// Stores tree links as vector indices
68struct TreeNodePosition {
69    node: usize,
70    parent: Option<NodePositionIdx>,
71    previous_sibling: Option<NodePositionIdx>,
72    child_index: NodePositionIdx,
73    children: Vec<NodePositionIdx>,
74    ancestor: NodePositionIdx,
75    thread: Option<NodePositionIdx>,
76    x: i32,
77    y: f32,
78    r#mod: f32,
79    change: f32,
80    shift: f32,
81}
82
83impl TreeNodePosition {
84    /// Adds displayed children traversing through nodes without a display
85    fn add_child(
86        nodes: &mut Vec<TreeNodePosition>,
87        tree: &mut AdvancementRegistry,
88        parent_idx: NodePositionIdx,
89        adv_node_idx: usize,
90        mut previous_idx: Option<NodePositionIdx>,
91    ) -> Option<NodePositionIdx> {
92        let adv_node = tree.adv_nodes.get(adv_node_idx)?;
93        if adv_node.has_display() {
94            let child_idx = nodes.len();
95            let node = &mut nodes[parent_idx];
96            let next_child_index = node.children.len() + 1;
97            node.children.push(child_idx);
98
99            nodes.push(TreeNodePosition {
100                node: adv_node_idx,
101                parent: Some(parent_idx),
102                previous_sibling: previous_idx,
103                child_index: next_child_index,
104                children: Vec::new(),
105                ancestor: child_idx,
106                thread: None,
107                x: 0,
108                y: -1.0,
109                r#mod: 0.0,
110                change: 0.0,
111                shift: 0.0,
112            });
113
114            let mut child_prev = None;
115            for child in adv_node.children() {
116                child_prev = Self::add_child(nodes, tree, child_idx, child, child_prev);
117            }
118
119            Some(child_idx)
120        } else {
121            for grandchild in adv_node.children() {
122                previous_idx = Self::add_child(nodes, tree, parent_idx, grandchild, previous_idx);
123            }
124            previous_idx
125        }
126    }
127
128    /// Traverse every node and compute each y position relative to its siblings and children.
129    fn first_walk(nodes: &mut [TreeNodePosition], idx: NodePositionIdx) {
130        let num_children = nodes[idx].children.len();
131        if num_children == 0 {
132            if let Some(prev_sib) = nodes[idx].previous_sibling {
133                nodes[idx].y = nodes[prev_sib].y + 1.0;
134            } else {
135                nodes[idx].y = 0.0;
136            }
137        } else {
138            let mut default_ancestor: Option<NodePositionIdx> = None;
139            for i in 0..num_children {
140                let child_idx = nodes[idx].children[i];
141                Self::first_walk(nodes, child_idx);
142                let arg_ancestor = default_ancestor.unwrap_or(child_idx);
143                default_ancestor = Some(Self::apportion(nodes, child_idx, arg_ancestor));
144            }
145
146            Self::execute_shifts(nodes, idx);
147
148            let node = &mut nodes[idx];
149            let first_child_idx = node.children[0];
150            let last_child_idx = node.children[num_children - 1];
151            let midpoint = f32::midpoint(nodes[first_child_idx].y, nodes[last_child_idx].y);
152
153            if let Some(prev_sib) = nodes[idx].previous_sibling {
154                nodes[idx].y = nodes[prev_sib].y + 1.0;
155                nodes[idx].r#mod = nodes[idx].y - midpoint;
156            } else {
157                nodes[idx].y = midpoint;
158            }
159        }
160    }
161
162    fn second_walk(
163        nodes: &mut [TreeNodePosition],
164        idx: NodePositionIdx,
165        mod_sum: f32,
166        depth: i32,
167        mut min: f32,
168    ) -> f32 {
169        let node = &mut nodes[idx];
170        node.y += mod_sum;
171        node.x = depth;
172
173        if node.y < min {
174            min = node.y;
175        }
176
177        let num_children = node.children.len();
178        let current_mod = node.r#mod;
179
180        for i in 0..num_children {
181            let child_idx = nodes[idx].children[i];
182            min = Self::second_walk(nodes, child_idx, mod_sum + current_mod, depth + 1, min);
183        }
184
185        min
186    }
187
188    /// just make the minimal y coord to be 0
189    fn normalize_y(nodes: &mut [TreeNodePosition], offset: f32) {
190        for node in nodes.iter_mut() {
191            node.y += offset;
192        }
193    }
194
195    fn execute_shifts(nodes: &mut [TreeNodePosition], idx: NodePositionIdx) {
196        let mut shift = 0.0;
197        let mut change = 0.0;
198
199        for &child_idx in nodes[idx].children.iter().rev() {
200            nodes[child_idx].y += shift;
201            nodes[child_idx].r#mod += shift;
202            change += nodes[child_idx].change;
203            shift += nodes[child_idx].shift + change;
204        }
205    }
206
207    #[inline]
208    fn previous_or_thread(
209        nodes: &[TreeNodePosition],
210        idx: NodePositionIdx,
211    ) -> Option<NodePositionIdx> {
212        nodes[idx]
213            .thread
214            .or_else(|| nodes[idx].children.first().copied())
215    }
216
217    #[inline]
218    fn next_or_thread(nodes: &[TreeNodePosition], idx: NodePositionIdx) -> Option<NodePositionIdx> {
219        nodes[idx]
220            .thread
221            .or_else(|| nodes[idx].children.last().copied())
222    }
223
224    fn apportion(
225        nodes: &mut [TreeNodePosition],
226        idx: NodePositionIdx,
227        mut default_ancestor: NodePositionIdx,
228    ) -> NodePositionIdx {
229        let Some(prev_sib) = nodes[idx].previous_sibling else {
230            return default_ancestor;
231        };
232        let parent_idx = nodes[idx].parent.expect("Tree invariant broken: no parent");
233        let mut inner_right = idx;
234        let mut outer_right = idx;
235        let mut inner_left = prev_sib;
236        let mut outer_left = nodes[parent_idx].children[0];
237
238        let mod_field = nodes[idx].r#mod;
239        let mut shift_inner_right = mod_field;
240        let mut shift_outer_right = mod_field;
241        let mut shift_inner_left = nodes[inner_left].r#mod;
242        let mut shift_outer_left = nodes[outer_left].r#mod;
243        while let Some(next_inner_left) = Self::next_or_thread(nodes, inner_left)
244            && let Some(next_inner_right) = Self::previous_or_thread(nodes, inner_right)
245        {
246            inner_left = next_inner_left;
247            inner_right = next_inner_right;
248            outer_left =
249                Self::previous_or_thread(nodes, outer_left).expect("Tree invariant broken");
250            outer_right = Self::next_or_thread(nodes, outer_right).expect("Tree invariant broken");
251
252            nodes[outer_right].ancestor = idx;
253
254            let shift = (nodes[inner_left].y + shift_inner_left)
255                - (nodes[inner_right].y + shift_inner_right)
256                + 1.0;
257            if shift > 0.0 {
258                let ancestor_idx = Self::get_ancestor(nodes, inner_left, idx, default_ancestor);
259                Self::move_subtree(nodes, ancestor_idx, idx, shift);
260                shift_inner_right += shift;
261                shift_outer_right += shift;
262            }
263
264            shift_inner_left += nodes[inner_left].r#mod;
265            shift_inner_right += nodes[inner_right].r#mod;
266            shift_outer_left += nodes[outer_left].r#mod;
267            shift_outer_right += nodes[outer_right].r#mod;
268        }
269
270        if let Some(next_inner_left) = Self::next_or_thread(nodes, inner_left)
271            && Self::next_or_thread(nodes, outer_right).is_none()
272        {
273            nodes[outer_right].thread = Some(next_inner_left);
274            nodes[outer_right].r#mod += shift_inner_left - shift_outer_right;
275        } else {
276            if let Some(next_inner_right) = Self::previous_or_thread(nodes, inner_right)
277                && Self::previous_or_thread(nodes, outer_left).is_none()
278            {
279                nodes[outer_left].thread = Some(next_inner_right);
280                nodes[outer_left].r#mod += shift_inner_right - shift_outer_left;
281            }
282            default_ancestor = idx;
283        }
284        default_ancestor
285    }
286
287    fn move_subtree(
288        nodes: &mut [TreeNodePosition],
289        left: NodePositionIdx,
290        right: NodePositionIdx,
291        shift: f32,
292    ) {
293        let subtrees = (nodes[right].child_index - nodes[left].child_index) as f32;
294        if subtrees != 0.0 {
295            nodes[right].change -= shift / subtrees;
296            nodes[left].change += shift / subtrees;
297        }
298        nodes[right].shift += shift;
299        nodes[right].y += shift;
300        nodes[right].r#mod += shift;
301    }
302
303    fn get_ancestor(
304        nodes: &[TreeNodePosition],
305        idx: NodePositionIdx,
306        other: NodePositionIdx,
307        default_ancestor: NodePositionIdx,
308    ) -> NodePositionIdx {
309        let ancestor = nodes[idx].ancestor;
310        let parent_idx = nodes[other].parent.expect("Tree invariant broken");
311
312        if nodes[parent_idx].children.contains(&ancestor) {
313            ancestor
314        } else {
315            default_ancestor
316        }
317    }
318
319    fn set_position(
320        tree: &mut AdvancementRegistry,
321        nodes: &[TreeNodePosition],
322        idx: NodePositionIdx,
323    ) {
324        tree.adv_nodes[nodes[idx].node].set_location(nodes[idx].x as f32, nodes[idx].y);
325        for &child_idx in &nodes[idx].children {
326            Self::set_position(tree, nodes, child_idx);
327        }
328    }
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334    use crate::advancement::registry::{AdvancementNode, AdvancementRef};
335    use crate::vanilla_advancements::*;
336    use steel_utils::Identifier;
337
338    fn get_location(registry: &AdvancementRegistry, key: &Identifier) -> (f32, f32) {
339        let loc: Option<(f32, f32)> = registry.by_key(key).map(AdvancementNode::get_location);
340        loc.unwrap_or_else(|| panic!("unbale to get the location of {key}"))
341    }
342
343    #[test]
344    fn single_root_no_children() {
345        let mut registry = AdvancementRegistry::default();
346        registry.register_without_load(&[&STORY_ROOT]);
347        let res = run(&mut registry, 0);
348        assert!(res.is_ok());
349        let location = get_location(&registry, &STORY_ROOT.key);
350        assert_eq!(location, (0f32, 0f32));
351    }
352
353    #[test]
354    fn root_with_linear_children() {
355        let mut registry = AdvancementRegistry::default();
356        let list: &[AdvancementRef] = &[
357            &STORY_ROOT,
358            &STORY_MINE_STONE,
359            &STORY_UPGRADE_TOOLS,
360            &STORY_SMELT_IRON,
361        ];
362        registry.register_without_load(list);
363        let idx = registry.by_key[&list[0].key];
364        let res = run(&mut registry, idx);
365        assert!(res.is_ok());
366        for (i, adv) in list.iter().enumerate() {
367            let loc = get_location(&registry, &adv.key);
368            assert_eq!(
369                loc,
370                (i as f32, 0f32),
371                "node {} isn't at the right location",
372                adv.key
373            );
374        }
375    }
376
377    #[test]
378    fn root_with_branching_children() {
379        let mut registry = AdvancementRegistry::default();
380        let list: &[AdvancementRef] = &[
381            &END_ROOT,
382            &END_KILL_DRAGON,
383            &END_ENTER_END_GATEWAY,
384            &END_RESPAWN_DRAGON,
385            &END_DRAGON_EGG,
386            &END_DRAGON_BREATH,
387        ];
388        registry.register_without_load(list);
389        let idx = registry.by_key[&list[0].key];
390        let res = run(&mut registry, idx);
391        assert!(res.is_ok());
392        let mut locs = [(0f32, 0f32); 6];
393        for (i, adv) in list.iter().enumerate() {
394            locs[i] = get_location(&registry, &adv.key);
395        }
396        let expected = [
397            (0f32, 1.5f32),
398            (1f32, 1.5f32),
399            (2f32, 0f32),
400            (2f32, 1f32),
401            (2f32, 2f32),
402            (2f32, 3f32),
403        ];
404        assert_eq!(locs, expected);
405    }
406}
steel_registry/advancement/positioner.rs

steel_registry/advancement/
positioner.rs
