Skip to main content

oximo_expr/
arena.rs

1use smallvec::SmallVec;
2
3#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
4pub struct ExprId(pub u32);
5
6impl ExprId {
7    #[inline]
8    pub fn index(self) -> usize {
9        self.0 as usize
10    }
11}
12
13#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
14pub struct VarId(pub u32);
15
16impl VarId {
17    #[inline]
18    pub fn index(self) -> usize {
19        self.0 as usize
20    }
21}
22
23#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
24pub struct ParamId(pub u32);
25
26impl ParamId {
27    #[inline]
28    pub fn index(self) -> usize {
29        self.0 as usize
30    }
31}
32
33pub type Children = SmallVec<[ExprId; 4]>;
34
35/// Here we use a linear fast-path: `sum(coeff * var) + constant`.
36/// Built by the operator overloads when all children are linear,
37/// so LP/MILP construction never walks an `Add(Mul(Const, Var), ...)` tree.
38
39#[derive(Clone, Debug)]
40pub enum ExprNode {
41    Const(f64),
42    Var(VarId),
43    Param(ParamId),
44    Add(Children),
45    Mul(Children),
46    Neg(ExprId),
47    Pow(ExprId, ExprId),
48    Sin(ExprId),
49    Cos(ExprId),
50    Exp(ExprId),
51    Log(ExprId),
52    Linear { coeffs: Vec<(VarId, f64)>, constant: f64 },
53}
54
55#[derive(Clone, Debug, Default)]
56pub struct ExprArena {
57    nodes: Vec<ExprNode>,
58}
59
60impl ExprArena {
61    pub fn new() -> Self {
62        Self::default()
63    }
64
65    pub fn with_capacity(cap: usize) -> Self {
66        Self { nodes: Vec::with_capacity(cap) }
67    }
68
69    #[inline]
70    pub fn len(&self) -> usize {
71        self.nodes.len()
72    }
73
74    #[inline]
75    pub fn is_empty(&self) -> bool {
76        self.nodes.is_empty()
77    }
78
79    /// # Panics
80    ///
81    /// Panics if the number of expressions exceeds `u32::MAX` (expression arena overflow).
82    pub fn push(&mut self, node: ExprNode) -> ExprId {
83        let id = ExprId(u32::try_from(self.nodes.len()).expect("expression arena overflow"));
84        self.nodes.push(node);
85        id
86    }
87
88    #[inline]
89    pub fn get(&self, id: ExprId) -> &ExprNode {
90        &self.nodes[id.index()]
91    }
92
93    #[inline]
94    pub fn get_mut(&mut self, id: ExprId) -> &mut ExprNode {
95        &mut self.nodes[id.index()]
96    }
97
98    pub fn nodes(&self) -> &[ExprNode] {
99        &self.nodes
100    }
101
102    pub fn constant(&mut self, v: f64) -> ExprId {
103        self.push(ExprNode::Const(v))
104    }
105
106    pub fn var(&mut self, v: VarId) -> ExprId {
107        self.push(ExprNode::Var(v))
108    }
109
110    pub fn param(&mut self, p: ParamId) -> ExprId {
111        self.push(ExprNode::Param(p))
112    }
113
114    pub fn linear(&mut self, coeffs: Vec<(VarId, f64)>, constant: f64) -> ExprId {
115        self.push(ExprNode::Linear { coeffs, constant })
116    }
117}