oximo_core/
set.rs

1use std::ops::Mul;
2
3use num_traits::PrimInt;
4use rayon::prelude::*;
5
6// Note: I used `Box<[IndexKey]>` over:
7//   - `Vec<IndexKey>`: saves one word
8//     (no capacity field) since tuples are never mutated after construction.
9//
10//   - `SmallVec<[IndexKey; N]>` is rejected by the compiler: recursive size
11//      cycle
12
13/// Heap-allocated tuple key payload. Immutable after construction.
14pub type IndexTuple = Box<[IndexKey]>;
15
16/// A finite, ordered index set.
17///
18/// Supports integer ranges, string lists, and arbitrary tuple lists (built via
19/// [`Set::product`] / the `&a * &b` operator, or constructed directly).
20#[derive(Clone, Debug)]
21pub enum Set {
22    Range(Vec<i64>),
23    Strings(Vec<String>),
24    Tuples(Vec<IndexTuple>),
25}
26
27impl Set {
28    /// Build an integer index set from a range over any primitive integer
29    /// type. Accepts `Range<i64>`, `Range<i32>`, `Range<usize>`, etc.
30    ///
31    /// The untyped literal form `Set::range(0..5)` defaults to `i32` and
32    /// works without annotation.
33    ///
34    /// # Panics
35    /// Panics if either range bound does not fit in `i64`.
36    pub fn range<T: PrimInt>(r: std::ops::Range<T>) -> Self {
37        let start = r.start.to_i64().expect("range start out of i64 range");
38        let end = r.end.to_i64().expect("range end out of i64 range");
39        Self::Range((start..end).collect())
40    }
41
42    /// Build an integer index set from an iterator of any primitive integer
43    /// type. Useful when keys are sparse or computed (e.g. only even indices).
44    ///
45    /// # Panics
46    /// Panics if any element does not fit in `i64`.
47    pub fn from_ints<T, I>(iter: I) -> Self
48    where
49        T: PrimInt,
50        I: IntoIterator<Item = T>,
51    {
52        Self::Range(
53            iter.into_iter().map(|v| v.to_i64().expect("element out of i64 range")).collect(),
54        )
55    }
56
57    pub fn strings<I, S>(iter: I) -> Self
58    where
59        I: IntoIterator<Item = S>,
60        S: Into<String>,
61    {
62        Self::Strings(iter.into_iter().map(Into::into).collect())
63    }
64
65    /// Build a tuple set directly from an iterator of keys.
66    pub fn tuples<I, T>(iter: I) -> Self
67    where
68        I: IntoIterator<Item = T>,
69        T: Into<IndexTuple>,
70    {
71        Self::Tuples(iter.into_iter().map(Into::into).collect())
72    }
73
74    /// Cartesian product of two sets. Inner tuple keys are flattened so a
75    /// product `(a * b) * c` yields 3-element tuples, not nested 2-tuples.
76    pub fn product(a: &Set, b: &Set) -> Self {
77        let mut out = Vec::new();
78        if let Some(capacity) = a.len().checked_mul(b.len()) {
79            // best-effort, fall back to dynamic growth on OOM
80            let _ = out.try_reserve_exact(capacity);
81        }
82        for ka in a {
83            for kb in b {
84                let mut parts: Vec<IndexKey> = Vec::new();
85                push_flat(&mut parts, ka.clone());
86                push_flat(&mut parts, kb);
87                out.push(parts.into_boxed_slice());
88            }
89        }
90        Self::Tuples(out)
91    }
92
93    /// Filter keys with a predicate. Preserves the original variant where
94    /// possible, filtered `Range`/`Strings` stay in their native variant.
95    pub fn filter<F>(&self, mut f: F) -> Self
96    where
97        F: FnMut(&IndexKey) -> bool,
98    {
99        match self {
100            Self::Range(v) => {
101                Self::Range(v.iter().copied().filter(|i| f(&IndexKey::Int(*i))).collect())
102            }
103            Self::Strings(v) => Self::Strings(
104                v.iter()
105                    .filter_map(|s| {
106                        let key = IndexKey::Str(s.clone());
107                        f(&key).then(|| match key {
108                            IndexKey::Str(owned) => owned,
109                            _ => unreachable!(),
110                        })
111                    })
112                    .collect(),
113            ),
114            Self::Tuples(v) => Self::Tuples(
115                v.iter()
116                    .filter_map(|t| {
117                        let key = IndexKey::Tuple(t.clone());
118                        f(&key).then(|| match key {
119                            IndexKey::Tuple(owned) => owned,
120                            _ => unreachable!(),
121                        })
122                    })
123                    .collect(),
124            ),
125        }
126    }
127
128    pub fn len(&self) -> usize {
129        match self {
130            Self::Range(v) => v.len(),
131            Self::Strings(v) => v.len(),
132            Self::Tuples(v) => v.len(),
133        }
134    }
135
136    pub fn is_empty(&self) -> bool {
137        self.len() == 0
138    }
139}
140
141fn push_flat(dst: &mut Vec<IndexKey>, k: IndexKey) {
142    match k {
143        IndexKey::Tuple(inner) => dst.extend(inner.into_vec()),
144        other => dst.push(other),
145    }
146}
147
148fn make_tuple<I: IntoIterator<Item = IndexKey>>(items: I) -> IndexTuple {
149    let mut v: Vec<IndexKey> = Vec::new();
150    for k in items {
151        push_flat(&mut v, k);
152    }
153    v.into_boxed_slice()
154}
155
156impl Mul<&Set> for &Set {
157    type Output = Set;
158    fn mul(self, rhs: &Set) -> Set {
159        Set::product(self, rhs)
160    }
161}
162
163/// A serializable index key from a [`Set`].
164#[derive(Clone, Debug, PartialEq, Eq, Hash)]
165pub enum IndexKey {
166    Int(i64),
167    Str(String),
168    Tuple(IndexTuple),
169}
170
171impl IndexKey {
172    /// Build a tuple key from any iterable of convertible items. Nested tuple
173    /// keys are flattened.
174    pub fn tuple<I, T>(iter: I) -> Self
175    where
176        I: IntoIterator<Item = T>,
177        T: Into<IndexKey>,
178    {
179        Self::Tuple(make_tuple(iter.into_iter().map(Into::into)))
180    }
181
182    pub fn as_i64(&self) -> Option<i64> {
183        if let Self::Int(v) = self { Some(*v) } else { None }
184    }
185
186    pub fn as_str(&self) -> Option<&str> {
187        if let Self::Str(s) = self { Some(s.as_str()) } else { None }
188    }
189
190    pub fn as_tuple(&self) -> Option<&[IndexKey]> {
191        if let Self::Tuple(t) = self { Some(&t[..]) } else { None }
192    }
193}
194
195impl From<i64> for IndexKey {
196    fn from(v: i64) -> Self {
197        Self::Int(v)
198    }
199}
200
201impl From<i32> for IndexKey {
202    fn from(v: i32) -> Self {
203        Self::Int(i64::from(v))
204    }
205}
206
207impl From<usize> for IndexKey {
208    fn from(v: usize) -> Self {
209        Self::Int(i64::try_from(v).expect("usize -> i64 overflow"))
210    }
211}
212
213impl From<&str> for IndexKey {
214    fn from(s: &str) -> Self {
215        Self::Str(s.to_owned())
216    }
217}
218
219impl From<String> for IndexKey {
220    fn from(s: String) -> Self {
221        Self::Str(s)
222    }
223}
224
225impl<A, B> From<(A, B)> for IndexKey
226where
227    A: Into<IndexKey>,
228    B: Into<IndexKey>,
229{
230    fn from(t: (A, B)) -> Self {
231        Self::Tuple(make_tuple([t.0.into(), t.1.into()]))
232    }
233}
234
235impl<A, B, C> From<(A, B, C)> for IndexKey
236where
237    A: Into<IndexKey>,
238    B: Into<IndexKey>,
239    C: Into<IndexKey>,
240{
241    fn from(t: (A, B, C)) -> Self {
242        Self::Tuple(make_tuple([t.0.into(), t.1.into(), t.2.into()]))
243    }
244}
245
246impl<A, B, C, D> From<(A, B, C, D)> for IndexKey
247where
248    A: Into<IndexKey>,
249    B: Into<IndexKey>,
250    C: Into<IndexKey>,
251    D: Into<IndexKey>,
252{
253    fn from(t: (A, B, C, D)) -> Self {
254        Self::Tuple(make_tuple([t.0.into(), t.1.into(), t.2.into(), t.3.into()]))
255    }
256}
257
258/// Typed projection of an [`IndexKey`]. Implementations panic when the
259/// key's shape does not match the target type, the same contract as
260/// [`crate::indexed::IndexedVar`] indexing on a missing key.
261///
262/// Used by [`crate::model::Model::add_constraints_over`] (and similar rule
263/// helpers) to give the closure typed indices directly:
264///
265/// ```ignore
266/// m.add_constraints_over("supply", &(&plants * &markets), |(p, m): (String, String)| {
267///     // p, m are native String, no manual unpack
268///     ...
269/// });
270/// ```
271pub trait FromIndexKey: Sized {
272    fn from_index_key(k: &IndexKey) -> Self;
273}
274
275impl FromIndexKey for IndexKey {
276    fn from_index_key(k: &IndexKey) -> Self {
277        k.clone()
278    }
279}
280
281impl FromIndexKey for i64 {
282    fn from_index_key(k: &IndexKey) -> Self {
283        k.as_i64().unwrap_or_else(|| panic!("expected Int key, got {k:?}"))
284    }
285}
286
287impl FromIndexKey for i32 {
288    fn from_index_key(k: &IndexKey) -> Self {
289        let v = i64::from_index_key(k);
290        i32::try_from(v).unwrap_or_else(|_| panic!("key {v} out of i32 range"))
291    }
292}
293
294impl FromIndexKey for usize {
295    fn from_index_key(k: &IndexKey) -> Self {
296        let v = i64::from_index_key(k);
297        usize::try_from(v).unwrap_or_else(|_| panic!("key {v} out of usize range"))
298    }
299}
300
301impl FromIndexKey for String {
302    fn from_index_key(k: &IndexKey) -> Self {
303        k.as_str().unwrap_or_else(|| panic!("expected Str key, got {k:?}")).to_owned()
304    }
305}
306
307fn tuple_parts<'a>(k: &'a IndexKey, expected: usize) -> &'a [IndexKey] {
308    let p = k.as_tuple().unwrap_or_else(|| panic!("expected Tuple key, got {k:?}"));
309    assert_eq!(p.len(), expected, "expected tuple of arity {expected}, got arity {}", p.len());
310    p
311}
312
313impl<A, B> FromIndexKey for (A, B)
314where
315    A: FromIndexKey,
316    B: FromIndexKey,
317{
318    fn from_index_key(k: &IndexKey) -> Self {
319        let p = tuple_parts(k, 2);
320        (A::from_index_key(&p[0]), B::from_index_key(&p[1]))
321    }
322}
323
324impl<A, B, C> FromIndexKey for (A, B, C)
325where
326    A: FromIndexKey,
327    B: FromIndexKey,
328    C: FromIndexKey,
329{
330    fn from_index_key(k: &IndexKey) -> Self {
331        let p = tuple_parts(k, 3);
332        (A::from_index_key(&p[0]), B::from_index_key(&p[1]), C::from_index_key(&p[2]))
333    }
334}
335
336impl<A, B, C, D> FromIndexKey for (A, B, C, D)
337where
338    A: FromIndexKey,
339    B: FromIndexKey,
340    C: FromIndexKey,
341    D: FromIndexKey,
342{
343    fn from_index_key(k: &IndexKey) -> Self {
344        let p = tuple_parts(k, 4);
345        (
346            A::from_index_key(&p[0]),
347            B::from_index_key(&p[1]),
348            C::from_index_key(&p[2]),
349            D::from_index_key(&p[3]),
350        )
351    }
352}
353
354impl<'a> IntoIterator for &'a Set {
355    type Item = IndexKey;
356    type IntoIter = SetIter<'a>;
357    fn into_iter(self) -> Self::IntoIter {
358        self.iter()
359    }
360}
361
362impl Set {
363    pub fn iter(&self) -> SetIter<'_> {
364        SetIter { set: self, pos: 0 }
365    }
366
367    pub fn par_iter(&self) -> impl ParallelIterator<Item = IndexKey> + '_ {
368        match self {
369            Self::Range(v) => v.par_iter().map(|i| IndexKey::Int(*i)).collect::<Vec<_>>(),
370            Self::Strings(v) => v.par_iter().map(|s| IndexKey::Str(s.clone())).collect::<Vec<_>>(),
371            Self::Tuples(v) => v.par_iter().map(|t| IndexKey::Tuple(t.clone())).collect::<Vec<_>>(),
372        }
373        .into_par_iter()
374    }
375}
376
377#[derive(Debug)]
378pub struct SetIter<'a> {
379    set: &'a Set,
380    pos: usize,
381}
382
383impl<'a> Iterator for SetIter<'a> {
384    type Item = IndexKey;
385    fn next(&mut self) -> Option<Self::Item> {
386        let out = match self.set {
387            Set::Range(v) => v.get(self.pos).copied().map(IndexKey::Int),
388            Set::Strings(v) => v.get(self.pos).cloned().map(IndexKey::Str),
389            Set::Tuples(v) => v.get(self.pos).cloned().map(IndexKey::Tuple),
390        };
391        if out.is_some() {
392            self.pos += 1;
393        }
394        out
395    }
396}
oximo_core/set.rs

oximo_core/
set.rs
