{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE CPP              #-}
{-# LANGUAGE RecordWildCards  #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

-- | Insert end witnesses in an L2 program by changing function types,
-- and updating expressions to pass (second-class) end locations around
-- via RetE in tail position and an extended binding form in LetE.
-- Assumes that expressions are flattened and in ANF, and that location
-- symbols are all unique! Failure to meet these assumptions will cause
-- this pass to fail or possibly produce wrong output.
--
--- Steps:
---
--- 1. For each function type, inspect its input parameter type and traversal
---    effect to determine which packed arguments are completely traversed,
---    and update the type to indicate that the EndOf witness for each
---    of these traversed arguments is returned.
---
--- 2. For each function body, walk through the let spine, updating bindings
---    to include locations returned as end witnesses, and add coersions from
---    EndOf witnesses to expected locations. Upon reaching tail position,
---    emit a RetE form and compute what the proper EndOf locations are.
---
--- 3. For the main body, do the same thing minus inserting the RetE in tail
---    position.

module Gibbon.Passes.RouteEnds
    ( routeEnds ) where

import qualified Data.List as L
import Data.Maybe ( fromJust )
import Data.Map as M
import Data.Set as S
import Control.Monad

import Gibbon.Common
import Gibbon.L2.Syntax as L2
import Gibbon.L1.Syntax as L1


-- | Data structure that accumulates what we know about the relationship
-- between locations and EndOf witnesses.
--
-- Performing a lookup (to find the end of a given location) will first check
-- if there exists a mapping in the endOf map for that location, then will check
-- if there exists a mapping in the equivTo map, and if so will recur to find
-- the end of that location.
--
-- This is used for when we perform pattern matching. The end of some binary tree
-- (for example) is the same as the end of its second node, so we want to record
-- that knowledge as we traverse the AST.
data EndOfRel = EndOfRel
    {
      EndOfRel -> Map LocVar LocVar
endOf :: M.Map LocVar LocVar -- ^ Map a location to it's EndOf witness
    , EndOfRel -> Map LocVar LocVar
equivTo :: M.Map LocVar LocVar -- ^ Map of a location to a known equivalent location
    }
  deriving (EndOfRel -> EndOfRel -> Bool
(EndOfRel -> EndOfRel -> Bool)
-> (EndOfRel -> EndOfRel -> Bool) -> Eq EndOfRel
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: EndOfRel -> EndOfRel -> Bool
== :: EndOfRel -> EndOfRel -> Bool
$c/= :: EndOfRel -> EndOfRel -> Bool
/= :: EndOfRel -> EndOfRel -> Bool
Eq, Eq EndOfRel
Eq EndOfRel
-> (EndOfRel -> EndOfRel -> Ordering)
-> (EndOfRel -> EndOfRel -> Bool)
-> (EndOfRel -> EndOfRel -> Bool)
-> (EndOfRel -> EndOfRel -> Bool)
-> (EndOfRel -> EndOfRel -> Bool)
-> (EndOfRel -> EndOfRel -> EndOfRel)
-> (EndOfRel -> EndOfRel -> EndOfRel)
-> Ord EndOfRel
EndOfRel -> EndOfRel -> Bool
EndOfRel -> EndOfRel -> Ordering
EndOfRel -> EndOfRel -> EndOfRel
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: EndOfRel -> EndOfRel -> Ordering
compare :: EndOfRel -> EndOfRel -> Ordering
$c< :: EndOfRel -> EndOfRel -> Bool
< :: EndOfRel -> EndOfRel -> Bool
$c<= :: EndOfRel -> EndOfRel -> Bool
<= :: EndOfRel -> EndOfRel -> Bool
$c> :: EndOfRel -> EndOfRel -> Bool
> :: EndOfRel -> EndOfRel -> Bool
$c>= :: EndOfRel -> EndOfRel -> Bool
>= :: EndOfRel -> EndOfRel -> Bool
$cmax :: EndOfRel -> EndOfRel -> EndOfRel
max :: EndOfRel -> EndOfRel -> EndOfRel
$cmin :: EndOfRel -> EndOfRel -> EndOfRel
min :: EndOfRel -> EndOfRel -> EndOfRel
Ord, ReadPrec [EndOfRel]
ReadPrec EndOfRel
Int -> ReadS EndOfRel
ReadS [EndOfRel]
(Int -> ReadS EndOfRel)
-> ReadS [EndOfRel]
-> ReadPrec EndOfRel
-> ReadPrec [EndOfRel]
-> Read EndOfRel
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS EndOfRel
readsPrec :: Int -> ReadS EndOfRel
$creadList :: ReadS [EndOfRel]
readList :: ReadS [EndOfRel]
$creadPrec :: ReadPrec EndOfRel
readPrec :: ReadPrec EndOfRel
$creadListPrec :: ReadPrec [EndOfRel]
readListPrec :: ReadPrec [EndOfRel]
Read, Int -> EndOfRel -> ShowS
[EndOfRel] -> ShowS
EndOfRel -> String
(Int -> EndOfRel -> ShowS)
-> (EndOfRel -> String) -> ([EndOfRel] -> ShowS) -> Show EndOfRel
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> EndOfRel -> ShowS
showsPrec :: Int -> EndOfRel -> ShowS
$cshow :: EndOfRel -> String
show :: EndOfRel -> String
$cshowList :: [EndOfRel] -> ShowS
showList :: [EndOfRel] -> ShowS
Show)

-- | Create an empty EndOfRel
emptyRel :: EndOfRel
emptyRel :: EndOfRel
emptyRel = Map LocVar LocVar -> Map LocVar LocVar -> EndOfRel
EndOfRel Map LocVar LocVar
forall k a. Map k a
M.empty Map LocVar LocVar
forall k a. Map k a
M.empty

-- | Assert that one location's end is equivalent to another's end.
-- Order is important here: we expect to look up the end of the first
-- location argument, not the second.
mkEqual :: LocVar -> LocVar -> EndOfRel -> EndOfRel
mkEqual :: LocVar -> LocVar -> EndOfRel -> EndOfRel
mkEqual LocVar
l1 LocVar
l2 EndOfRel{Map LocVar LocVar
endOf :: EndOfRel -> Map LocVar LocVar
endOf :: Map LocVar LocVar
endOf,Map LocVar LocVar
equivTo :: EndOfRel -> Map LocVar LocVar
equivTo :: Map LocVar LocVar
equivTo} =
    EndOfRel{endOf :: Map LocVar LocVar
endOf=Map LocVar LocVar
endOf,equivTo :: Map LocVar LocVar
equivTo=LocVar -> LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
l1 LocVar
l2 Map LocVar LocVar
equivTo}

-- | Assert that we have found an EndOf relation between two locations.
mkEnd :: LocVar -> LocVar -> EndOfRel -> EndOfRel
mkEnd :: LocVar -> LocVar -> EndOfRel -> EndOfRel
mkEnd LocVar
l1 LocVar
l2 EndOfRel{Map LocVar LocVar
endOf :: EndOfRel -> Map LocVar LocVar
endOf :: Map LocVar LocVar
endOf,Map LocVar LocVar
equivTo :: EndOfRel -> Map LocVar LocVar
equivTo :: Map LocVar LocVar
equivTo} =
    EndOfRel{endOf :: Map LocVar LocVar
endOf=LocVar -> LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
l1 LocVar
l2 Map LocVar LocVar
endOf,equivTo :: Map LocVar LocVar
equivTo=Map LocVar LocVar
equivTo}

-- | Look up the end of a location.
findEnd :: LocVar -> EndOfRel -> LocVar
findEnd :: LocVar -> EndOfRel -> LocVar
findEnd LocVar
l EndOfRel{Map LocVar LocVar
endOf :: EndOfRel -> Map LocVar LocVar
endOf :: Map LocVar LocVar
endOf,Map LocVar LocVar
equivTo :: EndOfRel -> Map LocVar LocVar
equivTo :: Map LocVar LocVar
equivTo} =
    case LocVar -> Map LocVar LocVar -> Maybe LocVar
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup LocVar
l Map LocVar LocVar
endOf of -- Can we immediately look up the end of l?
      Maybe LocVar
Nothing -> case LocVar -> Map LocVar LocVar -> Maybe LocVar
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup LocVar
l Map LocVar LocVar
equivTo of -- Is there an equivalent location to use?
                   Maybe LocVar
Nothing -> String -> LocVar
forall a. HasCallStack => String -> a
error (String -> LocVar) -> String -> LocVar
forall a b. (a -> b) -> a -> b
$ String
"Failed finding the end of " String -> ShowS
forall a. [a] -> [a] -> [a]
++ (LocVar -> String
forall a. Show a => a -> String
show LocVar
l)
                   Just LocVar
leq -> LocVar -> EndOfRel -> LocVar
findEnd LocVar
leq EndOfRel{Map LocVar LocVar
endOf :: Map LocVar LocVar
endOf :: Map LocVar LocVar
endOf,Map LocVar LocVar
equivTo :: Map LocVar LocVar
equivTo :: Map LocVar LocVar
equivTo}
      Just LocVar
lend -> LocVar
lend

eorAppend :: EndOfRel -> EndOfRel -> EndOfRel
eorAppend :: EndOfRel -> EndOfRel -> EndOfRel
eorAppend EndOfRel
eor1 EndOfRel
eor2 =
    EndOfRel { endOf :: Map LocVar LocVar
endOf   = (EndOfRel -> Map LocVar LocVar
endOf EndOfRel
eor1) Map LocVar LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` (EndOfRel -> Map LocVar LocVar
endOf EndOfRel
eor2)
             , equivTo :: Map LocVar LocVar
equivTo = (EndOfRel -> Map LocVar LocVar
equivTo EndOfRel
eor1) Map LocVar LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` (EndOfRel -> Map LocVar LocVar
equivTo EndOfRel
eor2) }

instance Monoid EndOfRel where
  mempty :: EndOfRel
mempty  = EndOfRel
emptyRel

#if !(MIN_VERSION_base(4,11,0))
  mappend = eorAppend
#endif

instance Semigroup EndOfRel where
  <> :: EndOfRel -> EndOfRel -> EndOfRel
(<>) = EndOfRel -> EndOfRel -> EndOfRel
eorAppend

bindReturns :: Exp2 -> PassM Exp2
bindReturns :: Exp2 -> PassM Exp2
bindReturns Exp2
ex =
  case Exp2
ex of
    VarE LocVar
v -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
v)
    LitE{} -> Exp2 -> (Exp2 -> Exp2) -> PassM Exp2
handleScalarRet Exp2
ex Exp2 -> Exp2
forall a. a -> a
id
    CharE{} -> Exp2 -> (Exp2 -> Exp2) -> PassM Exp2
handleScalarRet Exp2
ex Exp2 -> Exp2
forall a. a -> a
id
    FloatE{} -> Exp2 -> (Exp2 -> Exp2) -> PassM Exp2
handleScalarRet Exp2
ex Exp2 -> Exp2
forall a. a -> a
id
    LitSymE{} -> Exp2 -> (Exp2 -> Exp2) -> PassM Exp2
handleScalarRet Exp2
ex Exp2 -> Exp2
forall a. a -> a
id
    AppE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    PrimAppE Prim Ty2
p [Exp2]
_ ->
      case Prim Ty2
p of
        Prim Ty2
MkTrue -> Exp2 -> (Exp2 -> Exp2) -> PassM Exp2
handleScalarRet Exp2
ex Exp2 -> Exp2
forall a. a -> a
id
        Prim Ty2
MkFalse -> Exp2 -> (Exp2 -> Exp2) -> PassM Exp2
handleScalarRet Exp2
ex Exp2 -> Exp2
forall a. a -> a
id
        Prim Ty2
_ -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    LetE (LocVar
v,[LocVar]
locs,Ty2
ty,Exp2
rhs) Exp2
bod | Exp2 -> Bool
isScalar Exp2
bod -> do
      Exp2 -> (Exp2 -> Exp2) -> PassM Exp2
handleScalarRet Exp2
bod (\Exp2
bod' -> (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[LocVar]
locs,Ty2
ty,Exp2
rhs) Exp2
bod')

    LetE (LocVar
v,[LocVar]
locs,Ty2
ty,Exp2
rhs) Exp2
bod -> do
      -- rhs' <- bindReturns rhs
      Exp2
bod' <- Exp2 -> PassM Exp2
bindReturns Exp2
bod
      Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[LocVar]
locs,Ty2
ty,Exp2
rhs) Exp2
bod'

    IfE Exp2
