{-# LANGUAGE FlexibleContexts #-}

module Gibbon.Passes.Simplifier
  ( simplifyL1, simplifyLocBinds, lateInlineTriv )
  where

import Data.Functor.Foldable as Foldable
import qualified Data.Set as S
import qualified Data.Map as M
import qualified Data.List as L ( isPrefixOf )

import Gibbon.Common
import Gibbon.Language
import Gibbon.L1.Syntax
import Gibbon.L2.Syntax
import qualified Gibbon.L4.Syntax as L4
import Gibbon.Passes.Freshen (freshNames1, freshFun1)

--------------------------------------------------------------------------------

markRecFns :: Prog1 -> PassM Prog1
markRecFns :: Prog1 -> PassM Prog1
markRecFns (Prog DDefs (TyOf Exp1)
ddefs FunDefs Exp1
fundefs Maybe (Exp1, TyOf Exp1)
main) = do
    let fundefs' :: FunDefs Exp1
fundefs' = (FunDef Exp1 -> FunDef Exp1) -> FunDefs Exp1 -> FunDefs Exp1
forall a b k. (a -> b) -> Map k a -> Map k b
M.map
                     (\fn :: FunDef Exp1
fn@FunDef{Var
funName :: Var
funName :: forall ex. FunDef ex -> Var
funName,Exp1
funBody :: Exp1
funBody :: forall ex. FunDef ex -> ex
funBody,FunMeta
funMeta :: FunMeta
funMeta :: forall ex. FunDef ex -> FunMeta
funMeta} ->
                          if Var
funName Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` (Exp1 -> Set Var
forall a. FreeVars a => a -> Set Var
gFreeVars Exp1
funBody)
                          then FunDef Exp1
fn { funMeta :: FunMeta
funMeta = FunMeta
funMeta { funRec :: FunRec
funRec = FunRec
Rec } }
                          else FunDef Exp1
fn)
                     FunDefs Exp1
fundefs
    Prog1 -> PassM Prog1
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DDefs (TyOf Exp1)
-> FunDefs Exp1 -> Maybe (Exp1, TyOf Exp1) -> Prog1
forall ex.
DDefs (TyOf ex) -> FunDefs ex -> Maybe (ex, TyOf ex) -> Prog ex
Prog DDefs (TyOf Exp1)
ddefs FunDefs Exp1
fundefs' Maybe (Exp1, TyOf Exp1)
main)

inlineFuns :: Prog1 -> PassM Prog1
inlineFuns :: Prog1 -> PassM Prog1
inlineFuns (Prog DDefs (TyOf Exp1)
ddefs FunDefs Exp1
fundefs Maybe (Exp1, TyOf Exp1)
main) = do
    Maybe (Exp1, UrTy ())
main' <- case Maybe (Exp1, TyOf Exp1)
main of
               Maybe (Exp1, TyOf Exp1)
Nothing -> Maybe (Exp1, UrTy ()) -> PassM (Maybe (Exp1, UrTy ()))
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Exp1, UrTy ())
forall a. Maybe a
Nothing
               Just (Exp1
e,TyOf Exp1
ty) -> do
                 Exp1
e' <- ((Base Exp1 Exp1 -> PassM Exp1) -> Exp1 -> PassM Exp1
forall (m :: * -> *) t a.
(Monad m, Traversable (Base t), Recursive t) =>
(Base t a -> m a) -> t -> m a
cataM Base Exp1 Exp1 -> PassM Exp1
PreExpF E1Ext () (UrTy ()) Exp1 -> PassM Exp1
go) Exp1
e
                 Maybe (Exp1, UrTy ()) -> PassM (Maybe (Exp1, UrTy ()))
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Exp1, UrTy ()) -> PassM (Maybe (Exp1, UrTy ())))
-> Maybe (Exp1, UrTy ()) -> PassM (Maybe (Exp1, UrTy ()))
forall a b. (a -> b) -> a -> b
$ (Exp1, UrTy ()) -> Maybe (Exp1, UrTy ())
forall a. a -> Maybe a
Just (Exp1
e', TyOf Exp1
UrTy ()
ty)
    FunDefs Exp1
fundefs' <- (FunDef Exp1 -> PassM (FunDef Exp1))
-> FunDefs Exp1 -> PassM (FunDefs Exp1)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Map Var a -> m (Map Var b)
mapM (\FunDef Exp1
fn -> do
                           Exp1
bod <- ((Base Exp1 Exp1 -> PassM Exp1) -> Exp1 -> PassM Exp1
forall (m :: * -> *) t a.
(Monad m, Traversable (Base t), Recursive t) =>
(Base t a -> m a) -> t -> m a
cataM Base Exp1 Exp1 -> PassM Exp1
PreExpF E1Ext () (UrTy ()) Exp1 -> PassM Exp1
go) (FunDef Exp1 -> Exp1
forall ex. FunDef ex -> ex
funBody FunDef Exp1
fn)
                           FunDef Exp1 -> PassM (FunDef Exp1)
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FunDef Exp1 -> PassM (FunDef Exp1))
-> FunDef Exp1 -> PassM (FunDef Exp1)
forall a b. (a -> b) -> a -> b
$ FunDef Exp1
fn { funBody :: Exp1
funBody = Exp1
bod})
                     FunDefs Exp1
fundefs
    Prog1 -> PassM Prog1
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DDefs (TyOf Exp1)
-> FunDefs Exp1 -> Maybe (Exp1, TyOf Exp1) -> Prog1
forall ex.
DDefs (TyOf ex) -> FunDefs ex -> Maybe (ex, TyOf ex) -> Prog ex
Prog DDefs (TyOf Exp1)
ddefs FunDefs Exp1
fundefs' Maybe (Exp1, TyOf Exp1)
Maybe (Exp1, UrTy ())
main')
  where
    go :: PreExpF E1Ext () (UrTy ()) Exp1 -> PassM Exp1
    go :: PreExpF E1Ext () (UrTy ()) Exp1 -> PassM Exp1
go PreExpF E1Ext () (UrTy ()) Exp1
ex =
      case PreExpF E1Ext () (UrTy ()) Exp1
ex of
        AppEF Var
f [] [Exp1]
args -> do
            let fn :: FunDef Exp1
fn = FunDefs Exp1
fundefs FunDefs Exp1 -> Var -> FunDef Exp1
forall k a. Ord k => Map k a -> k -> a
M.! Var
f
            if FunMeta -> FunInline
funInline (FunDef Exp1 -> FunMeta
forall ex. FunDef ex -> FunMeta
funMeta FunDef Exp1
fn) FunInline -> FunInline -> Bool
forall a. Eq a => a -> a -> Bool
== FunInline
Inline Bool -> Bool -> Bool
&& FunMeta -> FunRec
funRec (FunDef Exp1 -> FunMeta
forall ex. FunDef ex -> FunMeta
funMeta FunDef Exp1
fn) FunRec -> FunRec -> Bool
forall a. Eq a => a -> a -> Bool
== FunRec
NotRec
              then do
                FunDef{[Var]
funArgs :: [Var]
funArgs :: forall ex. FunDef ex -> [Var]
funArgs,ArrowTy (TyOf Exp1)
funTy :: ArrowTy (TyOf Exp1)
funTy :: forall ex. FunDef ex -> ArrowTy (TyOf ex)
funTy,Exp1
funBody :: forall ex. FunDef ex -> ex
funBody :: Exp1
funBody} <- FunDef Exp1 -> PassM (FunDef Exp1)
freshFun1 FunDef Exp1
fn
                let in_tys :: [UrTy ()]
in_tys = ([UrTy ()], UrTy ()) -> [UrTy ()]
forall a b. (a, b) -> a
fst ([UrTy ()], UrTy ())
ArrowTy (TyOf Exp1)
funTy
                    binds :: [(Var, [()], UrTy (), Exp1)]
binds = ((Var, UrTy (), Exp1) -> (Var, [()], UrTy (), Exp1))
-> [(Var, UrTy (), Exp1)] -> [(Var, [()], UrTy (), Exp1)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Var
v,UrTy ()
ty,Exp1
e) -> (Var
v,[],UrTy ()
ty,Exp1
e)) ([Var] -> [UrTy ()] -> [Exp1] -> [(Var, UrTy (), Exp1)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Var]
funArgs [UrTy ()]
in_tys [Exp1]
args)
                Exp1 -> PassM Exp1
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp1 -> PassM Exp1) -> Exp1 -> PassM Exp1
forall a b. (a -> b) -> a -> b
$ [(Var, [()], UrTy (), Exp1)] -> Exp1 -> Exp1
forall loc dec (ext :: * -> * -> *).
[(Var, [loc], dec, PreExp ext loc dec)]
-> PreExp ext loc dec -> PreExp ext loc dec
mkLets [(Var, [()], UrTy (), Exp1)]
binds Exp1
funBody
              else do
                [Exp1]
args' <- (Exp1 -> PassM Exp1) -> [Exp1] -> PassM [Exp1]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (PreExpF E1Ext () (UrTy ()) Exp1 -> PassM Exp1
go (PreExpF E1Ext () (UrTy ()) Exp1 -> PassM Exp1)
-> (Exp1 -> PreExpF E1Ext () (UrTy ()) Exp1) -> Exp1 -> PassM Exp1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp1 -> Base Exp1 Exp1
Exp1 -> PreExpF E1Ext () (UrTy ()) Exp1
forall t. Recursive t => t -> Base t t
project) [Exp1]
args
                Exp1 -> PassM Exp1
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp1 -> PassM Exp1) -> Exp1 -> PassM Exp1
forall a b. (a -> b) -> a -> b
$ Var -> [()] -> [Exp1] -> Exp1
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f [] [Exp1]
args'
        PreExpF E1Ext () (UrTy ()) Exp1
_ -> Exp1 -> PassM Exp1
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp1 -> PassM Exp1) -> Exp1 -> PassM Exp1
forall a b. (a -> b) -> a -> b
$ Base Exp1 Exp1 -> Exp1
forall t. Corecursive t => Base t t -> t
embed Base Exp1 Exp1
PreExpF E1Ext () (UrTy ()) Exp1
ex

deadFunElim :: Prog1 -> PassM Prog1
deadFunElim :: Prog1 -> PassM Prog1
deadFunElim (Prog DDefs (TyOf Exp1)
ddefs FunDefs Exp1
fundefs Maybe (Exp1, TyOf Exp1)
main) = do
    let used :: Set Var
used = case Maybe (Exp1, TyOf Exp1)
main of
                 Maybe (Exp1, TyOf Exp1)
Nothing -> Set Var
forall a. Set a
S.empty
                 Just (Exp1
e,TyOf Exp1
_) -> Exp1 -> Set Var
forall a. FreeVars a => a -> Set Var
gFreeVars Exp1
e
    let used' :: Set Var
used' = Set Var -> Set Var -> Set Var -> Set Var
getUsedFns Set Var
used Set Var
forall a. Set a
S.empty Set Var
used
    let (FunDefs Exp1
fundefs',[Var]
deleted) =
            (FunDef Exp1 -> (FunDefs Exp1, [Var]) -> (FunDefs Exp1, [Var]))
-> (FunDefs Exp1, [Var]) -> FunDefs Exp1 -> (FunDefs Exp1, [Var])
forall a b k. (a -> b -> b) -> b -> Map k a -> b
M.foldr (\FunDef Exp1
fn (FunDefs Exp1
acc1,[Var]
acc2) ->
                         let f :: Var
f = FunDef Exp1 -> Var
forall ex. FunDef ex -> Var
funName FunDef Exp1
fn in
                           if Var
f Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Var
used' Bool -> Bool -> Bool
|| [Char] -> [Char] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
L.isPrefixOf [Char]
"_" (Var -> [Char]
fromVar Var
f)
                           then (Var -> FunDef Exp1 -> FunDefs Exp1 -> FunDefs Exp1
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
f FunDef Exp1
fn FunDefs Exp1
acc1, [Var]
acc2)
                           else (FunDefs Exp1
acc1, Var
fVar -> [Var] -> [Var]
forall a. a -> [a] -> [a]
:[Var]
acc2))
                    (FunDefs Exp1
forall k a. Map k a
M.empty,[])
                    FunDefs Exp1
fundefs
    Int -> [Char] -> PassM Prog1 -> PassM Prog1
forall a. Int -> [Char] -> a -> a
dbgTrace Int
3 ([Char]
"Removed unused functions: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Var] -> [Char]
forall a. Show a => a -> [Char]
show [Var]
deleted) (PassM Prog1 -> PassM Prog1) -> PassM Prog1 -> PassM Prog1
forall a b. (a -> b) -> a -> b
$
      Prog1 -> PassM Prog1
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DDefs (TyOf Exp1)
-> FunDefs Exp1 -> Maybe (Exp1, TyOf Exp1) -> Prog1
forall ex.
DDefs (TyOf ex) -> FunDefs ex -> Maybe (ex, TyOf ex) -> Prog ex
Prog DDefs (TyOf Exp1)
ddefs FunDefs Exp1
fundefs' Maybe (Exp1, TyOf Exp1)
main)
  where
    getUsedFns :: S.Set Var -> S.Set Var -> S.Set Var -> S.Set Var
    getUsedFns :: Set Var -> Set Var -> Set Var -> Set Var
getUsedFns Set Var
todo Set Var
inspected Set Var
acc =
      if Set Var -> Bool
forall a. Set a -> Bool
S.null Set Var
todo
      then Set Var
acc
      else
        let f :: Var
f = Int -> Set Var -> Var
forall a. Int -> Set a -> a
S.elemAt Int
0 Set Var
todo
            todo' :: Set Var
todo' = Int -> Set Var -> Set Var
forall a. Int -> Set a -> Set a
S.deleteAt Int
0 Set Var
todo
        in if Var
f Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Var
inspected
           then Set Var -> Set Var -> Set Var -> Set Var
getUsedFns Set Var
todo' Set Var
inspected Set Var
acc
           else
             let FunDef{[Var]
funArgs :: forall ex. FunDef ex -> [Var]
funArgs :: [Var]
funArgs, Var
funName :: forall ex. FunDef ex -> Var
funName :: Var
funName,Exp1
funBody :: forall ex. FunDef ex -> ex
funBody :: Exp1
funBody} = FunDefs Exp1
fundefs FunDefs Exp1 -> Var -> FunDef Exp1
forall k a. Ord k => Map k a -> k -> a
M.! Var
f
                 free :: Set Var
free = (Exp1 -> Set Var
forall a. FreeVars a => a -> Set Var
gFreeVars Exp1
funBody) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` ([Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList (Var
funName Var -> [Var] -> [Var]
forall a. a -> [a] -> [a]
: [Var]
funArgs))
             in Set Var -> Set Var -> Set Var -> Set Var
getUsedFns (Set Var
free Set Var -> Set Var -> Set Var
forall a. Semigroup a => a -> a -> a
<> Set Var
todo') (Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
f Set Var
inspected) (Set Var
acc Set Var -> Set Var -> Set Var
forall a. Semigroup a => a -> a -> a
<> Set Var
free)

simplifyL1 :: Prog1 -> PassM Prog1
simplifyL1 :: Prog1 -> PassM Prog1
simplifyL1 Prog1
p0 = do
    Prog1
p0' <- Prog1 -> PassM Prog1
freshNames1 Prog1
p0
    Prog1
p1 <- Prog1 -> PassM Prog1
markRecFns Prog1
p0'
    Prog1
p2 <- Prog1 -> PassM Prog1
inlineFuns Prog1
p1
    Prog1
p3 <- Prog1 -> PassM Prog1
deadFunElim Prog1
p2
    Prog1 -> PassM Prog1
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Prog1
p3

--------------------------------------------------------------------------------

simplifyLocBinds :: Bool -> Prog2 -> PassM Prog2
simplifyLocBinds :: Bool -> Prog2 -> PassM Prog2
simplifyLocBinds Bool
only_cse (Prog DDefs (TyOf Exp2)
ddefs FunDefs Exp2
fundefs Maybe (Exp2, TyOf Exp2)
mainExp) = do
    let fundefs' :: FunDefs Exp2
fundefs' = (FunDef Exp2 -> FunDef Exp2) -> FunDefs Exp2 -> FunDefs Exp2
forall a b k. (a -> b) -> Map k a -> Map k b
M.map FunDef Exp2 -> FunDef Exp2
gofun FunDefs Exp2
fundefs
    let mainExp' :: Maybe (Exp2, UrTy Var)
mainExp' = case Maybe (Exp2, TyOf Exp2)
mainExp of
                     Just (Exp2
e,TyOf Exp2
ty) -> (Exp2, UrTy Var) -> Maybe (Exp2, UrTy Var)
forall a. a -> Maybe a
Just (Exp2 -> Exp2
simpl Exp2
e, TyOf Exp2
UrTy Var
ty)
                     Maybe (Exp2, TyOf Exp2)
Nothing     -> Maybe (Exp2, UrTy Var)
forall a. Maybe a
Nothing
    Prog2 -> PassM Prog2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Prog2 -> PassM Prog2) -> Prog2 -> PassM Prog2
forall a b. (a -> b) -> a -> b
$ DDefs (TyOf Exp2)
-> FunDefs Exp2 -> Maybe (Exp2, TyOf Exp2) -> Prog2
forall ex.
DDefs (TyOf ex) -> FunDefs ex -> Maybe (ex, TyOf ex) -> Prog ex
Prog DDefs (TyOf Exp2)
ddefs FunDefs Exp2
fundefs' Maybe (Exp2, TyOf Exp2)
Maybe (Exp2, UrTy Var)
mainExp'

  where
    simpl :: Exp2 -> Exp2
simpl = if Bool
only_cse
              then Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
forall k a. Map k a
M.empty Map Var Var
forall k a. Map k a
M.empty
              else Exp2 -> Exp2
go2 (Exp2 -> Exp2) -> (Exp2 -> Exp2) -> Exp2 -> Exp2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
forall k a. Map k a
M.empty (Exp2 -> Exp2) -> (Exp2 -> Exp2) -> Exp2 -> Exp2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
forall k a. Map k a
M.empty Map Var Var
forall k a. Map k a
M.empty

    gofun :: FunDef Exp2 -> FunDef Exp2
gofun f :: FunDef Exp2
f@FunDef{Exp2
funBody :: forall ex. FunDef ex -> ex
funBody :: Exp2
funBody} =
        let funBody' :: Exp2
funBody' = Exp2 -> Exp2
simpl Exp2
funBody
        in FunDef Exp2
f { funBody :: Exp2
funBody = Exp2
funBody' }

    -- partially evaluate location arithmetic
    go :: M.Map LocVar (LocVar,Int) -> Exp2 -> Exp2
    go :: Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
ex =
      case Exp2
ex of
        AppE Var
f [Var]
locs [Exp2]
args -> Var -> [Var] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f [Var]
locs ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env) [Exp2]
args)
        PrimAppE Prim (UrTy Var)
p [Exp2]
args -> Prim (UrTy Var) -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim (UrTy Var)
p ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env) [Exp2]
args)
        LetE (Var
v,[Var]
locs,UrTy Var
ty,Exp2
rhs) Exp2
bod -> (Var, [Var], UrTy Var, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
v,[Var]
locs,UrTy Var
ty,(Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
rhs)) (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
bod)
        IfE Exp2
