{-|
  Copyright  :  (C) 2015-2016, University of Twente
  License    :  BSD2 (see the file LICENSE)
  Maintainer :  Christiaan Baaij <christiaan.baaij@gmail.com>

  Helper functions for the 'disjointExpressionConsolidation' transformation

  The 'disjointExpressionConsolidation' transformation lifts applications of
  global binders out of alternatives of case-statements.

  e.g. It converts:

  > case x of
  >   A -> f 3 y
  >   B -> f x x
  >   C -> h x

  into:

  > let f_arg0 = case x of {A -> 3; B -> x}
  >     f_arg1 = case x of {A -> y; B -> x}
  >     f_out  = f f_arg0 f_arg1
  > in  case x of
  >       A -> f_out
  >       B -> f_out
  >       C -> h x
-}

{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecursiveDo #-}
{-# LANGUAGE TemplateHaskell #-}

module Clash.Normalize.DEC
  (collectGlobals
  ,collectGlobalsArgs
  ,isDisjoint
  ,mkDisjointGroup
  )
where

-- external
import           Control.Concurrent.Supply        (splitSupply)
import qualified Control.Lens                     as Lens
import           Data.Bits                        ((.&.),complement)
import           Data.Coerce                      (coerce)
import qualified Data.Either                      as Either
import qualified Data.Foldable                    as Foldable
import qualified Data.Graph                       as Graph
import qualified Data.IntMap.Strict               as IM
import qualified Data.IntSet                      as IntSet
import qualified Data.List                        as List
import qualified Data.List.Extra                  as List
import qualified Data.Map.Strict                  as Map
import qualified Data.Maybe                       as Maybe
import           Data.Monoid                      (All (..))

#if MIN_VERSION_ghc(8,10,0)
import           GHC.Hs.Utils                     (chunkify,mkChunkified)
#else
import           HsUtils                          (chunkify,mkChunkified)
#endif

#if MIN_VERSION_ghc(9,0,0)
import           GHC.Settings.Constants           (mAX_TUPLE_SIZE)
#else
import           Constants                        (mAX_TUPLE_SIZE)
#endif

#if EXPERIMENTAL_EVALUATOR
import           System.IO.Unsafe
#endif

-- internal
import Clash.Core.DataCon    (DataCon)

#if EXPERIMENTAL_EVALUATOR
import Clash.Core.PartialEval
#else
import Clash.Core.Evaluator.Types  (whnf')
#endif

import Clash.Core.FreeVars
  (termFreeVars', typeFreeVars', localVarsDoNotOccurIn)
import Clash.Core.Literal    (Literal (..))
import Clash.Core.Term
  (LetBinding, Pat (..), PrimInfo (..), Term (..), TickInfo (..), collectArgs,
   collectArgsTicks, mkApps, mkTicks, patIds)
import Clash.Core.TermInfo   (termType)
import Clash.Core.TyCon      (TyConMap, TyConName, tyConDataCons)
import Clash.Core.Type       (Type, isPolyFunTy, mkTyConApp, splitFunForallTy)
import Clash.Core.Util       (mkInternalVar, mkSelectorCase, sccLetBindings)
import Clash.Core.Var        (isGlobalId, isLocalId)
import Clash.Core.VarEnv
  (InScopeSet, elemInScopeSet, extendInScopeSetList, notElemInScopeSet, unionInScope)
import Clash.Normalize.Types (NormalizeState)
import Clash.Rewrite.Types
import Clash.Rewrite.Util    (isUntranslatableType)
import Clash.Rewrite.WorkFree (isConstant)
import Clash.Unique          (lookupUniqMap)
import Clash.Util

data CaseTree a
  = Leaf a
  | LB [LetBinding] (CaseTree a)
  | Branch Term [(Pat,CaseTree a)]
  deriving (CaseTree a -> CaseTree a -> Bool
(CaseTree a -> CaseTree a -> Bool)
-> (CaseTree a -> CaseTree a -> Bool) -> Eq (CaseTree a)
forall a. Eq a => CaseTree a -> CaseTree a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: CaseTree a -> CaseTree a -> Bool
$c/= :: forall a. Eq a => CaseTree a -> CaseTree a -> Bool
== :: CaseTree a -> CaseTree a -> Bool
$c== :: forall a. Eq a => CaseTree a -> CaseTree a -> Bool
Eq,Int -> CaseTree a -> ShowS
[CaseTree a] -> ShowS
CaseTree a -> String
(Int -> CaseTree a -> ShowS)
-> (CaseTree a -> String)
-> ([CaseTree a] -> ShowS)
-> Show (CaseTree a)
forall a. Show a => Int -> CaseTree a -> ShowS
forall a. Show a => [CaseTree a] -> ShowS
forall a. Show a => CaseTree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CaseTree a] -> ShowS
$cshowList :: forall a. Show a => [CaseTree a] -> ShowS
show :: CaseTree a -> String
$cshow :: forall a. Show a => CaseTree a -> String
showsPrec :: Int -> CaseTree a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> CaseTree a -> ShowS
Show,a -> CaseTree b -> CaseTree a
(a -> b) -> CaseTree a -> CaseTree b
(forall a b. (a -> b) -> CaseTree a -> CaseTree b)
-> (forall a b. a -> CaseTree b -> CaseTree a) -> Functor CaseTree
forall a b. a -> CaseTree b -> CaseTree a
forall a b. (a -> b) -> CaseTree a -> CaseTree b
forall (f :: Type -> Type).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> CaseTree b -> CaseTree a
$c<$ :: forall a b. a -> CaseTree b -> CaseTree a
fmap :: (a -> b) -> CaseTree a -> CaseTree b
$cfmap :: forall a b. (a -> b) -> CaseTree a -> CaseTree b
Functor,CaseTree a -> Bool
(a -> m) -> CaseTree a -> m
(a -> b -> b) -> b -> CaseTree a -> b
(forall m. Monoid m => CaseTree m -> m)
-> (forall m a. Monoid m => (a -> m) -> CaseTree a -> m)
-> (forall m a. Monoid m => (a -> m) -> CaseTree a -> m)
-> (forall a b. (a -> b -> b) -> b -> CaseTree a -> b)
-> (forall a b. (a -> b -> b) -> b -> CaseTree a -> b)
-> (forall b a. (b -> a -> b) -> b -> CaseTree a -> b)
-> (forall b a. (b -> a -> b) -> b -> CaseTree a -> b)
-> (forall a. (a -> a -> a) -> CaseTree a -> a)
-> (forall a. (a -> a -> a) -> CaseTree a -> a)
-> (forall a. CaseTree a -> [a])
-> (forall a. CaseTree a -> Bool)
-> (forall a. CaseTree a -> Int)
-> (forall a. Eq a => a -> CaseTree a -> Bool)
-> (forall a. Ord a => CaseTree a -> a)
-> (forall a. Ord a => CaseTree a -> a)
-> (forall a. Num a => CaseTree a -> a)
-> (forall a. Num a => CaseTree a -> a)
-> Foldable CaseTree
forall a. Eq a => a -> CaseTree a -> Bool
forall a. Num a => CaseTree a -> a
forall a. Ord a => CaseTree a -> a
forall m. Monoid m => CaseTree m -> m
forall a. CaseTree a -> Bool
forall a. CaseTree a -> Int
forall a. CaseTree a -> [a]
forall a. (a -> a -> a) -> CaseTree a -> a
forall m a. Monoid m => (a -> m) -> CaseTree a -> m
forall b a. (b -> a -> b) -> b -> CaseTree a -> b
forall a b. (a -> b -> b) -> b -> CaseTree a -> b
forall (t :: Type -> Type).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: CaseTree a -> a
$cproduct :: forall a. Num a => CaseTree a -> a
sum :: CaseTree a -> a
$csum :: forall a. Num a => CaseTree a -> a
minimum :: CaseTree a -> a
$cminimum :: forall a. Ord a => CaseTree a -> a
maximum :: CaseTree a -> a
$cmaximum :: forall a. Ord a => CaseTree a -> a
elem :: a -> CaseTree a -> Bool
$celem :: forall a. Eq a => a -> CaseTree a -> Bool
length :: CaseTree a -> Int
$clength :: forall a. CaseTree a -> Int
null :: CaseTree a -> Bool
$cnull :: forall a. CaseTree a -> Bool
toList :: CaseTree a -> [a]
$ctoList :: forall a. CaseTree a -> [a]
foldl1 :: (a -> a -> a) -> CaseTree a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> CaseTree a -> a
foldr1 :: (a -> a -> a) -> CaseTree a -> a
$cfoldr1 :: forall a. (a -> a -> a) -> CaseTree a -> a
foldl' :: (b -> a -> b) -> b -> CaseTree a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> CaseTree a -> b
foldl :: (b -> a -> b) -> b -> CaseTree a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> CaseTree a -> b
foldr' :: (a -> b -> b) -> b -> CaseTree a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> CaseTree a -> b
foldr :: (a -> b -> b) -> b -> CaseTree a -> b
$cfoldr :: forall a b. (a -> b -> b) -> b -> CaseTree a -> b
foldMap' :: (a -> m) -> CaseTree a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> CaseTree a -> m
foldMap :: (a -> m) -> CaseTree a -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> CaseTree a -> m
fold :: CaseTree m -> m
$cfold :: forall m. Monoid m => CaseTree m -> m
Foldable)

-- | Test if a 'CaseTree' collected from an expression indicates that
-- application of a global binder is disjoint: occur in separate branches of a
-- case-expression.
isDisjoint :: CaseTree ([Either Term Type])
           -> Bool
isDisjoint :: CaseTree [Either Term Type] -> Bool
isDisjoint (Branch Term
_ [(Pat, CaseTree [Either Term Type])
_]) = Bool
False
isDisjoint CaseTree [Either Term Type]
ct = CaseTree [Either Term Type] -> Bool
forall b a. Eq b => CaseTree [Either a b] -> Bool
go CaseTree [Either Term Type]
ct
  where
    go :: CaseTree [Either a b] -> Bool
go (Leaf [Either a b]
_)             = Bool
False
    go (LB [LetBinding]
_ CaseTree [Either a b]
ct')           = CaseTree [Either a b] -> Bool
go CaseTree [Either a b]
ct'
    go (Branch Term
_ [])        = Bool
False
    go (Branch Term
_ [(Pat
_,CaseTree [Either a b]
x)])   = CaseTree [Either a b] -> Bool
go CaseTree [Either a b]
x
    go b :: CaseTree [Either a b]
b@(Branch Term
_ ((Pat, CaseTree [Either a b])
_:(Pat, CaseTree [Either a b])
_:[(Pat, CaseTree [Either a b])]
_)) = [[b]] -> Bool
forall a. Eq a => [a] -> Bool
allEqual (([Either a b] -> [b]) -> [[Either a b]] -> [[b]]
forall a b. (a -> b) -> [a] -> [b]
map [Either a b] -> [b]
forall a b. [Either a b] -> [b]
Either.rights (CaseTree [Either a b] -> [[Either a b]]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList CaseTree [Either a b]
b))

-- Remove empty branches from a 'CaseTree'
removeEmpty :: Eq a => CaseTree [a] -> CaseTree [a]
removeEmpty :: CaseTree [a] -> CaseTree [a]
removeEmpty l :: CaseTree [a]
l@(Leaf [a]
_) = CaseTree [a]
l
removeEmpty (LB [LetBinding]
lb CaseTree [a]
ct) =
  case CaseTree [a] -> CaseTree [a]
forall a. Eq a => CaseTree [a] -> CaseTree [a]
removeEmpty CaseTree [a]
ct of
    Leaf [] -> [a] -> CaseTree [a]
forall a. a -> CaseTree a
Leaf []
    CaseTree [a]
ct'     -> [LetBinding] -> CaseTree [a] -> CaseTree [a]
forall a. [LetBinding] -> CaseTree a -> CaseTree a
LB [LetBinding]
lb CaseTree [a]
ct'
removeEmpty (Branch Term
s [(Pat, CaseTree [a])]
bs) =
  case ((Pat, CaseTree [a]) -> Bool)
-> [(Pat, CaseTree [a])] -> [(Pat, CaseTree [a])]
forall a. (a -> Bool) -> [a] -> [a]
filter ((CaseTree [a] -> CaseTree [a] -> Bool
forall a. Eq a => a -> a -> Bool
/= ([a] -> CaseTree [a]
forall a. a -> CaseTree a
Leaf [])) (CaseTree [a] -> Bool)
-> ((Pat, CaseTree [a]) -> CaseTree [a])
-> (Pat, CaseTree [a])
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Pat, CaseTree [a]) -> CaseTree [a]
forall a b. (a, b) -> b
snd) (((Pat, CaseTree [a]) -> (Pat, CaseTree [a]))
-> [(Pat, CaseTree [a])] -> [(Pat, CaseTree [a])]
forall a b. (a -> b) -> [a] -> [b]
map ((CaseTree [a] -> CaseTree [a])
-> (Pat, CaseTree [a]) -> (Pat, CaseTree [a])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second CaseTree [a] -> CaseTree [a]
forall a. Eq a => CaseTree [a] -> CaseTree [a]
removeEmpty) [(Pat, CaseTree [a])]
bs) of
    []  -> [a] -> CaseTree [a]
forall a. a -> CaseTree a
Leaf []
    [(Pat, CaseTree [a])]
bs' -> Term -> [(Pat, CaseTree [a])] -> CaseTree [a]
forall a. Term -> [(Pat, CaseTree a)] -> CaseTree a
Branch Term
s [(Pat, CaseTree [a])]
bs'

-- | Test if all elements in a list are equal to each other.
allEqual :: Eq a => [a] -> Bool
allEqual :: [a] -> Bool
allEqual []     = Bool
True
allEqual (a
x:[a]
xs) = (a -> Bool) -> [a] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
all (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x) [a]
xs

-- | Collect 'CaseTree's for (potentially) disjoint applications of globals out
-- of an expression. Also substitute truly disjoint applications of globals by a
-- reference to a lifted out application.
collectGlobals'
  :: InScopeSet
  -> [(Term,Term)]
  -- ^ Substitution of (applications of) a global binder by a reference to a
  -- lifted term.
  -> [Term]
  -- ^ List of already seen global binders
  -> Term
  -- ^ The expression
  -> Bool
  -- ^ Whether expression is constant
  -> RewriteMonad
      NormalizeState
      (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobals' :: InScopeSet
-> [(Term, Term)]
-> [Term]
-> Term
-> Bool
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobals' InScopeSet
is0 [(Term, Term)]
substitution [Term]
seen (Case Term
scrut Type
ty [Alt]
alts) Bool
_eIsConstant = do
  rec ([Alt]
alts1, InScopeSet
isAlts, [(Term, ([Term], CaseTree [Either Term Type]))]
collectedAlts) <-
        InScopeSet
-> [(Term, Term)]
-> [Term]
-> Term
-> [Alt]
-> RewriteMonad
     NormalizeState
     ([Alt], InScopeSet,
      [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobalsAlts InScopeSet
is0 [(Term, Term)]
substitution [Term]
seen Term
scrut1 [Alt]
alts
      (Term
scrut1, InScopeSet
isScrut, [(Term, ([Term], CaseTree [Either Term Type]))]
collectedScrut) <-
        InScopeSet
-> [(Term, Term)]
-> [Term]
-> Term
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobals InScopeSet
is0 [(Term, Term)]
substitution (((Term, ([Term], CaseTree [Either Term Type])) -> Term)
-> [(Term, ([Term], CaseTree [Either Term Type]))] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map (Term, ([Term], CaseTree [Either Term Type])) -> Term
forall a b. (a, b) -> a
fst [(Term, ([Term], CaseTree [Either Term Type]))]
collectedAlts [Term] -> [Term] -> [Term]
forall a. [a] -> [a] -> [a]
++ [Term]
seen) Term
scrut
  (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
forall (m :: Type -> Type) a. Monad m => a -> m a
return ( Term -> Type -> [Alt] -> Term
Case Term
scrut1 Type
ty [Alt]
alts1
         , InScopeSet -> InScopeSet -> InScopeSet
unionInScope InScopeSet
isAlts InScopeSet
isScrut
         , [(Term, ([Term], CaseTree [Either Term Type]))]
collectedAlts [(Term, ([Term], CaseTree [Either Term Type]))]
-> [(Term, ([Term], CaseTree [Either Term Type]))]
-> [(Term, ([Term], CaseTree [Either Term Type]))]
forall a. [a] -> [a] -> [a]
++ [(Term, ([Term], CaseTree [Either Term Type]))]
collectedScrut )

collectGlobals' InScopeSet
is0 [(Term, Term)]
substitution [Term]
seen e :: Term
e@(Term -> (Term, [Either Term Type], [TickInfo])
collectArgsTicks -> (Term
fun, args :: [Either Term Type]
args@(Either Term Type
_:[Either Term Type]
_), [TickInfo]
ticks)) Bool
eIsconstant
  | Bool -> Bool
not Bool
eIsconstant = 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
Lens' RewriteEnv TyConMap
tcCache
    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
    Evaluator
evaluate <- Getting Evaluator RewriteEnv Evaluator
-> RewriteMonad NormalizeState Evaluator
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting Evaluator RewriteEnv Evaluator
Lens' RewriteEnv Evaluator
evaluator
    Supply
ids <- Getting Supply (RewriteState NormalizeState) Supply
-> RewriteMonad NormalizeState Supply
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting Supply (RewriteState NormalizeState) Supply
forall extra. Lens' (RewriteState extra) Supply
uniqSupply
    let (Supply
ids1,Supply
ids2) = Supply -> (Supply, Supply)
splitSupply Supply
ids
    (Supply -> Identity Supply)
-> RewriteState NormalizeState
-> Identity (RewriteState NormalizeState)
forall extra. Lens' (RewriteState extra) Supply
uniqSupply ((Supply -> Identity Supply)
 -> RewriteState NormalizeState
 -> Identity (RewriteState NormalizeState))
-> Supply -> RewriteMonad NormalizeState ()
forall s (m :: Type -> Type) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
Lens..= Supply
ids2
#if EXPERIMENTAL_EVALUATOR
    (i,_) <- Lens.use curFun
    heap <- Lens.use ioHeap
    addr <- Lens.use ioAddr
    fuel <- Lens.view fuelLimit
    let genv = mkGlobalEnv bndrs tcm is0 ids1 fuel heap addr
    let eval = fst . unsafePerformIO . nf evaluate genv False i
#else
    PrimHeap
gh <- Getting PrimHeap (RewriteState NormalizeState) PrimHeap
-> RewriteMonad NormalizeState PrimHeap
forall s (m :: Type -> Type) a.
MonadState s m =>
Getting a s a -> m a
Lens.use Getting PrimHeap (RewriteState NormalizeState) PrimHeap
forall extra. Lens' (RewriteState extra) PrimHeap
globalHeap
    let eval :: Term -> Term
eval = (Getting Term (PrimHeap, PureHeap, Term) Term
-> (PrimHeap, PureHeap, Term) -> Term
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting Term (PrimHeap, PureHeap, Term) Term
forall s t a b. Field3 s t a b => Lens s t a b
Lens._3) ((PrimHeap, PureHeap, Term) -> Term)
-> (Term -> (PrimHeap, PureHeap, Term)) -> Term -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Evaluator
-> BindingMap
-> TyConMap
-> PrimHeap
-> Supply
-> InScopeSet
-> Bool
-> Term
-> (PrimHeap, PureHeap, Term)
whnf' Evaluator
evaluate BindingMap
bndrs TyConMap
tcm PrimHeap
gh Supply
ids1 InScopeSet
is0 Bool
False
#endif
    let eTy :: Type
eTy  = TyConMap -> Term -> Type
termType TyConMap
tcm Term
e
    Bool
untran <- Bool -> Type -> RewriteMonad NormalizeState Bool
forall extra. Bool -> Type -> RewriteMonad extra Bool
isUntranslatableType Bool
False Type
eTy
    case Bool
untran of
      -- Don't lift out non-representable values, because they cannot be let-bound
      -- in our desired normal form.
      Bool
False -> do
        -- Look for, and substitute by, disjoint applications of globals in
        -- the arguments first before considering the current term in function
        -- position. Doing it in the other order (this term in function position
        -- first, followed by arguments) resulted in issue #1322
        ([Either Term Type]
args1,InScopeSet
isArgs,[(Term, ([Term], CaseTree [Either Term Type]))]
collectedArgs) <-
          InScopeSet
-> [(Term, Term)]
-> [Term]
-> [Either Term Type]
-> RewriteMonad
     NormalizeState
     ([Either Term Type], InScopeSet,
      [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobalsArgs InScopeSet
is0 [(Term, Term)]
substitution [Term]
seen [Either Term Type]
args
        let seenInArgs :: [Term]
seenInArgs = ((Term, ([Term], CaseTree [Either Term Type])) -> Term)
-> [(Term, ([Term], CaseTree [Either Term Type]))] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map (Term, ([Term], CaseTree [Either Term Type])) -> Term
forall a b. (a, b) -> a
fst [(Term, ([Term], CaseTree [Either Term Type]))]
collectedArgs [Term] -> [Term] -> [Term]
forall a. [a] -> [a] -> [a]
++ [Term]
seen
        Maybe Term
isInteresting <- InScopeSet
-> (Term -> Term)
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad NormalizeState (Maybe Term)
forall extra.
InScopeSet
-> (Term -> Term)
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad extra (Maybe Term)
interestingToLift InScopeSet
is0 Term -> Term
eval Term
fun [Either Term Type]
args [TickInfo]
ticks
        case Maybe Term
isInteresting of
          Just Term
fun1 | Term
fun1 Term -> [Term] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`notElem` [Term]
seenInArgs -> do
            let e1 :: Term
e1 = Term -> Maybe Term -> Term
forall a. a -> Maybe a -> a
Maybe.fromMaybe (Term -> [Either Term Type] -> Term
mkApps Term
fun1 [Either Term Type]
args1) (Term -> [(Term, Term)] -> Maybe Term
forall a b. Eq a => a -> [(a, b)] -> Maybe b
List.lookup Term
fun1 [(Term, Term)]
substitution)
            -- This function is lifted out an environment with the currently 'seen'
            -- binders. When we later apply substitution, we need to start with this
            -- environment, otherwise we perform incorrect substitutions in the
            -- arguments.
            (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Term
e1,InScopeSet
isArgs,(Term
fun1,([Term]
seen,[Either Term Type] -> CaseTree [Either Term Type]
forall a. a -> CaseTree a
Leaf [Either Term Type]
args1))(Term, ([Term], CaseTree [Either Term Type]))
-> [(Term, ([Term], CaseTree [Either Term Type]))]
-> [(Term, ([Term], CaseTree [Either Term Type]))]
forall a. a -> [a] -> [a]
:[(Term, ([Term], CaseTree [Either Term Type]))]
collectedArgs)
          Maybe Term
_ -> (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
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]
args1, InScopeSet
isArgs, [(Term, ([Term], CaseTree [Either Term Type]))]
collectedArgs)
      Bool
_ -> (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Term
e,InScopeSet
is0,[])

-- FIXME: This duplicates A LOT of let-bindings, where I just pray that after
-- the ANF, CSE, and DeadCodeRemoval pass all duplicates are removed.
--
-- I think we should be able to do better, but perhaps we cannot fix it here.
collectGlobals' InScopeSet
is0 [(Term, Term)]
substitution [Term]
seen (Letrec [LetBinding]
lbs Term
body) Bool
_eIsConstant = do
  let is1 :: InScopeSet
is1 = InScopeSet -> [Var Term] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 ((LetBinding -> Var Term) -> [LetBinding] -> [Var Term]
forall a b. (a -> b) -> [a] -> [b]
map LetBinding -> Var Term
forall a b. (a, b) -> a
fst [LetBinding]
lbs)
  (Term
body1,InScopeSet
isBody,[(Term, ([Term], CaseTree [Either Term Type]))]
collectedBody) <-
    InScopeSet
-> [(Term, Term)]
-> [Term]
-> Term
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobals InScopeSet
is1 [(Term, Term)]
substitution [Term]
seen Term
body
  ([LetBinding]
lbs1,InScopeSet
isBndrs,[(Term, ([Term], CaseTree [Either Term Type]))]
collectedBndrs) <-
    InScopeSet
-> [(Term, Term)]
-> [Term]
-> [LetBinding]
-> RewriteMonad
     NormalizeState
     ([LetBinding], InScopeSet,
      [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobalsLbs InScopeSet
is1 [(Term, Term)]
substitution (((Term, ([Term], CaseTree [Either Term Type])) -> Term)
-> [(Term, ([Term], CaseTree [Either Term Type]))] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map (Term, ([Term], CaseTree [Either Term Type])) -> Term
forall a b. (a, b) -> a
fst [(Term, ([Term], CaseTree [Either Term Type]))]
collectedBody [Term] -> [Term] -> [Term]
forall a. [a] -> [a] -> [a]
++ [Term]
seen) [LetBinding]
lbs
  (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
forall (m :: Type -> Type) a. Monad m => a -> m a
return ( [LetBinding] -> Term -> Term
Letrec [LetBinding]
lbs1 Term
body1
         , InScopeSet -> InScopeSet -> InScopeSet
unionInScope InScopeSet
isBody InScopeSet
isBndrs
         , ((Term, ([Term], CaseTree [Either Term Type]))
 -> (Term, ([Term], CaseTree [Either Term Type])))
-> [(Term, ([Term], CaseTree [Either Term Type]))]
-> [(Term, ([Term], CaseTree [Either Term Type]))]
forall a b. (a -> b) -> [a] -> [b]
map ((([Term], CaseTree [Either Term Type])
 -> ([Term], CaseTree [Either Term Type]))
-> (Term, ([Term], CaseTree [Either Term Type]))
-> (Term, ([Term], CaseTree [Either Term Type]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ((CaseTree [Either Term Type] -> CaseTree [Either Term Type])
-> ([Term], CaseTree [Either Term Type])
-> ([Term], CaseTree [Either Term Type])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ([LetBinding]
-> CaseTree [Either Term Type] -> CaseTree [Either Term Type]
forall a. [LetBinding] -> CaseTree a -> CaseTree a
LB [LetBinding]
lbs1))) ([(Term, ([Term], CaseTree [Either Term Type]))]
collectedBody [(Term, ([Term], CaseTree [Either Term Type]))]
-> [(Term, ([Term], CaseTree [Either Term Type]))]
-> [(Term, ([Term], CaseTree [Either Term Type]))]
forall a. [a] -> [a] -> [a]
++ [(Term, ([Term], CaseTree [Either Term Type]))]
collectedBndrs)
         )

collectGlobals' InScopeSet
is0 [(Term, Term)]
substitution [Term]
seen (Tick TickInfo
t Term
e) Bool
eIsConstant = do
  (Term
e1,InScopeSet
is1,[(Term, ([Term], CaseTree [Either Term Type]))]
collected) <- InScopeSet
-> [(Term, Term)]
-> [Term]
-> Term
-> Bool
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobals' InScopeSet
is0 [(Term, Term)]
substitution [Term]
seen Term
e Bool
eIsConstant
  (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (TickInfo -> Term -> Term
Tick TickInfo
t Term
e1, InScopeSet
is1, [(Term, ([Term], CaseTree [Either Term Type]))]
collected)

collectGlobals' InScopeSet
is0 [(Term, Term)]
_ [Term]
_ Term
e Bool
_ = (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Term
e,InScopeSet
is0,[])

-- | Collect 'CaseTree's for (potentially) disjoint applications of globals out
-- of an expression. Also substitute truly disjoint applications of globals by a
-- reference to a lifted out application.
collectGlobals
  :: InScopeSet
  -> [(Term,Term)]
  -- ^ Substitution of (applications of) a global binder by a reference to
  -- a lifted term.
  -> [Term]
  -- ^ List of already seen global binders
  -> Term
  -- ^ The expression
  -> RewriteMonad
      NormalizeState
      (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobals :: InScopeSet
-> [(Term, Term)]
-> [Term]
-> Term
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobals InScopeSet
inScope [(Term, Term)]
substitution [Term]
seen Term
e =
  InScopeSet
-> [(Term, Term)]
-> [Term]
-> Term
-> Bool
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobals' InScopeSet
inScope [(Term, Term)]
substitution [Term]
seen Term
e (Term -> Bool
isConstant Term
e)

-- | Collect 'CaseTree's for (potentially) disjoint applications of globals out
-- of a list of application arguments. Also substitute truly disjoint
-- applications of globals by a reference to a lifted out application.
collectGlobalsArgs ::
     InScopeSet
  -> [(Term,Term)] -- ^ Substitution of (applications of) a global
                   -- binder by a reference to a lifted term.
  -> [Term] -- ^ List of already seen global binders
  -> [Either Term Type] -- ^ The list of arguments
  -> RewriteMonad NormalizeState
                  ([Either Term Type]
                  ,InScopeSet
                  ,[(Term,([Term],CaseTree [(Either Term Type)]))]
                  )
collectGlobalsArgs :: InScopeSet
-> [(Term, Term)]
-> [Term]
-> [Either Term Type]
-> RewriteMonad
     NormalizeState
     ([Either Term Type], InScopeSet,
      [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobalsArgs InScopeSet
is0 [(Term, Term)]
substitution [Term]
seen [Either Term Type]
args = do
    ((InScopeSet
is1,[Term]
_),([Either Term Type]
args',[[(Term, ([Term], CaseTree [Either Term Type]))]]
collected)) <- ([(Either Term Type,
   [(Term, ([Term], CaseTree [Either Term Type]))])]
 -> ([Either Term Type],
     [[(Term, ([Term], CaseTree [Either Term Type]))]]))
-> ((InScopeSet, [Term]),
    [(Either Term Type,
      [(Term, ([Term], CaseTree [Either Term Type]))])])
-> ((InScopeSet, [Term]),
    ([Either Term Type],
     [[(Term, ([Term], CaseTree [Either Term Type]))]]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [(Either Term Type,
  [(Term, ([Term], CaseTree [Either Term Type]))])]
-> ([Either Term Type],
    [[(Term, ([Term], CaseTree [Either Term Type]))]])
forall a b. [(a, b)] -> ([a], [b])
unzip (((InScopeSet, [Term]),
  [(Either Term Type,
    [(Term, ([Term], CaseTree [Either Term Type]))])])
 -> ((InScopeSet, [Term]),
     ([Either Term Type],
      [[(Term, ([Term], CaseTree [Either Term Type]))]])))
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      [(Either Term Type,
        [(Term, ([Term], CaseTree [Either Term Type]))])])
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      ([Either Term Type],
       [[(Term, ([Term], CaseTree [Either Term Type]))]]))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ((InScopeSet, [Term])
 -> Either Term Type
 -> RewriteMonad
      NormalizeState
      ((InScopeSet, [Term]),
       (Either Term Type,
        [(Term, ([Term], CaseTree [Either Term Type]))])))
-> (InScopeSet, [Term])
-> [Either Term Type]
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      [(Either Term Type,
        [(Term, ([Term], CaseTree [Either Term Type]))])])
forall (m :: Type -> Type) acc x y.
Monad m =>
(acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
List.mapAccumLM (InScopeSet, [Term])
-> Either Term Type
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      (Either Term Type,
       [(Term, ([Term], CaseTree [Either Term Type]))]))
forall b.
(InScopeSet, [Term])
-> Either Term b
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      (Either Term b, [(Term, ([Term], CaseTree [Either Term Type]))]))
go (InScopeSet
is0,[Term]
seen) [Either Term Type]
args
    ([Either Term Type], InScopeSet,
 [(Term, ([Term], CaseTree [Either Term Type]))])
-> RewriteMonad
     NormalizeState
     ([Either Term Type], InScopeSet,
      [(Term, ([Term], CaseTree [Either Term Type]))])
forall (m :: Type -> Type) a. Monad m => a -> m a
return ([Either Term Type]
args',InScopeSet
is1,[[(Term, ([Term], CaseTree [Either Term Type]))]]
-> [(Term, ([Term], CaseTree [Either Term Type]))]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat [[(Term, ([Term], CaseTree [Either Term Type]))]]
collected)
  where
    go :: (InScopeSet, [Term])
-> Either Term b
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      (Either Term b, [(Term, ([Term], CaseTree [Either Term Type]))]))
go (InScopeSet
isN0,[Term]
s) (Left Term
tm) = do
      (Term
tm',InScopeSet
isN1,[(Term, ([Term], CaseTree [Either Term Type]))]
collected) <- InScopeSet
-> [(Term, Term)]
-> [Term]
-> Term
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobals InScopeSet
isN0 [(Term, Term)]
substitution [Term]
s Term
tm
      ((InScopeSet, [Term]),
 (Either Term b, [(Term, ([Term], CaseTree [Either Term Type]))]))
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      (Either Term b, [(Term, ([Term], CaseTree [Either Term Type]))]))
forall (m :: Type -> Type) a. Monad m => a -> m a
return ((InScopeSet
isN1,((Term, ([Term], CaseTree [Either Term Type])) -> Term)
-> [(Term, ([Term], CaseTree [Either Term Type]))] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map (Term, ([Term], CaseTree [Either Term Type])) -> Term
forall a b. (a, b) -> a
fst [(Term, ([Term], CaseTree [Either Term Type]))]
collected [Term] -> [Term] -> [Term]
forall a. [a] -> [a] -> [a]
++ [Term]
s),(Term -> Either Term b
forall a b. a -> Either a b
Left Term
tm',[(Term, ([Term], CaseTree [Either Term Type]))]
collected))
    go (InScopeSet
isN,[Term]
s) (Right b
ty) = ((InScopeSet, [Term]),
 (Either Term b, [(Term, ([Term], CaseTree [Either Term Type]))]))
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      (Either Term b, [(Term, ([Term], CaseTree [Either Term Type]))]))
forall (m :: Type -> Type) a. Monad m => a -> m a
return ((InScopeSet
isN,[Term]
s),(b -> Either Term b
forall a b. b -> Either a b
Right b
ty,[]))

-- | Collect 'CaseTree's for (potentially) disjoint applications of globals out
-- of a list of alternatives. Also substitute truly disjoint applications of
-- globals by a reference to a lifted out application.
collectGlobalsAlts ::
     InScopeSet
  -> [(Term,Term)] -- ^ Substitution of (applications of) a global
                   -- binder by a reference to a lifted term.
  -> [Term] -- ^ List of already seen global binders
  -> Term -- ^ The subject term
  -> [(Pat,Term)] -- ^ The list of alternatives
  -> RewriteMonad NormalizeState
                  ([(Pat,Term)]
                  ,InScopeSet
                  ,[(Term,([Term],CaseTree [(Either Term Type)]))]
                  )
collectGlobalsAlts :: InScopeSet
-> [(Term, Term)]
-> [Term]
-> Term
-> [Alt]
-> RewriteMonad
     NormalizeState
     ([Alt], InScopeSet,
      [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobalsAlts InScopeSet
is0 [(Term, Term)]
substitution [Term]
seen Term
scrut [Alt]
alts = do
    (InScopeSet
is1,([Alt]
alts',[[(Term, ([Term], (Pat, CaseTree [Either Term Type])))]]
collected)) <- ([(Alt, [(Term, ([Term], (Pat, CaseTree [Either Term Type])))])]
 -> ([Alt],
     [[(Term, ([Term], (Pat, CaseTree [Either Term Type])))]]))
-> (InScopeSet,
    [(Alt, [(Term, ([Term], (Pat, CaseTree [Either Term Type])))])])
-> (InScopeSet,
    ([Alt], [[(Term, ([Term], (Pat, CaseTree [Either Term Type])))]]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [(Alt, [(Term, ([Term], (Pat, CaseTree [Either Term Type])))])]
-> ([Alt],
    [[(Term, ([Term], (Pat, CaseTree [Either Term Type])))]])
forall a b. [(a, b)] -> ([a], [b])
unzip ((InScopeSet,
  [(Alt, [(Term, ([Term], (Pat, CaseTree [Either Term Type])))])])
 -> (InScopeSet,
     ([Alt], [[(Term, ([Term], (Pat, CaseTree [Either Term Type])))]])))
-> RewriteMonad
     NormalizeState
     (InScopeSet,
      [(Alt, [(Term, ([Term], (Pat, CaseTree [Either Term Type])))])])
-> RewriteMonad
     NormalizeState
     (InScopeSet,
      ([Alt], [[(Term, ([Term], (Pat, CaseTree [Either Term Type])))]]))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (InScopeSet
 -> Alt
 -> RewriteMonad
      NormalizeState
      (InScopeSet,
       (Alt, [(Term, ([Term], (Pat, CaseTree [Either Term Type])))])))
-> InScopeSet
-> [Alt]
-> RewriteMonad
     NormalizeState
     (InScopeSet,
      [(Alt, [(Term, ([Term], (Pat, CaseTree [Either Term Type])))])])
forall (m :: Type -> Type) acc x y.
Monad m =>
(acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
List.mapAccumLM InScopeSet
-> Alt
-> RewriteMonad
     NormalizeState
     (InScopeSet,
      (Alt, [(Term, ([Term], (Pat, CaseTree [Either Term Type])))]))
go InScopeSet
is0 [Alt]
alts
    let collectedM :: [Map Term ([Term], [(Pat, CaseTree [Either Term Type])])]
collectedM  = ([(Term, ([Term], (Pat, CaseTree [Either Term Type])))]
 -> Map Term ([Term], [(Pat, CaseTree [Either Term Type])]))
-> [[(Term, ([Term], (Pat, CaseTree [Either Term Type])))]]
-> [Map Term ([Term], [(Pat, CaseTree [Either Term Type])])]
forall a b. (a -> b) -> [a] -> [b]
map ([(Term, ([Term], [(Pat, CaseTree [Either Term Type])]))]
-> Map Term ([Term], [(Pat, CaseTree [Either Term Type])])
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Term, ([Term], [(Pat, CaseTree [Either Term Type])]))]
 -> Map Term ([Term], [(Pat, CaseTree [Either Term Type])]))
-> ([(Term, ([Term], (Pat, CaseTree [Either Term Type])))]
    -> [(Term, ([Term], [(Pat, CaseTree [Either Term Type])]))])
-> [(Term, ([Term], (Pat, CaseTree [Either Term Type])))]
-> Map Term ([Term], [(Pat, CaseTree [Either Term Type])])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Term, ([Term], (Pat, CaseTree [Either Term Type])))
 -> (Term, ([Term], [(Pat, CaseTree [Either Term Type])])))
-> [(Term, ([Term], (Pat, CaseTree [Either Term Type])))]
-> [(Term, ([Term], [(Pat, CaseTree [Either Term Type])]))]
forall a b. (a -> b) -> [a] -> [b]
map ((([Term], (Pat, CaseTree [Either Term Type]))
 -> ([Term], [(Pat, CaseTree [Either Term Type])]))
-> (Term, ([Term], (Pat, CaseTree [Either Term Type])))
-> (Term, ([Term], [(Pat, CaseTree [Either Term Type])]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (((Pat, CaseTree [Either Term Type])
 -> [(Pat, CaseTree [Either Term Type])])
-> ([Term], (Pat, CaseTree [Either Term Type]))
-> ([Term], [(Pat, CaseTree [Either Term Type])])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ((Pat, CaseTree [Either Term Type])
-> [(Pat, CaseTree [Either Term Type])]
-> [(Pat, CaseTree [Either Term Type])]
forall a. a -> [a] -> [a]
:[])))) [[(Term, ([Term], (Pat, CaseTree [Either Term Type])))]]
collected
        collectedUN :: Map Term ([Term], [(Pat, CaseTree [Either Term Type])])
collectedUN = (([Term], [(Pat, CaseTree [Either Term Type])])
 -> ([Term], [(Pat, CaseTree [Either Term Type])])
 -> ([Term], [(Pat, CaseTree [Either Term Type])]))
-> [Map Term ([Term], [(Pat, CaseTree [Either Term Type])])]
-> Map Term ([Term], [(Pat, CaseTree [Either Term Type])])
forall (f :: Type -> Type) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith (\([Term]
l1,[(Pat, CaseTree [Either Term Type])]
r1) ([Term]
l2,[(Pat, CaseTree [Either Term Type])]
r2) -> ([Term] -> [Term]
forall a. Eq a => [a] -> [a]
List.nub ([Term]
l1 [Term] -> [Term] -> [Term]
forall a. [a] -> [a] -> [a]
++ [Term]
l2),[(Pat, CaseTree [Either Term Type])]
r1 [(Pat, CaseTree [Either Term Type])]
-> [(Pat, CaseTree [Either Term Type])]
-> [(Pat, CaseTree [Either Term Type])]
forall a. [a] -> [a] -> [a]
++ [(Pat, CaseTree [Either Term Type])]
r2)) [Map Term ([Term], [(Pat, CaseTree [Either Term Type])])]
collectedM
        collected' :: [(Term, ([Term], CaseTree [Either Term Type]))]
collected'  = ((Term, ([Term], [(Pat, CaseTree [Either Term Type])]))
 -> (Term, ([Term], CaseTree [Either Term Type])))
-> [(Term, ([Term], [(Pat, CaseTree [Either Term Type])]))]
-> [(Term, ([Term], CaseTree [Either Term Type]))]
forall a b. (a -> b) -> [a] -> [b]
map ((([Term], [(Pat, CaseTree [Either Term Type])])
 -> ([Term], CaseTree [Either Term Type]))
-> (Term, ([Term], [(Pat, CaseTree [Either Term Type])]))
-> (Term, ([Term], CaseTree [Either Term Type]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (([(Pat, CaseTree [Either Term Type])]
 -> CaseTree [Either Term Type])
-> ([Term], [(Pat, CaseTree [Either Term Type])])
-> ([Term], CaseTree [Either Term Type])
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (Term
-> [(Pat, CaseTree [Either Term Type])]
-> CaseTree [Either Term Type]
forall a. Term -> [(Pat, CaseTree a)] -> CaseTree a
Branch Term
scrut))) (Map Term ([Term], [(Pat, CaseTree [Either Term Type])])
-> [(Term, ([Term], [(Pat, CaseTree [Either Term Type])]))]
forall k a. Map k a -> [(k, a)]
Map.toList Map Term ([Term], [(Pat, CaseTree [Either Term Type])])
collectedUN)
    ([Alt], InScopeSet,
 [(Term, ([Term], CaseTree [Either Term Type]))])
-> RewriteMonad
     NormalizeState
     ([Alt], InScopeSet,
      [(Term, ([Term], CaseTree [Either Term Type]))])
forall (m :: Type -> Type) a. Monad m => a -> m a
return ([Alt]
alts',InScopeSet
is1,[(Term, ([Term], CaseTree [Either Term Type]))]
collected')
  where
    go :: InScopeSet
-> Alt
-> RewriteMonad
     NormalizeState
     (InScopeSet,
      (Alt, [(Term, ([Term], (Pat, CaseTree [Either Term Type])))]))
go InScopeSet
isN0 (Pat
p,Term
e) = do
      let isN1 :: InScopeSet
isN1 = InScopeSet -> [Var Term] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
isN0 (([TyVar], [Var Term]) -> [Var Term]
forall a b. (a, b) -> b
snd (Pat -> ([TyVar], [Var Term])
patIds Pat
p))
      (Term
e',InScopeSet
isN2,[(Term, ([Term], CaseTree [Either Term Type]))]
collected) <- InScopeSet
-> [(Term, Term)]
-> [Term]
-> Term
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobals InScopeSet
isN1 [(Term, Term)]
substitution [Term]
seen Term
e
      (InScopeSet,
 (Alt, [(Term, ([Term], (Pat, CaseTree [Either Term Type])))]))
-> RewriteMonad
     NormalizeState
     (InScopeSet,
      (Alt, [(Term, ([Term], (Pat, CaseTree [Either Term Type])))]))
forall (m :: Type -> Type) a. Monad m => a -> m a
return (InScopeSet
isN2,((Pat
p,Term
e'),((Term, ([Term], CaseTree [Either Term Type]))
 -> (Term, ([Term], (Pat, CaseTree [Either Term Type]))))
-> [(Term, ([Term], CaseTree [Either Term Type]))]
-> [(Term, ([Term], (Pat, CaseTree [Either Term Type])))]
forall a b. (a -> b) -> [a] -> [b]
map ((([Term], CaseTree [Either Term Type])
 -> ([Term], (Pat, CaseTree [Either Term Type])))
-> (Term, ([Term], CaseTree [Either Term Type]))
-> (Term, ([Term], (Pat, CaseTree [Either Term Type])))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ((CaseTree [Either Term Type] -> (Pat, CaseTree [Either Term Type]))
-> ([Term], CaseTree [Either Term Type])
-> ([Term], (Pat, CaseTree [Either Term Type]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (Pat
p,))) [(Term, ([Term], CaseTree [Either Term Type]))]
collected))

-- | Collect 'CaseTree's for (potentially) disjoint applications of globals out
-- of a list of let-bindings. Also substitute truly disjoint applications of
-- globals by a reference to a lifted out application.
collectGlobalsLbs ::
     InScopeSet
  -> [(Term,Term)] -- ^ Substitution of (applications of) a global
                   -- binder by a reference to a lifted term.
  -> [Term] -- ^ List of already seen global binders
  -> [LetBinding] -- ^ The list let-bindings
  -> RewriteMonad NormalizeState
                  ([LetBinding]
                  ,InScopeSet
                  ,[(Term,([Term],CaseTree [(Either Term Type)]))]
                  )
collectGlobalsLbs :: InScopeSet
-> [(Term, Term)]
-> [Term]
-> [LetBinding]
-> RewriteMonad
     NormalizeState
     ([LetBinding], InScopeSet,
      [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobalsLbs InScopeSet
is0 [(Term, Term)]
substitution [Term]
seen [LetBinding]
lbs = do
    let lbsSCCs :: [SCC LetBinding]
lbsSCCs = HasCallStack => [LetBinding] -> [SCC LetBinding]
[LetBinding] -> [SCC LetBinding]
sccLetBindings [LetBinding]
lbs
    ((InScopeSet
is1,[Term]
_),([SCC LetBinding]
lbsSCCs1,[[(Term, ([Term], CaseTree [Either Term Type]))]]
collected)) <-
      ([(SCC LetBinding,
   [(Term, ([Term], CaseTree [Either Term Type]))])]
 -> ([SCC LetBinding],
     [[(Term, ([Term], CaseTree [Either Term Type]))]]))
-> ((InScopeSet, [Term]),
    [(SCC LetBinding,
      [(Term, ([Term], CaseTree [Either Term Type]))])])
-> ((InScopeSet, [Term]),
    ([SCC LetBinding],
     [[(Term, ([Term], CaseTree [Either Term Type]))]]))
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [(SCC LetBinding, [(Term, ([Term], CaseTree [Either Term Type]))])]
-> ([SCC LetBinding],
    [[(Term, ([Term], CaseTree [Either Term Type]))]])
forall a b. [(a, b)] -> ([a], [b])
unzip (((InScopeSet, [Term]),
  [(SCC LetBinding,
    [(Term, ([Term], CaseTree [Either Term Type]))])])
 -> ((InScopeSet, [Term]),
     ([SCC LetBinding],
      [[(Term, ([Term], CaseTree [Either Term Type]))]])))
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      [(SCC LetBinding,
        [(Term, ([Term], CaseTree [Either Term Type]))])])
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      ([SCC LetBinding],
       [[(Term, ([Term], CaseTree [Either Term Type]))]]))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ((InScopeSet, [Term])
 -> SCC LetBinding
 -> RewriteMonad
      NormalizeState
      ((InScopeSet, [Term]),
       (SCC LetBinding, [(Term, ([Term], CaseTree [Either Term Type]))])))
-> (InScopeSet, [Term])
-> [SCC LetBinding]
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      [(SCC LetBinding,
        [(Term, ([Term], CaseTree [Either Term Type]))])])
forall (m :: Type -> Type) acc x y.
Monad m =>
(acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
List.mapAccumLM (InScopeSet, [Term])
-> SCC LetBinding
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      (SCC LetBinding, [(Term, ([Term], CaseTree [Either Term Type]))]))
go (InScopeSet
is0,[Term]
seen) [SCC LetBinding]
lbsSCCs
    ([LetBinding], InScopeSet,
 [(Term, ([Term], CaseTree [Either Term Type]))])
-> RewriteMonad
     NormalizeState
     ([LetBinding], InScopeSet,
      [(Term, ([Term], CaseTree [Either Term Type]))])
forall (m :: Type -> Type) a. Monad m => a -> m a
return ([SCC LetBinding] -> [LetBinding]
forall a. [SCC a] -> [a]
Graph.flattenSCCs [SCC LetBinding]
lbsSCCs1,InScopeSet
is1,[[(Term, ([Term], CaseTree [Either Term Type]))]]
-> [(Term, ([Term], CaseTree [Either Term Type]))]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat [[(Term, ([Term], CaseTree [Either Term Type]))]]
collected)
  where
    go :: (InScopeSet,[Term]) -> Graph.SCC LetBinding
       -> RewriteMonad NormalizeState
                  ((InScopeSet, [Term])
                  ,(Graph.SCC LetBinding
                   ,[(Term,([Term],CaseTree [(Either Term Type)]))]
                   )
                  )
    go :: (InScopeSet, [Term])
-> SCC LetBinding
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      (SCC LetBinding, [(Term, ([Term], CaseTree [Either Term Type]))]))
go (InScopeSet
isN0,[Term]
s) (Graph.AcyclicSCC (Var Term
id_, Term
e)) = do
      (Term
e',InScopeSet
isN1,[(Term, ([Term], CaseTree [Either Term Type]))]
collected) <- InScopeSet
-> [(Term, Term)]
-> [Term]
-> Term
-> RewriteMonad
     NormalizeState
     (Term, InScopeSet, [(Term, ([Term], CaseTree [Either Term Type]))])
collectGlobals InScopeSet
isN0 [(Term, Term)]
substitution [Term]
s Term
e
      ((InScopeSet, [Term]),
 (SCC LetBinding, [(Term, ([Term], CaseTree [Either Term Type]))]))
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      (SCC LetBinding, [(Term, ([Term], CaseTree [Either Term Type]))]))
forall (m :: Type -> Type) a. Monad m => a -> m a
return ((InScopeSet
isN1,((Term, ([Term], CaseTree [Either Term Type])) -> Term)
-> [(Term, ([Term], CaseTree [Either Term Type]))] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map (Term, ([Term], CaseTree [Either Term Type])) -> Term
forall a b. (a, b) -> a
fst [(Term, ([Term], CaseTree [Either Term Type]))]
collected [Term] -> [Term] -> [Term]
forall a. [a] -> [a] -> [a]
++ [Term]
s),(LetBinding -> SCC LetBinding
forall vertex. vertex -> SCC vertex
Graph.AcyclicSCC (Var Term
id_,Term
e'),[(Term, ([Term], CaseTree [Either Term Type]))]
collected))
    -- TODO: This completely skips recursive let-bindings in the collection of
    -- potentially disjoint applications of globals; and skips substituting truly
    -- disjoint applications of globals by a reference to a lifted out application.
    --
    -- This is to prevent the creation of combinational loops that have occurred
    -- "in the wild", but for which we have not been able to a create small
    -- unit test that triggers this creation-of-combinational-loops bug.
    -- Completely skipping recursive let-bindings is taking the hammer to
    -- solving this bug, without knowing whether a scalpel even existed and what
    -- it might look like. We should at some point think hard how traversing
    -- recursive let-bindings can introduce combinational loops, and whether
    -- there exists a solution that can traverse recursive let-bindings,
    -- finding more opportunities for DEC, while not introducing combinational
    -- loops.
    go (InScopeSet, [Term])
acc scc :: SCC LetBinding
scc@(Graph.CyclicSCC {}) = ((InScopeSet, [Term]),
 (SCC LetBinding, [(Term, ([Term], CaseTree [Either Term Type]))]))
-> RewriteMonad
     NormalizeState
     ((InScopeSet, [Term]),
      (SCC LetBinding, [(Term, ([Term], CaseTree [Either Term Type]))]))
forall (m :: Type -> Type) a. Monad m => a -> m a
return ((InScopeSet, [Term])
acc,(SCC LetBinding
scc,[]))

-- | Given a case-tree corresponding to a disjoint interesting \"term-in-a-
-- function-position\", return a let-expression: where the let-binding holds
-- a case-expression selecting between the distinct arguments of the case-tree,
-- and the body is an application of the term applied to the shared arguments of
-- the case tree, and projections of let-binding corresponding to the distinct
-- argument positions.
mkDisjointGroup
  :: InScopeSet
  -- ^ Variables in scope at the very top of the case-tree, i.e., the original
  -- expression
  -> (Term,([Term],CaseTree [(Either Term Type)]))
  -- ^ Case-tree of arguments belonging to the applied term.
  -> RewriteMonad NormalizeState (Term,[Term])
mkDisjointGroup :: InScopeSet
-> (Term, ([Term], CaseTree [Either Term Type]))
-> RewriteMonad NormalizeState (Term, [Term])
mkDisjointGroup InScopeSet
inScope (Term
fun,([Term]
seen,CaseTree [Either Term Type]
cs)) = do
    let argss :: [[Either Term Type]]
argss    = CaseTree [Either Term Type] -> [[Either Term Type]]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList CaseTree [Either Term Type]
cs
        argssT :: [(Int, [Either Term Type])]
argssT   = [Int] -> [[Either Term Type]] -> [(Int, [Either Term Type])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] ([[Either Term Type]] -> [[Either Term Type]]
forall a. [[a]] -> [[a]]
List.transpose [[Either Term Type]]
argss)
        ([(Int, [Either Term Type])]
sharedT,[(Int, [Either Term Type])]
distinctT) = ((Int, [Either Term Type]) -> Bool)
-> [(Int, [Either Term Type])]
-> ([(Int, [Either Term Type])], [(Int, [Either Term Type])])
forall a. (a -> Bool) -> [a] -> ([a], [a])
List.partition (InScopeSet -> [Either Term Type] -> Bool
areShared InScopeSet
inScope ([Either Term Type] -> Bool)
-> ((Int, [Either Term Type]) -> [Either Term Type])
-> (Int, [Either Term Type])
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, [Either Term Type]) -> [Either Term Type]
forall a b. (a, b) -> b
snd) [(Int, [Either Term Type])]
argssT
        shared :: [(Int, Either Term Type)]
shared   = ((Int, [Either Term Type]) -> (Int, Either Term Type))
-> [(Int, [Either Term Type])] -> [(Int, Either Term Type)]
forall a b. (a -> b) -> [a] -> [b]
map (([Either Term Type] -> Either Term Type)
-> (Int, [Either Term Type]) -> (Int, Either Term Type)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [Either Term Type] -> Either Term Type
forall a. [a] -> a
head) [(Int, [Either Term Type])]
sharedT
        distinct :: [[Term]]
distinct = ([Either Term Type] -> [Term]) -> [[Either Term Type]] -> [[Term]]
forall a b. (a -> b) -> [a] -> [b]
map ([Either Term Type] -> [Term]
forall a b. [Either a b] -> [a]
Either.lefts) ([[Either Term Type]] -> [[Either Term Type]]
forall a. [[a]] -> [[a]]
List.transpose (((Int, [Either Term Type]) -> [Either Term Type])
-> [(Int, [Either Term Type])] -> [[Either Term Type]]
forall a b. (a -> b) -> [a] -> [b]
map (Int, [Either Term Type]) -> [Either Term Type]
forall a b. (a, b) -> b
snd [(Int, [Either Term Type])]
distinctT))
        cs' :: CaseTree [(Int, Either Term Type)]
cs'      = ([Either Term Type] -> [(Int, Either Term Type)])
-> CaseTree [Either Term Type]
-> CaseTree [(Int, Either Term Type)]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Int] -> [Either Term Type] -> [(Int, Either Term Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..]) CaseTree [Either Term Type]
cs
        cs'' :: CaseTree [Term]
cs''     = CaseTree [Term] -> CaseTree [Term]
forall a. Eq a => CaseTree [a] -> CaseTree [a]
removeEmpty
                 (CaseTree [Term] -> CaseTree [Term])
-> CaseTree [Term] -> CaseTree [Term]
forall a b. (a -> b) -> a -> b
$ ([(Int, Either Term Type)] -> [Term])
-> CaseTree [(Int, Either Term Type)] -> CaseTree [Term]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Either Term Type] -> [Term]
forall a b. [Either a b] -> [a]
Either.lefts ([Either Term Type] -> [Term])
-> ([(Int, Either Term Type)] -> [Either Term Type])
-> [(Int, Either Term Type)]
-> [Term]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Either Term Type) -> Either Term Type)
-> [(Int, Either Term Type)] -> [Either Term Type]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Either Term Type) -> Either Term Type
forall a b. (a, b) -> b
snd)
                        (if [(Int, Either Term Type)] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
null [(Int, Either Term Type)]
shared
                           then CaseTree [(Int, Either Term Type)]
cs'
                           else ([(Int, Either Term Type)] -> [(Int, Either Term Type)])
-> CaseTree [(Int, Either Term Type)]
-> CaseTree [(Int, Either Term Type)]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (((Int, Either Term Type) -> Bool)
-> [(Int, Either Term Type)] -> [(Int, Either Term Type)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int, Either Term Type) -> [(Int, Either Term Type)] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`notElem` [(Int, Either Term Type)]
shared)) CaseTree [(Int, Either Term Type)]
cs')
    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
Lens' RewriteEnv TyConMap
tcCache
    (Maybe LetBinding
distinctCaseM,[Term]
distinctProjections) <- case [[Term]]
distinct of
      -- only shared arguments: do nothing.
      [] -> (Maybe LetBinding, [Term])
-> RewriteMonad NormalizeState (Maybe LetBinding, [Term])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Maybe LetBinding
forall a. Maybe a
Nothing,[])
      -- Create selectors and projections
      ([Term]
uc:[[Term]]
_) -> do
        let argTys :: [Type]
argTys = (Term -> Type) -> [Term] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TyConMap -> Term -> Type
termType TyConMap
tcm) [Term]
uc
        InScopeSet
-> [Type]
-> CaseTree [Term]
-> RewriteMonad NormalizeState (Maybe LetBinding, [Term])
disJointSelProj InScopeSet
inScope [Type]
argTys CaseTree [Term]
cs''
    let newArgs :: [Either Term Type]
newArgs = Int -> [(Int, Either Term Type)] -> [Term] -> [Either Term Type]
mkDJArgs Int
0 [(Int, Either Term Type)]
shared [Term]
distinctProjections
    case Maybe LetBinding
distinctCaseM of
      Just LetBinding
lb -> (Term, [Term]) -> RewriteMonad NormalizeState (Term, [Term])
forall (m :: Type -> Type) a. Monad m => a -> m a
return ([LetBinding] -> Term -> Term
Letrec [LetBinding
lb] (Term -> [Either Term Type] -> Term
mkApps Term
fun [Either Term Type]
newArgs), [Term]
seen)
      Maybe LetBinding
Nothing -> (Term, [Term]) -> RewriteMonad NormalizeState (Term, [Term])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Term -> [Either Term Type] -> Term
mkApps Term
fun [Either Term Type]
newArgs, [Term]
seen)

-- | Create a single selector for all the representable distinct arguments by
-- selecting between tuples. This selector is only ('Just') created when the
-- number of representable uncommmon arguments is larger than one, otherwise it
-- is not ('Nothing').
--
-- It also returns:
--
-- * For all the non-representable distinct arguments: a selector
-- * For all the representable distinct arguments: a projection out of the tuple
--   created by the larger selector. If this larger selector does not exist, a
--   single selector is created for the single representable distinct argument.
disJointSelProj
  :: InScopeSet
  -> [Type]
  -- ^ Types of the arguments
  -> CaseTree [Term]
  -- The case-tree of arguments
  -> RewriteMonad NormalizeState (Maybe LetBinding,[Term])
disJointSelProj :: InScopeSet
-> [Type]
-> CaseTree [Term]
-> RewriteMonad NormalizeState (Maybe LetBinding, [Term])
disJointSelProj InScopeSet
_ [Type]
_ (Leaf []) = (Maybe LetBinding, [Term])
-> RewriteMonad NormalizeState (Maybe LetBinding, [Term])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Maybe LetBinding
forall a. Maybe a
Nothing,[])
disJointSelProj InScopeSet
inScope [Type]
argTys CaseTree [Term]
cs = 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
Lens' RewriteEnv TyConMap
tcCache
    IntMap TyConName
tupTcm <- Getting (IntMap TyConName) RewriteEnv (IntMap TyConName)
-> RewriteMonad NormalizeState (IntMap TyConName)
forall s (m :: Type -> Type) a.
MonadReader s m =>
Getting a s a -> m a
Lens.view Getting (IntMap TyConName) RewriteEnv (IntMap TyConName)
Lens' RewriteEnv (IntMap TyConName)
tupleTcCache
    let maxIndex :: Int
maxIndex = [Type] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Type]
argTys Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        css :: [CaseTree [Term]]
css = (Int -> CaseTree [Term]) -> [Int] -> [CaseTree [Term]]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
i -> ([Term] -> [Term]) -> CaseTree [Term] -> CaseTree [Term]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Term -> [Term] -> [Term]
forall a. a -> [a] -> [a]
:[]) (Term -> [Term]) -> ([Term] -> Term) -> [Term] -> [Term]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Term] -> Int -> Term
forall a. [a] -> Int -> a
!!Int
i)) CaseTree [Term]
cs) [Int
0..Int
maxIndex]
    ([(Int, Type)]
untran,[(Int, Type)]
tran) <- ((Int, Type) -> RewriteMonad NormalizeState Bool)
-> [(Int, Type)]
-> RewriteMonad NormalizeState ([(Int, Type)], [(Int, Type)])
forall (m :: Type -> Type) a.
Monad m =>
(a -> m Bool) -> [a] -> m ([a], [a])
List.partitionM (Bool -> Type -> RewriteMonad NormalizeState Bool
forall extra. Bool -> Type -> RewriteMonad extra Bool
isUntranslatableType Bool
False (Type -> RewriteMonad NormalizeState Bool)
-> ((Int, Type) -> Type)
-> (Int, Type)
-> RewriteMonad NormalizeState Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Type) -> Type
forall a b. (a, b) -> b
snd) ([Int] -> [Type] -> [(Int, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [Type]
argTys)
    let untranCs :: [CaseTree [Term]]
untranCs   = (Int -> CaseTree [Term]) -> [Int] -> [CaseTree [Term]]
forall a b. (a -> b) -> [a] -> [b]
map ([CaseTree [Term]]
css[CaseTree [Term]] -> Int -> CaseTree [Term]
forall a. [a] -> Int -> a
!!) (((Int, Type) -> Int) -> [(Int, Type)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Type) -> Int
forall a b. (a, b) -> a
fst [(Int, Type)]
untran)
        untranSels :: [Term]
untranSels = ((Int, Type) -> CaseTree [Term] -> Term)
-> [(Int, Type)] -> [CaseTree [Term]] -> [Term]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(Int
_,Type
ty) CaseTree [Term]
cs' -> TyConMap
-> IntMap TyConName -> Type -> [Type] -> CaseTree [Term] -> Term
genCase TyConMap
tcm IntMap TyConName
tupTcm Type
ty [Type
ty]  CaseTree [Term]
cs')
                             [(Int, Type)]
untran [CaseTree [Term]]
untranCs
    (Maybe LetBinding
lbM,[Term]
projs) <- case [(Int, Type)]
tran of
      []       -> (Maybe LetBinding, [Term])
-> RewriteMonad NormalizeState (Maybe LetBinding, [Term])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Maybe LetBinding
forall a. Maybe a
Nothing,[])
      [(Int
i,Type
ty)] -> (Maybe LetBinding, [Term])
-> RewriteMonad NormalizeState (Maybe LetBinding, [Term])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Maybe LetBinding
forall a. Maybe a
Nothing,[TyConMap
-> IntMap TyConName -> Type -> [Type] -> CaseTree [Term] -> Term
genCase TyConMap
tcm IntMap TyConName
tupTcm Type
ty [Type
ty] ([CaseTree [Term]]
css[CaseTree [Term]] -> Int -> CaseTree [Term]
forall a. [a] -> Int -> a
!!Int
i)])
      [(Int, Type)]
