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

  Transformations for specialization
-}

{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE NondecreasingIndentation #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskellQuotes #-}

module Clash.Normalize.Transformations.Specialize
  ( appProp
  , constantSpec
  , specialize
  , nonRepSpec
  , typeSpec
  , zeroWidthSpec
  ) where

import Control.Arrow ((***), (&&&))
import Control.DeepSeq (deepseq)
import Control.Exception (throw)
import Control.Lens ((%=))
import qualified Control.Lens as Lens
import qualified Control.Monad as Monad
import Control.Monad.Extra (orM)
import qualified Control.Monad.Writer as Writer (listen)
import Data.Bifunctor (bimap)
import Data.Coerce (coerce)
import qualified Data.Either as Either
import Data.Functor.Const (Const(..))
import qualified Data.Map.Strict as Map
import qualified Data.Monoid as Monoid (getAny)
import qualified Data.Set.Ordered as OSet
import qualified Data.Set.Ordered.Extra as OSet
import qualified Data.Text as Text
import qualified Data.Text.Extra as Text
import GHC.BasicTypes.Extra (isNoInline)
import GHC.Stack (HasCallStack)

#if MIN_VERSION_ghc(9,0,0)
import GHC.Types.Basic (InlineSpec (..))
#else
import BasicTypes (InlineSpec (..))
#endif

import qualified Clash.Sized.Internal.BitVector as BV (BitVector, fromInteger#)
import qualified Clash.Sized.Internal.Index as I (Index, fromInteger#)
import qualified Clash.Sized.Internal.Signed as S (Signed, fromInteger#)
import qualified Clash.Sized.Internal.Unsigned as U (Unsigned, fromInteger#)

import Clash.Core.DataCon (DataCon(dcArgTys))
import Clash.Core.FreeVars (freeLocalVars, termFreeTyVars, typeFreeVars)
import Clash.Core.HasType
import Clash.Core.Literal (Literal(..))
import Clash.Core.Name
  (NameSort(..), Name(..), appendToName, mkUnsafeInternalName, mkUnsafeSystemName)
import Clash.Core.Pretty (showPpr)
import Clash.Core.Subst
import Clash.Core.Term
  ( Term(..), TickInfo, collectArgs, collectArgsTicks, mkApps, mkTmApps, mkTicks, patIds, Bind(..)
  , patVars, mkAbstraction, PrimInfo(..), WorkInfo(..), IsMultiPrim(..), PrimUnfolding(..), stripAllTicks)
import Clash.Core.TermInfo (isLocalVar, isVar, isPolyFun)
import Clash.Core.TyCon (TyConMap, tyConDataCons)
import Clash.Core.Type
  (LitTy(NumTy), Type(LitTy,VarTy), applyFunTy, splitTyConAppM, normalizeType
  , mkPolyFunTy, mkTyConApp)
import Clash.Core.TysPrim
import Clash.Core.Util (listToLets)
import Clash.Core.Var (Var(..), Id, TyVar, mkTyVar)
import Clash.Core.VarEnv
  ( InScopeSet, extendInScopeSet, extendInScopeSetList, lookupVarEnv
  , mkInScopeSet, mkVarSet, unionInScope, elemVarSet)
import qualified Clash.Data.UniqMap as UniqMap
import Clash.Debug (traceIf, traceM)
import Clash.Driver.Types (Binding(..), TransformationInfo(..), hasTransformationInfo)
import Clash.Netlist.Util (representableType)
import Clash.Rewrite.Combinators (topdownR)
import Clash.Rewrite.Types
  ( TransformContext(..), bindings, censor, curFun, customReprs, extra, tcCache
  , typeTranslator, workFreeBinders, debugOpts, topEntities, specializationLimit)
import Clash.Rewrite.Util
  ( mkBinderFor, mkDerivedName, mkFunction, mkTmBinderFor, setChanged, changed
  , normalizeTermTypes, normalizeId)
import Clash.Rewrite.WorkFree (isWorkFree)
import Clash.Normalize.Types
  ( NormRewrite, NormalizeSession, specialisationCache, specialisationHistory)
import Clash.Normalize.Util
  (constantSpecInfo, csrFoundConstant, csrNewBindings, csrNewTerm)
import Clash.Util (ClashException(..))

-- | Propagate arguments of application inwards; except for 'Lam' where the
-- argument becomes let-bound. 'appProp' tries to propagate as many arguments
-- as possible, down as many levels as possible; and should be called in a
-- top-down traversal.
--
-- The idea is that this reduces the number of traversals, which hopefully leads
-- to shorter compile times.
--
-- Note [AppProp no shadowing]
--
-- Case 1.
--
-- Imagine:
--
-- @
-- (case x of
--    D a b -> h a) (f x y)
-- @
--
-- rewriting this to:
--
-- @
-- let b = f x y
-- in  case x of
--       D a b -> h a b
-- @
--
-- is very bad because @b@ in @h a b@ is now bound by the pattern instead of the
-- newly introduced let-binding
--
-- instead we must deshadow w.r.t. the new variable and rewrite to:
--
-- @
-- let b = f x y
-- in  case x of
--       D a b1 -> h a b
-- @
--
-- Case 2.
--
-- Imagine
--
-- @
-- (\\x -> e) u
-- @
--
-- where @u@ has a free variable named @x@, rewriting this to:
--
-- @
-- let x = u
-- in  e
-- @
--
-- would be very bad, because the let-binding suddenly captures the free
-- variable in @u@. To prevent this from happening we over-approximate and check
-- whether @x@ is in the current InScopeSet, and deshadow if that's the case,
-- i.e. we then rewrite to:
--
-- @
-- let x1 = u
-- in  e [x:=x1]
-- @
--
-- Case 3.
--
-- The same for:
--
-- @
-- (let x = w in e) u
-- @
--
-- where @u@ again has a free variable @x@, rewriting this to:
--
-- @
-- let x = w in (e u)
-- @
--
-- would be bad because the let-binding now captures the free variable in @u@.
--
-- To prevent this from happening, we unconditionally deshadow the function part
-- of the application w.r.t. the free variables in the argument part of the
-- application. It is okay to over-approximate in this case and deshadow w.r.t
-- the current InScopeSet.
appProp :: HasCallStack => NormRewrite
appProp :: NormRewrite
appProp ctx :: TransformContext
ctx@(TransformContext InScopeSet
is Context
_) = \case
  e :: Term
e@App {}
    | let (Term
fun,[Either Term Type]
args,[TickInfo]
ticks) = Term -> (Term, [Either Term Type], [TickInfo])
collectArgsTicks Term
e
    -> do (Term
eN,Any
hasChanged) <- RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState (Term, Any)
forall w (m :: Type -> Type) a. MonadWriter w m => m a -> m (a, w)
Writer.listen (InScopeSet
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad NormalizeState Term
go InScopeSet
is (HasCallStack => InScopeSet -> Term -> Term
InScopeSet -> Term -> Term
deShadowTerm InScopeSet
is Term
fun) [Either Term Type]
args [TickInfo]
ticks)
          if Any -> Bool
Monoid.getAny Any
hasChanged
            then Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
eN
            else Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e
  e :: Term
e@TyApp {}
    | let (Term
fun,[Either Term Type]
args,[TickInfo]
ticks) = Term -> (Term, [Either Term Type], [TickInfo])
collectArgsTicks Term
e
    -> do (Term
eN,Any
hasChanged) <- RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState (Term, Any)
forall w (m :: Type -> Type) a. MonadWriter w m => m a -> m (a, w)
Writer.listen (InScopeSet
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad NormalizeState Term
go InScopeSet
is (HasCallStack => InScopeSet -> Term -> Term
InScopeSet -> Term -> Term
deShadowTerm InScopeSet
is Term
fun) [Either Term Type]
args [TickInfo]
ticks)
          if Any -> Bool
Monoid.getAny Any
hasChanged
            then Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
eN
            else Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e
  Term
e          -> Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e
 where
  go :: InScopeSet -> Term -> [Either Term Type] -> [TickInfo] -> NormalizeSession Term
  go :: InScopeSet
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad NormalizeState Term
go InScopeSet
is0 (Term -> (Term, [Either Term Type], [TickInfo])
collectArgsTicks -> (Term
fun,args0 :: [Either Term Type]
args0@(Either Term Type
_:[Either Term Type]
_),[TickInfo]
ticks0)) [Either Term Type]
args1 [TickInfo]
ticks1 =
    InScopeSet
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad NormalizeState Term
go InScopeSet
is0 Term
fun ([Either Term Type]
args0 [Either Term Type] -> [Either Term Type] -> [Either Term Type]
forall a. [a] -> [a] -> [a]
++ [Either Term Type]
args1) ([TickInfo]
ticks0 [TickInfo] -> [TickInfo] -> [TickInfo]
forall a. [a] -> [a] -> [a]
++ [TickInfo]
ticks1)

  go InScopeSet
is0 (Lam Id
v Term
e) (Left Term
arg:[Either Term Type]
args) [TickInfo]
ticks = do
    RewriteMonad NormalizeState ()
forall extra. RewriteMonad extra ()
setChanged
    BindingMap
bndrs <- Getting BindingMap (RewriteState NormalizeState) BindingMap
-> RewriteMonad NormalizeState BindingMap
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting BindingMap (RewriteState NormalizeState) BindingMap
forall extra. Lens' (RewriteState extra) BindingMap
bindings
    [RewriteMonad NormalizeState Bool]
-> RewriteMonad NormalizeState Bool
forall (m :: Type -> Type). Monad m => [m Bool] -> m Bool
orM [Bool -> RewriteMonad NormalizeState Bool
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Bool
isVar Term
arg), Lens' (RewriteState NormalizeState) (VarEnv Bool)
-> BindingMap -> Term -> RewriteMonad NormalizeState Bool
forall s (m :: Type -> Type).
(HasCallStack, MonadState s m) =>
Lens' s (VarEnv Bool) -> BindingMap -> Term -> m Bool
isWorkFree forall extra. Lens' (RewriteState extra) (VarEnv Bool)
Lens' (RewriteState NormalizeState) (VarEnv Bool)
workFreeBinders BindingMap
bndrs Term
arg] RewriteMonad NormalizeState Bool
-> (Bool -> RewriteMonad NormalizeState Term)
-> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Bool
True ->
        let subst :: Subst
subst = Subst -> Id -> Term -> Subst
extendIdSubst (InScopeSet -> Subst
mkSubst InScopeSet
is0) Id
v Term
arg in
        (Term -> [TickInfo] -> Term
`mkTicks` [TickInfo]
ticks) (Term -> Term)
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> InScopeSet
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad NormalizeState Term
go InScopeSet
is0 (HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
"appProp.AppLam" Subst
subst Term
e) [Either Term Type]
args []
      Bool
False ->
        let is1 :: InScopeSet
is1 = InScopeSet -> Id -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
extendInScopeSet InScopeSet
is0 Id
v in
        Bind Term -> Term -> Term
Let (Id -> Term -> Bind Term
forall a. Id -> a -> Bind a
NonRec Id
v Term
arg) (Term -> Term)
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> InScopeSet
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad NormalizeState Term
go InScopeSet
is1 (HasCallStack => InScopeSet -> Term -> Term
InScopeSet -> Term -> Term
deShadowTerm InScopeSet
is1 Term
e) [Either Term Type]
args [TickInfo]
ticks

  go InScopeSet
is0 (Let (NonRec Id
i Term
x) Term
e) args :: [Either Term Type]
args@(Either Term Type
_:[Either Term Type]
_) [TickInfo]
ticks = do
    RewriteMonad NormalizeState ()
forall extra. RewriteMonad extra ()
setChanged
    let is1 :: InScopeSet
is1 = InScopeSet -> Id -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
extendInScopeSet InScopeSet
is0 Id
i
    -- XXX: binding should already be deshadowed w.r.t. 'is0'
    Bind Term -> Term -> Term
Let (Id -> Term -> Bind Term
forall a. Id -> a -> Bind a
NonRec Id
i Term
x) (Term -> Term)
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> InScopeSet
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad NormalizeState Term
go InScopeSet
is1 Term
e [Either Term Type]
args [TickInfo]
ticks

  go InScopeSet
is0 (Let (Rec [(Id, Term)]
vs) Term
e) args :: [Either Term Type]
args@(Either Term Type
_:[Either Term Type]
_) [TickInfo]
ticks = do
    RewriteMonad NormalizeState ()
forall extra. RewriteMonad extra ()
setChanged
    let vbs :: [Id]
vbs  = ((Id, Term) -> Id) -> [(Id, Term)] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map (Id, Term) -> Id
forall a b. (a, b) -> a
fst [(Id, Term)]
vs
        is1 :: InScopeSet
is1  = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 [Id]
vbs
    -- XXX: 'vs' should already be deshadowed w.r.t. 'is0'
    Bind Term -> Term -> Term
Let ([(Id, Term)] -> Bind Term
forall a. [(Id, a)] -> Bind a
Rec [(Id, Term)]
vs) (Term -> Term)
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> InScopeSet
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad NormalizeState Term
go InScopeSet
is1 Term
e [Either Term Type]
args [TickInfo]
ticks

  go InScopeSet
is0 (TyLam TyVar
tv Term
e) (Right Type
t:[Either Term Type]
args) [TickInfo]
ticks = do
    RewriteMonad NormalizeState ()
forall extra. RewriteMonad extra ()
setChanged
    let subst :: Subst
subst = Subst -> TyVar -> Type -> Subst
extendTvSubst (InScopeSet -> Subst
mkSubst InScopeSet
is0) TyVar
tv Type
t
    (Term -> [TickInfo] -> Term
`mkTicks` [TickInfo]
ticks) (Term -> Term)
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> InScopeSet
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad NormalizeState Term
go InScopeSet
is0 (HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
"appProp.TyAppTyLam" Subst
subst Term
e) [Either Term Type]
args []

  go InScopeSet
is0 (Case Term
scrut Type
ty0 [Alt]
alts) args0 :: [Either Term Type]
args0@(Either Term Type
_:[Either Term Type]
_) [TickInfo]
ticks = do
    RewriteMonad NormalizeState ()
forall extra. RewriteMonad extra ()
setChanged
    let isA1 :: InScopeSet
isA1 = InScopeSet -> InScopeSet -> InScopeSet
unionInScope
                 InScopeSet
is0
                 ((VarSet -> InScopeSet
mkInScopeSet (VarSet -> InScopeSet) -> ([Alt] -> VarSet) -> [Alt] -> InScopeSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Var Any] -> VarSet
forall a. [Var a] -> VarSet
mkVarSet ([Var Any] -> VarSet) -> ([Alt] -> [Var Any]) -> [Alt] -> VarSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Alt -> [Var Any]) -> [Alt] -> [Var Any]
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> [b]) -> t a -> [b]
concatMap (Pat -> [Var Any]
forall a. Pat -> [Var a]
patVars (Pat -> [Var Any]) -> (Alt -> Pat) -> Alt -> [Var Any]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alt -> Pat
forall a b. (a, b) -> a
fst)) [Alt]
alts)
    (Type
ty1,[(Id, Term)]
vs,[Either Term Type]
args1) <- InScopeSet
-> Type
-> [(Id, Term)]
-> [Either Term Type]
-> RewriteMonad
     NormalizeState (Type, [(Id, Term)], [Either Term Type])
forall extra (m :: Type -> Type).
(MonadState (RewriteState extra) m, MonadReader RewriteEnv m,
 MonadUnique m) =>
InScopeSet
-> Type
-> [(Id, Term)]
-> [Either Term Type]
-> m (Type, [(Id, Term)], [Either Term Type])
goCaseArg InScopeSet
isA1 Type
ty0 [] [Either Term Type]
args0
    case [(Id, Term)]
vs of
      [] -> (Term -> [TickInfo] -> Term
`mkTicks` [TickInfo]
ticks) (Term -> Term) -> ([Alt] -> Term) -> [Alt] -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> Type -> [Alt] -> Term
Case Term
scrut Type
ty1 ([Alt] -> Term)
-> RewriteMonad NormalizeState [Alt]
-> RewriteMonad NormalizeState Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (Alt -> RewriteMonad NormalizeState Alt)
-> [Alt] -> RewriteMonad NormalizeState [Alt]
forall (t :: Type -> Type) (m :: Type -> Type) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (InScopeSet
-> [Either Term Type] -> Alt -> RewriteMonad NormalizeState Alt
goAlt InScopeSet
is0 [Either Term Type]
args1) [Alt]
alts
      [(Id, Term)]
_  -> do
        let vbs :: [Id]
vbs   = ((Id, Term) -> Id) -> [(Id, Term)] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map (Id, Term) -> Id
forall a b. (a, b) -> a
fst [(Id, Term)]
vs
            is1 :: InScopeSet
is1   = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 [Id]
vbs
            alts1 :: [Alt]
alts1 = (Alt -> Alt) -> [Alt] -> [Alt]
forall a b. (a -> b) -> [a] -> [b]
map (HasCallStack => InScopeSet -> Alt -> Alt
InScopeSet -> Alt -> Alt
deShadowAlt InScopeSet
is1) [Alt]
alts
        -- TODO I should have a mkNonRecLets :: [LetBinding] -> Term -> Term
        -- function which makes a chain of non-recursive let expressions without
        -- needing to first take the SCCs of all the binders.
        [(Id, Term)] -> Term -> Term
listToLets [(Id, Term)]
vs (Term -> Term) -> ([Alt] -> Term) -> [Alt] -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term -> [TickInfo] -> Term
`mkTicks` [TickInfo]
ticks) (Term -> Term) -> ([Alt] -> Term) -> [Alt] -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> Type -> [Alt] -> Term
Case Term
scrut Type
ty1 ([Alt] -> Term)
-> RewriteMonad NormalizeState [Alt]
-> RewriteMonad NormalizeState Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (Alt -> RewriteMonad NormalizeState Alt)
-> [Alt] -> RewriteMonad NormalizeState [Alt]
forall (t :: Type -> Type) (m :: Type -> Type) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (InScopeSet
-> [Either Term Type] -> Alt -> RewriteMonad NormalizeState Alt
goAlt InScopeSet
is1 [Either Term Type]
args1) [Alt]
alts1

  go InScopeSet
