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

Random generation of type constructors.
-}

{-# LANGUAGE LambdaCase #-}

module Clash.Hedgehog.Core.TyCon
  ( genTyConMap
  ) where

import Control.Monad (forM)
import Data.Coerce (coerce)
import Data.Either (rights)
import Hedgehog (Range)
import qualified Hedgehog.Gen as Gen

import Clash.Core.DataCon
import Clash.Core.HasType
import Clash.Core.Name (nameUniq)
import Clash.Core.Subst
import Clash.Core.TyCon
import Clash.Core.Type (Kind, Type(VarTy), mkTyConApp, splitFunForallTy)
import Clash.Core.TysPrim (liftedTypeKind, tysPrimMap)
import Clash.Core.Var
import Clash.Core.VarEnv
import Clash.Data.UniqMap (UniqMap)
import qualified Clash.Data.UniqMap as UniqMap

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

{-
Note [order of generation]
~~~~~~~~~~~~~~~~~~~~~~~~~~
In Clash core (as with GHC core), there is a degree of circularity:

  * A type can be a type constructor (referenced by name)
  * A type constructor has a type (kind)
  * A data constructor has a type
  * A type constructor may contain data constructors

This makes it impossible to naively write a generator for these parts of the
core IR, as such a generator may never terminate. However not everything
generated is completely random, e.g. the codomain of a data constructor is
always the type constructor the data constructor belongs to.

A reasonable approximation of a real program can be made by first generating
an environment, then generating things which exist within that environment. To
that end, the first thing (typically) to generate is a TyConMap, which contains
information about all type constructors that exist. This leads to a sensible
order of generation, e.g.

  * generate a TyConMap, `tcm`
  * generate a kind that exists in `tcm`, `k`
  * generate a type of kind `k` that exists in `tcm`, `a`
  * generate a term of type `a` that exists in `tcm`

By generating the TyConMap first, the generation of complete programs becomes
a more manageable process of running increasingly more constrained generators.
This helps ensure that generated data is well-formed, as a constrained
generator produces random hole-fits instead of completely arbitrary values.
-}

arityOf :: Kind -> Int
arityOf :: Kind -> Int
arityOf = [Either TyVar Kind] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length ([Either TyVar Kind] -> Int)
-> (Kind -> [Either TyVar Kind]) -> Kind -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Either TyVar Kind], Kind) -> [Either TyVar Kind]
forall a b. (a, b) -> a
fst (([Either TyVar Kind], Kind) -> [Either TyVar Kind])
-> (Kind -> ([Either TyVar Kind], Kind))
-> Kind
-> [Either TyVar Kind]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Kind -> ([Either TyVar Kind], Kind)
splitFunForallTy

-- | A TyConMap contains all the algebraic data types and type families that
-- are used in a program. This is typically the first thing that should be
-- generated, as calls to other generators like @genKind@ or @genTypeFrom@ will
-- likely want to use the type constructors added to the TyConMap.
--
-- TODO It would be nice if this also included types from @clash-prelude@ like
-- Signal and the sized number types. Maybe we want to hook into @clash-ghc@
-- to load type constructors and primitives from @Clash.Prelude@.
--
genTyConMap
  :: forall m
   . (Alternative m, MonadGen m)
  => Range Int
  -> CoreGenT m TyConMap
genTyConMap :: Range Int -> CoreGenT m TyConMap
genTyConMap Range Int
numDcs = TyConMap -> CoreGenT m TyConMap
go TyConMap
tysPrimMap
 where
  -- Either stop adding new items to the TyConMap, or generate a new item.
  -- 'Gen.recursive' is necessary to ensure termination.
  go :: TyConMap -> CoreGenT m TyConMap
go TyConMap
tcm =
    ([CoreGenT m TyConMap] -> CoreGenT m TyConMap)
