{-|
Copyright   : (C) 2021, QBayLogic B.V.
License     : BSD2 (see the file LICENSE)
Maintainer  : QBayLogic B.V. <devops@qbaylogic.com>

Random type-directed generation of data constructors.
-}

module Clash.Hedgehog.Core.DataCon
  ( genDataConsFrom
  ) where

import Control.Monad (replicateM, zipWithM)
import Control.Monad.Morph (hoist)
import Data.Either (partitionEithers)
import Data.Functor.Identity (Identity(runIdentity))
import Data.Text (Text)
import qualified Faker.Lorem as Fake
import Hedgehog (GenT, Range)
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Gen.Faker as Gen

import Clash.Core.DataCon
import Clash.Core.Name
import Clash.Core.TyCon
import Clash.Core.Type
import Clash.Core.TysPrim (liftedTypeKind)
import qualified Clash.Data.UniqMap as UniqMap

import Clash.Hedgehog.Core.Monad
import Clash.Hedgehog.Core.Name
import Clash.Hedgehog.Core.Type
import Clash.Hedgehog.Core.Var

-- | Generate a list of data constructors for a type. This biases towards
-- creating constructors which match some common form seen in code, such as
-- simple enums with no fields, or records.
--
genDataConsFrom
  :: forall m
   . (Alternative m, MonadGen m)
  => Range Int
  -- ^ The number of constructors to create for the data type
  -> TyConMap
  -- ^ The types already in scope when defining this type
  -> TyConName
  -- ^ The name of the @AlgTyCon@ the constructors belong to
  -> Kind
  -- ^ The kind of the @AlgTyCon@ the constructors belong to
  -> CoreGenT m [DataCon]
genDataConsFrom :: Range Int -> TyConMap -> TyConName -> Kind -> CoreGenT m [DataCon]
genDataConsFrom Range Int
range TyConMap
tcm TyConName
tcn Kind
kn = do
  -- We want to bias towards sometimes just having a single constructor. This
  -- is pretty common, e.g. for record types and non-GADT existential types.
  Int
numConstructors <- [CoreGenT m Int] -> CoreGenT m Int
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice [Int -> CoreGenT m Int
forall (m :: Type -> Type) a. MonadGen m => a -> m a
Gen.constant Int
1, Range Int -> CoreGenT m Int
forall (m :: Type -> Type). MonadGen m => Range Int -> m Int
Gen.int Range Int
range]
  [Name DataCon]
names <- Int -> CoreGenT m (Name DataCon) -> CoreGenT m [Name DataCon]
forall (m :: Type -> Type) a.
MonadGen m =>
Int -> m (Name a) -> m [Name a]
genNames Int
numConstructors CoreGenT m (Name DataCon)
forall (m :: Type -> Type). MonadGen m => m (Name DataCon)
genDataConName

  -- Universal tyvars are generated now so they can be shared between each
  -- data constructor. This matches what GHC would produce.
  let ([TyVar]
knTvs, [Kind]
knTys) = [Either TyVar Kind] -> ([TyVar], [Kind])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either TyVar Kind] -> ([TyVar], [Kind]))
-> [Either TyVar Kind] -> ([TyVar], [Kind])
forall a b. (a -> b) -> a -> b
$ ([Either TyVar Kind], Kind) -> [Either TyVar Kind]
forall a b. (a, b) -> a
fst (Kind -> ([Either TyVar Kind], Kind)
splitFunForallTy Kind
kn)
  [TyVar]
