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

      The eta-expansion transformation.
    -}

    {-# LANGUAGE LambdaCase #-}
    {-# LANGUAGE OverloadedStrings #-}
    {-# LANGUAGE TemplateHaskell #-}

    module Clash.Normalize.Transformations.EtaExpand
      ( etaExpandSyn
      , etaExpansionTL
      ) where

    import qualified Control.Lens as Lens
    import qualified Data.Maybe as Maybe
    import GHC.Stack (HasCallStack)

    import Clash.Core.HasType
    import Clash.Core.Term (Bind(..), CoreContext(..), Term(..), collectArgs, mkLams)
    import Clash.Core.TermInfo (isFun)
    import Clash.Core.Type (splitFunTy)
    import Clash.Core.Util (mkInternalVar)
    import Clash.Core.Var (Id)
    import Clash.Core.VarEnv (elemVarSet, extendInScopeSet, extendInScopeSetList)
    import Clash.Normalize.Types (NormRewrite)
    import Clash.Rewrite.Types (TransformContext(..), tcCache, topEntities)
    import Clash.Rewrite.Util (changed)
    import Clash.Util (curLoc)

    -- | Eta-expand functions with a Synthesize annotation, needed to allow such
    -- functions to appear as arguments to higher-order primitives.
    etaExpandSyn :: HasCallStack => NormRewrite
    etaExpandSyn :: HasCallStack => NormRewrite
    etaExpandSyn (TransformContext InScopeSet
    is0 Context
    ctx) e :: Term
    e@(Term -> (Term, [Either Term Type])
    collectArgs -> (Var Id
    f, [Either Term Type]
    _)) = do
      UniqMap (Var Any)
    topEnts <- Getting (UniqMap (Var Any)) RewriteEnv (UniqMap (Var Any))
    -> RewriteMonad NormalizeState (UniqMap (Var Any))
    forall s (m :: Type -> Type) a.
    MonadReader s m =>
    Getting a s a -> m a
    Lens.view Getting (UniqMap (Var Any)) RewriteEnv (UniqMap (Var Any))
    Lens' RewriteEnv (UniqMap (Var Any))
    topEntities
      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
      let isTopEnt :: Bool
    isTopEnt = Id
    f Id -> UniqMap (Var Any) -> Bool
    forall a. Var a -> UniqMap (Var Any) -> Bool
    `elemVarSet` UniqMap (Var Any)
    topEnts
          isAppFunCtx :: Context -> Bool
    isAppFunCtx =
            \case
              CoreContext
    AppFun:Context
    _ -> Bool
    True
              TickC TickInfo
    _:Context
    c -> Context -> Bool
    isAppFunCtx Context
    c
              Context
    _ -> Bool
    False
          argTyM :: Maybe Type
    argTyM = ((Type, Type) -> Type) -> Maybe (Type, Type) -> Maybe Type
    forall a b. (a -> b) -> Maybe a -> Maybe b
    forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
    fmap (Type, Type) -> Type
    forall a b. (a, b) -> a
    fst (TyConMap -> Type -> Maybe (Type, Type)
    splitFunTy TyConMap
    tcm (TyConMap -> Term -> Type
    forall a. InferType a => TyConMap -> a -> Type
    inferCoreTypeOf TyConMap
    tcm Term
    e))
      case Maybe Type
    argTyM of
        Just Type
    argTy | Bool
    isTopEnt Bool -> Bool -> Bool
    && Bool -> Bool
    not (Context -> Bool
    isAppFunCtx Context
    ctx) -> do
          Id
    newId <- InScopeSet -> OccName -> Type -> RewriteMonad NormalizeState Id
    forall (m :: Type -> Type).
    MonadUnique m =>
    InScopeSet -> OccName -> Type -> m Id
    mkInternalVar InScopeSet
    is0 OccName
    "arg" Type
    argTy
          Term -> RewriteMonad NormalizeState Term
    forall a extra. a -> RewriteMonad extra a
    changed (Id -> Term -> Term
    Lam Id
    newId (Term -> Term -> Term
    App Term
    e (Id -> Term
    Var Id
    newId)))
        Maybe Type
    _ -> Term -> RewriteMonad NormalizeState Term
    forall a. a -> RewriteMonad NormalizeState a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return Term
    e

    etaExpandSyn 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
    {-# SCC etaExpandSyn #-}

    stripLambda :: Term -> ([Id], Term)
    stripLambda :: Term -> ([Id], Term)
    stripLambda (Lam Id
    bndr Term
    e) =
      let ([Id]
    bndrs, Term
    e') = Term -> ([Id], Term)
    stripLambda Term
    e
       in (Id
    bndr Id -> [Id] -> [Id]
    forall a. a -> [a] -> [a]
    : [Id]
    bndrs, Term
    e')

    stripLambda Term
    e = ([], Term
    e)

    -- | Eta-expand top-level lambda's (DON'T use in a traversal!)
    etaExpansionTL :: HasCallStack => NormRewrite
    etaExpansionTL :: HasCallStack => NormRewrite
    etaExpansionTL (TransformContext InScopeSet
    is0 Context
    ctx) (Lam Id
    bndr Term
    e) = do
      let ctx' :: TransformContext
    ctx' = InScopeSet -> Context -> TransformContext
    TransformContext (InScopeSet -> Id -> InScopeSet
    forall a. InScopeSet -> Var a -> InScopeSet
    extendInScopeSet InScopeSet
    is0 Id
    bndr) (Id -> CoreContext
    LamBody Id
    bndr CoreContext -> Context -> Context
    forall a. a -> [a] -> [a]
    : Context
    ctx)
      Term
    e' <- HasCallStack => NormRewrite
    NormRewrite
    etaExpansionTL TransformContext
    ctx' Term
    e
      Term -> RewriteMonad NormalizeState Term
    forall a. a -> RewriteMonad NormalizeState a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return (Term -> RewriteMonad NormalizeState Term)
    -> Term -> RewriteMonad NormalizeState Term
    forall a b. (a -> b) -> a -> b
    $ Id -> Term -> Term
    Lam Id
    bndr Term
    e'

    etaExpansionTL (TransformContext InScopeSet
    is0 Context
    ctx) (Let (NonRec Id
    i Term
    x) Term
    e) = do
      let ctx' :: TransformContext
    ctx' = InScopeSet -> Context -> TransformContext
    TransformContext (InScopeSet -> Id -> InScopeSet
    forall a. InScopeSet -> Var a -> InScopeSet
    extendInScopeSet InScopeSet
    is0 Id
    i) ([LetBinding] -> CoreContext
    LetBody [(Id
    i,Term
    x)] CoreContext -> Context -> Context
    forall a. a -> [a] -> [a]
    : Context
    ctx)
      Term
    e' <- HasCallStack => NormRewrite
    NormRewrite
    etaExpansionTL TransformContext
    ctx' Term
    e
      case Term -> ([Id], Term)
    stripLambda Term
    e' of
        (bs :: [Id]
    bs@(Id
    _:[Id]
    _),Term
    e2) -> do
          let e3 :: Term
    e3 = Bind Term -> Term -> Term
    Let (Id -> Term -> Bind Term
    forall a. Id -> a -> Bind a
    NonRec Id
    i Term
    x) Term
    e2
          Term -> RewriteMonad NormalizeState Term
    forall a extra. a -> RewriteMonad extra a
    changed (Term -> [Id] -> Term
    mkLams Term
    e3 [Id]
    bs)
        ([Id], Term)
    _ -> Term -> RewriteMonad NormalizeState Term
    forall a. a -> RewriteMonad NormalizeState a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return (Bind Term -> Term -> Term
    Let (Id -> Term -> Bind Term
    forall a. Id -> a -> Bind a
    NonRec Id
    i Term
    x) Term
    e')

    etaExpansionTL (TransformContext InScopeSet
    is0 Context
    ctx) (Let (Rec [LetBinding]
    xes) Term
    e) = do
      let bndrs :: [Id]
    bndrs = (LetBinding -> Id) -> [LetBinding] -> [Id]
    forall a b. (a -> b) -> [a] -> [b]
    map LetBinding -> Id
    forall a b. (a, b) -> a
    fst [LetBinding]
    xes
          ctx' :: TransformContext
    ctx' = InScopeSet -> Context -> TransformContext
    TransformContext (InScopeSet -> [Id] -> InScopeSet
    forall a. InScopeSet -> [Var a] -> InScopeSet
    extendInScopeSetList InScopeSet
    is0 [Id]
    bndrs) ([LetBinding] -> CoreContext
    LetBody [LetBinding]
    xes CoreContext -> Context -> Context
    forall a. a -> [a] -> [a]
    : Context
    ctx)
      Term
    e' <- HasCallStack => NormRewrite
    NormRewrite
    etaExpansionTL TransformContext
    ctx' Term
    e
      case Term -> ([Id], Term)
    stripLambda Term
    e' of
        (bs :: [Id]
    bs@(Id
    _:[Id]
    _),Term
    e2) -> do
          let e3 :: Term
    e3 = Bind Term -> Term -> Term
    Let ([LetBinding] -> Bind Term
    forall a. [(Id, a)] -> Bind a
    Rec [LetBinding]
    xes) Term
    e2
          Term -> RewriteMonad NormalizeState Term
    forall a extra. a -> RewriteMonad extra a
    changed (Term -> [Id] -> Term
    mkLams Term
    e3 [Id]
    bs)
        ([Id], Term)
    _ -> Term -> RewriteMonad NormalizeState Term
    forall a. a -> RewriteMonad NormalizeState a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return (Bind Term -> Term -> Term
    Let ([LetBinding] -> Bind Term
    forall a. [(Id, a)] -> Bind a
    Rec [LetBinding]
    xes) Term
    e')

    etaExpansionTL (TransformContext InScopeSet
    is0 Context
    ctx) Term
    e
      = 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
        if TyConMap -> Term -> Bool
    isFun TyConMap
    tcm Term
    e
          then do
            let argTy :: Type
    argTy = ( (Type, Type) -> Type
    forall a b. (a, b) -> a
    fst
                        ((Type, Type) -> Type) -> (Term -> (Type, Type)) -> Term -> Type
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Type, Type) -> Maybe (Type, Type) -> (Type, Type)
    forall a. a -> Maybe a -> a
    Maybe.fromMaybe ([Char] -> (Type, Type)
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> (Type, Type)) -> [Char] -> (Type, Type)
    forall a b. (a -> b) -> a -> b
    $ $(curLoc) [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "etaExpansion splitFunTy")
                        (Maybe (Type, Type) -> (Type, Type))
    -> (Term -> Maybe (Type, Type)) -> Term -> (Type, Type)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TyConMap -> Type -> Maybe (Type, Type)
    splitFunTy TyConMap
    tcm
                        (Type -> Maybe (Type, Type))
    -> (Term -> Type) -> Term -> Maybe (Type, Type)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TyConMap -> Term -> Type
    forall a. InferType a => TyConMap -> a -> Type
    inferCoreTypeOf TyConMap
    tcm
                        ) Term
    e
            Id
    newId <- InScopeSet -> OccName -> Type -> RewriteMonad NormalizeState Id
    forall (m :: Type -> Type).
    MonadUnique m =>
    InScopeSet -> OccName -> Type -> m Id
    mkInternalVar InScopeSet
    is0 OccName
    "arg" Type
    argTy
            let ctx' :: TransformContext
    ctx' = InScopeSet -> Context -> TransformContext
    TransformContext (InScopeSet -> Id -> InScopeSet
    forall a. InScopeSet -> Var a -> InScopeSet
    extendInScopeSet InScopeSet
    is0 Id
    newId) (Id -> CoreContext
    LamBody Id
    newId CoreContext -> Context -> Context
    forall a. a -> [a] -> [a]
    : Context
    ctx)
            Term
    e' <- HasCallStack => NormRewrite
    NormRewrite
    etaExpansionTL TransformContext
    ctx' (Term -> Term -> Term
    App Term
    e (Id -> Term
    Var Id
    newId))
            Term -> RewriteMonad NormalizeState Term
    forall a extra. a -> RewriteMonad extra a
    changed (Id -> Term -> Term
    Lam Id
    newId Term
    e')
          else Term -> RewriteMonad NormalizeState Term
    forall a. a -> RewriteMonad NormalizeState a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return Term
    e
    {-# SCC etaExpansionTL #-}