{-|
  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>

  Transformations for converting to A-Normal Form.
-}

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

module Clash.Normalize.Transformations.ANF
  ( makeANF
  , nonRepANF
  ) where

import Control.Arrow ((***))
import Control.Lens (_2)
import qualified Control.Lens as Lens
import qualified Control.Monad as Monad
import Control.Monad.State (StateT, lift, modify, runStateT)
import Control.Monad.Writer (listen)
import Data.Bifunctor (second)
import qualified Data.Monoid as Monoid (Any(..))
import qualified Data.Text.Extra as Text (showt)
import GHC.Stack (HasCallStack)

import Clash.Signal.Internal (Signal(..))

import Clash.Core.DataCon (DataCon(..))
import Clash.Core.HasFreeVars (disjointFreeVars)
import Clash.Core.HasType
import Clash.Core.Name (mkUnsafeSystemName, nameOcc)
import Clash.Core.Subst (deshadowLetExpr, freshenTm)
import Clash.Core.Term
  ( Alt, CoreContext(..), LetBinding, Pat(..), PrimInfo(..), Term(..)
  , collectArgs, collectTicks, mkTicks, partitionTicks, stripTicks)
import Clash.Core.TermInfo (isCon, isLocalVar, isPrim, isVar)
import Clash.Core.TyCon (TyConMap)
import Clash.Core.Type (Type, TypeView(..), coreView, tyView)
import Clash.Core.Util (mkSelectorCase)
import Clash.Core.Var (Id)
import Clash.Core.VarEnv (InScopeSet, extendInScopeSet, extendInScopeSetList, mkVarSet)
import Clash.Netlist.Util (bindsExistentials)
import Clash.Normalize.Transformations.Specialize (specialize)
import Clash.Normalize.Types (NormRewrite, NormalizeSession)
import Clash.Rewrite.Combinators (bottomupR)
import Clash.Rewrite.Types
  (Transform, TransformContext(..), tcCache)
import Clash.Rewrite.Util
  (changed, isUntranslatable, mkDerivedName, mkTmBinderFor)
import Clash.Rewrite.WorkFree (isConstant, isConstantNotClockReset)
import Clash.Util (curLoc)

{- [Note: ANF in Clash]
ANF suitable for use in Clash can be described with the given types:

  data ATerm
    = ALam !Id ATerm
    | ATyLam !TyVar ATerm
    | ALetrec [(Id, CTerm)] !ITerm

  data CTerm
    = CApp !Id [Either ITerm Type]
    | CCase !ITerm !Type [(Pat, ITerm)]
    | CCast !ITerm !Type !Type
    | CPrim !PrimInfo [Either ITerm Type]
    | CTick !TickInfo CTerm

  data ITerm
    = IVar !Id
    | ILiteral !Literal
    | IData !DataCon [Either ITerm Type]
    | IPrim !PrimInfo [Either ITerm Type]
    | ITick !TickInfo ITerm

where ATerm is a term in A-normal form, CTerm is a compound term (i.e. one
which can only appear let-bound in ANF) and ITerm is an immediate term (i.e.
one which represents some simple term).

There are two constructors for primtiives, CPrim and IPrim. The difference
between these are whether the primitive performs work or not. Primitives which
perform work should be shared, but work-free primitives can be used directly.

These types help codify some invariants that must hold for the result of ANF:

  * terms start with (ty)lambdas, lambdas do not occur in let bindings or the
    the body of a letrec expression

  * there are no nested letrec expressions, only a single letrec which may
    occur after all lambdas

  * an ANF term may not have a letrec expression if the definition is already
    an immediate term, e.g. where there is no benefit in sharing the result

  * only compound terms are let-bound, as there is no benefit from let binding
    an immediate term (there is no benefit to sharing immediate terms)

  * arguments to functions / data constructors / primitives are not let bound
    if they correspond are immediate, but are if they are compound (to produce
    a variable which is an immediate term)

  * the leftmost innermost term in a function application is always an
    identifier, lambdas should have been removed by application propagation

  * the right-hand side of a case alternative is an immediate term

  * the body of the letrec expression is an immediate term

Some invariants are not captured by these types:

  * non-representable terms and terms in IO are not let-bound, instead they are
    pushed down as far as possible

  * if a let binding is created for the result, the name of the Id is "result"

TODO: The best way to enforce that Clash implements ANF compatible with these
types is to implement ANF using these types. However, as currently implemented
ANF is mostly defined using the bottom-up transformation 'collectANF'. This
would be some amount of effort to replace currently, perhaps it would be better
to convert the result of partial evaluation to these data types when it is
implemented more, then use these Anf types directly in the conversion to
netlist, i.e. Term -> Value -> Normal -> Anf -> Netlist.
-}

