oximo_expr/
handle.rs

1use crate::arena::{Children, ExprArenaCell, ExprId, ExprNode, ModelId, ParamId, UnaryOp, VarId};
2use crate::classify::{ExprClass, classify_access};
3
4/// Lightweight handle to a node in an [`ExprArenaCell`].
5///
6/// Carries a borrow of the arena cell so operator overloads can push new nodes
7/// during arithmetic. `Expr` is `Copy`, so users freely reuse a variable
8/// handle in many constraints.
9#[derive(Copy, Clone)]
10pub struct Expr<'a> {
11    pub id: ExprId,
12    model_id: ModelId,
13    pub arena: &'a ExprArenaCell,
14}
15
16impl std::fmt::Debug for Expr<'_> {
17    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
18        f.debug_struct("Expr").field("id", &self.id).field("model_id", &self.model_id).finish()
19    }
20}
21
22impl<'a> Expr<'a> {
23    #[inline]
24    pub fn new(id: ExprId, arena: &'a ExprArenaCell) -> Self {
25        Self { id, model_id: arena.model_id(), arena }
26    }
27
28    /// Identity of the model/expression arena that created this handle.
29    #[inline]
30    #[must_use]
31    pub const fn model_id(self) -> ModelId {
32        self.model_id
33    }
34
35    pub fn constant(arena: &'a ExprArenaCell, v: f64) -> Self {
36        let id = arena.with_mut(|arena| arena.constant(v));
37        Self::new(id, arena)
38    }
39
40    pub fn from_var(arena: &'a ExprArenaCell, v: VarId) -> Self {
41        let id = arena.with_mut(|arena| arena.var(v));
42        Self::new(id, arena)
43    }
44
45    #[inline]
46    pub(crate) fn assert_same_arena(self, other: Self) {
47        assert!(std::ptr::eq(self.arena, other.arena), "expressions belong to different arenas");
48    }
49
50    /// If this handle is a bare variable, return its [`VarId`].
51    /// `None` for compound expressions (sums, products, constants, ...).
52    pub fn var_id(self) -> Option<VarId> {
53        self.arena.with_ref(|arena| match arena.get(self.id) {
54            ExprNode::Var(id) => Some(*id),
55            _ => None,
56        })
57    }
58
59    /// If this handle is a bare parameter, return its [`ParamId`].
60    /// `None` for compound expressions.
61    pub fn param_id(self) -> Option<ParamId> {
62        self.arena.with_ref(|arena| match arena.get(self.id) {
63            ExprNode::Param(id) => Some(*id),
64            _ => None,
65        })
66    }
67
68    /// Re-bind the parameter this handle references to `value`. Takes effect on
69    /// the next extraction/evaluation, which read the value straight from the
70    /// arena.
71    ///
72    /// # Panics
73    /// Panics if this handle is not a bare parameter (see [`Self::param_id`]).
74    pub fn set_param_value(self, value: f64) {
75        let id = self.param_id().expect("set_param_value expects a bare parameter handle");
76        self.arena.borrow_mut().set_param_value(id, value);
77    }
78
79    pub fn pow(self, exponent: Self) -> Self {
80        self.assert_same_arena(exponent);
81        let id = self.arena.with_mut(|arena| arena.push(ExprNode::Pow(self.id, exponent.id)));
82        Self::new(id, self.arena)
83    }
84
85    pub fn powi(self, n: i32) -> Self {
86        let id = self.arena.with_mut(|arena| {
87            let exp_id = arena.constant(f64::from(n));
88            arena.push(ExprNode::Pow(self.id, exp_id))
89        });
90        Self::new(id, self.arena)
91    }
92
93    pub fn powf(self, n: f64) -> Self {
94        let id = self.arena.with_mut(|arena| {
95            let exp_id = arena.constant(n);
96            arena.push(ExprNode::Pow(self.id, exp_id))
97        });
98        Self::new(id, self.arena)
99    }
100
101    fn unary(self, op: UnaryOp) -> Self {
102        let id = self.arena.with_mut(|arena| arena.push(ExprNode::Unary(op, self.id)));
103        Self::new(id, self.arena)
104    }
105
106    /// Unary negation as an explicit expression node. The `-expr` operator
107    /// keeps its affine fast path. Use this method when the public
108    /// [`UnaryOp`] node is required.
109    #[expect(
110        clippy::should_implement_trait,
111        reason = "explicit node constructor complements Neg::neg"
112    )]
113    pub fn neg(self) -> Self {
114        self.unary(UnaryOp::Neg)
115    }
116
117    pub fn abs(self) -> Self {
118        self.unary(UnaryOp::Abs)
119    }
120    pub fn sqrt(self) -> Self {
121        self.unary(UnaryOp::Sqrt)
122    }
123    pub fn cbrt(self) -> Self {
124        self.unary(UnaryOp::Cbrt)
125    }
126    pub fn exp(self) -> Self {
127        self.unary(UnaryOp::Exp)
128    }
129    pub fn exp2(self) -> Self {
130        self.unary(UnaryOp::Exp2)
131    }
132    pub fn expm1(self) -> Self {
133        self.unary(UnaryOp::Expm1)
134    }
135    /// Alias for [`Expr::expm1`], matching Rust's `f64::exp_m1` spelling.
136    pub fn exp_m1(self) -> Self {
137        self.expm1()
138    }
139    pub fn log(self) -> Self {
140        self.unary(UnaryOp::Log)
141    }
142    /// Alias for [`Expr::log`], matching Rust's `f64::ln` spelling.
143    pub fn ln(self) -> Self {
144        self.log()
145    }
146    pub fn log2(self) -> Self {
147        self.unary(UnaryOp::Log2)
148    }
149    pub fn log10(self) -> Self {
150        self.unary(UnaryOp::Log10)
151    }
152    pub fn log1p(self) -> Self {
153        self.unary(UnaryOp::Log1p)
154    }
155    /// Alias for [`Expr::log1p`], matching Rust's `f64::ln_1p` spelling.
156    pub fn ln_1p(self) -> Self {
157        self.log1p()
158    }
159    pub fn sin(self) -> Self {
160        self.unary(UnaryOp::Sin)
161    }
162    pub fn cos(self) -> Self {
163        self.unary(UnaryOp::Cos)
164    }
165    pub fn tan(self) -> Self {
166        self.unary(UnaryOp::Tan)
167    }
168    pub fn asin(self) -> Self {
169        self.unary(UnaryOp::Asin)
170    }
171    pub fn acos(self) -> Self {
172        self.unary(UnaryOp::Acos)
173    }
174    pub fn atan(self) -> Self {
175        self.unary(UnaryOp::Atan)
176    }
177    pub fn sinh(self) -> Self {
178        self.unary(UnaryOp::Sinh)
179    }
180    pub fn cosh(self) -> Self {
181        self.unary(UnaryOp::Cosh)
182    }
183    pub fn tanh(self) -> Self {
184        self.unary(UnaryOp::Tanh)
185    }
186    pub fn asinh(self) -> Self {
187        self.unary(UnaryOp::Asinh)
188    }
189    pub fn acosh(self) -> Self {
190        self.unary(UnaryOp::Acosh)
191    }
192    pub fn atanh(self) -> Self {
193        self.unary(UnaryOp::Atanh)
194    }
195
196    /// Two-argument arctangent in Rust's `y.atan2(x)` argument order.
197    pub fn atan2(self, x: Self) -> Self {
198        self.assert_same_arena(x);
199        let id = self.arena.with_mut(|arena| arena.push(ExprNode::Atan2(self.id, x.id)));
200        Self::new(id, self.arena)
201    }
202
203    fn extrema(self, other: Self, is_min: bool) -> Self {
204        self.assert_same_arena(other);
205        let id = self.arena.with_mut(|arena| {
206            let mut children = Children::new();
207            let left = match arena.get(self.id) {
208                ExprNode::Min(existing) if is_min => Some(existing.as_slice()),
209                ExprNode::Max(existing) if !is_min => Some(existing.as_slice()),
210                _ => None,
211            };
212            if let Some(existing) = left {
213                children.extend_from_slice(existing);
214            } else {
215                children.push(self.id);
216            }
217            let right = match arena.get(other.id) {
218                ExprNode::Min(existing) if is_min => Some(existing.as_slice()),
219                ExprNode::Max(existing) if !is_min => Some(existing.as_slice()),
220                _ => None,
221            };
222            if let Some(existing) = right {
223                children.extend_from_slice(existing);
224            } else {
225                children.push(other.id);
226            }
227            arena.push(if is_min { ExprNode::Min(children) } else { ExprNode::Max(children) })
228        });
229        Self::new(id, self.arena)
230    }
231
232    /// Pairwise minimum, flattening nested minima into deterministic n-ary nodes.
233    pub fn min(self, other: Self) -> Self {
234        self.extrema(other, true)
235    }
236
237    /// Pairwise maximum, flattening nested maxima into deterministic n-ary nodes.
238    pub fn max(self, other: Self) -> Self {
239        self.extrema(other, false)
240    }
241
242    #[doc(hidden)]
243    pub fn __class(self) -> ExprClass {
244        self.arena.with_ref(|arena| classify_access(arena, self.id))
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::Expr;
251    use crate::arena::{ExprArena, ExprArenaCell};
252
253    #[test]
254    fn set_param_value_rebinds_through_handle() {
255        let arena = ExprArenaCell::new(ExprArena::new());
256        let pid = arena.borrow_mut().new_param(0.05);
257        let node = arena.borrow_mut().param(pid);
258        let p = Expr::new(node, &arena);
259
260        p.set_param_value(0.2);
261        assert!((arena.borrow().param_value(pid) - 0.2).abs() < f64::EPSILON);
262    }
263
264    #[test]
265    #[should_panic(expected = "bare parameter handle")]
266    fn set_param_value_panics_on_non_param() {
267        let arena = ExprArenaCell::new(ExprArena::new());
268        let c = Expr::constant(&arena, 1.0);
269        c.set_param_value(3.0);
270    }
271
272    #[test]
273    #[should_panic(expected = "different arenas")]
274    fn combining_different_arenas_is_rejected() {
275        let left_arena = ExprArenaCell::new(ExprArena::new());
276        let right_arena = ExprArenaCell::new(ExprArena::new());
277        let left = Expr::constant(&left_arena, 1.0);
278        let right = Expr::constant(&right_arena, 2.0);
279        let _ = left + right;
280    }
281}
oximo_expr/handle.rs

oximo_expr/
handle.rs