univTvs <- [TyVar] -> [TyVar] -> [TyVar]
forall a. Monoid a => a -> a -> a
mappend [TyVar]
knTvs ([TyVar] -> [TyVar]) -> CoreGenT m [TyVar] -> CoreGenT m [TyVar]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (Kind -> CoreGenT m (Name Kind) -> CoreGenT m TyVar)
-> [Kind] -> CoreGenT m (Name Kind) -> CoreGenT m [TyVar]
forall (m :: Type -> Type) a.
MonadGen m =>
(Kind -> m (Name a) -> m (Var a))
-> [Kind] -> m (Name a) -> m [Var a]
genVars Kind -> CoreGenT m (Name Kind) -> CoreGenT m TyVar
forall (m :: Type -> Type).
MonadGen m =>
Kind -> m (Name Kind) -> m TyVar
genTyVar [Kind]
knTys CoreGenT m (Name Kind)
forall (m :: Type -> Type) a. MonadGen m => m (Name a)
genVarName

  [CoreGenT m [DataCon]] -> CoreGenT m [DataCon]
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice
    [ TyConName -> [TyVar] -> [Name DataCon] -> CoreGenT m [DataCon]
forall (m :: Type -> Type).
Applicative m =>
TyConName -> [TyVar] -> [Name DataCon] -> CoreGenT m [DataCon]
genSimpleDataCons TyConName
tcn [TyVar]
univTvs [Name DataCon]
names
    , TyConMap
-> TyConName -> [TyVar] -> [Name DataCon] -> CoreGenT m [DataCon]
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap
-> TyConName -> [TyVar] -> [Name DataCon] -> CoreGenT m [DataCon]
genRecordDataCons TyConMap
tcm TyConName
tcn [TyVar]
univTvs [Name DataCon]
names
    , TyConMap
-> TyConName -> [TyVar] -> [Name DataCon] -> CoreGenT m [DataCon]
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap
-> TyConName -> [TyVar] -> [Name DataCon] -> CoreGenT m [DataCon]
genAnyDataCons TyConMap
tcm TyConName
tcn [TyVar]
univTvs [Name DataCon]
names
    ]

-- | Generate data constructors for a type
--
--   data D a1 a2 ... an = C1 | C2 | ... | CK
--
-- where every constructor is nullary, but the type constructor may have an
-- arbitrary number of phantom type parameters.
--
genSimpleDataCons
  :: forall m
   . Applicative m
  => TyConName
  -> [TyVar]
  -> [DcName]
  -> CoreGenT m [DataCon]
genSimpleDataCons :: TyConName -> [TyVar] -> [Name DataCon] -> CoreGenT m [DataCon]
genSimpleDataCons TyConName
tcn [TyVar]
univTvs =
  [DataCon] -> CoreGenT m [DataCon]
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ([DataCon] -> CoreGenT m [DataCon])
-> ([Name DataCon] -> [DataCon])
-> [Name DataCon]
-> CoreGenT m [DataCon]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Name DataCon -> DataCon)
-> [Int] -> [Name DataCon] -> [DataCon]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Name DataCon -> DataCon
go [Int
1..]
 where
  go :: ConTag -> DcName -> DataCon
  go :: Int -> Name DataCon -> DataCon
go Int
tag Name DataCon
name = MkData :: Name DataCon
-> Int
-> Int
-> Kind
-> [TyVar]
-> [TyVar]
-> [Kind]
-> [DcStrictness]
-> [Text]
-> DataCon
MkData
    { dcName :: Name DataCon
dcName = Name DataCon
name
    , dcUniq :: Int
dcUniq = Name DataCon -> Int
forall a. Name a -> Int
nameUniq Name DataCon
name
    , dcTag :: Int
dcTag = Int
tag
    , dcType :: Kind
dcType = TyConName -> [Kind] -> Kind
mkTyConApp TyConName
tcn ((TyVar -> Kind) -> [TyVar] -> [Kind]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap TyVar -> Kind
VarTy [TyVar]
univTvs)
    , dcUnivTyVars :: [TyVar]
dcUnivTyVars = [TyVar]
univTvs
    , dcExtTyVars :: [TyVar]
dcExtTyVars = []
    , dcArgTys :: [Kind]
dcArgTys = []
    , dcArgStrict :: [DcStrictness]
dcArgStrict = []
    , dcFieldLabels :: [Text]
dcFieldLabels = []
    }

