module General.Bag(
Bag, Randomly,
emptyPure, emptyRandom,
insert, remove
) where
import qualified Data.HashMap.Strict as Map
import System.Random
type Randomly a = IO a
data Bag a
= BagPure [a]
| BagRandom {-# UNPACK #-} !Int (Map.HashMap Int a)
emptyPure :: Bag a
emptyPure :: Bag a
emptyPure = [a] -> Bag a
forall a. [a] -> Bag a
BagPure []
emptyRandom :: Bag a
emptyRandom :: Bag a
emptyRandom = Int -> HashMap Int a -> Bag a
forall a. Int -> HashMap Int a -> Bag a
BagRandom 0 HashMap Int a
forall k v. HashMap k v
Map.empty
insert :: a -> Bag a -> Bag a
insert :: a -> Bag a -> Bag a
insert x :: a
x (BagPure xs :: [a]
xs) = [a] -> Bag a
forall a. [a] -> Bag a
BagPure ([a] -> Bag a) -> [a] -> Bag a
forall a b. (a -> b) -> a -> b
$ a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs
insert x :: a
x (BagRandom n :: Int
n mp :: HashMap Int a
mp) = Int -> HashMap Int a -> Bag a
forall a. Int -> HashMap Int a -> Bag a
BagRandom (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) (HashMap Int a -> Bag a) -> HashMap Int a -> Bag a
forall a b. (a -> b) -> a -> b
$ Int -> a -> HashMap Int a -> HashMap Int a
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert Int
n a
x HashMap Int a
mp
remove :: Bag a -> Maybe (Randomly (a, Bag a))
remove :: Bag a -> Maybe (Randomly (a, Bag a))
remove (BagPure []) = Maybe (Randomly (a, Bag a))
forall a. Maybe a
Nothing
remove (BagPure (x :: a
x:xs :: [a]
xs)) = Randomly (a, Bag a) -> Maybe (Randomly (a, Bag a))
forall a. a -> Maybe a
Just (Randomly (a, Bag a) -> Maybe (Randomly (a, Bag a)))
-> Randomly (a, Bag a) -> Maybe (Randomly (a, Bag a))
forall a b. (a -> b) -> a -> b
$ (a, Bag a) -> Randomly (a, Bag a)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
x, [a] -> Bag a
forall a. [a] -> Bag a
BagPure [a]
xs)
remove (BagRandom n :: Int
n mp :: HashMap Int a
mp)
| Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 0 = Maybe (Randomly (a, Bag a))
forall a. Maybe a
Nothing
| Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 1 = Randomly (a, Bag a) -> Maybe (Randomly (a, Bag a))
forall a. a -> Maybe a
Just (Randomly (a, Bag a) -> Maybe (Randomly (a, Bag a)))
-> Randomly (a, Bag a) -> Maybe (Randomly (a, Bag a))
forall a b. (a -> b) -> a -> b
$ (a, Bag a) -> Randomly (a, Bag a)
forall (m :: * -> *) a. Monad m => a -> m a
return (HashMap Int a
mp HashMap Int a -> Int -> a
forall k v. (Eq k, Hashable k) => HashMap k v -> k -> v
Map.! 0, Bag a
forall a. Bag a
emptyRandom)
| Bool
otherwise = Randomly (a, Bag a) -> Maybe (Randomly (a, Bag a))
forall a. a -> Maybe a
Just (Randomly (a, Bag a) -> Maybe (Randomly (a, Bag a)))
-> Randomly (a, Bag a) -> Maybe (Randomly (a, Bag a))
forall a b. (a -> b) -> a -> b
$ do
Int
i <- (Int, Int) -> IO Int
forall a. Random a => (a, a) -> IO a
randomRIO (0, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1)
let mp2 :: HashMap Int a
mp2 | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1 = Int -> HashMap Int a -> HashMap Int a
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
Map.delete Int
i HashMap Int a
mp
| Bool
otherwise = Int -> a -> HashMap Int a -> HashMap Int a
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert Int
i (HashMap Int a
mp HashMap Int a -> Int -> a
forall k v. (Eq k, Hashable k) => HashMap k v -> k -> v
Map.! (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1)) (HashMap Int a -> HashMap Int a) -> HashMap Int a -> HashMap Int a
forall a b. (a -> b) -> a -> b
$ Int -> HashMap Int a -> HashMap Int a
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
Map.delete (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) HashMap Int a
mp
(a, Bag a) -> Randomly (a, Bag a)
forall (m :: * -> *) a. Monad m => a -> m a
return (HashMap Int a
mp HashMap Int a -> Int -> a
forall k v. (Eq k, Hashable k) => HashMap k v -> k -> v
Map.! Int
i, Int -> HashMap Int a -> Bag a
forall a. Int -> HashMap Int a -> Bag a
BagRandom (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) HashMap Int a
mp2)