tys      -> do
        let m :: Int
m            = [(Int, Type)] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [(Int, Type)]
tys
            ([Int]
tyIxs,[Type]
tys') = [(Int, Type)] -> ([Int], [Type])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Int, Type)]
tys
            tupTy :: Type
tupTy        = TyConMap -> IntMap TyConName -> [Type] -> Type
mkBigTupTy TyConMap
tcm IntMap TyConName
tupTcm [Type]
tys'
            cs' :: CaseTree [Term]
cs'          = ([Term] -> [Term]) -> CaseTree [Term] -> CaseTree [Term]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (\[Term]
es -> (Int -> Term) -> [Int] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map ([Term]
es [Term] -> Int -> Term
forall a. [a] -> Int -> a
!!) [Int]
tyIxs) CaseTree [Term]
cs
            djCase :: Term
djCase       = TyConMap
-> IntMap TyConName -> Type -> [Type] -> CaseTree [Term] -> Term
genCase TyConMap
tcm IntMap TyConName
tupTcm Type
tupTy [Type]
tys' CaseTree [Term]
cs'
        Var Term
scrutId <- InScopeSet
-> OccName -> Type -> RewriteMonad NormalizeState (Var Term)
forall (m :: Type -> Type).
MonadUnique m =>
InScopeSet -> OccName -> Type -> m (Var Term)
mkInternalVar InScopeSet
inScope OccName
"tupIn" Type
tupTy
        [Term]
