oximo_core/
sum.rs

1use oximo_expr::Expr;
2
3use crate::set::{FromIndexKey, Set};
4
5/// Domain over which the `sum!` macro iterates. Lets a single sum accept either
6/// a [`Set`] (with typed key decoding via [`FromIndexKey`]) or a borrowed slice
7/// of `Copy` keys, without intermediate conversions.
8///
9/// Returns an iterator (rather than taking a callback) so the trait method
10/// monomorphizes through to the loop body in [`__sum_over`], allowing inlining
11/// in hot sums. Implementations are typically one line.
12#[diagnostic::on_unimplemented(
13    message = "`{Self}` is not a valid index domain over key type `{K}`",
14    label = "the domain's keys are not `{K}`",
15    note = "the loop/closure binding type must match the domain's keys",
16    note = "for a `Set<T>` write `for x in set` (the key type is inferred) or annotate `for x: T in set`. Integer ranges yield `usize`/`i64`/`i32`. A slice/`Vec`/array yields its element type"
17)]
18pub trait SumDomain<K> {
19    fn keys(&self) -> impl Iterator<Item = K> + '_;
20}
21
22// A typed set yields exactly its own key type.
23// The single `SumDomain` impl for `Set<K>`, so `sum!`/`constraint!`
24// can infer the closure parameter without an annotation
25// (the erased `Set` defaulted to `Set<IndexKey>`).
26impl<K: FromIndexKey> SumDomain<K> for Set<K> {
27    fn keys(&self) -> impl Iterator<Item = K> + '_ {
28        self.iter().map(|k| K::from_index_key(&k))
29    }
30}
31
32impl<K: Copy> SumDomain<K> for [K] {
33    fn keys(&self) -> impl Iterator<Item = K> + '_ {
34        self.iter().copied()
35    }
36}
37
38impl<K: Copy> SumDomain<K> for Vec<K> {
39    fn keys(&self) -> impl Iterator<Item = K> + '_ {
40        self.iter().copied()
41    }
42}
43
44impl<K: Copy, const N: usize> SumDomain<K> for [K; N] {
45    fn keys(&self) -> impl Iterator<Item = K> + '_ {
46        self.iter().copied()
47    }
48}
49
50// Forward through a reference, so a domain that is itself a reference (e.g. a
51// `&Set` function parameter passed to `sum!`/`constraint!`) is accepted.
52impl<K, D: SumDomain<K> + ?Sized> SumDomain<K> for &D {
53    fn keys(&self) -> impl Iterator<Item = K> + '_ {
54        (**self).keys()
55    }
56}
57
58// Integer ranges as sum domains. Iteration is lazy, so `sum!(x[i] for i in 0..n)`
59// allocates nothing. Provided for the common integer types the `sum!`/`constraint!`
60// macros default to.
61impl SumDomain<usize> for std::ops::Range<usize> {
62    fn keys(&self) -> impl Iterator<Item = usize> + '_ {
63        self.clone()
64    }
65}
66
67impl SumDomain<i64> for std::ops::Range<i64> {
68    fn keys(&self) -> impl Iterator<Item = i64> + '_ {
69        self.clone()
70    }
71}
72
73impl SumDomain<i32> for std::ops::Range<i32> {
74    fn keys(&self) -> impl Iterator<Item = i32> + '_ {
75        self.clone()
76    }
77}
78
79/// Sum an expression over every element of a domain.
80///
81/// Reads as the mathematical `sum_{k in domain} f(k)`. The closure parameter is
82/// either decoded from the domain's [`crate::set::IndexKey`] via [`FromIndexKey`] (when
83/// the domain is a [`Set`]) or yielded directly (when the domain is a slice
84/// of `Copy` keys).
85///
86/// # Panics
87/// Panics if `domain` is empty, the resulting expression has no arena to
88/// attach to.
89/// Macro-facing entry point backing the `sum!` macro. Not part of the stable
90/// public API.
91#[doc(hidden)]
92pub fn __sum_over<'a, K, D, F>(domain: &D, f: F) -> Expr<'a>
93where
94    D: SumDomain<K> + ?Sized,
95    F: FnMut(K) -> Expr<'a>,
96{
97    let terms: Vec<Expr<'a>> = domain.keys().map(f).collect();
98    assert!(!terms.is_empty(), "sum_over on empty domain");
99    terms.into_iter().sum()
100}
101
102#[cfg(test)]
103mod tests {
104    use oximo_expr::extract_linear;
105
106    use super::*;
107    use crate::model::Model;
108
109    #[test]
110    fn sum_over_scalar_set() {
111        let m = Model::new("scalar");
112        let items = Set::range(0..4);
113        let x = m.__indexed_var("x", &items).lb(0.0).build();
114
115        let total = __sum_over(&items, |i: usize| x[i]);
116        let arena = m.arena();
117        let terms = extract_linear(&arena, total.id).expect("linear");
118        assert_eq!(terms.coeffs.len(), 4);
119        assert!(terms.coeffs.iter().all(|(_, c)| (c - 1.0).abs() < f64::EPSILON));
120    }
121
122    #[test]
123    fn sum_over_tuple_set() {
124        let m = Model::new("tuple");
125        let plants = Set::strings(["seattle", "san-diego"]);
126        let markets = Set::strings(["nyc", "chicago", "topeka"]);
127        let routes = &plants * &markets;
128        let x = m.__indexed_var("x", &routes).lb(0.0).build();
129
130        let total = __sum_over(&routes, |(p, q): (String, String)| x[(p, q)]);
131        let arena = m.arena();
132        let terms = extract_linear(&arena, total.id).expect("linear");
133        assert_eq!(terms.coeffs.len(), 6);
134    }
135
136    #[test]
137    fn nested_sum_over_double_sum() {
138        let m = Model::new("nested");
139        let plants = Set::strings(["a", "b"]);
140        let markets = Set::strings(["x", "y", "z"]);
141        let routes = &plants * &markets;
142        let x = m.__indexed_var("x", &routes).lb(0.0).build();
143
144        let total = __sum_over(&plants, |p: String| __sum_over(&markets, |q: String| x[(&p, q)]));
145        let arena = m.arena();
146        let terms = extract_linear(&arena, total.id).expect("linear");
147        assert_eq!(terms.coeffs.len(), 6);
148    }
149
150    #[test]
151    fn sum_over_passes_typed_usize_key() {
152        let m = Model::new("usizekey");
153        let items = Set::range(0..3);
154        let x = m.__indexed_var("x", &items).lb(0.0).build();
155
156        let total = __sum_over(&items, |i: usize| x[i]);
157        let arena = m.arena();
158        let terms = extract_linear(&arena, total.id).expect("linear");
159        assert_eq!(terms.coeffs.len(), 3);
160    }
161
162    #[test]
163    fn sum_over_slice_of_usize() {
164        let m = Model::new("slice");
165        let items = Set::range(0..5);
166        let x = m.__indexed_var("x", &items).lb(0.0).build();
167
168        let picked: &[usize] = &[0, 2, 4];
169        let total = __sum_over(picked, |i: usize| x[i]);
170        let arena = m.arena();
171        let terms = extract_linear(&arena, total.id).expect("linear");
172        assert_eq!(terms.coeffs.len(), 3);
173    }
174
175    #[test]
176    fn sum_over_vec_of_usize() {
177        let m = Model::new("vec");
178        let items = Set::range(0..5);
179        let x = m.__indexed_var("x", &items).lb(0.0).build();
180
181        let picked: Vec<usize> = vec![1, 3];
182        let total = __sum_over(&picked, |i: usize| x[i]);
183        let arena = m.arena();
184        let terms = extract_linear(&arena, total.id).expect("linear");
185        assert_eq!(terms.coeffs.len(), 2);
186    }
187
188    #[test]
189    fn sum_over_array_of_usize() {
190        let m = Model::new("array");
191        let items = Set::range(0..5);
192        let x = m.__indexed_var("x", &items).lb(0.0).build();
193
194        let picked: [usize; 4] = [0, 1, 2, 3];
195        let total = __sum_over(&picked, |i: usize| x[i]);
196        let arena = m.arena();
197        let terms = extract_linear(&arena, total.id).expect("linear");
198        assert_eq!(terms.coeffs.len(), 4);
199    }
200
201    #[test]
202    #[should_panic(expected = "sum_over on empty domain")]
203    fn sum_over_empty_set_panics() {
204        let m = Model::new("empty");
205        let empty = Set::range(0..0);
206        let _x = m.__indexed_var("x", &Set::range(0..1)).lb(0.0).build();
207        let _ = __sum_over(&empty, |_: usize| panic!("closure should not run"));
208    }
209
210    #[test]
211    #[should_panic(expected = "sum_over on empty domain")]
212    fn sum_over_empty_slice_panics() {
213        let m = Model::new("empty_slice");
214        let _x = m.__indexed_var("x", &Set::range(0..1)).lb(0.0).build();
215        let empty: &[usize] = &[];
216        let _ = __sum_over(empty, |_: usize| panic!("closure should not run"));
217    }
218}
oximo_core/sum.rs

oximo_core/
sum.rs