a Exp2
b Exp2
c -> Exp2 -> Exp2 -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> PreExp ext loc dec -> PreExp ext loc dec -> PreExp ext loc dec
IfE (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
a) (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
b) (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
c)
        MkProdE [Exp2]
args -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
[PreExp ext loc dec] -> PreExp ext loc dec
MkProdE ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env) [Exp2]
args)
        ProjE Int
i Exp2
bod -> Int -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Int -> PreExp ext loc dec -> PreExp ext loc dec
ProjE Int
i (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
bod)
        CaseE Exp2
scrt [([Char], [(Var, Var)], Exp2)]
brs -> Exp2 -> [([Char], [(Var, Var)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [([Char], [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
scrt) ((([Char], [(Var, Var)], Exp2) -> ([Char], [(Var, Var)], Exp2))
-> [([Char], [(Var, Var)], Exp2)] -> [([Char], [(Var, Var)], Exp2)]
forall a b. (a -> b) -> [a] -> [b]
map (\([Char]
a,[(Var, Var)]
b,Exp2
c) -> ([Char]
a,[(Var, Var)]
b,Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
c)) [([Char], [(Var, Var)], Exp2)]
brs)
        DataConE Var
loc [Char]
dcon [Exp2]
args -> Var -> [Char] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
loc -> [Char] -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE Var
loc [Char]
dcon ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env) [Exp2]
args)
        TimeIt Exp2
e UrTy Var
ty Bool
b -> Exp2 -> UrTy Var -> Bool -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> Bool -> PreExp ext loc dec
TimeIt (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
e) UrTy Var
ty Bool
b
        WithArenaE Var
v Exp2
bod -> Var -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE Var
v (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
bod)
        SpawnE Var
f [Var]
locs [Exp2]
args -> Var -> [Var] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
SpawnE Var
f [Var]
locs ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env) [Exp2]
args)
        Ext E2Ext Var (UrTy Var)
ext ->
          case E2Ext Var (UrTy Var)
ext of
            LetRegionE Region
reg RegionSize
sz Maybe RegionType
ty Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
reg RegionSize
sz Maybe RegionType
ty (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
bod))
            LetParRegionE Region