a Exp2
b Exp2
c  -> do
      Exp2
a' <- Exp2 -> PassM Exp2
bindReturns Exp2
a
      Exp2
b' <- Exp2 -> PassM Exp2
bindReturns Exp2
b
      Exp2
c' <- Exp2 -> PassM Exp2
bindReturns Exp2
c
      Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ 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
a' Exp2
b' Exp2
c'
    MkProdE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    ProjE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    CaseE Exp2
scrt [(String, [(LocVar, LocVar)], Exp2)]
brs -> do
      Exp2
scrt' <- Exp2 -> PassM Exp2
bindReturns Exp2
scrt
      Exp2 -> [(String, [(LocVar, LocVar)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(String, [(LocVar, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE Exp2
scrt' ([(String, [(LocVar, LocVar)], Exp2)] -> Exp2)
-> PassM [(String, [(LocVar, LocVar)], Exp2)] -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (((String, [(LocVar, LocVar)], Exp2)
 -> PassM (String, [(LocVar, LocVar)], Exp2))
-> [(String, [(LocVar, LocVar)], Exp2)]
-> PassM [(String, [(LocVar, LocVar)], Exp2)]
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 (\(String
a,[(LocVar, LocVar)]
b,Exp2
c) -> (String
a,[(LocVar, LocVar)]
b,) (Exp2 -> (String, [(LocVar, LocVar)], Exp2))
-> PassM Exp2 -> PassM (String, [(LocVar, LocVar)], Exp2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
bindReturns Exp2
c) [(String, [(LocVar, LocVar)], Exp2)]
brs)
    DataConE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    TimeIt{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    SpawnE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    Exp2
SyncE -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    WithArenaE LocVar
v Exp2
e -> do
      Exp2
e' <- Exp2 -> PassM Exp2
bindReturns Exp2
e
      Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ LocVar -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
LocVar -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE LocVar
v Exp2
e'
    Ext E2Ext LocVar Ty2
ext ->
      case E2Ext LocVar Ty2
ext of
        LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod -> do
          Exp2
bod' <- Exp2 -> PassM Exp2
bindReturns Exp2
bod
          Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ E2Ext LocVar Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocVar Ty2 -> Exp2) -> E2Ext LocVar Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext LocVar Ty2
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod'
        LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod -> do
          Exp2
bod' <- Exp2 -> PassM Exp2
bindReturns Exp2
bod
          Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ E2Ext LocVar Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocVar Ty2 -> Exp2) -> E2Ext LocVar Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext LocVar Ty2
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod'
        LetLocE LocVar
loc PreLocExp LocVar
locexp Exp2
bod -> do
          Exp2
bod' <- Exp2 -> PassM Exp2
bindReturns Exp2
bod
          Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ E2Ext LocVar Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocVar Ty2 -> Exp2) -> E2Ext LocVar Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ LocVar -> PreLocExp LocVar -> Exp2 -> E2Ext LocVar Ty2
forall loc dec.
LocVar -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE LocVar
loc PreLocExp LocVar
locexp Exp2
bod'
        L2.StartOfPkdCursor{}-> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        L2.TagCursor{}-> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        RetE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        L2.AddFixed{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        FromEndE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        BoundsCheck{}  -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        IndirectionE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        E2Ext LocVar Ty2
GetCilkWorkerNum-> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        LetAvail [LocVar]
a Exp2
bod  -> do
          Exp2
bod' <- Exp2 -> PassM Exp2
bindReturns Exp2
bod
          Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ E2Ext LocVar Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocVar Ty2 -> Exp2) -> E2Ext LocVar Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ [LocVar] -> Exp2 -> E2Ext LocVar Ty2
forall loc dec. [LocVar] -> E2 loc dec -> E2Ext loc dec
LetAvail [LocVar]
a Exp2
bod'
        AllocateTagHere{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        AllocateScalarsHere{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        SSPush{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        SSPop{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    MapE{}  -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error (String -> PassM Exp2) -> String -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ String
"bindReturns: TODO MapE"
    FoldE{} -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error (String -> PassM Exp2) -> String -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ String
"bindReturns: TODO FoldE"

handleScalarRet :: Exp2 -> (Exp2 -> Exp2) -> PassM Exp2
handleScalarRet :: Exp2 -> (Exp2 -> Exp2) -> PassM Exp2
handleScalarRet Exp2
bod Exp2 -> Exp2
fn = do
  let bind_and_recur :: Exp2 -> Ty2 -> PassM Exp2
bind_and_recur Exp2
bind_e Ty2
bind_ty = do
        LocVar
tmp <- LocVar -> PassM LocVar
forall (m :: * -> *). MonadState Int m => LocVar -> m LocVar
gensym LocVar
"fltScalar"
        let e1 :: Exp2
e1 = ((LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
tmp,[],Ty2
bind_ty,Exp2
bind_e) (LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
tmp))
        Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ Exp2 -> Exp2
fn Exp2
e1
  case Exp2
bod of
    LitE Int
n -> Exp2 -> Ty2 -> PassM Exp2
bind_and_recur (Int -> Exp2
forall (ext :: * -> * -> *) loc dec. Int -> PreExp ext loc dec
LitE Int
n) Ty2
forall loc. UrTy loc
IntTy
    CharE Char
n -> Exp2 -> Ty2 -> PassM Exp2
bind_and_recur (Char -> Exp2
forall (ext :: * -> * -> *) loc dec. Char -> PreExp ext loc dec
CharE Char
n) Ty2
forall loc. UrTy loc
CharTy
    FloatE Double
n -> Exp2 -> Ty2 -> PassM Exp2
bind_and_recur (Double -> Exp2
forall (ext :: * -> * -> *) loc dec. Double -> PreExp ext loc dec
FloatE Double
n) Ty2
forall loc. UrTy loc
FloatTy
    LitSymE LocVar
n -> Exp2 -> Ty2 -> PassM Exp2
bind_and_recur (LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
LitSymE LocVar
n) Ty2
forall loc. UrTy loc
SymTy
    PrimAppE Prim Ty2
MkTrue [] -> Exp2 -> Ty2 -> PassM Exp2
bind_and_recur (Prim Ty2 -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim Ty2
forall ty. Prim ty
MkTrue []) Ty2
forall loc. UrTy loc
BoolTy
    PrimAppE Prim Ty2
MkFalse [] -> Exp2 -> Ty2 -> PassM Exp2
bind_and_recur (Prim Ty2 -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim Ty2
forall ty. Prim ty
MkFalse []) Ty2
forall loc. UrTy loc
BoolTy
    Exp2
_ -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ Exp2 -> Exp2
fn Exp2
bod
    -- _ -> error $ "RouteEnds: Not scalar " ++ sdoc bod

isScalar :: Exp2 -> Bool
isScalar :: Exp2 -> Bool
isScalar Exp2
e1 =
  case Exp2
e1 of
    LitE{} -> Bool
True
    CharE{} -> Bool
True
    FloatE{} -> Bool
True
    LitSymE{} -> Bool
True
    PrimAppE Prim Ty2
MkTrue [] -> Bool
True
    PrimAppE Prim Ty2
MkFalse [] -> Bool
True
    Exp2
_ -> Bool
False

-- | Process an L2 Prog and thread through explicit end-witnesses.
-- Requires Gensym and runs in PassM. Assumes the Prog has been flattened.
routeEnds :: Prog2 -> PassM Prog2
routeEnds :: Prog2 -> PassM Prog2
routeEnds prg :: Prog2
prg@Prog{DDefs (TyOf Exp2)
ddefs :: DDefs (TyOf Exp2)
ddefs :: forall ex. Prog ex -> DDefs (TyOf ex)
ddefs,FunDefs Exp2
fundefs :: FunDefs Exp2
fundefs :: forall ex. Prog ex -> FunDefs ex
fundefs,Maybe (Exp2, TyOf Exp2)
mainExp :: Maybe (Exp2, TyOf Exp2)
mainExp :: forall ex. Prog ex -> Maybe (ex, TyOf ex)
mainExp} = do

  -- Handle functions in two steps (to account for mutual recursion):
  --
  -- First, compute the new types, and build a new fundefs structure:
  [FunDef2]
fds' <- (FunDef2 -> PassM FunDef2) -> [FunDef2] -> PassM [FunDef2]
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 FunDef2 -> PassM FunDef2
fdty ([FunDef2] -> PassM [FunDef2]) -> [FunDef2] -> PassM [FunDef2]
forall a b. (a -> b) -> a -> b
$ FunDefs Exp2 -> [FunDef2]
forall k a. Map k a -> [a]
M.elems FunDefs Exp2
fundefs
  let fundefs' :: FunDefs Exp2
fundefs' = [(LocVar, FunDef2)] -> FunDefs Exp2
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(LocVar, FunDef2)] -> FunDefs Exp2)
-> [(LocVar, FunDef2)] -> FunDefs Exp2
forall a b. (a -> b) -> a -> b
$ (FunDef2 -> (LocVar, FunDef2)) -> [FunDef2] -> [(LocVar, FunDef2)]
forall a b. (a -> b) -> [a] -> [b]
L.map (\FunDef2
f -> (FunDef2 -> LocVar
forall ex. FunDef ex -> LocVar
funName FunDef2
f,FunDef2
f)) [FunDef2]
fds'
  -- Then process the actual function bodies using the new fundefs structure:
  [FunDef2]
fds'' <- (FunDef2 -> PassM FunDef2) -> [FunDef2] -> PassM [FunDef2]
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 (FunDefs Exp2 -> FunDef2 -> PassM FunDef2
fd FunDefs Exp2
fundefs') [FunDef2]
fds'
  let fundefs'' :: FunDefs Exp2
fundefs'' = [(LocVar, FunDef2)] -> FunDefs Exp2
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(LocVar, FunDef2)] -> FunDefs Exp2)
-> [(LocVar, FunDef2)] -> FunDefs Exp2
forall a b. (a -> b) -> a -> b
$ (FunDef2 -> (LocVar, FunDef2)) -> [FunDef2] -> [(LocVar, FunDef2)]
forall a b. (a -> b) -> [a] -> [b]
L.map (\FunDef2
f -> (FunDef2 -> LocVar
forall ex. FunDef ex -> LocVar
funName FunDef2
f,FunDef2
f)) [FunDef2]
fds''
      env2 :: Env2 (TyOf Exp2)
env2 = Prog2 -> Env2 (TyOf Exp2)
forall a. Prog a -> Env2 (TyOf a)
progToEnv Prog2
prg
  -- Handle the main expression (if it exists):
  Maybe (Exp2, Ty2)
mainExp' <- case Maybe (Exp2, TyOf Exp2)
mainExp of
                Maybe (Exp2, TyOf Exp2)
Nothing -> Maybe (Exp2, Ty2) -> PassM (Maybe (Exp2, Ty2))
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Exp2, Ty2)
forall a. Maybe a
Nothing
                Just (Exp2
e,TyOf Exp2
t) -> do Exp2
e' <- Exp2 -> PassM Exp2
bindReturns Exp2
e
                                 Exp2
e'' <- FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fundefs'' [] EndOfRel
emptyRel Map LocVar LocVar
forall k a. Map k a
M.empty Map LocVar LocVar
forall k a. Map k a
M.empty Env2 (TyOf Exp2)
Env2 Ty2
env2 Exp2
e'
                                 Maybe (Exp2, Ty2) -> PassM (Maybe (Exp2, Ty2))
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Exp2, Ty2) -> PassM (Maybe (Exp2, Ty2)))
-> Maybe (Exp2, Ty2) -> PassM (Maybe (Exp2, Ty2))
forall a b. (a -> b) -> a -> b
$ (Exp2, Ty2) -> Maybe (Exp2, Ty2)
forall a. a -> Maybe a
Just (Exp2
e'',TyOf Exp2
Ty2
t)

  -- Return the updated Prog
  Prog2 -> PassM Prog2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (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, Ty2)
mainExp'


  where
    -- Helper functions:

    -- | Process function types (but don't handle bodies)
    fdty :: L2.FunDef2 -> PassM L2.FunDef2
    fdty :: FunDef2 -> PassM FunDef2
fdty FunDef{LocVar
funName :: forall ex. FunDef ex -> LocVar
funName :: LocVar
funName,ArrowTy (TyOf Exp2)
funTy :: ArrowTy (TyOf Exp2)
funTy :: forall ex. FunDef ex -> ArrowTy (TyOf ex)
funTy,[LocVar]
funArgs :: [LocVar]
funArgs :: forall ex. FunDef ex -> [LocVar]
funArgs,Exp2
funBody :: Exp2
funBody :: forall ex. FunDef ex -> ex
funBody,FunMeta
funMeta :: FunMeta
funMeta :: forall ex. FunDef ex -> FunMeta
funMeta} =
        do let (ArrowTy2 [LRM]
locin [Ty2]
tyin Set Effect
eff Ty2
tyout [LocRet]
_locout Bool
isPar) = ArrowTy (TyOf Exp2)
funTy
               handleLoc :: LRM -> [LRM] -> [LRM]
handleLoc (LRM LocVar
l Region
r Modality
m) [LRM]
ls = if Effect -> Set Effect -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member (LocVar -> Effect
Traverse LocVar
l) Set Effect
eff then (LocVar -> Region -> Modality -> LRM
LRM LocVar
l Region
r Modality
m)LRM -> [LRM] -> [LRM]
forall a. a -> [a] -> [a]
:[LRM]
ls else [LRM]
ls
               locout' :: [LocRet]
locout' = (LRM -> LocRet) -> [LRM] -> [LocRet]
forall a b. (a -> b) -> [a] -> [b]
L.map LRM -> LocRet
EndOf ([LRM] -> [LocRet]) -> [LRM] -> [LocRet]
forall a b. (a -> b) -> a -> b
$ (LRM -> [LRM] -> [LRM]) -> [LRM] -> [LRM] -> [LRM]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
L.foldr LRM -> [LRM] -> [LRM]
handleLoc [] [LRM]
locin
           FunDef2 -> PassM FunDef2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return FunDef{LocVar
funName :: LocVar
funName :: LocVar
funName,funTy :: ArrowTy (TyOf Exp2)
funTy=([LRM]
-> [Ty2] -> Set Effect -> Ty2 -> [LocRet] -> Bool -> ArrowTy2 Ty2
forall ty2.
[LRM]
-> [ty2] -> Set Effect -> ty2 -> [LocRet] -> Bool -> ArrowTy2 ty2
ArrowTy2 [LRM]
locin [Ty2]
tyin Set Effect
eff Ty2
tyout [LocRet]
locout' Bool
isPar),[LocVar]
funArgs :: [LocVar]
funArgs :: [LocVar]
funArgs,Exp2
funBody :: Exp2
funBody :: Exp2
funBody,FunMeta
funMeta :: FunMeta
funMeta :: FunMeta
funMeta}


    -- | Process function bodies
    fd :: FunDefs2 -> L2.FunDef2 -> PassM L2.FunDef2
    fd :: FunDefs Exp2 -> FunDef2 -> PassM FunDef2
fd FunDefs Exp2
fns FunDef{LocVar
funName :: forall ex. FunDef ex -> LocVar
funName :: LocVar
funName,ArrowTy (TyOf Exp2)
funTy :: forall ex. FunDef ex -> ArrowTy (TyOf ex)
funTy :: ArrowTy (TyOf Exp2)
funTy,[LocVar]
funArgs :: forall ex. FunDef ex -> [LocVar]
funArgs :: [LocVar]
funArgs,Exp2
funBody :: forall ex. FunDef ex -> ex
funBody :: Exp2
funBody,FunMeta
funMeta :: forall ex. FunDef ex -> FunMeta
funMeta :: FunMeta
funMeta} =
        do let (ArrowTy2 [LRM]
locin [Ty2]
tyins Set Effect
eff Ty2
_tyout [LocRet]
_locout Bool
_isPar) = ArrowTy (TyOf Exp2)
funTy
               handleLoc :: LRM -> [LocVar] -> [LocVar]
handleLoc (LRM LocVar
l Region
_r Modality
_m) [LocVar]
ls = if Effect -> Set Effect -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member (LocVar -> Effect
Traverse LocVar
l) Set Effect
eff then LocVar
lLocVar -> [LocVar] -> [LocVar]
forall a. a -> [a] -> [a]
:[LocVar]
ls else [LocVar]
ls
               retlocs :: [LocVar]
retlocs = (LRM -> [LocVar] -> [LocVar]) -> [LocVar] -> [LRM] -> [LocVar]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
L.foldr LRM -> [LocVar] -> [LocVar]
handleLoc [] [LRM]
locin
               lenv :: Map LocVar LocVar
lenv = ((LocVar, Ty2) -> Map LocVar LocVar -> Map LocVar LocVar)
-> Map LocVar LocVar -> [(LocVar, Ty2)] -> Map LocVar LocVar
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
L.foldr
                        (\(LocVar
a,Ty2
t) Map LocVar LocVar
acc -> case Ty2
t of
                                         PackedTy String
_ LocVar
loc -> LocVar -> LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
a LocVar
loc Map LocVar LocVar
acc
                                         Ty2
_ -> Map LocVar LocVar
acc)
                        Map LocVar LocVar
forall k a. Map k a
M.empty ([LocVar] -> [Ty2] -> [(LocVar, Ty2)]
forall a b. [a] -> [b] -> [(a, b)]
zip [LocVar]
funArgs [Ty2]
tyins)
               pakdLocs :: [UrTy a] -> [(a, UrTy a)]
pakdLocs = (UrTy a -> [(a, UrTy a)]) -> [UrTy a] -> [(a, UrTy a)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap
                            (\UrTy a
t -> case UrTy a
t of
                                     PackedTy String
_ a
loc -> [(a
loc, UrTy a
t)]
                                     ProdTy [UrTy a]
ys -> [UrTy a] -> [(a, UrTy a)]
pakdLocs [UrTy a]
ys
                                     UrTy a
_ -> [])
               initVEnv :: Map LocVar Ty2
initVEnv = [(LocVar, Ty2)] -> Map LocVar Ty2
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(LocVar, Ty2)] -> Map LocVar Ty2)
-> [(LocVar, Ty2)] -> Map LocVar Ty2
forall a b. (a -> b) -> a -> b
$ [Ty2] -> [(LocVar, Ty2)]
forall {a}. [UrTy a] -> [(a, UrTy a)]
pakdLocs [Ty2]
tyins [(LocVar, Ty2)] -> [(LocVar, Ty2)] -> [(LocVar, Ty2)]
forall a. [a] -> [a] -> [a]
++ [LocVar] -> [Ty2] -> [(LocVar, Ty2)]
forall a b. [a] -> [b] -> [(a, b)]
zip [LocVar]
funArgs [Ty2]
tyins
               env2 :: Env2 Ty2
env2 = Map LocVar Ty2 -> TyEnv (ArrowTy Ty2) -> Env2 Ty2
forall a. TyEnv a -> TyEnv (ArrowTy a) -> Env2 a
Env2 Map LocVar Ty2
initVEnv (FunDefs Exp2 -> TyEnv (ArrowTy (TyOf Exp2))
forall a. FunDefs a -> TyEnv (ArrowTy (TyOf a))
initFunEnv FunDefs Exp2
fundefs)
           Exp2
funBody' <- Exp2 -> PassM Exp2
bindReturns Exp2
funBody
           Exp2
funBody'' <- FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
emptyRel Map LocVar LocVar
lenv Map LocVar LocVar
forall k a. Map k a
M.empty Env2 Ty2
env2 Exp2
funBody'
           FunDef2 -> PassM FunDef2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return FunDef{LocVar
funName :: LocVar
funName :: LocVar
funName,ArrowTy (TyOf Exp2)
funTy :: ArrowTy (TyOf Exp2)
funTy :: ArrowTy (TyOf Exp2)
funTy,[LocVar]
funArgs :: [LocVar]
funArgs :: [LocVar]
funArgs,funBody :: Exp2
funBody=Exp2
funBody'',FunMeta
funMeta :: FunMeta
funMeta :: FunMeta
funMeta}


    -- | Process expressions.
    -- Takes the following arguments:
    -- 1. a function environment
    -- 2. a list of locations we need to return the ends of
    -- 3. an end-of relation
    -- 4. a map of var to location
    -- 5. a map from location to location after it
    -- 6. the expression to process
    exp :: FunDefs2 -> [LocVar] -> EndOfRel -> M.Map Var LocVar ->
           M.Map LocVar LocVar -> Env2 Ty2 -> Exp2 -> PassM Exp2
    exp :: FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor Map LocVar LocVar
lenv Map LocVar LocVar
afterenv Env2 Ty2
env2 Exp2
e =
        case Exp2
e of

          -- Variable case, *should* be the base case assuming our expression was
          -- properly put in ANF.
          -- We generate our RetE form here. By this point we should know the ends
          -- of each of the locactions in relocs.

          -- we fmap location at the top-level case expression
          VarE LocVar
v -> [LocVar] -> Exp2 -> PassM Exp2
mkRet [LocVar]
retlocs (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
v

          -- -- If a function has a literal as it's tail, it cannot return any
          -- -- end witnesses (since only VarE forms are converted to RetE's).
          -- -- We therefore bind that literal to a variable, and recur. We need
          -- -- to handle LetRegionE and LetLocE similarly.
          -- LetE (v,ls,ty,rhs) bod | isScalar bod -> do
          --   handleScalarRet bod (\bod' -> LetE (v,ls,ty,rhs) bod')

          -- Ext (LetRegionE r bod) | isScalar bod -> do
          --   handleScalarRet bod (\bod' -> Ext (LetRegionE r bod'))

          -- Ext (LetLocE v locexp bod) | isScalar bod -> do
          --   handleScalarRet bod (\bod' -> Ext (LetLocE v locexp bod'))

          -- This is the most interesting case: a let bound function application.
          -- We need to update the let binding's extra location binding list with
          -- the end witnesses returned from the function.
          LetE (LocVar
v,[LocVar]
_ls,Ty2
ty,(AppE LocVar
f [LocVar]
lsin [Exp2]
e1)) Exp2
e2 -> do
                 let lenv' :: Map LocVar LocVar
lenv' = case Ty2
ty of
                               PackedTy String
_n LocVar
l -> LocVar -> LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
v LocVar
l Map LocVar LocVar
lenv
                               Ty2
_ -> Map LocVar LocVar
lenv

                 ([LocVar]
outlocs,[(LocVar, LocVar)]
newls,EndOfRel
eor') <- LocVar
-> [LocVar] -> PassM ([LocVar], [(LocVar, LocVar)], EndOfRel)
doBoundApp LocVar
f [LocVar]
lsin
                 Exp2
e2' <- FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor' Map LocVar LocVar
lenv' Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v Ty2
ty Env2 Ty2
env2) Exp2
e2
                 Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[LocVar]
outlocs,Ty2
ty, LocVar -> [LocVar] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
LocVar -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE LocVar
f [LocVar]
lsin [Exp2]
e1)
                               (Exp2 -> [(LocVar, LocVar)] -> Exp2
wrapBody Exp2
e2' [(LocVar, LocVar)]
newls)

          -- Exactly like AppE.
          LetE (LocVar
v,[LocVar]
_ls,Ty2
ty,(SpawnE LocVar
f [LocVar]
lsin [Exp2]
e1)) Exp2
e2 -> do
                 let lenv' :: Map LocVar LocVar
lenv' = case Ty2
ty of
                               PackedTy String
_n LocVar
l -> LocVar -> LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
v LocVar
l Map LocVar LocVar
lenv
                               Ty2
_ -> Map LocVar LocVar
lenv
                 ([LocVar]
outlocs,[(LocVar, LocVar)]
newls,EndOfRel
eor') <- LocVar
-> [LocVar] -> PassM ([LocVar], [(LocVar, LocVar)], EndOfRel)
doBoundApp LocVar
f [LocVar]
lsin
                 Exp2
e2' <- FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor' Map LocVar LocVar
lenv' Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v Ty2
ty Env2 Ty2
env2) Exp2
e2
                 Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[LocVar]
outlocs,Ty2
ty, LocVar -> [LocVar] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
LocVar -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
SpawnE LocVar
f [LocVar]
lsin [Exp2]
e1)
                               (Exp2 -> [(LocVar, LocVar)] -> Exp2
wrapBody Exp2
e2' [(LocVar, LocVar)]
newls)

          SpawnE{} -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error String
"routeEnds: Unbound SpawnE"
          Exp2
SyncE    -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
e

          CaseE (VarE LocVar
x) [(String, [(LocVar, LocVar)], Exp2)]
brs -> do
                 -- We will need to gensym while processing the case clauses, so
                 -- it has to be in the PassM monad
                 [(String, [(LocVar, LocVar)], Exp2)]
brs' <-
                     [(String, [(LocVar, LocVar)], Exp2)]
-> ((String, [(LocVar, LocVar)], Exp2)
    -> PassM (String, [(LocVar, LocVar)], Exp2))
-> PassM [(String, [(LocVar, LocVar)], Exp2)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(String, [(LocVar, LocVar)], Exp2)]
brs (((String, [(LocVar, LocVar)], Exp2)
  -> PassM (String, [(LocVar, LocVar)], Exp2))
 -> PassM [(String, [(LocVar, LocVar)], Exp2)])
-> ((String, [(LocVar, LocVar)], Exp2)
    -> PassM (String, [(LocVar, LocVar)], Exp2))
-> PassM [(String, [(LocVar, LocVar)], Exp2)]
forall a b. (a -> b) -> a -> b
$ \(String
dc, [(LocVar, LocVar)]
vls, Exp2
e) ->
                       case [(LocVar, LocVar)]
vls of
                         [] ->
                           case (LocVar -> Map LocVar LocVar -> Maybe LocVar
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup LocVar
x Map LocVar LocVar
lenv) of
                             Just LocVar
l1 -> do
                               LocVar
l2 <- LocVar -> PassM LocVar
forall (m :: * -> *). MonadState Int m => LocVar -> m LocVar
gensym LocVar
"jump"
                               let eor' :: EndOfRel
eor' = LocVar -> LocVar -> EndOfRel -> EndOfRel
mkEnd LocVar
l1 LocVar
l2 EndOfRel
eor
                                   e' :: Exp2
e' = E2Ext LocVar Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocVar Ty2 -> Exp2) -> E2Ext LocVar Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ LocVar -> PreLocExp LocVar -> Exp2 -> E2Ext LocVar Ty2
forall loc dec.
LocVar -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE LocVar
l2 (Int -> LocVar -> PreLocExp LocVar
forall loc. Int -> loc -> PreLocExp loc
AfterConstantLE Int
1 LocVar
l1) Exp2
e
                               Exp2
e'' <- FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor' Map LocVar LocVar
lenv (LocVar -> LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
l1 LocVar
l2 Map LocVar LocVar
lenv) Env2 Ty2
env2 Exp2
e'
                               (String, [(LocVar, LocVar)], Exp2)
-> PassM (String, [(LocVar, LocVar)], Exp2)
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (String
dc, [(LocVar, LocVar)]
vls, Exp2
e'')
                             Maybe LocVar
Nothing -> String -> PassM (String, [(LocVar, LocVar)], Exp2)
forall a. HasCallStack => String -> a
error (String -> PassM (String, [(LocVar, LocVar)], Exp2))
-> String -> PassM (String, [(LocVar, LocVar)], Exp2)
forall a b. (a -> b) -> a -> b
$ String
"Failed to find " String -> ShowS
forall a. [a] -> [a] -> [a]
++ LocVar -> String
forall a. Out a => a -> String
sdoc LocVar
x String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" in " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Map LocVar LocVar -> String
forall a. Out a => a -> String
sdoc Map LocVar LocVar
lenv
                         [(LocVar, LocVar)]
_ -> do
                           let need :: LocVar
need = (LocVar, LocVar) -> LocVar
forall a b. (a, b) -> b
snd ((LocVar, LocVar) -> LocVar) -> (LocVar, LocVar) -> LocVar
forall a b. (a -> b) -> a -> b
$ [(LocVar, LocVar)] -> (LocVar, LocVar)
forall a. HasCallStack => [a] -> a
last [(LocVar, LocVar)]
vls
                               argtys :: [Ty2]
argtys = DDefs Ty2 -> String -> [Ty2]
forall a. Out a => DDefs a -> String -> [a]
lookupDataCon DDefs (TyOf Exp2)
DDefs Ty2
ddefs String
dc
                               lx :: LocVar
lx = case LocVar -> Map LocVar LocVar -> Maybe LocVar
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup LocVar
x Map LocVar LocVar
lenv of
                                      Maybe LocVar
Nothing -> String -> LocVar
forall a. HasCallStack => String -> a
error (String -> LocVar) -> String -> LocVar
forall a b. (a -> b) -> a -> b
$ String
"Failed to find " String -> ShowS
forall a. [a] -> [a] -> [a]
++ (LocVar -> String
forall a. Show a => a -> String
show LocVar
x)
                                      Just LocVar
l -> LocVar
l
                               -- we know lx and need have the same end, since
                               -- lx is the whole packed thing and need is its
                               -- last field, so when we look up the end of lx
                               -- what we really want is the end of need.
                               eor' :: EndOfRel
eor' = LocVar -> LocVar -> EndOfRel -> EndOfRel
mkEqual LocVar
lx LocVar
need EndOfRel
eor
                               f :: (k, a) -> Map k a -> Map k a
f (k
l1,a
l2) Map k a
env = k -> a -> Map k a -> Map k a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
l1 a
l2 Map k a
env
                               afterenv' :: Map LocVar LocVar
afterenv' = ((LocVar, LocVar) -> Map LocVar LocVar -> Map LocVar LocVar)
-> Map LocVar LocVar -> [(LocVar, LocVar)] -> Map LocVar LocVar
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
L.foldr (LocVar, LocVar) -> Map LocVar LocVar -> Map LocVar LocVar
forall {k} {a}. Ord k => (k, a) -> Map k a -> Map k a
f Map LocVar LocVar
afterenv ([(LocVar, LocVar)] -> Map LocVar LocVar)
-> [(LocVar, LocVar)] -> Map LocVar LocVar
forall a b. (a -> b) -> a -> b
$ [LocVar] -> [LocVar] -> [(LocVar, LocVar)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((LocVar, LocVar) -> LocVar) -> [(LocVar, LocVar)] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map (LocVar, LocVar) -> LocVar
forall a b. (a, b) -> b
snd [(LocVar, LocVar)]
vls) ([LocVar] -> [LocVar]
forall a. HasCallStack => [a] -> [a]
tail ([LocVar] -> [LocVar]) -> [LocVar] -> [LocVar]
forall a b. (a -> b) -> a -> b
$ ((LocVar, LocVar) -> LocVar) -> [(LocVar, LocVar)] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map (LocVar, LocVar) -> LocVar
forall a b. (a, b) -> b
snd [(LocVar, LocVar)]
vls)
                               -- two cases here for handing bound parameters:
                               -- we have a packed type:
                               handleLoc :: (EndOfRel, PreExp E2Ext LocVar dec)
-> (LocVar, UrTy a) -> m (EndOfRel, PreExp E2Ext LocVar dec)
handleLoc (EndOfRel
eor,PreExp E2Ext LocVar dec
e) (LocVar
_,(PackedTy String
_ a
_)) = (EndOfRel, PreExp E2Ext LocVar dec)
-> m (EndOfRel, PreExp E2Ext LocVar dec)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (EndOfRel
eor,PreExp E2Ext LocVar dec
e)
                               -- or we have a non-packed type, and we need to "jump" over it and
                               -- bind a location to after it
                               handleLoc (EndOfRel
eor,PreExp E2Ext LocVar dec
e) (LocVar
l1,UrTy a
ty) = do
                                    LocVar
l2 <- LocVar -> m LocVar
forall (m :: * -> *). MonadState Int m => LocVar -> m LocVar
gensym LocVar
"jump"
                                    let eor' :: EndOfRel
eor' = LocVar -> LocVar -> EndOfRel -> EndOfRel
mkEnd LocVar
l1 LocVar
l2 EndOfRel
eor
                                        (Just Int
jump) = UrTy a -> Maybe Int
forall a. UrTy a -> Maybe Int
L1.sizeOfTy UrTy a
ty
                                        e' :: PreExp E2Ext LocVar dec
e' = E2Ext LocVar dec -> PreExp E2Ext LocVar dec
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocVar dec -> PreExp E2Ext LocVar dec)
-> E2Ext LocVar dec -> PreExp E2Ext LocVar dec
forall a b. (a -> b) -> a -> b
$ LocVar
-> PreLocExp LocVar -> PreExp E2Ext LocVar dec -> E2Ext LocVar dec
forall loc dec.
LocVar -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE LocVar
l2 (Int -> LocVar -> PreLocExp LocVar
forall loc. Int -> loc -> PreLocExp loc
AfterConstantLE Int
jump LocVar
l1) PreExp E2Ext LocVar dec
e
                                    (EndOfRel, PreExp E2Ext LocVar dec)
