    {-|
      Copyright  :  (C) 2012-2016, University of Twente,
                        2016-2017, Myrtle Software Ltd,
                        2017-2018, Google Inc.
                        2022     , QBayLogic B.V.
      License    :  BSD2 (see the file LICENSE)
      Maintainer :  QBayLogic B.V. <devops@qbaylogic.com>

      Transformations on primitives with multiple results.
    -}

    {-# LANGUAGE NamedFieldPuns #-}
    {-# LANGUAGE OverloadedStrings #-}

    module Clash.Normalize.Transformations.MultiPrim
      ( setupMultiResultPrim
      ) where

    import qualified Control.Lens as Lens
    import qualified Data.Either as Either
    import Data.Text.Extra (showt)
    import GHC.Stack (HasCallStack)

    import Clash.Annotations.Primitive (extractPrim)
    import Clash.Core.Name (mkUnsafeInternalName)
    import Clash.Core.Term
      ( IsMultiPrim(..), MultiPrimInfo(..), PrimInfo(..), Term(..), WorkInfo(..)
      , mkAbstraction, mkApps, mkTmApps, mkTyApps, PrimUnfolding(..))
    import Clash.Core.TermInfo (multiPrimInfo')
    import Clash.Core.TyCon (TyConMap)
    import Clash.Core.Type (Type(..), mkPolyFunTy, splitFunForallTy)
    import Clash.Core.Util (listToLets)
    import Clash.Core.Var (mkLocalId)
    import Clash.Normalize.Types (NormRewrite)
    import Clash.Primitives.Types (Primitive(..))
    import Clash.Rewrite.Types (tcCache, primitives)
    import Clash.Rewrite.Util (changed)

    -- Note [MultiResult type]
    --
    -- A multi result primitive assigns its results to multiple result variables
    -- instead of one. Besides producing nicer HDL it works around issues with
    -- synthesis tooling described in:
    --
    --   https://github.com/clash-lang/clash-compiler/issues/1555
    --
    -- This transformation rewrites primitives indicating they can assign their
    -- results to multiple signals, such that netlist can easily render it. This
    -- involves inserting additional arguments for each of the result values, and
    -- then using the c$multiPrimSelect primitive to select individual results.
    --
    -- Example:
    --
    -- @
    -- prim :: forall a b c. a -> (b, c)
    -- @
    --
    -- will be rewritten to:
    --
    -- @
    --   \(x :: a) ->
    --         let
    --            r  = prim @a @b @c x r0 r1 -- With 'Clash.Core.Term.MultiPrim'
    --            r0 = c$multiPrimSelect r0 r
    --            r1 = c$multiPrimSelect r1 r
    --          in
    --            (r0, r1)
    -- @
    --
    -- Netlist will not render any @multiPrimSelect@ primitives. Similar to
    -- primitives having a /void/ return type, /r/ is not rendered either.
    --
    -- This transformation is currently hardcoded to recognize tuples as return
    -- types, not any product type. It will error if it sees a multi result primitive
    -- with a non-tuple return type.
    --
    setupMultiResultPrim :: HasCallStack => NormRewrite
    setupMultiResultPrim :: HasCallStack => NormRewrite
    setupMultiResultPrim TransformContext
    _ctx e :: Term
    e@(Prim pInfo :: PrimInfo
    pInfo@PrimInfo{primMultiResult :: PrimInfo -> IsMultiPrim
    primMultiResult=IsMultiPrim
    SingleResult}) = do
      TyConMap
    tcm <- Getting TyConMap RewriteEnv TyConMap
    -> RewriteMonad NormalizeState TyConMap
    forall s (m :: Type -> Type) a.
    MonadReader s m =>
    Getting a s a -> m a
    Lens.view Getting TyConMap RewriteEnv TyConMap
    Getter RewriteEnv TyConMap
    tcCache
      Maybe
      (PrimitiveGuard
         (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction)))
    prim <- Getting
      (Maybe
         (PrimitiveGuard
            (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))))
      RewriteEnv
      (Maybe
         (PrimitiveGuard
            (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))))
    -> RewriteMonad
         NormalizeState
         (Maybe
            (PrimitiveGuard
               (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))))
    forall s (m :: Type -> Type) a.
    MonadReader s m =>
    Getting a s a -> m a
    Lens.view ((CompiledPrimMap
     -> Const
          (Maybe
             (PrimitiveGuard
                (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))))
          CompiledPrimMap)
    -> RewriteEnv
    -> Const
         (Maybe
            (PrimitiveGuard
               (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))))
         RewriteEnv
    Getter RewriteEnv CompiledPrimMap
    primitives ((CompiledPrimMap
      -> Const
           (Maybe
              (PrimitiveGuard
                 (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))))
           CompiledPrimMap)
     -> RewriteEnv
     -> Const
          (Maybe
             (PrimitiveGuard
                (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))))
          RewriteEnv)
    -> ((Maybe
           (PrimitiveGuard
              (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction)))
         -> Const
              (Maybe
                 (PrimitiveGuard
                    (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))))
              (Maybe
                 (PrimitiveGuard
                    (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction)))))
        -> CompiledPrimMap
        -> Const
             (Maybe
                (PrimitiveGuard
                   (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))))
             CompiledPrimMap)
    -> Getting
         (Maybe
            (PrimitiveGuard
               (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))))
         RewriteEnv
         (Maybe
            (PrimitiveGuard
               (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Index CompiledPrimMap
    -> Lens' CompiledPrimMap (Maybe (IxValue CompiledPrimMap))
    forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
    Lens.at (PrimInfo -> Text
    primName PrimInfo
    pInfo))

      case Maybe
      (PrimitiveGuard
         (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction)))
    prim Maybe
      (PrimitiveGuard
         (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction)))
    -> (PrimitiveGuard
          (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))
        -> Maybe
             (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction)))
    -> Maybe
         (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))
    forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
    forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
    >>= PrimitiveGuard
      (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))
    -> Maybe
         (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))
    forall a. PrimitiveGuard a -> Maybe a
    extractPrim of
        Just (BlackBoxHaskell{multiResult :: forall a b c d. Primitive a b c d -> Bool
    multiResult=Bool
    True}) ->
          Term -> RewriteMonad NormalizeState Term
    forall a extra. a -> RewriteMonad extra a
    changed (HasCallStack => TyConMap -> PrimInfo -> Term
    TyConMap -> PrimInfo -> Term
    setupMultiResultPrim' TyConMap
    tcm PrimInfo
    pInfo)
        Just (BlackBox{multiResult :: forall a b c d. Primitive a b c d -> Bool
    multiResult=Bool
    True}) ->
          Term -> RewriteMonad NormalizeState Term
    forall a extra. a -> RewriteMonad extra a
    changed (HasCallStack => TyConMap -> PrimInfo -> Term
    TyConMap -> PrimInfo -> Term
    setupMultiResultPrim' TyConMap
    tcm PrimInfo
    pInfo)
        Maybe
      (Primitive BlackBoxTemplate BlackBox () (Int, BlackBoxFunction))
    _ ->
          Term -> RewriteMonad NormalizeState Term
    forall a. a -> RewriteMonad NormalizeState a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return Term
    e

    setupMultiResultPrim TransformContext
    _ Term
    e = Term -> RewriteMonad NormalizeState Term
    forall a. a -> RewriteMonad NormalizeState a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return Term
    e

    setupMultiResultPrim' :: HasCallStack => TyConMap -> PrimInfo -> Term
    setupMultiResultPrim' :: HasCallStack => TyConMap -> PrimInfo -> Term
    setupMultiResultPrim' TyConMap
    tcm primInfo :: PrimInfo
    primInfo@PrimInfo{Type
    primType :: Type
    primType :: PrimInfo -> Type
    primType} =
      Term -> [Either Id TyVar] -> Term
    mkAbstraction Term
    letTerm ((TyVar -> Either Id TyVar) -> [TyVar] -> [Either Id TyVar]
    forall a b. (a -> b) -> [a] -> [b]
    map TyVar -> Either Id TyVar
    forall a b. b -> Either a b
    Right [TyVar]
    typeVars [Either Id TyVar] -> [Either Id TyVar] -> [Either Id TyVar]
    forall a. Semigroup a => a -> a -> a
    <> (Id -> Either Id TyVar) -> [Id] -> [Either Id TyVar]
    forall a b. (a -> b) -> [a] -> [b]
    map Id -> Either Id TyVar
    forall a b. a -> Either a b
    Left [Id]
    argIds)
     where
      typeVars :: [TyVar]
    typeVars = [Either TyVar Type] -> [TyVar]
    forall a b. [Either a b] -> [a]
    Either.lefts [Either TyVar Type]
    pArgs

      internalNm :: Text -> Int -> Name a
    internalNm Text
    prefix Int
    n = Text -> Int -> Name a
    forall a. Text -> Int -> Name a
    mkUnsafeInternalName (Text
    prefix Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Int -> Text
    forall a. Show a => a -> Text
    showt Int
    n) Int
    n
      internalId :: Text -> Type -> Int -> Id
    internalId Text
    prefix Type
    typ Int
    n = Type -> TmName -> Id
    mkLocalId Type
    typ (Text -> Int -> TmName
    forall a. Text -> Int -> Name a
    internalNm Text
    prefix Int
    n)

      nTermArgs :: Int
    nTermArgs = [Type] -> Int
    forall a. [a] -> Int
    forall (t :: Type -> Type) a. Foldable t => t a -> Int
    length ([Either TyVar Type] -> [Type]
    forall a b. [Either a b] -> [b]
    Either.rights [Either TyVar Type]
    pArgs)
      argIds :: [Id]
    argIds = (Type -> Int -> Id) -> [Type] -> [Int] -> [Id]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith (Text -> Type -> Int -> Id
    internalId Text
    "a") ([Either TyVar Type] -> [Type]
    forall a b. [Either a b] -> [b]
    Either.rights [Either TyVar Type]
    pArgs) [Int
    1..Int
    nTermArgs]
      resIds :: [Id]
    resIds = (Type -> Int -> Id) -> [Type] -> [Int] -> [Id]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith (Text -> Type -> Int -> Id
    internalId Text
    "r") [Type]
    resTypes [Int
    nTermArgsInt -> Int -> Int
    forall a. Num a => a -> a -> a
    +Int
    1..Int
    nTermArgsInt -> 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]
    resTypes]
      resId :: Id
    resId = Type -> TmName -> Id
    mkLocalId Type
    pResTy (Text -> Int -> TmName
    forall a. Text -> Int -> Name a
    mkUnsafeInternalName Text
    "r" (Int
    nTermArgsInt -> 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]
    resTypesInt -> Int -> Int
    forall a. Num a => a -> a -> a
    +Int
    1))

      ([Either TyVar Type]
    pArgs, Type
    pResTy) = Type -> ([Either TyVar Type], Type)
    splitFunForallTy Type
    primType
      MultiPrimInfo{mpi_resultDc :: MultiPrimInfo -> DataCon
    mpi_resultDc=DataCon
    tupTc, mpi_resultTypes :: MultiPrimInfo -> [Type]
    mpi_resultTypes=[Type]
    resTypes} =
        HasCallStack => TyConMap -> PrimInfo -> MultiPrimInfo
    TyConMap -> PrimInfo -> MultiPrimInfo
    multiPrimInfo' TyConMap
    tcm PrimInfo
    primInfo

      multiPrimSelect :: Id -> Type -> (Id, Term)
    multiPrimSelect Id
    r Type
    t = (Id
    r, Term -> [Term] -> Term
    mkTmApps (PrimInfo -> Term
    Prim (Type -> PrimInfo
    multiPrimSelectInfo Type
    t)) [Id -> Term
    Var Id
    r, Id -> Term
    Var Id
    resId])
      multiPrimSelectBinds :: [(Id, Term)]
    multiPrimSelectBinds = (Id -> Type -> (Id, Term)) -> [Id] -> [Type] -> [(Id, Term)]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith Id -> Type -> (Id, Term)
    multiPrimSelect  [Id]
    resIds [Type]
    resTypes
      multiPrimTermArgs :: [Either Term b]
    multiPrimTermArgs = (Id -> Either Term b) -> [Id] -> [Either Term b]
    forall a b. (a -> b) -> [a] -> [b]
    map (Term -> Either Term b
    forall a b. a -> Either a b
    Left (Term -> Either Term b) -> (Id -> Term) -> Id -> Either Term b
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Id -> Term
    Var) ([Id]
    argIds [Id] -> [Id] -> [Id]
    forall a. Semigroup a => a -> a -> a
    <> [Id]
    resIds)
      multiPrimTypeArgs :: [Either a Type]
    multiPrimTypeArgs = (TyVar -> Either a Type) -> [TyVar] -> [Either a Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Type -> Either a Type
    forall a b. b -> Either a b
    Right (Type -> Either a Type)
    -> (TyVar -> Type) -> TyVar -> Either a Type
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TyVar -> Type
    VarTy) [TyVar]
    typeVars
      multiPrimBind :: Term
    multiPrimBind =
        Term -> [Either Term Type] -> Term
    mkApps
          (PrimInfo -> Term
    Prim PrimInfo
    primInfo{primMultiResult :: IsMultiPrim
    primMultiResult=IsMultiPrim
    MultiResult})
          ([Either Term Type]
    forall {a}. [Either a Type]
    multiPrimTypeArgs [Either Term Type] -> [Either Term Type] -> [Either Term Type]
    forall a. Semigroup a => a -> a -> a
    <> [Either Term Type]
    forall {b}. [Either Term b]
    multiPrimTermArgs)

      multiPrimSelectInfo :: Type -> PrimInfo
    multiPrimSelectInfo Type
    t = PrimInfo
        { primName :: Text
    primName = Text
    "c$multiPrimSelect"
        , primType :: Type
    primType = Type -> [Either TyVar Type] -> Type
    mkPolyFunTy Type
    pResTy [Type -> Either TyVar Type
    forall a b. b -> Either a b
    Right Type
    pResTy, Type -> Either TyVar Type
    forall a b. b -> Either a b
    Right Type
    t]
        , primWorkInfo :: WorkInfo
    primWorkInfo = WorkInfo
    WorkAlways
        , primMultiResult :: IsMultiPrim
    primMultiResult = IsMultiPrim
    SingleResult
        , primUnfolding :: PrimUnfolding
    primUnfolding = PrimUnfolding
    NoUnfolding
        }

      letTerm :: Term
    letTerm =
        [(Id, Term)] -> Term -> Term
    listToLets
          ((Id
    resId,Term
    multiPrimBind)(Id, Term) -> [(Id, Term)] -> [(Id, Term)]
    forall a. a -> [a] -> [a]
    :[(Id, Term)]
    multiPrimSelectBinds)
          (Term -> [Term] -> Term
    mkTmApps (Term -> [Type] -> Term
    mkTyApps (DataCon -> Term
    Data DataCon
    tupTc) [Type]
    resTypes) ((Id -> Term) -> [Id] -> [Term]
    forall a b. (a -> b) -> [a] -> [b]
    map Id -> Term
    Var [Id]
    resIds))