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

    Utility class to extract type information from data which has a type.
    -}

    {-# LANGUAGE CPP #-}
    {-# LANGUAGE LambdaCase #-}
    {-# LANGUAGE OverloadedStrings #-}
    {-# LANGUAGE QuasiQuotes #-}
    {-# LANGUAGE TemplateHaskell #-}

    module Clash.Core.HasType
      ( HasType(..)
      , coreKindOf
      , InferType(..)
      , inferCoreKindOf
      , applyTypeToArgs
      , piResultTy
      , piResultTys
      ) where

    #if MIN_VERSION_prettyprinter(1,7,0)
    import Prettyprinter (line)
    #else
    import Data.Text.Prettyprint.Doc (line)
    #endif

    import GHC.Stack (HasCallStack)

    import Clash.Core.DataCon (DataCon(dcType))
    import Clash.Core.HasFreeVars
    import Clash.Core.Literal (Literal(..))
    import Clash.Core.Pretty
    import Clash.Core.Subst
    import Clash.Core.Term (Term(..), IsMultiPrim(..), PrimInfo(..), collectArgs)
    import Clash.Core.TyCon (TyCon(tyConKind), TyConMap, isTupleTyConLike)
    import Clash.Core.Type
    import Clash.Core.TysPrim
    import Clash.Core.Var (Var(varType))
    import Clash.Core.VarEnv
    import qualified Clash.Data.UniqMap as UniqMap
    import Clash.Debug (debugIsOn)
    import Clash.Util (curLoc, pprPanic)
    import qualified Clash.Util.Interpolate as I

    class HasType a where
      coreTypeOf :: a -> Type

    coreKindOf :: (HasType a) => a -> Kind
    coreKindOf :: forall a. HasType a => a -> Kind
    coreKindOf = a -> Kind
    forall a. HasType a => a -> Kind
    coreTypeOf
    {-# INLINE coreKindOf #-}

    instance HasType DataCon where
      coreTypeOf :: DataCon -> Kind
    coreTypeOf = DataCon -> Kind
    dcType

    instance HasType Literal where
      coreTypeOf :: Literal -> Kind
    coreTypeOf = \case
        IntegerLiteral Integer
    _ -> Kind
    integerPrimTy
        IntLiteral Integer
    _ -> Kind
    intPrimTy
        WordLiteral Integer
    _ -> Kind
    wordPrimTy
        StringLiteral String
    _ -> Kind
    stringPrimTy
        FloatLiteral Word32
    _ -> Kind
    floatPrimTy
        DoubleLiteral Word64
    _ -> Kind
    doublePrimTy
        CharLiteral Char
    _ -> Kind
    charPrimTy
        Int64Literal Integer
    _ -> Kind
    int64PrimTy
        Word64Literal Integer
    _ -> Kind
    word64PrimTy
    #if MIN_VERSION_ghc(8,8,0)
        Int8Literal Integer
    _ -> Kind
    int8PrimTy
        Int16Literal Integer
    _ -> Kind
    int16PrimTy
        Int32Literal Integer
    _ -> Kind
    int32PrimTy
        Word8Literal Integer
    _ -> Kind
    word8PrimTy
        Word16Literal Integer
    _ -> Kind
    word16PrimTy
        Word32Literal Integer
    _ -> Kind
    word32PrimTy
    #endif
        NaturalLiteral Integer
    _ -> Kind
    naturalPrimTy
        ByteArrayLiteral ByteArray
    _ -> Kind
    byteArrayPrimTy

    instance HasType PrimInfo where
      coreTypeOf :: PrimInfo -> Kind
    coreTypeOf PrimInfo
    pr =
        case PrimInfo -> IsMultiPrim
    primMultiResult PrimInfo
    pr of
          IsMultiPrim
    SingleResult -> PrimInfo -> Kind
    primType PrimInfo
    pr

          -- See Note [MultiResult type] in Clash.Normalize.Transformations.MultiPrim
          IsMultiPrim
    MultiResult
            | let ([Either TyVar Kind]
    primArgs, Kind
    primResTy) = Kind -> ([Either TyVar Kind], Kind)
    splitFunForallTy (PrimInfo -> Kind
    primType PrimInfo
    pr)
            , TyConApp TyConName
    tupTcNm [Kind]
    tupArgs <- Kind -> TypeView
    tyView Kind
    primResTy
            , TyConName -> Bool
    isTupleTyConLike TyConName
    tupTcNm
            -> Kind -> [Either TyVar Kind] -> Kind
    mkPolyFunTy Kind
    primResTy ([Either TyVar Kind]
    primArgs [Either TyVar Kind] -> [Either TyVar Kind] -> [Either TyVar Kind]
    forall a. Semigroup a => a -> a -> a
    <> (Kind -> Either TyVar Kind) -> [Kind] -> [Either TyVar Kind]
    forall a b. (a -> b) -> [a] -> [b]
    forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
    fmap Kind -> Either TyVar Kind
    forall a b. b -> Either a b
    Right [Kind]
    tupArgs)

            | Bool
    otherwise
            -> String -> Kind
    forall a. HasCallStack => String -> a
    error String
    "PrimInfo.coreTypeOf: MultiResult primitive without tuple type"

    instance HasType TyCon where
      coreTypeOf :: TyCon -> Kind
    coreTypeOf = TyCon -> Kind
    tyConKind

    instance HasType Type where
      coreTypeOf :: Kind -> Kind
    coreTypeOf = Kind -> Kind
    forall a. a -> a
    id

    instance HasType (Var a) where
      coreTypeOf :: Var a -> Kind
    coreTypeOf = Var a -> Kind
    forall a. Var a -> Kind
    varType

    class InferType a where
      inferCoreTypeOf :: TyConMap -> a -> Type

    inferCoreKindOf :: (InferType a) => TyConMap -> a -> Kind
    inferCoreKindOf :: forall a. InferType a => TyConMap -> a -> Kind
    inferCoreKindOf = TyConMap -> a -> Kind
    forall a. InferType a => TyConMap -> a -> Kind
    inferCoreTypeOf
    {-# INLINE inferCoreKindOf #-}

    instance InferType Type where
      inferCoreTypeOf :: TyConMap -> Kind -> Kind
    inferCoreTypeOf TyConMap
    tcm Kind
    ty =
        case Kind -> TypeView
    tyView Kind
    ty of
          FunTy{} ->
            Kind
    liftedTypeKind

          TyConApp TyConName
    tc [Kind]
    args ->
            HasCallStack => TyConMap -> Kind -> [Kind] -> Kind
    TyConMap -> Kind -> [Kind] -> Kind
    piResultTys TyConMap
    tcm (TyCon -> Kind
    tyConKind (TyConName -> TyConMap -> TyCon
    forall a b. Uniquable a => a -> UniqMap b -> b
    UniqMap.find TyConName
    tc TyConMap
    tcm)) [Kind]
    args

          OtherType{} ->
            case Kind
    ty of
              ConstTy ConstTy
    c -> String -> Kind
    forall a. HasCallStack => String -> a
    error (String -> Kind) -> String -> Kind
    forall a b. (a -> b) -> a -> b
    $ $(curLoc) String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    "inferCoreTypeOf: naked ConstTy: " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ ConstTy -> String
    forall a. Show a => a -> String
    show ConstTy
    c
              VarTy TyVar
    k -> TyVar -> Kind
    forall a. Var a -> Kind
    varType TyVar
    k
              ForAllTy TyVar
    _ Kind
    a -> TyConMap -> Kind -> Kind
    forall a. InferType a => TyConMap -> a -> Kind
    inferCoreTypeOf TyConMap
    tcm Kind
    a
              LitTy NumTy{} -> Kind
    typeNatKind
              LitTy SymTy{} -> Kind
    typeSymbolKind
              LitTy CharTy{} -> Kind
    charPrimTy
              AnnType [Attr Text]
    _ Kind
    a -> TyConMap -> Kind -> Kind
    forall a. InferType a => TyConMap -> a -> Kind
    inferCoreTypeOf TyConMap
    tcm Kind
    a
              AppTy Kind
    a Kind
    b -> Kind -> [Kind] -> Kind
    go Kind
    a [Kind
    b]
               where
                go :: Kind -> [Kind] -> Kind
    go (AppTy Kind
    c Kind
    d) [Kind]
    args = Kind -> [Kind] -> Kind
    go Kind
    c (Kind
    d Kind -> [Kind] -> [Kind]
    forall a. a -> [a] -> [a]
    : [Kind]
    args)
                go Kind
    c [Kind]
    args = HasCallStack => TyConMap -> Kind -> [Kind] -> Kind
    TyConMap -> Kind -> [Kind] -> Kind
    piResultTys TyConMap
    tcm (TyConMap -> Kind -> Kind
    forall a. InferType a => TyConMap -> a -> Kind
    inferCoreTypeOf TyConMap
    tcm Kind
    c) [Kind]
    args

    instance InferType Term where
      inferCoreTypeOf :: TyConMap -> Term -> Kind
    inferCoreTypeOf TyConMap
    tcm = Term -> Kind
    go
       where
        go :: Term -> Kind
    go = \case
          Var Id
    i -> Id -> Kind
    forall a. HasType a => a -> Kind
    coreTypeOf Id
    i
          Data DataCon
    dc -> DataCon -> Kind
    forall a. HasType a => a -> Kind
    coreTypeOf DataCon
    dc
          Literal Literal
    l -> Literal -> Kind
    forall a. HasType a => a -> Kind
    coreTypeOf Literal
    l
          Prim PrimInfo
    pr -> PrimInfo -> Kind
    forall a. HasType a => a -> Kind
    coreTypeOf PrimInfo
    pr
          Lam Id
    i Term
    x -> Kind -> Kind -> Kind
    mkFunTy (Id -> Kind
    forall a. HasType a => a -> Kind
    coreTypeOf Id
    i) (Term -> Kind
    go Term
    x)
          TyLam TyVar
    i Term
    x -> TyVar -> Kind -> Kind
    ForAllTy TyVar
    i (Term -> Kind
    go Term
    x)

          x :: Term
    x@App{} ->
            case Term -> (Term, [Either Term Kind])
    collectArgs Term
    x of
              (Term
    fun, [Either Term Kind]
    args) -> Term -> TyConMap -> Kind -> [Either Term Kind] -> Kind
    applyTypeToArgs Term
    x TyConMap
    tcm (Term -> Kind
    go Term
    fun) [Either Term Kind]
    args

          x :: Term
    x@TyApp{} ->
            case Term -> (Term, [Either Term Kind])
    collectArgs Term
    x of
              (Term
    fun, [Either Term Kind]
    args) -> Term -> TyConMap -> Kind -> [Either Term Kind] -> Kind
    applyTypeToArgs Term
    x TyConMap
    tcm (Term -> Kind
    go Term
    fun) [Either Term Kind]
    args

          Let Bind Term
    _ Term
    x -> Term -> Kind
    go Term
    x
          Case Term
    _ Kind
    ty [Alt]
    _ -> Kind
    ty
          Cast Term
    _ Kind
    _ Kind
    a -> Kind
    a
          Tick TickInfo
    _ Term
    x -> Term -> Kind
    go Term
    x

    -- | Get the result type of a polymorphic function given a list of arguments
    applyTypeToArgs
      :: Term
      -- ^ The complete term, used for error messages.
      -> TyConMap
      -> Type
      -> [Either Term Type]
      -> Type
    applyTypeToArgs :: Term -> TyConMap -> Kind -> [Either Term Kind] -> Kind
    applyTypeToArgs Term
    e TyConMap
    m Kind
    opTy [Either Term Kind]
    args = Kind -> [Either Term Kind] -> Kind
    forall {p}. PrettyPrec p => Kind -> [Either p Kind] -> Kind
    go Kind
    opTy [Either Term Kind]
    args
     where
      go :: Kind -> [Either p Kind] -> Kind
    go Kind
    opTy' []               = Kind
    opTy'
      go Kind
    opTy' (Right Kind
    ty:[Either p Kind]
    args') = Kind -> [Kind] -> [Either p Kind] -> Kind
    goTyArgs Kind
    opTy' [Kind
    ty] [Either p Kind]
    args'
      go Kind
    opTy' (Left p
    a:[Either p Kind]
    args')   = case TyConMap -> Kind -> Maybe (Kind, Kind)
    splitFunTy TyConMap
    m Kind
    opTy' of
        Just (Kind
    _,Kind
    resTy) -> Kind -> [Either p Kind] -> Kind
    go Kind
    resTy [Either p Kind]
    args'
        Maybe (Kind, Kind)
    _ -> String -> Kind
    forall a. HasCallStack => String -> a
    error [I.i|
            Unexpected application. The term

              #{showPpr e}

            applied an argument

              #{showPpr a}

            to something with the non-function type

              #{showPpr opTy'}
          |]

      goTyArgs :: Kind -> [Kind] -> [Either p Kind] -> Kind
    goTyArgs Kind
    opTy' [Kind]
    revTys (Right Kind
    ty:[Either p Kind]
    args') = Kind -> [Kind] -> [Either p Kind] -> Kind
    goTyArgs Kind
    opTy' (Kind
    tyKind -> [Kind] -> [Kind]
    forall a. a -> [a] -> [a]
    :[Kind]
    revTys) [Either p Kind]
    args'
      goTyArgs Kind
    opTy' [Kind]
    revTys [Either p Kind]
    args'            = Kind -> [Either p Kind] -> Kind
    go (HasCallStack => TyConMap -> Kind -> [Kind] -> Kind
    TyConMap -> Kind -> [Kind] -> Kind
    piResultTys TyConMap
    m Kind
    opTy' ([Kind] -> [Kind]
    forall a. [a] -> [a]
    reverse [Kind]
    revTys)) [Either p Kind]
    args'

    -- | Like 'piResultTys', but only applies a single type. If multiple types are
    -- being applied use 'piResultTys', as it is more efficient to only substitute
    -- once with many types.
    piResultTy
      :: HasCallStack
      => TyConMap
      -> Type
      -> Type
      -> Type
    piResultTy :: HasCallStack => TyConMap -> Kind -> Kind -> Kind
    piResultTy TyConMap
    m Kind
    ty Kind
    arg =
      HasCallStack => TyConMap -> Kind -> [Kind] -> Kind
    TyConMap -> Kind -> [Kind] -> Kind
    piResultTys TyConMap
    m Kind
    ty [Kind
    arg]

    -- | @(piResultTys f_ty [ty1, ..., tyn])@ gives the type of @(f ty1 .. tyn)@
    -- where @f :: f_ty@
    --
    -- 'piResultTys' is interesting because:
    --
    --    1. 'f_ty' may have more foralls than there are args
    --    2. Less obviously, it may have fewer foralls
    --
    -- Fore case 2. think of:
    --
    --   piResultTys (forall a . a) [forall b.b, Int]
    --
    -- This really can happen, such as situations involving 'undefined's type:
    --
    --   undefined :: forall a. a
    --
    --   undefined (forall b. b -> b) Int
    --
    -- This term should have the type @(Int -> Int)@, but notice that there are
    -- more type args than foralls in 'undefined's type.
    --
    -- For efficiency reasons, when there are no foralls, we simply drop arrows from
    -- a function type/kind.
    piResultTys
      :: HasCallStack
      => TyConMap
      -> Type
      -> [Type]
      -> Type
    piResultTys :: HasCallStack => TyConMap -> Kind -> [Kind] -> Kind
    piResultTys TyConMap
    _ Kind
    ty [] = Kind
    ty
    piResultTys TyConMap
    m Kind
    ty origArgs :: [Kind]
    origArgs@(Kind
    arg:[Kind]
    args)
      | Just Kind
    ty' <- TyConMap -> Kind -> Maybe Kind
    coreView1 TyConMap
    m Kind
    ty
      = HasCallStack => TyConMap -> Kind -> [Kind] -> Kind
    TyConMap -> Kind -> [Kind] -> Kind
    piResultTys TyConMap
    m Kind
    ty' [Kind]
    origArgs
      | FunTy Kind
    a Kind
    res <- Kind -> TypeView
    tyView Kind
    ty
      -- TODO coreView is used here because the partial evaluator will sometimes
      -- encounter / not encounter a Signal as an argument unexpectedly. When PR
      -- #1064 is merged the coreView calls should be removed again.
      = if Bool
    debugIsOn Bool -> Bool -> Bool
    && Bool -> Bool
    not (Kind -> Kind -> Bool
    aeqType (TyConMap -> Kind -> Kind
    coreView TyConMap
    m Kind
    a) (TyConMap -> Kind -> Kind
    coreView TyConMap
    m Kind
    arg)) then String -> Kind
    forall a. HasCallStack => String -> a
    error [I.i|
          Unexpected application. A function with type:

            #{showPpr ty}

          Got applied to an argument of type:

            #{showPpr arg}
        |]
        else
          HasCallStack => TyConMap -> Kind -> [Kind] -> Kind
    TyConMap -> Kind -> [Kind] -> Kind
    piResultTys TyConMap
    m Kind
    res [Kind]
    args
      | ForAllTy TyVar
    tv Kind
    res <- Kind
    ty
      = VarEnv Kind -> Kind -> [Kind] -> Kind
    go (TyVar -> Kind -> VarEnv Kind -> VarEnv Kind
    forall b a. Var b -> a -> VarEnv a -> VarEnv a
    extendVarEnv TyVar
    tv Kind
    arg VarEnv Kind
    forall a. VarEnv a
    emptyVarEnv) Kind
    res [Kind]
    args
      | Bool
    otherwise
      = String -> Doc ClashAnnotation -> Kind
    forall ann a. String -> Doc ann -> a
    pprPanic String
    "piResultTys1" (Kind -> Doc ClashAnnotation
    forall p. PrettyPrec p => p -> Doc ClashAnnotation
    ppr Kind
    ty Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
    forall a. Semigroup a => a -> a -> a
    <> Doc ClashAnnotation
    forall ann. Doc ann
    line Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
    forall a. Semigroup a => a -> a -> a
    <> [Kind] -> Doc ClashAnnotation
    forall p. PrettyPrec p => p -> Doc ClashAnnotation
    ppr [Kind]
    origArgs)
     where
      inScope :: InScopeSet
    inScope = VarSet -> InScopeSet
    mkInScopeSet ([Kind] -> VarSet
    forall a. HasFreeVars a => a -> VarSet
    freeVarsOf (Kind
    tyKind -> [Kind] -> [Kind]
    forall a. a -> [a] -> [a]
    :[Kind]
    origArgs))

      go :: VarEnv Kind -> Kind -> [Kind] -> Kind
    go VarEnv Kind
    env Kind
    ty' [] = HasCallStack => Subst -> Kind -> Kind
    Subst -> Kind -> Kind
    substTy (InScopeSet -> VarEnv Kind -> Subst
    mkTvSubst InScopeSet
    inScope VarEnv Kind
    env) Kind
    ty'
      go VarEnv Kind
    env Kind
    ty' allArgs :: [Kind]
    allArgs@(Kind
    arg':[Kind]
    args')
        | Just Kind
    ty'' <- TyConMap -> Kind -> Maybe Kind
    coreView1 TyConMap
    m Kind
    ty'
        = VarEnv Kind -> Kind -> [Kind] -> Kind
    go VarEnv Kind
    env Kind
    ty'' [Kind]
    allArgs
        | FunTy Kind
    _ Kind
    res <- Kind -> TypeView
    tyView Kind
    ty'
        = VarEnv Kind -> Kind -> [Kind] -> Kind
    go VarEnv Kind
    env Kind
    res [Kind]
    args'
        | ForAllTy TyVar
    tv Kind
    res <- Kind
    ty'
        = VarEnv Kind -> Kind -> [Kind] -> Kind
    go (TyVar -> Kind -> VarEnv Kind -> VarEnv Kind
    forall b a. Var b -> a -> VarEnv a -> VarEnv a
    extendVarEnv TyVar
    tv Kind
    arg' VarEnv Kind
    env) Kind
    res [Kind]
    args'
        | VarTy TyVar
    tv <- Kind
    ty'
        , Just Kind
    ty'' <- TyVar -> VarEnv Kind -> Maybe Kind
    forall b a. Var b -> VarEnv a -> Maybe a
    lookupVarEnv TyVar
    tv VarEnv Kind
    env
          -- Deals with (piResultTys  (forall a.a) [forall b.b, Int])
        = HasCallStack => TyConMap -> Kind -> [Kind] -> Kind
    TyConMap -> Kind -> [Kind] -> Kind
    piResultTys TyConMap
    m Kind
    ty'' [Kind]
    allArgs
        | Bool
    otherwise
        = String -> Doc ClashAnnotation -> Kind
    forall ann a. String -> Doc ann -> a
    pprPanic String
    "piResultTys2" (Kind -> Doc ClashAnnotation
    forall p. PrettyPrec p => p -> Doc ClashAnnotation
    ppr Kind
    ty' Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
    forall a. Semigroup a => a -> a -> a
    <> Doc ClashAnnotation
    forall ann. Doc ann
    line Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
    forall a. Semigroup a => a -> a -> a
    <> [Kind] -> Doc ClashAnnotation
    forall p. PrettyPrec p => p -> Doc ClashAnnotation
    ppr [Kind]
    origArgs Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
    forall a. Semigroup a => a -> a -> a
    <> Doc ClashAnnotation
    forall ann. Doc ann
    line Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
    forall a. Semigroup a => a -> a -> a
    <> [Kind] -> Doc ClashAnnotation
    forall p. PrettyPrec p => p -> Doc ClashAnnotation
    ppr [Kind]
    allArgs)