{-# 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
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
applyTypeToArgs
:: Term
-> 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'
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
:: 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
= 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
= 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)