projections <- (Int -> RewriteMonad NormalizeState Term)
-> [Int] -> RewriteMonad NormalizeState [Term]
forall (t :: Type -> Type) (m :: Type -> Type) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (InScopeSet
-> TyConMap
-> IntMap TyConName
-> Term
-> [Type]
-> Int
-> RewriteMonad NormalizeState Term
forall (m :: Type -> Type).
MonadUnique m =>
InScopeSet
-> TyConMap -> IntMap TyConName -> Term -> [Type] -> Int -> m Term
mkBigTupSelector InScopeSet
inScope TyConMap
tcm IntMap TyConName
tupTcm (Var Term -> Term
Var Var Term
scrutId) [Type]
tys') [Int
0..Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]
        (Maybe LetBinding, [Term])
-> RewriteMonad NormalizeState (Maybe LetBinding, [Term])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (LetBinding -> Maybe LetBinding
forall a. a -> Maybe a
Just (Var Term
scrutId,Term
djCase),[Term]
projections)
    let selProjs :: [Term]
selProjs = Int -> [(Int, Term)] -> [Term] -> [Term]
forall a b. (Eq a, Num a) => a -> [(a, b)] -> [b] -> [b]
tranOrUnTran Int
0 ([Int] -> [Term] -> [(Int, Term)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Int, Type) -> Int) -> [(Int, Type)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Type) -> Int
forall a b. (a, b) -> a
fst [(Int, Type)]
untran) [Term]
untranSels) [Term]
projs

    (Maybe LetBinding, [Term])
