content: {-|
    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 (TransformContext is0 ctx) e@(collectArgs -> (Var f, _)) = do
    topEnts <- Lens.view topEntities
    tcm <- Lens.view tcCache
    let isTopEnt = f `elemVarSet` topEnts
        isAppFunCtx =
          \case
            AppFun:_ -> True
            TickC _:c -> isAppFunCtx c
            _ -> False
        argTyM = fmap fst (splitFunTy tcm (inferCoreTypeOf tcm e))
    case argTyM of
      Just argTy | isTopEnt && not (isAppFunCtx ctx) -> do
        newId <- mkInternalVar is0 "arg" argTy
        changed (Lam newId (App e (Var newId)))
      _ -> return e
  
  etaExpandSyn _ e = return e
  {-# SCC etaExpandSyn #-}
  
  stripLambda :: Term -> ([Id], Term)
  stripLambda (Lam bndr e) =
    let (bndrs, e') = stripLambda e
     in (bndr : bndrs, e')
  
  stripLambda e = ([], e)
  
  -- | Eta-expand top-level lambda's (DON'T use in a traversal!)
  etaExpansionTL :: HasCallStack => NormRewrite
  etaExpansionTL (TransformContext is0 ctx) (Lam bndr e) = do
    let ctx' = TransformContext (extendInScopeSet is0 bndr) (LamBody bndr : ctx)
    e' <- etaExpansionTL ctx' e
    return $ Lam bndr e'
  
  etaExpansionTL (TransformContext is0 ctx) (Let (NonRec i x) e) = do
    let ctx' = TransformContext (extendInScopeSet is0 i) (LetBody [(i,x)] : ctx)
    e' <- etaExpansionTL ctx' e
    case stripLambda e' of
      (bs@(_:_),e2) -> do
        let e3 = Let (NonRec i x) e2
        changed (mkLams e3 bs)
      _ -> return (Let (NonRec i x) e')
  
  etaExpansionTL (TransformContext is0 ctx) (Let (Rec xes) e) = do
    let bndrs = map fst xes
        ctx' = TransformContext (extendInScopeSetList is0 bndrs) (LetBody xes : ctx)
    e' <- etaExpansionTL ctx' e
    case stripLambda e' of
      (bs@(_:_),e2) -> do
        let e3 = Let (Rec xes) e2
        changed (mkLams e3 bs)
      _ -> return (Let (Rec xes) e')
  
  etaExpansionTL (TransformContext is0 ctx) e
    = do
      tcm <- Lens.view tcCache
      if isFun tcm e
        then do
          let argTy = ( fst
                      . Maybe.fromMaybe (error $ $(curLoc) ++ "etaExpansion splitFunTy")
                      . splitFunTy tcm
                      . inferCoreTypeOf tcm
                      ) e
          newId <- mkInternalVar is0 "arg" argTy
          let ctx' = TransformContext (extendInScopeSet is0 newId) (LamBody newId : ctx)
          e' <- etaExpansionTL ctx' (App e (Var newId))
          changed (Lam newId e')
        else return e
  {-# SCC etaExpansionTL #-}
  