{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE UndecidableInstances #-}

module Test.QuickCheck.StateModel.Variables (
  Var,
  Any (..),
  HasVariables (..),
  HasNoVariables,
  VarContext,
  mkVar,
  ctxAtType,
  arbitraryVar,
  shrinkVar,
  extendContext,
  isWellTyped,
  allVariables,
  unsafeCoerceVar,
) where

import Data.Data
import Data.List
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Ord
import Data.Set (Set)
import Data.Set qualified as Set
import GHC.Generics
import Test.QuickCheck as QC

-- | A symbolic variable for a value of type `a`
newtype Var a = Var Int
  deriving (Var a -> Var a -> Bool
forall a. Var a -> Var a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Var a -> Var a -> Bool
$c/= :: forall a. Var a -> Var a -> Bool
== :: Var a -> Var a -> Bool
$c== :: forall a. Var a -> Var a -> Bool
Eq, Var a -> Var a -> Bool
Var a -> Var a -> Ordering
Var a -> Var a -> Var a
forall a. Eq (Var a)
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Var a -> Var a -> Bool
forall a. Var a -> Var a -> Ordering
forall a. Var a -> Var a -> Var a
min :: Var a -> Var a -> Var a
$cmin :: forall a. Var a -> Var a -> Var a
max :: Var a -> Var a -> Var a
$cmax :: forall a. Var a -> Var a -> Var a
>= :: Var a -> Var a -> Bool
$c>= :: forall a. Var a -> Var a -> Bool
> :: Var a -> Var a -> Bool
$c> :: forall a. Var a -> Var a -> Bool
<= :: Var a -> Var a -> Bool
$c<= :: forall a. Var a -> Var a -> Bool
< :: Var a -> Var a -> Bool
$c< :: forall a. Var a -> Var a -> Bool
compare :: Var a -> Var a -> Ordering
$ccompare :: forall a. Var a -> Var a -> Ordering
Ord, Typeable, Var a -> DataType
Var a -> Constr
forall {a}. Data a => Typeable (Var a)
forall a. Data a => Var a -> DataType
forall a. Data a => Var a -> Constr
forall a. Data a => (forall b. Data b => b -> b) -> Var a -> Var a
forall a u.
Data a =>
Int -> (forall d. Data d => d -> u) -> Var a -> u
forall a u. Data a => (forall d. Data d => d -> u) -> Var a -> [u]
forall a r r'.
Data a =>
(r -> r' -> r) -> r -> (forall d. Data d => d -> r') -> Var a -> r
forall a r r'.
Data a =>
(r' -> r -> r) -> r -> (forall d. Data d => d -> r') -> Var a -> r
forall a (m :: * -> *).
(Data a, Monad m) =>
(forall d. Data d => d -> m d) -> Var a -> m (Var a)
forall a (m :: * -> *).
(Data a, MonadPlus m) =>
(forall d. Data d => d -> m d) -> Var a -> m (Var a)
forall a (c :: * -> *).
Data a =>
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Var a)
forall a (c :: * -> *).
Data a =>
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Var a -> c (Var a)
forall a (t :: * -> *) (c :: * -> *).
(Data a, Typeable t) =>
(forall d. Data d => c (t d)) -> Maybe (c (Var a))
forall a (t :: * -> * -> *) (c :: * -> *).
(Data a, Typeable t) =>
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Var a))
forall a.
Typeable a
-> (forall (c :: * -> *).
    (forall d b. Data d => c (d -> b) -> d -> c b)
    -> (forall g. g -> c g) -> a -> c a)
-> (forall (c :: * -> *).
    (forall b r. Data b => c (b -> r) -> c r)
    -> (forall r. r -> c r) -> Constr -> c a)
-> (a -> Constr)
-> (a -> DataType)
-> (forall (t :: * -> *) (c :: * -> *).
    Typeable t =>
    (forall d. Data d => c (t d)) -> Maybe (c a))
-> (forall (t :: * -> * -> *) (c :: * -> *).
    Typeable t =>
    (forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c a))
-> ((forall b. Data b => b -> b) -> a -> a)
-> (forall r r'.
    (r -> r' -> r) -> r -> (forall d. Data d => d -> r') -> a -> r)
-> (forall r r'.
    (r' -> r -> r) -> r -> (forall d. Data d => d -> r') -> a -> r)
