1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

//! Nice to have utilities that aren't core to dice
//! manipulation itself, just handy for some reason.
#[cfg(feature = "thread_rng")]
use crate::parse::{Expression, Term};
use crate::{Error, ExpressionResult, parse::ParseError};
use thiserror::Error;

#[derive(Debug, Clone, Copy, Error)]
pub enum UtilError {
    #[error("tried to DOS me.")]
    ExceededCap,
    #[error("{0}")]
    RollError(Error),
}
impl From<Error> for UtilError {
    fn from(e: Error) -> Self {
        UtilError::RollError(e)
    }
}
impl From<ParseError> for UtilError {
    fn from(e: ParseError) -> Self {
        UtilError::RollError(e.into())
    }
}

type UResult = Result<ExpressionResult, UtilError>;

#[cfg(feature = "thread_rng")]
fn exceeds_cap(dice: &Expression, cap: i64) -> bool {
    let mut roll_count = 0;
    for term in dice.terms() {
        match term {
            Term::Dice(d) => if d.size > 1 {
                roll_count += d.number;
            } else {
                roll_count += 1;
            },
            // This branch only saves time
            // in the worst case - when there's
            // a truly obscene number of terms.
            Term::Constant(_) => roll_count += 1,
        }
        // Prevent worst case performance
        if roll_count > cap {
            return true;
        }
    }
    roll_count > cap
}

mod private {
    pub trait Sealed {}
    impl Sealed for crate::parse::Expression {}
}

/// An extension trait for `Expression`.
/// The idea here is to provide convenience methods
/// that could eventually be implemented solely in terms of public API.
///
/// This trait is sealed. It cannot be implemented outside this crate.
pub trait ExpressionExt: private::Sealed {
    /// Counts the number of evaluation steps
    /// it will take to compute a result of an expression,
    /// and returns whether that sum exceeds the given `cap`.
    ///
    /// Constant terms take one step, dice terms take one step for each die.
    /// An `Nd1` dice term counts as a constant term for the purpose of this sum.
    fn exceeds_cap(&self, cap: i64) -> bool;
}
impl ExpressionExt for Expression {
    // Note that this method does not necessarily iterate over the whole expression.
    // In the case that the expression exceeds the given cap,
    // this method will return immediately upon discovering that.
    // Still, that does not make this suitable for dealing with truly absurd
    // sizes of untrusted input.
    // One might want integration with the parser for early termination, in such a case.
    // In my case, parsing and evaluating commands from Discord messages,
    // I know that the inputs I handle will not be large enough for this to matter.
    // (Due to Discord messages being limited to a maximum of 2k bytes.)
    // I would consider merging a PR that adds such functionality, if anyone is interested
    // in implementing it.
    fn exceeds_cap(&self, cap: i64) -> bool {
        exceeds_cap(self, cap)
    }
}

#[cfg(feature = "thread_rng")]
#[derive(Error, Debug)]
#[error("tried to DOS me.")]
pub struct ExceededCap;

#[cfg(feature = "thread_rng")]
impl From<ExceededCap> for UtilError {
    fn from(_: ExceededCap) -> Self {
        Self::ExceededCap
    }
}

#[cfg(feature = "thread_rng")]
pub fn roll_exp_capped(dice: Expression, cap: i64) ->
    Result<Result<ExpressionResult, crate::Error>, ExceededCap>
{
    if !exceeds_cap(&dice, cap) {
        Ok(crate::roll_expr_iter_with(&mut rand::thread_rng(), dice.into_iter()))
    } else {
        Err(ExceededCap)
    }
}

/// For providing access to mice to irresponsible users
#[cfg(feature = "thread_rng")]
pub fn roll_capped(input: &str, cap: i64) -> UResult {
    let dice = crate::parse::wrap_dice(input)?;
    match roll_exp_capped(dice, cap) {
        Ok(Ok(x)) => Ok(x),
        Ok(Err(e)) => Err(e.into()),
        Err(e) => Err(e.into()),
    }
}