Partial implementation of the Roll20 grammar in Rust-PEG

dice.rs

use rand::distributions::{Distribution, Uniform};

peg::parser!( grammar roller() for str {
    pub rule expr() -> Atom
        = l:term() op:$(['+'|'-']) r:term() {
            match op {
                "+" => l + r,
                "-" => l - r,
                _ => unreachable!(),
            }
        }

    rule term() -> Atom
        = l:call() op:$(['*'|'\\'|'/']) r:call() {
            match op {
                "*"        => l * r,
                "\\" | "/" => l / r,
                _ => unreachable!(),
            }
        }

    rule call() -> Atom
        = fun:$("floor"/"round"/"ceil"/"abs") "(" arg:expr() ")" {
            match fun {
                "floor" => arg.floor(),
                "round" => arg.round(),
                "ceil"  => arg.ceil(),
                "abs"   => arg.abs(),
                _ => unreachable!(),
            }
        }
        / atom()

    rule atom() -> Atom
        = dice()
        / number()
        / "(" e:expr() ")" { e }

    rule number() -> Atom
        = n:$(['-']?(['0']/['1'..='9']['0'..='9']*(['.']['0'..='9']+)?)) {
            Atom::Number(n.parse().unwrap())
        }

    rule dice() -> Atom
        = c:count() "d" s:sides() m:modifiers()* {
            Atom::Dice(Die::roll_from_parts(c, s, m))
        }

    rule modifiers() -> Mod
        = compounding()
        / penetrating()
        / exploding()
        / targets()
        / keep_and_drop()
        / rerolling()

    rule rerolling() -> Mod
        = "r" once:"o"? cmp:compare_point()? {
            if once.is_some() {
                Mod::RerollOnce(cmp)
            } else {
                Mod::Rerolling(cmp)
            }
        }

    rule keep_and_drop() -> Mod
        = "kh" c:int() { Mod::KeepHighest(c) }
        / "kl" c:int() { Mod::KeepLowest(c) }
        / "k" c:int()  { Mod::KeepHighest(c) }
        / "dh" c:int() { Mod::DropHighest(c) }
        / "dl" c:int() { Mod::DropLowest(c) }
        / "d" c:int()  { Mod::DropLowest(c) }

    rule targets() -> Mod
        = s:success() f:failure()* { let mut m = vec![s]; m.extend(f); Mod::Targeting(m) }

    rule success() -> Target
        = cmp:compare_point() { Target::Success(cmp) }

    rule failure() -> Target
        = "f" cmp:compare_point() { Target::Failure(cmp) }

    rule compounding() -> Mod
        = "!!" cmp:compare_point()? { Mod::Compounding(cmp) }

    rule penetrating() -> Mod
        = "!p" cmp:compare_point()? { Mod::Penetrating(cmp) }

    rule exploding() -> Mod
        = "!" cmp:compare_point()? { Mod::Exploding(cmp) }

    rule compare_point() -> Compare
        = "<" c:int() { Compare::Less(c) }
        / "=" c:int() { Compare::Equal(c) }
        / ">" c:int() { Compare::Greater(c) }
        /     c:int() { Compare::Equal(c) }

    rule count() -> Count
        = "(" e:expr() ")" { Count::Atom(Box::new(e)) }
        / i:int()          { Count::Int(i) }

    rule sides() -> Sides
        = "(" e:expr() ")" { Sides::Atom(Box::new(e)) }
        / i:int()          { Sides::Int(i) }
        / "F"              { Sides::Fudge }

    rule int() -> i64
        = n:$(['0'..='9']+) { n.parse().unwrap() }
});

#[derive(Clone, Debug)]
pub enum Sides {
    Fudge,
    Atom(Box<Atom>),
    Int(i64),
}

#[derive(Clone, Debug)]
pub enum Count {
    Atom(Box<Atom>),
    Int(i64),
}

