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

Random, type-directed generation of Term.
-}

{-# LANGUAGE TupleSections #-}

module Clash.Hedgehog.Core.Term
  ( genTermFrom
  ) where

import Control.Monad (forM)
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range

import Clash.Core.DataCon
import Clash.Core.HasType
import Clash.Core.Pretty (showPpr)
import Clash.Core.Term
import Clash.Core.TyCon
import Clash.Core.Type
import Clash.Core.TysPrim (liftedTypeKind, typeSymbolKind)
import Clash.Core.Util (listToLets)
import Clash.Core.Var
import Clash.Data.UniqMap (UniqMap)
import qualified Clash.Data.UniqMap as UniqMap

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

-- | Sample a data constructor from the environment, potentially partially
-- applying it so that the type fits the hole. If there are no possible fits
-- for the hole in the environment, an alternative generator is used instead.
sampleDataConOr
  :: forall m
   . (Alternative m, MonadGen m)
  => TyConMap
  -- ^ The types in scope while generating
  -> Type
  -- ^ The hole to generate a fit for
  -> (Type -> CoreGenT m Term)
  -- ^ A generator for sub-holes (used when partially applying a fit)
  -> CoreGenT m Term
  -- ^ A generator to use if there are no hole fits
  -> CoreGenT m Term
sampleDataConOr :: TyConMap
-> Type
-> (Type -> CoreGenT m Term)
-> CoreGenT m Term
-> CoreGenT m Term
sampleDataConOr TyConMap
tcm Type
hole Type -> CoreGenT m Term
genSub CoreGenT m Term
genOr =
  CoreGenT m Term
sampleDataCon CoreGenT m Term -> CoreGenT m Term -> CoreGenT m Term
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> CoreGenT m Term
genOr
 where
  sampleDataCon :: CoreGenT m Term
sampleDataCon = do
    -- TODO We cannot fill in any datacon where the type is polymorphic. This
    -- is because sampleUniqMap will never pick one because it cannot see the
    -- fit is valid without unification. See the TODO by sampleUniqMap.
    --
    -- We mitigate this by not generating algebraic type constructors with
    -- kinds other than Type, i.e. no type params / poly kinds. See the
    -- 'genAlgTyCon' function and NOTE [finding more complex fits].
    let dcs :: [DataCon]
dcs = (TyCon -> [DataCon]) -> TyConMap -> [DataCon]
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> [b]) -> t a -> [b]
concatMap TyCon -> [DataCon]
tyConDataCons TyConMap
tcm
    let dcm :: UniqMap DataCon
dcm = [(DataCon, DataCon)] -> UniqMap DataCon
forall a b. Uniquable a => [(a, b)] -> UniqMap b
UniqMap.fromList ([DataCon] -> [DataCon] -> [(DataCon, DataCon)]
forall a b. [a] -> [b] -> [(a, b)]
zip [DataCon]
dcs [DataCon]
dcs)
    (DataCon
dc, [Type]
holes) <- (DataCon -> Bool)
-> Type -> UniqMap DataCon -> CoreGenT m (DataCon, [Type])
forall (m :: Type -> Type) v.
(Alternative m, MonadGen m, HasType v) =>
(v -> Bool) -> Type -> UniqMap v -> m (v, [Type])
sampleUniqMap (Bool -> DataCon -> Bool
forall a b. a -> b -> a
const Bool
True) Type
hole UniqMap DataCon
dcm
    [Term]
holeFills <- (Type -> CoreGenT m Term) -> [Type] -> CoreGenT m [Term]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Type -> CoreGenT m Term
genSub [Type]
holes

    Term -> CoreGenT m Term
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> [Term] -> Term
mkTmApps (DataCon -> Term
Data DataCon
dc) [Term]
holeFills)

-- | Attempt to sample an identifier which can be made to fit a hole of the
-- desired type. If this is not possible (due to nothing in the environment
-- matching) then the given alternative generator is used instead.
--
sampleIdOr
  :: forall m
   . (Alternative m, MonadGen m)
  => UniqMap (Either TyVar Id)
  -- ^ The currently bound type and term variables
  -> Type
  -- ^ The hole to generate a fit for
  -> (Type -> CoreGenT m Term)
  -- ^ A generator for sub-holes (used when partially applying a fit)
  -> CoreGenT m Term
  -- ^ A generator to use if there are no hole fits
  -> CoreGenT m Term
