Skip to main content

oximo_core/
model.rs

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