Assemblies: use lookup maps when building the world state (#10497)

* claim: Assemblies: use lookup maps when building the world state

* Assemblies: use lookup maps when building the world state

world_state_reusing, made again after every geometry command in a design
with components, filtered all bodies for every placement, searched the
timeline for every sketch and construction of a placed part, and searched
all bodies twice per tap and placement. It now builds the feature to
component map, the bodies of each component (in body order) and each
tap's component and body offset once per call. The result and its order
are unchanged.

On a 100-part assembly placed 5 times (800 sketches, 200 taps) one world
state drops from 124 ms to 9.5 ms (release, median).

A test-only counter of searches through all features or all bodies
checks that the work no longer grows with the sketches and taps.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Markus Kaltenbrunner <kaltenbrunnermarkus@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

crates/doc/src/document.rs

@@ -3500,6 +3500,8 @@ impl Document {
     }
 
     pub fn feature(&self, id: u64) -> Option<&Feature> {
+        #[cfg(test)]
+        crate::eval::cost_tests::count_feature_search();
         self.features.iter().find(|f| f.id == id)
     }
     pub fn feature_mut(&mut self, id: u64) -> Option<&mut Feature> {

crates/doc/src/eval.rs

@@ -8604,7 +8604,10 @@ pub fn world_state_reusing(doc: &Document, st: &ModelState, placed_before: &mut
         return st.clone();
     }
     let mut placed_now = PlacedBodies::default();
-    let comp_of_feature = |id: u64| doc.feature(id).map(|f| f.component).unwrap_or(0);
+    // Looked up once: every sketch, construction and tap asked the timeline, and every placement
+    // went through all bodies, each time searching (a large assembly paid this every command).
+    let feature_components: std::collections::HashMap<u64, u64> = doc.features.iter().rev().map(|f| (f.id, f.component)).collect();
+    let comp_of_feature = |id: u64| feature_components.get(&id).copied().unwrap_or(0);
     let first_path = |comp: u64| -> Vec<u64> {
         let mut path = Vec::new();
         let mut c = comp;
@@ -8627,7 +8630,24 @@ pub fn world_state_reusing(doc: &Document, st: &ModelState, placed_before: &mut
     let mut jobs = Vec::new();
     let owners = doc.model_body_components(st);
     let names = crate::WorldNames::new(doc, st);
-    let bodies_of = |comp: u64| st.bodies.iter().zip(&owners).filter(move |(_, owner)| **owner == comp).map(|(b, _)| b);
+    let mut by_component: std::collections::HashMap<u64, Vec<&ModelBody>> = std::collections::HashMap::new();
+    for (b, owner) in st.bodies.iter().zip(&owners) {
+        by_component.entry(*owner).or_default().push(b);
+    }
+    let bodies_of = |comp: u64| by_component.get(&comp).into_iter().flatten().copied();
+    // Each tap's component and the offset of the body it is on (the first body of that name).
+    let mut first_body: std::collections::HashMap<&str, (&ModelBody, u64)> = std::collections::HashMap::new();
+    for (b, owner) in st.bodies.iter().zip(&owners).rev() {
+        first_body.insert(b.name.as_str(), (b, *owner));
+    }
+    let tap_frames: Vec<(u64, Option<&crate::Mat>)> = st
+        .threads
+        .iter()
+        .map(|t| match first_body.get(t.body.as_str()) {
+            Some((body, owner)) => (*owner, doc.model_body_offset(body)),
+            None => (comp_of_feature(t.feature), None),
+        })
+        .collect();
     for (comp, path, m) in &placements {
         for b in bodies_of(*comp) {
             let bm = body_matrix(b, m);
@@ -8707,14 +8727,10 @@ pub fn world_state_reusing(doc: &Document, st: &ModelState, placed_before: &mut
             }
         }
         // Unlike editable sketch guides, cosmetic surfaces belong to every occurrence.
-        for t in st
-            .threads
-            .iter()
-            .filter(|t| st.body(&t.body).map(|body| doc.model_body_component(body)).unwrap_or_else(|| comp_of_feature(t.feature)) == comp)
-        {
+        for (t, (_, offset)) in st.threads.iter().zip(&tap_frames).filter(|(_, (owner, _))| *owner == comp) {
             let mut t2 = t.clone();
             t2.body = names.name(&t.body, comp, &path);
-            let frame = st.body(&t.body).and_then(|body| doc.model_body_offset(body)).map(|offset| crate::mat_mul(&m, offset)).unwrap_or(m);
+            let frame = offset.map(|offset| crate::mat_mul(&m, offset)).unwrap_or(m);
             t2.axis_point = crate::apply_point(&frame, t.axis_point);
             t2.axis = crate::apply_vector(&frame, *t.axis).unit().unwrap_or(t.axis);
             out.threads.push(t2);
@@ -8899,7 +8915,7 @@ mod chain_tests {
 
 #[cfg(test)]
 #[path = "eval_cost_tests.rs"]
-mod cost_tests;
+pub(crate) mod cost_tests;
 
 #[cfg(test)]
 #[path = "eval_whole_face_tests.rs"]

crates/doc/src/eval_cost_tests.rs

@@ -16,7 +16,7 @@ use crate::eval::{PlacementReads, body_world, hash_placement, placement_digest};
 use crate::*;
 
 /// Work counted on this thread: bodies moved, STEP texts digested, joint faces and bodies
-/// measured, bodies searched for by name.
+/// measured, bodies searched for by name, features searched for by id.
 #[derive(Clone, Copy, Debug, Default, PartialEq)]
 pub(super) struct Counts {
     pub moved: usize,
@@ -24,10 +24,12 @@ pub(super) struct Counts {
     pub joint_faces: usize,
     pub measures: usize,
     pub body_searches: usize,
+    pub feature_searches: usize,
 }
 
 thread_local! {
-    static COUNTS: Cell<Counts> = const { Cell::new(Counts { moved: 0, digested: 0, joint_faces: 0, measures: 0, body_searches: 0 }) };
+    static COUNTS: Cell<Counts> =
+        const { Cell::new(Counts { moved: 0, digested: 0, joint_faces: 0, measures: 0, body_searches: 0, feature_searches: 0 }) };
 }
 
 pub(super) fn count(f: impl FnOnce(&mut Counts)) {
@@ -54,6 +56,17 @@ fn written(f: impl FnOnce()) -> usize {
     WRITTEN.with(Cell::get)
 }
 
+/// A search on this thread through all of a design's features for one.
+pub(crate) fn count_feature_search() {
+    count(|c| c.feature_searches += 1);
+}
+
+/// The searches `f` makes on this thread through all features or all bodies.
+fn searched(f: impl FnOnce()) -> usize {
+    let n = counted(f);
+    n.feature_searches + n.body_searches
+}
+
 /// The work `f` does on this thread.
 fn counted(f: impl FnOnce()) -> Counts {
     COUNTS.with(|c| c.set(Counts::default()));
@@ -967,3 +980,87 @@ fn looking_bodies_up_once_keeps_every_placement_digest() {
         assert_eq!(digest, digest_searching_each_body(&doc, f, &st), "cut {}", f.id);
     }
 }
+
+/// `comps` parts, each with `bodies` boxes, `sketches` sketches and `taps` taps, placed `copies`
+/// times: a design and its state, made without evaluating it.
+fn placed_parts(comps: usize, bodies: usize, sketches: usize, taps: usize, copies: usize) -> (Document, ModelState) {
+    let mut doc = Document::new("placed parts");
+    let mut one = Document::new("one sketch");
+    one.add_feature(line_sketch(PlaneRef::Origin { name: "XY".into() }, Vec2::new(3.0, 4.0)), None).unwrap();
+    let mut m = Model::new();
+    m.evaluate(&one);
+    let solved = m.state().sketches.first().unwrap().clone();
+    let mut st = ModelState::default();
+    let solid = kernel::box_solid(Vec3::ZERO, Vec3::new(2.0, 3.0, 4.0)).unwrap();
+    for c in 0..comps {
+        let comp = doc.add_component(Some(&format!("C{c}")), 0).unwrap();
+        let o = doc.occurrences.iter_mut().find(|o| o.component == comp).unwrap();
+        o.transform = rigid(Vec3::new(10.0 * c as f64, 0.0, 0.0), Vec3::ZERO, Vec3::Z, 0.01 * c as f64);
+        for k in 1..copies {
+            doc.add_occurrence(comp, 0, rigid(Vec3::new(10.0 * c as f64, 20.0 * k as f64, 0.0), Vec3::ZERO, Vec3::Z, 0.02 * k as f64)).unwrap();
+        }
+        let mut names = Vec::new();
+        for b in 0..bodies {
+            let kind =
+                serde_json::from_value(serde_json::json!({"type": "box", "corner": [0, 0, 0], "length": "2", "width": "3", "height": "4"})).unwrap();
+            let name = format!("B{c}_{b}");
+            let id = doc.add_feature(kind, Some(&name)).unwrap();
+            doc.feature_mut(id).unwrap().component = comp;
+            st.bodies.push(ModelBody::new(name.clone(), solid.clone(), id));
+            names.push((name, id));
+        }
+        for s in 0..sketches {
+            let id = doc.add_feature(line_sketch(PlaneRef::Origin { name: "XY".into() }, Vec2::new(1.0 + s as f64, 2.0)), None).unwrap();
+            doc.feature_mut(id).unwrap().component = comp;
+            let mut s2 = SolvedSketch::clone(&solved);
+            s2.feature = id;
+            st.sketches.push(Arc::new(s2));
+        }
+        for t in 0..taps {
+            let (body, feature) = names[t % names.len()].clone();
+            st.threads.push(ThreadInfo {
+                feature,
+                body,
+                designation: "M6".into(),
+                major_diameter: 6.0,
+                radius: 3.0,
+                pitch: 1.0,
+                internal: true,
+                handedness: Default::default(),
+                axis_point: Vec3::new(1.0, 1.0, 0.0),
+                axis: solidcraft_geom::UnitVec3::Z,
+                start: 0.0,
+                end: 4.0,
+            });
+        }
+    }
+    (doc, st)
+}
+
+#[test]
+fn placing_parts_looks_each_sketch_and_tap_up_once() {
+    // A large assembly searched all features for every sketch and construction, and all bodies
+    // twice for every tap and placement, each time its world state was made (every command).
+    let searches = |sketches: usize, taps: usize| {
+        let (doc, state) = placed_parts(4, 2, sketches, taps, 3);
+        let mut world = ModelState::default();
+        let n = searched(|| world = world_state(&doc, &state));
+        assert_eq!(world.bodies.len(), 4 * 2 * 3);
+        assert_eq!(world.sketches.len(), 4 * sketches);
+        assert_eq!(world.threads.len(), 4 * taps * 3);
+        n
+    };
+    assert_eq!(searches(8, 6), searches(1, 1));
+}
+
+#[test]
+fn a_placed_tap_follows_its_body_offset() {
+    let (mut doc, state) = placed_parts(1, 1, 0, 1, 1);
+    let comp = doc.components[0].id;
+    doc.body_offsets.insert("B0_0".into(), rigid(Vec3::new(0.0, 0.0, 7.0), Vec3::ZERO, Vec3::Z, 0.0));
+    let world = world_state(&doc, &state);
+    let tap = world.threads.first().unwrap();
+    let expected =
+        apply_point(&mat_mul(&doc.component_transform(comp), &rigid(Vec3::new(0.0, 0.0, 7.0), Vec3::ZERO, Vec3::Z, 0.0)), Vec3::new(1.0, 1.0, 0.0));
+    assert!(tap.axis_point.dist(expected) < 1e-9, "{:?} {expected:?}", tap.axis_point);
+}