-> [CoreGenT m TyConMap]
-> [CoreGenT m TyConMap]
-> CoreGenT m TyConMap
forall (m :: Type -> Type) a.
MonadGen m =>
([m a] -> m a) -> [m a] -> [m a] -> m a
Gen.recursive [CoreGenT m TyConMap] -> CoreGenT m TyConMap
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice
      [TyConMap -> CoreGenT m TyConMap
forall (m :: Type -> Type) a. MonadGen m => a -> m a
Gen.constant TyConMap
tcm]
      [CoreGenT m TyConMap
-> (TyConMap -> CoreGenT m TyConMap) -> CoreGenT m TyConMap
forall (m :: Type -> Type) a.
MonadGen m =>
m a -> (a -> m a) -> m a
Gen.subtermM (TyConMap -> CoreGenT m TyConMap
extendTyConMap TyConMap
tcm) TyConMap -> CoreGenT m TyConMap
go]

  extendTyConMap :: TyConMap -> CoreGenT m TyConMap
extendTyConMap TyConMap
tcm = do
    -- We return new UniqMap instead of individual TyCon, because for AlgTyCon
    -- we may also generate PromotedDataCon for -XDataKinds.
    TyConMap
new <- CoreGenT m Bool
forall (m :: Type -> Type). Monad m => CoreGenT m Bool
canGenTypeFamilies CoreGenT m Bool
-> (Bool -> CoreGenT m TyConMap) -> CoreGenT m TyConMap
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Bool
True -> [CoreGenT m TyConMap] -> CoreGenT m TyConMap
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice
        [ Range Int -> TyConMap -> CoreGenT m TyConMap
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
Range Int -> TyConMap -> CoreGenT m TyConMap
genAlgTyConFrom Range Int
numDcs TyConMap
tcm
        , TyConMap -> CoreGenT m TyConMap
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> CoreGenT m TyConMap
genFunTyConFrom TyConMap
tcm CoreGenT m TyConMap -> CoreGenT m TyConMap -> CoreGenT m TyConMap
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> Range Int -> TyConMap -> CoreGenT m TyConMap
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
Range Int -> TyConMap -> CoreGenT m TyConMap
genAlgTyConFrom Range Int
numDcs TyConMap
tcm
        ]
      Bool
False -> [CoreGenT m TyConMap] -> CoreGenT m TyConMap
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice [Range Int -> TyConMap -> CoreGenT m TyConMap
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
Range Int -> TyConMap -> CoreGenT m TyConMap
genAlgTyConFrom Range Int
numDcs TyConMap
tcm]

    TyConMap -> CoreGenT m TyConMap
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (TyConMap
tcm TyConMap -> TyConMap -> TyConMap
forall a. Semigroup a => a -> a -> a
<> TyConMap
new)

-- | Generate a new algebraic type constructor using the types that are already
-- in scope. This will also promote data constructors if the configuration
-- supports @-XDataKinds@.
--
genAlgTyConFrom
  :: forall m
   . (Alternative m, MonadGen m)
  => Range Int
  -> TyConMap
  -> CoreGenT m TyConMap
genAlgTyConFrom :: Range Int -> TyConMap -> CoreGenT m TyConMap
genAlgTyConFrom Range Int
range TyConMap
tcm = do
  let used :: UniqMap Int
used = (TyCon -> Int) -> TyConMap -> UniqMap Int
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap TyCon -> Int
tyConUniq TyConMap
tcm
  Name TyCon
name <- UniqMap Int -> CoreGenT m (Name TyCon) -> CoreGenT m (Name TyCon)
forall (m :: Type -> Type) a b.
MonadGen m =>
UniqMap b -> m (Name a) -> m (Name a)
genFreshName UniqMap Int
used CoreGenT m (Name TyCon)
forall (m :: Type -> Type). MonadGen m => m (Name TyCon)
genTyConName

  -- TODO We want to use this, but we cannot sample polymorphic constructors
  -- when making Term / Type, so we avoid generating polymorphic data now.
  --
  -- See NOTE [finding more complex fits] in Clash.Hedgehog.Unique.
  --
  -- let argGen = genClosedKindFrom tcm liftedTypeKind
  -- kn  <- genWithCodomain liftedTypeKind argGen

  -- All ADTs live in the kind Type
  let kn :: Kind
kn = Kind
liftedTypeKind
  let arity :: Int
arity = Kind -> Int
arityOf Kind
kn

  AlgTyConRhs
