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))
-                    .take(count - (dh.max(count - kl)))
+                    .skip(dl.max(rolls_len - kh))
+                    .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");
                 };