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

import qualified Data.Text as Text (isInfixOf)

#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.Name (Name(nameOcc))
import Clash.Core.Pretty
import Clash.Core.Subst
import Clash.Core.Term (Term(..), IsMultiPrim(..), PrimInfo(..), collectArgs)
import Clash.Core.TyCon (TyCon(tyConKind), TyConMap)
import Clash.Core.Type
import Clash.Core.TysPrim
import Clash.Core.Var (Var(varType))
import Clash.Core.VarEnv
import Clash.Debug (debugIsOn)
import Clash.Unique (lookupUniqMap')
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 :: 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
    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
        , Text -> Text -> Bool
Text.isInfixOf Text
"GHC.Tuple.(" (TyConName -> Text
forall a. Name a -> Text
nameOcc 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 (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 :: 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 (TyConMap -> TyConName -> TyCon
forall a b. (HasCallStack, Uniquable a) => UniqMap b -> a -> b
lookupUniqMap' TyConMap
tcm TyConName
tc)) [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
$ $(String
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
          AnnType [Attr']
_ 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 :: 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 :: 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)