is0 (Tick TickInfo
sp Term
e) [Either Term Type]
args [TickInfo]
ticks = do
    RewriteMonad NormalizeState ()
forall extra. RewriteMonad extra ()
setChanged
    InScopeSet
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad NormalizeState Term
go InScopeSet
is0 Term
e [Either Term Type]
args (TickInfo
spTickInfo -> [TickInfo] -> [TickInfo]
forall a. a -> [a] -> [a]
:[TickInfo]
ticks)

  go InScopeSet
_ Term
fun [Either Term Type]
args [TickInfo]
ticks = Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Term -> [Either Term Type] -> Term
mkApps (Term -> [TickInfo] -> Term
mkTicks Term
fun [TickInfo]
ticks) [Either Term Type]
args)

  goAlt :: InScopeSet
-> [Either Term Type] -> Alt -> RewriteMonad NormalizeState Alt
goAlt InScopeSet
is0 [Either Term Type]
args0 (Pat
p,Term
e) = do
    let ([TyVar]
tvs,[Id]
ids) = Pat -> ([TyVar], [Id])
patIds Pat
p
        is1 :: InScopeSet
is1       = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList (InScopeSet -> [TyVar] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 [TyVar]
tvs) [Id]
ids
    (Pat
p,) (Term -> Alt)
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState Alt
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> InScopeSet
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad NormalizeState Term
go InScopeSet
is1 Term
e [Either Term Type]
args0 []

  goCaseArg :: InScopeSet
-> Type
-> [(Id, Term)]
-> [Either Term Type]
-> m (Type, [(Id, Term)], [Either Term Type])
goCaseArg InScopeSet
isA Type
ty0 [(Id, Term)]
ls0 (Right Type
t:[Either Term Type]
args0) = do
    TyConMap
tcm <- Getting TyConMap RewriteEnv TyConMap -> m 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 ty1 :: Type
ty1 = HasCallStack => TyConMap -> Type -> Type -> Type
TyConMap -> Type -> Type -> Type
piResultTy TyConMap
tcm Type
ty0 Type
t
    (Type
ty2,[(Id, Term)]
ls1,[Either Term Type]
args1) <- InScopeSet
-> Type
-> [(Id, Term)]
-> [Either Term Type]
-> m (Type, [(Id, Term)], [Either Term Type])
goCaseArg InScopeSet
isA Type
ty1 [(Id, Term)]
ls0 [Either Term Type]
args0
    (Type, [(Id, Term)], [Either Term Type])
-> m (Type, [(Id, Term)], [Either Term Type])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Type
ty2,[(Id, Term)]
ls1,Type -> Either Term Type
forall a b. b -> Either a b
Right Type
tEither Term Type -> [Either Term Type] -> [Either Term Type]
forall a. a -> [a] -> [a]
:[Either Term Type]
args1)

  goCaseArg InScopeSet
isA0 Type
ty0 [(Id, Term)]
ls0 (Left Term
arg:[Either Term Type]
args0) = do
    TyConMap
tcm <- Getting TyConMap RewriteEnv TyConMap -> m 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
    BindingMap
bndrs <- Getting BindingMap (RewriteState extra) BindingMap -> m BindingMap
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting BindingMap (RewriteState extra) BindingMap
forall extra. Lens' (RewriteState extra) BindingMap
bindings
    let argTy :: Type
argTy = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
arg
        ty1 :: Type
ty1   = TyConMap -> Type -> Type -> Type
applyFunTy TyConMap
tcm Type
ty0 Type
argTy
    [m Bool] -> m Bool
forall (m :: Type -> Type). Monad m => [m Bool] -> m Bool
orM [Bool -> m Bool
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Bool
isVar Term
arg), Lens' (RewriteState extra) (VarEnv Bool)
-> BindingMap -> Term -> m Bool
forall s (m :: Type -> Type).
(HasCallStack, MonadState s m) =>
Lens' s (VarEnv Bool) -> BindingMap -> Term -> m Bool
isWorkFree forall extra. Lens' (RewriteState extra) (VarEnv Bool)
Lens' (RewriteState extra) (VarEnv Bool)
workFreeBinders BindingMap
bndrs Term
arg] m Bool
-> (Bool -> m (Type, [(Id, Term)], [Either Term Type]))
-> m (Type, [(Id, Term)], [Either Term Type])
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Bool
True -> do
        (Type
ty2,[(Id, Term)]
ls1,[Either Term Type]
args1) <- InScopeSet
-> Type
-> [(Id, Term)]
-> [Either Term Type]
-> m (Type, [(Id, Term)], [Either Term Type])
goCaseArg InScopeSet
isA0 Type
ty1 [(Id, Term)]
ls0 [Either Term Type]
args0
        (Type, [(Id, Term)], [Either Term Type])