-- | Generate data constructors for a type
--
--   data D a1 a2 ... an
--     = C1 { ... }
--     | C2 { ... }
--     | ...
--     | CK { ... }
--
-- where every constructor is either nullary, or a record.
--
genRecordDataCons
  :: forall m
   . (Alternative m, MonadGen m)
  => TyConMap
  -> TyConName
  -> [TyVar]
  -> [DcName]
  -> CoreGenT m [DataCon]
genRecordDataCons :: TyConMap
-> TyConName -> [TyVar] -> [Name DataCon] -> CoreGenT m [DataCon]
genRecordDataCons TyConMap
tcm TyConName
tcn [TyVar]
univTvs =
  (Int -> Name DataCon -> CoreGenT m DataCon)
-> [Int] -> [Name DataCon] -> CoreGenT m [DataCon]
forall (m :: Type -> Type) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Int -> Name DataCon -> CoreGenT m DataCon
go [Int
1..]
 where
  go :: ConTag -> DcName -> CoreGenT m DataCon
  go :: Int -> Name DataCon -> CoreGenT m DataCon
go Int
tag Name DataCon
name = do
    let resTy :: Kind
resTy = TyConName -> [Kind] -> Kind
mkTyConApp TyConName
tcn ((TyVar -> Kind) -> [TyVar] -> [Kind]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap TyVar -> Kind
VarTy [TyVar]
univTvs)
    let bound :: UniqMap TyVar
bound = [(TyVar, TyVar)] -> UniqMap TyVar
forall a b. Uniquable a => [(a, b)] -> UniqMap b
UniqMap.fromList ([TyVar] -> [TyVar] -> [(TyVar, TyVar)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TyVar]
univTvs [TyVar]
univTvs)
    let argGen :: CoreGenT m Kind
argGen = TyConMap -> UniqMap TyVar -> Kind -> CoreGenT m Kind
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap TyVar -> Kind -> CoreGenT m Kind
genMonoTypeFrom TyConMap
tcm UniqMap TyVar
bound Kind
liftedTypeKind -- TODO Make polymorphic
    Kind
ty <- Kind -> CoreGenT m Kind -> CoreGenT m Kind
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
Kind -> CoreGenT m Kind -> CoreGenT m Kind
genWithCodomain Kind
resTy CoreGenT m Kind
argGen

    let argTys :: [Kind]
argTys = case [Either TyVar Kind] -> ([TyVar], [Kind])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either TyVar Kind] -> ([TyVar], [Kind]))
-> [Either TyVar Kind] -> ([TyVar], [Kind])
forall a b. (a -> b) -> a -> b
$ ([Either TyVar Kind], Kind) -> [Either TyVar Kind]
forall a b. (a, b) -> a
fst (Kind -> ([Either TyVar Kind], Kind)
splitFunForallTy Kind
ty) of
          ([],[Kind]
_) -> [Char] -> [Kind]
forall a. HasCallStack => [Char] -> a
error [Char]
"getMonoTypeFrom is wrong, there are type variables"
          ([TyVar]
_,[Kind]
vs) -> [Kind]
vs
    [DcStrictness]
bangs <- (Kind -> CoreGenT m DcStrictness)
-> [Kind] -> CoreGenT m [DcStrictness]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (TyConMap -> Kind -> CoreGenT m DcStrictness
forall (m :: Type -> Type).
MonadGen m =>
TyConMap -> Kind -> m DcStrictness
genStrictness TyConMap
tcm) [Kind]
argTys
    [Text]