#[derive(Clone, Debug)]
pub enum Compare {
    Less(i64),
    Equal(i64),
    Greater(i64),
}

#[derive(Clone, Debug)]
pub enum Target {
    Success(Compare),
    Failure(Compare),
}

#[derive(Clone, Debug)]
pub enum Mod {
    Compounding(Option<Compare>),
    Penetrating(Option<Compare>),
    Exploding(Option<Compare>),
    Targeting(Vec<Target>),
    KeepHighest(i64),
    KeepLowest(i64),
    DropHighest(i64),
    DropLowest(i64),
    Rerolling(Option<Compare>),
    RerollOnce(Option<Compare>),
}

#[derive(Clone, Debug)]
pub enum Atom {
    Number(f64),
    Dice(Die),
}

#[derive(Clone, Debug)]
pub struct Die {
    count: Count,
    sides: Sides,
    mods: Vec<Mod>,
    rolls: Vec<f64>,
    ind_total: f64,
    total: f64,
    hits: i64,
    boxes: Vec<Die>,
}

impl Die {
    fn roll_from_parts(count: Count, sides: Sides, mods: Vec<Mod>) -> Self {
        let mut hits = 0i64;
        let mut boxes = Vec::new();

        if let Count::Atom(ref atom) = count {
            if let Atom::Dice(ref d) = **atom {
                boxes.push(d.clone());
            }
        }

        if let Sides::Atom(ref atom) = sides {
            if let Atom::Dice(ref d) = **atom {
                boxes.push(d.clone());
            }
        }

        let s = match sides {
            Sides::Atom(ref a) => match **a {
                Atom::Dice(ref d) => d.total.trunc() as i64,
                Atom::Number(f) => f.trunc() as i64,
            },
            Sides::Fudge => 1,
            Sides::Int(i) => i,
        };

        let c = match count {
            Count::Atom(ref a) => match **a {
                Atom::Dice(ref d) => d.total.trunc() as usize,
                Atom::Number(f) => f.trunc() as usize,
            },
            Count::Int(i) => i as usize,
        };

        let die = match sides {
            Sides::Atom(_) | Sides::Int(_) => Uniform::new_inclusive(1, s),
            Sides::Fudge => Uniform::new_inclusive(-1, 1),
        };

        let mut rng = rand::thread_rng();
        let initial_rolls: Vec<f64> = die
            .sample_iter(&mut rng)
            .take(c)
            .map(|i| i as f64)
            .collect();
        let mut wrolls = initial_rolls.clone();

        for modifier in mods.iter() {
            use Mod::*;
            match modifier {
                Exploding(cmp) => {
                    let default = Compare::Equal(s);
                    let cmp = cmp.as_ref().unwrap_or(&default);
                    for roll in wrolls.clone() {
                        match cmp {
                            Compare::Less(i) => {
                                let mut exploded = Vec::new();
                                let mut wroll = roll;
                                while wroll <= *i as f64 {
                                    wroll = die.sample(&mut rng) as f64;
                                    exploded.push(roll);
                                }
                                wrolls.extend(exploded);
                            }
                            Compare::Equal(i) => {
                                let mut exploded = Vec::new();
                                let mut wroll = roll;
                                while wroll == *i as f64 {
                                    wroll = die.sample(&mut rng) as f64;
                                    exploded.push(roll);
                                }
                                wrolls.extend(exploded);
                            }
                            Compare::Greater(i) => {
                                let mut exploded = Vec::new();
                                let mut wroll = roll;
                                while wroll >= *i as f64 {
                                    wroll = die.sample(&mut rng) as f64;
                                    exploded.push(roll);
                                }
                                wrolls.extend(exploded);
                            }
                        }
                    }
                }
                Compounding(cmp) => {
                    let default = Compare::Equal(s);
                    let cmp = cmp.as_ref().unwrap_or(&default);
                    for roll in &mut wrolls {
                        match cmp {
                            Compare::Less(i) => {
                                let mut wroll = *roll;
                                while wroll <= *i as f64 {
                                    wroll = die.sample(&mut rng) as f64;
                                    *roll += wroll;
                                }
                            }
                            Compare::Equal(i) => {
                                let mut wroll = *roll;
                                while wroll == *i as f64 {
                                    wroll = die.sample(&mut rng) as f64;
                                    *roll += wroll;
                                }
                            }
                            Compare::Greater(i) => {
                                let mut wroll = *roll;
                                while wroll >= *i as f64 {
                                    wroll = die.sample(&mut rng) as f64;
                                    *roll += wroll;
                                }
                            }
                        }
                    }
                }
                Penetrating(cmp) => {
                    let default = Compare::Equal(s);
                    let cmp = cmp.as_ref().unwrap_or(&default);
                    for roll in wrolls.clone() {
                        match cmp {
                            Compare::Less(i) => {
                                let mut exploded = Vec::new();
                                let mut wroll = roll;
                                while wroll <= *i as f64 {
                                    wroll = (die.sample(&mut rng) - 1) as f64;
                                    exploded.push(roll);
                                }
                                wrolls.extend(exploded);
                            }
                            Compare::Equal(i) => {
                                let mut exploded = Vec::new();
                                let mut wroll = roll;
                                while wroll == *i as f64 {
                                    wroll = (die.sample(&mut rng) - 1) as f64;
                                    exploded.push(roll);
                                }
                                wrolls.extend(exploded);
                            }
                            Compare::Greater(i) => {
                                let mut exploded = Vec::new();
                                let mut wroll = roll;
                                while wroll >= *i as f64 {
                                    wroll = (die.sample(&mut rng) - 1) as f64;
                                    exploded.push(roll);
                                }
                                wrolls.extend(exploded);
                            }
                        }
                    }
                }
                Targeting(targets) => {
                    for target in targets {
                        match target {
                            Target::Success(cmp) => {
                                for roll in wrolls.iter().map(|f| f.trunc() as i64) {
                                    match cmp {
                                        Compare::Less(i) if roll <= *i => hits += 1,
                                        Compare::Equal(i) if roll == *i => hits += 1,
                                        Compare::Greater(i) if roll >= *i => hits += 1,
                                        _ => (),
                                    }
                                }
                            }
                            Target::Failure(cmp) => {
                                for roll in wrolls.iter().map(|f| f.trunc() as i64) {
                                    match cmp {
                                        Compare::Less(i) if roll <= *i => hits -= 1,
                                        Compare::Equal(i) if roll == *i => hits -= 1,
                                        Compare::Greater(i) if roll >= *i => hits -= 1,
                                        _ => (),
                                    }
                                }
                            }
                        }
                    }
                }
                KeepHighest(i) => {
                    wrolls.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
                    wrolls = wrolls.into_iter().rev().take(*i as usize).collect();
                }
                KeepLowest(i) => {
                    wrolls.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
                    wrolls = wrolls.into_iter().take(*i as usize).collect();
                }
                DropHighest(i) => {
                    wrolls.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
                    wrolls = wrolls.into_iter().rev().skip(*i as usize).collect();
                }
                DropLowest(i) => {
                    wrolls.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
                    wrolls = wrolls.into_iter().skip(*i as usize).collect();
                }
                Rerolling(cmp) => {
                    let default = Compare::Equal(1);
                    let cmp = cmp.as_ref().unwrap_or(&default);
                    for roll in &mut wrolls {
                        match cmp {
                            Compare::Less(i) => {
                                while *roll <= *i as f64 {
                                    *roll = die.sample(&mut rng) as f64;
                                }
                            }
                            Compare::Equal(i) => {
                                while *roll == *i as f64 {
                                    *roll = die.sample(&mut rng) as f64;
                                }
                            }
                            Compare::Greater(i) => {
                                while *roll >= *i as f64 {
                                    *roll = die.sample(&mut rng) as f64;
                                }
                            }
                        }
                    }
                }
                RerollOnce(cmp) => {
                    let default = Compare::Equal(1);
                    let cmp = cmp.as_ref().unwrap_or(&default);
                    for roll in &mut wrolls {
                        match cmp {
                            Compare::Less(i) => {
                                if *roll <= *i as f64 {
                                    *roll = die.sample(&mut rng) as f64;
                                }
                            }
                            Compare::Equal(i) => {
                                if *roll == *i as f64 {
                                    *roll = die.sample(&mut rng) as f64;
                                }
                            }
                            Compare::Greater(i) => {
                                if *roll >= *i as f64 {
                                    *roll = die.sample(&mut rng) as f64;
                                }
                            }
                        }
                    }
                }
            }
        }

        let total = wrolls.iter().sum();

        Die {
            count,
            sides,
            mods,
            rolls: initial_rolls,
            ind_total: total,
            total,
            hits,
            boxes,
        }
    }
}