reg RegionSize
sz Maybe RegionType
ty Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE Region
reg RegionSize
sz Maybe RegionType
ty (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
bod))
            LetLocE Var
loc (AfterConstantLE Int
i Var
loc2) Exp2
bod ->
              case (Var -> Map Var (Var, Int) -> Maybe (Var, Int)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
loc2 Map Var (Var, Int)
env) of
                Maybe (Var, Int)
Nothing ->
                  E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2) -> E2Ext Var (UrTy Var) -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> PreLocExp Var -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc (Int -> Var -> PreLocExp Var
forall loc. Int -> loc -> PreLocExp loc
AfterConstantLE Int
i Var
loc2) (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> E2Ext Var (UrTy Var)
forall a b. (a -> b) -> a -> b
$
                        Map Var (Var, Int) -> Exp2 -> Exp2
go (Var -> (Var, Int) -> Map Var (Var, Int) -> Map Var (Var, Int)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc (Var
loc2,Int
i) Map Var (Var, Int)
env) Exp2
bod
                Just (Var
loc3,Int
j) ->
                  E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2) -> E2Ext Var (UrTy Var) -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> PreLocExp Var -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc (Int -> Var -> PreLocExp Var
forall loc. Int -> loc -> PreLocExp loc
AfterConstantLE (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
j) Var
loc3) (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> E2Ext Var (UrTy Var)
forall a b. (a -> b) -> a -> b
$
                        Map Var (Var, Int) -> Exp2 -> Exp2
go (Var -> (Var, Int) -> Map Var (Var, Int) -> Map Var (Var, Int)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc (Var
loc3,Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
j) Map Var (Var, Int)
env) Exp2
bod
            LetLocE Var
loc PreLocExp Var
rhs Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (Var -> PreLocExp Var -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc PreLocExp Var
rhs (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
bod))
            LetAvail [Var]
vars Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext ([Var] -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. [Var] -> E2 loc dec -> E2Ext loc dec
LetAvail [Var]
vars (Map Var (Var, Int) -> Exp2 -> Exp2
go Map Var (Var, Int)
env Exp2
bod))
            E2Ext Var (UrTy Var)
_ -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext E2Ext Var (UrTy Var)
ext
        Exp2
_ -> Exp2
ex

    -- drop dead bindings
    go2 :: Exp2 -> Exp2
    go2 :: Exp2 -> Exp2
go2 Exp2
ex =
      case Exp2
ex of
        AppE Var
f [Var]
locs [Exp2]
args -> Var -> [Var] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f [Var]
locs ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go2 [Exp2]
args)
        PrimAppE Prim (UrTy Var)