fields <- Int -> CoreGenT m Text -> CoreGenT m [Text]
forall (m :: Type -> Type) a. Applicative m => Int -> m a -> m [a]
replicateM ([Kind] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Kind]
argTys) CoreGenT m Text
forall (m :: Type -> Type). MonadGen m => m Text
genFieldLabel

    DataCon -> CoreGenT m DataCon
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MkData :: Name DataCon
-> Int
-> Int
-> Kind
-> [TyVar]
-> [TyVar]
-> [Kind]
-> [DcStrictness]
-> [Text]
-> DataCon
MkData
      { dcName :: Name DataCon
dcName = Name DataCon
name
      , dcUniq :: Int
dcUniq = Name DataCon -> Int
forall a. Name a -> Int
nameUniq Name DataCon
name
      , dcTag :: Int
dcTag = Int
tag
      , dcType :: Kind
dcType = Kind
ty
      , dcUnivTyVars :: [TyVar]
dcUnivTyVars = [TyVar]
univTvs
      , dcExtTyVars :: [TyVar]
dcExtTyVars = []
      , dcArgTys :: [Kind]
dcArgTys = [Kind]
argTys
      , dcArgStrict :: [DcStrictness]
dcArgStrict = [DcStrictness]
bangs
      , dcFieldLabels :: [Text]
dcFieldLabels = [Text]
fields
      }

-- | Generate data constructors for a type which does not match any common
-- idiom. Since this can generate any possible data constructor, it can
-- sometimes produce less representative results.
genAnyDataCons
  :: forall m
   . (Alternative m, MonadGen m)
  => TyConMap
  -> TyConName
  -> [TyVar]
  -> [DcName]
  -> CoreGenT m [DataCon]
genAnyDataCons :: TyConMap
-> TyConName -> [TyVar] -> [Name DataCon] -> CoreGenT m [DataCon]
genAnyDataCons TyConMap
tcm TyConName
tcn [TyVar]
univTvs =
  (Int -> Name DataCon -> CoreGenT m DataCon)
-> [Int] -> [Name DataCon] -> CoreGenT m [DataCon]
forall (m :: Type -> Type) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Int -> Name DataCon -> CoreGenT m DataCon
go [Int
1..]
 where
  go :: ConTag -> DcName -> CoreGenT m DataCon
  go :: Int -> Name DataCon -> CoreGenT m DataCon
go Int
tag Name DataCon
name = do
    let resTy :: Kind
resTy = TyConName -> [Kind] -> Kind
mkTyConApp TyConName
tcn ((TyVar -> Kind) -> [TyVar] -> [Kind]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap TyVar -> Kind
VarTy [TyVar]
univTvs)
    let bound :: UniqMap TyVar
bound = [(TyVar, TyVar)] -> UniqMap TyVar
forall a b. Uniquable a => [(a, b)] -> UniqMap b
UniqMap.fromList ([TyVar] -> [TyVar] -> [(TyVar, TyVar)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TyVar]
univTvs [TyVar]
univTvs)
    let argGen :: CoreGenT m Kind
argGen = TyConMap -> UniqMap TyVar -> Kind -> CoreGenT m Kind
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap TyVar -> Kind -> CoreGenT m Kind
genMonoTypeFrom TyConMap
tcm UniqMap TyVar
bound Kind
liftedTypeKind -- TODO Make polymorphic.
    Kind
ty <- Kind -> CoreGenT m Kind -> CoreGenT m Kind
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
Kind -> CoreGenT m Kind -> CoreGenT m Kind
genWithCodomain Kind
resTy CoreGenT m Kind
argGen

    -- Determine the argument types from the data constructor type
    -- Generate strictness and field labels from the argument types
    let ([TyVar]
extTvs, [Kind]
argTys) = [Either TyVar Kind] -> ([TyVar], [Kind])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either TyVar Kind] -> ([TyVar], [Kind]))
-> [Either TyVar Kind] -> ([TyVar], [Kind])
forall a b. (a -> b) -> a -> b
$ ([Either TyVar Kind], Kind) -> [Either TyVar Kind]
forall a b. (a, b) -> a
fst (Kind -> ([Either TyVar Kind], Kind)
splitFunForallTy Kind
ty)
    [DcStrictness]
