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

module Clash.Core.TermInfo where

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

import Clash.Core.HasType
import Clash.Core.Name
import Clash.Core.Term
import Clash.Core.TyCon (tyConDataCons, TyConMap)
import Clash.Core.Type
import Clash.Core.Var
import Clash.Unique (lookupUniqMap)
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
  , Text
"GHC.Tuple.(," Text -> Text -> Bool
`isInfixOf` TyConName -> Text
forall a. Name a -> Text
nameOcc TyConName
tupTcNm
  , Just TyCon
tupTc <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap 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