78 lines
3.2 KiB
Haskell
78 lines
3.2 KiB
Haskell
module Jobs.Handler.Intervals.Utils
|
|
( mkIntervals, mkIntervalsCached
|
|
, getCurrentInterval
|
|
, currentIntervalCached
|
|
) where
|
|
|
|
import Import hiding (init, maximumBy, cached)
|
|
|
|
import Control.Monad.Random.Lazy (evalRand, mkStdGen)
|
|
import System.Random.Shuffle (shuffleM)
|
|
import Data.List ((!!), unfoldr, maximumBy, init, genericLength)
|
|
import qualified Data.ByteString as ByteString
|
|
import Data.Bits (Bits(shiftR))
|
|
|
|
import qualified Data.Map.Strict as Map
|
|
|
|
|
|
mkIntervals :: forall a. Integer -> (ByteString -> Maybe a) -> Natural -> [(Maybe a, Maybe a)]
|
|
mkIntervals bytes fromBS numIterations = zip (Nothing : init intervals') intervals'
|
|
where
|
|
bits = bytes * 8
|
|
base :: Integer
|
|
base = 2 ^ bits
|
|
|
|
-- | Exclusive upper bounds
|
|
intervals
|
|
| numIterations <= 0 = pure base
|
|
| otherwise = go protoIntervals ^.. folded . _1
|
|
where
|
|
go [] = []
|
|
go ints
|
|
| maximumOf (folded . _1) ints == Just base = ints
|
|
| otherwise = go $ lts ++ over _1 succ (over _2 (subtract $ toInteger numIterations) closest) : map (over _1 succ) gts
|
|
where
|
|
closest = maximumBy (comparing $ view _2) ints
|
|
(lts, geqs) = partition (((>) `on` view _1) closest) ints
|
|
gts = filter (((<) `on` view _1) closest) geqs
|
|
-- | Exclusive upper bounds
|
|
protoIntervals :: [(Integer, Integer)]
|
|
protoIntervals = [ over _1 (i *) $ base `divMod` toInteger numIterations
|
|
| i <- [1 .. toInteger numIterations]
|
|
]
|
|
|
|
intervals' = map (fromBS' <=< assertM' (> 0)) intervals
|
|
|
|
fromBS' :: Integer -> Maybe a
|
|
fromBS' = fromBS . pad . ByteString.pack . reverse . unfoldr step
|
|
where step i
|
|
| i <= 0 || i >= base = Nothing
|
|
| otherwise = Just (fromIntegral i, i `shiftR` 8)
|
|
pad bs
|
|
| toInteger (ByteString.length bs) >= bytes = bs
|
|
| otherwise = pad $ ByteString.cons 0 bs
|
|
|
|
getCurrentInterval :: forall a. Natural -> Natural -> [a] -> a
|
|
getCurrentInterval epoch iteration intervals = permIntervals !! fromIntegral (toInteger iteration `mod` genericLength permIntervals)
|
|
where permIntervals = shuffleM intervals `evalRand` mkStdGen (hash epoch)
|
|
|
|
|
|
mkIntervalsCached :: forall m a. (NFData a, MonadIO m)
|
|
=> TVar (Map Natural [(Maybe a, Maybe a)])
|
|
-> Integer -> (ByteString -> Maybe a) -> Natural -> m [(Maybe a, Maybe a)]
|
|
mkIntervalsCached cacheTVar bytes fromBS numIterations = atomically $ do
|
|
cached <- readTVar cacheTVar
|
|
case Map.lookup numIterations cached of
|
|
Just c -> return c
|
|
Nothing -> do
|
|
modifyTVar' cacheTVar $ force . Map.insert numIterations intervals'
|
|
return intervals'
|
|
where intervals' = mkIntervals bytes fromBS numIterations
|
|
|
|
currentIntervalCached :: forall m a. (NFData a, MonadIO m)
|
|
=> TVar (Map Natural [(Maybe a, Maybe a)])
|
|
-> Integer -> (ByteString -> Maybe a)
|
|
-> Natural -> Natural -> Natural -> m (Maybe a, Maybe a)
|
|
currentIntervalCached cacheTVar bytes fromBS numIterations epoch iteration
|
|
= getCurrentInterval epoch iteration <$> mkIntervalsCached cacheTVar bytes fromBS numIterations
|