{- [Note: Name re-creation]
The names of heap bound variables are safely generate with mkUniqSystemId in
Clash.Core.Evaluator.newLetBinding. But only their uniqs end up in the heap,
not the complete names. So we use mkUnsafeSystemName to recreate the same Name.
-}

-- | Turn an expression into a modified ANF-form. As opposed to standard ANF,
-- constants do not become let-bound.
makeANF :: HasCallStack => NormRewrite
makeANF :: NormRewrite
makeANF (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
makeANF TransformContext
ctx' Term
e
  Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Id -> Term -> Term
Lam Id
bndr Term
e')

makeANF TransformContext
_ e :: Term
e@(TyLam {}) = Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e

makeANF ctx :: TransformContext
ctx@(TransformContext InScopeSet
is0 Context
_) Term
e0 = do
    -- We need to freshen all binders in `e` because we're shuffling them around
    -- into a single let-binder, because even when binders don't shadow, they
    -- don't have to be unique within an expression. And so lifting them all
    -- to a single let-binder will cause issues when they're not unique.
    --
    -- We cannot make freshening part of collectANF, because when we generate
    -- new binders, we need to make sure those names do not conflict with _any_
    -- of the existing binders in the expression.
    --
    -- See also Note [ANF InScopeSet]
    let (InScopeSet
is2,Term
e1) = InScopeSet -> Term -> (InScopeSet, Term)
freshenTm InScopeSet
is0 Term
e0
    ((Term
e2,([LetBinding]
bndrs,InScopeSet
_)),Any -> Bool
Monoid.getAny -> Bool
hasChanged) <-
      NormalizeSession (Term, ([LetBinding], InScopeSet))
-> NormalizeSession ((Term, ([LetBinding], InScopeSet)), Any)
forall w (m :: Type -> Type) a. MonadWriter w m => m a -> m (a, w)
listen (StateT ([LetBinding], InScopeSet) NormalizeSession Term
-> ([LetBinding], InScopeSet)
-> NormalizeSession (Term, ([LetBinding], InScopeSet))
forall s (m :: Type -> Type) a. StateT s m a -> s -> m (a, s)
runStateT (Transform (StateT ([LetBinding], InScopeSet) NormalizeSession)
-> Transform (StateT ([LetBinding], InScopeSet) NormalizeSession)
forall (m :: Type -> Type). Monad m => Transform m -> Transform m
bottomupR HasCallStack =>
Transform (StateT ([LetBinding], InScopeSet) NormalizeSession)
Transform (StateT ([LetBinding], InScopeSet) NormalizeSession)
collectANF TransformContext
ctx Term
e1) ([],InScopeSet
is2))
    case [LetBinding]
bndrs of
      [] -> if Bool
hasChanged then Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e2 else Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e0
      [LetBinding]
_  -> do
        let (Term
e3,[TickInfo]
ticks) = Term -> (Term, [TickInfo])
collectTicks Term
e2
            ([TickInfo]
srcTicks,[TickInfo]
nmTicks) = [TickInfo] -> ([TickInfo], [TickInfo])
partitionTicks [TickInfo]
ticks
        -- Ensure that `AppendName` ticks still scope over the entire expression
        Term -> RewriteMonad NormalizeState Term
forall a extra. a -> RewriteMonad extra a
changed (Term -> [TickInfo] -> Term
mkTicks ([LetBinding] -> Term -> Term
Letrec [LetBinding]
bndrs (Term -> [TickInfo] -> Term
mkTicks Term
e3 [TickInfo]
srcTicks)) [TickInfo]
nmTicks)
{-# SCC makeANF #-}

type NormRewriteW = Transform (StateT ([LetBinding],InScopeSet) NormalizeSession)

-- | See Note [ANF InScopeSet]
tellBinders :: [LetBinding] -> StateT ([LetBinding],InScopeSet) NormalizeSession ()
tellBinders :: [LetBinding]
-> StateT ([LetBinding], InScopeSet) NormalizeSession ()
tellBinders [LetBinding]
bs = (([LetBinding], InScopeSet) -> ([LetBinding], InScopeSet))
-> StateT ([LetBinding], InScopeSet) NormalizeSession ()
forall s (m :: Type -> Type). MonadState s m => (s -> s) -> m ()
modify (([LetBinding]
bs [LetBinding] -> [LetBinding] -> [LetBinding]
forall a. [a] -> [a] -> [a]
++) ([LetBinding] -> [LetBinding])
-> (InScopeSet -> InScopeSet)
-> ([LetBinding], InScopeSet)
-> ([LetBinding], InScopeSet)
forall (a :: Type -> Type -> Type) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** (InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
`extendInScopeSetList` ((LetBinding -> Id) -> [LetBinding] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map LetBinding -> Id
forall a b. (a, b) -> a
fst [LetBinding]
bs)))

-- | See Note [ANF InScopeSet]; only extends the inscopeset
notifyBinders :: Monad m => [LetBinding] -> StateT ([LetBinding],InScopeSet) m ()
notifyBinders :: [LetBinding] -> StateT ([LetBinding], InScopeSet) m ()
notifyBinders [LetBinding]
bs = (([LetBinding], InScopeSet) -> ([LetBinding], InScopeSet))
-> StateT ([LetBinding], InScopeSet) m ()
forall s (m :: Type -> Type). MonadState s m => (s -> s) -> m ()
modify ((InScopeSet -> InScopeSet)
-> ([LetBinding], InScopeSet) -> ([LetBinding], InScopeSet)
forall (p :: Type -> Type -> Type) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
`extendInScopeSetList` ((LetBinding -> Id) -> [LetBinding] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map LetBinding -> Id
forall a b. (a, b) -> a
fst [LetBinding]
bs)))

-- | Is the given type IO-like
isSimIOTy
  :: TyConMap
  -> Type
  -- ^ Type to check for IO-likeness
  -> Bool
isSimIOTy :: TyConMap -> Type -> Bool
isSimIOTy TyConMap
tcm Type
ty = case Type -> TypeView
tyView (TyConMap -> Type -> Type
coreView TyConMap
tcm Type
ty) of
  TyConApp TyConName
tcNm [Type]
args
    | TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
tcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Explicit.SimIO.SimIO"
    -> Bool
True
    | TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
tcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"GHC.Prim.(#,#)"
    , [Type
_,Type
_,Type
st,Type
_] <- [Type]
args
    -> TyConMap -> Type -> Bool
isStateTokenTy TyConMap
tcm Type
st
  FunTy Type
_ Type
res -> TyConMap -> Type -> Bool
isSimIOTy TyConMap
tcm Type
res
  TypeView
_ -> Bool
False

-- | Is the given type the state token
isStateTokenTy
  :: TyConMap
  -> Type
  -- ^ Type to check for state tokenness
  -> Bool
isStateTokenTy :: TyConMap -> Type -> Bool
isStateTokenTy TyConMap
tcm Type
ty = case Type -> TypeView
tyView (TyConMap -> Type -> Type
coreView TyConMap
tcm Type
ty) of
  TyConApp TyConName
tcNm [Type]
_ -> TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
tcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"GHC.Prim.State#"
  TypeView
_ -> Bool
False

-- | Note [ANF InScopeSet]
--
-- The InScopeSet contains:
--
--    1. All the free variables of the expression we are traversing
--
--    2. All the bound variables of the expression we are traversing
--
--    3. The newly created let-bindings as we recurse back up the traversal
--
-- All of these are needed to created let-bindings that
--
--    * Do not shadow
--    * Are not shadowed
--    * Nor conflict with each other (i.e. have the same unique)
--
-- Initially we start with the local InScopeSet and add the global variables:
--
-- @
-- is1 <- unionInScope is0 <$> Lens.use globalInScope
-- @
--
-- Which will gives us the (superset of) free variables of the expression. Then
-- we call  'freshenTm'
--
-- @
-- let (is2,e1) = freshenTm is1 e0
-- @
--
-- Which extends the InScopeSet with all the bound variables in 'e1', the
-- version of 'e0' where all binders are unique (not just deshadowed).
--
-- So we start out with an InScopeSet that satisfies points 1 and 2, now every
-- time we create a new binder we must add it to the InScopeSet to satisfy
-- point 3.
--
-- Note [ANF no let-bind]
--
-- | Do not let-bind:
--
-- 1. Arguments with an untranslatable type: untranslatable expressions
--    should be propagated down as far as possible
--
-- 2. Local variables or constants: they don't add any work, so no reason
--    to let-bind to enable sharing
--
-- 3. IO actions, the translation of IO actions to sequential HDL constructs
--    depends on IO actions to be propagated down as far as possible.
collectANF :: HasCallStack => NormRewriteW
collectANF :: Transform (StateT ([LetBinding], InScopeSet) NormalizeSession)
collectANF TransformContext
ctx e :: Term
e@(App Term
appf Term
arg)
  | (Term
conVarPrim, [Either Term Type]
_) <- Term -> (Term, [Either Term Type])
collectArgs Term
e
  , Term -> Bool
isCon Term
conVarPrim Bool -> Bool -> Bool
|| Term -> Bool
isPrim Term
conVarPrim Bool -> Bool -> Bool
|| Term -> Bool
isVar Term
conVarPrim
  = do
    TyConMap
tcm <- Getting TyConMap RewriteEnv TyConMap
-> StateT ([LetBinding], InScopeSet) NormalizeSession 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
    Bool
untranslatable <- RewriteMonad NormalizeState Bool
-> StateT ([LetBinding], InScopeSet) NormalizeSession Bool
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Bool -> Term -> RewriteMonad NormalizeState Bool
forall extra. Bool -> Term -> RewriteMonad extra Bool
isUntranslatable Bool
False Term
arg)
    let localVar :: Bool
