oximo_core/
indexed.rs

1use std::marker::PhantomData;
2use std::ops::Index;
3
4use oximo_expr::Expr;
5use rustc_hash::FxHashMap;
6
7use crate::set::{Axis, FromIndexKey, IndexKey};
8
9/// Backing storage for an [`IndexedFamily`].
10///
11/// `Dense` is used when the domain is a contiguous integer grid (a range, or a
12/// product of ranges). `Sparse` (string sets, sparse `from_ints`, or any
13/// `filter`ed family) keeps the original hash map.
14#[derive(Clone)]
15pub(crate) enum Storage<'a> {
16    Dense { data: Vec<Expr<'a>>, keys: Vec<IndexKey>, axes: Box<[Axis]> },
17    Sparse(FxHashMap<IndexKey, Expr<'a>>),
18}
19
20mod sealed {
21    pub trait Sealed {}
22}
23
24/// Marker selecting which kind of indexed family an [`IndexedFamily`] is.
25///
26/// Sealed implementation detail: implemented only for [`VarFamily`] and
27/// [`ParamFamily`]. It exists so [`IndexedVar`] and [`IndexedParam`] are distinct
28/// types (only a parameter family can be re-bound) while sharing one
29/// implementation.
30#[doc(hidden)]
31pub trait Family: sealed::Sealed {
32    /// Type name used in [`Debug`](std::fmt::Debug) output.
33    const NAME: &'static str;
34}
35
36/// Marker for an indexed family of decision variables ([`IndexedVar`]).
37#[doc(hidden)]
38#[derive(Debug)]
39pub struct VarFamily;
40
41/// Marker for an indexed family of parameters ([`IndexedParam`]).
42#[doc(hidden)]
43#[derive(Debug)]
44pub struct ParamFamily;
45
46impl sealed::Sealed for VarFamily {}
47impl sealed::Sealed for ParamFamily {}
48impl Family for VarFamily {
49    const NAME: &'static str = "IndexedVar";
50}
51impl Family for ParamFamily {
52    const NAME: &'static str = "IndexedParam";
53}
54
55/// Indexed family: maps an `IndexKey` to a single-element `Expr` (a variable or a
56/// parameter), tagged with the key type `K` its domain decodes to and the family
57/// kind `F`.
58///
59/// You normally name this through the [`IndexedVar`]/[`IndexedParam`] aliases,
60/// constructed by the indexed form of the `variable!`/`param!` macros.
61///
62/// When the domain is a contiguous integer range (or a Cartesian product of
63/// ranges) the family is stored densely (see the internal `Storage`).
64/// String, sparse, and `filter`ed families fall back to a hash map.
65pub struct IndexedFamily<'a, K = IndexKey, F = VarFamily> {
66    pub(crate) storage: Storage<'a>,
67    pub(crate) _marker: PhantomData<fn() -> (K, F)>,
68}
69
70/// Indexed family of decision variables; see [`IndexedFamily`].
71pub type IndexedVar<'a, K = IndexKey> = IndexedFamily<'a, K, VarFamily>;
72
73/// Indexed family of re-bindable parameters; see [`IndexedFamily`].
74///
75/// Re-bind a single entry with [`Model::set_param_idx`](crate::Model::set_param_idx).
76pub type IndexedParam<'a, K = IndexKey> = IndexedFamily<'a, K, ParamFamily>;
77
78impl<'a, K, F> Clone for IndexedFamily<'a, K, F> {
79    fn clone(&self) -> Self {
80        Self { storage: self.storage.clone(), _marker: PhantomData }
81    }
82}
83
84impl<'a, K, F: Family> std::fmt::Debug for IndexedFamily<'a, K, F> {
85    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
86        f.debug_struct(F::NAME).field("len", &self.len()).field("dense", &self.is_dense()).finish()
87    }
88}
89
90impl<'a, K, F> IndexedFamily<'a, K, F> {
91    pub fn len(&self) -> usize {
92        match &self.storage {
93            Storage::Dense { data, .. } => data.len(),
94            Storage::Sparse(m) => m.len(),
95        }
96    }
97
98    pub fn is_empty(&self) -> bool {
99        self.len() == 0
100    }
101
102    /// Whether this family is stored densely (domain was a range or product of
103    /// ranges).
104    pub fn is_dense(&self) -> bool {
105        matches!(self.storage, Storage::Dense { .. })
106    }
107
108    /// Per-axis lengths when stored densely, else `None`.
109    pub fn shape(&self) -> Option<Box<[usize]>> {
110        match &self.storage {
111            Storage::Dense { axes, .. } => Some(axes.iter().map(|a| a.len).collect()),
112            Storage::Sparse(_) => None,
113        }
114    }
115
116    pub fn iter(&self) -> impl Iterator<Item = (&IndexKey, &Expr<'a>)> + '_ {
117        let it: Box<dyn Iterator<Item = (&IndexKey, &Expr<'a>)>> = match &self.storage {
118            Storage::Dense { data, keys, .. } => Box::new(keys.iter().zip(data.iter())),
119            Storage::Sparse(m) => Box::new(m.iter()),
120        };
121        it
122    }
123
124    pub fn get<Q: Into<IndexKey>>(&self, key: Q) -> Option<Expr<'a>> {
125        match &self.storage {
126            Storage::Sparse(m) => m.get(&key.into()).copied(),
127            Storage::Dense { data, axes, .. } => {
128                grid_offset(axes, &key.into()).map(|off| data[off])
129            }
130        }
131    }
132
133    /// Zero-allocation typed index by integer coordinates.
134    /// On a dense family this maps straight to a flat offset with no
135    /// `IndexKey` built. On a sparse family it falls back to building a key.
136    ///
137    /// # Panics
138    /// Panics if the coordinates are out of range/not present.
139    pub fn at<const N: usize>(&self, coords: [usize; N]) -> Expr<'a> {
140        *self.get_ref(&coords).expect("indexed family: coordinates not present")
141    }
142
143    /// Fallible form of [`Self::at`].
144    pub fn get_at<const N: usize>(&self, coords: [usize; N]) -> Option<Expr<'a>> {
145        self.get_ref(&coords).copied()
146    }
147
148    fn get_ref(&self, coords: &[usize]) -> Option<&Expr<'a>> {
149        match &self.storage {
150            Storage::Dense { data, axes, .. } => {
151                grid_offset_coords(axes, coords).map(|off| &data[off])
152            }
153            Storage::Sparse(m) => m.get(&coords_to_key(coords)),
154        }
155    }
156}
157
158impl<'a, K: FromIndexKey, F> IndexedFamily<'a, K, F> {
159    /// Iterate the family's entries with each key decoded to the typed `K`.
160    pub fn keys(&self) -> impl Iterator<Item = (K, Expr<'a>)> + '_ {
161        self.iter().map(|(k, e)| (K::from_index_key(k), *e))
162    }
163}
164
165impl<'a, K, F, Q: Into<IndexKey>> Index<Q> for IndexedFamily<'a, K, F> {
166    type Output = Expr<'a>;
167    fn index(&self, key: Q) -> &Self::Output {
168        match &self.storage {
169            Storage::Sparse(m) => m.get(&key.into()).expect("indexed family: key not present"),
170            Storage::Dense { data, axes, .. } => {
171                let off = grid_offset(axes, &key.into()).expect("indexed family: key not present");
172                &data[off]
173            }
174        }
175    }
176}
177
178impl<'a, K, F> Index<&IndexKey> for IndexedFamily<'a, K, F> {
179    type Output = Expr<'a>;
180    fn index(&self, key: &IndexKey) -> &Self::Output {
181        match &self.storage {
182            Storage::Sparse(m) => m.get(key).expect("indexed family: key not present"),
183            Storage::Dense { data, axes, .. } => {
184                let off = grid_offset(axes, key).expect("indexed family: key not present");
185                &data[off]
186            }
187        }
188    }
189}
190
191impl<'a, K, F, const N: usize> Index<[usize; N]> for IndexedFamily<'a, K, F> {
192    type Output = Expr<'a>;
193    fn index(&self, coords: [usize; N]) -> &Self::Output {
194        self.get_ref(&coords).expect("indexed family: coordinates not present")
195    }
196}
197
198/// Build [`Storage`] from a family's keys and optional dense axes, registering
199/// each element through `make`.
200pub(crate) fn build_storage<'a>(
201    keys: Vec<IndexKey>,
202    axes: Option<Box<[Axis]>>,
203    mut make: impl FnMut(&IndexKey) -> Expr<'a>,
204) -> Storage<'a> {
205    if let Some(axes) = axes {
206        let total = keys.len();
207        let mut data: Vec<Option<Expr<'a>>> = vec![None; total];
208        let mut kept: Vec<Option<IndexKey>> = vec![None; total];
209        for key in keys {
210            let expr = make(&key);
211            let off = grid_offset(&axes, &key).expect("dense grid key out of range");
212            data[off] = Some(expr);
213            kept[off] = Some(key);
214        }
215        let data = data.into_iter().map(|o| o.expect("dense grid had a gap")).collect();
216        let kept = kept.into_iter().map(|o| o.expect("dense grid had a gap")).collect();
217        Storage::Dense { data, keys: kept, axes }
218    } else {
219        let mut entries = FxHashMap::default();
220        for key in keys {
221            let expr = make(&key);
222            entries.insert(key, expr);
223        }
224        Storage::Sparse(entries)
225    }
226}
227
228/// Position of a key value along one axis, or `None` if out of `[start, start+len)`.
229fn axis_index(a: &Axis, v: i64) -> Option<usize> {
230    let d = v.checked_sub(a.start)?;
231    let u = usize::try_from(d).ok()?;
232    (u < a.len).then_some(u)
233}
234
235/// Row-major flat offset (axis 0 outermost) of an `IndexKey` in a dense grid, or
236/// `None` if the key's shape does not match the axes or is out of range.
237pub(crate) fn grid_offset(axes: &[Axis], key: &IndexKey) -> Option<usize> {
238    match (axes, key) {
239        ([a], IndexKey::Int(v)) => axis_index(a, *v),
240        (axes, IndexKey::Tuple(parts)) if parts.len() == axes.len() => {
241            let mut off = 0usize;
242            for (a, p) in axes.iter().zip(parts.iter()) {
243                off = off.checked_mul(a.len)?.checked_add(axis_index(a, p.as_i64()?)?)?;
244            }
245            Some(off)
246        }
247        _ => None,
248    }
249}
250
251/// Row-major flat offset from raw integer coordinates (key values).
252fn grid_offset_coords(axes: &[Axis], coords: &[usize]) -> Option<usize> {
253    if coords.len() != axes.len() {
254        return None;
255    }
256    let mut off = 0usize;
257    for (a, &c) in axes.iter().zip(coords) {
258        off = off * a.len + axis_index(a, i64::try_from(c).ok()?)?;
259    }
260    Some(off)
261}
262
263/// Build the `IndexKey` a coordinate array would hash to (sparse fallback for
264/// [`IndexedFamily::get_ref`]).
265fn coords_to_key(coords: &[usize]) -> IndexKey {
266    if let [single] = coords {
267        IndexKey::from(*single)
268    } else {
269        IndexKey::Tuple(coords.iter().map(|&c| IndexKey::from(c)).collect())
270    }
271}
oximo_core/indexed.rs

oximo_core/
indexed.rs