-> m (Type, [(Id, Term)], [Either Term Type])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Type
ty2,[(Id, Term)]
ls1,Term -> Either Term Type
forall a b. a -> Either a b
Left Term
argEither Term Type -> [Either Term Type] -> [Either Term Type]
forall a. a -> [a] -> [a]
:[Either Term Type]
args1)
      Bool
False -> do
        Id
boundArg <- InScopeSet -> TyConMap -> Name Term -> Term -> m Id
forall (m :: Type -> Type) a.
MonadUnique m =>
InScopeSet -> TyConMap -> Name a -> Term -> m Id
mkTmBinderFor InScopeSet
isA0 TyConMap
tcm (TransformContext -> OccName -> Name Term
mkDerivedName TransformContext
ctx OccName
"app_arg") Term
arg
        let isA1 :: InScopeSet
isA1 = InScopeSet -> Id -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
extendInScopeSet InScopeSet
isA0 Id
boundArg
        (Type
ty2,[(Id, Term)]
ls1,[Either Term Type]
args1) <- InScopeSet
-> Type
-> [(Id, Term)]
-> [Either Term Type]
-> m (Type, [(Id, Term)], [Either Term Type])
goCaseArg InScopeSet
isA1 Type
ty1 [(Id, Term)]
ls0 [Either Term Type]
args0
        (Type, [(Id, Term)], [Either Term Type])
-> m (Type, [(Id, Term)], [Either Term Type])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Type
ty2,(Id
boundArg,Term
arg)(Id, Term) -> [(Id, Term)] -> [(Id, Term)]
forall a. a -> [a] -> [a]
:[(Id, Term)]
ls1,Term -> Either Term Type
forall a b. a -> Either a b
Left (Id -> Term
Var Id
boundArg)Either Term Type -> [Either Term Type] -> [Either Term Type]
forall a. a -> [a] -> [a]
:[Either Term Type]
args1)

  goCaseArg InScopeSet
_ Type
ty [(Id, Term)]
ls [] = (Type, [(Id, Term)], [Either Term Type])
-> m (Type, [(Id, Term)], [Either Term Type])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Type
ty,[(Id, Term)]
ls,[])
{-# SCC appProp #-}

-- | Specialize functions on arguments which are constant, except when they
-- are clock, reset generators.
constantSpec :: HasCallStack => NormRewrite
constantSpec :: NormRewrite
constantSpec ctx :: TransformContext
ctx@(TransformContext InScopeSet
is0 Context
tfCtx) e :: Term
e@(App Term
e1 Term
e2)
  | (Var {}, [Either Term Type]
args) <- Term -> (Term, [Either Term Type])
collectArgs Term
e1
  , ([Term]
_, []) <- [Either Term Type] -> ([Term], [Type])
forall a b. [Either a b] -> ([a], [b])
Either.partitionEithers [Either Term Type]
args
  , [TyVar] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
null ([TyVar] -> Bool) -> [TyVar] -> Bool
forall a b. (a -> b) -> a -> b
$ Getting (Endo [TyVar]) Term TyVar -> Term -> [TyVar]
forall a s. Getting (Endo [a]) s a -> s -> [a]
Lens.toListOf Getting (Endo [TyVar]) Term TyVar
Fold Term TyVar
termFreeTyVars Term
e2
  = do ConstantSpecInfo
specInfo<- TransformContext
-> Term -> RewriteMonad NormalizeState ConstantSpecInfo
constantSpecInfo TransformContext
ctx Term
e2
       if ConstantSpecInfo -> Bool
csrFoundConstant ConstantSpecInfo
specInfo then
         let newBindings :: [(Id, Term)]
newBindings = ConstantSpecInfo -> [(Id, Term)]
csrNewBindings ConstantSpecInfo
specInfo in
         if [(Id, Term)] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
null [(Id, Term)]
newBindings then
           -- Whole of e2 is constant
           NormRewrite
specialize TransformContext
ctx (Term -> Term -> Term
App Term
e1 Term
e2)
         else do
           -- Parts of e2 are constant
           let is1 :: InScopeSet
is1 = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 ((Id, Term) -> Id
forall a b. (a, b) -> a
fst ((Id, Term) -> Id) -> [(Id, Term)] -> [Id]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ConstantSpecInfo -> [(Id, Term)]
csrNewBindings ConstantSpecInfo
specInfo)
           (Term
body, Any
isSpec) <- RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState (Term, Any)
forall w (m :: Type -> Type) a. MonadWriter w m => m a -> m (a, w)
Writer.listen (RewriteMonad NormalizeState Term
 -> RewriteMonad NormalizeState (Term, Any))
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState (Term, Any)
forall a b. (a -> b) -> a -> b
$ NormRewrite
specialize
             (InScopeSet -> Context -> TransformContext
TransformContext InScopeSet
is1 Context
tfCtx)
             (Term -> Term -> Term
App Term
e1 (ConstantSpecInfo -> Term
csrNewTerm ConstantSpecInfo
specInfo))

           if Any -> Bool
Monoid.getAny Any
isSpec
             then Term -> RewriteMonad NormalizeState Term
forall a extra. a -> RewriteMonad extra a
changed ([(Id, Term)] -> Term -> Term
listToLets [(Id, Term)]
newBindings Term
body)
             else Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e
       else
        -- e2 has no constant parts
        Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e
constantSpec TransformContext
_ Term
e = Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e
{-# SCC constantSpec #-}

-- | Specialize an application on its argument
specialize :: NormRewrite
specialize :: NormRewrite
specialize TransformContext
ctx Term
e = case Term
e of
  (TyApp Term
e1 Type
ty) -> TransformContext
-> Term
-> (Term, [Either Term Type], [TickInfo])
-> Either Term Type
-> RewriteMonad NormalizeState Term
specialize' TransformContext
ctx Term
e (Term -> (Term, [Either Term Type], [TickInfo])
collectArgsTicks Term
e1) (Type -> Either Term Type
forall a b. b -> Either a b
Right Type
ty)
  (App Term
e1 Term
e2)   -> TransformContext
-> Term
-> (Term, [Either Term Type], [TickInfo])
-> Either Term Type
-> RewriteMonad NormalizeState Term
specialize' TransformContext
ctx Term
e (Term -> (Term, [Either Term Type], [TickInfo])
collectArgsTicks Term
e1) (Term -> Either Term Type
forall a b. a -> Either a b
Left  Term
e2)
  Term
_             -> Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e

{-
Note [ticks and specialization]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
As Clash now distinguishes between ticks in expressions when comparing for
alpha equality, this has a knock-on effect when accessing the specialization
cache. Consider these applications which differ only by ticks:

    f[GlobalId] (\x -> ... x[LocalId])
    f[GlobalId] <tick>(\x -> ... x[LocalId])
    f[GlobalId] (\x -> ... <tick>x[LocalId])

If one of these had been specialized, the other two would hit that term in the
specialization cache, saving Clash from having to re-do work which is in effect
the same. To preserve this behavior, we use 'stripAllTicks' on the keys for
the specialization cache.

TODO While this preserves the old behavior, the old behavior is likely not
quite what we want. Using a value from the specialization cache may change the
ticks present, which can affect naming / debugging information in generated HDL.
We may also not want to look at ticks, as then the specialization cache will
miss on virtually every lookup which could add to normalization time.
-}

-- | Given two 'InlineSpec's, return the \"strongest\" one. I.e., the one that's
-- closest to @NoInline@ (or @Opaque@ for newer GHCs).
preferNoInline :: InlineSpec -> InlineSpec -> InlineSpec
preferNoInline :: InlineSpec -> InlineSpec -> InlineSpec
preferNoInline InlineSpec
is0 InlineSpec
is1
  | InlineSpec -> Int
enumInlineSpec InlineSpec
is0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= InlineSpec -> Int
enumInlineSpec InlineSpec
is1 = InlineSpec
is0
  | Bool
otherwise                                = InlineSpec
is1
 where
  enumInlineSpec :: InlineSpec -> Int
  enumInlineSpec :: InlineSpec -> Int
enumInlineSpec = \case
#if MIN_VERSION_ghc(9,2,0)
    NoUserInlinePrag {} -> -1
#else
    NoUserInline {} -> -Int
1
#endif
    Inline {} -> Int
0
    Inlinable {} -> Int
1
    NoInline {} -> Int
2
#if MIN_VERSION_ghc(9,4,0)
    Opaque {} -> 3
#endif

-- | Specialize an application on its argument
specialize'
  :: TransformContext
  -- ^ Transformation context
  -> Term
  -- ^ Original term
  -> (Term, [Either Term Type], [TickInfo])
  -- ^ Function part of the term, split into root and applied arguments
  -> Either Term Type
  -- ^ Argument to specialize on
  -> NormalizeSession Term
specialize' :: TransformContext
-> Term
-> (Term, [Either Term Type], [TickInfo])
-> Either Term Type
-> RewriteMonad NormalizeState Term
specialize' (TransformContext InScopeSet
is0 Context
_) Term
e (Var Id
f, [Either Term Type]
args, [TickInfo]
ticks) Either Term Type
specArgIn = do
  DebugOpts
opts <- Getting DebugOpts RewriteEnv DebugOpts
-> RewriteMonad NormalizeState DebugOpts
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting DebugOpts RewriteEnv DebugOpts
Getter RewriteEnv DebugOpts
debugOpts
  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

  -- Don't specialize TopEntities
  VarSet
topEnts <- Getting VarSet RewriteEnv VarSet
-> RewriteMonad NormalizeState VarSet
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting VarSet RewriteEnv VarSet
Lens' RewriteEnv VarSet
topEntities
  if Id
f Id -> VarSet -> Bool
forall a. Var a -> VarSet -> Bool
`elemVarSet` VarSet
topEnts
  then do
    case Either Term Type
specArgIn of
      Left Term
_ -> do
        String -> RewriteMonad NormalizeState ()
forall (f :: Type -> Type). Applicative f => String -> f ()
traceM (String
"Not specializing TopEntity: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name Term -> String
forall p. PrettyPrec p => p -> String
showPpr (Id -> Name Term
forall a. Var a -> Name a
varName Id
f))
        Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e
      Right Type
tyArg ->
        Bool
-> String
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState Term
forall a. Bool -> String -> a -> a
traceIf (TransformationInfo -> DebugOpts -> Bool
hasTransformationInfo TransformationInfo
AppliedTerm DebugOpts
opts) (String
"Dropping type application on TopEntity: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name Term -> String
forall p. PrettyPrec p => p -> String
showPpr (Id -> Name Term
forall a. Var a -> Name a
varName Id
f) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"\ntype:\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall p. PrettyPrec p => p -> String
showPpr Type
tyArg) (RewriteMonad NormalizeState Term
 -> RewriteMonad NormalizeState Term)
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState Term
forall a b. (a -> b) -> a -> b
$
        -- TopEntities aren't allowed to be semantically polymorphic.
        -- But using type equality constraints they may be syntactically polymorphic.
        -- > topEntity :: forall dom . (dom ~ "System") => Signal dom Bool -> Signal dom Bool
        -- The TyLam's in the body will have been removed by 'Clash.Normalize.Util.substWithTyEq'.
        -- So we drop the TyApp ("specializing" on it) and change the varType to match.
        let newVarTy :: Type
newVarTy = HasCallStack => TyConMap -> Type -> Type -> Type
TyConMap -> Type -> Type -> Type
piResultTy TyConMap
tcm (Id -> Type
forall a. HasType a => a -> Type
coreTypeOf Id
f) Type
tyArg
        in  Term -> RewriteMonad NormalizeState Term
forall a extra. a -> RewriteMonad extra a
changed (Term -> [Either Term Type] -> Term
mkApps (Term -> [TickInfo] -> Term
mkTicks (Id -> Term
Var Id
f{varType :: Type
varType = Type
newVarTy}) [TickInfo]
ticks) [Either Term Type]
args)
  else do -- NondecreasingIndentation

  let specArg :: Either Term Type
specArg = (Term -> Term)
-> (Type -> Type) -> Either Term Type -> Either Term Type
forall (p :: Type -> Type -> Type) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (TyConMap -> Term -> Term
normalizeTermTypes TyConMap
tcm) (TyConMap -> Type -> Type
normalizeType TyConMap
tcm) Either Term Type
specArgIn
      -- Create binders and variable references for free variables in 'specArg'
      -- (specBndrsIn,specVars) :: ([Either Id TyVar], [Either Term Type])
      ([Either Id TyVar]
specBndrsIn,[Either Term Type]
specVars) = Either Term Type -> ([Either Id TyVar], [Either Term Type])
specArgBndrsAndVars Either Term Type
specArg
      argLen :: Int
argLen  = [Either Term Type] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Either Term Type]
args
      specBndrs :: [Either Id TyVar]
      specBndrs :: [Either Id TyVar]