rhs <- [CoreGenT m AlgTyConRhs] -> CoreGenT m AlgTyConRhs
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice
           [ [DataCon] -> AlgTyConRhs
DataTyCon ([DataCon] -> AlgTyConRhs)
-> CoreGenT m [DataCon] -> CoreGenT m AlgTyConRhs
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Range Int -> TyConMap -> Name TyCon -> Kind -> CoreGenT m [DataCon]
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
Range Int -> TyConMap -> Name TyCon -> Kind -> CoreGenT m [DataCon]
genDataConsFrom Range Int
range TyConMap
tcm Name TyCon
name Kind
kn
             -- TODO Generate NewTyCon
           ]

  let tc :: TyCon
tc = Int -> Name TyCon -> Kind -> Int -> AlgTyConRhs -> Bool -> TyCon
AlgTyCon (Name TyCon -> Int
forall a. Name a -> Int
nameUniq Name TyCon
name) Name TyCon
name Kind
kn Int
arity AlgTyConRhs
rhs Bool
False

  CoreGenT m Bool
forall (m :: Type -> Type). Monad m => CoreGenT m Bool
canGenDataKinds CoreGenT m Bool
-> (Bool -> CoreGenT m TyConMap) -> CoreGenT m TyConMap
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Bool
True ->
      -- Promote all the data constructors in the TyCon.
      let dcs :: [DataCon]
dcs = TyCon -> [DataCon]
tyConDataCons TyCon
tc
       in TyConMap -> CoreGenT m TyConMap
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ([(Name TyCon, TyCon)] -> TyConMap
forall a b. Uniquable a => [(a, b)] -> UniqMap b
UniqMap.fromList ((Name TyCon
name, TyCon
tc) (Name TyCon, TyCon)
-> [(Name TyCon, TyCon)] -> [(Name TyCon, TyCon)]
forall a. a -> [a] -> [a]
: (DataCon -> (Name TyCon, TyCon))
-> [DataCon] -> [(Name TyCon, TyCon)]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap DataCon -> (Name TyCon, TyCon)
promoteDataCon [DataCon]
dcs))

    Bool
False ->
      TyConMap -> CoreGenT m TyConMap
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Name TyCon -> TyCon -> TyConMap
forall a b. Uniquable a => a -> b -> UniqMap b
UniqMap.singleton Name TyCon
name TyCon
tc)
 where
   promoteDataCon :: DataCon -> (Name TyCon, TyCon)
promoteDataCon DataCon
dc =
     let tcn :: Name TyCon
tcn = DcName -> Name TyCon
coerce (DataCon -> DcName
dcName DataCon
dc)
         arity :: Int
arity = Kind -> Int
arityOf (DataCon -> Kind
dcType DataCon
dc)
      in (Name TyCon
tcn, Int -> Name TyCon -> Kind -> Int -> DataCon -> TyCon
PromotedDataCon (DataCon -> Int
dcUniq DataCon
dc) Name TyCon
tcn (DataCon -> Kind
dcType DataCon
dc) Int
arity DataCon
dc)

-- TODO In the future we may want to also generate indirectly recursive type
-- families. For example:
--
--   Even 0 = 'True         Odd 0 = 'False
--   Even n = Odd (n - 1)   Odd n = Even (n - 1)

-- | Generate a new type family, using the types that are already in scope.
--
genFunTyConFrom
  :: forall m
   . (Alternative m, MonadGen m)
  => TyConMap
  -> CoreGenT m TyConMap
genFunTyConFrom :: TyConMap -> CoreGenT m TyConMap
genFunTyConFrom TyConMap
tcm = do
  let used :: UniqMap Int
used = (TyCon -> Int) -> TyConMap -> UniqMap Int
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap TyCon -> Int
tyConUniq TyConMap
tcm
  Name TyCon
name <- UniqMap Int -> CoreGenT m (Name TyCon) -> CoreGenT m (Name TyCon)
forall (m :: Type -> Type) a b.
MonadGen m =>
UniqMap b -> m (Name a) -> m (Name a)
genFreshName UniqMap Int
used CoreGenT m (Name TyCon)
forall (m :: Type -> Type). MonadGen m => m (Name TyCon)
genTyConName

  Kind
