{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}

module Clash.Core.TermInfo where

import Data.Maybe (fromMaybe)
import GHC.Stack (HasCallStack)

import Clash.Core.HasType
import Clash.Core.Term
import Clash.Core.TyCon (tyConDataCons, isTupleTyConLike, TyConMap)
import Clash.Core.Type
import Clash.Core.Var
import qualified Clash.Data.UniqMap as UniqMap
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 (Let (NonRec Id
_ Term
x) Term
e) = Term -> Word
termSize Term
x Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Term -> Word
termSize Term
e
termSize (Let (Rec [(Id, Term)]
xs) 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 = ((Id, Term) -> Word) -> [(Id, Term)] -> [Word]
forall a b. (a -> b) -> [a] -> [b]
map (Term -> Word
termSize (Term -> Word) -> ((Id, Term) -> Term) -> (Id, Term) -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id, Term) -> Term
forall a b. (a, b) -> b
snd) [(Id, Term)]
xs
  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

multPrimErr :: PrimInfo -> String
multPrimErr :: PrimInfo -> String
multPrimErr PrimInfo
primInfo =  [I.i|
  Internal error in multiPrimInfo': could not produce MultiPrimInfo. This
  probably means a multi result blackbox's result type was not a tuple.
  PrimInfo:

    #{primInfo}
|]

splitMultiPrimArgs ::
  HasCallStack =>
  MultiPrimInfo ->
  [Either Term Type] ->
  ([Either Term Type], [Id])
splitMultiPrimArgs :: MultiPrimInfo -> [Either Term Type] -> ([Either Term Type], [Id])
splitMultiPrimArgs MultiPrimInfo{[Type]
mpi_resultTypes :: MultiPrimInfo -> [Type]
mpi_resultTypes :: [Type]
mpi_resultTypes} [Either Term Type]
args0 = ([Either Term Type]
args1, [Id]
resArgs1)
 where
  resArgs1 :: [Id]
resArgs1 = [Id
id_ | Left (Var Id
id_) <- [Either Term Type]
resArgs0]
  ([Either Term Type]
args1, [Either Term Type]
resArgs0) = Int
-> [Either Term Type] -> ([Either Term Type], [Either Term Type])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Either Term Type] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Either Term Type]
args0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Type] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Type]
mpi_resultTypes) [Either Term Type]
args0

-- | Same as 'multiPrimInfo', but produced an error if it could not produce a
-- 'MultiPrimInfo'.
multiPrimInfo' :: HasCallStack => TyConMap -> PrimInfo -> MultiPrimInfo
multiPrimInfo' :: TyConMap -> PrimInfo -> MultiPrimInfo
multiPrimInfo' TyConMap
tcm PrimInfo
primInfo =
  MultiPrimInfo -> Maybe MultiPrimInfo -> MultiPrimInfo
forall a. a -> Maybe a -> a
fromMaybe (String -> MultiPrimInfo
forall a. HasCallStack => String -> a
error (PrimInfo -> String
multPrimErr PrimInfo
primInfo)) (TyConMap -> PrimInfo -> Maybe MultiPrimInfo
multiPrimInfo TyConMap
tcm PrimInfo
primInfo)

-- | Produce 'MutliPrimInfo' for given primitive
multiPrimInfo :: TyConMap -> PrimInfo -> Maybe MultiPrimInfo
multiPrimInfo :: TyConMap -> PrimInfo -> Maybe MultiPrimInfo
multiPrimInfo TyConMap
tcm PrimInfo
primInfo
  | ([Either TyVar Type]
_primArgs, Type
primResTy) <- Type -> ([Either TyVar Type], Type)
splitFunForallTy (PrimInfo -> Type
primType PrimInfo
primInfo)
  , TyConApp TyConName
tupTcNm [Type]
tupEls <- Type -> TypeView
tyView Type
primResTy
    -- XXX: Hardcoded for tuples
  , TyConName -> Bool
isTupleTyConLike TyConName
tupTcNm
  , Just TyCon
tupTc <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
tupTcNm TyConMap
tcm
  , [DataCon
tupDc] <- TyCon -> [DataCon]
tyConDataCons TyCon
tupTc
  = MultiPrimInfo -> Maybe MultiPrimInfo
forall a. a -> Maybe a
Just (MultiPrimInfo -> Maybe MultiPrimInfo)
-> MultiPrimInfo -> Maybe MultiPrimInfo
forall a b. (a -> b) -> a -> b
$ MultiPrimInfo :: PrimInfo -> DataCon -> [Type] -> MultiPrimInfo
MultiPrimInfo
    { mpi_primInfo :: PrimInfo
mpi_primInfo = PrimInfo
primInfo
    , mpi_resultDc :: DataCon
mpi_resultDc = DataCon
tupDc
    , mpi_resultTypes :: [Type]
mpi_resultTypes = [Type]
tupEls }
multiPrimInfo TyConMap
_ PrimInfo
_ = Maybe MultiPrimInfo
forall a. Maybe a
Nothing

-- | 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
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf 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
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
m Term
t)

-- | Is a term a recursive let-binding?
isLet :: Term -> Bool
isLet :: Term -> Bool
isLet Let{} = 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

-- | Is a term a tick?
isTick :: Term -> Bool
isTick :: Term -> Bool
isTick Tick{} = Bool
True
isTick Term
_ = Bool
False

-- | Is a term a cast?
isCast :: Term -> Bool
isCast :: Term -> Bool
isCast (Cast {}) = Bool
True
isCast Term
_         = Bool
False