initial impl

.gitignore

@@ -0,0 +1,2 @@
+/target
+.idea/

Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "krait"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+dotenv = "0.15.0"
+poise = "0.6.2"
+serenity = { version = "0.12.5", features = ["collector"] }
+tokio = { version = "1.53.1", features = ["rt-multi-thread"] }
+tracing = "0.1.44"
+tracing-subscriber = "0.3.23"
+anyhow = "1.0.104"
+nom = "8.0.0"
+thiserror = "2.0.20"
+nom-language = "0.1.0"
+id-arena = "2.3.0"
+rand_xoshiro = "0.8.1"
+rand = "0.10.2"
+itertools = "0.15.0"

src/dice.rs

@@ -0,0 +1,354 @@
+use anyhow::bail;
+use id_arena::Arena;
+use rand::{Rng, RngExt};
+use std::fmt::Write;
+
+#[derive(Debug)]
+pub struct Expression {
+    pub(crate) arena: Arena<AstNode>,
+    pub(crate) root: AstNodeId,
+}
+
+pub type AstNodeId = id_arena::Id<AstNode>;
+
+#[derive(Debug, PartialEq, Clone)]
+pub enum AstNode {
+    Const(f64),
+    Dice(DiceFormula),
+    Neg(AstNodeId),
+    Add(AstNodeId, AstNodeId),
+    Sub(AstNodeId, AstNodeId),
+    Mul(AstNodeId, AstNodeId),
+    Div(AstNodeId, AstNodeId),
+}
+
+#[derive(Debug, PartialEq, Clone)]
+pub struct DiceFormula {
+    pub(crate) count: Option<AstNodeId>,
+    pub(crate) size: AstNodeId,
+    pub(crate) kh: Option<AstNodeId>,
+    pub(crate) kl: Option<AstNodeId>,
+    pub(crate) dh: Option<AstNodeId>,
+    pub(crate) dl: Option<AstNodeId>,
+}
+
+impl DiceFormula {
+    pub fn new(size: AstNodeId) -> Self {
+        Self {
+            count: None,
+            size,
+            kh: None,
+            kl: None,
+            dh: None,
+            dl: None,
+        }
+    }
+}
+
+impl Expression {
+    pub fn evaluate(&self, mut rng: impl Rng) -> anyhow::Result<String> {
+        let mut result = String::new();
+        let total = self.node_evaluate(&mut rng, self.root, &mut result, 100)?;
+        Ok(format!("**{total}** = {result}"))
+    }
+
+    pub fn average(&self, rng: &mut impl Rng) -> anyhow::Result<f64> {
+        if let Ok(avg) = self.node_average(self.root) {
+            return Ok(avg);
+        }
+        let mut average = self.node_sample(rng, self.root)?;
+        for idx in 1..1000 {
+            average =
+                (average * idx as f64 + self.node_sample(rng, self.root)?) / (idx as f64 + 1f64);
+        }
+        Ok(average)
+    }
+
+    fn node_evaluate(
+        &self,
+        rng: &mut impl Rng,
+        node_id: AstNodeId,
+        output: &mut String,
+        outer_precedence: i64,
+    ) -> anyhow::Result<f64> {
+        use AstNode::*;
+        let node = &self.arena[node_id];
+        let precedence = node.precedence();
+        let needs_parens = precedence > outer_precedence;
+        if needs_parens {
+            output.push('(')
+        }
+        let result = match node {
+            Const(x) => {
+                write!(output, "{}", x).unwrap();
+                *x
+            }
+            Dice(DiceFormula {
+                count: count_node,
+                size: size_node,
+                kh,
+                kl,
+                dh,
+                dl,
+            }) => {
+                let count = if let Some(count_node) = count_node {
+                    self.node_evaluate(rng, *count_node, output, precedence)? as usize
+                } else {
+                    1
+                };
+                if count > 1_000_000 {
+                    anyhow::bail!("Too many dice.")
+                };
+                output.push_str("d");
+                let size = self.node_evaluate(rng, *size_node, output, precedence)? as i64;
+
+                let rolls = (0..count)
+                    .map(|_| rng.random_range(1..=size))
+                    .collect::<Vec<_>>();
+
+                let kh = if let Some(id) = kh {
+                    output.push_str("kh");
+                    Some(self.node_evaluate(rng, *id, output, precedence)? as usize)
+                } else {
+                    None
+                };
+                let kl = if let Some(id) = kl {
+                    output.push_str("kl");
+                    Some(self.node_evaluate(rng, *id, output, precedence)? as usize)
+                } else {
+                    None
+                };
+                let dh = if let Some(id) = dh {
+                    output.push_str("dh");
+                    Some(self.node_evaluate(rng, *id, output, precedence)? as usize)
+                } else {
+                    None
+                };
+                let dl = if let Some(id) = dl {
+                    output.push_str("dl");
+                    Some(self.node_evaluate(rng, *id, output, precedence)? as usize)
+                } else {
+                    None
+                };
+
+                let mut sorted_rolls = (0..count).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() {
+                    ranks[roll_idx] = rank;
+                }
+                let mut result_set = Vec::with_capacity(rolls.len());
+
+                let print_all_rolls = count <= 100;
+
+                if print_all_rolls {
+                    output.push_str(" [ ");
+                }
+                for idx in 0..count {
+                    let mut skip = false;
+                    let roll = rolls[idx];
+                    let rank = ranks[idx];
+                    if let Some(kh) = kh {
+                        if count - 1 - rank >= kh {
+                            skip = true;
+                        }
+                    }
+                    if let Some(dh) = dh {
+                        if count - 1 - rank <= dh {
+                            skip = true;
+                        }
+                    }
+                    if let Some(kl) = kl {
+                        if rank >= kl {
+                            skip = true;
+                        }
+                    }
+                    if let Some(dl) = dl {
+                        if rank <= dl {
+                            skip = true;
+                        }
+                    }
+                    if skip {
+                        if print_all_rolls {
+                            write!(output, "~~`{}`~~", roll)?;
+                        }
+                        continue;
+                    }
+                    result_set.push(roll as f64);
+                    if print_all_rolls {
+                        let bold = roll == 1 || roll == size;
+                        if bold {
+                            write!(output, "**`{}`**", roll)?;
+                        } else {
+                            write!(output, "`{}`", roll)?;
+                        }
+                        if idx != count - 1 {
+                            output.push_str(", ");
+                        }
+                    }
+                }
+                if print_all_rolls {
+                    output.push_str(" ]");
+                }
+
+                result_set.iter().sum()
+            }
+            Neg(inner_node) => {
+                output.push_str("-");
+                let inner = self.node_evaluate(rng, *inner_node, output, precedence)?;
+                -inner
+            }
+            Add(lhs_node, rhs_node) => {
+                let lhs = self.node_evaluate(rng, *lhs_node, output, precedence)?;
+                output.push_str(" + ");
+                let rhs = self.node_evaluate(rng, *rhs_node, output, precedence)?;
+                lhs + rhs
+            }
+            Sub(lhs_node, rhs_node) => {
+                let lhs = self.node_evaluate(rng, *lhs_node, output, precedence)?;
+                output.push_str(" - ");
+                let rhs = self.node_evaluate(rng, *rhs_node, output, precedence)?;
+                lhs - rhs
+            }
+            Mul(lhs_node, rhs_node) => {
+                let lhs = self.node_evaluate(rng, *lhs_node, output, precedence)?;
+                output.push_str(" × ");
+                let rhs = self.node_evaluate(rng, *rhs_node, output, precedence)?;
+                lhs * rhs
+            }
+            Div(lhs_node, rhs_node) => {
+                let lhs = self.node_evaluate(rng, *lhs_node, output, precedence)?;
+                output.push_str(" ÷ ");
+                let rhs = self.node_evaluate(rng, *rhs_node, output, precedence)?;
+                lhs / rhs
+            }
+        };
+        if needs_parens {
+            output.push(')')
+        }
+        Ok(result)
+    }
+
+    fn node_sample(&self, rng: &mut impl Rng, node_id: AstNodeId) -> anyhow::Result<f64> {
+        let node = &self.arena[node_id];
+        let result = match node {
+            AstNode::Const(x) => *x,
+            AstNode::Dice(DiceFormula {
+                count,
+                size,
+                kh,
+                kl,
+                dh,
+                dl,
+            }) => {
+                let count = if let Some(count) = count {
+                    self.node_sample(rng, *count)? as usize
+                } else {
+                    1
+                };
+                if count > 1_000_000 {
+                    anyhow::bail!("Too many dice.")
+                };
+                let size = self.node_sample(rng, *size)? as i64;
+
+                let mut rolls = (0..count)
+                    .map(|_| rng.random_range(1..=size))
+                    .collect::<Vec<_>>();
+                rolls.sort();
+
+                let kh = if let Some(id) = kh {
+                    self.node_sample(rng, *id)? as usize
+                } else {
+                    count
+                };
+                let kl = if let Some(id) = kl {
+                    self.node_sample(rng, *id)? as usize
+                } else {
+                    count
+                };
+                let dh = if let Some(id) = dh {
+                    self.node_sample(rng, *id)? as usize
+                } else {
+                    0
+                };
+                let dl = if let Some(id) = dl {
+                    self.node_sample(rng, *id)? as usize
+                } else {
+                    0
+                };
+
+                rolls
+                    .into_iter()
+                    .skip(dl.max(count - kh))
+                    .take(count - (dh.max(count - kl)))
+                    .map(|x| x as f64)
+                    .sum::<f64>()
+            }
+            AstNode::Neg(id) => -self.node_sample(rng, *id)?,
+            AstNode::Add(lhs, rhs) => self.node_sample(rng, *lhs)? + self.node_sample(rng, *rhs)?,
+            AstNode::Sub(lhs, rhs) => self.node_sample(rng, *lhs)? - self.node_sample(rng, *rhs)?,
+            AstNode::Mul(lhs, rhs) => self.node_sample(rng, *lhs)? * self.node_sample(rng, *rhs)?,
+            AstNode::Div(lhs, rhs) => self.node_sample(rng, *lhs)? / self.node_sample(rng, *rhs)?,
+        };
+        Ok(result)
+    }
+
+    fn node_average(&self, node_id: AstNodeId) -> anyhow::Result<f64> {
+        let node = &self.arena[node_id];
+        let result = match node {
+            AstNode::Const(x) => *x,
+            AstNode::Dice(DiceFormula {
+                count,
+                size,
+                kh,
+                kl,
+                dh,
+                dl,
+            }) => {
+                let count = if let Some(count) = count {
+                    self.node_average(*count)? as usize
+                } else {
+                    1
+                };
+                if count > 1_000_000 {
+                    bail!("Too many dice.")
+                };
+                let size = self.node_average(*size)? as i64;
+
+                if let Some(_) = kh {
+                    bail!("Not implemented yet");
+                };
+                if let Some(_) = kl {
+                    bail!("Not implemented yet");
+                };
+                if let Some(_) = dh {
+                    bail!("Not implemented yet");
+                };
+                if let Some(_) = dl {
+                    bail!("Not implemented yet");
+                };
+
+                count as f64 * (size as f64 + 1f64) * 0.5
+            }
+            AstNode::Neg(id) => -self.node_average(*id)?,
+            AstNode::Add(lhs, rhs) => self.node_average(*lhs)? + self.node_average(*rhs)?,
+            AstNode::Sub(lhs, rhs) => self.node_average(*lhs)? - self.node_average(*rhs)?,
+            AstNode::Mul(lhs, rhs) => self.node_average(*lhs)? * self.node_average(*rhs)?,
+            AstNode::Div(lhs, rhs) => self.node_average(*lhs)? / self.node_average(*rhs)?,
+        };
+        Ok(result)
+    }
+}
+
+impl AstNode {
+    fn precedence(&self) -> i64 {
+        use AstNode::*;
+        match self {
+            Const(_) => 1,
+            Dice { .. } => 2,
+            Neg(_) => 5,
+            Mul(_, _) | Div(_, _) => 7,
+            Add(_, _) | Sub(_, _) => 8,
+        }
+    }
+}