p [Exp2]
args -> Prim (UrTy Var) -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim (UrTy Var)
p ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go2 [Exp2]
args)
        LetE (Var
v,[Var]
locs,UrTy Var
ty,Exp2
rhs) Exp2
bod -> (Var, [Var], UrTy Var, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
v,[Var]
locs,UrTy Var
ty,(Exp2 -> Exp2
go2 Exp2
rhs)) (Exp2 -> Exp2
go2 Exp2
bod)
        IfE Exp2
a Exp2
b Exp2
c -> Exp2 -> Exp2 -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> PreExp ext loc dec -> PreExp ext loc dec -> PreExp ext loc dec
IfE (Exp2 -> Exp2
go2 Exp2
a) (Exp2 -> Exp2
go2 Exp2
b) (Exp2 -> Exp2
go2 Exp2
c)
        MkProdE [Exp2]
args -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
[PreExp ext loc dec] -> PreExp ext loc dec
MkProdE ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go2 [Exp2]
args)
        ProjE Int
i Exp2
bod -> Int -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Int -> PreExp ext loc dec -> PreExp ext loc dec
ProjE Int
i (Exp2 -> Exp2
go2 Exp2
bod)
        CaseE Exp2
scrt [([Char], [(Var, Var)], Exp2)]
brs -> Exp2 -> [([Char], [(Var, Var)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [([Char], [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE (Exp2 -> Exp2
go2 Exp2
scrt) ((([Char], [(Var, Var)], Exp2) -> ([Char], [(Var, Var)], Exp2))
-> [([Char], [(Var, Var)], Exp2)] -> [([Char], [(Var, Var)], Exp2)]
forall a b. (a -> b) -> [a] -> [b]
map (\([Char]
a,[(Var, Var)]
b,Exp2
c) -> ([Char]
a,[(Var, Var)]
b,Exp2 -> Exp2
go2 Exp2
c)) [([Char], [(Var, Var)], Exp2)]
brs)
        DataConE Var
loc [Char]
dcon [Exp2]
args -> Var -> [Char] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
loc -> [Char] -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE Var
loc [Char]
dcon ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go2 [Exp2]
args)
        TimeIt Exp2
e UrTy Var
ty Bool
b -> Exp2 -> UrTy Var -> Bool -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> Bool -> PreExp ext loc dec
TimeIt (Exp2 -> Exp2
go2 Exp2
e) UrTy Var
ty Bool
b
        WithArenaE Var
v Exp2
bod -> Var -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE Var
v (Exp2 -> Exp2
go2 Exp2
bod)
        SpawnE Var
f [Var]
locs [Exp2]
args -> Var -> [Var] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
SpawnE Var
f [Var]
locs ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go2 [Exp2]
args)
        Ext E2Ext Var (UrTy Var)
ext ->
          case E2Ext Var (UrTy Var)
ext of
            LetRegionE Region
reg RegionSize
sz Maybe RegionType
ty Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
reg RegionSize
sz Maybe RegionType
ty (Exp2 -> Exp2
go2 Exp2
bod))
            LetParRegionE Region
reg RegionSize
sz Maybe RegionType
ty Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE Region
reg RegionSize
sz Maybe RegionType
ty (Exp2 -> Exp2
go2 Exp2
bod))
            LetLocE Var
loc PreLocExp Var
rhs Exp2
bod ->
              let bod' :: Exp2
bod' = Exp2 -> Exp2
go2 Exp2
bod
                  free_vars :: Set Var
free_vars = (Exp2 -> Set Var
allFreeVars Exp2
bod')
              in
                if (Var
loc Var -> Set Var -> Bool
forall a. Eq a => a -> Set a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Set Var
free_vars)
                then E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (Var -> PreLocExp Var -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc PreLocExp Var
rhs Exp2
bod')
                else Exp2
bod'
            LetAvail [Var]
vars Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext ([Var] -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. [Var] -> E2 loc dec -> E2Ext loc dec
LetAvail [Var]
vars (Exp2 -> Exp2
go2 Exp2
bod))
            E2Ext Var (UrTy Var)
_ -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext E2Ext Var (UrTy Var)
ext
        Exp2
_ -> Exp2
ex

    -- partially evaluate location arithmetic
    go0 :: M.Map LocExp LocVar -> M.Map LocVar LocVar -> Exp2 -> Exp2
    go0 :: Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
ex =
      case Exp2
ex of
        AppE Var
f [Var]
locs [Exp2]
args -> Var -> [Var] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f ((Var -> Var) -> [Var] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (Map Var Var -> Var -> Var
forall {k}. Ord k => Map k k -> k -> k
substloc Map Var Var
env2) [Var]
locs) ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2) [Exp2]
args)
        PrimAppE Prim (UrTy Var)
p [Exp2]
args -> Prim (UrTy Var) -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim (UrTy Var)
p ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2) [Exp2]
args)
        LetE (Var
v,[Var]
locs,UrTy Var
ty,Exp2
rhs) Exp2
bod -> (Var, [Var], UrTy Var, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
v,[Var]
locs,Map Var Var -> UrTy Var -> UrTy Var
substLoc Map Var Var
env2 UrTy Var
ty,(Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
rhs)) (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
bod)
        IfE Exp2
