Skip to main content

oximo_expr/
handle.rs

1use std::cell::RefCell;
2
3use crate::arena::{ExprArena, ExprId, ExprNode, 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    pub fn pow(self, exponent: Self) -> Self {
39        let id = self.arena.borrow_mut().push(ExprNode::Pow(self.id, exponent.id));
40        Self::new(id, self.arena)
41    }
42
43    pub fn powi(self, n: i32) -> Self {
44        let id = {
45            let mut a = self.arena.borrow_mut();
46            let exp_id = a.constant(f64::from(n));
47            a.push(ExprNode::Pow(self.id, exp_id))
48        };
49        Self::new(id, self.arena)
50    }
51
52    pub fn powf(self, n: f64) -> Self {
53        let id = {
54            let mut a = self.arena.borrow_mut();
55            let exp_id = a.constant(n);
56            a.push(ExprNode::Pow(self.id, exp_id))
57        };
58        Self::new(id, self.arena)
59    }
60
61    pub fn sin(self) -> Self {
62        let id = self.arena.borrow_mut().push(ExprNode::Sin(self.id));
63        Self::new(id, self.arena)
64    }
65
66    pub fn cos(self) -> Self {
67        let id = self.arena.borrow_mut().push(ExprNode::Cos(self.id));
68        Self::new(id, self.arena)
69    }
70
71    pub fn exp(self) -> Self {
72        let id = self.arena.borrow_mut().push(ExprNode::Exp(self.id));
73        Self::new(id, self.arena)
74    }
75
76    pub fn log(self) -> Self {
77        let id = self.arena.borrow_mut().push(ExprNode::Log(self.id));
78        Self::new(id, self.arena)
79    }
80}