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