a Exp2
b Exp2
c -> Exp2 -> Exp2 -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> PreExp ext loc dec -> PreExp ext loc dec -> PreExp ext loc dec
IfE (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
a) (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
b) (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
c)
        MkProdE [Exp2]
args -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
[PreExp ext loc dec] -> PreExp ext loc dec
MkProdE ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2) [Exp2]
args)
        ProjE Int
i Exp2
bod -> Int -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Int -> PreExp ext loc dec -> PreExp ext loc dec
ProjE Int
i (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
bod)
        CaseE Exp2
scrt [([Char], [(Var, Var)], Exp2)]
brs -> Exp2 -> [([Char], [(Var, Var)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [([Char], [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
scrt) ((([Char], [(Var, Var)], Exp2) -> ([Char], [(Var, Var)], Exp2))
-> [([Char], [(Var, Var)], Exp2)] -> [([Char], [(Var, Var)], Exp2)]
forall a b. (a -> b) -> [a] -> [b]
map (\([Char]
a,[(Var, Var)]
b,Exp2
c) -> ([Char]
a,[(Var, Var)]
b,Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
c)) [([Char], [(Var, Var)], Exp2)]
brs)
        DataConE Var
loc [Char]
dcon [Exp2]
args -> Var -> [Char] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
loc -> [Char] -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE (Map Var Var -> Var -> Var
forall {k}. Ord k => Map k k -> k -> k
substloc Map Var Var
env2 Var
loc) [Char]
dcon ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2) [Exp2]
args)
        TimeIt Exp2
e UrTy Var
ty Bool
b -> Exp2 -> UrTy Var -> Bool -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> Bool -> PreExp ext loc dec
TimeIt (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
e) UrTy Var
ty Bool
b
        WithArenaE Var
v Exp2
bod -> Var -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE Var
v (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
bod)
        SpawnE Var
f [Var]
locs [Exp2]
args -> Var -> [Var] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
SpawnE Var
f ((Var -> Var) -> [Var] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (Map Var Var -> Var -> Var
forall {k}. Ord k => Map k k -> k -> k
substloc Map Var Var
env2) [Var]
locs) ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2) [Exp2]
args)
        Ext E2Ext Var (UrTy Var)
ext ->
          case E2Ext Var (UrTy Var)
ext of
            LetRegionE Region
reg RegionSize
sz Maybe RegionType
ty Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
reg RegionSize
sz Maybe RegionType
ty (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
bod))
            LetParRegionE Region