src/main.rs

@@ -0,0 +1,93 @@
+use std::fmt::Write;
+mod dice;
+mod parsing;
+
+use std::env;
+
+use dotenv::dotenv;
+use poise::CreateReply;
+use serenity::{client::ClientBuilder, model::prelude::*};
+
+type Context<'a> = poise::Context<'a, Data, anyhow::Error>;
+pub struct Data;
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    dotenv().ok();
+    tracing_subscriber::fmt::init();
+
+    let token = env::var("DISCORD_TOKEN").expect("Expected a token in the environment");
+
+    let intents = GatewayIntents::non_privileged();
+
+    let framework = poise::Framework::builder()
+        .options(poise::FrameworkOptions {
+            commands: vec![help(), roll()],
+            ..Default::default()
+        })
+        .setup(move |ctx, _ready, framework| {
+            Box::pin(async move {
+                poise::builtins::register_globally(ctx, &framework.options().commands).await?;
+                Ok(Data {})
+            })
+        })
+        .build();
+
+    let client = ClientBuilder::new(token, intents)
+        .framework(framework)
+        .await;
+
+    client.unwrap().start().await?;
+    Ok(())
+}
+
+#[poise::command(slash_command)]
+pub async fn help(ctx: Context<'_>, command: Option<String>) -> anyhow::Result<()> {
+    let configuration = poise::builtins::HelpConfiguration {
+        ..Default::default()
+    };
+    poise::builtins::help(ctx, command.as_deref(), configuration).await?;
+    Ok(())
+}
+
+#[poise::command(slash_command)]
+pub async fn roll(
+    ctx: Context<'_>,
+    formula: String,
+    #[description = "Roll the formula several times"] multi: Option<u32>,
+    #[description = "Make the roll visible only to you"]
+    #[flag]
+    private: bool,
+) -> anyhow::Result<()> {
+    let result = parsing::parse(&formula);
+    let expression = match result {
+        Ok(expression) => expression,
+        Err(error) => {
+            ctx.send(
+                CreateReply::default()
+                    .content(format!("Dice formula could not be parsed.\n{}", error))
+                    .ephemeral(true)
+                    .reply(true),
+            )
+            .await?;
+            return Ok(());
+        }
+    };
+    if private {
+        ctx.defer_ephemeral().await?;
+    } else {
+        ctx.defer().await?;
+    }
+    let multi = multi.unwrap_or(1).min(32);
+    let mut content = format!("**{} rolled:**\n", ctx.author().mention());
+    for _ in 0..multi {
+        write!(content, "{}\n", expression.evaluate(rand::rng())?)?;
+    }
+    write!(content, "-# Expected Value: {}", (expression.average(&mut rand::rng())? * 10.0).round() / 10.0)?;
+    let reply = CreateReply::default()
+        .content(content)
+        .ephemeral(private)
+        .reply(true);
+    ctx.send(reply).await?;
+    Ok(())
+}