specBndrs = (Either Id TyVar -> Either Id TyVar)
-> [Either Id TyVar] -> [Either Id TyVar]
forall a b. (a -> b) -> [a] -> [b]
map (ASetter (Either Id TyVar) (Either Id TyVar) Id Id
-> (Id -> Id) -> Either Id TyVar -> Either Id TyVar
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
Lens.over ASetter (Either Id TyVar) (Either Id TyVar) Id Id
forall a c b. Prism (Either a c) (Either b c) a b
Lens._Left (TyConMap -> Id -> Id
normalizeId TyConMap
tcm)) [Either Id TyVar]
specBndrsIn

      -- See Note [ticks and specialization]
      specAbs :: Either Term Type
      specAbs :: Either Term Type
specAbs = (Term -> Either Term Type)
-> (Type -> Either Term Type)
-> Either Term Type
-> Either Term Type
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Term -> Either Term Type
forall a b. a -> Either a b
Left (Term -> Either Term Type)
-> (Term -> Term) -> Term -> Either Term Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> Term
stripAllTicks (Term -> Term) -> (Term -> Term) -> Term -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term -> [Either Id TyVar] -> Term
`mkAbstraction` [Either Id TyVar]
specBndrs)) (Type -> Either Term Type
forall a b. b -> Either a b
Right (Type -> Either Term Type)
-> (Type -> Type) -> Type -> Either Term Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type
forall a. a -> a
id) Either Term Type
specArg
  -- Determine if 'f' has already been specialized on (a type-normalized) 'specArg'
  Maybe Id
specM <- (Id, Int, Either Term Type)
-> Map (Id, Int, Either Term Type) Id -> Maybe Id
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Id
f,Int
argLen,Either Term Type
specAbs) (Map (Id, Int, Either Term Type) Id -> Maybe Id)
-> RewriteMonad NormalizeState (Map (Id, Int, Either Term Type) Id)
-> RewriteMonad NormalizeState (Maybe Id)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Getting
  (Map (Id, Int, Either Term Type) Id)
  (RewriteState NormalizeState)
  (Map (Id, Int, Either Term Type) Id)
-> RewriteMonad NormalizeState (Map (Id, Int, Either Term Type) Id)
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use ((NormalizeState
 -> Const (Map (Id, Int, Either Term Type) Id) NormalizeState)
-> RewriteState NormalizeState
-> Const
     (Map (Id, Int, Either Term Type) Id) (RewriteState NormalizeState)
forall extra extra2.
Lens (RewriteState extra) (RewriteState extra2) extra extra2
extra((NormalizeState
  -> Const (Map (Id, Int, Either Term Type) Id) NormalizeState)
 -> RewriteState NormalizeState
 -> Const
      (Map (Id, Int, Either Term Type) Id) (RewriteState NormalizeState))
-> ((Map (Id, Int, Either Term Type) Id
     -> Const
          (Map (Id, Int, Either Term Type) Id)
          (Map (Id, Int, Either Term Type) Id))
    -> NormalizeState
    -> Const (Map (Id, Int, Either Term Type) Id) NormalizeState)
-> Getting
     (Map (Id, Int, Either Term Type) Id)
     (RewriteState NormalizeState)
     (Map (Id, Int, Either Term Type) Id)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Map (Id, Int, Either Term Type) Id
 -> Const
      (Map (Id, Int, Either Term Type) Id)
      (Map (Id, Int, Either Term Type) Id))
-> NormalizeState
-> Const (Map (Id, Int, Either Term Type) Id) NormalizeState
Lens' NormalizeState (Map (Id, Int, Either Term Type) Id)
specialisationCache)
  case Maybe Id
specM of
    -- Use previously specialized function
    Just Id
f' ->
      Bool
-> String
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState Term
forall a. Bool -> String -> a -> a
traceIf (TransformationInfo -> DebugOpts -> Bool
hasTransformationInfo TransformationInfo
AppliedTerm DebugOpts
opts)
        (String
"Using previous specialization of " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name Term -> String
forall p. PrettyPrec p => p -> String
showPpr (Id -> Name Term
forall a. Var a -> Name a
varName Id
f) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" on " String -> String -> String
forall a. [a] -> [a] -> [a]
++
          ((Term -> String) -> (Type -> String) -> Either Term Type -> String
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either Term -> String
forall p. PrettyPrec p => p -> String
showPpr Type -> String
forall p. PrettyPrec p => p -> String
showPpr) Either Term Type
specAbs String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name Term -> String
forall p. PrettyPrec p => p -> String
showPpr (Id -> Name Term
forall a. Var a -> Name a
varName Id
f')) (RewriteMonad NormalizeState Term
 -> RewriteMonad NormalizeState Term)
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState Term
forall a b. (a -> b) -> a -> b
$
        Term -> RewriteMonad NormalizeState Term
forall a extra. a -> RewriteMonad extra a
changed (Term -> RewriteMonad NormalizeState Term)
-> Term -> RewriteMonad NormalizeState Term
forall a b. (a -> b) -> a -> b
$ Term -> [Either Term Type] -> Term
mkApps (Term -> [TickInfo] -> Term
mkTicks (Id -> Term
Var Id
f') [TickInfo]
ticks) ([Either Term Type]
args [Either Term Type] -> [Either Term Type] -> [Either Term Type]
forall a. [a] -> [a] -> [a]
++ [Either Term Type]
specVars)
    -- Create new specialized function
    Maybe Id
Nothing -> do
      -- Determine if we can specialize f
      Maybe (Binding Term)
bodyMaybe <- (BindingMap -> Maybe (Binding Term))
-> RewriteMonad NormalizeState BindingMap
-> RewriteMonad NormalizeState (Maybe (Binding Term))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (Id -> BindingMap -> Maybe (Binding Term)
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup Id
f) (RewriteMonad NormalizeState BindingMap
 -> RewriteMonad NormalizeState (Maybe (Binding Term)))
-> RewriteMonad NormalizeState BindingMap
-> RewriteMonad NormalizeState (Maybe (Binding Term))
forall a b. (a -> b) -> a -> b
$ Getting BindingMap (RewriteState NormalizeState) BindingMap
-> RewriteMonad NormalizeState BindingMap
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting BindingMap (RewriteState NormalizeState) BindingMap
forall extra. Lens' (RewriteState extra) BindingMap
bindings
      case Maybe (Binding Term)
bodyMaybe of
        Just (Binding Id
_ SrcSpan
sp InlineSpec
inl IsPrim
_ Term
bodyTm Bool
_) -> do
          -- Determine if we see a sequence of specializations on a growing argument
          Maybe Int
specHistM <- Id -> UniqMap Int -> Maybe Int
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup Id
f (UniqMap Int -> Maybe Int)
-> RewriteMonad NormalizeState (UniqMap Int)
-> RewriteMonad NormalizeState (Maybe Int)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Getting (UniqMap Int) (RewriteState NormalizeState) (UniqMap Int)
-> RewriteMonad NormalizeState (UniqMap Int)
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use ((NormalizeState -> Const (UniqMap Int) NormalizeState)
-> RewriteState NormalizeState
-> Const (UniqMap Int) (RewriteState NormalizeState)
forall extra extra2.
Lens (RewriteState extra) (RewriteState extra2) extra extra2
extra((NormalizeState -> Const (UniqMap Int) NormalizeState)
 -> RewriteState NormalizeState
 -> Const (UniqMap Int) (RewriteState NormalizeState))
-> ((UniqMap Int -> Const (UniqMap Int) (UniqMap Int))
    -> NormalizeState -> Const (UniqMap Int) NormalizeState)
-> Getting
     (UniqMap Int) (RewriteState NormalizeState) (UniqMap Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(UniqMap Int -> Const (UniqMap Int) (UniqMap Int))
-> NormalizeState -> Const (UniqMap Int) NormalizeState
Lens' NormalizeState (UniqMap Int)
specialisationHistory)
          Int
specLim   <- Getting Int RewriteEnv Int -> RewriteMonad NormalizeState Int
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting Int RewriteEnv Int
Getter RewriteEnv Int
specializationLimit
          if Bool -> (Int -> Bool) -> Maybe Int -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
specLim) Maybe Int
specHistM
            then ClashException -> RewriteMonad NormalizeState Term
forall a e. Exception e => e -> a
throw (SrcSpan -> String -> Maybe String -> ClashException
ClashException
                        SrcSpan
sp
                        ([String] -> String
unlines [ String
"Hit specialization limit " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
specLim String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" on function `" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name Term -> String
forall p. PrettyPrec p => p -> String
showPpr (Id -> Name Term
forall a. Var a -> Name a
varName Id
f) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"'.\n"
                                 , String
"The function `" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Id -> String
forall p. PrettyPrec p => p -> String
showPpr Id
f String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"' is most likely recursive, and looks like it is being indefinitely specialized on a growing argument.\n"
                                 , String
"Body of `" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Id -> String
forall p. PrettyPrec p => p -> String
showPpr Id
f String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"':\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Term -> String
forall p. PrettyPrec p => p -> String
showPpr Term
bodyTm String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"\n"
                                 , String
"Argument (in position: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
argLen String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
") that triggered termination:\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ ((Term -> String) -> (Type -> String) -> Either Term Type -> String
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either Term -> String
forall p. PrettyPrec p => p -> String
showPpr Type -> String
forall p. PrettyPrec p => p -> String
showPpr) Either Term Type
specArg
                                 , String
"Run with '-fclash-spec-limit=N' to increase the specialization limit to N."
                                 ])
                        Maybe String
forall a. Maybe a
Nothing)
            else do
              let existingNames :: [Name a]
existingNames = Term -> [Name a]
forall a. Term -> [Name a]
collectBndrsMinusApps Term
bodyTm
                  newNames :: [Name a]
newNames      = [ OccName -> Int -> Name a
forall a. OccName -> Int -> Name a
mkUnsafeInternalName (OccName
"pTS" OccName -> OccName -> OccName
`Text.append` String -> OccName
Text.pack (Int -> String
forall a. Show a => a -> String
show Int
n)) Int
n
                                  | Int
n <- [(Int
0::Int)..]
                                  ]
              -- Make new binders for existing arguments
              ([Either Id TyVar]
boundArgs,[Either Term Type]
argVars) <- ([Either Id TyVar] -> ([Either Id TyVar], [Either Term Type]))
-> RewriteMonad NormalizeState [Either Id TyVar]
-> RewriteMonad
     NormalizeState ([Either Id TyVar], [Either Term Type])
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap ([(Either Id TyVar, Either Term Type)]
-> ([Either Id TyVar], [Either Term Type])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Either Id TyVar, Either Term Type)]
 -> ([Either Id TyVar], [Either Term Type]))