reg RegionSize
sz Maybe RegionType
ty Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE Region
reg RegionSize
sz Maybe RegionType
ty (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
bod))
            LetLocE Var
loc PreLocExp Var
rhs Exp2
bod ->
              let rhs' :: PreLocExp Var
rhs' = case PreLocExp Var
rhs of
                           AfterConstantLE Int
i Var
loc2 -> Int -> Var -> PreLocExp Var
forall loc. Int -> loc -> PreLocExp loc
AfterConstantLE Int
i (Map Var Var -> Var -> Var
forall {k}. Ord k => Map k k -> k -> k
substloc Map Var Var
env2 Var
loc2)
                           AfterVariableLE Var
v Var
loc2 Bool
b -> Var -> Var -> Bool -> PreLocExp Var
forall loc. Var -> loc -> Bool -> PreLocExp loc
AfterVariableLE Var
v (Map Var Var -> Var -> Var
forall {k}. Ord k => Map k k -> k -> k
substloc Map Var Var
env2 Var
loc2) Bool
b
                           PreLocExp Var
_ -> PreLocExp Var
rhs
              in case PreLocExp Var -> Map (PreLocExp Var) Var -> Maybe Var
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup PreLocExp Var
rhs' Map (PreLocExp Var) Var
env1 of
                Maybe Var
Nothing  -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (Var -> PreLocExp Var -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc PreLocExp Var
rhs' (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 (PreLocExp Var
-> Var -> Map (PreLocExp Var) Var -> Map (PreLocExp Var) Var
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert PreLocExp Var
rhs' Var
loc Map (PreLocExp Var) Var
env1) Map Var Var
env2 Exp2
bod))
                Just Var
new -> Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 (Var -> Var -> Map Var Var -> Map Var Var
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc Var
new Map Var Var
env2) Exp2
bod
            LetAvail [Var]
