refactor: Per path wavelengths and solves

crates/cherry-rs/src/core/sequential_model/mod.rs

@@ -69,6 +69,9 @@ struct OpticalPath {
     stop_surface: Option<usize>,
     /// Step-indexed cursor state, parallel to `surface_indices`.
     steps: Vec<CursorPlacement>,
+    /// Wavelengths this path's `submodels` were built from, in the same
+    /// order (`submodels[i]` corresponds to `wavelengths[i]`).
+    wavelengths: Vec<Float>,
 }
 
 /// A gap between two surfaces in a sequential system.
@@ -90,7 +93,6 @@ pub struct Gap {
 pub struct SequentialModel {
     store: SurfaceStore,
     paths: Vec<OpticalPath>,
-    wavelengths: Vec<Float>,
 }
 
 /// A submodel of a sequential optical system.
@@ -329,6 +331,43 @@ pub fn reversed_surface_id(num_surfaces: usize, surf_id: usize) -> usize {
 /// Only meaningful for `Reversed` paths; callers must gate on orientation
 /// before using the result (see
 /// [`SequentialModel::from_path_specs_with_builder`]).
+/// Locates the `PathSpec`/step that introduced store index `store_index` via
+/// `PathSurfaceRef::New`, using the identical counting rule
+/// `from_path_specs_with_builder` uses to assign store indices: `New` and
+/// `ObjectLinkedTo` steps each consume one index, in path/step order;
+/// `Shared` steps do not. Returns `None` if `store_index` is out of range.
+///
+/// An `ObjectLinkedTo` step also consumes a store index by this rule but
+/// never owns a mutable `SurfaceSpec` (it always synthesizes a plain
+/// `SurfaceSpec::Object` at build time) — callers needing a mutable target
+/// must match the result against `PathSurfaceRef::New` and treat any other
+/// variant as ineligible; no special-casing is done here.
+///
+/// `paths`' topology (which steps are `New`/`Shared`/`ObjectLinkedTo`, and
+/// their order) never changes across a builder's solve-rebuild iterations —
+/// only the `SurfaceSpec` values inside `New` steps do — so a caller
+/// applying multiple solves may compute this once and reuse it.
+pub(crate) fn locate_surface_owner(
+    paths: &[PathSpec],
+    store_index: usize,
+) -> Option<(usize, usize)> {
+    let mut counter = 0;
+    for (path_idx, ps) in paths.iter().enumerate() {
+        for (step_idx, sref) in ps.surface_refs.iter().enumerate() {
+            match sref {
+                PathSurfaceRef::New(_) | PathSurfaceRef::ObjectLinkedTo { .. } => {
+                    if counter == store_index {
+                        return Some((path_idx, step_idx));
+                    }
+                    counter += 1;
+                }
+                PathSurfaceRef::Shared(_) => {}
+            }
+        }
+    }
+    None
+}
+
 fn count_leading_shared(refs: &[PathSurfaceRef]) -> usize {
     refs.iter()
         .skip(1)
@@ -403,11 +442,11 @@ impl SequentialModel {
                 submodels,
                 stop_surface,
                 steps: cursor_placements,
+                wavelengths: wavelengths.to_vec(),
             };
             Ok(Self {
                 store,
                 paths: vec![path],
-                wavelengths: wavelengths.to_vec(),
             })
         }
     }
@@ -447,11 +486,11 @@ impl SequentialModel {
             submodels,
             stop_surface,
             steps: cursor_placements,
+            wavelengths: wavelengths.to_vec(),
         };
         Ok(Self {
             store,
             paths: vec![path],
-            wavelengths: wavelengths.to_vec(),
         })
     }
 
@@ -518,11 +557,11 @@ impl SequentialModel {
             submodels,
             stop_surface,
             steps: cursor_placements,
+            wavelengths: wavelengths.to_vec(),
         };
         Ok(Self {
             store,
             paths: vec![path],
-            wavelengths: wavelengths.to_vec(),
         })
     }
 