-> ([Either Id TyVar] -> [(Either Id TyVar, Either Term Type)])
-> [Either Id TyVar]
-> ([Either Id TyVar], [Either Term Type])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Either Id TyVar -> (Either Id TyVar, Either Term Type))
-> [Either Id TyVar] -> [(Either Id TyVar, Either Term Type)]
forall a b. (a -> b) -> [a] -> [b]
map ((Id -> (Either Id TyVar, Either Term Type))
-> (TyVar -> (Either Id TyVar, Either Term Type))
-> Either Id TyVar
-> (Either Id TyVar, Either Term Type)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Id -> Either Id TyVar
forall a b. a -> Either a b
Left (Id -> Either Id TyVar)
-> (Id -> Either Term Type)
-> Id
-> (Either Id TyVar, Either Term Type)
forall (a :: Type -> Type -> Type) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& Term -> Either Term Type
forall a b. a -> Either a b
Left (Term -> Either Term Type)
-> (Id -> Term) -> Id -> Either Term Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> Term
Var) (TyVar -> Either Id TyVar
forall a b. b -> Either a b
Right (TyVar -> Either Id TyVar)
-> (TyVar -> Either Term Type)
-> TyVar
-> (Either Id TyVar, Either Term Type)
forall (a :: Type -> Type -> Type) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& Type -> Either Term Type
forall a b. b -> Either a b
Right (Type -> Either Term Type)
-> (TyVar -> Type) -> TyVar -> Either Term Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyVar -> Type
VarTy))) (RewriteMonad NormalizeState [Either Id TyVar]
 -> RewriteMonad
      NormalizeState ([Either Id TyVar], [Either Term Type]))
-> RewriteMonad NormalizeState [Either Id TyVar]
-> RewriteMonad
     NormalizeState ([Either Id TyVar], [Either Term Type])
forall a b. (a -> b) -> a -> b
$
                                     (Name Any
 -> Either Term Type
 -> RewriteMonad NormalizeState (Either Id TyVar))
-> [Name Any]
-> [Either Term Type]
-> RewriteMonad NormalizeState [Either Id TyVar]
forall (m :: Type -> Type) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
Monad.zipWithM
                                       (InScopeSet
-> TyConMap
-> Name Any
-> Either Term Type
-> RewriteMonad NormalizeState (Either Id TyVar)
forall (m :: Type -> Type) a.
MonadUnique m =>
InScopeSet
-> TyConMap -> Name a -> Either Term Type -> m (Either Id TyVar)
mkBinderFor InScopeSet
is0 TyConMap
tcm)
                                       ([Name Any]
forall a. [Name a]
existingNames [Name Any] -> [Name Any] -> [Name Any]
forall a. [a] -> [a] -> [a]
++ [Name Any]
forall a. [Name a]
newNames)
                                       [Either Term Type]
args
              -- Determine name the resulting specialized function, and the
              -- form of the specialized-on argument
              (Name Term
newName, InlineSpec
inl', Either Term Type
specArg') <- case Either Term Type
specArg of
                Left a :: Term
a@(Term -> (Term, [Either Term Type], [TickInfo])
collectArgsTicks -> (Var Id
g,[Either Term Type]
gArgs,[TickInfo]
_gTicks)) -> if TyConMap -> Term -> Bool
isPolyFun TyConMap
tcm Term
a
                    then do
                      -- In case we are specializing on an argument that is a
                      -- global function then we use that function's name as part
                      -- of the name of the specialized higher-order function.
                      -- Additionally, we will return the body of the global
                      -- function, instead of a variable reference to the
                      -- global function.
                      --
                      -- This will turn things like @mealy g@ into a new
                      -- binding where both the body of @mealy@ and @g@
                      -- are inlined, meaning the state-transition-function
                      -- and the memory element will be in a single function.
                      --
                      -- The name of the new function depends on NOINLINE / OPAQUE
                      -- annotations.
                      --
                      --    OPAQUE    | Name for @f g@
                      --   -----------|---------------
                      --    f and g   | f_g
                      --    f         | f
                      --    g         | g
                      --    !f and !g | f_g
                      --
                      Maybe (Binding Term)
gTmM <- (BindingMap -> Maybe (Binding Term))
-> RewriteMonad NormalizeState BindingMap
-> RewriteMonad NormalizeState (Maybe (Binding Term))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (Id -> BindingMap -> Maybe (Binding Term)
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup Id
g) (RewriteMonad NormalizeState BindingMap
 -> RewriteMonad NormalizeState (Maybe (Binding Term)))
-> RewriteMonad NormalizeState BindingMap
-> RewriteMonad NormalizeState (Maybe (Binding Term))
forall a b. (a -> b) -> a -> b
$ Getting BindingMap (RewriteState NormalizeState) BindingMap
-> RewriteMonad NormalizeState BindingMap
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting BindingMap (RewriteState NormalizeState) BindingMap
forall extra. Lens' (RewriteState extra) BindingMap
bindings
                      (Name Term, InlineSpec, Either Term Type)
-> RewriteMonad
     NormalizeState (Name Term, InlineSpec, Either Term Type)
forall (m :: Type -> Type) a. Monad m => a -> m a
return
                        ( (InlineSpec, Name Term)
-> (Maybe InlineSpec, Name Term) -> Name Term
specializeName (InlineSpec
inl, Id -> Name Term
forall a. Var a -> Name a
varName Id
f) (Binding Term -> InlineSpec
forall a. Binding a -> InlineSpec
bindingSpec (Binding Term -> InlineSpec)
-> Maybe (Binding Term) -> Maybe InlineSpec
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Binding Term)
gTmM, Id -> Name Term
forall a. Var a -> Name a
varName Id
g)
                        , InlineSpec -> InlineSpec -> InlineSpec
preferNoInline InlineSpec
inl (InlineSpec
-> (Binding Term -> InlineSpec)
-> Maybe (Binding Term)
-> InlineSpec
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InlineSpec
noUserInline Binding Term -> InlineSpec
forall a. Binding a -> InlineSpec
bindingSpec Maybe (Binding Term)
gTmM)
                        , Either Term Type
-> (Binding Term -> Either Term Type)
-> Maybe (Binding Term)
-> Either Term Type
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Either Term Type
specArg (Term -> Either Term Type
forall a b. a -> Either a b
Left (Term -> Either Term Type)
-> (Binding Term -> Term) -> Binding Term -> Either Term Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term -> [Either Term Type] -> Term
`mkApps` [Either Term Type]
gArgs) (Term -> Term) -> (Binding Term -> Term) -> Binding Term -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Binding Term -> Term
forall a. Binding a -> a
bindingTerm) Maybe (Binding Term)
gTmM
                        )
                    else (Name Term, InlineSpec, Either Term Type)
-> RewriteMonad
     NormalizeState (Name Term, InlineSpec, Either Term Type)
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Id -> Name Term
forall a. Var a -> Name a
varName Id
f, InlineSpec
inl, Either Term Type
specArg)
                Either Term Type
_ -> (Name Term, InlineSpec, Either Term Type)
-> RewriteMonad
     NormalizeState (Name Term, InlineSpec, Either Term Type)
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Id -> Name Term
forall a. Var a -> Name a
varName Id
f, InlineSpec
inl, Either Term Type
specArg)
              -- Create specialized functions
              let newBody :: Term