-> m (EndOfRel, PreExp E2Ext LocVar dec)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (EndOfRel
eor', PreExp E2Ext LocVar dec
e')
                               vars :: [LocVar]
vars = ((LocVar, LocVar) -> LocVar) -> [(LocVar, LocVar)] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map (LocVar, LocVar) -> LocVar
forall a b. (a, b) -> a
fst [(LocVar, LocVar)]
vls
                               locs :: [LocVar]
locs = ((LocVar, LocVar) -> LocVar) -> [(LocVar, LocVar)] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map (LocVar, LocVar) -> LocVar
forall a b. (a, b) -> b
snd [(LocVar, LocVar)]
vls
                               env2' :: Env2 Ty2
env2' = Map LocVar Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. Map LocVar a -> Env2 a -> Env2 a
extendsVEnv ([(LocVar, Ty2)] -> Map LocVar Ty2
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([LocVar] -> [Ty2] -> [(LocVar, Ty2)]
forall a b. [a] -> [b] -> [(a, b)]
zip [LocVar]
locs [Ty2]
argtys [(LocVar, Ty2)] -> [(LocVar, Ty2)] -> [(LocVar, Ty2)]
forall a. [a] -> [a] -> [a]
++ [LocVar] -> [Ty2] -> [(LocVar, Ty2)]
forall a b. [a] -> [b] -> [(a, b)]
zip [LocVar]
vars [Ty2]
argtys)) Env2 Ty2
env2
                               lenv' :: Map LocVar LocVar
lenv' = Map LocVar LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Map LocVar LocVar
lenv (Map LocVar LocVar -> Map LocVar LocVar)
-> Map LocVar LocVar -> Map LocVar LocVar
forall a b. (a -> b) -> a -> b
$ [(LocVar, LocVar)] -> Map LocVar LocVar
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(LocVar, LocVar)]
vls
                           (EndOfRel
eor'',Exp2
e') <- ((EndOfRel, Exp2) -> (LocVar, Ty2) -> PassM (EndOfRel, Exp2))
-> (EndOfRel, Exp2) -> [(LocVar, Ty2)] -> PassM (EndOfRel, Exp2)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (EndOfRel, Exp2) -> (LocVar, Ty2) -> PassM (EndOfRel, Exp2)
forall {m :: * -> *} {dec} {a}.
MonadState Int m =>
(EndOfRel, PreExp E2Ext LocVar dec)
-> (LocVar, UrTy a) -> m (EndOfRel, PreExp E2Ext LocVar dec)
handleLoc (EndOfRel
eor',Exp2
e) ([(LocVar, Ty2)] -> PassM (EndOfRel, Exp2))
-> [(LocVar, Ty2)] -> PassM (EndOfRel, Exp2)
forall a b. (a -> b) -> a -> b
$ [LocVar] -> [Ty2] -> [(LocVar, Ty2)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((LocVar, LocVar) -> LocVar) -> [(LocVar, LocVar)] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map (LocVar, LocVar) -> LocVar
forall a b. (a, b) -> b
snd [(LocVar, LocVar)]
vls) [Ty2]
argtys
                           Exp2
e'' <- FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor'' Map LocVar LocVar
lenv' Map LocVar LocVar
afterenv' Env2 Ty2
env2' Exp2
e'
                           (String, [(LocVar, LocVar)], Exp2)
-> PassM (String, [(LocVar, LocVar)], Exp2)
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (String
dc, [(LocVar, LocVar)]
vls, Exp2
e'')
                 Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ Exp2 -> [(String, [(LocVar, LocVar)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(String, [(LocVar, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE (LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
x) [(String, [(LocVar, LocVar)], Exp2)]
brs'



          CaseE Exp2
complex [(String, [(LocVar, LocVar)], Exp2)]
brs -> do
            let ty :: TyOf Exp2
ty = DDefs (TyOf Exp2) -> Env2 (TyOf Exp2) -> Exp2 -> TyOf Exp2
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf Exp2)
ddefs Env2 (TyOf Exp2)
Env2 Ty2
env2 Exp2
complex
            LocVar
v <- LocVar -> PassM LocVar
forall (m :: * -> *). MonadState Int m => LocVar -> m LocVar
gensym LocVar
"flt_RE"
            let ex :: Exp2
ex = [(LocVar, [LocVar], Ty2, Exp2)] -> Exp2 -> Exp2
forall loc dec (ext :: * -> * -> *).
[(LocVar, [loc], dec, PreExp ext loc dec)]
-> PreExp ext loc dec -> PreExp ext loc dec
L1.mkLets [(LocVar
v,[],TyOf Exp2
Ty2
ty,Exp2
complex)] (Exp2 -> [(String, [(LocVar, LocVar)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(String, [(LocVar, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE (LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
v) [(String, [(LocVar, LocVar)], Exp2)]
brs)
            FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor Map LocVar LocVar
lenv Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v TyOf Exp2
Ty2
ty Env2 Ty2
env2) Exp2
ex


          -- This shouldn't happen, but as a convenience we can ANF-ify this AppE
          -- by gensyming a new variable, sticking the AppE in a LetE, and recuring.
          -- Question: should this fail instead? I'm not sure.
          AppE LocVar
v [LocVar]
args [Exp2]
arg -> do
                 LocVar
v' <- LocVar -> PassM LocVar
forall (m :: * -> *). MonadState Int m => LocVar -> m LocVar
gensym LocVar
"tailapp"
                 let ty :: TyOf Exp2
ty = DDefs (TyOf Exp2) -> Env2 (TyOf Exp2) -> Exp2 -> TyOf Exp2
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf Exp2)
ddefs Env2 (TyOf Exp2)
Env2 Ty2
env2 Exp2
e
                     e' :: Exp2
e' = (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v',[], TyOf Exp2
Ty2
ty, LocVar -> [LocVar] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
LocVar -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE LocVar
v [LocVar]
args [Exp2]
arg) (LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
v')
                 Exp2 -> PassM Exp2
go (Exp2
e')

          PrimAppE (DictInsertP Ty2
dty) [(VarE LocVar
a),Exp2
d,Exp2
k,Exp2
v] -> do
                 LocVar
v' <- LocVar -> PassM LocVar
forall (m :: * -> *). MonadState Int m => LocVar -> m LocVar
gensym LocVar
"tailprim"
                 let e' :: Exp2
e' = (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v',[],Maybe LocVar -> UrTy () -> Ty2
forall loc. Maybe LocVar -> UrTy () -> UrTy loc
SymDictTy (LocVar -> Maybe LocVar
forall a. a -> Maybe a
Just LocVar
a) (UrTy () -> Ty2) -> UrTy () -> Ty2
forall a b. (a -> b) -> a -> b
$ Ty2 -> UrTy ()
forall a. UrTy a -> UrTy ()
stripTyLocs Ty2
dty, Prim Ty2 -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE (Ty2 -> Prim Ty2
forall ty. ty -> Prim ty
DictInsertP Ty2
dty) [(LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
a),Exp2
d,Exp2
k,Exp2
v]) (LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
v')
                 -- we fmap location at the top-level case expression
                 Exp2 -> PassM Exp2
go (Exp2
e')

          -- Same AppE as above. This could just fail, instead of trying to repair
          -- the program.
          PrimAppE Prim Ty2
pr [Exp2]
es -> do
                 LocVar
v' <- LocVar -> PassM LocVar
forall (m :: * -> *). MonadState Int m => LocVar -> m LocVar
gensym LocVar
"tailprim"
                 let ty :: Ty2
ty = Prim Ty2 -> Ty2
forall a. Prim (UrTy a) -> UrTy a
L1.primRetTy Prim Ty2
pr
                     e' :: Exp2
e' = (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v',[],Ty2
ty, Prim Ty2 -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim Ty2
pr [Exp2]
es) (LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
v')
                 -- we fmap location at the top-level case expression
                 Exp2 -> PassM Exp2
go (Exp2
e')

          -- RouteEnds creates let bindings for such expressions (see those cases below).
          -- Processing the RHS here would cause an infinite loop.

          LetE (LocVar
v,[LocVar]
ls,ty :: Ty2
ty@(PackedTy String
_ LocVar
loc),e1 :: Exp2
e1@DataConE{}) Exp2
e2 -> do
            Exp2
e2' <- FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor (LocVar -> LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
v LocVar
loc Map LocVar LocVar
lenv) Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v Ty2
ty Env2 Ty2
env2) Exp2
e2
            Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[LocVar]
ls,Ty2
ty,Exp2
e1) Exp2
e2'

          LetE (LocVar
v,[LocVar]
ls,ty :: Ty2
ty@(PackedTy String
_ LocVar
loc),e1 :: Exp2
e1@(PrimAppE (ReadPackedFile{}) [])) Exp2
e2 -> do
            Exp2
e2' <- FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor (LocVar -> LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
v LocVar
loc Map LocVar LocVar
lenv) Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v Ty2
ty Env2 Ty2
env2) Exp2
e2
            Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[LocVar]
ls,Ty2
ty,Exp2
e1) Exp2
e2'

          LetE (LocVar
v,[LocVar]
ls,Ty2
ty,e1 :: Exp2
e1@ProjE{}) Exp2
e2 -> do
            let lenv' :: Map LocVar LocVar
lenv' = case Ty2
ty of
                          PackedTy String
_ LocVar
loc -> LocVar -> LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
v LocVar
loc Map LocVar LocVar
lenv
                          Ty2
_ -> Map LocVar LocVar
lenv
            Exp2
e2' <- FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor Map LocVar LocVar
lenv' Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v Ty2
ty Env2 Ty2
env2) Exp2
e2
            Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[LocVar]
ls,Ty2
ty,Exp2
e1) Exp2
e2'

          LetE (LocVar
v,[LocVar]
ls,Ty2
ty,e1 :: Exp2
e1@MkProdE{}) Exp2
e2 -> do
            (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[LocVar]
ls,Ty2
ty,Exp2
e1) (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor Map LocVar LocVar
lenv Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v Ty2
ty Env2 Ty2
env2) Exp2
e2

          LetE (LocVar
v,[LocVar]
ls,Ty2
ty,e1 :: Exp2
e1@(PrimAppE (DictLookupP Ty2
_) [Exp2]
_)) Exp2
e2 -> do
            (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[LocVar]
ls,Ty2
ty,Exp2
e1) (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor Map LocVar LocVar
lenv Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v Ty2
ty Env2 Ty2
env2) Exp2
e2

          --

          LetE (LocVar
v,[LocVar]
ls,ty :: Ty2
ty@(PackedTy String
n LocVar
l),Exp2
e1) Exp2
e2 -> do
                 Exp2
e1' <- Exp2 -> PassM Exp2
go Exp2
e1
                 Exp2
e2' <- FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor (LocVar -> LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
v LocVar
l Map LocVar LocVar
lenv) Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v Ty2
ty Env2 Ty2
env2) Exp2
e2
                 Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[LocVar]
ls,String -> LocVar -> Ty2
forall loc. String -> loc -> UrTy loc
PackedTy String
n LocVar
l,Exp2
e1') Exp2
e2'

          LetE (LocVar
v,[LocVar]
ls,Ty2
ty,e1 :: Exp2
e1@TimeIt{}) Exp2
e2 -> do
                 Exp2
e1' <- Exp2 -> PassM Exp2
go Exp2
e1
                 Exp2
e2' <- FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor Map LocVar LocVar
lenv Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v Ty2
ty Env2 Ty2
env2) Exp2
e2
                 Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[LocVar]
ls,Ty2
ty,Exp2
e1') Exp2
e2'

          -- Most boring LetE case, just recur on body
          LetE (LocVar
v,[LocVar]
ls,Ty2
ty,Exp2
rhs) Exp2
bod -> do
            Exp2
bod' <- FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor Map LocVar LocVar
lenv Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v Ty2
ty Env2 Ty2
env2) Exp2
bod
            Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[LocVar]
ls,Ty2
ty,Exp2
rhs) Exp2
bod'

          IfE Exp2