@@ -532,14 +571,9 @@ impl SequentialModel {
     /// from a [`SurfaceSpec`]. It is injected so that the serde and non-serde
     /// versions can supply the appropriate `surface_from_spec` variant.
     fn from_path_specs_with_builder(
-        paths: Vec<PathSpec>,
-        wavelengths: &[Float],
+        paths: &[PathSpec],
         mut build_surface: impl FnMut(&SurfaceSpec) -> Result<Box<dyn Surface>>,
     ) -> Result<Self> {
-        if wavelengths.is_empty() {
-            return Err(anyhow!("At least one wavelength must be specified."));
-        }
-
         let mut store_surfaces: Vec<Box<dyn Surface>> = Vec::new();
         let mut store_placements: Vec<SurfacePlacement> = Vec::new();
         let mut optical_paths: Vec<OpticalPath> = Vec::new();
@@ -655,7 +689,7 @@ impl SequentialModel {
                 };
 
             // Build the dense arm vec by consuming beam_splitter_arms in step order.
-            let mut bs_arms_iter = ps.beam_splitter_arms.into_iter();
+            let mut bs_arms_iter = ps.beam_splitter_arms.iter().copied();
             let mut dense_bs_arms: Vec<Option<BeamSplitterPathKind>> = Vec::with_capacity(n_refs);
 
             let mut cursor = match linked {
@@ -845,8 +879,13 @@ impl SequentialModel {
                 .cloned()
                 .collect();
 
+            if ps.wavelengths.is_empty() {
+                return Err(anyhow!(
+                    "each PathSpec must have at least one wavelength; this path has none"
+                ));
+            }
             let mut submodels: Vec<SequentialSubModelBase> = Vec::new();
-            for &wavelength in wavelengths.iter() {
+            for &wavelength in ps.wavelengths.iter() {
                 let gaps = Self::gap_specs_to_gaps(&all_gap_specs, wavelength)?;
                 submodels.push(SequentialSubModelBase::new(gaps));
             }
@@ -864,6 +903,7 @@ impl SequentialModel {
                 submodels,
                 stop_surface: ps.stop_surface,
                 steps: path_steps,
+                wavelengths: ps.wavelengths.clone(),
             });
         }
 
@@ -874,26 +914,22 @@ impl SequentialModel {
         Ok(Self {
             store,
             paths: optical_paths,
-            wavelengths: wavelengths.to_vec(),
         })
     }
 
     /// Builds a multipath model from `PathSpec`s (serde + registry variant).
     #[cfg(feature = "serde")]
     pub(crate) fn from_path_specs(
-        paths: Vec<PathSpec>,
-        wavelengths: &[Float],
+        paths: &[PathSpec],
         registry: Option<&SurfaceRegistry>,
     ) -> Result<Self> {
-        Self::from_path_specs_with_builder(paths, wavelengths, |spec| {
-            surface_from_spec(spec, registry)
-        })
+        Self::from_path_specs_with_builder(paths, |spec| surface_from_spec(spec, registry))
     }
 
     /// Builds a multipath model from `PathSpec`s (non-serde variant).
     #[cfg(not(feature = "serde"))]
-    pub(crate) fn from_path_specs(paths: Vec<PathSpec>, wavelengths: &[Float]) -> Result<Self> {
-        Self::from_path_specs_with_builder(paths, wavelengths, surface_from_spec)
+    pub(crate) fn from_path_specs(paths: &[PathSpec]) -> Result<Self> {
+        Self::from_path_specs_with_builder(paths, surface_from_spec)
     }
 
     /// Number of optical paths in the model.
@@ -1019,8 +1055,15 @@ impl SequentialModel {
     }
 
     /// Returns the wavelengths at which the system is modeled.
+    ///
+    /// Single-path shorthand; delegates to `paths[0]`.
     pub fn wavelengths(&self) -> &[Float] {
-        &self.wavelengths
+        self.wavelengths_for_path(0)
+    }
+
+    /// Returns the wavelengths for path `path_id`, in submodel order.
+    pub fn wavelengths_for_path(&self, path_id: usize) -> &[Float] {
+        &self.paths[path_id].wavelengths
     }
 
     fn gap_specs_to_gaps(gap_specs: &[GapSpec], wavelength: Float) -> Result<Vec<Gap>> {