newBody = Term -> [Either Id TyVar] -> Term
mkAbstraction (Term -> [Either Term Type] -> Term
mkApps Term
bodyTm ([Either Term Type]
argVars [Either Term Type] -> [Either Term Type] -> [Either Term Type]
forall a. [a] -> [a] -> [a]
++ [Either Term Type
specArg'])) ([Either Id TyVar]
boundArgs [Either Id TyVar] -> [Either Id TyVar] -> [Either Id TyVar]
forall a. [a] -> [a] -> [a]
++ [Either Id TyVar]
specBndrs)
              Id
newf <- Name Term
-> SrcSpan -> InlineSpec -> Term -> RewriteMonad NormalizeState Id
forall extra.
Name Term -> SrcSpan -> InlineSpec -> Term -> RewriteMonad extra Id
mkFunction Name Term
newName SrcSpan
sp InlineSpec
inl' Term
newBody
              -- Remember specialization
              ((NormalizeState -> Identity NormalizeState)
-> RewriteState NormalizeState
-> Identity (RewriteState NormalizeState)
forall extra extra2.
Lens (RewriteState extra) (RewriteState extra2) extra extra2
extra((NormalizeState -> Identity NormalizeState)
 -> RewriteState NormalizeState
 -> Identity (RewriteState NormalizeState))
-> ((UniqMap Int -> Identity (UniqMap Int))
    -> NormalizeState -> Identity NormalizeState)
-> (UniqMap Int -> Identity (UniqMap Int))
-> RewriteState NormalizeState
-> Identity (RewriteState NormalizeState)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(UniqMap Int -> Identity (UniqMap Int))
-> NormalizeState -> Identity NormalizeState
Lens' NormalizeState (UniqMap Int)
specialisationHistory) ((UniqMap Int -> Identity (UniqMap Int))
 -> RewriteState NormalizeState
 -> Identity (RewriteState NormalizeState))
-> (UniqMap Int -> UniqMap Int) -> RewriteMonad NormalizeState ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
%= (Int -> Int -> Int) -> Id -> Int -> UniqMap Int -> UniqMap Int
forall a b.
Uniquable a =>
(b -> b -> b) -> a -> b -> UniqMap b -> UniqMap b
UniqMap.insertWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Id
f Int
1
              ((NormalizeState -> Identity NormalizeState)
-> RewriteState NormalizeState
-> Identity (RewriteState NormalizeState)
forall extra extra2.
Lens (RewriteState extra) (RewriteState extra2) extra extra2
extra((NormalizeState -> Identity NormalizeState)
 -> RewriteState NormalizeState
 -> Identity (RewriteState NormalizeState))
-> ((Map (Id, Int, Either Term Type) Id
     -> Identity (Map (Id, Int, Either Term Type) Id))
    -> NormalizeState -> Identity NormalizeState)
-> (Map (Id, Int, Either Term Type) Id
    -> Identity (Map (Id, Int, Either Term Type) Id))
-> RewriteState NormalizeState
-> Identity (RewriteState NormalizeState)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Map (Id, Int, Either Term Type) Id
 -> Identity (Map (Id, Int, Either Term Type) Id))
-> NormalizeState -> Identity NormalizeState
Lens' NormalizeState (Map (Id, Int, Either Term Type) Id)
specialisationCache)  ((Map (Id, Int, Either Term Type) Id
  -> Identity (Map (Id, Int, Either Term Type) Id))
 -> RewriteState NormalizeState
 -> Identity (RewriteState NormalizeState))
-> (Map (Id, Int, Either Term Type) Id
    -> Map (Id, Int, Either Term Type) Id)
-> RewriteMonad NormalizeState ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
%= (Id, Int, Either Term Type)
-> Id
-> Map (Id, Int, Either Term Type) Id
-> Map (Id, Int, Either Term Type) Id
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (Id
f,Int
argLen,Either Term Type
specAbs) Id
newf
              -- use specialized function
              let newExpr :: Term
newExpr = Term -> [Either Term Type] -> Term
mkApps (Term -> [TickInfo] -> Term
mkTicks (Id -> Term
Var Id
newf) [TickInfo]
ticks) ([Either Term Type]
args [Either Term Type] -> [Either Term Type] -> [Either Term Type]
forall a. [a] -> [a] -> [a]
++ [Either Term Type]
specVars)
              Id
newf Id
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState Term
forall a b. NFData a => a -> b -> b
`deepseq` Term -> RewriteMonad NormalizeState Term
forall a extra. a -> RewriteMonad extra a
changed Term
newExpr
        Maybe (Binding Term)
Nothing -> Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e
  where
    noUserInline :: InlineSpec
#if MIN_VERSION_ghc(9,2,0)
    noUserInline = NoUserInlinePrag
#else
    noUserInline :: InlineSpec
noUserInline = InlineSpec
NoUserInline
#endif

    specializeName :: (InlineSpec, Name Term) -> (Maybe InlineSpec, Name Term) -> Name Term
    specializeName :: (InlineSpec, Name Term)
-> (Maybe InlineSpec, Name Term) -> Name Term
specializeName (InlineSpec
spec0, Name Term
n0) (Maybe InlineSpec
spec1, Name Term
n1)
      | (Bool
ann0 Bool -> Bool -> Bool
&& Bool
ann1) Bool -> Bool -> Bool
|| (Bool -> Bool
not Bool
ann0 Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
ann1)
      = Name Term
n0{nameOcc :: OccName
nameOcc=Name Term -> OccName
forall a. Name a -> OccName
nameOcc Name Term
n0 OccName -> OccName -> OccName
forall a. Semigroup a => a -> a -> a
<> OccName
"_" OccName -> OccName -> OccName
forall a. Semigroup a => a -> a -> a
<> OccName -> OccName
stripMods (Name Term -> OccName
forall a. Name a -> OccName
nameOcc Name Term
n1)}
      | Bool
ann0 = Name Term
n0
      | Bool
otherwise = Name Term
n1
     where
      ann0 :: Bool
ann0 = InlineSpec -> Bool
isNoInline InlineSpec
spec0
      ann1 :: Bool
ann1 = (InlineSpec -> Bool) -> Maybe InlineSpec -> Maybe Bool
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap InlineSpec -> Bool
isNoInline Maybe InlineSpec
spec1 Maybe Bool -> Maybe Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
True

    stripMods :: Text.Text -> Text.Text
    stripMods :: OccName -> OccName
stripMods = (Char -> Bool) -> OccName -> OccName
Text.takeWhileEnd (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'.')

    collectBndrsMinusApps :: Term -> [Name a]
    collectBndrsMinusApps :: Term -> [Name a]
collectBndrsMinusApps = [Name a] -> [Name a]
forall a. [a] -> [a]
reverse ([Name a] -> [Name a]) -> (Term -> [Name a]) -> Term -> [Name a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Name a] -> Term -> [Name a]
forall a.
(Coercible a (Name Term), Coercible a (Name Type)) =>
[a] -> Term -> [a]
go []
      where
        go :: [a] -> Term -> [a]
go [a]
bs (Lam Id
v Term
e')    = [a] -> Term -> [a]
go (Name Term -> a
coerce (Id -> Name Term
forall a. Var a -> Name a
varName Id
v)a -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
bs)  Term
e'
        go [a]
bs (TyLam TyVar
tv Term
e') = [a] -> Term -> [a]
go (Name Type -> a
coerce (TyVar -> Name Type
forall a. Var a -> Name a
varName TyVar
tv)a -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
bs) Term
e'
        go [a]
bs (App Term
e' Term
_) = case [a] -> Term -> [a]
go [] Term
e' of
          []  -> [a]
bs
          [a]
bs' -> [a] -> [a]
forall a. [a] -> [a]
init [a]
bs' [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
bs
        go [a]
bs (TyApp Term
e' Type
_) = case [a] -> Term -> [a]
go [] Term
e' of
          []  -> [a]
bs
          [a]
bs' -> [a] -> [a]
forall a. [a] -> [a]
init [a]
bs' [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
bs
        go [a]
bs Term
_ = [a]
bs

-- Specializing non Var's is used by nonRepANF
specialize' TransformContext
_ctx Term
_ (Term
appE,[Either Term Type]
args,[TickInfo]
ticks) (Left Term
specArg) = do
  -- Create binders and variable references for free variables in 'specArg'
  let ([Either Id TyVar]
specBndrs,[Either Term Type]
specVars) = Either Term Type -> ([Either Id TyVar], [Either Term Type])
specArgBndrsAndVars (Term -> Either Term Type
forall a b. a -> Either a b
Left Term
specArg)
  -- Create specialized function
      newBody :: Term
newBody = Term -> [Either Id TyVar] -> Term
mkAbstraction Term
specArg [Either Id TyVar]
specBndrs
  -- See if there's an existing binder that's alpha-equivalent to the
  -- specialized function
  BindingMap
existing <- (Binding Term -> Bool) -> BindingMap -> BindingMap
forall b. (b -> Bool) -> UniqMap b -> UniqMap b
UniqMap.filter ((Term -> Term -> Bool
`aeqTerm` Term
newBody) (Term -> Bool) -> (Binding Term -> Term) -> Binding Term -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Binding Term -> Term
forall a. Binding a -> a
bindingTerm) (BindingMap -> BindingMap)
-> RewriteMonad NormalizeState BindingMap
-> RewriteMonad NormalizeState BindingMap
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Getting BindingMap (RewriteState NormalizeState) BindingMap
-> RewriteMonad NormalizeState BindingMap
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting BindingMap (RewriteState NormalizeState) BindingMap
forall extra. Lens' (RewriteState extra) BindingMap
bindings
  -- Create a new function if an alpha-equivalent binder doesn't exist
  Id
newf <- case BindingMap -> [Binding Term]
forall b. UniqMap b -> [b]
UniqMap.elems BindingMap
existing of
    [] -> do (Id
cf,SrcSpan
sp) <- Getting (Id, SrcSpan) (RewriteState NormalizeState) (Id, SrcSpan)
-> RewriteMonad NormalizeState (Id, SrcSpan)
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting (Id, SrcSpan) (RewriteState NormalizeState) (Id, SrcSpan)
forall extra. Lens' (RewriteState extra) (Id, SrcSpan)
curFun
#if MIN_VERSION_ghc(9,2,0)
             mkFunction (appendToName (varName cf) "_specF") sp NoUserInlinePrag newBody
#else
             Name Term
-> SrcSpan -> InlineSpec -> Term -> RewriteMonad NormalizeState Id
forall extra.
Name Term -> SrcSpan -> InlineSpec -> Term -> RewriteMonad extra Id
mkFunction (Name Term -> OccName -> Name Term
forall a. Name a -> OccName -> Name a
appendToName (Id -> Name Term
forall a. Var a -> Name a
varName Id
cf) OccName
"_specF") SrcSpan
sp InlineSpec
NoUserInline Term
newBody
#endif
    (Binding Term
b:[Binding Term]
_) -> Id -> RewriteMonad NormalizeState Id
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Binding Term -> Id
forall a. Binding a -> Id
bindingId Binding Term
b)
  -- Create specialized argument
  let newArg :: Either Term b
newArg  = Term -> Either Term b
forall a b. a -> Either a b
Left (Term -> Either Term b) -> Term -> Either Term b
forall a b. (a -> b) -> a -> b
$ Term -> [Either Term Type] -> Term
mkApps (Id -> Term
Var Id
newf) [Either Term Type]
specVars
  -- Use specialized argument
  let newExpr :: Term
newExpr = Term -> [Either Term Type] -> Term
mkApps (Term -> [TickInfo] -> Term
mkTicks Term
appE [TickInfo]
ticks) ([Either Term Type]
args [Either Term Type] -> [Either Term Type] -> [Either Term Type]
forall a. [a] -> [a] -> [a]
++ [Either Term Type
forall b. Either Term b
newArg])
  Term -> RewriteMonad NormalizeState Term
forall a extra. a -> RewriteMonad extra a
changed Term
newExpr

specialize' TransformContext
_ Term
e (Term, [Either Term Type], [TickInfo])
_ Either Term Type
_ = Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e

-- Note [Collect free-variables in an insertion-ordered set]
--
-- In order for the specialization cache to work, 'specArgBndrsAndVars' should
-- yield (alpha equivalent) results for the same specialization. While collecting
-- free variables in a given term or type it should therefore keep a stable
-- ordering based on the order in which it finds free vars. To see why,
-- consider the following two pseudo-code calls to 'specialize':
--
--     specialize {f ('a', x[123], y[456])}
--     specialize {f ('b', x[456], y[123])}
--
-- Collecting the binders in a VarSet would yield the following (unique ordered)
-- sets:
--
--     {x[123], y[456]}
--     {y[123], x[456]}
--
-- ..and therefore breaking specializing caching. We now track them in insert-
-- ordered sets, yielding:
--
--     {x[123], y[456]}
--     {x[456], y[123]}
--

-- | Create binders and variable references for free variables in 'specArg'
specArgBndrsAndVars
  :: Either Term Type
  -> ([Either Id TyVar], [Either Term Type])
specArgBndrsAndVars :: Either Term Type -> ([Either Id TyVar], [Either Term Type])
specArgBndrsAndVars Either Term Type
specArg =
  -- See Note [Collect free-variables in an insertion-ordered set]
  let unitFV :: Var a -> Const (OSet.OLSet TyVar, OSet.OLSet Id) (Var a)
      unitFV :: Var a -> Const (OLSet TyVar, OLSet Id) (Var a)
unitFV v :: Var a
v@(Id {}) = (OLSet TyVar, OLSet Id) -> Const (OLSet TyVar, OLSet Id) (Var a)
forall k a (b :: k). a -> Const a b
Const (OLSet TyVar
forall a. Monoid a => a
mempty, OSet Id -> OLSet Id
coerce (Id -> OSet Id
forall a. a -> OSet a
OSet.singleton (Var a -> Id
coerce Var a
v)))
      unitFV v :: Var a
v@(TyVar {}) = (OLSet TyVar, OLSet Id) -> Const (OLSet TyVar, OLSet Id) (Var a)
forall k a (b :: k). a -> Const a b
Const (OSet TyVar -> OLSet TyVar
coerce (TyVar -> OSet TyVar
forall a. a -> OSet a
OSet.singleton (Var a -> TyVar
coerce Var a
v)), OLSet Id
forall a. Monoid a => a
mempty)

      ([TyVar]
specFTVs,[Id]
specFVs) = case Either Term Type
specArg of
        Left Term
tm  -> (OLSet TyVar -> [TyVar]
forall a. OLSet a -> [a]
OSet.toListL (OLSet TyVar -> [TyVar])
-> (OLSet Id -> [Id]) -> (OLSet TyVar, OLSet Id) -> ([TyVar], [Id])
forall (a :: Type -> Type -> Type) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** OLSet Id -> [Id]
forall a. OLSet a -> [a]
OSet.toListL) ((OLSet TyVar, OLSet Id) -> ([TyVar], [Id]))
-> (Const (OLSet TyVar, OLSet Id) (Var Any)
    -> (OLSet TyVar, OLSet Id))
-> Const (OLSet TyVar, OLSet Id) (Var Any)
-> ([TyVar], [Id])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Const (OLSet TyVar, OLSet Id) (Var Any) -> (OLSet TyVar, OLSet Id)
forall a k (b :: k). Const a b -> a
getConst (Const (OLSet TyVar, OLSet Id) (Var Any) -> ([TyVar], [Id]))
-> Const (OLSet TyVar, OLSet Id) (Var Any) -> ([TyVar], [Id])
forall a b. (a -> b) -> a -> b
$
                    Getting (Const (OLSet TyVar, OLSet Id) (Var Any)) Term (Var Any)
-> (Var Any -> Const (OLSet TyVar, OLSet Id) (Var Any))
-> Term
-> Const (OLSet TyVar, OLSet Id) (Var Any)
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting (Const (OLSet TyVar, OLSet Id) (Var Any)) Term (Var Any)
forall a. Fold Term (Var a)
freeLocalVars Var Any -> Const (OLSet TyVar, OLSet Id) (Var Any)
forall a. Var a -> Const (OLSet TyVar, OLSet Id) (Var a)
unitFV Term
tm
        Right Type
ty -> ( UniqMap TyVar -> [TyVar]
forall b. UniqMap b -> [b]
UniqMap.elems
                        (Getting (UniqMap TyVar) Type TyVar
-> (TyVar -> UniqMap TyVar) -> Type -> UniqMap TyVar
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting (UniqMap TyVar) Type TyVar
Fold Type TyVar
typeFreeVars (\TyVar
x -> TyVar -> UniqMap TyVar
forall a. Uniquable a => a -> UniqMap a
UniqMap.singletonUnique (TyVar -> TyVar
coerce TyVar
x)) Type
ty)
                    , [] :: [Id])

      specTyBndrs :: [Either a TyVar]
specTyBndrs = (TyVar -> Either a TyVar) -> [TyVar] -> [Either a TyVar]
forall a b. (a -> b) -> [a] -> [b]
map TyVar -> Either a TyVar
forall a b. b -> Either a b
Right [TyVar]
specFTVs
      specTmBndrs :: [Either Id b]
specTmBndrs = (Id -> Either Id b) -> [Id] -> [Either Id b]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Either Id b
forall a b. a -> Either a b
Left  [Id]
specFVs

      specTyVars :: [Either a Type]
specTyVars  = (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]
specFTVs
      specTmVars :: [Either Term b]
specTmVars  = (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]
specFVs

  in  ([Either Id TyVar]
forall a. [Either a TyVar]
specTyBndrs [Either Id TyVar] -> [Either Id TyVar] -> [Either Id TyVar]
forall a. [a] -> [a] -> [a]
++ [Either Id TyVar]
forall b. [Either Id b]
specTmBndrs,[Either Term Type]
forall a. [Either a Type]
specTyVars [Either Term Type] -> [Either Term Type] -> [Either Term Type]
forall a. [a] -> [a] -> [a]
++ [Either Term Type]
forall b. [Either Term b]
specTmVars)

-- | Specialize functions on their non-representable argument
nonRepSpec :: HasCallStack => NormRewrite
nonRepSpec :: NormRewrite
nonRepSpec TransformContext
ctx e :: Term
e@(App Term
e1 Term
e2)
  | (Var {}, [Either Term Type]
args) <- Term -> (Term, [Either Term Type])
collectArgs Term
e1
  , ([Term]
_, [])     <- [Either Term Type] -> ([Term], [Type])
forall a b. [Either a b] -> ([a], [b])
Either.partitionEithers [Either Term Type]
args
  , [TyVar] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
null ([TyVar] -> Bool) -> [TyVar] -> Bool
forall a b. (a -> b) -> a -> b
$ Getting (Endo [TyVar]) Term TyVar -> Term -> [TyVar]
forall a s. Getting (Endo [a]) s a -> s -> [a]
Lens.toListOf Getting (Endo [TyVar]) Term TyVar
Fold Term TyVar
termFreeTyVars Term
e2
  = 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
       let e2Ty :: Type
e2Ty = TyConMap -> Term -> Type
forall a. InferType a => TyConMap -> a -> Type
inferCoreTypeOf TyConMap
tcm Term
e2
       let localVar :: Bool
localVar = Term -> Bool
isLocalVar Term
e2
       Bool
nonRepE2 <- Bool -> Bool
not (Bool -> Bool)
-> RewriteMonad NormalizeState Bool
-> RewriteMonad NormalizeState Bool
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ((CustomReprs
 -> TyConMap
 -> Type
 -> State HWMap (Maybe (Either String FilteredHWType)))
-> CustomReprs -> Bool -> TyConMap -> Type -> Bool
representableType ((CustomReprs
  -> TyConMap
  -> Type
  -> State HWMap (Maybe (Either String FilteredHWType)))
 -> CustomReprs -> Bool -> TyConMap -> Type -> Bool)
-> RewriteMonad
     NormalizeState
     (CustomReprs
      -> TyConMap
      -> Type
      -> State HWMap (Maybe (Either String FilteredHWType)))
-> RewriteMonad
     NormalizeState (CustomReprs -> Bool -> TyConMap -> Type -> Bool)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Getting
  (CustomReprs
   -> TyConMap
   -> Type
   -> State HWMap (Maybe (Either String FilteredHWType)))
  RewriteEnv
  (CustomReprs
   -> TyConMap
   -> Type
   -> State HWMap (Maybe (Either String FilteredHWType)))
-> RewriteMonad
     NormalizeState
     (CustomReprs
      -> TyConMap
      -> Type
      -> State HWMap (Maybe (Either String FilteredHWType)))
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting
  (CustomReprs
   -> TyConMap
   -> Type
   -> State HWMap (Maybe (Either String FilteredHWType)))
  RewriteEnv
  (CustomReprs
   -> TyConMap
   -> Type
   -> State HWMap (Maybe (Either String FilteredHWType)))
Lens'
  RewriteEnv
  (CustomReprs
   -> TyConMap
   -> Type
   -> State HWMap (Maybe (Either String FilteredHWType)))
typeTranslator
                                              RewriteMonad
  NormalizeState (CustomReprs -> Bool -> TyConMap -> Type -> Bool)
-> RewriteMonad NormalizeState CustomReprs
-> RewriteMonad NormalizeState (Bool -> TyConMap -> Type -> Bool)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Getting CustomReprs RewriteEnv CustomReprs
-> RewriteMonad NormalizeState CustomReprs
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting CustomReprs RewriteEnv CustomReprs
Getter RewriteEnv CustomReprs
customReprs
                                              RewriteMonad NormalizeState (Bool -> TyConMap -> Type -> Bool)
-> RewriteMonad NormalizeState Bool
-> RewriteMonad NormalizeState (TyConMap -> Type -> Bool)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Bool -> RewriteMonad NormalizeState Bool
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Bool
False
                                              RewriteMonad NormalizeState (TyConMap -> Type -> Bool)
-> RewriteMonad NormalizeState TyConMap
-> RewriteMonad NormalizeState (Type -> Bool)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> 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
                                              RewriteMonad NormalizeState (Type -> Bool)
-> RewriteMonad NormalizeState Type
-> RewriteMonad NormalizeState Bool
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Type -> RewriteMonad NormalizeState Type
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Type
e2Ty)
       if Bool
nonRepE2 Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
localVar
         then do
           Term
e2' <- Term -> RewriteMonad NormalizeState Term
inlineInternalSpecialisationArgument Term
e2
           NormRewrite
specialize TransformContext
ctx (Term -> Term -> Term
App Term
e1 Term
e2')
         else Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e
  where
    -- | If the argument on which we're specializing is an internal function,
    -- one created by the compiler, then inline that function before we
    -- specialize.
    --
    -- We need to do this because otherwise the specialization history won't
    -- recognize the new specialization argument as something the function has
    -- already been specialized on
    inlineInternalSpecialisationArgument
      :: Term
      -> NormalizeSession Term
    inlineInternalSpecialisationArgument :: Term -> RewriteMonad NormalizeState Term
inlineInternalSpecialisationArgument Term
app
      | (Var Id
f,[Either Term Type]
fArgs,[TickInfo]
ticks) <- Term -> (Term, [Either Term Type], [TickInfo])
collectArgsTicks Term
app
      = do
        Maybe (Binding Term)
fTmM <- Id -> BindingMap -> Maybe (Binding Term)
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv Id
f (BindingMap -> Maybe (Binding Term))
-> RewriteMonad NormalizeState BindingMap
-> RewriteMonad NormalizeState (Maybe (Binding Term))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Getting BindingMap (RewriteState NormalizeState) BindingMap
-> RewriteMonad NormalizeState BindingMap
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting BindingMap (RewriteState NormalizeState) BindingMap
forall extra. Lens' (RewriteState extra) BindingMap
bindings
        case Maybe (Binding Term)
fTmM of
          Just Binding Term
b
            | Name Term -> NameSort
forall a. Name a -> NameSort
nameSort (Id -> Name Term
forall a. Var a -> Name a
varName (Binding Term -> Id
forall a. Binding a -> Id
bindingId Binding Term
b)) NameSort -> NameSort -> Bool
forall a. Eq a => a -> a -> Bool
== NameSort
Internal
            -> (Any -> Any)
-> RewriteMonad NormalizeState Term
-> RewriteMonad NormalizeState Term
forall extra a.
(Any -> Any) -> RewriteMonad extra a -> RewriteMonad extra a
censor (Any -> Any -> Any
forall a b. a -> b -> a
const Any
forall a. Monoid a => a
mempty)
                      (NormRewrite -> NormRewrite
forall m. Rewrite m -> Rewrite m
topdownR HasCallStack => NormRewrite
NormRewrite
appProp TransformContext
ctx
                        (Term -> [Either Term Type] -> Term
mkApps (Term -> [TickInfo] -> Term
mkTicks (Binding Term -> Term
forall a. Binding a -> a
bindingTerm Binding Term
b) [TickInfo]
ticks) [Either Term Type]
fArgs))
          Maybe (Binding Term)
_ -> Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
app
      | Bool
otherwise = Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
app

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

-- | Specialize functions on their type
typeSpec :: HasCallStack => NormRewrite
typeSpec :: NormRewrite
typeSpec TransformContext
ctx e :: Term
e@(TyApp Term
e1 Type
ty)
  | (Var {},  [Either Term Type]
args) <- Term -> (Term, [Either Term Type])
collectArgs Term
e1
  , [TyVar] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
null ([TyVar] -> Bool) -> [TyVar] -> Bool
forall a b. (a -> b) -> a -> b
$ Getting (Endo [TyVar]) Type TyVar -> Type -> [TyVar]
forall a s. Getting (Endo [a]) s a -> s -> [a]
Lens.toListOf Getting (Endo [TyVar]) Type TyVar
Fold Type TyVar
typeFreeVars Type
ty
  , ([Term]
_, []) <- [Either Term Type] -> ([Term], [Type])
forall a b. [Either a b] -> ([a], [b])
Either.partitionEithers [Either Term Type]
args
  = NormRewrite
specialize TransformContext
ctx Term
e

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

-- | Specialize functions on arguments which are zero-width. These arguments
-- can have only one possible value, and specializing on this value may create
-- additional opportunities for transformations to fire.
--
-- As we can't remove zero-width arguements (as transformations cannot change
-- the type of a term), we instead substitute all occurances of a lambda-bound
-- variable with a zero-width type with the only value of that type.
--
zeroWidthSpec :: HasCallStack => NormRewrite
zeroWidthSpec :: NormRewrite
zeroWidthSpec (TransformContext InScopeSet
is Context
_) e :: Term
e@(Lam Id
i Term
x0) = 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
  let bndrTy :: Type
bndrTy = TyConMap -> Type -> Type
normalizeType TyConMap
tcm (Id -> Type
forall a. HasType a => a -> Type
coreTypeOf Id
i)

  case TyConMap -> Type -> Maybe Term
zeroWidthTypeElem TyConMap
tcm Type
bndrTy of
    Just Term
tm ->
      let subst :: Subst
subst = Subst -> Id -> Term -> Subst
extendIdSubst (InScopeSet -> Subst
mkSubst InScopeSet
is) Id
i Term
tm
          x1 :: Term
x1 = HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
"zeroWidthSpec" Subst
subst Term
x0
       in Term -> RewriteMonad NormalizeState Term
forall a extra. a -> RewriteMonad extra a
changed (Id -> Term -> Term
Lam Id
i Term
x1)

    Maybe Term
Nothing ->
      Term -> RewriteMonad NormalizeState Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
e

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

-- Get the only element of a type, if it is zero-width.
--
zeroWidthTypeElem :: TyConMap -> Type -> Maybe Term
zeroWidthTypeElem :: TyConMap -> Type -> Maybe Term
zeroWidthTypeElem TyConMap
tcm Type
ty = do
  (TyConName
tcNm, [Type]
args) <- Type -> Maybe (TyConName, [Type])
splitTyConAppM Type
ty

  if | TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
tcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== Name -> OccName
forall a. Show a => a -> OccName
Text.showt ''BV.BitVector
     , [LitTy (NumTy Integer
0)] <- [Type]
args
     -> Term -> Maybe Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return (TyConName -> [Type] -> Term
bitVectorZW TyConName
tcNm [Type]
args)

     | TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
tcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== Name -> OccName
forall a. Show a => a -> OccName
Text.showt ''I.Index
     , [LitTy (NumTy Integer
1)] <- [Type]
args
     -> Term -> Maybe Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return (TyConName -> [Type] -> Term
indexZW TyConName
tcNm [Type]
args)

     | TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
tcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== Name -> OccName
forall a. Show a => a -> OccName
Text.showt ''S.Signed
     , [LitTy (NumTy Integer
0)] <- [Type]
args
     -> Term -> Maybe Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return (TyConName -> [Type] -> Term
signedZW TyConName
tcNm [Type]
args)

     | TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
tcNm OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== Name -> OccName
forall a. Show a => a -> OccName
Text.showt ''U.Unsigned
     , [LitTy (NumTy Integer
0)] <- [Type]
args
     -> Term -> Maybe Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return (TyConName -> [Type] -> Term
unsignedZW TyConName
tcNm [Type]
args)

     -- Any other zero-width type should only have a single data constructor
     -- where all fields are also zero-width.
     | Bool
otherwise
     -> do
       TyCon
tc <- TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup TyConName
tcNm TyConMap
tcm

       case TyCon -> [DataCon]
tyConDataCons TyCon
tc of
         [DataCon
dc] -> do
           [Term]
zwArgs <- (Type -> Maybe Term) -> [Type] -> Maybe [Term]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (TyConMap -> Type -> Maybe Term
zeroWidthTypeElem TyConMap
tcm) (DataCon -> [Type]
dcArgTys DataCon
dc)
           Term -> Maybe Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Term -> [Term] -> Term
mkTmApps (DataCon -> Term
Data DataCon
dc) [Term]
zwArgs)

         [DataCon]
_ ->
           Maybe Term
forall a. Maybe a
Nothing
 where
  nNm :: Name a
nNm = OccName -> Int -> Name a
forall a. OccName -> Int -> Name a
mkUnsafeSystemName OccName
"n" Int
0
  nTv :: TyVar
nTv = Type -> Name Type -> TyVar
mkTyVar Type
typeNatKind Name Type
forall a. Name a
nNm

  mkBitVector :: TyConName -> PrimInfo
mkBitVector TyConName
tcNm =
    let prTy :: Type
prTy = Type -> [Either TyVar Type] -> Type
mkPolyFunTy (TyConName -> [Type] -> Type
mkTyConApp TyConName
tcNm [TyVar -> Type
VarTy TyVar
nTv])
                 [TyVar -> Either TyVar Type
forall a b. a -> Either a b
Left TyVar
nTv, Type -> Either TyVar Type
forall a b. b -> Either a b
Right Type
naturalPrimTy, Type -> Either TyVar Type
forall a b. b -> Either a b
Right Type
naturalPrimTy, Type -> Either TyVar Type
forall a b. b -> Either a b
Right Type
integerPrimTy]
     in OccName
-> Type -> WorkInfo -> IsMultiPrim -> PrimUnfolding -> PrimInfo
PrimInfo (Name -> OccName
forall a. Show a => a -> OccName
Text.showt 'BV.fromInteger#) Type
prTy WorkInfo
WorkNever IsMultiPrim
SingleResult PrimUnfolding
NoUnfolding

  bitVectorZW :: TyConName -> [Type] -> Term
bitVectorZW TyConName
tcNm [Type]
tyArgs =
    let pr :: PrimInfo
pr = TyConName -> PrimInfo
mkBitVector TyConName
tcNm
     in Term -> [Either Term Type] -> Term
mkApps (PrimInfo -> Term
Prim PrimInfo
pr) ([Either Term Type] -> Term) -> [Either Term Type] -> Term
forall a b. (a -> b) -> a -> b
$ (Type -> Either Term Type) -> [Type] -> [Either Term Type]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Either Term Type
forall a b. b -> Either a b
Right [Type]
tyArgs [Either Term Type] -> [Either Term Type] -> [Either Term Type]
forall a. Semigroup a => a -> a -> a
<>
          [ Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
NaturalLiteral Integer
0))
          , Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
NaturalLiteral Integer
0))
          , Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