impl Atom {
    fn abs(self) -> Self {
        match self {
            Atom::Number(f) => Atom::Number(f.abs()),
            Atom::Dice(mut die) => {
                die.total = die.total.abs();
                Atom::Dice(die)
            }
        }
    }

    fn ceil(self) -> Self {
        match self {
            Atom::Number(f) => Atom::Number(f.ceil()),
            Atom::Dice(mut die) => {
                die.total = die.total.ceil();
                Atom::Dice(die)
            }
        }
    }

    fn floor(self) -> Self {
        match self {
            Atom::Number(f) => Atom::Number(f.floor()),
            Atom::Dice(mut die) => {
                die.total = die.total.floor();
                Atom::Dice(die)
            }
        }
    }

    fn round(self) -> Self {
        match self {
            Atom::Number(f) => Atom::Number((f + 0.5).floor()),
            Atom::Dice(mut die) => {
                die.total = (die.total + 0.5).floor();
                Atom::Dice(die)
            }
        }
    }
}

macro_rules! impl_ops_trait_for_atom {
    ($trait:ident, $fun:ident, $op:tt) => {
        impl std::ops::$trait for Atom {
            type Output = Self;
            fn $fun(self, rhs: Self) -> Self {
                match (self, rhs) {
                    (Atom::Number(lhs), Atom::Number(rhs)) => Atom::Number(lhs $op rhs),
                    (Atom::Number(num), Atom::Dice(die)) => Atom::Dice(Die {
                        count: die.count,
                        sides: die.sides,
                        mods: die.mods,
                        rolls: die.rolls,
                        ind_total: die.ind_total,
                        total: num $op die.total,
                        hits: die.hits,
                        boxes: die.boxes,
                    }),
                    (Atom::Dice(die), Atom::Number(num)) => Atom::Dice(Die {
                        count: die.count,
                        sides: die.sides,
                        mods: die.mods,
                        rolls: die.rolls,
                        ind_total: die.ind_total,
                        total: num $op die.total,
                        hits: die.hits,
                        boxes: die.boxes,
                    }),
                    (Atom::Dice(mut lhs), Atom::Dice(mut rhs)) => {
                        let rhs_total = rhs.total;
                        lhs.boxes.extend(std::mem::take(&mut rhs.boxes));
                        lhs.boxes.push(rhs);

                        Atom::Dice(Die {
                            count: lhs.count,
                            sides: lhs.sides,
                            mods: lhs.mods,
                            rolls: lhs.rolls,
                            ind_total: lhs.ind_total,
                            total: lhs.total $op rhs_total,
                            hits: lhs.hits,
                            boxes: lhs.boxes,
                        })
                    }
                }
            }
        }
    }
}

impl_ops_trait_for_atom!(Add, add, +);
impl_ops_trait_for_atom!(Sub, sub, -);
impl_ops_trait_for_atom!(Mul, mul, *);
impl_ops_trait_for_atom!(Div, div, /);
impl_ops_trait_for_atom!(Rem, rem, %);