vars Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext ([Var] -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. [Var] -> E2 loc dec -> E2Ext loc dec
LetAvail [Var]
vars (Map (PreLocExp Var) Var -> Map Var Var -> Exp2 -> Exp2
go0 Map (PreLocExp Var) Var
env1 Map Var Var
env2 Exp2
bod))
            E2Ext Var (UrTy Var)
_ -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext E2Ext Var (UrTy Var)
ext
        Exp2
_ -> Exp2
ex
      where
        substloc :: Map k k -> k -> k
substloc Map k k
env k
loc = case k -> Map k k -> Maybe k
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
loc Map k k
env of
                             Maybe k
Nothing  -> k
loc
                             Just k
new -> k
new

--------------------------------------------------------------------------------

lateInlineTriv :: L4.Prog -> PassM L4.Prog
lateInlineTriv :: Prog -> PassM Prog
lateInlineTriv (L4.Prog InfoTable
info_tbl SymTable
sym_tbl [FunDecl]
fundefs Maybe MainExp
mainExp) = do
    let fundefs' :: [FunDecl]
fundefs' = (FunDecl -> FunDecl) -> [FunDecl] -> [FunDecl]
forall a b. (a -> b) -> [a] -> [b]
map FunDecl -> FunDecl
lateInlineTrivFn [FunDecl]
fundefs
        mainExp' :: Maybe MainExp
mainExp' = case Maybe MainExp
mainExp of
                       Just (L4.PrintExp Tail
tl) -> do
                           MainExp -> Maybe MainExp
forall a. a -> Maybe a
Just (Tail -> MainExp
L4.PrintExp (Map Var Triv -> Tail -> Tail
lateInlineTrivExp Map Var Triv
forall k a. Map k a
M.empty Tail
tl))
                       Maybe MainExp
Nothing -> Maybe MainExp
forall a. Maybe a
Nothing
    Prog -> PassM Prog
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Prog -> PassM Prog) -> Prog -> PassM Prog
forall a b. (a -> b) -> a -> b
$ InfoTable -> SymTable -> [FunDecl] -> Maybe MainExp -> Prog
L4.Prog InfoTable
info_tbl SymTable
sym_tbl [FunDecl]
fundefs' Maybe MainExp
mainExp'
  where
    lateInlineTrivFn :: L4.FunDecl -> L4.FunDecl
    lateInlineTrivFn :: FunDecl -> FunDecl
lateInlineTrivFn f :: FunDecl
f@L4.FunDecl{Tail
funBody :: Tail
funBody :: FunDecl -> Tail
L4.funBody} =
        let bod' :: Tail
bod' = Map Var Triv -> Tail -> Tail
lateInlineTrivExp Map Var Triv
forall k a. Map k a
M.empty Tail
funBody
        in FunDecl
f {funBody :: Tail
L4.funBody = Tail
bod'}

    lateInlineTrivExp :: M.Map Var L4.Triv -> L4.Tail -> L4.Tail
    lateInlineTrivExp :: Map Var Triv -> Tail -> Tail
lateInlineTrivExp = Map Var Triv -> Tail -> Tail
go
      where
        gotriv :: Map Var Triv -> Triv -> Triv
gotriv Map Var Triv
env Triv
trv =
            case Triv
trv of
                L4.VarTriv Var
v -> case Var -> Map Var Triv -> Maybe Triv
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
v Map Var Triv
env of
                                 Maybe Triv
Nothing -> Triv
trv
                                 Just Triv
t2 -> Triv
t2
                L4.ProdTriv [Triv]
ls -> [Triv] -> Triv
L4.ProdTriv ((Triv -> Triv) -> [Triv] -> [Triv]
forall a b. (a -> b) -> [a] -> [b]
map (Map Var Triv -> Triv -> Triv
gotriv Map Var Triv
env) [Triv]
ls)
                L4.ProjTriv Int
i Triv
t -> Int -> Triv -> Triv
L4.ProjTriv Int
i (Map Var Triv -> Triv -> Triv
gotriv Map Var Triv
env Triv
t)
                Triv
_ -> Triv
trv

        goalts :: Map Var Triv -> Alts -> Alts
goalts Map Var Triv
env Alts
alts =
            case Alts
alts of
                L4.TagAlts [(Tag, Tail)]
ls -> [(Tag, Tail)] -> Alts
L4.TagAlts ([(Tag, Tail)] -> Alts) -> [(Tag, Tail)] -> Alts
forall a b. (a -> b) -> a -> b
$ ((Tag, Tail) -> (Tag, Tail)) -> [(Tag, Tail)] -> [(Tag, Tail)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Tag
t,Tail
tl) -> (Tag
t,Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
tl)) [(Tag, Tail)]
ls
                L4.IntAlts [(Int64, Tail)]
ls -> [(Int64, Tail)] -> Alts
L4.IntAlts ([(Int64, Tail)] -> Alts) -> [(Int64, Tail)] -> Alts
forall a b. (a -> b) -> a -> b
$ ((Int64, Tail) -> (Int64, Tail))
-> [(Int64, Tail)] -> [(Int64, Tail)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Int64
t,Tail
tl) -> (Int64
t,Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
tl)) [(Int64, Tail)]
ls
        go :: Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
tl =
              case Tail
tl of
                   Tail
L4.EndOfMain -> Tail
L4.EndOfMain
                   L4.RetValsT [Triv]
trvs ->
                       [Triv] -> Tail
L4.RetValsT ((Triv -> Triv) -> [Triv] -> [Triv]
forall a b. (a -> b) -> [a] -> [b]
map (Map Var Triv -> Triv -> Triv
gotriv Map Var Triv
env) [Triv]
trvs)
                   L4.AssnValsT [(Var, Ty, Triv)]
