From 15163ac11de2f1edd98105cb76d95022c229ff55 Mon Sep 17 00:00:00 2001
From: Matthew Sparks <matthew@kyengineer.com>
Date: Tue, 7 Sep 2021 20:43:08 -0400
Subject: [PATCH] Adding calculations for keep, and adding validation on keep
 input

---
 dicebot/src/basic/parser.rs |  7 ++++++-
 dicebot/src/basic/roll.rs   | 27 ++++++++++++++++++++-------
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/dicebot/src/basic/parser.rs b/dicebot/src/basic/parser.rs
index f6f2863..7b11c6f 100644
--- a/dicebot/src/basic/parser.rs
+++ b/dicebot/src/basic/parser.rs
@@ -41,7 +41,12 @@ fn parse_dice(input: &str) -> IResult<&str, Dice> {
         // if ok, keep expression is present
         Ok(r) => {
             input = r.0;
-            keep  = r.1.1;
+            // don't allow keep greater than number of dice, and don't allow keep zero
+            if r.1.1 <= count && r.1.1 != "0" {
+                keep  = r.1.1;
+            } else {
+                keep  = count;
+            }
         }
         // otherwise absent and keep all dice
         Err(_) => keep = count,
diff --git a/dicebot/src/basic/roll.rs b/dicebot/src/basic/roll.rs
index 5614f73..f888eef 100644
--- a/dicebot/src/basic/roll.rs
+++ b/dicebot/src/basic/roll.rs
@@ -19,15 +19,21 @@ pub trait Rolled {
 }
 
 #[derive(Debug, PartialEq, Eq, Clone)]
-pub struct DiceRoll(pub Vec<u32>);
+// array of rolls in order, and how many dice to keep
+pub struct DiceRoll (pub Vec<u32>, usize);
 
 impl DiceRoll {
     pub fn rolls(&self) -> &[u32] {
         &self.0
     }
 
+    pub fn keep(&self) -> usize {
+        self.1
+    }
+
+    // only count kept dice in total
     pub fn total(&self) -> u32 {
-        self.0.iter().sum()
+        self.0[..=(self.1-1)].iter().sum()
     }
 }
 
@@ -41,11 +47,16 @@ impl fmt::Display for DiceRoll {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "{}", self.rolled_value())?;
         let rolls = self.rolls();
-        let mut iter = rolls.iter();
+        let keep  = self.keep();
+        let mut iter = rolls.iter().enumerate();
         if let Some(first) = iter.next() {
-            write!(f, " ({}", first)?;
+            write!(f, " ({}", first.1)?;
             for roll in iter {
-                write!(f, " + {}", roll)?;
+                if roll.0 < keep {
+                    write!(f, " + {}", roll.1)?;
+                } else {
+                    write!(f, " + [{}]", roll.1)?;
+                }
             }
             write!(f, ")")?;
         }
@@ -58,11 +69,13 @@ impl Roll for dice::Dice {
 
     fn roll(&self) -> DiceRoll {
         let mut rng = rand::thread_rng();
-        let rolls: Vec<_> = (0..self.count)
+        let mut rolls: Vec<_> = (0..self.count)
             .map(|_| rng.gen_range(1..=self.sides))
             .collect();
+        // sort rolls in descending order
+        rolls.sort_by(|a, b| b.cmp(a));
 
-        DiceRoll(rolls)
+        DiceRoll(rolls,self.keep as usize)
     }
 }