sampleIdOr :: UniqMap (Either TyVar Id)
-> Type
-> (Type -> CoreGenT m Term)
-> CoreGenT m Term
-> CoreGenT m Term
sampleIdOr UniqMap (Either TyVar Id)
env Type
hole Type -> CoreGenT m Term
genSub CoreGenT m Term
genOr =
  CoreGenT m Term
sampleId CoreGenT m Term -> CoreGenT m Term -> CoreGenT m Term
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> CoreGenT m Term
genOr
 where
  sampleId :: CoreGenT m Term
sampleId = do
    let tmEnv :: UniqMap Id
tmEnv = (Either TyVar Id -> Maybe Id)
-> UniqMap (Either TyVar Id) -> UniqMap Id
forall a b. (a -> Maybe b) -> UniqMap a -> UniqMap b
UniqMap.mapMaybe ((TyVar -> Maybe Id)
-> (Id -> Maybe Id) -> Either TyVar Id -> Maybe Id
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe Id -> TyVar -> Maybe Id
forall a b. a -> b -> a
const Maybe Id
forall a. Maybe a
Nothing) Id -> Maybe Id
forall a. a -> Maybe a
Just) UniqMap (Either TyVar Id)
env
    (Id
i, [Type]
holes) <- (Id -> Bool) -> Type -> UniqMap Id -> CoreGenT m (Id, [Type])
forall (m :: Type -> Type) v.
(Alternative m, MonadGen m, HasType v) =>
(v -> Bool) -> Type -> UniqMap v -> m (v, [Type])
sampleUniqMap (Bool -> Id -> Bool
forall a b. a -> b -> a
const Bool
True) Type
hole UniqMap Id
tmEnv
    [Term]
holeFills <- (Type -> CoreGenT m Term) -> [Type] -> CoreGenT m [Term]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Type -> CoreGenT m Term
genSub [Type]
holes

    Term -> CoreGenT m Term
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> [Term] -> Term
mkTmApps (Id -> Term
Var Id
i) [Term]
holeFills)

-- | Generate a term that is valid for the given type constructor map and
-- environment of free type and term variables. The term generated must have
-- the specified type.
--
genTermFrom
  :: forall m
   . (Alternative m, MonadGen m)
  => TyConMap
  -- ^ The types in scope while generating
  -> UniqMap (Either TyVar Id)
  -- ^ The currently bound type and term variables
  -> Type
  -- ^ The type of the term being generated
  -> CoreGenT m Term
genTermFrom :: TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genTermFrom TyConMap
tcm UniqMap (Either TyVar Id)
env Type
hole =
  let genSub :: Type -> CoreGenT m Term
genSub = TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genTermFrom TyConMap
tcm UniqMap (Either TyVar Id)
env
      genOr :: CoreGenT m Term
genOr = TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genFreshTerm TyConMap
tcm UniqMap (Either TyVar Id)
env Type
hole
   in [CoreGenT m Term] -> CoreGenT m Term
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice
        [ TyConMap
-> Type
-> (Type -> CoreGenT m Term)
-> CoreGenT m Term
-> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap
-> Type
-> (Type -> CoreGenT m Term)
-> CoreGenT m Term
-> CoreGenT m Term
sampleDataConOr TyConMap
tcm Type
hole Type -> CoreGenT m Term
genSub CoreGenT m Term
genOr
        , UniqMap (Either TyVar Id)
-> Type
-> (Type -> CoreGenT m Term)
-> CoreGenT m Term
-> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
UniqMap (Either TyVar Id)
-> Type
-> (Type -> CoreGenT m Term)
-> CoreGenT m Term
-> CoreGenT m Term
sampleIdOr UniqMap (Either TyVar Id)
env Type
hole Type -> CoreGenT m Term
genSub CoreGenT m Term
genOr
        ]

{-
NOTE [generated terms]
~~~~~~~~~~~~~~~~~~~~~~
Term generation is currently limited in some ways. This is for no particular
reason other than to make the generator easier to understand. For example

  * when the hole is a forall or a function, a lambda (or tick) is inserted
    instead of allowing expressions like letrec or case which could still have
    the desired type

  * primitives are not currently generated, as we likely want to build a
    collection of known primitives when we build the environments before making
    types and terms for tests

  * casts are currently not generated, as the majority are discarded by Clash
    during the GHC2Core stage. See PR #1064.
-}