upd Maybe Tail
mb_bod ->
                       [(Var, Ty, Triv)] -> Maybe Tail -> Tail
L4.AssnValsT [(Var, Ty, Triv)]
upd ((Tail -> Tail) -> Maybe Tail -> Maybe Tail
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Map Var Triv -> Tail -> Tail
go Map Var Triv
env) Maybe Tail
mb_bod)
                   L4.LetCallT Bool
async [(Var, Ty)]
binds Var
rator [Triv]
rands Tail
bod ->
                       Bool -> [(Var, Ty)] -> Var -> [Triv] -> Tail -> Tail
L4.LetCallT Bool
async [(Var, Ty)]
binds Var
rator ((Triv -> Triv) -> [Triv] -> [Triv]
forall a b. (a -> b) -> [a] -> [b]
map (Map Var Triv -> Triv -> Triv
gotriv Map Var Triv
env) [Triv]
rands) (Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
bod)
                   L4.LetPrimCallT [(Var, Ty)]
binds Prim
prim [Triv]
rands Tail
bod ->
                       [(Var, Ty)] -> Prim -> [Triv] -> Tail -> Tail
L4.LetPrimCallT [(Var, Ty)]
binds Prim
prim ((Triv -> Triv) -> [Triv] -> [Triv]
forall a b. (a -> b) -> [a] -> [b]
map (Map Var Triv -> Triv -> Triv
gotriv Map Var Triv
env) [Triv]
rands) (Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
bod)
                   L4.LetTrivT (Var
v,Ty
_ty,Triv
trv) Tail
bod ->
                       case Triv
trv of
                           L4.VarTriv Var
w -> case Var -> Map Var Triv -> Maybe Triv
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
w Map Var Triv
env of
                                            Maybe Triv
Nothing -> Map Var Triv -> Tail -> Tail
go (Var -> Triv -> Map Var Triv -> Map Var Triv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
v Triv
trv Map Var Triv
env) Tail
bod
                                            Just Triv
trv' -> Map Var Triv -> Tail -> Tail
go (Var -> Triv -> Map Var Triv -> Map Var Triv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
v Triv
trv' Map Var Triv
env) Tail
bod
                           Triv
_ -> Map Var Triv -> Tail -> Tail
go (Var -> Triv -> Map Var Triv -> Map Var Triv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
v (Map Var Triv -> Triv -> Triv
gotriv Map Var Triv
env Triv
trv) Map Var Triv
env) Tail
bod
                   L4.LetIfT [(Var, Ty)]
binds (Triv
trv,Tail
tl1,Tail
tl2) Tail
bod ->
                       [(Var, Ty)] -> (Triv, Tail, Tail) -> Tail -> Tail
L4.LetIfT [(Var, Ty)]
binds (Map Var Triv -> Triv -> Triv
gotriv Map Var Triv
env Triv
trv, Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
tl1, Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
tl2) (Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
bod)
                   L4.LetUnpackT [(Var, Ty)]
binds Var
ptr Tail
bod ->
                       [(Var, Ty)] -> Var -> Tail -> Tail
L4.LetUnpackT [(Var, Ty)]
binds Var
ptr (Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
bod)
                   L4.LetAllocT Var
lhs [(Ty, Triv)]
vals Tail
bod ->
                       Var -> [(Ty, Triv)] -> Tail -> Tail
L4.LetAllocT Var
lhs (((Ty, Triv) -> (Ty, Triv)) -> [(Ty, Triv)] -> [(Ty, Triv)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Ty
ty,Triv
trv) -> (Ty
ty,Map Var Triv -> Triv -> Triv
gotriv Map Var Triv
env Triv
trv)) [(Ty, Triv)]
vals) (Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
bod)
                   L4.LetAvailT [Var]
vars Tail
bod ->
                       [Var] -> Tail -> Tail
L4.LetAvailT [Var]
vars (Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
bod)
                   L4.IfT Triv
tst Tail
thn Tail
els ->
                       Triv -> Tail -> Tail -> Tail
L4.IfT (Map Var Triv -> Triv -> Triv
gotriv Map Var Triv
env Triv
tst) (Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
thn) (Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
els)
                   L4.ErrT [Char]
str -> [Char] -> Tail
L4.ErrT [Char]
str
                   L4.LetTimedT Bool
isiter [(Var, Ty)]
binds Tail
timed Tail
bod ->
                       Bool -> [(Var, Ty)] -> Tail -> Tail -> Tail
L4.LetTimedT Bool
isiter [(Var, Ty)]
binds (Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
timed) (Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
bod)
                   L4.Switch Var
lbl Triv
trv Alts
alts Maybe Tail
mb_tl ->
                       Var -> Triv -> Alts -> Maybe Tail -> Tail
L4.Switch Var
lbl Triv
trv (Map Var Triv -> Alts -> Alts
goalts Map Var Triv
env Alts
alts) ((Tail -> Tail) -> Maybe Tail -> Maybe Tail
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Map Var Triv -> Tail -> Tail
go Map Var Triv
env) Maybe Tail
mb_tl)
                   L4.TailCall Var
var [Triv]
trvs ->
                       Var -> [Triv] -> Tail
L4.TailCall Var
var ((Triv -> Triv) -> [Triv] -> [Triv]
forall a b. (a -> b) -> [a] -> [b]
map (Map Var Triv -> Triv -> Triv
gotriv Map Var Triv
env) [Triv]
trvs)
                   L4.Goto Var
lbl -> Var -> Tail
L4.Goto Var
lbl
                   L4.LetArenaT Var
lhs Tail
bod ->
                       Var -> Tail -> Tail
L4.LetArenaT Var
lhs (Map Var Triv -> Tail -> Tail
go Map Var Triv
env Tail
bod)