Skip to main content

oximo_core/
model.rs

1use std::cell::{Cell, Ref, RefCell};
2use std::marker::PhantomData;
3
4use oximo_expr::{Expr, ExprArena, ExprClass, ExprId, ParamId, VarId, classify};
5use rustc_hash::FxHashMap;
6use smol_str::SmolStr;
7
8use crate::constraint::{Constraint, ConstraintExpr, ConstraintId, IntoRhs, Relate, Sense};
9use crate::domain::Domain;
10use crate::error::{Error, Result};
11use crate::indexed::{IndexedFamily, IndexedParam, IndexedVar, build_storage};
12use crate::objective::{Objective, ObjectiveSense};
13use crate::param::Parameter;
14use crate::set::{Axis, FromIndexKey, IndexKey, Set};
15use crate::soc::{SocConstraint, SocConstraintId, detect_soc};
16use crate::var::{VarBuilder, Variable};
17
18/// The kind of mathematical program a `Model` represents.
19///
20/// This is inferred from the variables and expressions in the model, not set
21/// explicitly by the user. See [`Model::kind`] for the exact decision ladder.
22///
23/// The `MI*` variant of each class is picked when any variable has an integer
24/// domain. The continuous classes are, from most to least general:
25///
26/// - `NLP`: some expression is nonlinear (degree > 2, transcendental, division)
27/// - `QCP`: some constraint is quadratic and not recognized as a second-order
28///   cone
29/// - `SOCP`: second-order cone constraints are present (explicit
30///   [`crate::SocConstraint`]s or SOC-shaped quadratic constraints recognized
31///   by [`crate::detect_soc`]); the objective may be linear or quadratic
32/// - `QP`: quadratic objective, linear constraints
33/// - `LP`: everything linear
34#[derive(Copy, Clone, Debug, PartialEq, Eq)]
35pub enum ModelKind {
36    LP,
37    MILP,
38    QP,
39    MIQP,
40    QCP,
41    MIQCP,
42    SOCP,
43    MISOCP,
44    NLP,
45    MINLP,
46}
47
48/// The optimization model. Owns the expression arena, variable/parameter
49/// registries, constraints, and (optional) objective.
50///
51/// `Model` uses interior mutability so the builder API can take `&self`
52/// references.
53///
54/// Variables, constraints, and the objective are added through
55/// `RefCell`s under the hood.
56pub struct Model {
57    pub name: String,
58    pub(crate) arena: RefCell<ExprArena>,
59    pub(crate) variables: RefCell<Vec<Variable>>,
60    pub(crate) var_names: RefCell<FxHashMap<SmolStr, VarId>>,
61    pub(crate) parameters: RefCell<Vec<Parameter>>,
62    pub(crate) param_names: RefCell<FxHashMap<SmolStr, ParamId>>,
63    pub(crate) constraints: RefCell<Vec<Constraint>>,
64    pub(crate) constraint_names: RefCell<FxHashMap<SmolStr, ConstraintId>>,
65    pub(crate) soc_constraints: RefCell<Vec<SocConstraint>>,
66    pub(crate) soc_names: RefCell<FxHashMap<SmolStr, SocConstraintId>>,
67    pub(crate) objective: RefCell<Option<Objective>>,
68    objective_declared: Cell<bool>,
69    cached_kind: Cell<Option<ModelKind>>,
70    /// Monotonic counter for auto-naming anonymous constraints registered via
71    /// the `constraint!` macro.
72    auto_seq: Cell<u32>,
73}
74
75impl std::fmt::Debug for Model {
76    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        f.debug_struct("Model")
78            .field("name", &self.name)
79            .field("vars", &self.variables.borrow().len())
80            .field("params", &self.parameters.borrow().len())
81            .field("constraints", &self.constraints.borrow().len())
82            .field("soc_constraints", &self.soc_constraints.borrow().len())
83            .field("has_objective", &self.objective.borrow().is_some())
84            .field("feasibility", &self.is_feasibility())
85            .finish()
86    }
87}
88
89impl Model {
90    pub fn new(name: impl Into<String>) -> Self {
91        Self {
92            name: name.into(),
93            arena: RefCell::new(ExprArena::new()),
94            variables: RefCell::new(Vec::new()),
95            var_names: RefCell::new(FxHashMap::default()),
96            parameters: RefCell::new(Vec::new()),
97            param_names: RefCell::new(FxHashMap::default()),
98            constraints: RefCell::new(Vec::new()),
99            constraint_names: RefCell::new(FxHashMap::default()),
100            soc_constraints: RefCell::new(Vec::new()),
101            soc_names: RefCell::new(FxHashMap::default()),
102            objective: RefCell::new(None),
103            objective_declared: Cell::new(false),
104            cached_kind: Cell::new(None),
105            auto_seq: Cell::new(0),
106        }
107    }
108
109    // Variables
110
111    /// Macro-facing entry point backing the `variable!` macro. Not part of the
112    /// stable public API.
113    #[doc(hidden)]
114    pub fn __var(&self, name: impl Into<SmolStr>) -> VarBuilder<'_> {
115        VarBuilder {
116            model: self,
117            name: name.into(),
118            lb: f64::NEG_INFINITY,
119            ub: f64::INFINITY,
120            domain: Domain::Real,
121            initial: None,
122        }
123    }
124
125    /// Called by [`VarBuilder::build`]. Pushes the var into the registry and
126    /// returns its `Expr` handle.
127    pub(crate) fn register_var<'a>(&'a self, b: VarBuilder<'a>) -> Expr<'a> {
128        let mut names = self.var_names.borrow_mut();
129        assert!(!names.contains_key(&b.name), "variable name {:?} already registered", b.name);
130        let mut vars = self.variables.borrow_mut();
131        let id = VarId(u32::try_from(vars.len()).expect("variable count overflow"));
132        vars.push(Variable {
133            id,
134            name: b.name.clone(),
135            domain: b.domain,
136            lb: b.lb,
137            ub: b.ub,
138            initial: b.initial,
139        });
140        names.insert(b.name, id);
141        drop(vars);
142        drop(names);
143        self.cached_kind.set(None);
144        Expr::from_var(&self.arena, id)
145    }
146
147    /// Macro-facing entry point backing the indexed form of the `variable!`
148    /// macro. Not part of the stable public API.
149    #[doc(hidden)]
150    pub fn __indexed_var<'a, K>(
151        &'a self,
152        name: impl Into<String>,
153        set: &Set<K>,
154    ) -> IndexedVarBuilder<'a, K> {
155        IndexedVarBuilder {
156            model: self,
157            base_name: name.into(),
158            keys: set.iter().collect(),
159            axes: set.axes().map(Box::from),
160            lb: f64::NEG_INFINITY,
161            ub: f64::INFINITY,
162            lb_by: None,
163            ub_by: None,
164            domain: Domain::Real,
165            _k: PhantomData,
166        }
167    }
168
169    pub fn variable_id(&self, name: &str) -> Option<VarId> {
170        self.var_names.borrow().get(name).copied()
171    }
172
173    pub fn variables(&self) -> Ref<'_, Vec<Variable>> {
174        self.variables.borrow()
175    }
176
177    pub fn arena(&self) -> Ref<'_, ExprArena> {
178        self.arena.borrow()
179    }
180
181    pub fn num_variables(&self) -> usize {
182        self.variables.borrow().len()
183    }
184
185    /// Fix a single-variable expression to `value`.
186    /// Convenience over [`Self::fix_var`] for handles from the `variable!` macro
187    /// or [`crate::IndexedVar`] indexing.
188    ///
189    /// # Panics
190    ///
191    /// Panics if `e` is not a bare variable handle.
192    pub fn fix(&self, e: Expr<'_>, value: f64) {
193        let id = e.var_id().expect("Model::fix expects a single-variable expression");
194        self.fix_var(id, value);
195    }
196
197    /// Fix variable `id` to `value` by setting `lb = ub = value`.
198    pub fn fix_var(&self, id: VarId, value: f64) {
199        let mut vars = self.variables.borrow_mut();
200        let v = &mut vars[id.index()];
201        v.lb = value;
202        v.ub = value;
203        drop(vars);
204        self.cached_kind.set(None);
205    }
206
207    /// Set the initial (warm-start) value of a single-variable expression.
208    /// The macro API has no bound-style syntax for warm starts, so this is the
209    /// supported way to seed `variable!`-declared variables.
210    ///
211    /// # Panics
212    ///
213    /// Panics if `e` is not a bare variable handle.
214    pub fn set_initial(&self, e: Expr<'_>, value: f64) {
215        let id = e.var_id().expect("Model::set_initial expects a single-variable expression");
216        self.variables.borrow_mut()[id.index()].initial = Some(value);
217    }
218
219    /// Restore bounds on variable `id`. Pass `f64::NEG_INFINITY` / `f64::INFINITY`
220    /// to restore an unbounded direction.
221    pub fn unfix_var(&self, id: VarId, lb: f64, ub: f64) {
222        let mut vars = self.variables.borrow_mut();
223        let v = &mut vars[id.index()];
224        v.lb = lb;
225        v.ub = ub;
226        drop(vars);
227        self.cached_kind.set(None);
228    }
229
230    // Parameters
231
232    /// Macro-facing entry point backing the `param!` macro. Not part of the
233    /// stable public API.
234    ///
235    /// Registers a named scalar parameter initialized to `value`, returning an
236    /// [`Expr`] handle that references it symbolically. A parameter behaves like a
237    /// constant coefficient (`param * var` is linear) but stays symbolic so it can
238    /// be re-bound with [`Self::set_param`] / [`Self::set_param_id`] between solves
239    /// without rebuilding the model.
240    ///
241    /// # Panics
242    ///
243    /// Panics if a parameter with the same name is already registered.
244    #[doc(hidden)]
245    pub fn __param<'a>(&'a self, name: impl Into<SmolStr>, value: f64) -> Expr<'a> {
246        self.register_param(name.into(), value)
247    }
248
249    /// Register one scalar parameter named `name` initialized to `value` and
250    /// return its `Expr` handle. Shared by [`Self::__param`] and the indexed
251    /// builder.
252    ///
253    /// # Panics
254    ///
255    /// Panics if a parameter with the same name is already registered.
256    fn register_param(&self, name: SmolStr, value: f64) -> Expr<'_> {
257        assert!(
258            !self.param_names.borrow().contains_key(&name),
259            "parameter name {name:?} already registered"
260        );
261        let (id, node) = {
262            let mut a = self.arena.borrow_mut();
263            let id = a.new_param(value);
264            (id, a.param(id))
265        };
266        self.parameters.borrow_mut().push(Parameter { id, name: name.clone() });
267        self.param_names.borrow_mut().insert(name, id);
268        self.cached_kind.set(None);
269        Expr::new(node, &self.arena)
270    }
271
272    /// Macro-facing entry point backing the indexed form of the `param!` macro
273    /// (`param!(m, cost[i in items] = data[i])`). Registers one scalar parameter
274    /// per key, evaluating `value` on the typed key, and returns an
275    /// [`IndexedParam`]. Not part of the stable public API.
276    ///
277    /// # Panics
278    ///
279    /// Panics if a per-key parameter name collides with one already registered.
280    #[doc(hidden)]
281    pub fn __indexed_param<'a, K, F>(
282        &'a self,
283        name: impl Into<String>,
284        set: &Set<K>,
285        mut value: F,
286    ) -> IndexedParam<'a, K>
287    where
288        K: FromIndexKey,
289        F: FnMut(K) -> f64,
290    {
291        let base = name.into();
292        let axes = set.axes().map(Box::from);
293        let keys: Vec<IndexKey> = set.iter().collect();
294        let make = |key: &IndexKey| -> Expr<'a> {
295            let pname: SmolStr = format_index_name(&base, key).into();
296            let v = value(K::from_index_key(key));
297            self.register_param(pname, v)
298        };
299        let storage = build_storage(keys, axes, make);
300        IndexedFamily { storage, _marker: PhantomData }
301    }
302
303    /// Re-bind the parameter at `key` of an indexed family to `value`. Takes
304    /// effect on the next solve.
305    ///
306    /// # Panics
307    ///
308    /// Panics if `key` is not present in the family, or if `params` was built on
309    /// a different `Model`.
310    pub fn set_param_idx<K, Q: Into<IndexKey>>(
311        &self,
312        params: &IndexedParam<'_, K>,
313        key: Q,
314        value: f64,
315    ) {
316        let e = params.get(key).expect("set_param_idx: key not present in indexed parameter");
317        assert!(
318            std::ptr::eq(e.arena, std::ptr::from_ref(&self.arena)),
319            "set_param_idx: indexed parameter belongs to a different model"
320        );
321        let id = e.param_id().expect("indexed parameter entry is not a parameter handle");
322        self.set_param_id(id, value);
323    }
324
325    /// Current value bound to the parameter at `key` of an indexed family, or
326    /// `None` if the key is absent.
327    pub fn param_value_idx<K, Q: Into<IndexKey>>(
328        &self,
329        params: &IndexedParam<'_, K>,
330        key: Q,
331    ) -> Option<f64> {
332        params.get(key).and_then(|e| self.param_value_of(e))
333    }
334
335    /// Re-bind the parameter referenced by handle `p` to `value`.
336    ///
337    /// # Panics
338    ///
339    /// Panics if `p` is not a bare parameter handle (one returned by the `param!`
340    /// macro).
341    pub fn set_param(&self, p: Expr<'_>, value: f64) {
342        let id = p.param_id().expect("Model::set_param expects a single-parameter expression");
343        self.set_param_id(id, value);
344    }
345
346    /// Re-bind parameter `id` to `value`. Takes effect on the next solve.
347    ///
348    /// The value is stored only in the expression arena (its single source of
349    /// truth); extraction and evaluation read it from there.
350    pub fn set_param_id(&self, id: ParamId, value: f64) {
351        self.arena.borrow_mut().set_param_value(id, value);
352        self.cached_kind.set(None);
353    }
354
355    /// Current value bound to parameter `id`.
356    ///
357    /// # Panics
358    ///
359    /// Panics if `id` does not belong to a parameter registered on this model.
360    pub fn param_value(&self, id: ParamId) -> f64 {
361        self.arena.borrow().param_value(id)
362    }
363
364    /// Current value of the parameter referenced by handle `p`, or `None` if
365    /// `p` is not a bare parameter handle.
366    pub fn param_value_of(&self, p: Expr<'_>) -> Option<f64> {
367        p.param_id().map(|id| self.param_value(id))
368    }
369
370    pub fn parameter_id(&self, name: &str) -> Option<ParamId> {
371        self.param_names.borrow().get(name).copied()
372    }
373
374    pub fn parameters(&self) -> Ref<'_, Vec<Parameter>> {
375        self.parameters.borrow()
376    }
377
378    pub fn num_parameters(&self) -> usize {
379        self.parameters.borrow().len()
380    }
381
382    // Constraints
383
384    /// Macro-facing entry point backing the `constraint!` macro. Not part of the
385    /// stable public API.
386    ///
387    /// # Panics
388    ///
389    /// Panics if a constraint with the same name is already registered, or if
390    /// the constraint count exceeds `u32::MAX`.
391    #[doc(hidden)]
392    pub fn __add_constraint(
393        &self,
394        name: impl Into<SmolStr>,
395        c: ConstraintExpr<'_>,
396    ) -> ConstraintId {
397        let (lower, upper) = match c.sense {
398            Sense::Le => (f64::NEG_INFINITY, c.rhs),
399            Sense::Ge => (c.rhs, f64::INFINITY),
400            Sense::Eq => (c.rhs, c.rhs),
401        };
402        self.register_constraint(name.into(), c.lhs.id, lower, upper)
403    }
404
405    /// Push a constraint row `lower <= lhs <= upper` into the registry. Shared by
406    /// [`Self::__add_constraint`] and the range entry points.
407    ///
408    /// # Panics
409    ///
410    /// Panics if a constraint with the same name is already registered, if a
411    /// bound is NaN, or if the constraint count exceeds `u32::MAX`.
412    fn register_constraint(
413        &self,
414        name: SmolStr,
415        lhs: ExprId,
416        lower: f64,
417        upper: f64,
418    ) -> ConstraintId {
419        assert!(
420            !lower.is_nan() && !upper.is_nan(),
421            "constraint {name:?} has NaN bound (lower={lower}, upper={upper})"
422        );
423        let mut by_name = self.constraint_names.borrow_mut();
424        assert!(!by_name.contains_key(&name), "constraint name {name:?} already registered");
425        let mut all = self.constraints.borrow_mut();
426        let id = ConstraintId(u32::try_from(all.len()).expect("constraint count overflow"));
427        all.push(Constraint { name: name.clone(), lhs, lower, upper, active: true });
428        by_name.insert(name, id);
429        self.cached_kind.set(None);
430        id
431    }
432
433    /// A fresh unique auto-name `_c{n}`, skipping any a user already took.
434    fn next_auto_name(&self) -> SmolStr {
435        loop {
436            let n = self.auto_seq.get();
437            self.auto_seq.set(n + 1);
438            let candidate: SmolStr = format!("_c{n}").into();
439            if !self.constraint_names.borrow().contains_key(&candidate) {
440                break candidate;
441            }
442        }
443    }
444
445    /// Register an anonymous constraint, deriving a unique name `_c{n}` from an
446    /// internal counter. Backs the name-less form of the `constraint!` macro.
447    #[doc(hidden)]
448    pub fn __add_constraint_auto(&self, c: ConstraintExpr<'_>) -> ConstraintId {
449        self.__add_constraint(self.next_auto_name(), c)
450    }
451
452    /// Bulk-register constraints. Each entry is `(name, ConstraintExpr)`.
453    /// Useful with `.par_iter().map(...).collect()` style construction.
454    pub fn add_constraints<'a, I>(&'a self, items: I)
455    where
456        I: IntoIterator<Item = (SmolStr, ConstraintExpr<'a>)>,
457    {
458        for (name, c) in items {
459            self.__add_constraint(name, c);
460        }
461    }
462
463    /// Macro-facing entry point backing the indexed-family form of the
464    /// `constraint!` macro. The closure receives the index as a typed value `K`
465    /// (any [`FromIndexKey`]: `i64`, `i32`, `usize`, `String`, raw `IndexKey`, or
466    /// tuples up to arity 4). Not part of the stable public API.
467    #[doc(hidden)]
468    pub fn __add_constraints_over<'a, K, F>(&'a self, name_prefix: &str, set: &Set<K>, mut rule: F)
469    where
470        K: FromIndexKey,
471        F: FnMut(K) -> ConstraintExpr<'a>,
472    {
473        for key in set {
474            let typed = K::from_index_key(&key);
475            let c = rule(typed);
476            let name: SmolStr = format_index_name(name_prefix, &key).into();
477            self.__add_constraint(name, c);
478        }
479    }
480
481    /// Macro-facing entry point for a two-sided range `lo <= mid <= hi`.
482    ///
483    /// Collapses to a single interval [`Constraint`] named `name` only when both
484    /// bounds are pure constants and the body is linear (the condition under which
485    /// one two-sided row is representable).
486    #[doc(hidden)]
487    pub fn __add_range<'a, B1, B2>(&'a self, name: &str, mid: Expr<'a>, lo: B1, hi: B2)
488    where
489        B1: IntoRhs<'a>,
490        B2: IntoRhs<'a>,
491    {
492        if let Some((lower, upper)) = self.collapse_bounds(mid.id, &lo, &hi) {
493            self.register_constraint(name.into(), mid.id, lower, upper);
494        } else {
495            self.__add_constraint(format!("{name}_lo"), mid.ge(lo));
496            self.__add_constraint(format!("{name}_hi"), mid.le(hi));
497        }
498    }
499
500    /// Anonymous form of [`Self::__add_range`] (auto-named rows).
501    #[doc(hidden)]
502    pub fn __add_range_auto<'a, B1, B2>(&'a self, mid: Expr<'a>, lo: B1, hi: B2)
503    where
504        B1: IntoRhs<'a>,
505        B2: IntoRhs<'a>,
506    {
507        if let Some((lower, upper)) = self.collapse_bounds(mid.id, &lo, &hi) {
508            self.register_constraint(self.next_auto_name(), mid.id, lower, upper);
509        } else {
510            self.__add_constraint_auto(mid.ge(lo));
511            self.__add_constraint_auto(mid.le(hi));
512        }
513    }
514
515    /// The interval `(lower, upper)` a range collapses to, or `None` (keep two
516    /// rows). Requires both bounds to be literal constants and the body `mid` to
517    /// be linear.
518    fn collapse_bounds<'a>(
519        &self,
520        mid: ExprId,
521        lo: &impl IntoRhs<'a>,
522        hi: &impl IntoRhs<'a>,
523    ) -> Option<(f64, f64)> {
524        let lower = lo.const_bound()?;
525        let upper = hi.const_bound()?;
526        (classify(&self.arena.borrow(), mid) == ExprClass::Linear).then_some((lower, upper))
527    }
528
529    /// Macro-facing entry point for a two-sided range family. One row per key,
530    /// each collapsing to a single interval constraint when both bounds are
531    /// constant (see [`Self::__add_range`]).
532    #[doc(hidden)]
533    pub fn __add_range_constraints_over<'a, K, B1, B2, F>(
534        &'a self,
535        name: &str,
536        set: &Set<K>,
537        mut rule: F,
538    ) where
539        K: FromIndexKey,
540        B1: IntoRhs<'a>,
541        B2: IntoRhs<'a>,
542        F: FnMut(K) -> (Expr<'a>, B1, B2),
543    {
544        for key in set {
545            let (mid, lo, hi) = rule(K::from_index_key(&key));
546            let row_name = format_index_name(name, &key);
547            self.__add_range(&row_name, mid, lo, hi);
548        }
549    }
550
551    pub fn constraints(&self) -> Ref<'_, Vec<Constraint>> {
552        self.constraints.borrow()
553    }
554
555    pub fn num_constraints(&self) -> usize {
556        self.constraints.borrow().len()
557    }
558
559    pub fn constraint_id(&self, name: &str) -> Option<ConstraintId> {
560        self.constraint_names.borrow().get(name).copied()
561    }
562
563    // Second-order cone constraints
564
565    /// Register the explicit second-order cone constraint
566    /// `||terms||_2 <= bound`.
567    ///
568    /// Every member of `terms` and the `bound` must be affine; the bound is
569    /// additionally constrained to be nonnegative by the cone itself, so
570    /// backends emit a `bound >= 0` side condition where needed.
571    ///
572    /// # Panics
573    ///
574    /// Panics if a SOC constraint with the same name is already registered, if
575    /// `terms` is empty, if any term or the bound is not affine, or if the
576    /// count exceeds `u32::MAX`.
577    pub fn add_soc_constraint<'a>(
578        &'a self,
579        name: impl Into<SmolStr>,
580        terms: impl IntoIterator<Item = Expr<'a>>,
581        bound: Expr<'a>,
582    ) -> SocConstraintId {
583        let name = name.into();
584        let arena = self.arena.borrow();
585        let terms: Vec<ExprId> = terms
586            .into_iter()
587            .map(|e| {
588                assert!(
589                    classify(&arena, e.id) == ExprClass::Linear,
590                    "SOC constraint {name:?} has a non-affine term"
591                );
592                e.id
593            })
594            .collect();
595        assert!(!terms.is_empty(), "SOC constraint {name:?} has no terms");
596        assert!(
597            classify(&arena, bound.id) == ExprClass::Linear,
598            "SOC constraint {name:?} has a non-affine bound"
599        );
600        drop(arena);
601
602        let mut by_name = self.soc_names.borrow_mut();
603        assert!(!by_name.contains_key(&name), "SOC constraint name {name:?} already registered");
604        let mut all = self.soc_constraints.borrow_mut();
605        let id = SocConstraintId(u32::try_from(all.len()).expect("SOC constraint count overflow"));
606        all.push(SocConstraint { name: name.clone(), terms, bound: bound.id, active: true });
607        by_name.insert(name, id);
608        self.cached_kind.set(None);
609        id
610    }
611
612    /// A fresh unique auto-name `_soc{n}` in the SOC namespace, skipping any a
613    /// user already took. Shares `auto_seq` with [`Self::next_auto_name`]; the
614    /// prefixes differ, so the two namespaces never collide.
615    fn next_auto_soc_name(&self) -> SmolStr {
616        loop {
617            let n = self.auto_seq.get();
618            self.auto_seq.set(n + 1);
619            let candidate: SmolStr = format!("_soc{n}").into();
620            if !self.soc_names.borrow().contains_key(&candidate) {
621                break candidate;
622            }
623        }
624    }
625
626    /// Register an anonymous SOC constraint, deriving a unique name `_soc{n}`
627    /// from an internal counter. Backs the name-less form of the
628    /// `soc_constraint!` macro. Not part of the stable public API.
629    #[doc(hidden)]
630    pub fn __add_soc_constraint_auto<'a>(
631        &'a self,
632        terms: impl IntoIterator<Item = Expr<'a>>,
633        bound: Expr<'a>,
634    ) -> SocConstraintId {
635        self.add_soc_constraint(self.next_auto_soc_name(), terms, bound)
636    }
637
638    /// Macro-facing entry point backing the indexed-family form of the
639    /// `soc_constraint!` macro: one cone per key, named `{prefix}[{key}]`. The
640    /// closure returns the cone's `(terms, bound)` pair for each typed key.
641    /// Not part of the stable public API.
642    #[doc(hidden)]
643    pub fn __add_soc_constraints_over<'a, K, T, F>(
644        &'a self,
645        name_prefix: &str,
646        set: &Set<K>,
647        mut rule: F,
648    ) where
649        K: FromIndexKey,
650        T: IntoIterator<Item = Expr<'a>>,
651        F: FnMut(K) -> (T, Expr<'a>),
652    {
653        for key in set {
654            let typed = K::from_index_key(&key);
655            let (terms, bound) = rule(typed);
656            let name: SmolStr = format_index_name(name_prefix, &key).into();
657            self.add_soc_constraint(name, terms, bound);
658        }
659    }
660
661    pub fn soc_constraints(&self) -> Ref<'_, Vec<SocConstraint>> {
662        self.soc_constraints.borrow()
663    }
664
665    pub fn num_soc_constraints(&self) -> usize {
666        self.soc_constraints.borrow().len()
667    }
668
669    pub fn soc_constraint_id(&self, name: &str) -> Option<SocConstraintId> {
670        self.soc_names.borrow().get(name).copied()
671    }
672
673    /// Whether the model carries any explicit second-order cone constraints.
674    pub fn has_cones(&self) -> bool {
675        !self.soc_constraints.borrow().is_empty()
676    }
677
678    // Objective
679
680    /// Macro-facing entry point backing `objective!(m, Min, ..)`. Not part of the
681    /// stable public API.
682    #[doc(hidden)]
683    pub fn __minimize(&self, expr: Expr<'_>) {
684        self.set_objective(expr, ObjectiveSense::Minimize);
685    }
686
687    /// Macro-facing entry point backing `objective!(m, Max, ..)`. Not part of the
688    /// stable public API.
689    #[doc(hidden)]
690    pub fn __maximize(&self, expr: Expr<'_>) {
691        self.set_objective(expr, ObjectiveSense::Maximize);
692    }
693
694    /// Macro-facing entry point backing `objective!(m, Feasibility)`. Declares
695    /// the model a feasibility problem (no objective to optimize), clearing any
696    /// previously set objective. Not part of the stable public API.
697    #[doc(hidden)]
698    pub fn __feasibility(&self) {
699        *self.objective.borrow_mut() = None;
700        self.objective_declared.set(true);
701        self.cached_kind.set(None);
702    }
703
704    fn set_objective(&self, expr: Expr<'_>, sense: ObjectiveSense) {
705        *self.objective.borrow_mut() = Some(Objective { expr: expr.id, sense });
706        self.objective_declared.set(true);
707        self.cached_kind.set(None);
708    }
709
710    /// Whether feasibility was declared explicitly via `objective!(m, Feasibility)`,
711    /// as opposed to a model that simply has no objective set.
712    pub fn is_feasibility(&self) -> bool {
713        self.objective_declared.get() && self.objective.borrow().is_none()
714    }
715
716    /// Ensure the model has a solve direction declared: either an objective
717    /// (`Min`/`Max`) or an explicit feasibility problem.
718    ///
719    /// # Errors
720    ///
721    /// Returns [`Error::NoObjective`] if neither an objective nor
722    /// `objective!(m, Feasibility)` was declared.
723    pub fn ensure_objective_declared(&self) -> Result<()> {
724        if self.objective_declared.get() { Ok(()) } else { Err(Error::NoObjective) }
725    }
726
727    pub fn objective(&self) -> Ref<'_, Option<Objective>> {
728        self.objective.borrow()
729    }
730
731    /// Try to get a cloned copy of the objective.
732    ///
733    /// # Errors
734    ///
735    /// Returns [`Error::NoObjective`] if no objective is set on this model.
736    pub fn try_objective(&self) -> Result<Objective> {
737        self.objective.borrow().clone().ok_or(Error::NoObjective)
738    }
739
740    // Classification
741
742    /// Infer the [`ModelKind`] from current variables and expressions.
743    /// Result is cached and invalidated whenever variables, constraints, or the
744    /// objective change.
745    ///
746    /// The decision ladder, top-down (any integer variable picks the `MI*`
747    /// column):
748    ///
749    /// 1. any nonlinear expression (objective or constraint) -> `NLP`
750    /// 2. any quadratic constraint not recognized as SOC (see
751    ///    [`crate::detect_soc`]) -> `QCP`
752    /// 3. cones present (explicit or detected) -> `SOCP`
753    /// 4. quadratic objective -> `QP`
754    /// 5. otherwise -> `LP`
755    pub fn kind(&self) -> ModelKind {
756        if let Some(k) = self.cached_kind.get() {
757            return k;
758        }
759        let arena = self.arena.borrow();
760        let vars = self.variables.borrow();
761        let has_int = vars.iter().any(|v| v.domain.is_integer());
762        let obj_class = self
763            .objective
764            .borrow()
765            .as_ref()
766            .map_or(ExprClass::Linear, |o| classify(&arena, o.expr));
767
768        let mut any_nonlinear = obj_class == ExprClass::Nonlinear;
769        // A quadratic constraint that is not SOC-shaped.
770        let mut plain_quad_con = false;
771        let mut detected_soc = false;
772        if !any_nonlinear {
773            for c in self.constraints.borrow().iter() {
774                match classify(&arena, c.lhs) {
775                    ExprClass::Linear => {}
776                    ExprClass::Quadratic => {
777                        if detect_soc(&arena, &vars, c).is_some() {
778                            detected_soc = true;
779                        } else {
780                            plain_quad_con = true;
781                        }
782                    }
783                    ExprClass::Nonlinear => {
784                        any_nonlinear = true;
785                        break;
786                    }
787                }
788            }
789        }
790        let has_soc = detected_soc || !self.soc_constraints.borrow().is_empty();
791
792        let pick = |cont, int| if has_int { int } else { cont };
793        let k = if any_nonlinear {
794            pick(ModelKind::NLP, ModelKind::MINLP)
795        } else if plain_quad_con {
796            pick(ModelKind::QCP, ModelKind::MIQCP)
797        } else if has_soc {
798            pick(ModelKind::SOCP, ModelKind::MISOCP)
799        } else if obj_class == ExprClass::Quadratic {
800            pick(ModelKind::QP, ModelKind::MIQP)
801        } else {
802            pick(ModelKind::LP, ModelKind::MILP)
803        };
804        self.cached_kind.set(Some(k));
805        k
806    }
807}
808
809// IndexedVarBuilder
810
811/// Builder for a collection of scalar variables indexed by a [`Set`].
812///
813/// For example, `flow[i]` for `i in 0..3` registers `flow[0]`, `flow[1]`, and
814/// `flow[2]` as separate scalar variables in the model. Call `.build()` to get
815/// an [`IndexedVar`] that maps each key to its [`Expr`] handle. Bounds and
816/// domain set here apply uniformly to every scalar in the collection.
817type BoundFn<'a> = Box<dyn Fn(&IndexKey) -> f64 + 'a>;
818
819#[must_use = "IndexedVarBuilder does nothing until you call .build()"]
820pub struct IndexedVarBuilder<'a, K = IndexKey> {
821    model: &'a Model,
822    base_name: String,
823    keys: Vec<IndexKey>,
824    axes: Option<Box<[Axis]>>,
825    lb: f64,
826    ub: f64,
827    lb_by: Option<BoundFn<'a>>,
828    ub_by: Option<BoundFn<'a>>,
829    domain: Domain,
830    _k: PhantomData<fn() -> K>,
831}
832
833impl<'a, K> std::fmt::Debug for IndexedVarBuilder<'a, K> {
834    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
835        f.debug_struct("IndexedVarBuilder")
836            .field("base_name", &self.base_name)
837            .field("keys", &self.keys.len())
838            .field("lb", &self.lb)
839            .field("ub", &self.ub)
840            .field("per_key_lb", &self.lb_by.is_some())
841            .field("per_key_ub", &self.ub_by.is_some())
842            .field("domain", &self.domain)
843            .finish()
844    }
845}
846
847impl<'a, K> IndexedVarBuilder<'a, K> {
848    pub fn lb(mut self, v: f64) -> Self {
849        self.lb = v;
850        self
851    }
852    pub fn ub(mut self, v: f64) -> Self {
853        self.ub = v;
854        self
855    }
856    pub fn bounds(mut self, lb: f64, ub: f64) -> Self {
857        self.lb = lb;
858        self.ub = ub;
859        self
860    }
861    /// Per-key lower bound. Overrides [`Self::lb`] when both are set.
862    ///
863    /// The closure receives a typed index value via [`FromIndexKey`].
864    /// Annotate the argument to select the projection:
865    /// ```ignore
866    /// .lb_by(|(p, q): (String, String)| floor_for(&p, &q))
867    /// .lb_by(|i: usize| lower_bounds[i])
868    /// ```
869    pub fn lb_by<F>(mut self, f: F) -> Self
870    where
871        K: FromIndexKey,
872        F: Fn(K) -> f64 + 'a,
873    {
874        self.lb_by = Some(Box::new(move |k: &IndexKey| f(K::from_index_key(k))));
875        self
876    }
877    /// Per-key upper bound. Overrides [`Self::ub`] when both are set.
878    ///
879    /// The closure receives a typed index value via [`FromIndexKey`]; annotate
880    /// the argument to select the projection:
881    /// ```ignore
882    /// .ub_by(|(p, q): (String, String)| capacity_for(&p, &q))
883    /// .ub_by(|i: usize| upper_bounds[i])
884    /// ```
885    pub fn ub_by<F>(mut self, f: F) -> Self
886    where
887        K: FromIndexKey,
888        F: Fn(K) -> f64 + 'a,
889    {
890        self.ub_by = Some(Box::new(move |k: &IndexKey| f(K::from_index_key(k))));
891        self
892    }
893    pub fn domain(mut self, d: Domain) -> Self {
894        self.domain = d;
895        self
896    }
897    pub fn integer(mut self) -> Self {
898        self.domain = Domain::Integer;
899        self
900    }
901    pub fn binary(mut self) -> Self {
902        self.domain = Domain::Binary;
903        self.lb = 0.0;
904        self.ub = 1.0;
905        self
906    }
907
908    /// Register one scalar variable per key and return the [`IndexedVar`] handle.
909    ///
910    /// # Panics
911    /// Panics if a scalar variable name collides with one already registered.
912    pub fn build(self) -> IndexedVar<'a, K> {
913        let Self { model, base_name, keys, axes, lb, ub, lb_by, ub_by, domain, _k } = self;
914
915        let make = |key: &IndexKey| -> Expr<'a> {
916            let scalar_name: SmolStr = format_index_name(&base_name, key).into();
917            let lo = lb_by.as_ref().map_or(lb, |f| f(key));
918            let hi = ub_by.as_ref().map_or(ub, |f| f(key));
919            model.__var(scalar_name).lb(lo).ub(hi).domain(domain).build()
920        };
921
922        let storage = build_storage(keys, axes, make);
923        IndexedFamily { storage, _marker: PhantomData }
924    }
925}
926
927fn format_index_name(base: &str, key: &IndexKey) -> String {
928    let mut out = String::with_capacity(base.len() + 4);
929    out.push_str(base);
930    out.push('[');
931    write_key_parts(&mut out, key);
932    out.push(']');
933    out
934}
935
936fn write_key_parts(out: &mut String, key: &IndexKey) {
937    use std::fmt::Write;
938    match key {
939        IndexKey::Int(i) => write!(out, "{i}").unwrap(),
940        IndexKey::Str(s) => out.push_str(s),
941        IndexKey::Tuple(parts) => {
942            for (i, p) in parts.iter().enumerate() {
943                if i > 0 {
944                    out.push(',');
945                }
946                write_key_parts(out, p);
947            }
948        }
949    }
950}
951
952/// Public render of an `IndexKey`'s textual form, used when deriving
953/// auto-generated names for indexed-family constraints.
954pub fn display_index_key(key: &IndexKey) -> String {
955    let mut out = String::new();
956    write_key_parts(&mut out, key);
957    out
958}
959
960#[cfg(test)]
961mod tests {
962    use oximo_expr::extract_linear;
963
964    use super::*;
965    use crate::Set;
966    use crate::constraint::Relate;
967
968    #[test]
969    fn param_times_var_keeps_model_linear() {
970        let m = Model::new("p");
971        let param = m.__param("param", 4.0);
972        let x = m.__var("x").lb(0.0).build();
973        m.__minimize(param * x);
974        assert_eq!(m.kind(), ModelKind::LP);
975    }
976
977    #[test]
978    fn param_coeff_resolves_and_rebinds() {
979        let m = Model::new("p");
980        let param = m.__param("param", 4.0);
981        let x = m.__var("x").lb(0.0).build();
982        let obj = param * x;
983
984        let coeff = |m: &Model| {
985            let arena = m.arena();
986            extract_linear(&arena, obj.id).expect("linear").coeffs[0].1
987        };
988        assert!((coeff(&m) - 4.0).abs() < f64::EPSILON);
989
990        m.set_param(param, 9.0);
991        assert!((coeff(&m) - 9.0).abs() < f64::EPSILON);
992        assert_eq!(m.parameter_id("param"), Some(param.param_id().unwrap()));
993    }
994
995    #[test]
996    fn param_value_reads_live_arena_value() {
997        let m = Model::new("p");
998        let param = m.__param("param", 4.0);
999        let id = param.param_id().unwrap();
1000        assert!((m.param_value(id) - 4.0).abs() < f64::EPSILON);
1001        assert!((m.param_value_of(param).unwrap() - 4.0).abs() < f64::EPSILON);
1002
1003        m.set_param(param, 7.5);
1004        assert!((m.param_value(id) - 7.5).abs() < f64::EPSILON);
1005
1006        let x = m.__var("x").build();
1007        assert!(m.param_value_of(x).is_none());
1008    }
1009
1010    #[test]
1011    fn set_param_invalidates_kind_cache() {
1012        let m = Model::new("p");
1013        let p = m.__param("p", 1.0);
1014        let x = m.__var("x").lb(0.0).build();
1015        m.__add_constraint("c", (p * x).le(10.0));
1016        assert_eq!(m.kind(), ModelKind::LP);
1017        m.set_param(p, 2.0);
1018        assert_eq!(m.kind(), ModelKind::LP);
1019    }
1020
1021    #[test]
1022    #[should_panic(expected = "parameter name \"dup\" already registered")]
1023    fn duplicate_param_name_panics() {
1024        let m = Model::new("p");
1025        let _a = m.__param("dup", 1.0);
1026        let _b = m.__param("dup", 2.0);
1027    }
1028
1029    #[test]
1030    fn indexed_param_dense_value_and_per_key_rebind() {
1031        let m = Model::new("ip");
1032        let items = Set::range(0..3);
1033        let data = [10.0, 20.0, 30.0];
1034        let cost = m.__indexed_param("cost", &items, |i: usize| data[i]);
1035
1036        assert!(cost.is_dense());
1037        assert_eq!(cost.len(), 3);
1038        assert_eq!(m.num_parameters(), 3);
1039        assert!(m.parameter_id("cost[0]").is_some());
1040        assert!(m.parameter_id("cost[2]").is_some());
1041        assert!((m.param_value_idx(&cost, 1usize).unwrap() - 20.0).abs() < f64::EPSILON);
1042
1043        let x = m.__var("x").lb(0.0).build();
1044        let obj = cost.at([1]) * x;
1045        let coeff = |m: &Model| {
1046            let arena = m.arena();
1047            extract_linear(&arena, obj.id).expect("linear").coeffs[0].1
1048        };
1049        assert!((coeff(&m) - 20.0).abs() < f64::EPSILON);
1050
1051        m.set_param_idx(&cost, 1usize, 99.0);
1052        assert!((coeff(&m) - 99.0).abs() < f64::EPSILON);
1053        assert!((m.param_value_idx(&cost, 1usize).unwrap() - 99.0).abs() < f64::EPSILON);
1054        assert!((m.param_value_idx(&cost, 0usize).unwrap() - 10.0).abs() < f64::EPSILON);
1055    }
1056
1057    #[test]
1058    #[should_panic(expected = "different model")]
1059    fn set_param_idx_rejects_foreign_family() {
1060        let a = Model::new("a");
1061        let b = Model::new("b");
1062        let items = Set::range(0..2);
1063        let pa = a.__indexed_param("p", &items, |_i: usize| 1.0);
1064        b.set_param_idx(&pa, 0usize, 5.0);
1065    }
1066
1067    #[test]
1068    fn indexed_param_sparse_string_keyed() {
1069        let m = Model::new("ips");
1070        let plants = Set::strings(["a", "b"]);
1071        let price =
1072            m.__indexed_param("price", &plants, |p: String| if p == "a" { 1.5 } else { 2.5 });
1073        assert!(!price.is_dense());
1074        assert_eq!(price.len(), 2);
1075        assert!((m.param_value_idx(&price, "a").unwrap() - 1.5).abs() < f64::EPSILON);
1076        assert!((m.param_value_idx(&price, "b").unwrap() - 2.5).abs() < f64::EPSILON);
1077        assert!(m.param_value_idx(&price, "z").is_none());
1078    }
1079}