kn <- TyConMap -> Kind -> CoreGenT m Kind
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> Kind -> CoreGenT m Kind
genClosedKindFrom TyConMap
tcm Kind
liftedTypeKind
  let arity :: Int
arity = Kind -> Int
arityOf Kind
kn

  let ([Either TyVar Kind]
argKns, Kind
resKn) = Kind -> ([Either TyVar Kind], Kind)
splitFunForallTy Kind
kn
  [([Kind], Kind)]
substs <- Name TyCon -> [Kind] -> Kind -> CoreGenT m [([Kind], Kind)]
genSubsts Name TyCon
name ([Either TyVar Kind] -> [Kind]
forall a b. [Either a b] -> [b]
rights [Either TyVar Kind]
argKns) Kind
resKn

  let tc :: TyCon
tc = Int -> Name TyCon -> Kind -> Int -> [([Kind], Kind)] -> TyCon
FunTyCon (Name TyCon -> Int
forall a. Name a -> Int
nameUniq Name TyCon
name) Name TyCon
name Kind
kn Int
arity [([Kind], Kind)]
substs
  TyConMap -> CoreGenT m TyConMap
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Name TyCon -> TyCon -> TyConMap
forall a b. Uniquable a => a -> b -> UniqMap b
UniqMap.singleton Name TyCon
name TyCon
tc)
 where
  genSubsts :: TyConName -> [Kind] -> Kind -> CoreGenT m [([Type], Type)]
  genSubsts :: Name TyCon -> [Kind] -> Kind -> CoreGenT m [([Kind], Kind)]
genSubsts Name TyCon
_ [] Kind
rhsKn = do
    -- Nullary type family, we only need to generate the RHS type
    let tcm' :: TyConMap
tcm' = (TyCon -> Bool) -> TyConMap -> TyConMap
forall b. (b -> Bool) -> UniqMap b -> UniqMap b
UniqMap.filter (Bool -> Bool
not (Bool -> Bool) -> (TyCon -> Bool) -> TyCon -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyCon -> Bool
isPrimTc) TyConMap
tcm
    Kind
rhs <- 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
forall a. Monoid a => a
mempty Kind
rhsKn

    [([Kind], Kind)] -> CoreGenT m [([Kind], Kind)]
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure [([], Kind
rhs)]

  genSubsts Name TyCon
name [Kind]
argKns Kind
rhsKn = do
    let tcm' :: TyConMap
tcm' = (TyCon -> Bool) -> TyConMap -> TyConMap
forall b. (b -> Bool) -> UniqMap b -> UniqMap b
UniqMap.filter (Bool -> Bool
not (Bool -> Bool) -> (TyCon -> Bool) -> TyCon -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyCon -> Bool
isPrimTc) TyConMap
tcm

    [TyVar]
tvs <- (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]
argKns CoreGenT m (Name Kind)
forall (m :: Type -> Type) a. MonadGen m => m (Name a)
genVarName
    let acc :: [(UniqMap TyVar, Kind)]
acc = (TyVar -> (UniqMap TyVar, Kind))
-> [TyVar] -> [(UniqMap TyVar, Kind)]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (\TyVar
x -> (TyVar -> UniqMap TyVar
forall a. Uniquable a => a -> UniqMap a
UniqMap.singletonUnique TyVar
x, TyVar -> Kind
VarTy TyVar
x)) [TyVar]
tvs

    [[(UniqMap TyVar, Kind)]]
lhss <- TyConMap
-> [(UniqMap TyVar, Kind)] -> CoreGenT m [[(UniqMap TyVar, Kind)]]
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> [(UniqMap TyVar, Kind)] -> m [[(UniqMap TyVar, Kind)]]
refineArgs TyConMap
tcm [(UniqMap TyVar, Kind)]
acc

    [[(UniqMap TyVar, Kind)]]
-> ([(UniqMap TyVar, Kind)] -> CoreGenT m ([Kind], Kind))
-> CoreGenT m [([Kind], Kind)]
forall (t :: Type -> Type) (m :: Type -> Type) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [[(UniqMap TyVar, Kind)]]
lhss (([(UniqMap TyVar, Kind)] -> CoreGenT m ([Kind], Kind))
 -> CoreGenT m [([Kind], Kind)])