-- | Generate a "fresh" term, i.e. one which is randomly created according to
-- the type of the hole, rather than sampling from the known variables or data
-- constructors.
--
-- This generator will fail if there are no values for the given hole.
--
genFreshTerm
  :: forall m
   . (Alternative m, MonadGen m)
  => TyConMap
  -> UniqMap (Either TyVar Id)
  -> Type
  -> CoreGenT m Term
genFreshTerm :: TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genFreshTerm TyConMap
tcm UniqMap (Either TyVar Id)
env Type
hole =
  -- We need to normalize in case the hole is a type family.
  --
  -- TODO When casts are supported, this will need a cast between the
  -- normalized type and the given type.
  case TyConMap -> Type -> Type
normalizeType TyConMap
tcm Type
hole of
    -- Hole: forall i. a
    normHole :: Type
normHole@(ForAllTy TyVar
i Type
a) ->
      ([CoreGenT m Term] -> CoreGenT m Term)
-> [CoreGenT m Term] -> [CoreGenT m Term] -> CoreGenT m Term
forall (m :: Type -> Type) a.
MonadGen m =>
([m a] -> m a) -> [m a] -> [m a] -> m a
Gen.recursive [CoreGenT m Term] -> CoreGenT m Term
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice
        [TyVar -> Term -> Term
TyLam TyVar
i (Term -> Term) -> CoreGenT m Term -> CoreGenT m Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genTermFrom TyConMap
tcm (TyVar
-> Either TyVar Id
-> UniqMap (Either TyVar Id)
-> UniqMap (Either TyVar Id)
forall a b. Uniquable a => a -> b -> UniqMap b -> UniqMap b
UniqMap.insert TyVar
i (TyVar -> Either TyVar Id
forall a b. a -> Either a b
Left TyVar
i) UniqMap (Either TyVar Id)
env) Type
a]
        [TickInfo -> Term -> Term
Tick (TickInfo -> Term -> Term)
-> CoreGenT m TickInfo -> CoreGenT m (Term -> Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> TyConMap -> CoreGenT m TickInfo
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> CoreGenT m TickInfo
genTickInfo TyConMap
tcm CoreGenT m (Term -> Term) -> CoreGenT m Term -> CoreGenT m Term
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genFreshTerm TyConMap
tcm UniqMap (Either TyVar Id)
env Type
normHole]

    AnnType [Attr Text]
_ Type
a ->
      TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genTermFrom TyConMap
tcm UniqMap (Either TyVar Id)
env Type
a

    Type
normHole ->
      case Type -> TypeView
tyView Type
normHole of
        -- Hole: a -> b
        FunTy Type
a Type
b ->
          ([CoreGenT m Term] -> CoreGenT m Term)
-> [CoreGenT m Term] -> [CoreGenT m Term] -> CoreGenT m Term
forall (m :: Type -> Type) a.
MonadGen m =>
([m a] -> m a) -> [m a] -> [m a] -> m a
Gen.recursive [CoreGenT m Term] -> CoreGenT m Term
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice
            [do Id
i <- Type -> CoreGenT m TmName -> CoreGenT m Id
forall (m :: Type -> Type). MonadGen m => Type -> m TmName -> m Id
genLocalId Type
a (UniqMap (Either TyVar Id) -> CoreGenT m TmName -> CoreGenT m TmName
forall (m :: Type -> Type) a b.
MonadGen m =>
UniqMap b -> m (Name a) -> m (Name a)
genFreshName UniqMap (Either TyVar Id)
env CoreGenT m TmName
forall (m :: Type -> Type) a. MonadGen m => m (Name a)
genVarName)
                CoreGenT m Term -> (Term -> Term) -> CoreGenT m Term
forall (m :: Type -> Type) a. MonadGen m => m a -> (a -> a) -> m a
Gen.subterm (TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genTermFrom TyConMap
tcm (Id
-> Either TyVar Id
-> UniqMap (Either TyVar Id)
-> UniqMap (Either TyVar Id)
forall a b. Uniquable a => a -> b -> UniqMap b -> UniqMap b
UniqMap.insert Id
i (Id -> Either TyVar Id
forall a b. b -> Either a b
Right Id
i) UniqMap (Either TyVar Id)
env) Type
b) (Id -> Term -> Term
Lam Id
i)
            ]
            [TickInfo -> Term -> Term
Tick (TickInfo -> Term -> Term)
-> CoreGenT m TickInfo -> CoreGenT m (Term -> Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> TyConMap -> CoreGenT m TickInfo
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> CoreGenT m TickInfo
genTickInfo TyConMap
tcm CoreGenT m (Term -> Term) -> CoreGenT m Term -> CoreGenT m Term
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genFreshTerm TyConMap
tcm UniqMap (Either TyVar Id)
env Type
normHole]

        -- Hole: Primitive type constructor.
        TyConApp TyConName
tcn []
          |  Just PrimTyCon{} <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
tcn TyConMap
tcm
          -> ([CoreGenT m Term] -> CoreGenT m Term)
-> [CoreGenT m Term] -> [CoreGenT m Term] -> CoreGenT m Term
forall (m :: Type -> Type) a.
MonadGen m =>
([m a] -> m a) -> [m a] -> [m a] -> m a
Gen.recursive [CoreGenT m Term] -> CoreGenT m Term
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice
               [Literal -> Term
Literal (Literal -> Term) -> CoreGenT m Literal -> CoreGenT m Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> CoreGenT m Literal
forall (m :: Type -> Type). MonadGen m => Type -> m Literal
genLiteralFrom Type
normHole]
               -- We may fail to generate a case expression if there is nothing
               -- in scope to use as a subject. If this happens, let bindings
               -- are introduced so next time genCase is called it does not fail.
               [ TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genCase TyConMap
tcm UniqMap (Either TyVar Id)
env Type
normHole CoreGenT m Term -> CoreGenT m Term -> CoreGenT m Term
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genLet TyConMap
tcm UniqMap (Either TyVar Id)
env Type
normHole
               , TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genLet TyConMap
tcm UniqMap (Either TyVar Id)
env Type
normHole
               ]

        -- Hole: Algebraic type constructor.
        TyConApp TyConName
tcn [Type]
_
          |  Just AlgTyCon{} <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
tcn TyConMap
tcm
          -- We may have got here by trying to fill the hole with an identifier, so
          -- it makes sense to try again. If we got here by sampleDataConOr, the
          -- data constructor is isomorphic to Void, and we will hit the error.
          -> ([CoreGenT m Term] -> CoreGenT m Term)
-> [CoreGenT m Term] -> [CoreGenT m Term] -> CoreGenT m Term
forall (m :: Type -> Type) a.
MonadGen m =>
([m a] -> m a) -> [m a] -> [m a] -> m a
Gen.recursive [CoreGenT m Term] -> CoreGenT m Term
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice
               [TyConMap
-> Type
-> (Type -> CoreGenT m Term)
-> CoreGenT m Term
-> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap
-> Type
-> (Type -> CoreGenT m Term)
-> CoreGenT m Term
-> CoreGenT m Term
sampleDataConOr TyConMap
tcm Type
hole (TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genTermFrom TyConMap
tcm UniqMap (Either TyVar Id)
env)
                 ([Char] -> CoreGenT m Term
forall a. HasCallStack => [Char] -> a
error ([Char]
"No term level value for hole: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
hole))]
               -- We may fail to generate a case expression if there is nothing
               -- in scope to use as a subject. If this happens, let bindings
               -- are introduced so next time genCase is called it does not fail.
               [ TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genCase TyConMap
tcm UniqMap (Either TyVar Id)
env Type
normHole CoreGenT m Term -> CoreGenT m Term -> CoreGenT m Term
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genLet TyConMap
tcm UniqMap (Either TyVar Id)
env Type
normHole
               , TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genLet TyConMap
tcm UniqMap (Either TyVar Id)
env Type
normHole
               ]

        TypeView
_ ->
          [Char] -> CoreGenT m Term
forall a. HasCallStack => [Char] -> a
error ([Char]
"No term level value for hole: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
normHole)

-- TODO
-- genIsMultiPrim
-- genPrimInfo
-- genPrimUnfolding
-- genMultiPrimInfo
-- genWorkInfo

genLet
  :: forall m
   . (Alternative m, MonadGen m)
  => TyConMap
  -> UniqMap (Either TyVar Id)
  -> Type
  -> CoreGenT m Term
genLet :: TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genLet TyConMap
tcm UniqMap (Either TyVar Id)
env Type
hole = do
  [LetBinding]
binds <- TyConMap -> UniqMap (Either TyVar Id) -> CoreGenT m [LetBinding]
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> CoreGenT m [LetBinding]
genLetBindings TyConMap
tcm UniqMap (Either TyVar Id)
env
  let vars :: [Id]
vars = (LetBinding -> Id) -> [LetBinding] -> [Id]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap LetBinding -> Id
forall a b. (a, b) -> a
fst [LetBinding]
binds
  let env' :: UniqMap (Either TyVar Id)
env' = [(Id, Either TyVar Id)]
-> UniqMap (Either TyVar Id) -> UniqMap (Either TyVar Id)
forall a b. Uniquable a => [(a, b)] -> UniqMap b -> UniqMap b
UniqMap.insertMany ([Id] -> [Either TyVar Id] -> [(Id, Either TyVar Id)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
vars ((Id -> Either TyVar Id) -> [Id] -> [Either TyVar Id]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Id -> Either TyVar Id
forall a b. b -> Either a b
Right [Id]
vars)) UniqMap (Either TyVar Id)
env

  Term
body <- TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genTermFrom TyConMap
tcm UniqMap (Either TyVar Id)
env' Type
hole

  Term -> CoreGenT m Term
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ([LetBinding] -> Term -> Term
listToLets [LetBinding]
binds Term
body)

genLetBindings
  :: forall m
   . (Alternative m, MonadGen m)
  => TyConMap
  -> UniqMap (Either TyVar Id)
  -> CoreGenT m [LetBinding]
genLetBindings :: TyConMap -> UniqMap (Either TyVar Id) -> CoreGenT m [LetBinding]
genLetBindings TyConMap
tcm UniqMap (Either TyVar Id)
env = do
  let tyEnv :: UniqMap TyVar
tyEnv = (Either TyVar Id -> Maybe TyVar)
-> UniqMap (Either TyVar Id) -> UniqMap TyVar
forall a b. (a -> Maybe b) -> UniqMap a -> UniqMap b
UniqMap.mapMaybe ((TyVar -> Maybe TyVar)
-> (Id -> Maybe TyVar) -> Either TyVar Id -> Maybe TyVar
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either TyVar -> Maybe TyVar
forall a. a -> Maybe a
Just (Maybe TyVar -> Id -> Maybe TyVar
forall a b. a -> b -> a
const Maybe TyVar
forall a. Maybe a
Nothing)) UniqMap (Either TyVar Id)
env
  -- Limit the number of new bindings to 8 to prevent an explosion in the
  -- number of sub-holes to generate.
  [Type]
types <- Range Int -> CoreGenT m Type -> CoreGenT m [Type]
forall (m :: Type -> Type) a.
MonadGen m =>
Range Int -> m a -> m [a]
Gen.list (Int -> Int -> Range Int
forall a. Integral a => a -> a -> Range a
Range.linear Int
1 Int
8) (TyConMap -> UniqMap TyVar -> Type -> CoreGenT m Type
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap TyVar -> Type -> CoreGenT m Type
genMonoTypeFrom TyConMap
tcm UniqMap TyVar
tyEnv Type
liftedTypeKind)
  [Id]
vars <- (Type -> CoreGenT m TmName -> CoreGenT m Id)
-> [Type] -> CoreGenT m TmName -> CoreGenT m [Id]
forall (m :: Type -> Type) a.
MonadGen m =>
(Type -> m (Name a) -> m (Var a))
-> [Type] -> m (Name a) -> m [Var a]
genVars Type -> CoreGenT m TmName -> CoreGenT m Id
forall (m :: Type -> Type). MonadGen m => Type -> m TmName -> m Id
genLocalId [Type]
types CoreGenT m TmName
forall (m :: Type -> Type) a. MonadGen m => m (Name a)
genVarName

  [(Id, Type)]
-> ((Id, Type) -> CoreGenT m LetBinding) -> CoreGenT m [LetBinding]
forall (t :: Type -> Type) (m :: Type -> Type) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([Id] -> [Type] -> [(Id, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
vars [Type]
types) (((Id, Type) -> CoreGenT m LetBinding) -> CoreGenT m [LetBinding])
-> ((Id, Type) -> CoreGenT m LetBinding) -> CoreGenT m [LetBinding]
forall a b. (a -> b) -> a -> b
$ \(Id
v, Type
ty) ->
    -- Bindings can be indirectly recursive, but not directly recursive. This
    -- stops the generator from generating let x = x in ...
    let vars' :: [Id]
vars' = (Id -> Bool) -> [Id] -> [Id]
forall a. (a -> Bool) -> [a] -> [a]
filter (Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
/= Id
v) [Id]
vars
        env' :: UniqMap (Either TyVar Id)
env' = [(Id, Either TyVar Id)]
-> UniqMap (Either TyVar Id) -> UniqMap (Either TyVar Id)
forall a b. Uniquable a => [(a, b)] -> UniqMap b -> UniqMap b
UniqMap.insertMany ([Id] -> [Either TyVar Id] -> [(Id, Either TyVar Id)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
vars' ((Id -> Either TyVar Id) -> [Id] -> [Either TyVar Id]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Id -> Either TyVar Id
forall a b. b -> Either a b
Right [Id]
vars')) UniqMap (Either TyVar Id)
env
     in (Id
v,) (Term -> LetBinding) -> CoreGenT m Term -> CoreGenT m LetBinding
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genTermFrom TyConMap
tcm UniqMap (Either TyVar Id)
env' Type
ty

{-
NOTE [generating useful case expressions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When generating case expressions, it is perfectly valid to generate something
like the following:

  case C x_1 ... x_n of
    C p_1 ... p_n -> ...

However, we want to favour generating the more interesting case expressions
where the subject is an application on variables, like so

  case v x_1 ... x_n of
    C p_1 ... p_i -> ...
    D p_1 ... p_j -> ...
    E p_1 ... p_k -> ...

Likewise, we want to avoid generating case expressions with literals as the
subject, or function types (where the only viable pattern is DefaultPat). While
these are valid terms to generate, the generator will naturally bias towards
them, making too many of the tests too "artificial" in nature.

A downside of this approach is that generating a case expression is not
guaranteed to work, as if there is nothing in the environment that can be used
as the subject then the generator will return 'empty'. However, since letrec
always passes and introduces new bindings, we fallback to this if case fails.
This means the failure will happen at most once when generating.
-}

genCase
  :: forall m
   . (Alternative m, MonadGen m)
  => TyConMap
  -> UniqMap (Either TyVar Id)
  -> Type
  -> CoreGenT m Term
genCase :: TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genCase TyConMap
tcm UniqMap (Either TyVar Id)
env Type
altTy = do
  -- I need to select something as the subject. It can be any type, but should
  -- bias towards something from the environment where possible. It may not be
  -- possible though, so I should be able to fallback to just building a term.
  let tmEnv :: UniqMap Id
tmEnv = (Either TyVar Id -> Maybe Id)
-> UniqMap (Either TyVar Id) -> UniqMap Id
forall a b. (a -> Maybe b) -> UniqMap a -> UniqMap b
UniqMap.mapMaybe ((TyVar -> Maybe Id)
-> (Id -> Maybe Id) -> Either TyVar Id -> Maybe Id
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe Id -> TyVar -> Maybe Id
forall a b. a -> b -> a
const Maybe Id
forall a. Maybe a
Nothing) Id -> Maybe Id
forall a. a -> Maybe a
Just) UniqMap (Either TyVar Id)
env
  Term
subj <- UniqMap Id -> CoreGenT m Term
sampleSubjFrom UniqMap Id
tmEnv
  let subjTy :: Type
subjTy = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
subj

  case ((TyConName, [Type]) -> TyConName)
-> Maybe (TyConName, [Type]) -> Maybe TyConName
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (TyConName, [Type]) -> TyConName
forall a b. (a, b) -> a
fst (Type -> Maybe (TyConName, [Type])
splitTyConAppM Type
subjTy) of
    Just TyConName
tcn ->
      case TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
tcn TyConMap
tcm of
        Just tc :: TyCon
tc@AlgTyCon{} -> do
          [DataCon]
dcs <- [DataCon] -> CoreGenT m [DataCon]
forall (m :: Type -> Type) a. MonadGen m => [a] -> m [a]
Gen.subsequence (TyCon -> [DataCon]
tyConDataCons TyCon
tc)
          [Pat]
dcPats <- (DataCon -> CoreGenT m Pat) -> [DataCon] -> CoreGenT m [Pat]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse DataCon -> CoreGenT m Pat
genDataPatFrom [DataCon]
dcs
          [Alt]
alts <- (Pat -> CoreGenT m Alt) -> [Pat] -> CoreGenT m [Alt]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Pat -> CoreGenT m Alt
genAltFrom (Pat
DefaultPat Pat -> [Pat] -> [Pat]
forall a. a -> [a] -> [a]
: [Pat]
dcPats)

          Term -> CoreGenT m Term
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Type -> [Alt] -> Term
Case Term
subj Type
altTy [Alt]
alts)

        Just PrimTyCon{} -> do
          -- Upper bound is 8 to prevent explosion in number of sub-holes.
          [Pat]
litPats <- Range Int -> CoreGenT m Pat -> CoreGenT m [Pat]
forall (m :: Type -> Type) a.
MonadGen m =>
Range Int -> m a -> m [a]
Gen.list (Int -> Int -> Range Int
forall a. Integral a => a -> a -> Range a
Range.linear Int
0 Int
8) (Literal -> Pat
LitPat (Literal -> Pat) -> CoreGenT m Literal -> CoreGenT m Pat
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> CoreGenT m Literal
forall (m :: Type -> Type). MonadGen m => Type -> m Literal
genLiteralFrom Type
subjTy)
          [Alt]
alts <- (Pat -> CoreGenT m Alt) -> [Pat] -> CoreGenT m [Alt]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Pat -> CoreGenT m Alt
genAltFrom (Pat
DefaultPat Pat -> [Pat] -> [Pat]
forall a. a -> [a] -> [a]
: [Pat]
litPats)

          Term -> CoreGenT m Term
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Type -> [Alt] -> Term
Case Term
subj Type
altTy [Alt]
alts)

        Maybe TyCon
_ -> do
          Alt
alt <- Pat -> CoreGenT m Alt
genAltFrom Pat
DefaultPat
          Term -> CoreGenT m Term
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Type -> [Alt] -> Term
Case Term
subj Type
altTy [Alt
alt])

    Maybe TyConName
_ -> do
      Alt
alt <- Pat -> CoreGenT m Alt
genAltFrom Pat
DefaultPat
      Term -> CoreGenT m Term
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Type -> [Alt] -> Term
Case Term
subj Type
altTy [Alt
alt])
 where
  -- Subjects are applications on variables in the environment.
  -- See NOTE [generating useful case expressions].
  sampleSubjFrom :: UniqMap Id -> CoreGenT m Term
  sampleSubjFrom :: UniqMap Id -> CoreGenT m Term
sampleSubjFrom UniqMap Id
tmEnv = do
    (Id
v, [Type]
holes) <- UniqMap Id -> CoreGenT m (Id, [Type])
forall (m :: Type -> Type) v.
(Alternative m, MonadGen m, HasType v) =>
UniqMap v -> m (v, [Type])
sampleAnyUniqMap UniqMap Id
tmEnv
    [Term]
holeFills <- (Type -> CoreGenT m Term) -> [Type] -> CoreGenT m [Term]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genTermFrom TyConMap
tcm UniqMap (Either TyVar Id)
env) [Type]
holes

    Term -> CoreGenT m Term
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> [Term] -> Term
mkTmApps (Id -> Term
Var Id
v) [Term]
holeFills)

  genDataPatFrom :: DataCon -> CoreGenT m Pat
  genDataPatFrom :: DataCon -> CoreGenT m Pat
genDataPatFrom DataCon
dc = do
    [Id]
ids <- (Type -> CoreGenT m TmName -> CoreGenT m Id)
-> [Type] -> CoreGenT m TmName -> CoreGenT m [Id]
forall (m :: Type -> Type) a.
MonadGen m =>
(Type -> m (Name a) -> m (Var a))
-> [Type] -> m (Name a) -> m [Var a]
genVars Type -> CoreGenT m TmName -> CoreGenT m Id
forall (m :: Type -> Type). MonadGen m => Type -> m TmName -> m Id
genLocalId (DataCon -> [Type]
dcArgTys DataCon
dc) CoreGenT m TmName
forall (m :: Type -> Type) a. MonadGen m => m (Name a)
genVarName
    Pat -> CoreGenT m Pat
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (DataCon -> [TyVar] -> [Id] -> Pat
DataPat DataCon
dc (DataCon -> [TyVar]
dcExtTyVars DataCon
dc) [Id]
ids)

  genAltFrom :: Pat -> CoreGenT m Alt
  genAltFrom :: Pat -> CoreGenT m Alt
genAltFrom Pat
pat = do
    let ([TyVar]
tvs, [Id]
ids) = Pat -> ([TyVar], [Id])
patIds Pat
pat
    let toTvBind :: Var a -> (Int, Either (Var a) b)
toTvBind Var a
x = (Var a -> Int
forall a. Var a -> Int
varUniq Var a
x, Var a -> Either (Var a) b
forall a b. a -> Either a b
Left Var a
x)
    let toIdBind :: Var a -> (Int, Either a (Var a))
toIdBind Var a
x = (Var a -> Int
forall a. Var a -> Int
varUniq Var a
x, Var a -> Either a (Var a)
forall a b. b -> Either a b
Right Var a
x)

    -- Generate the terms in alternatives with the newly bound vars in scope.
    let env' :: UniqMap (Either TyVar Id)
env' = [(Int, Either TyVar Id)]
-> UniqMap (Either TyVar Id) -> UniqMap (Either TyVar Id)
forall a b. Uniquable a => [(a, b)] -> UniqMap b -> UniqMap b
UniqMap.insertMany ((TyVar -> (Int, Either TyVar Id))
-> [TyVar] -> [(Int, Either TyVar Id)]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap TyVar -> (Int, Either TyVar Id)
forall a b. Var a -> (Int, Either (Var a) b)
toTvBind [TyVar]
tvs [(Int, Either TyVar Id)]
-> [(Int, Either TyVar Id)] -> [(Int, Either TyVar Id)]
forall a. Semigroup a => a -> a -> a
<> (Id -> (Int, Either TyVar Id)) -> [Id] -> [(Int, Either TyVar Id)]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Id -> (Int, Either TyVar Id)
forall a a. Var a -> (Int, Either a (Var a))
toIdBind [Id]
ids) UniqMap (Either TyVar Id)
env
    Term
term <- TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> UniqMap (Either TyVar Id) -> Type -> CoreGenT m Term
genTermFrom TyConMap
tcm UniqMap (Either TyVar Id)
env' Type
altTy

    Alt -> CoreGenT m Alt
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Pat
pat, Term
term)

