{-# LANGUAGE QuasiQuotes #-}

module Clash.Core.TermInfo where

import Data.Text.Prettyprint.Doc (line)

import Clash.Core.DataCon
import Clash.Core.FreeVars
import Clash.Core.Literal
import Clash.Core.Pretty
import Clash.Core.Subst
import Clash.Core.Term
import Clash.Core.TyCon (TyConMap)
import Clash.Core.Type
import Clash.Core.Var
import Clash.Core.VarEnv
import Clash.Debug (debugIsOn)
import Clash.Util
import Clash.Util.Interpolate as I

termSize :: Term -> Word
termSize :: Term -> Word
termSize (Var {})     = Word
1
termSize (Data {})    = Word
1
termSize (Literal {}) = Word
1
termSize (Prim {})    = Word
1
termSize (Lam Id
_ Term
e)    = Term -> Word
termSize Term
e Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1
termSize (TyLam TyVar
_ Term
e)  = Term -> Word
termSize Term
e
termSize (App Term
e1 Term
e2)  = Term -> Word
termSize Term
e1 Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Term -> Word
termSize Term
e2
termSize (TyApp Term
e Type
_)  = Term -> Word
termSize Term
e
termSize (Cast Term
e Type
_ Type
_) = Term -> Word
termSize Term
e
termSize (Tick TickInfo
_ Term
e)   = Term -> Word
termSize Term
e
termSize (Letrec [LetBinding]
bndrs Term
e) = [Word] -> Word
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
sum (Word
bodySzWord -> [Word] -> [Word]
forall a. a -> [a] -> [a]
:[Word]
bndrSzs)
 where
  bndrSzs :: [Word]
bndrSzs = (LetBinding -> Word) -> [LetBinding] -> [Word]
forall a b. (a -> b) -> [a] -> [b]
map (Term -> Word
termSize (Term -> Word) -> (LetBinding -> Term) -> LetBinding -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LetBinding -> Term
forall a b. (a, b) -> b
snd) [LetBinding]
bndrs
  bodySz :: Word
bodySz  = Term -> Word
termSize Term
e
termSize (Case Term
subj Type
_ [Alt]
alts) = [Word] -> Word
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
sum (Word
subjSzWord -> [Word] -> [Word]
forall a. a -> [a] -> [a]
:[Word]
altSzs)
 where
  subjSz :: Word
subjSz = Term -> Word
termSize Term
subj
  altSzs :: [Word]
altSzs = (Alt -> Word) -> [Alt] -> [Word]
forall a b. (a -> b) -> [a] -> [b]
map (Term -> Word
termSize (Term -> Word) -> (Alt -> Term) -> Alt -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alt -> Term
forall a b. (a, b) -> b
snd) [Alt]
alts

-- | Determine the type of a term
termType :: TyConMap -> Term -> Type
termType :: TyConMap -> Term -> Type
termType TyConMap
m Term
e = case Term
e of
  Var Id
t          -> Id -> Type
forall a. Var a -> Type
varType Id
t
  Data DataCon
dc        -> DataCon -> Type
dcType DataCon
dc
  Literal Literal
l      -> Literal -> Type
literalType Literal
l
  Prim PrimInfo
t         -> PrimInfo -> Type
primType PrimInfo
t
  Lam Id
v Term
e'       -> Type -> Type -> Type
mkFunTy (Id -> Type
forall a. Var a -> Type
varType Id
v) (TyConMap -> Term -> Type
termType TyConMap
m Term
e')
  TyLam TyVar
tv Term
e'    -> TyVar -> Type -> Type
ForAllTy TyVar
tv (TyConMap -> Term -> Type
termType TyConMap
m Term
e')
  App Term
_ Term
_        -> case Term -> (Term, [Either Term Type])
collectArgs Term
e of
                      (Term
fun, [Either Term Type]
args) -> Term -> TyConMap -> Type -> [Either Term Type] -> Type
applyTypeToArgs Term
e TyConMap
m (TyConMap -> Term -> Type
termType TyConMap
m Term
fun) [Either Term Type]
args
  TyApp Term
_ Type
_      -> case Term -> (Term, [Either Term Type])
collectArgs Term
e of
                      (Term
fun, [Either Term Type]
args) -> Term -> TyConMap -> Type -> [Either Term Type] -> Type
applyTypeToArgs Term
e TyConMap
m (TyConMap -> Term -> Type
termType TyConMap
m Term
fun) [Either Term Type]
args
  Letrec [LetBinding]
_ Term
e'    -> TyConMap -> Term -> Type
termType TyConMap
m Term
e'
  Case Term
_ Type
ty [Alt]
_    -> Type
ty
  Cast Term
_ Type
_ Type
ty2   -> Type
ty2
  Tick TickInfo
_ Term
e'      -> TyConMap -> Term -> Type
termType TyConMap
m Term
e'

-- | Get the result type of a polymorphic function given a list of arguments
applyTypeToArgs
  :: Term
  -> TyConMap
  -> Type
  -> [Either Term Type]
  -> Type
applyTypeToArgs :: Term -> TyConMap -> Type -> [Either Term Type] -> Type
applyTypeToArgs Term
e TyConMap
m Type
opTy [Either Term Type]
args = Type -> [Either Term Type] -> Type
forall a. Type -> [Either a Type] -> Type
go Type
opTy [Either Term Type]
args
 where
  go :: Type -> [Either a Type] -> Type
go Type
opTy' []               = Type
opTy'
  go Type
opTy' (Right Type
ty:[Either a Type]
args') = Type -> [Type] -> [Either a Type] -> Type
goTyArgs Type
opTy' [Type
ty] [Either a Type]
args'
  go Type
opTy' (Left a
_:[Either a Type]
args')   = case TyConMap -> Type -> Maybe (Type, Type)
splitFunTy TyConMap
m Type
opTy' of
    Just (Type
_,Type
resTy) -> Type -> [Either a Type] -> Type
go Type
resTy [Either a Type]
args'
    Maybe (Type, Type)
_ -> [Char] -> Type
forall a. HasCallStack => [Char] -> a
error ([Char] -> Type) -> [Char] -> Type
forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Char]
unlines [[Char]
"applyTypeToArgs:"
                         ,[Char]
"Expression: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Term -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Term
e
                         ,[Char]
"Type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type
opTy
                         ,[Char]
"Args: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [[Char]] -> [Char]
unlines ((Either Term Type -> [Char]) -> [Either Term Type] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map ((Term -> [Char]) -> (Type -> [Char]) -> Either Term Type -> [Char]
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either Term -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr Type -> [Char]
forall p. PrettyPrec p => p -> [Char]
showPpr) [Either Term Type]
args)
                         ]

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

-- | Like 'piResultTyMaybe', but errors out when a type application is not
-- valid.
--
-- Do not iterate 'piResultTy', because it's inefficient to substitute one
-- variable at a time; instead use 'piResultTys'
piResultTy
  :: HasCallStack
  => TyConMap
  -> Type
  -> Type
  -> Type
piResultTy :: TyConMap -> Type -> Type -> Type
piResultTy TyConMap
m Type
ty Type
arg = case HasCallStack => TyConMap -> Type -> Type -> Maybe Type
TyConMap -> Type -> Type -> Maybe Type
piResultTyMaybe TyConMap
m Type
ty Type
arg of
  Just Type
res -> Type
res
  Maybe Type
Nothing  -> [Char] -> Doc ClashAnnotation -> Type
forall ann a. [Char] -> Doc ann -> a
pprPanic [Char]
"piResultTy" (Type -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr Type
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
<> Type -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr Type
arg)

-- | Like 'piResultTys' but for a single argument.
--
-- Do not iterate 'piResultTyMaybe', because it's inefficient to substitute one
-- variable at a time; instead use 'piResultTys'
piResultTyMaybe
  :: HasCallStack
  => TyConMap
  -> Type
  -> Type
  -> Maybe Type
piResultTyMaybe :: TyConMap -> Type -> Type -> Maybe Type
piResultTyMaybe TyConMap
m Type
ty Type
arg
  | Just Type
ty' <- TyConMap -> Type -> Maybe Type
coreView1 TyConMap
m Type
ty
  = HasCallStack => TyConMap -> Type -> Type -> Maybe Type
TyConMap -> Type -> Type -> Maybe Type
piResultTyMaybe TyConMap
m Type
ty' Type
arg
  | FunTy Type
a Type
res <- Type -> TypeView
tyView Type
ty
  = if Bool
debugIsOn Bool -> Bool -> Bool
&& Bool -> Bool
not (Type -> Type -> Bool
aeqType Type
a Type
arg) then [Char] -> Maybe Type
forall a. HasCallStack => [Char] -> a
error [I.i|
      Unexpected application. A function with type:

        #{showPpr ty}

      Got applied to an argument of type:

        #{showPpr arg}
    |]
    else
      Type -> Maybe Type
forall a. a -> Maybe a
Just Type
res
  | ForAllTy TyVar
tv Type
res <- Type
ty
  = let emptySubst :: Subst
emptySubst = InScopeSet -> Subst
mkSubst (VarSet -> InScopeSet
mkInScopeSet ([Type] -> VarSet
forall (f :: Type -> Type). Foldable f => f Type -> VarSet
tyFVsOfTypes [Type
arg,Type
res]))
    in  Type -> Maybe Type
forall a. a -> Maybe a
Just (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy (Subst -> TyVar -> Type -> Subst
extendTvSubst Subst
emptySubst TyVar
tv Type
arg) Type
res)
  | Bool
otherwise
  = Maybe Type
forall a. Maybe a
Nothing

-- | @(piResultTys f_ty [ty1, ..., tyn])@ give sthe 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 -> Type -> [Type] -> Type
piResultTys TyConMap
_ Type
ty [] = Type
ty
piResultTys TyConMap
m Type
ty origArgs :: [Type]
origArgs@(Type
arg:[Type]
args)
  | Just Type
ty' <- TyConMap -> Type -> Maybe Type
coreView1 TyConMap
m Type
ty
  = HasCallStack => TyConMap -> Type -> [Type] -> Type
TyConMap -> Type -> [Type] -> Type
piResultTys TyConMap
m Type
ty' [Type]
origArgs
  | FunTy Type
a Type
res <- Type -> TypeView
tyView Type
ty
  = if Bool
debugIsOn Bool -> Bool -> Bool
&& Bool -> Bool
not (Type -> Type -> Bool
aeqType Type
a Type
arg) then [Char] -> Type
forall a. HasCallStack => [Char] -> a
error [I.i|
      Unexpected application. A function with type:

        #{showPpr ty}

      Got applied to an argument of type:

        #{showPpr arg}
    |]
    else
      HasCallStack => TyConMap -> Type -> [Type] -> Type
TyConMap -> Type -> [Type] -> Type
piResultTys TyConMap
m Type
res [Type]
args
  | ForAllTy TyVar
tv Type
res <- Type
ty
  = VarEnv Type -> Type -> [Type] -> Type
go (TyVar -> Type -> VarEnv Type -> VarEnv Type
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv TyVar
tv Type
arg VarEnv Type
forall a. VarEnv a
emptyVarEnv) Type
res [Type]
args
  | Bool
otherwise
  = [Char] -> Doc ClashAnnotation -> Type
forall ann a. [Char] -> Doc ann -> a
pprPanic [Char]
"piResultTys1" (Type -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr Type
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
<> [Type] -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr [Type]
origArgs)
 where
  inScope :: InScopeSet
inScope = VarSet -> InScopeSet
mkInScopeSet ([Type] -> VarSet
forall (f :: Type -> Type). Foldable f => f Type -> VarSet
tyFVsOfTypes (Type
tyType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
origArgs))

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

-- | Does a term have a function type?
isFun :: TyConMap -> Term -> Bool
isFun :: TyConMap -> Term -> Bool
isFun TyConMap
m Term
t = TyConMap -> Type -> Bool
isFunTy TyConMap
m (TyConMap -> Term -> Type
termType TyConMap
m Term
t)

-- | Does a term have a function or polymorphic type?
isPolyFun :: TyConMap -> Term -> Bool
isPolyFun :: TyConMap -> Term -> Bool
isPolyFun TyConMap
m Term
t = TyConMap -> Type -> Bool
isPolyFunCoreTy TyConMap
m (TyConMap -> Term -> Type
termType TyConMap
m Term
t)

-- | Is a term a term-abstraction?
isLam :: Term -> Bool
isLam :: Term -> Bool
isLam (Lam {}) = Bool
True
isLam Term
_        = Bool
False

-- | Is a term a recursive let-binding?
isLet :: Term -> Bool
isLet :: Term -> Bool
isLet (Letrec {}) = Bool
True
isLet Term
_           = Bool
False

-- | Is a term a variable reference?
isVar :: Term -> Bool
isVar :: Term -> Bool
isVar (Var {}) = Bool
True
isVar Term
_        = Bool
False

isLocalVar :: Term -> Bool
isLocalVar :: Term -> Bool
isLocalVar (Var Id
v) = Id -> Bool
forall a. Var a -> Bool
isLocalId Id
v
isLocalVar Term
_ = Bool
False

-- | Is a term a datatype constructor?
isCon :: Term -> Bool
isCon :: Term -> Bool
isCon (Data {}) = Bool
True
isCon Term
_         = Bool
False

-- | Is a term a primitive?
isPrim :: Term -> Bool
isPrim :: Term -> Bool
isPrim (Prim {}) = Bool
True
isPrim Term
_         = Bool
False