Skip to main content

oximo_expr/
handle.rs

1use std::cell::RefCell;
2
3use crate::arena::{ExprArena, ExprId, ExprNode, ParamId, VarId};
4
5/// Lightweight handle to a node in an [`ExprArena`].
6///
7/// Carries a borrow of the arena (wrapped in `RefCell` so operator overloads
8/// can push new nodes during arithmetic). `Expr` is `Copy`, so users freely
9/// reuse a variable handle in many constraints.
10#[derive(Copy, Clone)]
11pub struct Expr<'a> {
12    pub id: ExprId,
13    pub arena: &'a RefCell<ExprArena>,
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).finish()
19    }
20}
21
22impl<'a> Expr<'a> {
23    #[inline]
24    pub fn new(id: ExprId, arena: &'a RefCell<ExprArena>) -> Self {
25        Self { id, arena }
26    }
27
28    pub fn constant(arena: &'a RefCell<ExprArena>, v: f64) -> Self {
29        let id = arena.borrow_mut().constant(v);
30        Self::new(id, arena)
31    }
32
33    pub fn from_var(arena: &'a RefCell<ExprArena>, v: VarId) -> Self {
34        let id = arena.borrow_mut().var(v);
35        Self::new(id, arena)
36    }
37
38    /// If this handle is a bare variable, return its [`VarId`].
39    /// `None` for compound expressions (sums, products, constants, ...).
40    pub fn var_id(self) -> Option<VarId> {
41        match self.arena.borrow().get(self.id) {
42            ExprNode::Var(id) => Some(*id),
43            _ => None,
44        }
45    }
46
47    /// If this handle is a bare parameter, return its [`ParamId`].
48    /// `None` for compound expressions.
49    pub fn param_id(self) -> Option<ParamId> {
50        match self.arena.borrow().get(self.id) {
51            ExprNode::Param(id) => Some(*id),
52            _ => None,
53        }
54    }
55
56    pub fn pow(self, exponent: Self) -> Self {
57        let id = self.arena.borrow_mut().push(ExprNode::Pow(self.id, exponent.id));
58        Self::new(id, self.arena)
59    }
60
61    pub fn powi(self, n: i32) -> Self {
62        let id = {
63            let mut a = self.arena.borrow_mut();
64            let exp_id = a.constant(f64::from(n));
65            a.push(ExprNode::Pow(self.id, exp_id))
66        };
67        Self::new(id, self.arena)
68    }
69
70    pub fn powf(self, n: f64) -> Self {
71        let id = {
72            let mut a = self.arena.borrow_mut();
73            let exp_id = a.constant(n);
74            a.push(ExprNode::Pow(self.id, exp_id))
75        };
76        Self::new(id, self.arena)
77    }
78
79    pub fn sin(self) -> Self {
80        let id = self.arena.borrow_mut().push(ExprNode::Sin(self.id));
81        Self::new(id, self.arena)
82    }
83
84    pub fn cos(self) -> Self {
85        let id = self.arena.borrow_mut().push(ExprNode::Cos(self.id));
86        Self::new(id, self.arena)
87    }
88
89    pub fn exp(self) -> Self {
90        let id = self.arena.borrow_mut().push(ExprNode::Exp(self.id));
91        Self::new(id, self.arena)
92    }
93
94    pub fn log(self) -> Self {
95        let id = self.arena.borrow_mut().push(ExprNode::Log(self.id));
96        Self::new(id, self.arena)
97    }
98
99    pub fn abs(self) -> Self {
100        let id = self.arena.borrow_mut().push(ExprNode::Abs(self.id));
101        Self::new(id, self.arena)
102    }
103}