-- TODO genCast

genTickInfo
  :: forall m
   . (Alternative m, MonadGen m)
  => TyConMap
  -> CoreGenT m TickInfo
genTickInfo :: TyConMap -> CoreGenT m TickInfo
genTickInfo TyConMap
tcm =
  [CoreGenT m TickInfo] -> CoreGenT m TickInfo
forall (m :: Type -> Type) a. MonadGen m => [m a] -> m a
Gen.choice
    [ NameMod -> Type -> TickInfo
NameMod (NameMod -> Type -> TickInfo)
-> CoreGenT m NameMod -> CoreGenT m (Type -> TickInfo)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> CoreGenT m NameMod
forall (m :: Type -> Type). MonadGen m => m NameMod
genNameMod CoreGenT m (Type -> TickInfo)
-> CoreGenT m Type -> CoreGenT m TickInfo
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> TyConMap -> Type -> CoreGenT m Type
forall (m :: Type -> Type).
(Alternative m, MonadGen m) =>
TyConMap -> Type -> CoreGenT m Type
genClosedKindFrom TyConMap
tcm Type
typeSymbolKind
    , TickInfo -> CoreGenT m TickInfo
forall (m :: Type -> Type) a. MonadGen m => a -> m a
Gen.constant TickInfo
DeDup
    , TickInfo -> CoreGenT m TickInfo
forall (m :: Type -> Type) a. MonadGen m => a -> m a
Gen.constant TickInfo
NoDeDup
    ]

genNameMod :: forall m. MonadGen m => m NameMod
genNameMod :: m NameMod
genNameMod = [NameMod] -> m NameMod
forall (m :: Type -> Type) a. MonadGen m => [a] -> m a
Gen.element [NameMod
PrefixName, NameMod
SuffixName, NameMod
SuffixNameP, NameMod
SetName]