exploding dice

src/dice.rs

@@ -30,6 +30,7 @@ pub struct DiceFormula {
     pub(crate) kl: Option<AstNodeId>,
     pub(crate) dh: Option<AstNodeId>,
     pub(crate) dl: Option<AstNodeId>,
+    pub(crate) x: bool,
 }
 
 impl DiceFormula {
@@ -41,6 +42,7 @@ impl DiceFormula {
             kl: None,
             dh: None,
             dl: None,
+            x: false,
         }
     }
 }
@@ -90,6 +92,7 @@ impl Expression {
                 kl,
                 dh,
                 dl,
+                x,
             }) => {
                 let count = if let Some(count_node) = count_node {
                     self.node_evaluate(rng, *count_node, output, precedence)? as usize
@@ -97,14 +100,31 @@ impl Expression {
                     1
                 };
                 if count > 1_000_000 {
-                    anyhow::bail!("Too many dice.")
+                    bail!("Too many dice.")
                 };
                 output.push_str("d");
                 let size = self.node_evaluate(rng, *size_node, output, precedence)? as i64;
+                if size < 1 {
+                    bail!("Invalid die size.")
+                }
 
-                let rolls = (0..count)
+                let mut rolls = (0..count)
                     .map(|_| rng.random_range(1..=size))
                     .collect::<Vec<_>>();
+                if *x {
+                    output.push_str("x");
+                    if size < 2 {
+                        bail!("Infinite explosion.")
+                    }
+                    for i in 0..1_000_000 {
+                        if i >= rolls.len() {
+                            break;
+                        }
+                        if rolls[i] == size {
+                            rolls.push(rng.random_range(1..=size));
+                        }
+                    }
+                }
 
                 let kh = if let Some(id) = kh {
                     output.push_str("kh");
@@ -131,7 +151,7 @@ impl Expression {
                     None
                 };
 
-                let mut sorted_rolls = (0..count).collect::<Vec<_>>();
+                let mut sorted_rolls = (0..rolls.len()).collect::<Vec<_>>();
                 sorted_rolls.sort_by_key(|idx| rolls[*idx]);
                 let mut ranks = vec![0usize; sorted_rolls.len()];
                 for (rank, roll_idx) in sorted_rolls.into_iter().enumerate() {
@@ -139,22 +159,22 @@ impl Expression {
                 }
                 let mut result_set = Vec::with_capacity(rolls.len());
 
-                let print_all_rolls = count <= 100;
+                let print_all_rolls = rolls.len() <= 100;
 
                 if print_all_rolls {
                     output.push_str(" [ ");
                 }
-                for idx in 0..count {
+                for idx in 0..rolls.len() {
                     let mut skip = false;
                     let roll = rolls[idx];
                     let rank = ranks[idx];
                     if let Some(kh) = kh {
-                        if count - 1 - rank >= kh {
+                        if rolls.len() - 1 - rank >= kh {
                             skip = true;
                         }
                     }
                     if let Some(dh) = dh {
-                        if count - 1 - rank <= dh {
+                        if rolls.len() - 1 - rank <= dh {
                             skip = true;
                         }
                     }
@@ -176,13 +196,13 @@ impl Expression {
                     }
                     result_set.push(roll as f64);
                     if print_all_rolls {
-                        let bold = roll == 1 || roll == size;
+                        let bold = size > 2 && (roll == 1 || roll == size);
                         if bold {
                             write!(output, "**`{}`**", roll)?;
                         } else {
                             write!(output, "`{}`", roll)?;
                         }
-                        if idx != count - 1 {
+                        if idx != rolls.len() - 1 {
                             output.push_str(", ");
                         }
                     }
@@ -240,6 +260,7 @@ impl Expression {
                 kl,
                 dh,
                 dl,
+                x,
             }) => {
                 let count = if let Some(count) = count {
                     self.node_sample(rng, *count)? as usize
@@ -247,24 +268,40 @@ impl Expression {
                     1
                 };
                 if count > 1_000_000 {
-                    anyhow::bail!("Too many dice.")
+                    bail!("Too many dice.")
                 };
                 let size = self.node_sample(rng, *size)? 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..1_000_000 {
+                        if i >= rolls.len() {
+                            break;
+                        }
+                        if rolls[i] == size {
+                            rolls.push(rng.random_range(1..=size));
+                        }
+                    }
+                }
                 rolls.sort();
 
                 let kh = if let Some(id) = kh {
                     self.node_sample(rng, *id)? as usize
                 } else {
-                    count
+                    rolls.len()
                 };
                 let kl = if let Some(id) = kl {
                     self.node_sample(rng, *id)? as usize
                 } else {
-                    count
+                    rolls.len()
                 };
                 let dh = if let Some(id) = dh {
                     self.node_sample(rng, *id)? as usize
@@ -277,10 +314,11 @@ impl Expression {
                     0
                 };
 
+                let rolls_len = rolls.len();
                 rolls
                     .into_iter()
-                    .skip(dl.max(count - kh))
+                    .skip(dl.max(rolls_len - kh))
-                    .take(count - (dh.max(count - kl)))
+                    .take(rolls_len - (dh.max(rolls_len - kl)))
                     .map(|x| x as f64)
                     .sum::<f64>()
             }
@@ -304,6 +342,7 @@ impl Expression {
                 kl,
                 dh,
                 dl,
+                x,
             }) => {
                 let count = if let Some(count) = count {
                     self.node_average(*count)? as usize
@@ -314,7 +353,13 @@ impl Expression {
                     bail!("Too many dice.")
                 };
                 let size = self.node_average(*size)? as i64;
+                if size < 1 {
+                    bail!("Invalid die size.")
+                }
 
+                if *x {
+                    bail!("Not implemented yet");
+                }
                 if let Some(_) = kh {
                     bail!("Not implemented yet");
                 };

src/main.rs

@@ -7,6 +7,7 @@ use std::env;
 use dotenv::dotenv;
 use poise::CreateReply;
 use serenity::{client::ClientBuilder, model::prelude::*};
+use serenity::builder::CreateEmbed;
 
 type Context<'a> = poise::Context<'a, Data, anyhow::Error>;
 pub struct Data;

src/parsing.rs

@@ -1,14 +1,18 @@
+use crate::dice::AstNode::Const;
 use crate::dice::{AstNode, AstNodeId, DiceFormula, Expression};
 use Assoc::Left;
 use id_arena::Arena;
-use nom::multi::many0;
 use nom::{
     Parser,
     branch::alt,
+    branch::permutation,
     bytes::tag,
     character::digit1,
+    combinator::opt,
     combinator::{all_consuming, complete, cut, fail, map_res},
+    multi::many0,
     sequence::delimited,
+    sequence::preceded,
 };
 use nom_language::{
     error::{VerboseError, convert_error},
@@ -19,7 +23,7 @@ use thiserror::Error;
 use tracing::{error, instrument};
 
 type InternalError<I> = VerboseError<I>;
-// type IResult<I, O> = nom::IResult<I, O, InternalError<I>>;
+type IResult<I, O> = nom::IResult<I, O, InternalError<I>>;
 
 #[derive(Error, Debug)]
 pub enum ParseError {
@@ -72,26 +76,36 @@ fn expr<'c, 'i>(
 ) -> impl Parser<&'i str, Output = AstNode, Error = InternalError<&'i str>> + use<'c, 'i> {
     use crate::dice::AstNode::*;
     |i| {
-        let number = map_res(digit1(), |s: &str| s.parse::<f64>()).map(Const);
+        let dice_formula = (
+            opt(basic_operand(ctx)),
+            preceded(tag("d"), basic_operand(ctx)),
+            permutation((
+                opt(preceded(tag("kh"), opt(basic_operand(ctx)))),
+                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)))),
+            )),
+        )
+            .map(|(count, size, (kh, kl, dh, dl, x))| {
+                Dice(DiceFormula {
+                    count: count.map(|n| ctx.alloc(n)),
+                    size: ctx.alloc(size),
+                    kh: kh.map(|n| ctx.alloc(n.unwrap_or(Const(1f64)))),
+                    kl: kl.map(|n| ctx.alloc(n.unwrap_or(Const(1f64)))),
+                    dh: dh.map(|n| ctx.alloc(n.unwrap_or(Const(1f64)))),
+                    dl: dl.map(|n| ctx.alloc(n.unwrap_or(Const(1f64)))),
+                    x: x.is_some(),
+                })
+            });
         let operand = alt((
-            number,
+            dice_formula,
-            delimited(tag("("), expr(ctx), tag(")")),
+            basic_operand(ctx),
         ));
         precedence(
-            complete(alt((unary_op(5, tag("-")), unary_op(2, tag("d"))))),
+            complete(unary_op(5, tag("-"))),
-            // complete(alt((
-            //     unary_op(4, tag("kh")),
-            //     unary_op(4, tag("kl")),
-            //     unary_op(4, tag("dh")),
-            //     unary_op(4, tag("dl")),
-            // ))),
             fail(),
             complete(alt((
-                binary_op(3, Left, tag("kh")),
-                binary_op(3, Left, tag("kl")),
-                binary_op(3, Left, tag("dh")),
-                binary_op(3, Left, tag("dl")),
-                binary_op(1, Left, tag("d")),
                 binary_op(7, Left, spaced_op("*")),
                 binary_op(7, Left, spaced_op("/")),
                 binary_op(8, Left, spaced_op("+")),
@@ -102,43 +116,6 @@ fn expr<'c, 'i>(
                 use nom_language::precedence::Operation::*;
 
                 Ok(match op {
-                    Prefix("d", x) => Dice(DiceFormula::new(ctx.alloc(x))),
-                    Binary(count, "d", size) => Dice(DiceFormula {
-                        count: Some(ctx.alloc(count)),
-                        ..DiceFormula::new(ctx.alloc(size))
-                    }),
-                    Binary(Dice(dice), "kh", x) => Dice(DiceFormula {
-                        kh: Some(ctx.alloc(x)),
-                        ..dice
-                    }),
-                    Binary(Dice(dice), "kl", x) => Dice(DiceFormula {
-                        kl: Some(ctx.alloc(x)),
-                        ..dice
-                    }),
-                    Binary(Dice(dice), "dh", x) => Dice(DiceFormula {
-                        dh: Some(ctx.alloc(x)),
-                        ..dice
-                    }),
-                    Binary(Dice(dice), "dl", x) => Dice(DiceFormula {
-                        dl: Some(ctx.alloc(x)),
-                        ..dice
-                    }),
-                    Postfix(Dice(dice), "kh") => Dice(DiceFormula {
-                        kh: Some(ctx.alloc(Const(1.0))),
-                        ..dice
-                    }),
-                    Postfix(Dice(dice), "kl") => Dice(DiceFormula {
-                        kl: Some(ctx.alloc(Const(1.0))),
-                        ..dice
-                    }),
-                    Postfix(Dice(dice), "dh") => Dice(DiceFormula {
-                        dh: Some(ctx.alloc(Const(1.0))),
-                        ..dice
-                    }),
-                    Postfix(Dice(dice), "dl") => Dice(DiceFormula {
-                        dl: Some(ctx.alloc(Const(1.0))),
-                        ..dice
-                    }),
                     Prefix("-", x) => Neg(ctx.alloc(x)),
                     Binary(lhs, "*", rhs) => Mul(ctx.alloc(lhs), ctx.alloc(rhs)),
                     Binary(lhs, "/", rhs) => Div(ctx.alloc(lhs), ctx.alloc(rhs)),
@@ -152,6 +129,22 @@ fn expr<'c, 'i>(
     }
 }
 
+fn number(i: &str) -> IResult<&str, AstNode> {
+    map_res(digit1(), |s: &str| s.parse::<f64>())
+        .map(Const)
+        .parse_complete(i)
+}
+
+fn basic_operand<'c, 'i>(
+    ctx: &'c Context,
+) -> impl Parser<&'i str, Output = AstNode, Error = InternalError<&'i str>> + use<'c, 'i> {
+    alt((
+        number,
+        delimited(tag("("), expr(ctx), tag(")")),
+    ))
+}
+
+
 fn spaced_op(t: &str) -> impl Parser<&str, Output = &str, Error = VerboseError<&str>> {
     delimited(many0(tag(" ")), tag(t), many0(tag(" ")))
 }
@@ -164,45 +157,48 @@ mod test {
     #[test]
     pub fn parse_basic() {
         let expression = parse("2d6").unwrap();
-        assert_matches!(
+        assert_matches!(expression.arena[expression.root], AstNode::Dice { .. });
-            expression.arena[expression.root],
-            AstNode::Dice {..}
-        );
     }
 
     #[test]
     pub fn parse_bare_dice() {
         let expression = parse("d8").unwrap();
-        assert_matches!(
+        assert_matches!(expression.arena[expression.root], AstNode::Dice { .. });
-            expression.arena[expression.root],
-            AstNode::Dice {..}
-        );
     }
 
     #[test]
     pub fn parse_math() {
         let expression = parse("8 * 2 + 2 * 5").unwrap();
-        assert_matches!(
+        assert_matches!(expression.arena[expression.root], AstNode::Add { .. });
-            expression.arena[expression.root],
+    }
-            AstNode::Add {..}
+
-        );
+    #[test]
+    pub fn parse_kh() {
+        let expression = parse("2d20kh").unwrap();
+        assert_matches!(expression.arena[expression.root], AstNode::Dice { .. });
     }
 
     #[test]
     pub fn parse_kh1() {
         let expression = parse("2d20kh1").unwrap();
-        assert_matches!(
+        assert_matches!(expression.arena[expression.root], AstNode::Dice { .. });
-            expression.arena[expression.root],
+    }
-            AstNode::Dice {..}
+
-        );
+    #[test]
+    pub fn parse_kh_nested() {
+        let expression = parse("2d(10*2)kh(2-3)").unwrap();
+        assert_matches!(expression.arena[expression.root], AstNode::Dice { .. });
+    }
+
+    #[test]
+    pub fn parse_dl() {
+        let expression = parse("2d20dl").unwrap();
+        assert_matches!(expression.arena[expression.root], AstNode::Dice { .. });
     }
 
     #[test]
     pub fn parse_dl1() {
         let expression = parse("2d20dl1").unwrap();
-        assert_matches!(
+        assert_matches!(expression.arena[expression.root], AstNode::Dice { .. });
-            expression.arena[expression.root],
-            AstNode::Dice {..}
-        );
     }
 }