localVar   = Term -> Bool
isLocalVar Term
arg
        constantNoCR :: Bool
constantNoCR = TyConMap -> Term -> Bool
isConstantNotClockReset TyConMap
tcm Term
arg
    -- See Note [ANF no let-bind]
    case (Bool
untranslatable,Bool
localVar Bool -> Bool -> Bool
|| Bool
constantNoCR, Term -> Bool
isSimBind Term
conVarPrim,Term
arg) of
      (Bool
False,Bool
False,Bool
False,Term
_) -> do
        -- See Note [ANF InScopeSet]
        InScopeSet
is1   <- Getting InScopeSet ([LetBinding], InScopeSet) InScopeSet
-> StateT ([LetBinding], InScopeSet) NormalizeSession InScopeSet
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting InScopeSet ([LetBinding], InScopeSet) InScopeSet
forall s t a b. Field2 s t a b => Lens s t a b
_2
        Id
argId <- NormalizeSession Id
-> StateT ([LetBinding], InScopeSet) NormalizeSession Id
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (InScopeSet -> TyConMap -> Name Term -> Term -> NormalizeSession Id
forall (m :: Type -> Type) a.
MonadUnique m =>
InScopeSet -> TyConMap -> Name a -> Term -> m Id
mkTmBinderFor InScopeSet
is1 TyConMap
tcm (TransformContext -> OccName -> Name Term
mkDerivedName TransformContext
ctx OccName
"app_arg") Term
arg)
        -- See Note [ANF InScopeSet]
        [LetBinding]
