polysub/

parser.rs

// Copyright 2025-2026 Cornell University
// released under MIT license
// author: Kevin Laeufer <laeufer@cornell.edu>

use crate::poly::VarIndex;
use crate::{Coef, Mod};
use num_bigint::BigInt;
use num_traits::{Num, ToPrimitive};
use std::ops::Neg;

/// Used to efficiently represent coefficients generated by the polynomial parser.
/// Needs to be converted into a type that implements `Coef` to actually perform polynomial
/// substitution.
pub enum IntCoef {
    I64(i64),
    Big(BigInt),
}

impl IntCoef {
    fn neg(&self) -> Self {
        match self {
            IntCoef::I64(v) => IntCoef::I64(-*v),
            IntCoef::Big(v) => IntCoef::Big(v.neg()),
        }
    }

    pub fn into_coef<C: Coef>(self, m: Mod) -> C {
        match self {
            IntCoef::I64(v) => C::from_i64(v, m),
            IntCoef::Big(v) => C::from_big(&v, m),
        }
    }

    pub fn as_i64(&self) -> Option<i64> {
        match self {
            IntCoef::I64(v) => Some(*v),
            IntCoef::Big(v) => v.to_i64(),
        }
    }
}

pub fn parse_poly(line: &[u8]) -> impl Iterator<Item = (IntCoef, Vec<VarIndex>)> {
    PolyParser::new(line)
}

struct PolyParser<'a> {
    line: &'a [u8],
}

impl<'a> PolyParser<'a> {
    fn new(line: &'a [u8]) -> Self {
        Self { line }
    }
}

impl<'a> Iterator for PolyParser<'a> {
    type Item = (IntCoef, Vec<VarIndex>);

    fn next(&mut self) -> Option<Self::Item> {
        use State::*;
        let mut state = LookingForSign(0);
        let mut coef: Option<IntCoef> = None;
        let mut vars = Vec::with_capacity(16);
        while !self.line.is_empty() {
            state = match state {
                LookingForSign(i) => {
                    // end of line special case
                    if i == self.line.len() {
                        self.line = &[];
                        continue;
                    }
                    match self.line[i] {
                        b'+' => ParsingCoef(Sign::Plus, i + 1, i + 1),
                        b'-' => ParsingCoef(Sign::Minus, i + 1, i + 1),
                        // sometimes there is no sign
                        b'0' | b'1' | b'2' | b'3' | b'4' | b'5' | b'6' | b'7' | b'8' | b'9' => {
                            ParsingCoef(Sign::Plus, i, i + 1)
                        }
                        _ => LookingForSign(i + 1),
                    }
                }
                ParsingCoef(sign, start, i) => {
                    // end of line special case
                    if i == self.line.len() {
                        let coef = parse_coef(&self.line[start..i]);
                        self.line = &[];
                        return Some((coef, vec![]));
                    }
                    match self.line[i] {
                        // skip whitespace and `[` at the beginning
                        b' ' | b'[' if start == i => ParsingCoef(sign, i + 1, i + 1),
                        b'*' | b']' => {
                            let cc = parse_coef(&self.line[start..i]);
                            if sign == Sign::Minus {
                                coef = Some(cc.neg());
                            } else {
                                coef = Some(cc);
                            }
                            ParsingVar(i + 1, i + 1)
                        }
                        _ => ParsingCoef(sign, start, i + 1),
                    }
                }
                ParsingVar(start, i) => {
                    // end of line special case
                    if i == self.line.len() {
                        let var: u32 = std::str::from_utf8(&self.line[start..i])
                            .unwrap()
                            .parse()
                            .unwrap();
                        self.line = &[];
                        vars.push(var.into());
                        return Some((coef.unwrap(), vars));
                    }
                    match self.line[i] {
                        b'x' => {
                            debug_assert_eq!(start, i);
                            ParsingVar(i + 1, i + 1)
                        }
                        b'0' | b'1' | b'2' | b'3' | b'4' | b'5' | b'6' | b'7' | b'8' | b'9' => {
                            ParsingVar(start, i + 1)
                        }
                        other => {
                            if start != i {
                                let var: u32 = std::str::from_utf8(&self.line[start..i])
                                    .unwrap()
                                    .parse()
                                    .unwrap();
                                vars.push(var.into());
                            }
                            // there is another var waiting!
                            if other == b'*' {
                                ParsingVar(i + 1, i + 1)
                            } else {
                                // done!
                                self.line = &self.line[i..];
                                return Some((coef.unwrap(), vars));
                            }
                        }
                    }
                }
            }
        }
        None
    }
}

fn parse_coef(token: &[u8]) -> IntCoef {
    let s = std::str::from_utf8(token).unwrap().trim();
    // first, try to see if we can parse as an i64
    if let Ok(v) = str::parse::<i64>(s) {
        IntCoef::I64(v)
    } else {
        // if that fails, we use a bigint
        IntCoef::Big(BigInt::from_str_radix(s, 10).unwrap())
    }
}

#[derive(Debug, Copy, Clone)]
enum State {
    LookingForSign(usize),
    ParsingCoef(Sign, usize, usize),
    ParsingVar(usize, usize),
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum Sign {
    Minus,
    Plus,
}

#[cfg(test)]
mod tests {
    use super::*;

    fn do_parse_poly(line: &[u8]) -> Vec<(i64, Vec<u32>)> {
        parse_poly(line)
            .map(|(c, vars)| {
                (
                    c.as_i64().unwrap(),
                    vars.into_iter().map(|v| u32::from(v)).collect(),
                )
            })
            .collect()
    }

    #[test]
    fn test_parse_poly() {
        assert_eq!(
            do_parse_poly(b"-1*x2817+1*x33"),
            [(-1, vec![2817]), (1, vec![33])]
        );
        assert_eq!(do_parse_poly(b"+2147483648*x2848+1073741824*x2847+536870912*x2846+268435456*x2845+134217728*x2844+67108864*x2843+33554432*x2842+16777216*x2841"),
                   [(2147483648i64, vec![2848]),
                       (1073741824i64, vec![2847]),
                       (536870912i64, vec![2846]),
                       (268435456i64, vec![2845]),
                       (134217728i64, vec![2844]),
                       (67108864i64, vec![2843]),
                       (33554432i64, vec![2842]),
                       (16777216i64, vec![2841]),
                   ]);

        // here there is a `+1` at the end
        assert_eq!(
            do_parse_poly(b"-1*x663-1*x493+1"),
            [(-1, vec![663]), (-1, vec![493]), (1, vec![])]
        );

        // from `b03_sp-ar-rc_64bit_step`
        assert_eq!(
            do_parse_poly(b"-2*x65*x2-1*x65*x1"),
            [(-2, vec![65, 2]), (-1, vec![65, 1]),]
        );
    }

    // make sure that we can parse polynoms that are turned into strings using our implementation of Display
    #[test]
    fn test_parse_polynom_display() {
        assert_eq!(
            do_parse_poly(b"[1*x2817] + [1*x33]"),
            [(1, vec![2817]), (1, vec![33])]
        );

        assert_eq!(
            do_parse_poly(b"[-1*x2817] + [1*x33]"),
            [(-1, vec![2817]), (1, vec![33])]
        );

        assert_eq!(
            do_parse_poly(b"[1*x2817] + [ -1 *x33]"),
            [(1, vec![2817]), (-1, vec![33])]
        );

        // parse a monom with empty term
        assert_eq!(
            do_parse_poly(b"[2*] + [4294967294*x1*x18]"),
            [(2, vec![]), (4294967294i64, vec![1, 18])]
        );
    }
}