feat(filter): dice filter in grammar parsing

.changes/filter.md

@@ -0,0 +1,7 @@
+---
+"noctisroll": feat:patch
+---
+
+Support dice filter in grammar parsing.
+
+For example, `6d6k3` will roll 6 six-sided dice and keep the highest 3 results and `6d6q3` will roll 6 six-sided dice and keep the lowest 3 results.

src/lib.rs

@@ -3,11 +3,45 @@ lalrpop_util::lalrpop_mod!(pub roll);
 use lalrpop_util::{lexer::Token, ParseError};
 use rand::{rngs::OsRng, Rng};
 
+pub enum Filter {
+    MaxN(i32),
+    MinN(i32),
+    None,
+}
+
+impl Filter {
+    pub fn max(n: i32) -> Filter {
+        Filter::MaxN(n)
+    }
+
+    pub fn min(n: i32) -> Filter {
+        Filter::MinN(n)
+    }
+
+    pub fn filter(&self, points: &mut [i32]) -> Vec<i32> {
+        let mut points = points.to_vec();
+        let length = points.len();
+        match self {
+            Filter::MaxN(n) => {
+                points.sort_by(|a, b| b.cmp(a));
+                points.split_at((*n as usize).min(length)).0.to_vec()
+            }
+            Filter::MinN(n) => {
+                points.sort();
+                points.split_at((*n as usize).min(length)).0.to_vec()
+            }
+            Filter::None => points,
+        }
+    }
+}
+
 pub struct Dice {
     count: usize,
     sides: i32,
     rolls: Vec<i32>,
+    filtered: Vec<i32>,
     result: i32,
+    filter: Filter,
 }
 
 impl Dice {
@@ -16,18 +50,27 @@ impl Dice {
             count,
             sides,
             rolls: vec![],
+            filtered: vec![],
             result: 0,
+            filter: Filter::None,
         }
     }
 
+    pub fn filter(mut self, filter: Filter) -> Self {
+        self.filter = filter;
+        self
+    }
+
     pub fn roll(mut self) -> Self {
         let mut rng = OsRng;
-        let mut sum = 0;
         for _ in 0..self.count {
             let roll = rng.gen_range(1..=self.sides);
-            sum += roll;
             self.rolls.push(roll);
         }
+
+        self.filtered = self.filter.filter(&mut self.rolls);
+
+        let sum = self.filtered.iter().sum();
         self.result = sum;
         self
     }
@@ -43,7 +86,22 @@ impl Dice {
             }
         }
         result.push(']');
-        result
+        if let Filter::None = self.filter {
+            result
+        } else {
+            result.push('(');
+            let filtered_count = self.filtered.len();
+            for (idx, roll) in self.filtered.iter().enumerate() {
+                if idx == filtered_count - 1 {
+                    result.push_str(&roll.to_string());
+                } else {
+                    result.push_str(&roll.to_string());
+                    result.push_str(", ");
+                }
+            }
+            result.push(')');
+            result
+        }
     }
 }
 
@@ -89,22 +147,34 @@ mod tests {
         assert_eq!(roll_inline("-6d1").unwrap(), -6);
         assert_eq!(roll_inline("+6d1").unwrap(), 6);
         assert_eq!(roll_inline("3d1+3d1").unwrap(), 6);
+        assert_eq!(roll_inline("3d1k2").unwrap(), 2);
     }
 
     #[test]
     fn test_roll() {
         assert_eq!(roll("6d1").unwrap(), (6, "[1, 1, 1, 1, 1, 1]".to_string()));
-
         assert_eq!(roll("3d1+3d1").unwrap().1, "[1, 1, 1] + [1, 1, 1]");
+        assert_eq!(roll("3d1k2").unwrap().1, "[1, 1, 1](1, 1)");
     }
 
     #[test]
     fn test_math() {
         assert_eq!(roll_inline("min(6d1, 3d1)").unwrap(), 3);
         assert_eq!(roll_inline("max(6d1, 3d1)").unwrap(), 6);
+        assert_eq!(roll_inline("-6d1").unwrap(), -6);
         assert_eq!(roll_inline("abs(-6d1)").unwrap(), 6);
+    }
 
-        // TODO: support high precision math functions
-        // assert_eq!(roll_inline("cos(6d1)").unwrap(), -0.5403023058681398);
+    #[test]
+    fn test_filter() {
+        let mut arr = vec![1, 2, 3, 4, 5, 6];
+        assert_eq!(Filter::min(3).filter(&mut arr), vec![1, 2, 3]);
+        assert_eq!(Filter::max(3).filter(&mut arr), vec![6, 5, 4]);
+        assert_eq!(Filter::min(0).filter(&mut arr), vec![]);
+
+        let dice = Dice::new(6, 1).filter(Filter::max(3)).roll();
+        assert_eq!(dice.rolls, vec![1, 1, 1, 1, 1, 1]);
+        assert_eq!(dice.filtered, vec![1, 1, 1]);
+        assert_eq!(dice.result, 3);
     }
 }

src/roll.lalrpop

@@ -1,6 +1,6 @@
 use std::str::FromStr;
 
-use crate::Dice;
+use crate::{Dice, Filter};
 
 grammar;
 
@@ -23,7 +23,7 @@ Factor: (i32, String) = {
 };
 
 Term: (i32, String) = {
-    BasicDice,
+    FacedDice,
     "(" <Expr> ")",
     "abs(" <Expr> ")" => {
         (<>.0.abs(), format!("abs({})", <>.1))
@@ -66,22 +66,28 @@ Term: (i32, String) = {
     },
 };
 
-BasicDice: (i32, String) = {
-    <l:Num> "d" <r:Num> => {
-        let dice = Dice::new(l.0 as usize, r.0).roll();
+FacedDice: (i32, String) = {
+    BaseDice => {
+        let dice = <>.roll();
         (dice.result, dice.roll_str())
     },
-    "d" <r:Num> => {
-        let dice = Dice::new(1, r.0).roll();
+    <dice:BaseDice> "k" <k:Num> => {
+        let dice = dice.filter(Filter::MaxN(k)).roll();
         (dice.result, dice.roll_str())
     },
-    "d" => {
-        let dice = Dice::new(1, 100).roll();
-        return (dice.result, dice.roll_str());
+    <dice:BaseDice> "q" <q:Num> => {
+        let dice = dice.filter(Filter::MinN(q)).roll();
+        (dice.result, dice.roll_str())
     },
-    Num,
+    Num => (<>, format!("{}", <>)),
+}
+
+BaseDice: Dice = {
+    <l:Num> "d" <r:Num> => Dice::new(l as usize, r),
+    "d" <r:Num> => Dice::new(1, r),
+    "d" => Dice::new(1, 100),
 }
 
-Num: (i32, String) = {
-    r"[0-9]+" => (i32::from_str(<>).unwrap(), <>.to_string()),
+Num: i32 = {
+    r"[0-9]+" => i32::from_str(<>).unwrap(),
 };