-- | 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 (Var(..))

{-------------------------------------------------------------------------------
  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 = 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) = forall (f :: * -> *). [EnvEntry f] -> EnvF f
EnvF (forall a (f :: * -> *). Typeable a => Var a -> f a -> EnvEntry f
EnvEntry Var a
x f a
fa 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) ->
    forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (\(EnvEntry Var a
var' f a
fa') -> 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 Int
x) (Var Int
y) f a'
fa' = do
        forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
x forall a. Eq a => a -> a -> Bool
== Int
y)
        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) = forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(EnvEntry Var a
var f a
_) -> forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Var a
var) [EnvEntry f]
env