e1 Exp2
e2 Exp2
e3 -> do
                 Exp2
e2' <- Exp2 -> PassM Exp2
go Exp2
e2
                 Exp2
e3' <- Exp2 -> PassM Exp2
go Exp2
e3
                 Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ 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
e1 Exp2
e2' Exp2
e3'

          MkProdE [Exp2]
ls -> do
            let tys :: [Ty2]
tys = (Exp2 -> Ty2) -> [Exp2] -> [Ty2]
forall a b. (a -> b) -> [a] -> [b]
L.map (DDefs (TyOf Exp2) -> Env2 (TyOf Exp2) -> Exp2 -> TyOf Exp2
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf Exp2)
ddefs Env2 (TyOf Exp2)
Env2 Ty2
env2) [Exp2]
ls
                prodty :: Ty2
prodty = [Ty2] -> Ty2
forall loc. [UrTy loc] -> UrTy loc
ProdTy [Ty2]
tys
            LocVar
v <- LocVar -> PassM LocVar
forall (m :: * -> *). MonadState Int m => LocVar -> m LocVar
gensym LocVar
"flt_RE"
            let ex :: Exp2
ex = [(LocVar, [LocVar], Ty2, Exp2)] -> Exp2 -> Exp2
forall loc dec (ext :: * -> * -> *).
[(LocVar, [loc], dec, PreExp ext loc dec)]
-> PreExp ext loc dec -> PreExp ext loc dec
L1.mkLets [(LocVar
v,[],Ty2
prodty,([Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
[PreExp ext loc dec] -> PreExp ext loc dec
MkProdE [Exp2]
ls))] (LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
v)
            FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor Map LocVar LocVar
lenv Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v Ty2
prodty Env2 Ty2
env2) Exp2
ex

          ProjE{} -> do
            LocVar