IntegerLiteral Integer
0))
          ]

  mkSizedNum :: TyConName -> OccName -> PrimInfo
mkSizedNum TyConName
tcNm OccName
n =
    let prTy :: Type
prTy = Type -> [Either TyVar Type] -> Type
mkPolyFunTy (TyConName -> [Type] -> Type
mkTyConApp TyConName
tcNm [TyVar -> Type
VarTy TyVar
nTv])
                 [TyVar -> Either TyVar Type
forall a b. a -> Either a b
Left TyVar
nTv, Type -> Either TyVar Type
forall a b. b -> Either a b
Right Type
naturalPrimTy, Type -> Either TyVar Type
forall a b. b -> Either a b
Right Type
integerPrimTy]
     in OccName
-> Type -> WorkInfo -> IsMultiPrim -> PrimUnfolding -> PrimInfo
PrimInfo OccName
n Type
prTy WorkInfo
WorkNever IsMultiPrim
SingleResult PrimUnfolding
NoUnfolding

  indexZW :: TyConName -> [Type] -> Term
indexZW TyConName
tcNm [Type]
tyArgs =
    let pr :: PrimInfo
pr = TyConName -> OccName -> PrimInfo
mkSizedNum TyConName
tcNm (Name -> OccName
forall a. Show a => a -> OccName
Text.showt 'I.fromInteger#)
     in Term -> [Either Term Type] -> Term
mkApps (PrimInfo -> Term
Prim PrimInfo
pr) ([Either Term Type] -> Term) -> [Either Term Type] -> Term
forall a b. (a -> b) -> a -> b
$ (Type -> Either Term Type) -> [Type] -> [Either Term Type]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Either Term Type
forall a b. b -> Either a b
Right [Type]
tyArgs [Either Term Type] -> [Either Term Type] -> [Either Term Type]
forall a. Semigroup a => a -> a -> a
<>
          [ Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
NaturalLiteral Integer
1))
          , Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
