Skip to main content

sirius/solver/poly/
mono.rs

1//! Monomial without coefficient, the variable part of a [`Poly`](super::Poly) term.
2
3use std::cmp::Ordering;
4
5/// A variable, represented by `u32`.
6pub type Var = u32;
7/// An exponent, represented by `u32`.
8pub type Pow = u32;
9
10/// A power product $\prod_i x_i^{e_i}$ (for example, $x^2y^5$). The variable part of a [`Poly`](super::Poly) term.
11/// [`Mono::cmp`] implements [graded-lex ordering](https://en.wikipedia.org/wiki/Monomial_order#Graded_lexicographic_order).
12/// ```
13/// # use sirius::solver::poly::mono::mono;
14/// let a = mono!(x^2);
15/// let b = mono!(x*y);
16/// let c = mono!(x*z^2);
17/// assert_eq!(a.mul(&b).mul(&c), mono!(x^4*y*z^2));
18/// ```
19#[derive(Clone, PartialEq, Eq, Hash, Debug)]
20pub struct Mono {
21    exps: Vec<(Var, Pow)>,
22}
23
24impl Mono {
25    pub fn new(exps: impl IntoIterator<Item = (Var, Pow)>) -> Self {
26        // Normalize an exp list: sort, combine like vars, and remove zero powers.
27        let mut pairs: Vec<_> = exps.into_iter().filter(|&(_, e)| e > 0).collect();
28        pairs.sort_unstable_by_key(|&(v, _)| v);
29        let mut exps = Vec::<(Var, Pow)>::with_capacity(pairs.len());
30        for (v, e) in pairs {
31            match exps.last_mut() {
32                Some((lv, le)) if *lv == v => *le += e,
33                _ => exps.push((v, e)),
34            }
35        }
36        Self { exps }
37    }
38
39    /// Unit monomial, 1.
40    pub fn unit() -> Self {
41        Self::new(vec![])
42    }
43
44    pub fn exps(&self) -> &[(Var, Pow)] {
45        &self.exps
46    }
47
48    pub fn is_unit(&self) -> bool {
49        self.exps.is_empty()
50    }
51
52    pub fn total_degree(&self) -> Pow {
53        self.exps.iter().map(|&(_, e)| e).sum()
54    }
55
56    pub fn always_nonneg(&self, nonneg: &dyn Fn(Var) -> bool) -> bool {
57        self.exps()
58            .iter()
59            .all(|&(v, pow)| nonneg(v) || pow % 2 == 0)
60    }
61
62    pub fn mul(&self, other: &Self) -> Self {
63        // Merge sorted `exps` and combine like vars to maintain canonical form.
64        let mut exps = Vec::with_capacity(self.exps.len() + other.exps.len());
65        let (mut a, mut b) = (self.exps.iter().peekable(), other.exps.iter().peekable());
66        loop {
67            match (a.peek(), b.peek()) {
68                (Some(&&(va, ea)), Some(&&(vb, eb))) => match va.cmp(&vb) {
69                    Ordering::Less => {
70                        exps.push((va, ea));
71                        a.next();
72                    }
73                    Ordering::Greater => {
74                        exps.push((vb, eb));
75                        b.next();
76                    }
77                    Ordering::Equal => {
78                        exps.push((va, ea + eb));
79                        a.next();
80                        b.next();
81                    }
82                },
83                (Some(_), None) => {
84                    exps.extend(a);
85                    break;
86                }
87                (None, Some(_)) => {
88                    exps.extend(b);
89                    break;
90                }
91                (None, None) => break,
92            }
93        }
94        Self { exps }.debug_checked()
95    }
96
97    pub fn div(&self, other: &Self) -> Option<Self> {
98        let mut lhs_var_iter = self.exps.iter().peekable();
99        let mut rhs_var_iter = other.exps.iter().peekable();
100        let mut vars = vec![];
101        while let Some((rhs_var, rhs_pow)) = rhs_var_iter.peek() {
102            if let Some((lhs_var, lhs_pow)) = lhs_var_iter.peek() {
103                match lhs_var.cmp(rhs_var) {
104                    Ordering::Equal => match lhs_pow.cmp(rhs_pow) {
105                        Ordering::Greater => {
106                            vars.push((*lhs_var, lhs_pow - rhs_pow));
107                            lhs_var_iter.next();
108                            rhs_var_iter.next();
109                            continue;
110                        }
111                        Ordering::Equal => {
112                            lhs_var_iter.next();
113                            rhs_var_iter.next();
114                            continue;
115                        }
116                        Ordering::Less => return None,
117                    },
118                    Ordering::Less => {
119                        vars.push((*lhs_var, *lhs_pow));
120                        lhs_var_iter.next();
121                        continue;
122                    }
123                    Ordering::Greater => {
124                        return None;
125                    }
126                }
127            }
128
129            return None;
130        }
131
132        for (lhs_var, lhs_pow) in lhs_var_iter {
133            vars.push((*lhs_var, *lhs_pow));
134        }
135
136        Some(Mono::new(vars))
137    }
138
139    pub fn degree_in<T: Into<Var>>(&self, v: T) -> Pow {
140        let v = v.into();
141        self.exps
142            .iter()
143            .find(|(va, _)| *va == v)
144            .map(|exp| exp.1)
145            .unwrap_or(0)
146    }
147
148    pub fn lcm(&self, rhs: &Self) -> Self {
149        let mut vars = vec![];
150
151        let mut lhs_vars = self.exps.iter().peekable();
152        let mut rhs_vars = rhs.exps.iter().peekable();
153
154        loop {
155            match (lhs_vars.peek(), rhs_vars.peek()) {
156                (Some(lhs_v), Some(rhs_v)) => match lhs_v.0.cmp(&rhs_v.0) {
157                    Ordering::Equal => {
158                        vars.push((lhs_v.0, lhs_v.1.max(rhs_v.1)));
159                        lhs_vars.next();
160                        rhs_vars.next();
161                    }
162                    Ordering::Greater => {
163                        vars.push(**rhs_v);
164                        rhs_vars.next();
165                    }
166                    Ordering::Less => {
167                        vars.push(**lhs_v);
168                        lhs_vars.next();
169                    }
170                },
171                (Some(lhs_v), None) => {
172                    vars.push(**lhs_v);
173                    lhs_vars.next();
174                }
175                (None, Some(rhs_v)) => {
176                    vars.push(**rhs_v);
177                    rhs_vars.next();
178                }
179                (None, None) => break,
180            }
181        }
182
183        Mono { exps: vars }
184    }
185
186    pub fn assert_canonical(&self) {
187        for w in self.exps.windows(2) {
188            assert!(
189                w[0].0 < w[1].0,
190                "monomial vars not strictly ascending: {:?}",
191                self.exps
192            );
193        }
194        for &(_, e) in &self.exps {
195            assert!(e >= 1, "monomial has zero exponent: {:?}", self.exps);
196        }
197    }
198
199    fn debug_checked(self) -> Self {
200        #[cfg(debug_assertions)]
201        self.assert_canonical();
202        self
203    }
204}
205
206impl Ord for Mono {
207    /// [Graded-lex](https://en.wikipedia.org/wiki/Monomial_order#Graded_lexicographic_order): total degree first;
208    /// ties broken lexicographically on exponent vectors with lower [`Var`] index more significant
209    /// (larger exponent on the first differing variable wins).
210    fn cmp(&self, other: &Self) -> Ordering {
211        self.total_degree()
212            .cmp(&other.total_degree())
213            .then_with(|| {
214                let (mut a, mut b) = (self.exps.iter(), other.exps.iter());
215                loop {
216                    match (a.next(), b.next()) {
217                        (Some(&(va, ea)), Some(&(vb, eb))) => match va.cmp(&vb) {
218                            Ordering::Less => return Ordering::Greater,
219                            Ordering::Greater => return Ordering::Less,
220                            Ordering::Equal => {
221                                if ea != eb {
222                                    return ea.cmp(&eb);
223                                }
224                            }
225                        },
226                        (Some(_), None) => return Ordering::Greater,
227                        (None, Some(_)) => return Ordering::Less,
228                        (None, None) => return Ordering::Equal,
229                    }
230                }
231            })
232    }
233}
234
235impl PartialOrd for Mono {
236    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
237        Some(self.cmp(other))
238    }
239}
240
241#[doc(hidden)]
242/// Helper function for the [`mono`](crate::mono) macro.
243pub const fn __read_var_name(name: &str) -> Var {
244    let bytes = name.as_bytes();
245    assert!(bytes.len() == 1, "variable name must have length 1");
246    bytes[0] as Var
247}
248
249/// Create a [`Mono`]. Accepts 1-character ASCII variable names; for example, `mono!(x*y^5)`.
250#[doc(hidden)]
251#[macro_export]
252macro_rules! __mono {
253    () => { $crate::solver::poly::mono::Mono::unit() };
254
255    ($var:ident ^ $pow:literal $(* $($rest:tt)*)?) => {{
256        use $crate::solver::poly::mono::{Mono, Var, __read_var_name};
257        const VAR_NAME: Var = __read_var_name(stringify!($var));
258        Mono::new(vec![(VAR_NAME, $pow)]).mul(&$crate::__mono!($($($rest)*)?))
259    }};
260    ($var:ident $(* $($rest:tt)*)?) => {{
261        use $crate::solver::poly::mono::{Mono, Var, __read_var_name};
262        const VAR_NAME: Var = __read_var_name(stringify!($var));
263        Mono::new(vec![(VAR_NAME, 1)]).mul(&$crate::__mono!($($($rest)*)?))
264    }};
265}
266
267#[doc(inline)]
268pub use crate::__mono as mono;
269
270#[cfg(test)]
271mod tests {
272    use super::mono;
273
274    #[test]
275    fn graded_lex() {
276        let descending = [
277            mono!(x ^ 2),
278            mono!(x * y),
279            mono!(x * z),
280            mono!(y ^ 2),
281            mono!(y * z),
282            mono!(z ^ 2),
283            mono!(x),
284            mono!(y),
285            mono!(z),
286            mono!(),
287        ];
288        for w in descending.windows(2) {
289            assert!(w[0] > w[1], "expected {:?} > {:?}", w[0], w[1]);
290        }
291    }
292
293    #[test]
294    fn multiplication() {
295        assert_eq!(mono!(x).mul(&mono!(x)), mono!(x ^ 2));
296        let a = mono!(x ^ 2 * y).mul(&mono!(x * z ^ 2));
297        a.assert_canonical();
298        assert_eq!(a, mono!(x ^ 3 * y * z ^ 2));
299    }
300}