-> (forall u. (forall d. Data d => d -> u) -> a -> [u])
-> (forall u. Int -> (forall d. Data d => d -> u) -> a -> u)
-> (forall (m :: * -> *).
    Monad m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> Data a
forall (c :: * -> *).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Var a)
forall (c :: * -> *).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Var a -> c (Var a)
forall (t :: * -> *) (c :: * -> *).
Typeable t =>
(forall d. Data d => c (t d)) -> Maybe (c (Var a))
gmapMo :: forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d) -> Var a -> m (Var a)
$cgmapMo :: forall a (m :: * -> *).
(Data a, MonadPlus m) =>
(forall d. Data d => d -> m d) -> Var a -> m (Var a)
gmapMp :: forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d) -> Var a -> m (Var a)
$cgmapMp :: forall a (m :: * -> *).
(Data a, MonadPlus m) =>
(forall d. Data d => d -> m d) -> Var a -> m (Var a)
gmapM :: forall (m :: * -> *).
Monad m =>
(forall d. Data d => d -> m d) -> Var a -> m (Var a)
$cgmapM :: forall a (m :: * -> *).
(Data a, Monad m) =>
(forall d. Data d => d -> m d) -> Var a -> m (Var a)
gmapQi :: forall u. Int -> (forall d. Data d => d -> u) -> Var a -> u
$cgmapQi :: forall a u.
Data a =>
Int -> (forall d. Data d => d -> u) -> Var a -> u
gmapQ :: forall u. (forall d. Data d => d -> u) -> Var a -> [u]
$cgmapQ :: forall a u. Data a => (forall d. Data d => d -> u) -> Var a -> [u]
gmapQr :: forall r r'.
(r' -> r -> r) -> r -> (forall d. Data d => d -> r') -> Var a -> r
$cgmapQr :: forall a r r'.
Data a =>
(r' -> r -> r) -> r -> (forall d. Data d => d -> r') -> Var a -> r
gmapQl :: forall r r'.
(r -> r' -> r) -> r -> (forall d. Data d => d -> r') -> Var a -> r
$cgmapQl :: forall a r r'.
Data a =>
(r -> r' -> r) -> r -> (forall d. Data d => d -> r') -> Var a -> r
gmapT :: (forall b. Data b => b -> b) -> Var a -> Var a
$cgmapT :: forall a. Data a => (forall b. Data b => b -> b) -> Var a -> Var a
dataCast2 :: forall (t :: * -> * -> *) (c :: * -> *).
Typeable t =>
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Var a))
$cdataCast2 :: forall a (t :: * -> * -> *) (c :: * -> *).
(Data a, Typeable t) =>
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Var a))
dataCast1 :: forall (t :: * -> *) (c :: * -> *).
Typeable t =>
(forall d. Data d => c (t d)) -> Maybe (c (Var a))
$cdataCast1 :: forall a (t :: * -> *) (c :: * -> *).
(Data a, Typeable t) =>
(forall d. Data d => c (t d)) -> Maybe (c (Var a))
dataTypeOf :: Var a -> DataType
$cdataTypeOf :: forall a. Data a => Var a -> DataType
toConstr :: Var a -> Constr
$ctoConstr :: forall a. Data a => Var a -> Constr
gunfold :: forall (c :: * -> *).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Var a)
$cgunfold :: forall a (c :: * -> *).
Data a =>
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Var a)
gfoldl :: forall (c :: * -> *).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Var a -> c (Var a)
$cgfoldl :: forall a (c :: * -> *).
Data a =>
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Var a -> c (Var a)
Data)

mkVar :: Int -> Var a
mkVar :: forall a. Int -> Var a
mkVar = forall a. Int -> Var a
Var

instance Show (Var a) where
  show :: Var a -> String
show (Var Int
i) = String
"var" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
i

-- | This type class gives you a way to get all the symbolic variables that
-- appear in a value.
class HasVariables a where
  getAllVariables :: a -> Set (Any Var)
  default getAllVariables :: (Generic a, GenericHasVariables (Rep a)) => a -> Set (Any Var)
  getAllVariables = forall (f :: * -> *) k.
GenericHasVariables f =>
f k -> Set (Any Var)
genericGetAllVariables forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a x. Generic a => a -> Rep a x
from

instance HasVariables a => HasVariables (Smart a) where
  getAllVariables :: Smart a -> Set (Any Var)
getAllVariables (Smart Int
_ a
a) = forall a. HasVariables a => a -> Set (Any Var)
getAllVariables a
a

instance Typeable a => HasVariables (Var a) where
  getAllVariables :: Var a -> Set (Any Var)
getAllVariables = forall a. a -> Set a
Set.singleton forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a (f :: * -> *). (Typeable a, Eq (f a)) => f a -> Any f
Some