v <- LocVar -> PassM LocVar
forall (m :: * -> *). MonadState Int m => LocVar -> m LocVar
gensym LocVar
"flt_RE"
            let ty :: TyOf Exp2
ty = DDefs (TyOf Exp2) -> Env2 (TyOf Exp2) -> Exp2 -> TyOf Exp2
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf Exp2)
ddefs Env2 (TyOf Exp2)
Env2 Ty2
env2 Exp2
e
                lenv' :: Map LocVar LocVar
lenv' = case TyOf Exp2
ty of
                          PackedTy String
_ LocVar
loc -> LocVar -> LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
v LocVar
loc Map LocVar LocVar
lenv
                          TyOf Exp2
_ -> Map LocVar LocVar
lenv
                ex :: Exp2
ex = [(LocVar, [LocVar], Ty2, Exp2)] -> Exp2 -> Exp2
forall loc dec (ext :: * -> * -> *).
[(LocVar, [loc], dec, PreExp ext loc dec)]
-> PreExp ext loc dec -> PreExp ext loc dec
L1.mkLets [(LocVar
v,[],TyOf Exp2
Ty2
ty,Exp2
e)] (LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
v)
            FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor Map LocVar LocVar
lenv' Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v TyOf Exp2
Ty2
ty Env2 Ty2
env2) Exp2
ex

          -- Could fail here, but try to fix the broken program
          DataConE LocVar