-> ([(UniqMap TyVar, Kind)] -> CoreGenT m ([Kind], Kind))
-> CoreGenT m [([Kind], Kind)]
forall a b. (a -> b) -> a -> b
$ \[(UniqMap TyVar, Kind)]
args -> do
      -- The RHS of each equation can use free vars in the types, since types
      -- in the LHS of a type family are treated more like patterns.
      let free :: UniqMap TyVar
free = [UniqMap TyVar] -> UniqMap TyVar
forall a. Monoid a => [a] -> a
mconcat (((UniqMap TyVar, Kind) -> UniqMap TyVar)
-> [(UniqMap TyVar, Kind)] -> [UniqMap TyVar]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (UniqMap TyVar, Kind) -> UniqMap TyVar
forall a b. (a, b) -> a
fst [(UniqMap TyVar, Kind)]
args)

      -- Direct recursion in type families requires -XUndecidableInstances.
      Kind
rhs <- CoreGenT m Bool
forall (m :: Type -> Type). Monad m => CoreGenT m Bool
canGenUndecidableInstances CoreGenT m Bool -> (Bool -> CoreGenT m Kind) -> CoreGenT m Kind
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Bool
True -> [CoreGenT m Kind] -> CoreGenT m Kind
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice
                   [ 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
free Kind
rhsKn
                   , Name TyCon -> [Kind] -> Kind
mkTyConApp Name TyCon
name
                       ([Kind] -> Kind) -> CoreGenT m [Kind] -> CoreGenT m Kind
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (Kind -> CoreGenT m Kind) -> [Kind] -> CoreGenT m [Kind]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (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
free) [Kind]
argKns
                   ]
        Bool
False -> 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
free Kind
rhsKn

      ([Kind], Kind) -> CoreGenT m ([Kind], Kind)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (((UniqMap TyVar, Kind) -> Kind)
-> [(UniqMap TyVar, Kind)] -> [Kind]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (UniqMap TyVar, Kind) -> Kind
forall a b. (a, b) -> b
snd [(UniqMap TyVar, Kind)]
args, Kind
rhs)

refineArgs
  :: forall m
   . (Alternative m, MonadGen m)
  => TyConMap
  -> [(UniqMap TyVar, Type)]
  -> m [[(UniqMap TyVar, Type)]]
refineArgs :: TyConMap -> [(UniqMap TyVar, Kind)] -> m [[(UniqMap TyVar, Kind)]]
refineArgs TyConMap
tcm [(UniqMap TyVar, Kind)]
args = [[(UniqMap TyVar, Kind)]] -> m [[(UniqMap TyVar, Kind)]]
go [[(UniqMap TyVar, Kind)]
args]
 where
  go :: [[(UniqMap TyVar, Kind)]] -> m [[(UniqMap TyVar, Kind)]]
go [[(UniqMap TyVar, Kind)]]
acc =
    ([m [[(UniqMap TyVar, Kind)]]] -> m [[(UniqMap TyVar, Kind)]])
-> [m [[(UniqMap TyVar, Kind)]]]
-> [m [[(UniqMap TyVar, Kind)]]]
-> m [[(UniqMap TyVar, Kind)]]
forall (m :: Type -> Type) a.
MonadGen m =>
([m a] -> m a) -> [m a] -> [m a] -> m a
Gen.recursive [m [[(UniqMap TyVar, Kind)]]] -> m [[(UniqMap TyVar, Kind)]]
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice
      [[[(UniqMap TyVar, Kind)]] -> m [[(UniqMap TyVar, Kind)]]
forall (m :: Type -> Type) a. MonadGen m => a -> m a
Gen.constant [[(UniqMap TyVar, Kind)]]
acc]
      [m [[(UniqMap TyVar, Kind)]]
-> ([[(UniqMap TyVar, Kind)]] -> m [[(UniqMap TyVar, Kind)]])
-> m [[(UniqMap TyVar, Kind)]]
forall (m :: Type -> Type) a.
MonadGen m =>
m a -> (a -> m a) -> m a
Gen.subtermM ([[(UniqMap TyVar, Kind)]] -> m [[(UniqMap TyVar, Kind)]]
refineAgain [[(UniqMap TyVar, Kind)]]
acc) [[(UniqMap TyVar, Kind)]] -> m [[(UniqMap TyVar, Kind)]]
go]

  refineAgain :: [[(UniqMap TyVar, Kind)]] -> m [[(UniqMap TyVar, Kind)]]
