{-# LANGUAGE BlockArguments  #-}
{-# LANGUAGE TemplateHaskell #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Disco.Effects.Store
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Polysemy effect for a memory store with integer keys.
--
-----------------------------------------------------------------------------

module Disco.Effects.Store where

import qualified Data.IntMap.Lazy      as IntMap
import           Data.IntSet           (IntSet)
import qualified Data.IntSet           as IntSet

import           Disco.Effects.Counter
import           Polysemy
import           Polysemy.State

data Store v m a where

  ClearStore  :: Store v m ()
  New         :: v -> Store v m Int
  LookupStore :: Int -> Store v m (Maybe v)
  InsertStore :: Int -> v -> Store v m ()
  MapStore    :: (v -> v) -> Store v m ()
  AssocsStore :: Store v m [(Int, v)]
  KeepKeys    :: IntSet -> Store v m ()

makeSem ''Store

-- | Dispatch a store effect.
runStore :: forall v r a. Sem (Store v ': r) a -> Sem r a
runStore :: Sem (Store v : r) a -> Sem r a
runStore
  = Sem (Counter : r) a -> Sem r a
forall (r :: [(* -> *) -> * -> *]) a.
Sem (Counter : r) a -> Sem r a
runCounter
  (Sem (Counter : r) a -> Sem r a)
-> (Sem (Store v : r) a -> Sem (Counter : r) a)
-> Sem (Store v : r) a
-> Sem r a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap v
-> Sem (State (IntMap v) : Counter : r) a -> Sem (Counter : r) a
forall s (r :: [(* -> *) -> * -> *]) a.
s -> Sem (State s : r) a -> Sem r a
evalState @(IntMap.IntMap v) IntMap v
forall a. IntMap a
IntMap.empty
  (Sem (State (IntMap v) : Counter : r) a -> Sem (Counter : r) a)
-> (Sem (Store v : r) a -> Sem (State (IntMap v) : Counter : r) a)
-> Sem (Store v : r) a
-> Sem (Counter : r) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (rInitial :: [(* -> *) -> * -> *]) x.
 Store v (Sem rInitial) x -> Sem (State (IntMap v) : Counter : r) x)
-> Sem (Store v : r) a -> Sem (State (IntMap v) : Counter : r) a
forall (e1 :: (* -> *) -> * -> *) (e2 :: (* -> *) -> * -> *)
       (e3 :: (* -> *) -> * -> *) (r :: [(* -> *) -> * -> *]) a.
FirstOrder e1 "reinterpret2" =>
(forall (rInitial :: [(* -> *) -> * -> *]) x.
 e1 (Sem rInitial) x -> Sem (e2 : e3 : r) x)
-> Sem (e1 : r) a -> Sem (e2 : e3 : r) a
reinterpret2 \case
      Store v (Sem rInitial) x
ClearStore      -> IntMap v -> Sem (State (IntMap v) : Counter : r) ()
forall s (r :: [(* -> *) -> * -> *]).
Member (State s) r =>
s -> Sem r ()
put IntMap v
forall a. IntMap a
IntMap.empty
      New v           -> do
        Int
loc <- Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int)
-> Sem (State (IntMap v) : Counter : r) Integer
-> Sem (State (IntMap v) : Counter : r) Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sem (State (IntMap v) : Counter : r) Integer
forall (r :: [(* -> *) -> * -> *]).
Member Counter r =>
Sem r Integer
next
        (IntMap v -> IntMap v) -> Sem (State (IntMap v) : Counter : r) ()
forall s (r :: [(* -> *) -> * -> *]).
Member (State s) r =>
(s -> s) -> Sem r ()
modify ((IntMap v -> IntMap v) -> Sem (State (IntMap v) : Counter : r) ())
-> (IntMap v -> IntMap v)
-> Sem (State (IntMap v) : Counter : r) ()
forall a b. (a -> b) -> a -> b
$ Int -> v -> IntMap v -> IntMap v
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
loc v
v
        Int -> Sem (State (IntMap v) : Counter : r) Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
loc
      LookupStore k   -> (IntMap v -> Maybe v)
-> Sem (State (IntMap v) : Counter : r) (Maybe v)
forall s a (r :: [(* -> *) -> * -> *]).
Member (State s) r =>
(s -> a) -> Sem r a
gets (Int -> IntMap v -> Maybe v
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
k)
      InsertStore k v -> (IntMap v -> IntMap v) -> Sem (State (IntMap v) : Counter : r) ()
forall s (r :: [(* -> *) -> * -> *]).
Member (State s) r =>
(s -> s) -> Sem r ()
modify (Int -> v -> IntMap v -> IntMap v
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
k v
v)
      MapStore f      -> (IntMap v -> IntMap v) -> Sem (State (IntMap v) : Counter : r) ()
forall s (r :: [(* -> *) -> * -> *]).
Member (State s) r =>
(s -> s) -> Sem r ()
modify ((v -> v) -> IntMap v -> IntMap v
forall a b. (a -> b) -> IntMap a -> IntMap b
IntMap.map v -> v
f)
      Store v (Sem rInitial) x
AssocsStore     -> (IntMap v -> [(Int, v)])
-> Sem (State (IntMap v) : Counter : r) [(Int, v)]
forall s a (r :: [(* -> *) -> * -> *]).
Member (State s) r =>
(s -> a) -> Sem r a
gets IntMap v -> [(Int, v)]
forall a. IntMap a -> [(Int, a)]
IntMap.assocs
      KeepKeys ks     -> (IntMap v -> IntMap v) -> Sem (State (IntMap v) : Counter : r) ()
forall s (r :: [(* -> *) -> * -> *]).
Member (State s) r =>
(s -> s) -> Sem r ()
modify (\IntMap v
m -> IntMap v -> IntSet -> IntMap v
forall a. IntMap a -> IntSet -> IntMap a
IntMap.withoutKeys IntMap v
m (IntMap v -> IntSet
forall a. IntMap a -> IntSet
IntMap.keysSet IntMap v
m IntSet -> IntSet -> IntSet
`IntSet.difference` IntSet
ks))