{-|
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.Unique

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 = length . fst . 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 numDcs = go tysPrimMap
 where
  -- Either stop adding new items to the TyConMap, or generate a new item.
  -- 'Gen.recursive' is necessary to ensure termination.
  go tcm =
    Gen.recursive Gen.choice
      [Gen.constant tcm]
      [Gen.subtermM (extendTyConMap tcm) go]

  extendTyConMap tcm = do
    -- We return new UniqMap instead of individual TyCon, because for AlgTyCon
    -- we may also generate PromotedDataCon for -XDataKinds.
    new <- canGenTypeFamilies >>= \case
      True -> Gen.choice
        [ genAlgTyConFrom numDcs tcm
        , genFunTyConFrom tcm <|> genAlgTyConFrom numDcs tcm
        ]
      False -> Gen.choice [genAlgTyConFrom numDcs tcm]

    pure (unionUniqMap tcm 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 tcm = do
  let used = uniqMapToUniqSet (fmap tyConUniq tcm)
  name <- genFreshName used 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 = liftedTypeKind
  let arity = arityOf kn

  rhs <- Gen.choice
           [ DataTyCon <$> genDataConsFrom range tcm name kn
             -- TODO Generate NewTyCon
           ]

  let tc = AlgTyCon (nameUniq name) name kn arity rhs False

  canGenDataKinds >>= \case
    True ->
      -- Promote all the data constructors in the TyCon.
      let dcs = tyConDataCons tc
       in pure (listToUniqMap ((name, tc) : fmap promoteDataCon dcs))

    False ->
      pure (unitUniqMap name tc)
 where
   promoteDataCon dc =
     let tcn = coerce (dcName dc)
         arity = arityOf (dcType dc)
      in (tcn, PromotedDataCon (dcUniq dc) tcn (dcType dc) arity 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 tcm = do
  let used = uniqMapToUniqSet (fmap tyConUniq tcm)
  name <- genFreshName used genTyConName

  kn <- genClosedKindFrom tcm liftedTypeKind
  let arity = arityOf kn

  let (argKns, resKn) = splitFunForallTy kn
  substs <- genSubsts name (rights argKns) resKn

  let tc = FunTyCon (nameUniq name) name kn arity substs
  pure (unitUniqMap name tc)
 where
  genSubsts :: TyConName -> [Kind] -> Kind -> CoreGenT m [([Type], Type)]
  genSubsts _ [] rhsKn = do
    -- Nullary type family, we only need to generate the RHS type
    let tcm' = filterUniqMap (not . isPrimTc) tcm
    rhs <- genMonoTypeFrom tcm' emptyUniqMap rhsKn

    pure [([], rhs)]

  genSubsts name argKns rhsKn = do
    let tcm' = filterUniqMap (not . isPrimTc) tcm

    tvs <- genVars genTyVar argKns genVarName
    let acc = fmap (\x -> (unitUniqMap x x, VarTy x)) tvs

    lhss <- refineArgs tcm acc

    forM lhss $ \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 = mconcat (fmap fst args)

      -- Direct recursion in type families requires -XUndecidableInstances.
      rhs <- canGenUndecidableInstances >>= \case
        True -> Gen.choice
                   [ genMonoTypeFrom tcm' free rhsKn
                   , mkTyConApp name
                       <$> traverse (genMonoTypeFrom tcm' free) argKns
                   ]
        False -> genMonoTypeFrom tcm' free rhsKn

      pure (fmap snd args, rhs)

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

  refineAgain acc@(xs:_) = do
    -- Every arg can be refined or left alone.
    let gen x = Gen.choice [uncurry (refineArg tcm) x, Gen.constant x]
    refined <- traverse gen xs
    pure (refined : acc)

  refineAgain [] =
    error "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 tcm free ty
  | nullUniqMap free
  = pure (free, ty)

  | otherwise
  = do -- Pick a free variable and remove it from free vars
       fv <- fst <$> sampleAnyUniqMap free
       let free' = delUniqMap free fv

       -- 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.
       (tc, holes) <- sampleUniqMapBiased (not . isPrimTc) (coreTypeOf fv) tcm

       -- Take any holes for that constructor and make them new free variables.
       holeVars <- genVars genTyVar holes genVarName
       let free'' = extendListUniqMap free' (zip holeVars holeVars)

       -- Substitute the removed free variable for the type constructor with
       -- any new free variables applied to it.
       let inScope = extendInScopeSetList emptyInScopeSet (eltsUniqMap free'')
       let substTv = unitVarEnv fv (mkTyConApp (tyConName tc) (fmap VarTy holeVars))
       let subst = mkTvSubst inScope substTv

       -- Return the refined type and free variable environment.
       pure (free'', substTy subst 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.
-}