instance (HasVariables k, HasVariables v) => HasVariables (Map k v) where
  getAllVariables :: Map k v -> Set (Any Var)
getAllVariables = forall a. HasVariables a => a -> Set (Any Var)
getAllVariables forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
Map.toList

instance HasVariables a => HasVariables (Set a) where
  getAllVariables :: Set a -> Set (Any Var)
getAllVariables = forall a. HasVariables a => a -> Set (Any Var)
getAllVariables forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Set a -> [a]
Set.toList

newtype HasNoVariables a = HasNoVariables a

instance HasVariables (HasNoVariables a) where
  getAllVariables :: HasNoVariables a -> Set (Any Var)
getAllVariables HasNoVariables a
_ = forall a. Monoid a => a
mempty

deriving via HasNoVariables Integer instance HasVariables Integer
deriving via HasNoVariables Int instance HasVariables Int
deriving via HasNoVariables Char instance HasVariables Char

data Any f where
  Some :: (Typeable a, Eq (f a)) => f a -> Any f

instance Eq (Any f) where
  Some (f a
a :: f a) == :: Any f -> Any f -> Bool
== Some (f a
b :: f b) =
    case forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT @a @b of
      Just a :~: a
Refl -> f a
a forall a. Eq a => a -> a -> Bool
== f a
b
      Maybe (a :~: a)
Nothing -> Bool
False

instance (forall a. Ord (f a)) => Ord (Any f) where
  compare :: Any f -> Any f -> Ordering
compare (Some (f a
a :: f a)) (Some (f a
a' :: f a')) =
    case forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT @a @a' of
      Just a :~: a
Refl -> forall a. Ord a => a -> a -> Ordering
compare f a
a f a
a'
      Maybe (a :~: a)
Nothing -> forall a. Ord a => a -> a -> Ordering
compare (forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep f a
a) (forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep f a
a')

newtype VarContext = VarCtx (Set (Any Var))
  deriving (NonEmpty VarContext -> VarContext
VarContext -> VarContext -> VarContext
forall b. Integral b => b -> VarContext -> VarContext
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: forall b. Integral b => b -> VarContext -> VarContext
$cstimes :: forall b. Integral b => b -> VarContext -> VarContext
sconcat :: NonEmpty VarContext -> VarContext
$csconcat :: NonEmpty VarContext -> VarContext
<> :: VarContext -> VarContext -> VarContext
$c<> :: VarContext -> VarContext -> VarContext
Semigroup, Semigroup VarContext
VarContext
[VarContext] -> VarContext
VarContext -> VarContext -> VarContext
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
mconcat :: [VarContext] -> VarContext
$cmconcat :: [VarContext] -> VarContext
mappend :: VarContext -> VarContext -> VarContext
$cmappend :: VarContext -> VarContext -> VarContext
mempty :: VarContext
$cmempty :: VarContext
Monoid) via Set (Any Var)

instance Show VarContext where
  show :: VarContext -> String
show (VarCtx Set (Any Var)
vs) =
    String
"[" forall a. [a] -> [a] -> [a]
++ forall a. [a] -> [[a]] -> [a]
intercalate String
", " (forall a b. (a -> b) -> [a] -> [b]
map Any Var -> String
showBinding forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing Any Var -> Int
getIdx) forall a b. (a -> b) -> a -> b
$ forall a. Set a -> [a]
Set.toList Set (Any Var)
vs) forall a. [a] -> [a] -> [a]
++ String
"]"
    where
      getIdx :: Any Var -> Int
getIdx (Some (Var Int
i)) = Int
i
      showBinding :: Any Var -> String
      -- The use of typeRep here is on purpose to avoid printing `Var` unnecessarily.
      showBinding :: Any Var -> String
showBinding (Some Var a
v) = forall a. Show a => a -> String
show Var a
v forall a. [a] -> [a] -> [a]
++ String
" :: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep Var a
v)

isWellTyped :: Typeable a => Var a -> VarContext -> Bool
isWellTyped :: forall a. Typeable a => Var a -> VarContext -> Bool
isWellTyped Var a
v (VarCtx Set (Any Var)
ctx) = forall a (f :: * -> *). (Typeable a, Eq (f a)) => f a -> Any f
Some Var a
v forall a. Ord a => a -> Set a -> Bool
`Set.member` Set (Any Var)
ctx