-> RewriteMonad NormalizeState (Maybe LetBinding, [Term])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Maybe LetBinding
lbM,[Term]
selProjs)
  where
    tranOrUnTran :: a -> [(a, b)] -> [b] -> [b]
tranOrUnTran a
_ []       [b]
projs     = [b]
projs
    tranOrUnTran a
_ [(a, b)]
sels     []        = ((a, b) -> b) -> [(a, b)] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> b
forall a b. (a, b) -> b
snd [(a, b)]
sels
    tranOrUnTran a
n ((a
ut,b
s):[(a, b)]
uts) (b
p:[b]
projs)
      | a
n a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
ut   = b
s b -> [b] -> [b]
forall a. a -> [a] -> [a]
: a -> [(a, b)] -> [b] -> [b]
tranOrUnTran (a
na -> a -> a
forall a. Num a => a -> a -> a
+a
1) [(a, b)]
uts          (b
pb -> [b] -> [b]
forall a. a -> [a] -> [a]
:[b]
projs)
      | Bool
otherwise = b
p b -> [b] -> [b]
forall a. a -> [a] -> [a]
: a -> [(a, b)] -> [b] -> [b]
tranOrUnTran (a
na -> a -> a
forall a. Num a => a -> a -> a
+a
1) ((a
ut,b
s)(a, b) -> [(a, b)] -> [(a, b)]
forall a. a -> [a] -> [a]
:[(a, b)]
uts) [b]
projs

