-- | Environment parameterised by functor @f@
--
-- Intended for qualified import:
--
-- > import Test.QuickCheck.StateModel.Lockstep.EnvF (EnvF)
-- > import Test.QuickCheck.StateModel.Lockstep.EnvF qualified as EnvF
module Test.QuickCheck.StateModel.Lockstep.EnvF (
    EnvF -- opaque
    -- * Construction
  , empty
  , insert
    -- * Query
  , lookup
  , keysOfType
  ) where

import Prelude hiding (lookup)

import Control.Monad
import Data.Foldable (asum)
import Data.Maybe (mapMaybe)
import Data.Typeable

import Test.QuickCheck.StateModel.Variables (Var, unsafeCoerceVar)

{-------------------------------------------------------------------------------
  Types
-------------------------------------------------------------------------------}

data EnvEntry f where
  EnvEntry :: Typeable a => Var a -> f a -> EnvEntry f

newtype EnvF f = EnvF [EnvEntry f]

{-------------------------------------------------------------------------------
  Construction
-------------------------------------------------------------------------------}

empty :: EnvF f
empty :: forall (f :: * -> *). EnvF f
empty = [EnvEntry f] -> EnvF f
forall (f :: * -> *). [EnvEntry f] -> EnvF f
EnvF []

insert :: Typeable a => Var a -> f a -> EnvF f -> EnvF f
insert :: forall a (f :: * -> *).
Typeable a =>
Var a -> f a -> EnvF f -> EnvF f
insert Var a
x f a
fa (EnvF [EnvEntry f]
env) = [EnvEntry f] -> EnvF f
forall (f :: * -> *). [EnvEntry f] -> EnvF f
EnvF (Var a -> f a -> EnvEntry f
forall a (f :: * -> *). Typeable a => Var a -> f a -> EnvEntry f
EnvEntry Var a
x f a
fa EnvEntry f -> [EnvEntry f] -> [EnvEntry f]
forall a. a -> [a] -> [a]
: [EnvEntry f]
env)

{-------------------------------------------------------------------------------
  Query
-------------------------------------------------------------------------------}

lookup :: forall f a. (Typeable f, Typeable a) => Var a -> EnvF f -> Maybe (f a)
lookup :: forall (f :: * -> *) a.
(Typeable f, Typeable a) =>
Var a -> EnvF f -> Maybe (f a)
lookup = \Var a
var (EnvF [EnvEntry f]
env) ->
    [Maybe (f a)] -> Maybe (f a)
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum ([Maybe (f a)] -> Maybe (f a)) -> [Maybe (f a)] -> Maybe (f a)
forall a b. (a -> b) -> a -> b
$ (EnvEntry f -> Maybe (f a)) -> [EnvEntry f] -> [Maybe (f a)]
forall a b. (a -> b) -> [a] -> [b]
map (\(EnvEntry Var a
var' f a
fa') -> Var a -> Var a -> f a -> Maybe (f a)
forall a'. Typeable a' => Var a -> Var a' -> f a' -> Maybe (f a)
aux Var a
var Var a
var' f a
fa') [EnvEntry f]
env
  where
    aux :: Typeable a' => Var a -> Var a' -> f a' -> Maybe (f a)
    aux :: forall a'. Typeable a' => Var a -> Var a' -> f a' -> Maybe (f a)
aux Var a
v1 Var a'
v2 f a'
fa' = do
        Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Var a
v1 Var a -> Var a -> Bool
forall a. Eq a => a -> a -> Bool
== Var a' -> Var a
forall a b. Var a -> Var b
unsafeCoerceVar Var a'
v2)
        f a' -> Maybe (f a)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast f a'
fa'

keysOfType :: Typeable a => EnvF f -> [Var a]
keysOfType :: forall a (f :: * -> *). Typeable a => EnvF f -> [Var a]
keysOfType (EnvF [EnvEntry f]
env) = (EnvEntry f -> Maybe (Var a)) -> [EnvEntry f] -> [Var a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(EnvEntry Var a
var f a
_) -> Var a -> Maybe (Var a)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Var a
var) [EnvEntry f]
env