added conditional fragments to exploding dice

src/dice.rs

@@ -1,6 +1,6 @@
 use anyhow::bail;
 use rand::{Rng, RngExt};
-use std::fmt::Write;
+use std::io::Write;
 
 pub type ExpBox = Box<Expression>;
 
@@ -23,24 +23,26 @@ pub struct DiceFormula {
     pub(crate) kl: Option<ExpBox>,
     pub(crate) dh: Option<ExpBox>,
     pub(crate) dl: Option<ExpBox>,
-    pub(crate) x: bool,
+    pub(crate) x: Option<Vec<CompareFragment>>,
 }
 
+#[derive(Debug, PartialEq, Clone)]
 pub enum CompareFragment {
     Eq(ExpBox),
     Gt(ExpBox),
-    Gte(ExpBox),
+    Ge(ExpBox),
     Lt(ExpBox),
-    Lte(ExpBox),
+    Le(ExpBox),
 }
 
 const DICE_POOL_LIMIT: usize = 100_000;
 
 impl Expression {
     pub fn collect_evaluation(&self, mut rng: impl Rng) -> anyhow::Result<String> {
-        let mut result = String::new();
+        let mut buffer = Vec::new();
-        let total = self.evaluate(&mut rng, &mut result, 100)?;
+        let total = self.evaluate(&mut rng, &mut buffer, 100)?;
-        Ok(format!("**{total}** = {result}"))
+        let result_text = String::from_utf8(buffer)?;
+        Ok(format!("**{total}** = {result_text}"))
     }
 
     pub fn average(&self, rng: &mut impl Rng) -> anyhow::Result<f64> {
@@ -49,8 +51,7 @@ impl Expression {
         }
         let mut average = self.sample(rng)?;
         for idx in 1..3000 {
-            average =
+            average = (average * idx as f64 + self.sample(rng)?) / (idx as f64 + 1f64);
-                (average * idx as f64 + self.sample(rng)?) / (idx as f64 + 1f64);
         }
         Ok(average)
     }
@@ -58,19 +59,18 @@ impl Expression {
     fn evaluate(
         &self,
         rng: &mut impl Rng,
-        output: &mut String,
+        writer: &mut impl Write,
         outer_precedence: i64,
     ) -> anyhow::Result<f64> {
         use Expression::*;
-        let node = self;
+        let precedence = self.precedence();
-        let precedence = node.precedence();
         let needs_parens = precedence > outer_precedence;
         if needs_parens {
-            output.push('(')
+            write!(writer, "(")?;
         }
-        let result = match node {
+        let result = match self {
             Const(x) => {
-                write!(output, "{}", x).unwrap();
+                write!(writer, "{}", x)?;
                 *x
             }
             Dice(DiceFormula {
@@ -83,15 +83,15 @@ impl Expression {
                 x,
             }) => {
                 let count = if let Some(count_node) = count_node {
-                    count_node.evaluate(rng, output, precedence)? as usize
+                    count_node.evaluate(rng, writer, precedence)? as usize
                 } else {
                     1
                 };
                 if count > DICE_POOL_LIMIT {
                     bail!("Too many dice.")
                 };
-                output.push_str("d");
+                write!(writer, "d")?;
-                let size = size_node.evaluate(rng, output, precedence)? as i64;
+                let size = size_node.evaluate(rng, writer, precedence)? as i64;
                 if size < 1 {
                     bail!("Invalid die size.")
                 }
@@ -99,42 +99,57 @@ impl Expression {
                 let mut rolls = (0..count)
                     .map(|_| rng.random_range(1..=size))
                     .collect::<Vec<_>>();
-                if *x {
+                if let Some(fragments) = x {
-                    output.push_str("x");
+                    write!(writer, "x")?;
                     if size < 2 {
                         bail!("Infinite explosion.")
                     }
+                    let mut comparers = Vec::with_capacity(fragments.len());
+                    for fragment in fragments {
+                        comparers.push(fragment.comparer(rng, writer)?);
+                    }
                     for i in 0..DICE_POOL_LIMIT {
                         if i >= rolls.len() {
                             break;
                         }
-                        if rolls[i] == size {
+                        if comparers.len() == 0 {
-                            rolls.push(rng.random_range(1..=size));
+                            // Explode on max size dice
+                            if rolls[i] == size {
+                                rolls.push(rng.random_range(1..=size));
+                            }
+                        } else {
+                            // Explode based on comparers provided
+                            for comparer in comparers.iter() {
+                                if comparer(rolls[i] as f64) {
+                                    rolls.push(rng.random_range(1..=size));
+                                    break;
+                                }
+                            }
                         }
                     }
                 }
 
                 let kh = if let Some(child) = kh {
-                    output.push_str("kh");
+                    write!(writer, "kh")?;
-                    Some(child.evaluate(rng, output, precedence)? as usize)
+                    Some(child.evaluate(rng, writer, precedence)? as usize)
                 } else {
                     None
                 };
                 let kl = if let Some(child) = kl {
-                    output.push_str("kl");
+                    write!(writer, "kl")?;
-                    Some(child.evaluate(rng, output, precedence)? as usize)
+                    Some(child.evaluate(rng, writer, precedence)? as usize)
                 } else {
                     None
                 };
                 let dh = if let Some(child) = dh {
-                    output.push_str("dh");
+                    write!(writer, "dh")?;
-                    Some(child.evaluate(rng, output, precedence)? as usize)
+                    Some(child.evaluate(rng, writer, precedence)? as usize)
                 } else {
                     None
                 };
                 let dl = if let Some(child) = dl {
-                    output.push_str("dl");
+                    write!(writer, "dl")?;
-                    Some(child.evaluate(rng, output, precedence)? as usize)
+                    Some(child.evaluate(rng, writer, precedence)? as usize)
                 } else {
                     None
                 };
@@ -150,7 +165,7 @@ impl Expression {
                 let print_all_rolls = rolls.len() <= 100;
 
                 if print_all_rolls {
-                    output.push_str(" ‹ ");
+                    write!(writer, " ‹ ")?;
                 }
                 for idx in 0..rolls.len() {
                     let mut skip = false;
@@ -181,140 +196,62 @@ impl Expression {
                     }
                     if print_all_rolls {
                         if skip {
-                            write!(output, "~~`{}`~~", roll)?;
+                            write!(writer, "~~`{}`~~", roll)?;
                         } else if size > 2 && (roll == 1 || roll == size) {
-                            write!(output, "**`{}`**", roll)?;
+                            write!(writer, "**`{}`**", roll)?;
                         } else {
-                            write!(output, "`{}`", roll)?;
+                            write!(writer, "`{}`", roll)?;
                         }
 
                         if idx != rolls.len() - 1 {
-                            output.push_str(" ");
+                            write!(writer, " ")?;
                         }
                     }
                 }
                 if print_all_rolls {
-                    output.push_str(" ›");
+                    write!(writer, " ›")?;
                 }
 
                 result_set.iter().sum()
             }
             Neg(inner_node) => {
-                output.push_str("-");
+                write!(writer, "-")?;
-                let inner = inner_node.evaluate(rng, output, precedence)?;
+                let inner = inner_node.evaluate(rng, writer, precedence)?;
                 -inner
             }
             Add(lhs_node, rhs_node) => {
-                let lhs = lhs_node.evaluate(rng, output, precedence)?;
+                let lhs = lhs_node.evaluate(rng, writer, precedence)?;
-                output.push_str(" + ");
+                write!(writer, " + ")?;
-                let rhs = rhs_node.evaluate(rng, output, precedence)?;
+                let rhs = rhs_node.evaluate(rng, writer, precedence)?;
                 lhs + rhs
             }
             Sub(lhs_node, rhs_node) => {
-                let lhs = lhs_node.evaluate(rng, output, precedence)?;
+                let lhs = lhs_node.evaluate(rng, writer, precedence)?;
-                output.push_str(" - ");
+                write!(writer, " - ")?;
-                let rhs = rhs_node.evaluate(rng, output, precedence)?;
+                let rhs = rhs_node.evaluate(rng, writer, precedence)?;
                 lhs - rhs
             }
             Mul(lhs_node, rhs_node) => {
-                let lhs = lhs_node.evaluate(rng, output, precedence)?;
+                let lhs = lhs_node.evaluate(rng, writer, precedence)?;
-                output.push_str(" × ");
+                write!(writer, " × ")?;
-                let rhs = rhs_node.evaluate(rng, output, precedence)?;
+                let rhs = rhs_node.evaluate(rng, writer, precedence)?;
                 lhs * rhs
             }
             Div(lhs_node, rhs_node) => {
-                let lhs = lhs_node.evaluate(rng, output, precedence)?;
+                let lhs = lhs_node.evaluate(rng, writer, precedence)?;
-                output.push_str(" ÷ ");
+                write!(writer, " ÷ ")?;
-                let rhs = rhs_node.evaluate(rng, output, precedence)?;
+                let rhs = rhs_node.evaluate(rng, writer, precedence)?;
                 lhs / rhs
             }
         };
         if needs_parens {
-            output.push(')')
+            write!(writer, ")")?;
         }
         Ok(result)
     }
 
     fn sample(&self, rng: &mut impl Rng) -> anyhow::Result<f64> {
-        let node = self;
+        self.evaluate(rng, &mut std::io::sink(), 100)
-        let result = match node {
-            Expression::Const(x) => *x,
-            Expression::Dice(DiceFormula {
-                count,
-                size,
-                kh,
-                kl,
-                dh,
-                dl,
-                x,
-            }) => {
-                let count = if let Some(count) = count {
-                    count.sample(rng)? as usize
-                } else {
-                    1
-                };
-                if count > DICE_POOL_LIMIT {
-                    bail!("Too many dice.")
-                };
-                let size = size.sample(rng)? as i64;
-                if size < 1 {
-                    bail!("Invalid die size.")
-                }
-
-                let mut rolls = (0..count)
-                    .map(|_| rng.random_range(1..=size))
-                    .collect::<Vec<_>>();
-                if *x {
-                    if size < 2 {
-                        bail!("Infinite explosion.")
-                    }
-                    for i in 0..DICE_POOL_LIMIT {
-                        if i >= rolls.len() {
-                            break;
-                        }
-                        if rolls[i] == size {
-                            rolls.push(rng.random_range(1..=size));
-                        }
-                    }
-                }
-                rolls.sort();
-
-                let kh = if let Some(child) = kh {
-                    child.sample(rng)? as usize
-                } else {
-                    rolls.len()
-                };
-                let kl = if let Some(child) = kl {
-                    child.sample(rng)? as usize
-                } else {
-                    rolls.len()
-                };
-                let dh = if let Some(child) = dh {
-                    child.sample(rng)? as usize
-                } else {
-                    0
-                };
-                let dl = if let Some(child) = dl {
-                    child.sample(rng)? as usize
-                } else {
-                    0
-                };
-
-                let rolls_len = rolls.len();
-                rolls
-                    .into_iter()
-                    .skip(dl.max(rolls_len - kh))
-                    .take(rolls_len - (dh.max(rolls_len - kl)))
-                    .map(|x| x as f64)
-                    .sum::<f64>()
-            }
-            Expression::Neg(id) => -id.sample(rng)?,
-            Expression::Add(lhs, rhs) => lhs.sample(rng)? + rhs.sample(rng)?,
-            Expression::Sub(lhs, rhs) => lhs.sample(rng)? - rhs.sample(rng)?,
-            Expression::Mul(lhs, rhs) => lhs.sample(rng)? * rhs.sample(rng)?,
-            Expression::Div(lhs, rhs) => lhs.sample(rng)? / rhs.sample(rng)?,
-        };
-        Ok(result)
     }
 
     fn avg(&self) -> anyhow::Result<f64> {
@@ -343,7 +280,7 @@ impl Expression {
                     bail!("Invalid die size.")
                 }
 
-                if *x {
+                if let Some(_) = x {
                     bail!("Not implemented yet");
                 }
                 if let Some(_) = kh {
@@ -371,6 +308,42 @@ impl Expression {
     }
 }
 
+impl CompareFragment {
+    fn comparer(
+        &self,
+        rng: &mut impl Rng,
+        writer: &mut impl Write,
+    ) -> anyhow::Result<Box<dyn Fn(f64) -> bool>> {
+        Ok(match self {
+            CompareFragment::Eq(inner) => {
+                write!(writer, "=")?;
+                let comp = inner.evaluate(rng, writer, 2)?;
+                Box::new(move |value| value == comp)
+            }
+            CompareFragment::Gt(inner) => {
+                write!(writer, ">")?;
+                let comp = inner.evaluate(rng, writer, 2)?;
+                Box::new(move |value| value > comp)
+            }
+            CompareFragment::Ge(inner) => {
+                write!(writer, ">=")?;
+                let comp = inner.evaluate(rng, writer, 2)?;
+                Box::new(move |value| value >= comp)
+            }
+            CompareFragment::Lt(inner) => {
+                write!(writer, "<")?;
+                let comp = inner.evaluate(rng, writer, 2)?;
+                Box::new(move |value| value < comp)
+            }
+            CompareFragment::Le(inner) => {
+                write!(writer, "<=")?;
+                let comp = inner.evaluate(rng, writer, 2)?;
+                Box::new(move |value| value <= comp)
+            }
+        })
+    }
+}
+
 impl Expression {
     fn precedence(&self) -> i64 {
         use Expression::*;

src/parsing.rs

@@ -1,7 +1,4 @@
-use crate::dice::{
+use crate::dice::{CompareFragment, DiceFormula, Expression::{self, Const}};
-    DiceFormula,
-    Expression::{self, Const},
-};
 use Assoc::Left;
 use nom::{
     Parser,
@@ -73,7 +70,7 @@ fn expr<'c, 'i>(
                 opt(preceded(tag("kl"), opt(basic_operand(ctx)))),
                 opt(preceded(tag("dh"), opt(basic_operand(ctx)))),
                 opt(preceded(tag("dl"), opt(basic_operand(ctx)))),
-                opt(preceded(tag("x"), opt(basic_operand(ctx)))),
+                opt(preceded(tag("x"), many0(compare_fragment(ctx)))),
             )),
         )
             .map(|(count, size, (kh, kl, dh, dl, x))| {
@@ -84,7 +81,7 @@ fn expr<'c, 'i>(
                     kl: kl.map(|n| Box::new(n.unwrap_or(Const(1f64)))),
                     dh: dh.map(|n| Box::new(n.unwrap_or(Const(1f64)))),
                     dl: dl.map(|n| Box::new(n.unwrap_or(Const(1f64)))),
-                    x: x.is_some(),
+                    x,
                 })
             });
         let operand = alt((dice_formula, basic_operand(ctx)));
@@ -127,6 +124,19 @@ fn basic_operand<'c, 'i>(
     alt((number, delimited(tag("("), expr(ctx), tag(")"))))
 }
 
+fn compare_fragment<'c, 'i>(
+    ctx: &'c Context,
+) -> impl Parser<&'i str, Output = CompareFragment, Error = InternalError<&'i str>> + use<'c, 'i>  {
+    use CompareFragment::*;
+    alt((
+        preceded(tag("="), basic_operand(ctx)).map(|x| Eq(Box::new(x))),
+        preceded(tag(">"), basic_operand(ctx)).map(|x| Gt(Box::new(x))),
+        preceded(tag(">="), basic_operand(ctx)).map(|x| Ge(Box::new(x))),
+        preceded(tag("<"), basic_operand(ctx)).map(|x| Lt(Box::new(x))),
+        preceded(tag("<="), basic_operand(ctx)).map(|x| Le(Box::new(x))),
+    ))
+}
+
 fn spaced_op(t: &str) -> impl Parser<&str, Output = &str, Error = VerboseError<&str>> {
     delimited(many0(tag(" ")), tag(t), many0(tag(" ")))
 }
@@ -183,4 +193,16 @@ mod test {
         let expression = parse("2d20dl1").unwrap();
         assert_matches!(expression, Expression::Dice { .. });
     }
+
+    #[test]
+    pub fn parse_x() {
+        let expression = parse("3d4x").unwrap();
+        assert_matches!(expression, Expression::Dice { .. });
+    }
+
+    #[test]
+    pub fn parse_xgt() {
+        let expression = parse("3d4x>=3").unwrap();
+        assert_matches!(expression, Expression::Dice { .. });
+    }
 }