-> StateT ([LetBinding], InScopeSet) NormalizeSession ()
tellBinders [(Id
argId,Term
arg)]
        Term -> StateT ([LetBinding], InScopeSet) NormalizeSession Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Term -> Term -> Term
App Term
appf (Id -> Term
Var Id
argId))
      (Bool
True,Bool
False,Bool
_,Letrec [LetBinding]
binds Term
body) -> do
        [LetBinding]
-> StateT ([LetBinding], InScopeSet) NormalizeSession ()
tellBinders [LetBinding]
binds
        Term -> StateT ([LetBinding], InScopeSet) NormalizeSession Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Term -> Term -> Term
App Term
appf Term
body)
      (Bool, Bool, Bool, Term)
_ -> Term -> StateT ([LetBinding], InScopeSet) NormalizeSession Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e
 where
  isSimBind :: Term -> Bool
isSimBind (Prim PrimInfo
p) = PrimInfo -> OccName
primName PrimInfo
p OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Explicit.SimIO.bindSimIO#"
  isSimBind Term
_ = Bool
False

collectANF TransformContext
_ (Letrec [LetBinding]
binds Term
body) = do
  TyConMap
tcm <- Getting TyConMap RewriteEnv TyConMap
-> StateT ([LetBinding], InScopeSet) NormalizeSession 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 isSimIO :: Bool
isSimIO = TyConMap -> Type -> Bool
isSimIOTy TyConMap
tcm (TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
body)
  Bool
untranslatable <- RewriteMonad NormalizeState Bool
-> StateT ([LetBinding], InScopeSet) NormalizeSession Bool
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Bool -> Term -> RewriteMonad NormalizeState Bool
forall extra. Bool -> Term -> RewriteMonad extra Bool
isUntranslatable Bool
False Term
body)
  let localVar :: Bool
localVar = Term -> Bool
isLocalVar Term
body
  -- See Note [ANF no let-bind]
  if Bool