IntegerLiteral Integer
0))
          ]

  signedZW :: TyConName -> [Type] -> Term
signedZW TyConName
tcNm [Type]
tyArgs =
    let pr :: PrimInfo
pr = TyConName -> OccName -> PrimInfo
mkSizedNum TyConName
tcNm (Name -> OccName
forall a. Show a => a -> OccName
Text.showt 'S.fromInteger#)
     in Term -> [Either Term Type] -> Term
mkApps (PrimInfo -> Term
Prim PrimInfo
pr) ([Either Term Type] -> Term) -> [Either Term Type] -> Term
forall a b. (a -> b) -> a -> b
$ (Type -> Either Term Type) -> [Type] -> [Either Term Type]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Either Term Type
forall a b. b -> Either a b
Right [Type]
tyArgs [Either Term Type] -> [Either Term Type] -> [Either Term Type]
forall a. Semigroup a => a -> a -> a
<>
          [ Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
NaturalLiteral Integer
0))
          , Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
IntegerLiteral Integer
0))
          ]

  unsignedZW :: TyConName -> [Type] -> Term
unsignedZW TyConName
tcNm [Type]
tyArgs =
    let pr :: PrimInfo
pr = TyConName -> OccName -> PrimInfo
mkSizedNum TyConName
tcNm (Name -> OccName
forall a. Show a => a -> OccName
Text.showt 'U.fromInteger#)
     in Term -> [Either Term Type] -> Term
mkApps (PrimInfo -> Term
Prim PrimInfo
pr) ([Either Term Type] -> Term) -> [Either Term Type] -> Term
forall a b. (a -> b) -> a -> b
$ (Type -> Either Term Type) -> [Type] -> [Either Term Type]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Either Term Type
forall a b. b -> Either a b
Right [Type]
tyArgs [Either Term Type] -> [Either Term Type] -> [Either Term Type]
forall a. Semigroup a => a -> a -> a
<>
          [ Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
NaturalLiteral Integer
0))
          , Term -> Either Term Type
forall a b. a -> Either a b
Left (Literal -> Term
Literal (Integer -> Literal
IntegerLiteral Integer
0))
          ]