-- | Arguments are shared between invocations if:
--
-- * They contain _no_ references to locally-bound variables
-- * Are all equal
areShared :: InScopeSet -> [Either Term Type] -> Bool
areShared :: InScopeSet -> [Either Term Type] -> Bool
areShared InScopeSet
_       []       = Bool
True
areShared InScopeSet
inScope xs :: [Either Term Type]
xs@(Either Term Type
x:[Either Term Type]
_) = Bool
noFV1 Bool -> Bool -> Bool
&& [Either Term Type] -> Bool
forall a. Eq a => [a] -> Bool
allEqual [Either Term Type]
xs
 where
  noFV1 :: Bool
noFV1 = case Either Term Type
x of
    Right Type
ty -> All -> Bool
getAll (Getting All Type (Var Any) -> (Var Any -> All) -> Type -> All
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf ((forall b. Var b -> Bool) -> IntSet -> Getting All Type (Var Any)
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' forall b. Var b -> Bool
isLocallyBound IntSet
IntSet.empty)
                                       (All -> Var Any -> All
forall a b. a -> b -> a
const (Bool -> All
All Bool
False)) Type
ty)
    Left Term
tm  -> All -> Bool
getAll (Getting All Term (Var Any) -> (Var Any -> All) -> Term -> All
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf ((forall b. Var b -> Bool) -> Getting All Term (Var Any)
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool) -> (Var a -> f (Var a)) -> Term -> f Term
termFreeVars' forall b. Var b -> Bool
isLocallyBound)
                                       (All -> Var Any -> All
forall a b. a -> b -> a
const (Bool -> All
All Bool
False)) Term
tm)

  isLocallyBound :: Var a -> Bool