localVar Bool -> Bool -> Bool
|| Bool
untranslatable Bool -> Bool -> Bool
|| Bool
isSimIO
    then do
      [LetBinding]
-> StateT ([LetBinding], InScopeSet) NormalizeSession ()
tellBinders [LetBinding]
binds
      Term -> StateT ([LetBinding], InScopeSet) NormalizeSession Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
body
    else do
      -- See Note [ANF InScopeSet]
      InScopeSet
is1 <- Getting InScopeSet ([LetBinding], InScopeSet) InScopeSet
-> StateT ([LetBinding], InScopeSet) NormalizeSession InScopeSet
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting InScopeSet ([LetBinding], InScopeSet) InScopeSet
forall s t a b. Field2 s t a b => Lens s t a b
_2
      Id
argId <- NormalizeSession Id
-> StateT ([LetBinding], InScopeSet) NormalizeSession Id
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (InScopeSet -> TyConMap -> Name Any -> Term -> NormalizeSession Id
forall (m :: Type -> Type) a.
MonadUnique m =>
InScopeSet -> TyConMap -> Name a -> Term -> m Id
mkTmBinderFor InScopeSet
is1 TyConMap
tcm (OccName -> Unique -> Name Any
forall a. OccName -> Unique -> Name a
mkUnsafeSystemName OccName
"result" Unique
0) Term
body)
      -- See Note [ANF InScopeSet]
      [LetBinding]
-> StateT ([LetBinding], InScopeSet) NormalizeSession ()
tellBinders [(Id
argId,Term
body)]
      [LetBinding]
-> StateT ([LetBinding], InScopeSet) NormalizeSession ()
tellBinders [LetBinding]
binds
      Term -> StateT ([LetBinding], InScopeSet) NormalizeSession Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Id -> Term
Var Id
argId)

-- TODO: The code below special-cases ANF for the ':-' constructor for the
-- 'Signal' type. The 'Signal' type is essentially treated as a "transparent"
-- type by the Clash compiler, so observing its constructor leads to all kinds
-- of problems. In this case that "Clash.Rewrite.Util.mkSelectorCase" will
-- try to project the LHS and RHS of the ':-' constructor, however,
-- 'mkSelectorCase' uses 'coreView1' to find the "real" data-constructor.
-- 'coreView1' however looks through the 'Signal' type, and hence 'mkSelector'
-- finds the data constructors for the element type of Signal. This resulted in
-- error #24 (https://github.com/christiaanb/clash2/issues/24), where we
-- try to get the first field out of the 'Vec's 'Nil' constructor.
--
-- Ultimately we should stop treating Signal as a "transparent" type and deal
-- handling of the Signal type, and the involved co-recursive functions,
-- properly. At the moment, Clash cannot deal with this recursive type and the
-- recursive functions involved, hence the need for special-casing code. After
-- everything is done properly, we should remove the two lines below.
collectANF TransformContext
_ e :: Term
e@(Case Term
_ Type
_ [(DataPat DataCon
dc [TyVar]
_ [Id]
_,Term
_)])
  | Name DataCon -> OccName
forall a. Name a -> OccName
nameOcc (DataCon -> Name DataCon
dcName DataCon
dc) OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== Name -> OccName
forall a. Show a => a -> OccName
Text.showt '(:-) = Term -> StateT ([LetBinding], InScopeSet) NormalizeSession Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e

collectANF TransformContext
ctx (Case Term
subj Type
ty [Alt]
alts) = do
    let localVar :: Bool
localVar = Term -> Bool
isLocalVar Term
subj
    let isConstantSubj :: Bool
isConstantSubj = Term -> Bool
isConstant Term
subj

    (Term
subj',[LetBinding]
subjBinders) <- if Bool
localVar Bool -> Bool -> Bool
|| Bool
isConstantSubj
      then (Term, [LetBinding])
-> StateT
     ([LetBinding], InScopeSet) NormalizeSession (Term, [LetBinding])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Term
subj,[])
      else do
        TyConMap
tcm <- Getting TyConMap RewriteEnv TyConMap
-> StateT ([LetBinding], InScopeSet) NormalizeSession 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
        -- See Note [ANF InScopeSet]
        InScopeSet
is1 <- Getting InScopeSet ([LetBinding], InScopeSet) InScopeSet
-> StateT ([LetBinding], InScopeSet) NormalizeSession InScopeSet
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting InScopeSet ([LetBinding], InScopeSet) InScopeSet
forall s t a b. Field2 s t a b => Lens s t a b
_2
        Id
