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

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

import Data.Data
import Data.Kind
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 GHC.TypeLits
import GHC.Word
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)

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

deriving via a instance Show a => Show (HasNoVariables a)
deriving via a instance Eq a => Eq (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
deriving via HasNoVariables Word8 instance HasVariables Word8
deriving via HasNoVariables Word16 instance HasVariables Word16
deriving via HasNoVariables Word32 instance HasVariables Word32
deriving via HasNoVariables Word64 instance HasVariables Word64

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

unsafeNextVarIndex :: VarContext -> Int
unsafeNextVarIndex :: VarContext -> Int
unsafeNextVarIndex (VarCtx Set (Any Var)
vs) = Int
1 forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Int
0 forall a. a -> [a] -> [a]
: [Int
i | Some (Var Int
i) <- forall a. Set a -> [a]
Set.toList Set (Any Var)
vs])

-- NOTE: This trick is taken from this blog post:
-- https://blog.csongor.co.uk/report-stuck-families/
data Dummy x
type family Break (c :: Constraint) (rep :: Type -> Type) :: Constraint where
  Break _ Dummy = ((), ())
  Break _ _ = ()

instance
  {-# OVERLAPPABLE #-}
  ( Break
      (TypeError ( 'Text "Missing instance of HasVariables for non-Generic type " ':<>: 'ShowType a))
      (Rep a)
  , Generic a
  , GenericHasVariables (Rep a)
  ) =>
  HasVariables a
  where
  getAllVariables :: 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

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