{-# LANGUAGE NamedFieldPuns #-}
module Freckle.App.Random
( smallRandomSubsetOfLargeIntegerRange
, Range (..)
, NonEmptyRange (..)
, inclusiveRange
) where
import Freckle.App.Prelude
import Control.Monad.Random (MonadRandom (..), Random)
import Control.Monad.ST (runST)
import Control.Monad.State.Strict (execStateT, get, put)
import Data.Functor ((<&>))
import Data.STRef (newSTRef, readSTRef, writeSTRef)
import Numeric.Natural (Natural)
import qualified Data.Set as Set
data Range i
= RangeEmpty
| RangeNonEmpty (NonEmptyRange i)
data NonEmptyRange i = NonEmptyRange
{ forall i. NonEmptyRange i -> i
inclusiveMinBound :: i
, forall i. NonEmptyRange i -> Natural
offset :: Natural
}
inclusiveRange
:: Integral i
=> i
-> i
-> Range i
inclusiveRange :: forall i. Integral i => i -> i -> Range i
inclusiveRange i
a i
b =
if i
a forall a. Ord a => a -> a -> Bool
<= i
b
then forall i. NonEmptyRange i -> Range i
RangeNonEmpty (forall i. i -> Natural -> NonEmptyRange i
NonEmptyRange i
a (forall a b. (Integral a, Num b) => a -> b
fromIntegral (i
b forall a. Num a => a -> a -> a
- i
a)))
else forall i. Range i
RangeEmpty
smallRandomSubsetOfLargeIntegerRange
:: (MonadRandom m, Random i, Integral i)
=> Natural
-> Range i
-> m (Set i)
smallRandomSubsetOfLargeIntegerRange :: forall (m :: * -> *) i.
(MonadRandom m, Random i, Integral i) =>
Natural -> Range i -> m (Set i)
smallRandomSubsetOfLargeIntegerRange Natural
n = \case
Range i
RangeEmpty -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Set a
Set.empty
RangeNonEmpty NonEmptyRange i
r ->
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall i. RangeWithGaps i -> Set i
gaps forall a b. (a -> b) -> a -> b
$
forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT (forall i. NonEmptyRange i -> Set i -> RangeWithGaps i
RangeWithGaps NonEmptyRange i
r forall a. Set a
Set.empty) forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Natural
1 .. Natural
n] forall a b. (a -> b) -> a -> b
$ \Natural
_ -> do
forall s (m :: * -> *). MonadState s m => m s
get forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) i.
(MonadRandom m, Random i, Integral i) =>
RangeWithGaps i -> m (RangeWithGaps i)
randomlyRemove) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall s (m :: * -> *). MonadState s m => s -> m ()
put
data RangeWithGaps i = RangeWithGaps
{ forall i. RangeWithGaps i -> NonEmptyRange i
rangeWithoutGaps :: NonEmptyRange i
, forall i. RangeWithGaps i -> Set i
gaps :: Set i
}
randomlyRemove
:: (MonadRandom m, Random i, Integral i) => RangeWithGaps i -> m (RangeWithGaps i)
randomlyRemove :: forall (m :: * -> *) i.
(MonadRandom m, Random i, Integral i) =>
RangeWithGaps i -> m (RangeWithGaps i)
randomlyRemove RangeWithGaps i
rg =
forall (m :: * -> *) i.
(MonadRandom m, Random i, Integral i) =>
RangeWithGaps i -> m (Maybe i)
randomFromRangeWithGaps RangeWithGaps i
rg forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
Maybe i
Nothing -> RangeWithGaps i
rg
Just i
i -> RangeWithGaps i
rg {gaps :: Set i
gaps = forall a. Ord a => a -> Set a -> Set a
Set.insert i
i (forall i. RangeWithGaps i -> Set i
gaps RangeWithGaps i
rg)}
randomFromRangeWithGaps
:: (MonadRandom m, Random i, Integral i)
=> RangeWithGaps i
-> m (Maybe i)
randomFromRangeWithGaps :: forall (m :: * -> *) i.
(MonadRandom m, Random i, Integral i) =>
RangeWithGaps i -> m (Maybe i)
randomFromRangeWithGaps RangeWithGaps i
rg =
let
RangeWithGaps {NonEmptyRange i
rangeWithoutGaps :: NonEmptyRange i
rangeWithoutGaps :: forall i. RangeWithGaps i -> NonEmptyRange i
rangeWithoutGaps, Set i
gaps :: Set i
gaps :: forall i. RangeWithGaps i -> Set i
gaps} = RangeWithGaps i
rg
NonEmptyRange {i
inclusiveMinBound :: i
inclusiveMinBound :: forall i. NonEmptyRange i -> i
inclusiveMinBound, Natural
offset :: Natural
offset :: forall i. NonEmptyRange i -> Natural
offset} = NonEmptyRange i
rangeWithoutGaps
in
if forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Set a -> Int
Set.size Set i
gaps) forall a. Eq a => a -> a -> Bool
== Natural
offset forall a. Num a => a -> a -> a
+ Natural
1
then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
else
forall a. a -> Maybe a
Just
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> do
i
r <-
(i
inclusiveMinBound forall a. Num a => a -> a -> a
+)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. (MonadRandom m, Random a) => (a, a) -> m a
getRandomR (i
0, forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
offset forall a. Num a => a -> a -> a
- forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Set a -> Int
Set.size Set i
gaps))
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
STRef s i
xRef <- forall a s. a -> ST s (STRef s a)
newSTRef i
r
STRef s [i]
gapQueue <- forall a s. a -> ST s (STRef s a)
newSTRef forall a b. (a -> b) -> a -> b
$ forall a. Set a -> [a]
Set.toAscList Set i
gaps
let go :: ST s i
go = do
i
x <- forall s a. STRef s a -> ST s a
readSTRef STRef s i
xRef
forall s a. STRef s a -> ST s a
readSTRef STRef s [i]
gapQueue forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
i
g : [i]
gs | i
g forall a. Ord a => a -> a -> Bool
<= i
x -> do
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s i
xRef (i
x forall a. Num a => a -> a -> a
+ i
1)
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [i]
gapQueue [i]
gs
ST s i
go
[i]
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure i
x
ST s i
go