refineAgain acc :: [[(UniqMap TyVar, Kind)]]
acc@([(UniqMap TyVar, Kind)]
xs:[[(UniqMap TyVar, Kind)]]
_) = do
    -- Every arg can be refined or left alone.
    let gen :: (UniqMap TyVar, Kind) -> m (UniqMap TyVar, Kind)
gen (UniqMap TyVar, Kind)
x = [m (UniqMap TyVar, Kind)] -> m (UniqMap TyVar, Kind)
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice [(UniqMap TyVar -> Kind -> m (UniqMap TyVar, Kind))
-> (UniqMap TyVar, Kind) -> m (UniqMap TyVar, Kind)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (TyConMap -> UniqMap TyVar -> Kind -> m (UniqMap TyVar, Kind)
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap TyVar -> Kind -> m (UniqMap TyVar, Kind)
refineArg TyConMap
tcm) (UniqMap TyVar, Kind)
x, (UniqMap TyVar, Kind) -> m (UniqMap TyVar, Kind)
forall (m :: Type -> Type) a. MonadGen m => a -> m a
Gen.constant (UniqMap TyVar, Kind)
x]
    [(UniqMap TyVar, Kind)]
refined <- ((UniqMap TyVar, Kind) -> m (UniqMap TyVar, Kind))
-> [(UniqMap TyVar, Kind)] -> m [(UniqMap TyVar, Kind)]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (UniqMap TyVar, Kind) -> m (UniqMap TyVar, Kind)
gen [(UniqMap TyVar, Kind)]
xs
    [[(UniqMap TyVar, Kind)]] -> m [[(UniqMap TyVar, Kind)]]
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ([(UniqMap TyVar, Kind)]
refined [(UniqMap TyVar, Kind)]
-> [[(UniqMap TyVar, Kind)]] -> [[(UniqMap TyVar, Kind)]]
forall a. a -> [a] -> [a]
: [[(UniqMap TyVar, Kind)]]
acc)

  refineAgain [] =
    [Char] -> m [[(UniqMap TyVar, Kind)]]
forall a. HasCallStack => [Char] -> a
error [Char]
"refineArgs: No types to refine."

-- | Refine a type, selecting one of the free variables and substituting it
-- for a type constructor of the desired kind (filling in any holes with new
-- type variables). For example, successive calls may give
--
--   a ~> A b c ~> A (B b) c ~> A (B b) C ~> A (B D) C
--
refineArg
  :: forall m
   . (Alternative m, MonadGen m)
  => TyConMap
  -> UniqMap TyVar
  -> Type
  -> m (UniqMap TyVar, Type)
refineArg :: TyConMap -> UniqMap TyVar -> Kind -> m (UniqMap TyVar, Kind)
refineArg TyConMap
tcm UniqMap TyVar
free Kind
ty
  | UniqMap TyVar -> Bool
forall a. UniqMap a -> Bool
UniqMap.null UniqMap TyVar
free
  = (UniqMap TyVar, Kind) -> m (UniqMap TyVar, Kind)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (UniqMap TyVar
free, Kind
ty)

  | Bool
otherwise
  = do -- Pick a free variable and remove it from free vars
       TyVar
fv <- (TyVar, [Kind]) -> TyVar
forall a b. (a, b) -> a
fst ((TyVar, [Kind]) -> TyVar) -> m (TyVar, [Kind]) -> m TyVar
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> UniqMap TyVar -> m (TyVar, [Kind])
forall (m :: Type -> Type) v.
(Alternative m, MonadGen m, HasType v) =>
UniqMap v -> m (v, [Kind])
sampleAnyUniqMap UniqMap TyVar
free
       let free' :: UniqMap TyVar