isLocallyBound Var a
v = Var a -> Bool
forall b. Var b -> Bool
isLocalId Var a
v Bool -> Bool -> Bool
&& Var a
v Var a -> InScopeSet -> Bool
forall a. Var a -> InScopeSet -> Bool
`notElemInScopeSet` InScopeSet
inScope

-- | Create a list of arguments given a map of positions to common arguments,
-- and a list of arguments
mkDJArgs :: Int -- ^ Current position
         -> [(Int,Either Term Type)] -- ^ map from position to common argument
         -> [Term] -- ^ (projections for) distinct arguments
         -> [Either Term Type]
mkDJArgs :: Int -> [(Int, Either Term Type)] -> [Term] -> [Either Term Type]
mkDJArgs Int
_ [(Int, Either Term Type)]
cms []   = ((Int, Either Term Type) -> Either Term Type)
-> [(Int, Either Term Type)] -> [Either Term Type]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Either Term Type) -> Either Term Type
forall a b. (a, b) -> b
snd [(Int, Either Term Type)]
cms
mkDJArgs Int
_ [] [Term]
uncms = (Term -> Either Term Type) -> [Term] -> [Either Term Type]
forall a b. (a -> b) -> [a] -> [b]
map Term -> Either Term Type
forall a b. a -> Either a b
Left [Term]
uncms
mkDJArgs Int
n ((Int
m,Either Term Type
x):[(Int, Either Term Type)]
cms) (Term
y:[Term]
uncms)
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
m    = Either Term Type
x       Either Term Type -> [Either Term Type] -> [Either Term Type]
forall a. a -> [a] -> [a]
: Int -> [(Int, Either Term Type)] -> [Term] -> [Either Term Type]
mkDJArgs (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [(Int, Either Term Type)]
cms (Term
yTerm -> [Term] -> [Term]
forall a. a -> [a] -> [a]
:[Term]
uncms)
  | Bool
otherwise = Term -> Either Term Type
forall a b. a -> Either a b
Left Term
y  Either Term Type -> [Either Term Type] -> [Either Term Type]
forall a. a -> [a] -> [a]
: Int -> [(Int, Either Term Type)] -> [Term] -> [Either Term Type]
mkDJArgs (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) ((Int
m,Either Term Type
x)(Int, Either Term Type)
-> [(Int, Either Term Type)] -> [(Int, Either Term Type)]
forall a. a -> [a] -> [a]
:[(Int, Either Term Type)]
cms) [Term]
uncms

-- | Create a case-expression that selects between the distinct arguments given
-- a case-tree
genCase :: TyConMap
        -> IM.IntMap TyConName
        -> Type -- ^ Type of the alternatives
        -> [Type] -- ^ Types of the arguments
        -> CaseTree [Term] -- ^ CaseTree of arguments
        -> Term
genCase :: TyConMap
-> IntMap TyConName -> Type -> [Type] -> CaseTree [Term] -> Term
genCase TyConMap
tcm IntMap TyConName
tupTcm Type
ty [Type]
argTys = CaseTree [Term] -> Term
go
  where
    go :: CaseTree [Term] -> Term
go (Leaf [Term]
tms) =
      TyConMap -> IntMap TyConName -> [(Type, Term)] -> Term
mkBigTupTm TyConMap
tcm IntMap TyConName
tupTcm ([Type] -> [Term] -> [(Type, Term)]
forall a b. HasCallStack => [a] -> [b] -> [(a, b)]
List.zipEqual [Type]
argTys [Term]
tms)

    go (LB [LetBinding]
lb CaseTree [Term]
ct) =
      [LetBinding] -> Term -> Term
Letrec [LetBinding]
lb (CaseTree [Term] -> Term
go CaseTree [Term]
ct)

    go (Branch Term
scrut [(Pat
p,CaseTree [Term]
ct)]) =
      let ct' :: Term
ct' = CaseTree [Term] -> Term
go CaseTree [Term]
ct
          ([TyVar]
ptvs,[Var Term]
pids) = Pat -> ([TyVar], [Var Term])
patIds Pat
p
      in  if ([TyVar] -> [Var Any]
coerce [TyVar]
ptvs [Var Any] -> [Var Any] -> [Var Any]
forall a. [a] -> [a] -> [a]
++ [Var Term] -> [Var Any]
coerce [Var Term]
pids) [Var Any] -> Term -> Bool
forall a. [Var a] -> Term -> Bool
`localVarsDoNotOccurIn` Term
ct'
             then Term
ct'
             else Term -> Type -> [Alt] -> Term
Case Term
scrut Type
ty [(Pat
p,Term
ct')]

    go (Branch Term
scrut [(Pat, CaseTree [Term])]
pats) =
      Term -> Type -> [Alt] -> Term
Case Term
scrut Type
ty (((Pat, CaseTree [Term]) -> Alt)
-> [(Pat, CaseTree [Term])] -> [Alt]
forall a b. (a -> b) -> [a] -> [b]
map ((CaseTree [Term] -> Term) -> (Pat, CaseTree [Term]) -> Alt
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second CaseTree [Term] -> Term
go) [(Pat, CaseTree [Term])]
pats)

-- | Lookup the TyConName and DataCon for a tuple of size n
findTup :: TyConMap -> IM.IntMap TyConName -> Int -> (TyConName,DataCon)
findTup :: TyConMap -> IntMap TyConName -> Int -> (TyConName, DataCon)
findTup TyConMap
tcm IntMap TyConName
tupTcm Int
n = (TyConName
tupTcNm,DataCon
tupDc)
  where
    tupTcNm :: TyConName
tupTcNm      = TyConName -> Maybe TyConName -> TyConName
forall a. a -> Maybe a -> a
Maybe.fromMaybe (String -> TyConName
forall a. HasCallStack => String -> a
error (String -> TyConName) -> String -> TyConName
forall a b. (a -> b) -> a -> b
$ String
$curLoc String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"Can't find " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"-tuple") (Maybe TyConName -> TyConName) -> Maybe TyConName -> TyConName
forall a b. (a -> b) -> a -> b
$ Int -> IntMap TyConName -> Maybe TyConName
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
n IntMap TyConName
tupTcm
    Just TyCon
tupTc   = TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap TyConName
tupTcNm TyConMap
tcm
    [DataCon
tupDc]      = TyCon -> [DataCon]
tyConDataCons TyCon
tupTc

mkBigTupTm :: TyConMap -> IM.IntMap TyConName -> [(Type,Term)] -> Term
mkBigTupTm :: TyConMap -> IntMap TyConName -> [(Type, Term)] -> Term
mkBigTupTm TyConMap
tcm IntMap TyConName
tupTcm [(Type, Term)]
args = (Type, Term) -> Term
forall a b. (a, b) -> b
snd ((Type, Term) -> Term) -> (Type, Term) -> Term
forall a b. (a -> b) -> a -> b
$ TyConMap -> IntMap TyConName -> [(Type, Term)] -> (Type, Term)
mkBigTup TyConMap
tcm IntMap TyConName
tupTcm [(Type, Term)]
args

mkSmallTup,mkBigTup :: TyConMap -> IM.IntMap TyConName -> [(Type,Term)] -> (Type,Term)
mkSmallTup :: TyConMap -> IntMap TyConName -> [(Type, Term)] -> (Type, Term)
mkSmallTup TyConMap
_ IntMap TyConName
_ [] = String -> (Type, Term)
forall a. HasCallStack => String -> a
error (String -> (Type, Term)) -> String -> (Type, Term)
forall a b. (a -> b) -> a -> b
$ String
$curLoc String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"mkSmallTup: Can't create 0-tuple"
mkSmallTup TyConMap
_ IntMap TyConName
_ [(Type
ty,Term
tm)] = (Type
ty,Term
tm)
mkSmallTup TyConMap
tcm IntMap TyConName
tupTcm [(Type, Term)]
args = (Type
ty,Term
tm)
  where
    ([Type]
argTys,[Term]
tms) = [(Type, Term)] -> ([Type], [Term])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Type, Term)]
args
    (TyConName
tupTcNm,DataCon
tupDc) = TyConMap -> IntMap TyConName -> Int -> (TyConName, DataCon)
findTup TyConMap
tcm IntMap TyConName
tupTcm ([(Type, Term)] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [(Type, Term)]
args)
    tm :: Term
tm = Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
tupDc) ((Type -> Either Term Type) -> [Type] -> [Either Term Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Either Term Type
forall a b. b -> Either a b
Right [Type]
argTys [Either Term Type] -> [Either Term Type] -> [Either Term Type]
forall a. [a] -> [a] -> [a]
++ (Term -> Either Term Type) -> [Term] -> [Either Term Type]
forall a b. (a -> b) -> [a] -> [b]
map Term -> Either Term Type
forall a b. a -> Either a b
Left [Term]
tms)
    ty :: Type
ty = TyConName -> [Type] -> Type
mkTyConApp TyConName
tupTcNm [Type]
argTys

mkBigTup :: TyConMap -> IntMap TyConName -> [(Type, Term)] -> (Type, Term)
mkBigTup TyConMap
tcm IntMap TyConName
tupTcm = ([(Type, Term)] -> (Type, Term)) -> [(Type, Term)] -> (Type, Term)
forall a. ([a] -> a) -> [a] -> a
mkChunkified (TyConMap -> IntMap TyConName -> [(Type, Term)] -> (Type, Term)
mkSmallTup TyConMap
tcm IntMap TyConName
tupTcm)

mkSmallTupTy,mkBigTupTy
  :: TyConMap
  -> IM.IntMap TyConName
  -> [Type]
  -> Type
mkSmallTupTy :: TyConMap -> IntMap TyConName -> [Type] -> Type
mkSmallTupTy TyConMap
_ IntMap TyConName
_ [] = String -> Type
forall a. HasCallStack => String -> a
error (String -> Type) -> String -> Type
forall a b. (a -> b) -> a -> b
$ String
$curLoc String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"mkSmallTupTy: Can't create 0-tuple"
mkSmallTupTy TyConMap
_ IntMap TyConName
_ [Type
ty] = Type
ty
mkSmallTupTy TyConMap
tcm IntMap TyConName
tupTcm [Type]
tys = TyConName -> [Type] -> Type
mkTyConApp TyConName
tupTcNm [Type]
tys
  where
    m :: Int
m = [Type] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Type]
tys
    (TyConName
tupTcNm,DataCon
_) = TyConMap -> IntMap TyConName -> Int -> (TyConName, DataCon)
findTup TyConMap
tcm IntMap TyConName
tupTcm Int
m

