Store.hs

{-# 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))