loc String
dc [Exp2]
es -> do
                 LocVar
v' <- LocVar -> PassM LocVar
forall (m :: * -> *). MonadState Int m => LocVar -> m LocVar
gensym LocVar
"taildc"
                 let ty :: Ty2
ty = String -> LocVar -> Ty2
forall loc. String -> loc -> UrTy loc
PackedTy (DDefs Ty2 -> ShowS
forall a. Out a => DDefs a -> ShowS
getTyOfDataCon DDefs (TyOf Exp2)
DDefs Ty2
ddefs String
dc) LocVar
loc
                     e' :: Exp2
e' = (LocVar, [LocVar], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v',[],Ty2
ty, LocVar -> String -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
loc -> String -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE LocVar
loc String
dc [Exp2]
es)
                               (LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
v')
                 FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor (LocVar -> LocVar -> Map LocVar LocVar -> Map LocVar LocVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
v' LocVar
loc Map LocVar LocVar
lenv) Map LocVar LocVar
afterenv (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v' Ty2
ty Env2 Ty2
env2) (Exp2
e')

          LitE Int
i -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Exp2
forall (ext :: * -> * -> *) loc dec. Int -> PreExp ext loc dec
LitE Int
i)
          CharE Char
i -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Char -> Exp2
forall (ext :: * -> * -> *) loc dec. Char -> PreExp ext loc dec
CharE Char
i)
          FloatE Double
