D&D 5e stat rolls

Statrolls.hs

-- | Small library for probability calculations about D&D 5e attribute rolls.
-- The intended usage is via GHCi session, started with a command like:
--     ghci -fobject-code -O2 Statrolls.hs
{- HLINT ignore "Use camelCase" -}
module Statrolls where

import Data.Int
import Data.Bits
import Data.List ( sort, compareLength, sortOn )
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Control.Monad
import Data.Ratio

-- * Attribute arrays

-- | An array of attributes. The representation is six 5-bit unsigned integer fields,
-- packed into an 'Int'.
newtype Array = Array {unArray :: Int}

-- | Pack an attribute array into the compact representation, sorted high to low.
packArray :: [Int] -> Array
packArray vals
  | vals `compareLength` 6 == LT = error "packArray: too few values (must be 6)"
  | vals `compareLength` 6 == GT = error "packArray: too many values (must be 6)"
  | or [v < 0 || v > 31 | v <- vals] = error "packArray: value out of range (0..31)"
  | otherwise = Array $ foldl' (.|.) 0 [val `shiftL` i | (val, i) <- zip (sortOn negate vals) [0,5..30]]

infix 9 @@
-- | Get the nth-highest stat in the array.
(@@) :: Array -> Int -> Int
Array bf @@ ix
  | ix < 1 || ix > 6 = error "(@@): index out of range (1..6)"
  | otherwise = (bf .&. (2^(i+5)-2^i)) `shiftR` i where i = 5*(ix-1)

-- | Unpack an attribute array into the indivudual stats, sorted high to low.
unpackArray :: Array -> [Int]
unpackArray arr = [ arr @@ i | i <- [1..6]]

-- * Conditions on attribute arrays

infixr 3 &&&
infixr 2 |||
-- | Conjunction (AND) of two conditions.
(&&&) :: (a -> Bool) -> (a -> Bool) -> a -> Bool
(&&&) = liftA2 (&&)

-- | Disjunction (OR) of two conditions.
(|||) :: (a -> Bool) -> (a -> Bool) -> a -> Bool
(|||) = liftA2 (||)

-- | Negation (NOT) of a condition.
neg :: (a -> Bool) -> a -> Bool
neg = (not .)

-- Comparison operators form basic conditions. The left-hand argument is the index, and the right-hand argument is the value
-- to compare with.
-- For example, 1 @< 16 means "1st attribute less than 16".
infix 4 @<, @>, @<=, @>=, @==, @/=
(@<), (@>), (@<=), (@>=), (@==), (@/=) :: Int -> Int -> Array -> Bool
[(@<), (@>), (@<=), (@>=), (@==), (@/=)] = map cmpAt [(<), (>), (<=), (>=), (==), (/=)]
  where cmpAt cmp i other arr = (arr@@i) `cmp` other

-- * Probability calculation

{- | Calculate the exact probability of rolling an attribute array satisfying the given condition.

WARNING: the first call in a session will hang a little while as the 'outcomes_statroll' cache is populated.
On my machine, this takes about 20 seconds with -fobject-code -O2.
 
Examples:
  probabilityOf (1 @< 10)
      Probability that the highest attribute is below 10

  probabilityOf (1 @>= 17 &&& 2 @>= 16)
      Probability that the highest two attributes are at least 17 and 16, allowing to start with two 18s.
  
  probabilityOf (1 @== 3)
      Probability that the highest attribute is a 3, meaning you've rolled 1 on all 24d6. Very unlikely!
-}
probabilityOf :: Fractional a => (Array -> Bool) -> a
probabilityOf p = fromRational $ IntMap.foldlWithKey' step 0 outcomes_statroll % denominator_statroll
  where step acc arr count = if p (Array arr) then acc+count else acc

-- | The number of distinct results from rolling 24d6, and therefore the denominator of all probabilities.
denominator_statroll :: Integer
denominator_statroll = 6^24

-- | Cache counting all the ways to roll a certain result on 4d6kh3.
-- key = actual stat value (3..18)
outcomes_4d6kh3 :: IntMap Integer
outcomes_4d6kh3 = IntMap.fromListWith (+)
  [(sum highest, 1)
  | dice <- replicateM 4 [1..6]
  , let highest = drop 1 (sort dice)
  ]

-- | Cache counting all the ways to roll a certain stat array (order-insensitive).
-- key = packed stat array
outcomes_statroll :: IntMap Integer
outcomes_statroll = IntMap.fromListWith (+)
  [ (unArray $ packArray array, product [outcomes_4d6kh3 IntMap.! stat | stat <- array])
  | array <- replicateM 6 [3..18]
  ]