oximo_core/
indexed.rs

1use std::marker::PhantomData;
2use std::ops::Index;
3
4use oximo_expr::{Expr, ModelId};
5use rustc_hash::FxHashMap;
6
7use crate::constraint::{ConstraintHandle, RangeConstraintHandles};
8use crate::indicator::{IndicatorConstraintHandle, RangeIndicatorConstraintHandles};
9use crate::set::{Axis, FromIndexKey, IndexKey};
10
11/// Owned, ordered IDs with a domain-specific lookup index.
12#[derive(Clone, Debug)]
13struct ConstraintFamilyStorage<T> {
14    entries: Vec<(IndexKey, T)>,
15    lookup: ConstraintFamilyLookup,
16}
17
18#[derive(Clone, Debug)]
19enum ConstraintFamilyLookup {
20    Dense(Box<[Axis]>),
21    Sparse(FxHashMap<IndexKey, usize>),
22}
23
24impl<T: Copy> ConstraintFamilyStorage<T> {
25    fn new(keys: Vec<IndexKey>, axes: Option<&[Axis]>, values: Vec<T>) -> Self {
26        assert_eq!(keys.len(), values.len(), "constraint family key/value length mismatch");
27        let lookup = axes.map_or_else(
28            || {
29                ConstraintFamilyLookup::Sparse(
30                    keys.iter().cloned().enumerate().map(|(i, key)| (key, i)).collect(),
31                )
32            },
33            |axes| ConstraintFamilyLookup::Dense(axes.into()),
34        );
35        Self { entries: keys.into_iter().zip(values).collect(), lookup }
36    }
37
38    fn get(&self, key: &IndexKey) -> Option<T> {
39        let position = match &self.lookup {
40            ConstraintFamilyLookup::Dense(axes) => grid_offset(axes, key)?,
41            ConstraintFamilyLookup::Sparse(positions) => *positions.get(key)?,
42        };
43        self.entries.get(position).map(|(_, value)| *value)
44    }
45}
46
47macro_rules! constraint_family {
48    ($(#[$doc:meta])* $name:ident, $value:ty) => {
49        $(#[$doc])*
50        pub struct $name<K = IndexKey> {
51            storage: ConstraintFamilyStorage<$value>,
52            _marker: PhantomData<fn() -> K>,
53        }
54
55        impl<K> Clone for $name<K> {
56            fn clone(&self) -> Self {
57                Self { storage: self.storage.clone(), _marker: PhantomData }
58            }
59        }
60
61        impl<K> std::fmt::Debug for $name<K> {
62            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63                f.debug_struct(stringify!($name)).field("entries", &self.storage.entries).finish()
64            }
65        }
66
67        impl<K> $name<K> {
68            pub(crate) fn new(keys: Vec<IndexKey>, axes: Option<&[Axis]>, values: Vec<$value>) -> Self {
69                Self { storage: ConstraintFamilyStorage::new(keys, axes, values), _marker: PhantomData }
70            }
71
72            /// Number of domain entries (not the number of lowered rows).
73            pub fn len(&self) -> usize {
74                self.storage.entries.len()
75            }
76
77            pub fn is_empty(&self) -> bool {
78                self.storage.entries.is_empty()
79            }
80
81            /// Look up an entry, returning `None` for missing or filtered-out keys.
82            pub fn get<Q: Into<IndexKey>>(&self, key: Q) -> Option<$value> {
83                self.storage.get(&key.into())
84            }
85        }
86
87        impl<K: FromIndexKey> $name<K> {
88            /// Iterate typed keys and copied IDs in the original domain order.
89            pub fn iter(&self) -> impl Iterator<Item = (K, $value)> + '_ {
90                self.storage.entries.iter().map(|(key, value)| (K::from_index_key(key), *value))
91            }
92        }
93    };
94}
95
96constraint_family!(
97    /// Owned handle returned by an indexed single-relation `constraint!` declaration.
98    ///
99    /// IDs refer to rows of the originating model.
100    /// `get` supports integer, string, and tuple keys, including sparse domains.
101    IndexedConstraint, ConstraintHandle
102);
103
104/// Owned handles returned by an indexed single-relation indicator declaration.
105#[derive(Clone, Debug)]
106pub struct IndexedIndicatorConstraint<'a, K = IndexKey> {
107    storage: ConstraintFamilyStorage<IndicatorConstraintHandle<'a>>,
108    _marker: PhantomData<fn() -> K>,
109}
110
111impl<'a, K> IndexedIndicatorConstraint<'a, K> {
112    pub(crate) fn new(
113        keys: Vec<IndexKey>,
114        axes: Option<&[Axis]>,
115        values: Vec<IndicatorConstraintHandle<'a>>,
116    ) -> Self {
117        Self { storage: ConstraintFamilyStorage::new(keys, axes, values), _marker: PhantomData }
118    }
119    pub fn len(&self) -> usize {
120        self.storage.entries.len()
121    }
122    pub fn is_empty(&self) -> bool {
123        self.storage.entries.is_empty()
124    }
125    pub fn get<Q: Into<IndexKey>>(&self, key: Q) -> Option<IndicatorConstraintHandle<'a>> {
126        self.storage.get(&key.into())
127    }
128}
129
130impl<'a, K: FromIndexKey> IndexedIndicatorConstraint<'a, K> {
131    pub fn iter(&self) -> impl Iterator<Item = (K, IndicatorConstraintHandle<'a>)> + '_ {
132        self.storage.entries.iter().map(|(key, value)| (K::from_index_key(key), *value))
133    }
134}
135
136/// Owned handles returned by an indexed two-sided indicator declaration.
137#[derive(Clone, Debug)]
138pub struct IndexedRangeIndicatorConstraint<'a, K = IndexKey> {
139    storage: ConstraintFamilyStorage<RangeIndicatorConstraintHandles<'a>>,
140    _marker: PhantomData<fn() -> K>,
141}
142
143impl<'a, K> IndexedRangeIndicatorConstraint<'a, K> {
144    pub(crate) fn new(
145        keys: Vec<IndexKey>,
146        axes: Option<&[Axis]>,
147        values: Vec<RangeIndicatorConstraintHandles<'a>>,
148    ) -> Self {
149        Self { storage: ConstraintFamilyStorage::new(keys, axes, values), _marker: PhantomData }
150    }
151    pub fn len(&self) -> usize {
152        self.storage.entries.len()
153    }
154    pub fn is_empty(&self) -> bool {
155        self.storage.entries.is_empty()
156    }
157    pub fn get<Q: Into<IndexKey>>(&self, key: Q) -> Option<RangeIndicatorConstraintHandles<'a>> {
158        self.storage.get(&key.into())
159    }
160}
161
162impl<'a, K: FromIndexKey> IndexedRangeIndicatorConstraint<'a, K> {
163    pub fn iter(&self) -> impl Iterator<Item = (K, RangeIndicatorConstraintHandles<'a>)> + '_ {
164        self.storage.entries.iter().map(|(key, value)| (K::from_index_key(key), *value))
165    }
166}
167
168constraint_family!(
169    /// Owned handle returned by an indexed two-sided range `constraint!` declaration.
170    ///
171    /// Each key maps to one interval ID or separate lower/upper IDs. Entries can
172    /// have different lowering forms within the same family.
173    IndexedRangeConstraint, RangeConstraintHandles
174);
175
176/// Backing storage for an [`IndexedFamily`].
177///
178/// `Dense` is used when the domain is a contiguous integer grid (a range, or a
179/// product of ranges). `Sparse` (string sets, sparse `from_ints`, or any
180/// `filter`ed family) keeps the original hash map.
181#[derive(Clone)]
182pub(crate) enum Storage<'a> {
183    Dense { data: Vec<Expr<'a>>, keys: Vec<IndexKey>, axes: Box<[Axis]> },
184    Sparse(FxHashMap<IndexKey, Expr<'a>>),
185}
186
187mod sealed {
188    pub trait Sealed {}
189}
190
191/// Marker selecting which kind of indexed family an [`IndexedFamily`] is.
192///
193/// Sealed implementation detail: implemented only for [`VarFamily`] and
194/// [`ParamFamily`]. It exists so [`IndexedVar`] and [`IndexedParam`] are distinct
195/// types (only a parameter family can be re-bound) while sharing one
196/// implementation.
197#[doc(hidden)]
198pub trait Family: sealed::Sealed {
199    /// Type name used in [`Debug`](std::fmt::Debug) output.
200    const NAME: &'static str;
201}
202
203/// Marker for an indexed family of decision variables ([`IndexedVar`]).
204#[doc(hidden)]
205#[derive(Debug)]
206pub struct VarFamily;
207
208/// Marker for an indexed family of parameters ([`IndexedParam`]).
209#[doc(hidden)]
210#[derive(Debug)]
211pub struct ParamFamily;
212
213impl sealed::Sealed for VarFamily {}
214impl sealed::Sealed for ParamFamily {}
215impl Family for VarFamily {
216    const NAME: &'static str = "IndexedVar";
217}
218impl Family for ParamFamily {
219    const NAME: &'static str = "IndexedParam";
220}
221
222/// Indexed family: maps an `IndexKey` to a single-element `Expr` (a variable or a
223/// parameter), tagged with the key type `K` its domain decodes to and the family
224/// kind `F`.
225///
226/// You normally name this through the [`IndexedVar`]/[`IndexedParam`] aliases,
227/// constructed by the indexed form of the `variable!`/`param!` macros.
228///
229/// When the domain is a contiguous integer range (or a Cartesian product of
230/// ranges) the family is stored densely (see the internal `Storage`).
231/// String, sparse, and `filter`ed families fall back to a hash map.
232pub struct IndexedFamily<'a, K = IndexKey, F = VarFamily> {
233    pub(crate) storage: Storage<'a>,
234    pub(crate) model_id: ModelId,
235    pub(crate) _marker: PhantomData<fn() -> (K, F)>,
236}
237
238/// Indexed family of decision variables; see [`IndexedFamily`].
239pub type IndexedVar<'a, K = IndexKey> = IndexedFamily<'a, K, VarFamily>;
240
241/// Indexed family of re-bindable parameters; see [`IndexedFamily`].
242///
243/// Re-bind a single entry with [`Model::set_param_idx`](crate::Model::set_param_idx).
244pub type IndexedParam<'a, K = IndexKey> = IndexedFamily<'a, K, ParamFamily>;
245
246impl<'a, K, F> Clone for IndexedFamily<'a, K, F> {
247    fn clone(&self) -> Self {
248        Self { storage: self.storage.clone(), model_id: self.model_id, _marker: PhantomData }
249    }
250}
251
252impl<'a, K, F: Family> std::fmt::Debug for IndexedFamily<'a, K, F> {
253    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
254        f.debug_struct(F::NAME).field("len", &self.len()).field("dense", &self.is_dense()).finish()
255    }
256}
257
258impl<'a, K, F> IndexedFamily<'a, K, F> {
259    /// Identity of the model that created this indexed family.
260    #[must_use]
261    pub const fn model_id(&self) -> ModelId {
262        self.model_id
263    }
264
265    pub fn len(&self) -> usize {
266        match &self.storage {
267            Storage::Dense { data, .. } => data.len(),
268            Storage::Sparse(m) => m.len(),
269        }
270    }
271
272    pub fn is_empty(&self) -> bool {
273        self.len() == 0
274    }
275
276    /// Whether this family is stored densely (domain was a range or product of
277    /// ranges).
278    pub fn is_dense(&self) -> bool {
279        matches!(self.storage, Storage::Dense { .. })
280    }
281
282    /// Per-axis lengths when stored densely, else `None`.
283    pub fn shape(&self) -> Option<Box<[usize]>> {
284        match &self.storage {
285            Storage::Dense { axes, .. } => Some(axes.iter().map(|a| a.len).collect()),
286            Storage::Sparse(_) => None,
287        }
288    }
289
290    pub fn iter(&self) -> impl Iterator<Item = (&IndexKey, &Expr<'a>)> + '_ {
291        let it: Box<dyn Iterator<Item = (&IndexKey, &Expr<'a>)>> = match &self.storage {
292            Storage::Dense { data, keys, .. } => Box::new(keys.iter().zip(data.iter())),
293            Storage::Sparse(m) => Box::new(m.iter()),
294        };
295        it
296    }
297
298    pub fn get<Q: Into<IndexKey>>(&self, key: Q) -> Option<Expr<'a>> {
299        match &self.storage {
300            Storage::Sparse(m) => m.get(&key.into()).copied(),
301            Storage::Dense { data, axes, .. } => {
302                grid_offset(axes, &key.into()).map(|off| data[off])
303            }
304        }
305    }
306
307    /// Zero-allocation typed index by integer coordinates.
308    /// On a dense family this maps straight to a flat offset with no
309    /// `IndexKey` built. On a sparse family it falls back to building a key.
310    ///
311    /// # Panics
312    /// Panics if the coordinates are out of range/not present.
313    pub fn at<const N: usize>(&self, coords: [usize; N]) -> Expr<'a> {
314        *self.get_ref(&coords).expect("indexed family: coordinates not present")
315    }
316
317    /// Fallible form of [`Self::at`].
318    pub fn get_at<const N: usize>(&self, coords: [usize; N]) -> Option<Expr<'a>> {
319        self.get_ref(&coords).copied()
320    }
321
322    fn get_ref(&self, coords: &[usize]) -> Option<&Expr<'a>> {
323        match &self.storage {
324            Storage::Dense { data, axes, .. } => {
325                grid_offset_coords(axes, coords).map(|off| &data[off])
326            }
327            Storage::Sparse(m) => m.get(&coords_to_key(coords)),
328        }
329    }
330}
331
332impl<'a, K: FromIndexKey, F> IndexedFamily<'a, K, F> {
333    /// Iterate the family's entries with each key decoded to the typed `K`.
334    pub fn keys(&self) -> impl Iterator<Item = (K, Expr<'a>)> + '_ {
335        self.iter().map(|(k, e)| (K::from_index_key(k), *e))
336    }
337}
338
339impl<'a, K, F, Q: Into<IndexKey>> Index<Q> for IndexedFamily<'a, K, F> {
340    type Output = Expr<'a>;
341    fn index(&self, key: Q) -> &Self::Output {
342        match &self.storage {
343            Storage::Sparse(m) => m.get(&key.into()).expect("indexed family: key not present"),
344            Storage::Dense { data, axes, .. } => {
345                let off = grid_offset(axes, &key.into()).expect("indexed family: key not present");
346                &data[off]
347            }
348        }
349    }
350}
351
352impl<'a, K, F> Index<&IndexKey> for IndexedFamily<'a, K, F> {
353    type Output = Expr<'a>;
354    fn index(&self, key: &IndexKey) -> &Self::Output {
355        match &self.storage {
356            Storage::Sparse(m) => m.get(key).expect("indexed family: key not present"),
357            Storage::Dense { data, axes, .. } => {
358                let off = grid_offset(axes, key).expect("indexed family: key not present");
359                &data[off]
360            }
361        }
362    }
363}
364
365impl<'a, K, F, const N: usize> Index<[usize; N]> for IndexedFamily<'a, K, F> {
366    type Output = Expr<'a>;
367    fn index(&self, coords: [usize; N]) -> &Self::Output {
368        self.get_ref(&coords).expect("indexed family: coordinates not present")
369    }
370}
371
372/// Build [`Storage`] from ordered keys and their already-registered handles.
373pub(crate) fn build_storage<'a>(
374    keys: Vec<IndexKey>,
375    axes: Option<Box<[Axis]>>,
376    values: Vec<Expr<'a>>,
377) -> Storage<'a> {
378    assert_eq!(keys.len(), values.len(), "indexed storage key/value length mismatch");
379    if let Some(axes) = axes {
380        if keys.iter().enumerate().all(|(i, key)| grid_offset(&axes, key) == Some(i)) {
381            return Storage::Dense { data: values, keys, axes };
382        }
383        let total = keys.len();
384        let mut data: Vec<Option<Expr<'a>>> = vec![None; total];
385        let mut kept: Vec<Option<IndexKey>> = vec![None; total];
386        for (key, expr) in keys.into_iter().zip(values) {
387            let off = grid_offset(&axes, &key).expect("dense grid key out of range");
388            data[off] = Some(expr);
389            kept[off] = Some(key);
390        }
391        let data = data.into_iter().map(|o| o.expect("dense grid had a gap")).collect();
392        let kept = kept.into_iter().map(|o| o.expect("dense grid had a gap")).collect();
393        Storage::Dense { data, keys: kept, axes }
394    } else {
395        let entries = keys.into_iter().zip(values).collect::<FxHashMap<_, _>>();
396        Storage::Sparse(entries)
397    }
398}
399
400/// Position of a key value along one axis, or `None` if out of `[start, start+len)`.
401fn axis_index(a: &Axis, v: i64) -> Option<usize> {
402    let d = v.checked_sub(a.start)?;
403    let u = usize::try_from(d).ok()?;
404    (u < a.len).then_some(u)
405}
406
407/// Row-major flat offset (axis 0 outermost) of an `IndexKey` in a dense grid, or
408/// `None` if the key's shape does not match the axes or is out of range.
409pub(crate) fn grid_offset(axes: &[Axis], key: &IndexKey) -> Option<usize> {
410    match (axes, key) {
411        ([a], IndexKey::Int(v)) => axis_index(a, *v),
412        (axes, IndexKey::Tuple(parts)) if parts.len() == axes.len() => {
413            let mut off = 0usize;
414            for (a, p) in axes.iter().zip(parts.iter()) {
415                off = off.checked_mul(a.len)?.checked_add(axis_index(a, p.as_i64()?)?)?;
416            }
417            Some(off)
418        }
419        _ => None,
420    }
421}
422
423/// Row-major flat offset from raw integer coordinates (key values).
424fn grid_offset_coords(axes: &[Axis], coords: &[usize]) -> Option<usize> {
425    if coords.len() != axes.len() {
426        return None;
427    }
428    let mut off = 0usize;
429    for (a, &c) in axes.iter().zip(coords) {
430        off = off * a.len + axis_index(a, i64::try_from(c).ok()?)?;
431    }
432    Some(off)
433}
434
435/// Build the `IndexKey` a coordinate array would hash to (sparse fallback for
436/// [`IndexedFamily::get_ref`]).
437fn coords_to_key(coords: &[usize]) -> IndexKey {
438    if let [single] = coords {
439        IndexKey::from(*single)
440    } else {
441        IndexKey::Tuple(coords.iter().map(|&c| IndexKey::from(c)).collect())
442    }
443}
444
445#[cfg(test)]
446mod tests {
447    use super::*;
448    use crate::Model;
449
450    #[test]
451    fn dense_storage_keeps_key_value_pairs_when_ordered_or_shuffled() {
452        let model = Model::new("storage_order");
453        let x = model.__var("x").build();
454        let y = model.__var("y").build();
455        for (keys, values) in [
456            (vec![IndexKey::Int(-1), IndexKey::Int(0)], vec![x, y]),
457            (vec![IndexKey::Int(0), IndexKey::Int(-1)], vec![y, x]),
458        ] {
459            let storage = build_storage(keys, Some(Box::new([Axis { start: -1, len: 2 }])), values);
460            let Storage::Dense { data, keys, .. } = storage else {
461                panic!("expected dense storage")
462            };
463            assert_eq!(data.iter().map(|expr| expr.id).collect::<Vec<_>>(), [x.id, y.id]);
464            assert_eq!(keys, [IndexKey::Int(-1), IndexKey::Int(0)]);
465        }
466    }
467}
oximo_core/indexed.rs

oximo_core/
indexed.rs