-- TODO: check the invariant that no variable index is used
-- twice at different types. This is generally not an issue
-- because lookups take types into account (so it *shouldn't*
-- cause an issue, but it might be good practise to crash
-- if the invariant is violated anyway as it is evidence that
-- something is horribly broken at the use site).
extendContext :: Typeable a => VarContext -> Var a -> VarContext
extendContext :: forall a. Typeable a => VarContext -> Var a -> VarContext
extendContext (VarCtx Set (Any Var)
ctx) Var a
v = Set (Any Var) -> VarContext
VarCtx forall a b. (a -> b) -> a -> b
$ forall a. Ord a => a -> Set a -> Set a
Set.insert (forall a (f :: * -> *). (Typeable a, Eq (f a)) => f a -> Any f
Some Var a
v) Set (Any Var)
ctx

allVariables :: HasVariables a => a -> VarContext
allVariables :: forall a. HasVariables a => a -> VarContext
allVariables = Set (Any Var) -> VarContext
VarCtx forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. HasVariables a => a -> Set (Any Var)
getAllVariables

ctxAtType :: Typeable a => VarContext -> [Var a]
ctxAtType :: forall a. Typeable a => VarContext -> [Var a]
ctxAtType (VarCtx Set (Any Var)
vs) = [Var a
v | Some (forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast -> Just Var a
v) <- forall a. Set a -> [a]
Set.toList Set (Any Var)
vs]

arbitraryVar :: Typeable a => VarContext -> Gen (Var a)
arbitraryVar :: forall a. Typeable a => VarContext -> Gen (Var a)
arbitraryVar = forall a. [a] -> Gen a
elements forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Typeable a => VarContext -> [Var a]
ctxAtType

shrinkVar :: Typeable a => VarContext -> Var a -> [Var a]
shrinkVar :: forall a. Typeable a => VarContext -> Var a -> [Var a]
shrinkVar VarContext
ctx Var a
v = forall a. (a -> Bool) -> [a] -> [a]
filter (forall a. Ord a => a -> a -> Bool
< Var a
v) forall a b. (a -> b) -> a -> b
$ forall a. Typeable a => VarContext -> [Var a]
ctxAtType VarContext
ctx

unsafeCoerceVar :: Var a -> Var b
unsafeCoerceVar :: forall a b. Var a -> Var b
unsafeCoerceVar (Var Int
i) = forall a. Int -> Var a
Var Int
i

instance {-# OVERLAPPABLE #-} (Generic a, GenericHasVariables (Rep a)) => HasVariables a

class GenericHasVariables f where
  genericGetAllVariables :: f k -> Set (Any Var)

instance GenericHasVariables f => GenericHasVariables (M1 i c f) where
  genericGetAllVariables :: forall k. M1 i c f k -> Set (Any Var)
genericGetAllVariables = forall (f :: * -> *) k.
GenericHasVariables f =>
f k -> Set (Any Var)
genericGetAllVariables forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k i (c :: Meta) (f :: k -> *) (p :: k). M1 i c f p -> f p
unM1

instance HasVariables c => GenericHasVariables (K1 i c) where
  genericGetAllVariables :: forall k. K1 i c k -> Set (Any Var)
genericGetAllVariables = forall a. HasVariables a => a -> Set (Any Var)
getAllVariables forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k i c (p :: k). K1 i c p -> c
unK1

instance GenericHasVariables U1 where
  genericGetAllVariables :: forall k. U1 k -> Set (Any Var)
genericGetAllVariables U1 k
_ = forall a. Monoid a => a
mempty

instance (GenericHasVariables f, GenericHasVariables g) => GenericHasVariables (f :*: g) where
  genericGetAllVariables :: forall k. (:*:) f g k -> Set (Any Var)
genericGetAllVariables (f k
x :*: g k
y) = forall (f :: * -> *) k.
GenericHasVariables f =>
f k -> Set (Any Var)
genericGetAllVariables f k
x forall a. Semigroup a => a -> a -> a
<> forall (f :: * -> *) k.
GenericHasVariables f =>
f k -> Set (Any Var)
genericGetAllVariables g k
y

instance (GenericHasVariables f, GenericHasVariables g) => GenericHasVariables (f :+: g) where
  genericGetAllVariables :: forall k. (:+:) f g k -> Set (Any Var)
genericGetAllVariables (L1 f k
x) = forall (f :: * -> *) k.
GenericHasVariables f =>
f k -> Set (Any Var)
genericGetAllVariables f k
x
  genericGetAllVariables (R1 g k
x) = forall (f :: * -> *) k.
GenericHasVariables f =>
f k -> Set (Any Var)
genericGetAllVariables g k
x