free' = TyVar -> UniqMap TyVar -> UniqMap TyVar
forall a b. Uniquable a => a -> UniqMap b -> UniqMap b
UniqMap.delete TyVar
fv UniqMap TyVar
free

       -- Pick a type constructor that fits that free variable. This cannot be
       -- an unboxed primitive type, so for now all primitive types are excluded.
       -- This is slightly too strict, as Integer and Natural can be used.
       (TyCon
tc, [Kind]
holes) <- (TyCon -> Bool) -> Kind -> TyConMap -> m (TyCon, [Kind])
forall (m :: Type -> Type) v.
(Alternative m, MonadGen m, HasType v, Bias v) =>
(v -> Bool) -> Kind -> UniqMap v -> m (v, [Kind])
sampleUniqMapBiased (Bool -> Bool
not (Bool -> Bool) -> (TyCon -> Bool) -> TyCon -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyCon -> Bool
isPrimTc) (TyVar -> Kind
forall a. HasType a => a -> Kind
coreTypeOf TyVar
fv) TyConMap
tcm

       -- Take any holes for that constructor and make them new free variables.
       [TyVar]
holeVars <- (Kind -> m (Name Kind) -> m TyVar)
-> [Kind] -> m (Name Kind) -> 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 -> m (Name Kind) -> m TyVar
forall (m :: Type -> Type).
MonadGen m =>
Kind -> m (Name Kind) -> m TyVar
genTyVar [Kind]
holes m (Name Kind)
forall (m :: Type -> Type) a. MonadGen m => m (Name a)
genVarName
       let free'' :: UniqMap TyVar
free'' = [(TyVar, TyVar)] -> UniqMap TyVar -> UniqMap TyVar
forall a b. Uniquable a => [(a, b)] -> UniqMap b -> UniqMap b
UniqMap.insertMany ([TyVar] -> [TyVar] -> [(TyVar, TyVar)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TyVar]
holeVars [TyVar]
holeVars) UniqMap TyVar
free'

       -- Substitute the removed free variable for the type constructor with
       -- any new free variables applied to it.
       let inScope :: InScopeSet
inScope = InScopeSet -> [TyVar] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
emptyInScopeSet (UniqMap TyVar -> [TyVar]
forall b. UniqMap b -> [b]
UniqMap.elems UniqMap TyVar
free'')
       let substTv :: VarEnv Kind
substTv = TyVar -> Kind -> VarEnv Kind
forall b a. Var b -> a -> VarEnv a
unitVarEnv TyVar
fv (Name TyCon -> [Kind] -> Kind
mkTyConApp (TyCon -> Name TyCon
tyConName TyCon
tc) ((TyVar -> Kind) -> [TyVar] -> [Kind]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap TyVar -> Kind
VarTy [TyVar]
holeVars))
       let subst :: Subst
subst = InScopeSet -> VarEnv Kind -> Subst
mkTvSubst InScopeSet
inScope VarEnv Kind
substTv

       -- Return the refined type and free variable environment.
       (UniqMap TyVar, Kind) -> m (UniqMap TyVar, Kind)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (UniqMap TyVar
free'', HasCallStack => Subst -> Kind -> Kind
Subst -> Kind -> Kind
substTy Subst
subst Kind
ty)

{-
Note [generating substs]
~~~~~~~~~~~~~~~~~~~~~~~~
When generating a FunTyCon, we generate a sequence of alternatives like so:

  type family F t_1 t_2 ... t_n :: k where
    F a_1 a_2 ... a_n = k_1
    F b_1 b_2 ... b_n = k_2
    ...

In "real" code, these would typically be written in a way where no alternatives
become dead, i.e. if you start with

  F v_1 v_2 ... v_n

where all arguments v_1 are type variables, this pattern will always be used
and the other alternatives never considered. To ensure FunTyCon are more
realistic, alternatives should be ordered from more specific to more general.

We can achieve this by starting from the most general solution and working
towards a more specific solution. At each step we can either return the list
of alternatives, or make a more specific alternative based on the previous
most specific alternative and put this at the head of the list. The returned
list is then guaranteed to not have any dead alternatives. Multiple of these
lists can be merged provided the merge operation preserves the ordering of
more to less specific.

If there is need to test type families where some alternatives may be dead, we
can use Gen.shuffle to rearrange the substs before taking a subsequence.
-}