mkBigTupTy :: TyConMap -> IntMap TyConName -> [Type] -> Type
mkBigTupTy TyConMap
tcm IntMap TyConName
tupTcm = ([Type] -> Type) -> [Type] -> Type
forall a. ([a] -> a) -> [a] -> a
mkChunkified (TyConMap -> IntMap TyConName -> [Type] -> Type
mkSmallTupTy TyConMap
tcm IntMap TyConName
tupTcm)

mkSmallTupSelector,mkBigTupSelector
  :: MonadUnique m
  => InScopeSet
  -> TyConMap
  -> IM.IntMap TyConName
  -> Term
  -> [Type]
  -> Int
  -> m Term
mkSmallTupSelector :: InScopeSet
-> TyConMap -> IntMap TyConName -> Term -> [Type] -> Int -> m Term
mkSmallTupSelector InScopeSet
_ TyConMap
_ IntMap TyConName
_ Term
scrut [Type
_] Int
0 = Term -> m Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
scrut
mkSmallTupSelector InScopeSet
_ TyConMap
_ IntMap TyConName
_ Term
_     [Type
_] Int
n = String -> m Term
forall a. HasCallStack => String -> a
error (String -> m Term) -> String -> m Term
forall a b. (a -> b) -> a -> b
$ String
$curLoc String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"mkSmallTupSelector called with one type, but to select " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n
mkSmallTupSelector InScopeSet
inScope TyConMap
tcm IntMap TyConName
_ Term
scrut [Type]
_ Int
n = String -> InScopeSet -> TyConMap -> Term -> Int -> Int -> m Term
forall (m :: Type -> Type).
(HasCallStack, MonadUnique m) =>
String -> InScopeSet -> TyConMap -> Term -> Int -> Int -> m Term
mkSelectorCase (String
$curLoc String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"mkSmallTupSelector") InScopeSet
inScope TyConMap
tcm Term
scrut Int
1 Int
n

mkBigTupSelector :: InScopeSet
-> TyConMap -> IntMap TyConName -> Term -> [Type] -> Int -> m Term
mkBigTupSelector InScopeSet
inScope TyConMap
tcm IntMap TyConName
tupTcm Term
scrut [Type]
tys Int
n = [[Type]] -> m Term
forall (m :: Type -> Type). MonadUnique m => [[Type]] -> m Term
go ([Type] -> [[Type]]
forall a. [a] -> [[a]]
chunkify [Type]
tys)
  where
    go :: [[Type]] -> m Term
go [[Type]
_] = InScopeSet
-> TyConMap -> IntMap TyConName -> Term -> [Type] -> Int -> m Term
forall (m :: Type -> Type).
MonadUnique m =>
InScopeSet
-> TyConMap -> IntMap TyConName -> Term -> [Type] -> Int -> m Term
mkSmallTupSelector InScopeSet
inScope TyConMap
tcm IntMap TyConName
tupTcm Term
scrut [Type]
tys Int
n
    go [[Type]]
tyss = do
      let (Int
nOuter,Int
nInner) = Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
divMod Int
n Int
mAX_TUPLE_SIZE
          tyss' :: [Type]
tyss' = ([Type] -> Type) -> [[Type]] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TyConMap -> IntMap TyConName -> [Type] -> Type
mkSmallTupTy TyConMap
tcm IntMap TyConName
tupTcm) [[Type]]
tyss
      Term
outer <- InScopeSet
-> TyConMap -> IntMap TyConName -> Term -> [Type] -> Int -> m Term
forall (m :: Type -> Type).
MonadUnique m =>
InScopeSet
-> TyConMap -> IntMap TyConName -> Term -> [Type] -> Int -> m Term
mkSmallTupSelector InScopeSet
inScope TyConMap
tcm IntMap TyConName
tupTcm Term
scrut [Type]
tyss' Int
nOuter
      Term
inner <- InScopeSet
-> TyConMap -> IntMap TyConName -> Term -> [Type] -> Int -> m Term
forall (m :: Type -> Type).
MonadUnique m =>
InScopeSet
-> TyConMap -> IntMap TyConName -> Term -> [Type] -> Int -> m Term
mkSmallTupSelector InScopeSet
inScope TyConMap
tcm IntMap TyConName
tupTcm Term
outer ([[Type]]
tyss [[Type]] -> Int -> [Type]
forall a. [a] -> Int -> a
List.!! Int
nOuter) Int
nInner
      Term -> m Term
forall (m :: Type -> Type) a. Monad m => a -> m a
return Term
inner


-- | Determine if a term in a function position is interesting to lift out of
-- of a case-expression.
--
-- This holds for all global functions, and certain primitives. Currently those
-- primitives are:
--
-- * All non-power-of-two multiplications
-- * All division-like operations with a non-power-of-two divisor
interestingToLift
  :: InScopeSet
  -- ^ in scope
  -> (Term -> Term)
  -- ^ Evaluator
  -> Term
  -- ^ Term in function position
  -> [Either Term Type]
  -- ^ Arguments
  -> [TickInfo]
  -- ^ Tick annoations
  -> RewriteMonad extra (Maybe Term)