argId <- NormalizeSession Id
-> StateT ([LetBinding], InScopeSet) NormalizeSession Id
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (InScopeSet -> TyConMap -> Name Term -> Term -> NormalizeSession Id
forall (m :: Type -> Type) a.
MonadUnique m =>
InScopeSet -> TyConMap -> Name a -> Term -> m Id
mkTmBinderFor InScopeSet
is1 TyConMap
tcm (TransformContext -> OccName -> Name Term
mkDerivedName TransformContext
ctx OccName
"case_scrut") Term
subj)
        -- See Note [ANF InScopeSet]
        [LetBinding]
-> StateT ([LetBinding], InScopeSet) NormalizeSession ()
forall (m :: Type -> Type).
Monad m =>
[LetBinding] -> StateT ([LetBinding], InScopeSet) m ()
notifyBinders [(Id
argId,Term
subj)]
        (Term, [LetBinding])
-> StateT
     ([LetBinding], InScopeSet) NormalizeSession (Term, [LetBinding])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Id -> Term
Var Id
argId,[(Id
argId,Term
subj)])

    TyConMap
tcm <- Getting TyConMap RewriteEnv TyConMap
-> StateT ([LetBinding], InScopeSet) NormalizeSession 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 isSimIOAlt :: Bool
isSimIOAlt = TyConMap -> Type -> Bool
isSimIOTy TyConMap
tcm Type
ty

    [Alt]
alts' <- (Alt -> StateT ([LetBinding], InScopeSet) NormalizeSession Alt)
-> [Alt]
-> StateT ([LetBinding], InScopeSet) NormalizeSession [Alt]
forall (t :: Type -> Type) (m :: Type -> Type) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Bool
-> Term
-> Alt
-> StateT ([LetBinding], InScopeSet) NormalizeSession Alt
doAlt Bool
isSimIOAlt Term
subj') [Alt]
alts
    [LetBinding]
-> StateT ([LetBinding], InScopeSet) NormalizeSession ()
tellBinders [LetBinding]
subjBinders

    case [Alt]
alts' of
      [(DataPat DataCon
_ [] [Id]
xs,Term
altExpr)]
        | [Id] -> VarSet
forall a. [Var a] -> VarSet
mkVarSet [Id]
xs VarSet -> Term -> Bool
forall a. HasFreeVars a => VarSet -> a -> Bool
`disjointFreeVars` Term
altExpr Bool -> Bool -> Bool
|| Bool
isSimIOAlt
        -> Term -> StateT ([LetBinding], InScopeSet) NormalizeSession Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
altExpr
      [Alt]
_ -> Term -> StateT ([LetBinding], InScopeSet) NormalizeSession Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Term -> Type -> [Alt] -> Term
Case Term
subj' Type
ty [Alt]
alts')
  where
    doAlt :: Bool -> Term -> Alt -> StateT ([LetBinding],InScopeSet) NormalizeSession Alt
    doAlt :: Bool
-> Term
-> Alt
-> StateT ([LetBinding], InScopeSet) NormalizeSession Alt
doAlt Bool
isSimIOAlt Term
subj' alt :: Alt
alt@(DataPat DataCon
dc [TyVar]
exts [Id]
xs,Term
altExpr) | Bool -> Bool
not ([TyVar] -> [Id] -> Bool
forall a. [TyVar] -> [Var a] -> Bool
bindsExistentials [TyVar]
exts [Id]
xs) = do
      let lv :: Bool
lv = Term -> Bool
isLocalVar Term
altExpr
      [LetBinding]
patSels <- (Id
 -> Unique
 -> StateT ([LetBinding], InScopeSet) NormalizeSession LetBinding)
-> [Id]
-> [Unique]
-> StateT ([LetBinding], InScopeSet) NormalizeSession [LetBinding]
forall (m :: Type -> Type) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
Monad.zipWithM (Term
-> DataCon
-> Id
-> Unique
-> StateT ([LetBinding], InScopeSet) NormalizeSession LetBinding
doPatBndr Term
subj' DataCon
dc) [Id]
xs [Unique
0..]
      let altExprIsConstant :: Bool
altExprIsConstant = Term -> Bool
isConstant Term
altExpr
      let usesXs :: Term -> Bool
usesXs (Var Id
n) = (Id -> Bool) -> [Id] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any (Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
n) [Id]
xs
          usesXs Term
_       = Bool
False
      -- See [ANF no let-bind]
      if [Bool] -> Bool