src/parsing.rs

@@ -0,0 +1,208 @@
+use crate::dice::{AstNode, AstNodeId, DiceFormula, Expression};
+use Assoc::Left;
+use id_arena::Arena;
+use nom::multi::many0;
+use nom::{
+    Parser,
+    branch::alt,
+    bytes::tag,
+    character::digit1,
+    combinator::{all_consuming, complete, cut, fail, map_res},
+    sequence::delimited,
+};
+use nom_language::{
+    error::{VerboseError, convert_error},
+    precedence::{Assoc, Operation, binary_op, precedence, unary_op},
+};
+use std::cell::RefCell;
+use thiserror::Error;
+use tracing::{error, instrument};
+
+type InternalError<I> = VerboseError<I>;
+// type IResult<I, O> = nom::IResult<I, O, InternalError<I>>;
+
+#[derive(Error, Debug)]
+pub enum ParseError {
+    #[error("Parse error:\n{0}")]
+    ParseError(String),
+}
+
+struct Context {
+    arena: RefCell<Arena<AstNode>>,
+}
+impl Context {
+    fn new() -> Self {
+        Self {
+            arena: RefCell::new(Arena::new()),
+        }
+    }
+    fn alloc(&self, node: AstNode) -> AstNodeId {
+        self.arena.borrow_mut().alloc(node)
+    }
+}
+
+#[instrument]
+pub fn parse(formula: &str) -> Result<Expression, ParseError> {
+    let context = Context::new();
+    let result = root(&context).parse_complete(formula);
+    match result {
+        Ok((_, expr)) => {
+            let mut arena = context.arena.into_inner();
+            let root = arena.alloc(expr);
+            Ok(Expression { arena, root })
+        }
+        Err(nom::Err::Error(err) | nom::Err::Failure(err)) => {
+            Err(ParseError::ParseError(convert_error(formula, err)))
+        }
+        Err(nom::Err::Incomplete(_)) => {
+            error!("incomplete expression was entered!");
+            Err(ParseError::ParseError("Incomplete".to_owned()))
+        }
+    }
+}
+
+fn root<'c, 'i>(
+    ctx: &'c Context,
+) -> impl Parser<&'i str, Output = AstNode, Error = InternalError<&'i str>> + use<'c, 'i> {
+    cut(all_consuming(complete(expr(ctx))))
+}
+
+fn expr<'c, 'i>(
+    ctx: &'c Context,
+) -> 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 operand = alt((
+            number,
+            delimited(tag("("), expr(ctx), tag(")")),
+        ));
+        precedence(
+            complete(alt((unary_op(5, tag("-")), unary_op(2, tag("d"))))),
+            // 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("+")),
+                binary_op(8, Left, spaced_op("-")),
+            ))),
+            complete(operand),
+            |op: Operation<&str, &str, &str, AstNode>| {
+                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)),
+                    Binary(lhs, "+", rhs) => Add(ctx.alloc(lhs), ctx.alloc(rhs)),
+                    Binary(lhs, "-", rhs) => Sub(ctx.alloc(lhs), ctx.alloc(rhs)),
+                    _ => return Err("Invalid combination"),
+                })
+            },
+        )
+        .parse_complete(i)
+    }
+}
+
+fn spaced_op(t: &str) -> impl Parser<&str, Output = &str, Error = VerboseError<&str>> {
+    delimited(many0(tag(" ")), tag(t), many0(tag(" ")))
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use std::assert_matches;
+
+    #[test]
+    pub fn parse_basic() {
+        let expression = parse("2d6").unwrap();
+        assert_matches!(
+            expression.arena[expression.root],
+            AstNode::Dice {..}
+        );
+    }
+
+    #[test]
+    pub fn parse_bare_dice() {
+        let expression = parse("d8").unwrap();
+        assert_matches!(
+            expression.arena[expression.root],
+            AstNode::Dice {..}
+        );
+    }
+
+    #[test]
+    pub fn parse_math() {
+        let expression = parse("8 * 2 + 2 * 5").unwrap();
+        assert_matches!(
+            expression.arena[expression.root],
+            AstNode::Add {..}
+        );
+    }
+
+    #[test]
+    pub fn parse_kh1() {
+        let expression = parse("2d20kh1").unwrap();
+        assert_matches!(
+            expression.arena[expression.root],
+            AstNode::Dice {..}
+        );
+    }
+
+    #[test]
+    pub fn parse_dl1() {
+        let expression = parse("2d20dl1").unwrap();
+        assert_matches!(
+            expression.arena[expression.root],
+            AstNode::Dice {..}
+        );
+    }
+}