interestingToLift :: InScopeSet
-> (Term -> Term)
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad extra (Maybe Term)
interestingToLift InScopeSet
inScope Term -> Term
_ e :: Term
e@(Var Var Term
v) [Either Term Type]
_ [TickInfo]
ticks =
  if TickInfo
NoDeDup TickInfo -> [TickInfo] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`notElem` [TickInfo]
ticks Bool -> Bool -> Bool
&& (Var Term -> Bool
forall b. Var b -> Bool
isGlobalId Var Term
v Bool -> Bool -> Bool
||  Var Term
v Var Term -> InScopeSet -> Bool
forall a. Var a -> InScopeSet -> Bool
`elemInScopeSet` InScopeSet
inScope)
     then Maybe Term -> RewriteMonad extra (Maybe Term)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Maybe Term
forall a. a -> Maybe a
Just Term
e)
     else Maybe Term -> RewriteMonad extra (Maybe Term)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Maybe Term
forall a. Maybe a
Nothing
interestingToLift InScopeSet
inScope Term -> Term
eval e :: Term
e@(Prim PrimInfo
pInfo) [Either Term Type]
args [TickInfo]
ticks
  | TickInfo
NoDeDup TickInfo -> [TickInfo] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`notElem` [TickInfo]
ticks = do
  let anyArgNotConstant :: Bool
anyArgNotConstant = (Term -> Bool) -> [Term] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any (Bool -> Bool
not (Bool -> Bool) -> (Term -> Bool) -> Term -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> Bool
isConstant) [Term]
lArgs
  case OccName -> [(OccName, Bool)] -> Maybe Bool
forall a b. Eq a => a -> [(a, b)] -> Maybe b
List.lookup (PrimInfo -> OccName
primName PrimInfo
pInfo) [(OccName, Bool)]
interestingPrims of
    Just Bool
t | Bool
t Bool -> Bool -> Bool
|| Bool
anyArgNotConstant -> Maybe Term -> RewriteMonad extra (Maybe Term)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Maybe Term
forall a. a -> Maybe a
Just Term
e)
    Maybe Bool
_ | TickInfo
DeDup TickInfo -> [TickInfo] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`elem` [TickInfo]
ticks -> Maybe Term -> RewriteMonad extra (Maybe Term)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Maybe Term
forall a. a -> Maybe a
Just Term
e)
    Maybe Bool
_ -> do
      let isInteresting :: Term -> RewriteMonad extra (Maybe Term)
isInteresting = (\(Term
x, [Either Term Type]
y, [TickInfo]
z) -> InScopeSet
-> (Term -> Term)
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad extra (Maybe Term)
forall extra.
InScopeSet
-> (Term -> Term)
-> Term
-> [Either Term Type]
-> [TickInfo]
-> RewriteMonad extra (Maybe Term)
interestingToLift InScopeSet
inScope Term -> Term
eval Term
x [Either Term Type]
y [TickInfo]
z) ((Term, [Either Term Type], [TickInfo])
 -> RewriteMonad extra (Maybe Term))
-> (Term -> (Term, [Either Term Type], [TickInfo]))
-> Term
-> RewriteMonad extra (Maybe Term)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> (Term, [Either Term Type], [TickInfo])
collectArgsTicks
      if Type -> Bool
isHOTy (PrimInfo -> Type
primType PrimInfo
pInfo) then do
        Bool
anyInteresting <- (Term -> RewriteMonad extra Bool)
-> [Term] -> RewriteMonad extra Bool
forall (m :: Type -> Type) a.
Monad m =>
(a -> m Bool) -> [a] -> m Bool
List.anyM ((Maybe Term -> Bool)
-> RewriteMonad extra (Maybe Term) -> RewriteMonad extra Bool
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Maybe Term -> Bool
forall a. Maybe a -> Bool
Maybe.isJust (RewriteMonad extra (Maybe Term) -> RewriteMonad extra Bool)
-> (Term -> RewriteMonad extra (Maybe Term))
-> Term
-> RewriteMonad extra Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> RewriteMonad extra (Maybe Term)
forall extra. Term -> RewriteMonad extra (Maybe Term)
isInteresting) [Term]
lArgs
        if Bool
anyInteresting then Maybe Term -> RewriteMonad extra (Maybe Term)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Maybe Term
forall a. a -> Maybe a
Just Term
e) else Maybe Term -> RewriteMonad extra (Maybe Term)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Maybe Term
forall a. Maybe a
Nothing
      else
        Maybe Term -> RewriteMonad extra (Maybe Term)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Maybe Term
forall a. Maybe a
Nothing

  where
    interestingPrims :: [(OccName, Bool)]
interestingPrims =
      [(OccName
"Clash.Sized.Internal.BitVector.*#",Bool
tailNonPow2)
      ,(OccName
"Clash.Sized.Internal.BitVector.times#",Bool
tailNonPow2)
      ,(OccName
"Clash.Sized.Internal.BitVector.quot#",Bool
lastNotPow2)
      ,(OccName
"Clash.Sized.Internal.BitVector.rem#",Bool
lastNotPow2)
      ,(OccName
"Clash.Sized.Internal.Index.*#",Bool
tailNonPow2)
      ,(OccName
"Clash.Sized.Internal.Index.quot#",Bool
lastNotPow2)
      ,(OccName
"Clash.Sized.Internal.Index.rem#",Bool
lastNotPow2)
      ,(OccName
"Clash.Sized.Internal.Signed.*#",Bool
tailNonPow2)
      ,(OccName
"Clash.Sized.Internal.Signed.times#",Bool
tailNonPow2)
      ,(OccName
"Clash.Sized.Internal.Signed.rem#",Bool
lastNotPow2)
      ,(OccName
"Clash.Sized.Internal.Signed.quot#",Bool
lastNotPow2)
      ,(OccName
"Clash.Sized.Internal.Signed.div#",Bool
lastNotPow2)
      ,(OccName
"Clash.Sized.Internal.Signed.mod#",Bool
lastNotPow2)
      ,(OccName
"Clash.Sized.Internal.Unsigned.*#",Bool
tailNonPow2)
      ,(OccName
"Clash.Sized.Internal.Unsigned.times#",Bool
tailNonPow2)
      ,(OccName
"Clash.Sized.Internal.Unsigned.quot#",Bool
lastNotPow2)
      ,(OccName
"Clash.Sized.Internal.Unsigned.rem#",Bool
lastNotPow2)
      ,(OccName
"GHC.Base.quotInt",Bool
lastNotPow2)
      ,(OccName
"GHC.Base.remInt",Bool
lastNotPow2)
      ,(OccName
"GHC.Base.divInt",Bool
lastNotPow2)
      ,(OccName
"GHC.Base.modInt",Bool
lastNotPow2)
      ,(OccName
"GHC.Classes.divInt#",Bool
lastNotPow2)
      ,(OccName
"GHC.Classes.modInt#",Bool
lastNotPow2)
#if MIN_VERSION_base(4,15,0)
      ,("GHC.Num.Integer.integerMul",allNonPow2)
      ,("GHC.Num.Integer.integerDiv",lastNotPow2)
      ,("GHC.Num.Integer.integerMod",lastNotPow2)
      ,("GHC.Num.Integer.integerQuot",lastNotPow2)
      ,("GHC.Num.Integer.integerRem",lastNotPow2)
#else
      ,(OccName
"GHC.Integer.Type.timesInteger",Bool
allNonPow2)
      ,(OccName
"GHC.Integer.Type.divInteger",Bool
lastNotPow2)
      ,(OccName
"GHC.Integer.Type.modInteger",Bool
lastNotPow2)
      ,(OccName
"GHC.Integer.Type.quotInteger",Bool
lastNotPow2)
      ,(OccName
"GHC.Integer.Type.remInteger",Bool
lastNotPow2)
#endif
      ,(OccName
"GHC.Prim.*#",Bool
allNonPow2)
      ,(OccName
"GHC.Prim.quotInt#",Bool
lastNotPow2)
      ,(OccName
"GHC.Prim.remInt#",Bool
lastNotPow2)
      ]

    lArgs :: [Term]
lArgs       = [Either Term Type] -> [Term]
forall a b. [Either a b] -> [a]
Either.lefts [Either Term Type]
args

    allNonPow2 :: Bool
allNonPow2  = (Term -> Bool) -> [Term] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
all (Bool -> Bool
not (Bool -> Bool) -> (Term -> Bool) -> Term -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> Bool
termIsPow2) [Term]
lArgs
    tailNonPow2 :: Bool
tailNonPow2 = case [Term]
lArgs of
                    [] -> Bool
True
                    [Term]
_  -> (Term -> Bool) -> [Term] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
all (Bool -> Bool
not (Bool -> Bool) -> (Term -> Bool) -> Term -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> Bool
termIsPow2) ([Term] -> [Term]
forall a. [a] -> [a]
tail [Term]
lArgs)
    lastNotPow2 :: Bool
lastNotPow2 = case [Term]
lArgs of
                    [] -> Bool
True
                    [Term]
_  -> Bool -> Bool
not (Term -> Bool
termIsPow2 ([Term] -> Term
forall a. [a] -> a
last [Term]
lArgs))

    termIsPow2 :: Term -> Bool
termIsPow2 Term
e' = case Term -> Term
eval Term
e' of
      Literal (IntegerLiteral Integer
n) -> Integer -> Bool
forall a. (Num a, Bits a) => a -> Bool
isPow2 Integer
n
      Term
a -> case Term -> (Term, [Either Term Type])
collectArgs Term
a of
        (Prim PrimInfo
p,[Right Type
_,Left Term
_,Left (Literal (IntegerLiteral Integer
n))])
          | OccName -> Bool
forall a. (Eq a, IsString a) => a -> Bool
isFromInteger (PrimInfo -> OccName
primName PrimInfo
p) -> Integer -> Bool
forall a. (Num a, Bits a) => a -> Bool
isPow2 Integer
n
        (Prim PrimInfo
p,[Right Type
_,Left Term
_,Left Term
_,Left (Literal (IntegerLiteral Integer
n))])
          | PrimInfo -> OccName
primName PrimInfo
p OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Internal.BitVector.fromInteger#"  -> Integer -> Bool
forall a. (Num a, Bits a) => a -> Bool
isPow2 Integer
n
        (Prim PrimInfo
p,[Right Type
_,       Left Term
_,Left (Literal (IntegerLiteral Integer
n))])
          | PrimInfo -> OccName
primName PrimInfo
p OccName -> OccName -> Bool
forall a. Eq a => a -> a -> Bool
== OccName
"Clash.Sized.Internal.BitVector.fromInteger##" -> Integer -> Bool
forall a. (Num a, Bits a) => a -> Bool
isPow2 Integer
n

        (Term, [Either Term Type])
_ -> Bool
False

    isPow2 :: a -> Bool
isPow2 a
x = a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0 Bool -> Bool -> Bool
&& (a
x a -> a -> a
forall a. Bits a => a -> a -> a
.&. (a -> a
forall a. Bits a => a -> a
complement a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
1)) a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x

    isFromInteger :: a -> Bool
isFromInteger a
x = a
x a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`elem` [a
"Clash.Sized.Internal.BitVector.fromInteger#"
                               ,a
"Clash.Sized.Integer.Index.fromInteger"
                               ,a
"Clash.Sized.Internal.Signed.fromInteger#"
                               ,a
"Clash.Sized.Internal.Unsigned.fromInteger#"
                               ]

    isHOTy :: Type -> Bool
isHOTy Type
t = case Type -> ([Either TyVar Type], Type)
splitFunForallTy Type
t of
      ([Either TyVar Type]
args',Type
_) -> (Type -> Bool) -> [Type] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any Type -> Bool
isPolyFunTy ([Either TyVar Type] -> [Type]
forall a b. [Either a b] -> [b]
Either.rights [Either TyVar Type]
args')

interestingToLift InScopeSet
_ Term -> Term
_ Term
_ [Either Term Type]
_ [TickInfo]
_ = Maybe Term -> RewriteMonad extra (Maybe Term)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Maybe Term
forall a. Maybe a
Nothing