forall (t :: Type -> Type). Foldable t => t Bool -> Bool
or [Bool
isSimIOAlt, Bool
lv Bool -> Bool -> Bool
&& (Bool -> Bool
not (Term -> Bool
usesXs Term
altExpr) Bool -> Bool -> Bool
|| [Alt] -> Unique
forall (t :: Type -> Type) a. Foldable t => t a -> Unique
length [Alt]
alts Unique -> Unique -> Bool
forall a. Eq a => a -> a -> Bool
== Unique
1), Bool
altExprIsConstant]
        then do
          -- See Note [ANF InScopeSet]
          [LetBinding]
-> StateT ([LetBinding], InScopeSet) NormalizeSession ()
tellBinders [LetBinding]
patSels
          Alt -> StateT ([LetBinding], InScopeSet) NormalizeSession Alt
forall (m :: Type -> Type) a. Monad m => a -> m a
return Alt
alt
        else do
          TyConMap
tcm <- Getting TyConMap RewriteEnv TyConMap
-> StateT ([LetBinding], InScopeSet) NormalizeSession 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
          -- See Note [ANF InScopeSet]
          InScopeSet
is1 <- Getting InScopeSet ([LetBinding], InScopeSet) InScopeSet
-> StateT ([LetBinding], InScopeSet) NormalizeSession InScopeSet
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting InScopeSet ([LetBinding], InScopeSet) InScopeSet
forall s t a b. Field2 s t a b => Lens s t a b
_2
          Id
altId <- NormalizeSession Id
-> StateT ([LetBinding], InScopeSet) NormalizeSession Id
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (InScopeSet -> TyConMap -> Name Term -> Term -> NormalizeSession Id
forall (m :: Type -> Type) a.
MonadUnique m =>
InScopeSet -> TyConMap -> Name a -> Term -> m Id
mkTmBinderFor InScopeSet
is1 TyConMap
tcm (TransformContext -> OccName -> Name Term
mkDerivedName TransformContext
ctx OccName
"case_alt") Term
altExpr)
          -- See Note [ANF InScopeSet]
          [LetBinding]
-> StateT ([LetBinding], InScopeSet) NormalizeSession ()
tellBinders ([LetBinding]
patSels [LetBinding] -> [LetBinding] -> [LetBinding]
forall a. [a] -> [a] -> [a]
++ [(Id
altId,Term
altExpr)])
          Alt -> StateT ([LetBinding], InScopeSet) NormalizeSession Alt
forall (m :: Type -> Type) a. Monad m => a -> m a
return (DataCon -> [TyVar] -> [Id] -> Pat
DataPat DataCon
dc [TyVar]
exts [Id]
xs,Id -> Term
Var Id
altId)
    doAlt Bool
_ Term
_ alt :: Alt
alt@(DataPat {}, Term
_) = Alt -> StateT ([LetBinding], InScopeSet) NormalizeSession Alt
forall (m :: Type -> Type) a. Monad m => a -> m a
return Alt
alt
    doAlt Bool
isSimIOAlt Term
_ alt :: Alt
alt@(Pat
pat,Term
altExpr) = do
      let lv :: Bool
lv = Term -> Bool
isLocalVar Term
altExpr
      let altExprIsConstant :: Bool
altExprIsConstant = Term -> Bool
isConstant Term
altExpr
      -- See [ANF no let-bind]
      if Bool
isSimIOAlt Bool -> Bool -> Bool
|| Bool
lv Bool -> Bool -> Bool
|| Bool
altExprIsConstant
        then Alt -> StateT ([LetBinding], InScopeSet) NormalizeSession Alt
forall (m :: Type -> Type) a. Monad m => a -> m a
return Alt
alt
        else do
          TyConMap
tcm <- Getting TyConMap RewriteEnv TyConMap
-> StateT ([LetBinding], InScopeSet) NormalizeSession 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
          -- See Note [ANF InScopeSet]
          InScopeSet
is1 <- Getting InScopeSet ([LetBinding], InScopeSet) InScopeSet
-> StateT ([LetBinding], InScopeSet) NormalizeSession InScopeSet
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting InScopeSet ([LetBinding], InScopeSet) InScopeSet
forall s t a b. Field2 s t a b => Lens s t a b
_2
          Id
altId <- NormalizeSession Id
-> StateT ([LetBinding], InScopeSet) NormalizeSession Id
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (InScopeSet -> TyConMap -> Name Term -> Term -> NormalizeSession Id
forall (m :: Type -> Type) a.
MonadUnique m =>
InScopeSet -> TyConMap -> Name a -> Term -> m Id
mkTmBinderFor InScopeSet
is1 TyConMap
tcm (TransformContext -> OccName -> Name Term
mkDerivedName TransformContext
ctx OccName
"case_alt") Term
altExpr)
          [LetBinding]