bangs <- (Kind -> CoreGenT m DcStrictness)
-> [Kind] -> CoreGenT m [DcStrictness]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (TyConMap -> Kind -> CoreGenT m DcStrictness
forall (m :: Type -> Type).
MonadGen m =>
TyConMap -> Kind -> m DcStrictness
genStrictness TyConMap
tcm) [Kind]
argTys

    DataCon -> CoreGenT m DataCon
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MkData :: Name DataCon
-> Int
-> Int
-> Kind
-> [TyVar]
-> [TyVar]
-> [Kind]
-> [DcStrictness]
-> [Text]
-> DataCon
MkData
      { dcName :: Name DataCon
dcName = Name DataCon
name
      , dcUniq :: Int
dcUniq = Name DataCon -> Int
forall a. Name a -> Int
nameUniq Name DataCon
name
      , dcTag :: Int
dcTag = Int
tag
      , dcType :: Kind
dcType = Kind
ty
      , dcUnivTyVars :: [TyVar]
dcUnivTyVars = [TyVar]
univTvs
      , dcExtTyVars :: [TyVar]
dcExtTyVars = [TyVar]
extTvs
      , dcArgTys :: [Kind]
dcArgTys = [Kind]
argTys
      , dcArgStrict :: [DcStrictness]
dcArgStrict = [DcStrictness]
bangs
      , dcFieldLabels :: [Text]
dcFieldLabels = []
      }

-- TODO genGadt, which can insert ~# arguments after the existential type
-- variables are introduced. I may also want a `genConstraints` in
-- Clash.Hedgehog.Core.Type to generate any constraints for a type.

-- | Generate strictness annotations for data constructor arguments. This
-- ensures that any types which are always strict, e.g. Int#, are strict and
-- types which may be lazy have a random strictness assigned.
--
-- This generator shrinks towards choosing lazy by default for types where it
-- is possible.
genStrictness
  :: forall m. MonadGen m => TyConMap -> Kind -> m DcStrictness
genStrictness :: TyConMap -> Kind -> m DcStrictness
genStrictness TyConMap
tcm Kind
kn
  -- Assume that any primitive type constructor is always strict. This may
  -- overapproximate strictness, as it means Type, Nat and Symbol are strict.
  | TyConApp TyConName
tc [] <- Kind -> TypeView
tyView Kind
kn
  , Just PrimTyCon{} <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
tc TyConMap
tcm
  = DcStrictness -> m DcStrictness
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure DcStrictness
Strict

  -- Shrink towards laziness as this is the default in Haskell (assuming no
  -- extensions like -XStrict or -XStrictData are enabled).
  | Bool
otherwise
  = [DcStrictness] -> m DcStrictness
forall (m :: Type -> Type) a. MonadGen m => [a] -> m a
Gen.element [DcStrictness
Lazy, DcStrictness
Strict]

-- | Generate a field label for use in a record.
genFieldLabel :: forall m. MonadGen m => m Text
genFieldLabel :: m Text
genFieldLabel =
  GenT (GenBase m) Text -> m Text
forall (m :: Type -> Type) a.
MonadGen m =>
GenT (GenBase m) a -> m a
fromGenT (GenT (GenBase m) Text -> m Text)
-> GenT (GenBase m) Text -> m Text
forall a b. (a -> b) -> a -> b
$ (forall a. Identity a -> GenBase m a)
-> GenT Identity Text -> GenT (GenBase m) Text
forall k (t :: (Type -> Type) -> k -> Type) (m :: Type -> Type)
       (n :: Type -> Type) (b :: k).
(MFunctor t, Monad m) =>
(forall a. m a -> n a) -> t m b -> t n b
hoist @GenT @Identity @(GenBase m)
    (a -> GenBase m a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (a -> GenBase m a)
-> (Identity a -> a) -> Identity a -> GenBase m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Identity a -> a
forall a. Identity a -> a
runIdentity)
    (Fake Text -> GenT Identity Text
forall a. Fake a -> Gen a
Gen.fake Fake Text
Fake.words)