i -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> Exp2
forall (ext :: * -> * -> *) loc dec. Double -> PreExp ext loc dec
FloatE Double
i)

          LitSymE LocVar
v -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ LocVar -> Exp2
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
LitSymE LocVar
v

          TimeIt Exp2
e Ty2
ty Bool
b -> do
                 Exp2
e' <- Exp2 -> PassM Exp2
go Exp2
e
                 Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ Exp2 -> Ty2 -> Bool -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> Bool -> PreExp ext loc dec
TimeIt Exp2
e' Ty2
ty Bool
b

          WithArenaE LocVar
v Exp2
e -> LocVar -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
LocVar -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE LocVar
v (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
e

          Ext (LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
e) -> do
            Exp2
e' <- Exp2 -> PassM Exp2
go Exp2
e
            Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ E2Ext LocVar Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext LocVar Ty2
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
e')

          Ext (LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
e) -> do
            Exp2
e' <- Exp2 -> PassM Exp2
go Exp2
e
            Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ E2Ext LocVar Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext LocVar Ty2
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
e')

          Ext (LetLocE LocVar
v PreLocExp LocVar
locexp Exp2
bod) -> do
            let only_recur :: Exp2 -> PassM Exp2
only_recur Exp2
e = do
                  Exp2
e' <- Exp2 -> PassM Exp2
go Exp2
e
                  Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ E2Ext LocVar Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (LocVar -> PreLocExp LocVar -> Exp2 -> E2Ext LocVar Ty2
forall loc dec.
LocVar -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE LocVar
v PreLocExp LocVar
locexp Exp2
e')
            case PreLocExp LocVar
locexp of
              StartOfRegionLE{} -> Exp2 -> PassM Exp2
only_recur Exp2
bod
              AfterConstantLE{} -> Exp2 -> PassM Exp2
only_recur Exp2
bod
              AfterVariableLE{} -> Exp2 -> PassM Exp2
only_recur Exp2
bod
              InRegionLE{} -> Exp2 -> PassM Exp2
only_recur Exp2
bod
              FreeLE{} -> Exp2 -> PassM Exp2
only_recur Exp2
bod
              PreLocExp LocVar
_ -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error (String -> PassM Exp2) -> String -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ String
"RouteEnds: todo" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Exp2 -> String
forall a. Out a => a -> String
sdoc Exp2
e

          Ext (L2.StartOfPkdCursor{})-> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
e

          Ext (IndirectionE{}) -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
e

          Ext (LetAvail [LocVar]
vs Exp2
e)  -> E2Ext LocVar Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocVar Ty2 -> Exp2)
-> (Exp2 -> E2Ext LocVar Ty2) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [LocVar] -> Exp2 -> E2Ext LocVar Ty2
forall loc dec. [LocVar] -> E2 loc dec -> E2Ext loc dec
LetAvail [LocVar]
vs (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
e

          Ext E2Ext LocVar Ty2
ext -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error (String -> PassM Exp2) -> String -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ String
"RouteEnds: Shouldn't encounter " String -> ShowS
forall a. [a] -> [a] -> [a]
++ E2Ext LocVar Ty2 -> String
forall a. Out a => a -> String
sdoc E2Ext LocVar Ty2
ext

          MapE{} -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error String
"RouteEnds: todo MapE"
          FoldE{} -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error String
"RouteEnds: todo FoldE"

        where  mkRet :: [LocVar] -> Exp2 -> PassM Exp2
               mkRet :: [LocVar] -> Exp2 -> PassM Exp2
mkRet [LocVar]
ls (VarE LocVar
v) =
                 let ends :: [LocVar]
ends = (LocVar -> LocVar) -> [LocVar] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map (\LocVar
l -> LocVar -> EndOfRel -> LocVar
findEnd LocVar
l EndOfRel
eor) [LocVar]
ls
                 in Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ E2Ext LocVar Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext ([LocVar] -> LocVar -> E2Ext LocVar Ty2
forall loc dec. [loc] -> LocVar -> E2Ext loc dec
RetE [LocVar]
ends LocVar
v)
               mkRet [LocVar]
_ Exp2
e = String -> PassM Exp2
forall a. HasCallStack => String -> a
error (String -> PassM Exp2) -> String -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ String
"Expected variable reference in tail call, got "
                           String -> ShowS
forall a. [a] -> [a] -> [a]
++ (Exp2 -> String
forall a. Show a => a -> String
show Exp2
e)

               funtype :: Var -> ArrowTy2 Ty2
               funtype :: LocVar -> ArrowTy2 Ty2
funtype LocVar
v = case LocVar -> FunDefs Exp2 -> Maybe FunDef2
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup LocVar
v FunDefs Exp2
fns of
                             Maybe FunDef2
Nothing -> String -> ArrowTy2 Ty2
forall a. HasCallStack => String -> a
error (String -> ArrowTy2 Ty2) -> String -> ArrowTy2 Ty2
forall a b. (a -> b) -> a -> b
$ String
"Function " String -> ShowS
forall a. [a] -> [a] -> [a]
++ (LocVar -> String
forall a. Show a => a -> String
show LocVar
v) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" not found"
                             Just FunDef2
fundef -> FunDef2 -> ArrowTy (TyOf Exp2)
forall ex. FunDef ex -> ArrowTy (TyOf ex)
funTy FunDef2
fundef

               go :: Exp2 -> PassM Exp2
go = FunDefs Exp2
-> [LocVar]
-> EndOfRel
-> Map LocVar LocVar
-> Map LocVar LocVar
-> Env2 Ty2
-> Exp2
-> PassM Exp2
exp FunDefs Exp2
fns [LocVar]
retlocs EndOfRel
eor Map LocVar LocVar
lenv Map LocVar LocVar
afterenv Env2 Ty2
env2


               -- We may need to emit some additional let bindings if we've reached
               -- an end witness that is equivalent to the after location of something.
               wrapBody :: Exp2 -> [(LocVar, LocVar)] -> Exp2
wrapBody Exp2
e ((LocVar
l1,LocVar
l2):[(LocVar, LocVar)]
ls) =
                 case LocVar -> Map LocVar LocVar -> Maybe LocVar
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup LocVar
l1 Map LocVar LocVar
afterenv of
                   Maybe LocVar
Nothing -> Exp2 -> [(LocVar, LocVar)] -> Exp2
wrapBody Exp2
e [(LocVar, LocVar)]
ls
                   Just LocVar
la ->
                     let go :: LocVar -> [(LocVar, LocVar, Int)] -> [(LocVar, LocVar, Int)]
go LocVar
loc [(LocVar, LocVar, Int)]
acc =
                           case LocVar -> Map LocVar Ty2 -> Maybe Ty2
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup LocVar
loc (Env2 Ty2 -> Map LocVar Ty2
forall a. Env2 a -> TyEnv a
vEnv Env2 Ty2
env2) of
                             Just Ty2
ty
                               | Ty2 -> Bool
forall a. UrTy a -> Bool
isScalarTy Ty2
ty ->
                                 case LocVar -> Map LocVar LocVar -> Maybe LocVar
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup LocVar
loc Map LocVar LocVar
afterenv of
                                   Maybe LocVar
Nothing -> [(LocVar, LocVar, Int)]
acc
                                   Just LocVar
lb -> LocVar -> [(LocVar, LocVar, Int)] -> [(LocVar, LocVar, Int)]
go LocVar
lb ([(LocVar, LocVar, Int)]
acc [(LocVar, LocVar, Int)]
-> [(LocVar, LocVar, Int)] -> [(LocVar, LocVar, Int)]
forall a. [a] -> [a] -> [a]
++ [(LocVar
lb,LocVar
loc,Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ Ty2 -> Maybe Int
forall a. UrTy a -> Maybe Int
sizeOfTy Ty2
ty)])
                               | Bool
otherwise -> [(LocVar, LocVar, Int)]
acc
                             Maybe Ty2
Nothing -> [(LocVar, LocVar, Int)]
acc
                         scalar_witnesses :: [(LocVar, LocVar, Int)]
scalar_witnesses = LocVar -> [(LocVar, LocVar, Int)] -> [(LocVar, LocVar, Int)]
go LocVar
la []
                         bind_witnesses :: PreExp E2Ext loc dec
-> t (LocVar, loc, Int) -> PreExp E2Ext loc dec
bind_witnesses PreExp E2Ext loc dec
bod t (LocVar, loc, Int)
ls =
                           ((LocVar, loc, Int)
 -> PreExp E2Ext loc dec -> PreExp E2Ext loc dec)
-> PreExp E2Ext loc dec
-> t (LocVar, loc, Int)
-> PreExp E2Ext loc dec
forall a b. (a -> b -> b) -> b -> t a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
L.foldr (\(LocVar
v,loc
w,Int
sz) PreExp E2Ext loc dec
acc ->
                                     E2Ext loc dec -> PreExp E2Ext loc dec
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext loc dec -> PreExp E2Ext loc dec)
-> E2Ext loc dec -> PreExp E2Ext loc dec
forall a b. (a -> b) -> a -> b
$ LocVar -> PreLocExp loc -> PreExp E2Ext loc dec -> E2Ext loc dec
forall loc dec.
LocVar -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE LocVar
v (Int -> loc -> PreLocExp loc
forall loc. Int -> loc -> PreLocExp loc
AfterConstantLE Int
sz loc
w) PreExp E2Ext loc dec
acc)
                           PreExp E2Ext loc dec
bod t (LocVar, loc, Int)
ls
                         bod' :: Exp2
bod' = Exp2 -> [(LocVar, LocVar, Int)] -> Exp2
forall {t :: * -> *} {loc} {dec}.
Foldable t =>
PreExp E2Ext loc dec
-> t (LocVar, loc, Int) -> PreExp E2Ext loc dec
bind_witnesses Exp2
e [(LocVar, LocVar, Int)]
scalar_witnesses
                         bod'' :: Exp2
bod'' =  E2Ext LocVar Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (LocVar -> PreLocExp LocVar -> Exp2 -> E2Ext LocVar Ty2
forall loc dec.
LocVar -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE LocVar
la (LocVar -> PreLocExp LocVar
forall loc. loc -> PreLocExp loc
FromEndLE LocVar
l2) Exp2
bod')
                     in Exp2 -> [(LocVar, LocVar)] -> Exp2
wrapBody Exp2
bod'' [(LocVar, LocVar)]
ls
               wrapBody Exp2
e [] = Exp2
e

               -- Process a let bound fn app.
               doBoundApp :: Var -> [LocVar] -> PassM ([LocVar], [(LocVar, Var)], EndOfRel)
               doBoundApp :: LocVar
-> [LocVar] -> PassM ([LocVar], [(LocVar, LocVar)], EndOfRel)
doBoundApp LocVar
f [LocVar]
lsin = do
                 let fty :: ArrowTy2 Ty2
fty = LocVar -> ArrowTy2 Ty2
funtype LocVar
f
                     rets :: Set LocRet
rets = [LocRet] -> Set LocRet
forall a. Ord a => [a] -> Set a
S.fromList ([LocRet] -> Set LocRet) -> [LocRet] -> Set LocRet
forall a b. (a -> b) -> a -> b
$ ArrowTy2 Ty2 -> [LocRet]
forall ty2. ArrowTy2 ty2 -> [LocRet]
locRets ArrowTy2 Ty2
fty
                     -- The travlist is a list of pair (location, bool) where the bool is
                     -- if the location was traversed, and the location is from the
                     -- AppE call.
                     travlist :: [(LocVar, Bool)]
travlist = [LocVar] -> [Bool] -> [(LocVar, Bool)]
forall a b. [a] -> [b] -> [(a, b)]
zip [LocVar]
lsin ([Bool] -> [(LocVar, Bool)]) -> [Bool] -> [(LocVar, Bool)]
forall a b. (a -> b) -> a -> b
$ (LRM -> Bool) -> [LRM] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
L.map (\LRM
l -> LocRet -> Set LocRet -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member (LRM -> LocRet
EndOf LRM
l) Set LocRet
rets)  (ArrowTy2 Ty2 -> [LRM]
forall ty2. ArrowTy2 ty2 -> [LRM]
locVars ArrowTy2 Ty2
fty)

                 -- For each traversed location, gensym a new variable for its end,
                 -- and generate a list of (location, endof location) pairs.
                 let handleTravList :: [(a, LocVar)] -> (a, Bool) -> m [(a, LocVar)]
handleTravList [(a, LocVar)]
lst (a
_l,Bool
False) = [(a, LocVar)] -> m [(a, LocVar)]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return [(a, LocVar)]
lst
                     handleTravList [(a, LocVar)]
lst (a
l,Bool
True) = LocVar -> m LocVar
forall (m :: * -> *). MonadState Int m => LocVar -> m LocVar
gensym LocVar
"endof" m LocVar -> (LocVar -> m [(a, LocVar)]) -> m [(a, LocVar)]
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \LocVar
l' -> [(a, LocVar)] -> m [(a, LocVar)]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(a, LocVar)] -> m [(a, LocVar)])
-> [(a, LocVar)] -> m [(a, LocVar)]
forall a b. (a -> b) -> a -> b
$ (a
l,LocVar
l')(a, LocVar) -> [(a, LocVar)] -> [(a, LocVar)]
forall a. a -> [a] -> [a]
:[(a, LocVar)]
lst

                 -- Walk through our pairs of (location, endof location) and update the
                 -- endof relation.
                 let mkEor :: (LocVar, LocVar) -> EndOfRel -> EndOfRel
mkEor (LocVar
l1,LocVar
l2) EndOfRel
eor = LocVar -> LocVar -> EndOfRel -> EndOfRel
mkEnd LocVar
l1 LocVar
l2 EndOfRel
eor


                 [(LocVar, LocVar)]
newls <- [(LocVar, LocVar)] -> [(LocVar, LocVar)]
forall a. [a] -> [a]
reverse ([(LocVar, LocVar)] -> [(LocVar, LocVar)])
-> PassM [(LocVar, LocVar)] -> PassM [(LocVar, LocVar)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([(LocVar, LocVar)] -> (LocVar, Bool) -> PassM [(LocVar, LocVar)])
-> [(LocVar, LocVar)]
-> [(LocVar, Bool)]
-> PassM [(LocVar, LocVar)]
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM [(LocVar, LocVar)] -> (LocVar, Bool) -> PassM [(LocVar, LocVar)]
forall {m :: * -> *} {a}.
MonadState Int m =>
[(a, LocVar)] -> (a, Bool) -> m [(a, LocVar)]
handleTravList [] [(LocVar, Bool)]
travlist
                 let eor' :: EndOfRel
eor' = ((LocVar, LocVar) -> EndOfRel -> EndOfRel)
-> EndOfRel -> [(LocVar, LocVar)] -> EndOfRel
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
L.foldr (LocVar, LocVar) -> EndOfRel -> EndOfRel
mkEor EndOfRel
eor [(LocVar, LocVar)]
newls
                 let outlocs :: [LocVar]
outlocs = ((LocVar, LocVar) -> LocVar) -> [(LocVar, LocVar)] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map (LocVar, LocVar) -> LocVar
forall a b. (a, b) -> b
snd [(LocVar, LocVar)]
newls
                 ([LocVar], [(LocVar, LocVar)], EndOfRel)
-> PassM ([LocVar], [(LocVar, LocVar)], EndOfRel)
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return ([LocVar]
outlocs, [(LocVar, LocVar)]
newls, EndOfRel
eor')