    {-# 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 a. Num a => [a] -> a
    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 a. Num a => [a] -> a
    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 :: HasCallStack =>
    MultiPrimInfo -> [Either Term Type] -> ([Either Term Type], [Id])
    splitMultiPrimArgs MultiPrimInfo{[Type]
    mpi_resultTypes :: [Type]
    mpi_resultTypes :: MultiPrimInfo -> [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 a. [a] -> 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 a. [a] -> 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' :: HasCallStack => 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
        { 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