-> StateT ([LetBinding], InScopeSet) NormalizeSession ()
tellBinders [(Id
altId,Term
altExpr)]
          Alt -> StateT ([LetBinding], InScopeSet) NormalizeSession Alt
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Pat
pat,Id -> Term
Var Id
altId)

    doPatBndr :: Term -> DataCon -> Id -> Int -> StateT ([LetBinding],InScopeSet) NormalizeSession LetBinding
    doPatBndr :: Term
-> DataCon
-> Id
-> Unique
-> StateT ([LetBinding], InScopeSet) NormalizeSession LetBinding
doPatBndr Term
subj' DataCon
dc Id
pId Unique
i = do
      TyConMap
tcm <- Getting TyConMap RewriteEnv TyConMap
-> StateT ([LetBinding], InScopeSet) NormalizeSession 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
      -- See Note [ANF InScopeSet]
      InScopeSet
is1 <- Getting InScopeSet ([LetBinding], InScopeSet) InScopeSet
-> StateT ([LetBinding], InScopeSet) NormalizeSession InScopeSet
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting InScopeSet ([LetBinding], InScopeSet) InScopeSet
forall s t a b. Field2 s t a b => Lens s t a b
_2
      Term
patExpr <- RewriteMonad NormalizeState Term
-> StateT ([LetBinding], InScopeSet) NormalizeSession Term
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (String
-> InScopeSet
-> TyConMap
-> Term
-> Unique
-> Unique
-> RewriteMonad NormalizeState Term
forall (m :: Type -> Type).
(HasCallStack, MonadUnique m) =>
String
-> InScopeSet -> TyConMap -> Term -> Unique -> Unique -> m Term
mkSelectorCase ($(String
curLoc) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"doPatBndr") InScopeSet
is1 TyConMap
tcm Term
subj' (DataCon -> Unique
dcTag DataCon
dc) Unique
i)
      -- No need to 'tellBinders' here because 'pId' is already in the ANF
      -- InScopeSet.
      --
      -- See also Note [ANF InScopeSet]
      LetBinding
-> StateT ([LetBinding], InScopeSet) NormalizeSession LetBinding
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Id
pId,Term
patExpr)

collectANF TransformContext
_ Term
e = Term -> StateT ([LetBinding], InScopeSet) NormalizeSession Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e
{-# SCC collectANF #-}

-- | Bring an application of a DataCon or Primitive in ANF, when the argument is
-- is considered non-representable
nonRepANF :: HasCallStack => NormRewrite
nonRepANF :: NormRewrite
nonRepANF ctx :: TransformContext
ctx@(TransformContext InScopeSet
is0 Context
_) e :: Term
e@(App Term
appConPrim Term
arg)
  | (Term
conPrim, [Either Term Type]
_) <- Term -> (Term, [Either Term Type])
collectArgs Term
e
  , Term -> Bool
isCon Term
conPrim Bool -> Bool -> Bool
|| Term -> Bool
isPrim Term
conPrim
  = do
    Bool
untranslatable <- Bool -> Term -> RewriteMonad NormalizeState Bool
forall extra. Bool -> Term -> RewriteMonad extra Bool
isUntranslatable Bool
False Term
arg
    case (Bool
untranslatable,Term -> Term
stripTicks Term
arg) of
      (Bool
True,Let Bind Term
binds Term
body) ->
        -- This is a situation similar to Note [CaseLet deshadow]
        let (Bind Term
binds1,Term
body1) = HasCallStack =>
InScopeSet -> Bind Term -> Term -> (Bind Term, Term)
InScopeSet -> Bind Term -> Term -> (Bind Term, Term)
deshadowLetExpr InScopeSet
is0 Bind Term
binds Term
body
        in  Term -> RewriteMonad NormalizeState Term
forall a extra. a -> RewriteMonad extra a
changed (Bind Term -> Term -> Term
Let Bind Term
binds1 (Term -> Term -> Term
App Term
appConPrim Term
body1))
      (Bool
True,Case {})  -> NormRewrite
specialize TransformContext
ctx Term
e
      (Bool
True,Lam {})   -> NormRewrite
specialize TransformContext
ctx Term
e
      (Bool
True,TyLam {}) -> NormRewrite
specialize TransformContext
ctx Term
e
      (Bool, Term)
_               -> Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e

nonRepANF TransformContext
_ Term
e = Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e
{-# SCC nonRepANF #-}