content: {-# 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 =
    runCounter
      . evalState @(IntMap.IntMap v) IntMap.empty
      . reinterpret2 \case
        ClearStore -> put IntMap.empty
        New v -> do
          loc <- fromIntegral <$> next
          modify $ IntMap.insert loc v
          return loc
        LookupStore k -> gets (IntMap.lookup k)
        InsertStore k v -> modify (IntMap.insert k v)
        MapStore f -> modify (IntMap.map f)
        AssocsStore -> gets IntMap.assocs
        KeepKeys ks -> modify (\m -> IntMap.withoutKeys m (IntMap.keysSet m `IntSet.difference` ks))
  