{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}

{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# OPTIONS_GHC -fdefer-typed-holes #-}

-- | An intermediate language with an effect system that captures traversals.

module Gibbon.NewL2.Syntax
    (
    -- * Extended language L2 with location types.
      Old.E2Ext(..)
    , Prog2, DDefs2, DDef2, FunDef2, FunDefs2, Exp2, Ty2(..)
    , Old.Effect(..), Old.ArrowTy2(..) , Old.LocRet(..), LocArg(..), LocExp, Old.PreLocExp(..)

    -- * Regions and locations
    , LocVar, Old.Region(..), Old.Modality(..),  Old.LRM(..), LREM(..)
    , Old.Multiplicity(..), Old.RegionSize(..), Old.RegionType(..), Old.regionToVar

    -- * Operations on types
    , Old.allLocVars, Old.inLocVars, Old.outLocVars, Old.outRegVars, Old.inRegVars, Old.allRegVars
    , substLoc, substLocs, Old.substEff, Old.substEffs, extendPatternMatchEnv
    , locsInTy, Old.dummyTyLocs, allFreeVars, freeLocVars
    , toLocVar, fromLRM

    -- * Other helpers
    , revertToL1, Old.occurs, Old.mapPacked, Old.constPacked, depList, Old.changeAppToSpawn
    , toEndFromTaggedV, toTagV

    , module Gibbon.Language
    )
    where

import           Control.DeepSeq
import qualified Data.List as L
import qualified Data.Set as S
import qualified Data.Map as M
import           GHC.Stack (HasCallStack)
import           Text.PrettyPrint.GenericPretty

import           Gibbon.Common
import           Gibbon.Language
-- import           Text.PrettyPrint.HughesPJ
import           Gibbon.L1.Syntax hiding (AddFixed, StartOfPkdCursor)
import qualified Gibbon.L1.Syntax as L1

import qualified Gibbon.L2.Syntax as Old

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

type Prog2    = Prog Exp2
type DDef2    = DDef Ty2
type DDefs2   = DDefs Ty2
type FunDef2  = FunDef Exp2
type FunDefs2 = FunDefs Exp2

-- | Function types know about locations and traversal effects.
instance FunctionTy Ty2 where
  type ArrowTy Ty2 = Old.ArrowTy2 Ty2
  inTys :: ArrowTy Ty2 -> [Ty2]
inTys = ArrowTy Ty2 -> [Ty2]
ArrowTy2 Ty2 -> [Ty2]
forall ty2. ArrowTy2 ty2 -> [ty2]
Old.arrIns
  outTy :: ArrowTy Ty2 -> Ty2
outTy = ArrowTy Ty2 -> Ty2
ArrowTy2 Ty2 -> Ty2
forall ty2. ArrowTy2 ty2 -> ty2
Old.arrOut

-- | Extended expressions, L2.
--
--   By adding a `LocVar` decoration, all data constructors,
--   applications, and bindings gain a location annotation.
type Exp2   = PreExp Old.E2Ext LocArg Ty2
type LocExp = Old.PreLocExp LocArg

-- We need a newtype here to avoid overlapping type family instance for FunctionTy
-- | L1 Types extended with abstract Locations.
newtype Ty2 = MkTy2 { Ty2 -> UrTy LocVar
unTy2 :: (UrTy LocVar) }
  deriving (ReadPrec [Ty2]
ReadPrec Ty2
Int -> ReadS Ty2
ReadS [Ty2]
(Int -> ReadS Ty2)
-> ReadS [Ty2] -> ReadPrec Ty2 -> ReadPrec [Ty2] -> Read Ty2
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS Ty2
readsPrec :: Int -> ReadS Ty2
$creadList :: ReadS [Ty2]
readList :: ReadS [Ty2]
$creadPrec :: ReadPrec Ty2
readPrec :: ReadPrec Ty2
$creadListPrec :: ReadPrec [Ty2]
readListPrec :: ReadPrec [Ty2]
Read, Int -> Ty2 -> ShowS
[Ty2] -> ShowS
Ty2 -> String
(Int -> Ty2 -> ShowS)
-> (Ty2 -> String) -> ([Ty2] -> ShowS) -> Show Ty2
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Ty2 -> ShowS
showsPrec :: Int -> Ty2 -> ShowS
$cshow :: Ty2 -> String
show :: Ty2 -> String
$cshowList :: [Ty2] -> ShowS
showList :: [Ty2] -> ShowS
Show, Ty2 -> Ty2 -> Bool
(Ty2 -> Ty2 -> Bool) -> (Ty2 -> Ty2 -> Bool) -> Eq Ty2
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Ty2 -> Ty2 -> Bool
== :: Ty2 -> Ty2 -> Bool
$c/= :: Ty2 -> Ty2 -> Bool
/= :: Ty2 -> Ty2 -> Bool
Eq, Eq Ty2
Eq Ty2
-> (Ty2 -> Ty2 -> Ordering)
-> (Ty2 -> Ty2 -> Bool)
-> (Ty2 -> Ty2 -> Bool)
-> (Ty2 -> Ty2 -> Bool)
-> (Ty2 -> Ty2 -> Bool)
-> (Ty2 -> Ty2 -> Ty2)
-> (Ty2 -> Ty2 -> Ty2)
-> Ord Ty2
Ty2 -> Ty2 -> Bool
Ty2 -> Ty2 -> Ordering
Ty2 -> Ty2 -> Ty2
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 :: Ty2 -> Ty2 -> Ordering
compare :: Ty2 -> Ty2 -> Ordering
$c< :: Ty2 -> Ty2 -> Bool
< :: Ty2 -> Ty2 -> Bool
$c<= :: Ty2 -> Ty2 -> Bool
<= :: Ty2 -> Ty2 -> Bool
$c> :: Ty2 -> Ty2 -> Bool
> :: Ty2 -> Ty2 -> Bool
$c>= :: Ty2 -> Ty2 -> Bool
>= :: Ty2 -> Ty2 -> Bool
$cmax :: Ty2 -> Ty2 -> Ty2
max :: Ty2 -> Ty2 -> Ty2
$cmin :: Ty2 -> Ty2 -> Ty2
min :: Ty2 -> Ty2 -> Ty2
Ord, (forall x. Ty2 -> Rep Ty2 x)
-> (forall x. Rep Ty2 x -> Ty2) -> Generic Ty2
forall x. Rep Ty2 x -> Ty2
forall x. Ty2 -> Rep Ty2 x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Ty2 -> Rep Ty2 x
from :: forall x. Ty2 -> Rep Ty2 x
$cto :: forall x. Rep Ty2 x -> Ty2
to :: forall x. Rep Ty2 x -> Ty2
Generic)

instance Out Ty2
instance NFData Ty2

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

data LREM = LREM { LREM -> LocVar
lremLoc    :: LocVar
                 , LREM -> LocVar
lremReg    :: RegVar
                 , LREM -> LocVar
lremEndReg :: RegVar
                 , LREM -> Modality
lremMode   :: Old.Modality
                 }
  deriving (ReadPrec [LREM]
ReadPrec LREM
Int -> ReadS LREM
ReadS [LREM]
(Int -> ReadS LREM)
-> ReadS [LREM] -> ReadPrec LREM -> ReadPrec [LREM] -> Read LREM
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS LREM
readsPrec :: Int -> ReadS LREM
$creadList :: ReadS [LREM]
readList :: ReadS [LREM]
$creadPrec :: ReadPrec LREM
readPrec :: ReadPrec LREM
$creadListPrec :: ReadPrec [LREM]
readListPrec :: ReadPrec [LREM]
Read,Int -> LREM -> ShowS
[LREM] -> ShowS
LREM -> String
(Int -> LREM -> ShowS)
-> (LREM -> String) -> ([LREM] -> ShowS) -> Show LREM
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LREM -> ShowS
showsPrec :: Int -> LREM -> ShowS
$cshow :: LREM -> String
show :: LREM -> String
$cshowList :: [LREM] -> ShowS
showList :: [LREM] -> ShowS
Show,LREM -> LREM -> Bool
(LREM -> LREM -> Bool) -> (LREM -> LREM -> Bool) -> Eq LREM
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LREM -> LREM -> Bool
== :: LREM -> LREM -> Bool
$c/= :: LREM -> LREM -> Bool
/= :: LREM -> LREM -> Bool
Eq,Eq LREM
Eq LREM
-> (LREM -> LREM -> Ordering)
-> (LREM -> LREM -> Bool)
-> (LREM -> LREM -> Bool)
-> (LREM -> LREM -> Bool)
-> (LREM -> LREM -> Bool)
-> (LREM -> LREM -> LREM)
-> (LREM -> LREM -> LREM)
-> Ord LREM
LREM -> LREM -> Bool
LREM -> LREM -> Ordering
LREM -> LREM -> LREM
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 :: LREM -> LREM -> Ordering
compare :: LREM -> LREM -> Ordering
$c< :: LREM -> LREM -> Bool
< :: LREM -> LREM -> Bool
$c<= :: LREM -> LREM -> Bool
<= :: LREM -> LREM -> Bool
$c> :: LREM -> LREM -> Bool
> :: LREM -> LREM -> Bool
$c>= :: LREM -> LREM -> Bool
>= :: LREM -> LREM -> Bool
$cmax :: LREM -> LREM -> LREM
max :: LREM -> LREM -> LREM
$cmin :: LREM -> LREM -> LREM
min :: LREM -> LREM -> LREM
Ord,(forall x. LREM -> Rep LREM x)
-> (forall x. Rep LREM x -> LREM) -> Generic LREM
forall x. Rep LREM x -> LREM
forall x. LREM -> Rep LREM x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. LREM -> Rep LREM x
from :: forall x. LREM -> Rep LREM x
$cto :: forall x. Rep LREM x -> LREM
to :: forall x. Rep LREM x -> LREM
Generic)

instance Out LREM

instance NFData LREM where
  rnf :: LREM -> ()
rnf (LREM LocVar
a LocVar
b LocVar
c Modality
d)  = LocVar -> ()
forall a. NFData a => a -> ()
rnf LocVar
a () -> () -> ()
forall a b. a -> b -> b
`seq` LocVar -> ()
forall a. NFData a => a -> ()
rnf LocVar
b () -> () -> ()
forall a b. a -> b -> b
`seq` LocVar -> ()
forall a. NFData a => a -> ()
rnf LocVar
c () -> () -> ()
forall a b. a -> b -> b
`seq` Modality -> ()
forall a. NFData a => a -> ()
rnf Modality
d

fromLRM :: Old.LRM -> LREM
fromLRM :: LRM -> LREM
fromLRM (Old.LRM LocVar
loc Region
reg Modality
mode) = LocVar -> LocVar -> LocVar -> Modality -> LREM
LREM LocVar
loc (Region -> LocVar
Old.regionToVar Region
reg) (LocVar -> LocVar
toEndV (Region -> LocVar
Old.regionToVar Region
reg)) Modality
mode

data LocArg = Loc LREM
            | EndWitness LREM Var
            | Reg RegVar Old.Modality
            | EndOfReg RegVar Old.Modality RegVar
            | EndOfReg_Tagged RegVar
  deriving (ReadPrec [LocArg]
ReadPrec LocArg
Int -> ReadS LocArg
ReadS [LocArg]
(Int -> ReadS LocArg)
-> ReadS [LocArg]
-> ReadPrec LocArg
-> ReadPrec [LocArg]
-> Read LocArg
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS LocArg
readsPrec :: Int -> ReadS LocArg
$creadList :: ReadS [LocArg]
readList :: ReadS [LocArg]
$creadPrec :: ReadPrec LocArg
readPrec :: ReadPrec LocArg
$creadListPrec :: ReadPrec [LocArg]
readListPrec :: ReadPrec [LocArg]
Read, Int -> LocArg -> ShowS
[LocArg] -> ShowS
LocArg -> String
(Int -> LocArg -> ShowS)
-> (LocArg -> String) -> ([LocArg] -> ShowS) -> Show LocArg
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LocArg -> ShowS
showsPrec :: Int -> LocArg -> ShowS
$cshow :: LocArg -> String
show :: LocArg -> String
$cshowList :: [LocArg] -> ShowS
showList :: [LocArg] -> ShowS
Show, LocArg -> LocArg -> Bool
(LocArg -> LocArg -> Bool)
-> (LocArg -> LocArg -> Bool) -> Eq LocArg
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LocArg -> LocArg -> Bool
== :: LocArg -> LocArg -> Bool
$c/= :: LocArg -> LocArg -> Bool
/= :: LocArg -> LocArg -> Bool
Eq, Eq LocArg
Eq LocArg
-> (LocArg -> LocArg -> Ordering)
-> (LocArg -> LocArg -> Bool)
-> (LocArg -> LocArg -> Bool)
-> (LocArg -> LocArg -> Bool)
-> (LocArg -> LocArg -> Bool)
-> (LocArg -> LocArg -> LocArg)
-> (LocArg -> LocArg -> LocArg)
-> Ord LocArg
LocArg -> LocArg -> Bool
LocArg -> LocArg -> Ordering
LocArg -> LocArg -> LocArg
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 :: LocArg -> LocArg -> Ordering
compare :: LocArg -> LocArg -> Ordering
$c< :: LocArg -> LocArg -> Bool
< :: LocArg -> LocArg -> Bool
$c<= :: LocArg -> LocArg -> Bool
<= :: LocArg -> LocArg -> Bool
$c> :: LocArg -> LocArg -> Bool
> :: LocArg -> LocArg -> Bool
$c>= :: LocArg -> LocArg -> Bool
>= :: LocArg -> LocArg -> Bool
$cmax :: LocArg -> LocArg -> LocArg
max :: LocArg -> LocArg -> LocArg
$cmin :: LocArg -> LocArg -> LocArg
min :: LocArg -> LocArg -> LocArg
Ord, (forall x. LocArg -> Rep LocArg x)
-> (forall x. Rep LocArg x -> LocArg) -> Generic LocArg
forall x. Rep LocArg x -> LocArg
forall x. LocArg -> Rep LocArg x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. LocArg -> Rep LocArg x
from :: forall x. LocArg -> Rep LocArg x
$cto :: forall x. Rep LocArg x -> LocArg
to :: forall x. Rep LocArg x -> LocArg
Generic)

instance Out LocArg
instance NFData LocArg

toLocVar :: LocArg -> LocVar
toLocVar :: LocArg -> LocVar
toLocVar LocArg
arg =
  case LocArg
arg of
    Loc LREM
lrm        -> LREM -> LocVar
lremLoc LREM
lrm
    EndWitness LREM
_ LocVar
v -> LocVar
v
    Reg LocVar
v Modality
_        -> LocVar
v
    EndOfReg LocVar
_ Modality
_ LocVar
v -> LocVar
v
    EndOfReg_Tagged LocVar
v -> LocVar -> LocVar
toEndFromTaggedV LocVar
v

instance Out (Old.ArrowTy2 Ty2)

toTagV :: Var -> Var
toTagV :: LocVar -> LocVar
toTagV LocVar
v = (String -> LocVar
toVar String
"tag_") LocVar -> LocVar -> LocVar
`varAppend` LocVar
v

toEndFromTaggedV :: Var -> Var
toEndFromTaggedV :: LocVar -> LocVar
toEndFromTaggedV LocVar
v = (String -> LocVar
toVar String
"end_from_tagged_") LocVar -> LocVar -> LocVar
`varAppend` LocVar
v

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

instance FreeVars LocExp where
  gFreeVars :: LocExp -> Set LocVar
gFreeVars LocExp
e =
    case LocExp
e of
      Old.AfterConstantLE Int
_ LocArg
loc   -> LocVar -> Set LocVar
forall a. a -> Set a
S.singleton (LocArg -> LocVar
toLocVar LocArg
loc)
      Old.AfterVariableLE LocVar
v LocArg
loc Bool
_ -> [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar
v,LocArg -> LocVar
toLocVar LocArg
loc]
      LocExp
_ -> Set LocVar
forall a. Set a
S.empty


instance Typeable (Old.E2Ext LocArg Ty2) where
  gRecoverType :: DDefs (TyOf (E2Ext LocArg Ty2))
-> Env2 (TyOf (E2Ext LocArg Ty2))
-> E2Ext LocArg Ty2
-> TyOf (E2Ext LocArg Ty2)
gRecoverType DDefs (TyOf (E2Ext LocArg Ty2))
ddfs Env2 (TyOf (E2Ext LocArg Ty2))
env2 E2Ext LocArg Ty2
ex =
    case E2Ext LocArg Ty2
ex of
      Old.LetRegionE Region
_r RegionSize
_ Maybe RegionType
_ E2 LocArg Ty2
bod    -> DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
DDefs (TyOf (E2Ext LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
Env2 (TyOf (E2Ext LocArg Ty2))
env2 E2 LocArg Ty2
bod
      Old.LetParRegionE Region
_r RegionSize
_ Maybe RegionType
_ E2 LocArg Ty2
bod -> DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
DDefs (TyOf (E2Ext LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
Env2 (TyOf (E2Ext LocArg Ty2))
env2 E2 LocArg Ty2
bod
      Old.StartOfPkdCursor{}       -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ UrTy LocVar
forall loc. UrTy loc
CursorTy
      Old.TagCursor{}      -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ UrTy LocVar
forall loc. UrTy loc
CursorTy
      Old.LetLocE LocVar
_l LocExp
_rhs E2 LocArg Ty2
bod -> DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
DDefs (TyOf (E2Ext LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
Env2 (TyOf (E2Ext LocArg Ty2))
env2 E2 LocArg Ty2
bod
      Old.RetE [LocArg]
_loc LocVar
var       -> case LocVar -> Map LocVar Ty2 -> Maybe Ty2
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup LocVar
var (Env2 Ty2 -> Map LocVar Ty2
forall a. Env2 a -> TyEnv a
vEnv Env2 (TyOf (E2Ext LocArg Ty2))
Env2 Ty2
env2) of
                                   Just Ty2
ty -> TyOf (E2Ext LocArg Ty2)
Ty2
ty
                                   Maybe Ty2
Nothing -> String -> TyOf (E2Ext LocArg Ty2)
forall a. HasCallStack => String -> a
error (String -> TyOf (E2Ext LocArg Ty2))
-> String -> TyOf (E2Ext LocArg Ty2)
forall a b. (a -> b) -> a -> b
$ String
"gRecoverType: unbound variable " String -> ShowS
forall a. [a] -> [a] -> [a]
++ LocVar -> String
forall a. Out a => a -> String
sdoc LocVar
var
      Old.FromEndE LocArg
_loc       -> String -> Ty2
forall a. HasCallStack => String -> a
error String
"Shouldn't enconter FromEndE in tail position"
      Old.BoundsCheck{}       -> String -> Ty2
forall a. HasCallStack => String -> a
error String
"Shouldn't enconter BoundsCheck in tail position"
      Old.IndirectionE String
tycon String
_ (LocArg, LocArg)
_ (LocArg
to,LocArg
_) E2 LocArg Ty2
_ -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ String -> LocVar -> UrTy LocVar
forall loc. String -> loc -> UrTy loc
PackedTy String
tycon (LocArg -> LocVar
toLocVar LocArg
to)
      Old.AddFixed{}          -> String -> Ty2
forall a. HasCallStack => String -> a
error String
"Shouldn't enconter AddFixed in tail position"
      E2Ext LocArg Ty2
Old.GetCilkWorkerNum    -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ UrTy LocVar
forall loc. UrTy loc
IntTy
      Old.LetAvail [LocVar]
_ E2 LocArg Ty2
bod      -> DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
DDefs (TyOf (E2Ext LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
Env2 (TyOf (E2Ext LocArg Ty2))
env2 E2 LocArg Ty2
bod
      Old.AllocateTagHere{}   -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ [UrTy LocVar] -> UrTy LocVar
forall loc. [UrTy loc] -> UrTy loc
ProdTy []
      Old.AllocateScalarsHere{} -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ [UrTy LocVar] -> UrTy LocVar
forall loc. [UrTy loc] -> UrTy loc
ProdTy []
      Old.SSPush{}              -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ [UrTy LocVar] -> UrTy LocVar
forall loc. [UrTy loc] -> UrTy loc
ProdTy []
      Old.SSPop{}               -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ [UrTy LocVar] -> UrTy LocVar
forall loc. [UrTy loc] -> UrTy loc
ProdTy []


-- | The 'gRecoverType' instance defined in Language.Syntax is incorrect for L2.
-- For the AppE case, it'll just return the type with with the function was
-- defined. However, we want the recovered type to have the locations actually
-- used at the callsites! For example,
--
--     add1 :: Tree @ a -> Tree @ b
--     add1 = _
--
--     ... (add1 [loc1, loc2] tr1) ..
--
-- in this case, we want the type of (add1 tr1) to be (Tree @ loc2)
-- and NOT (Tree @ b). We have to do something similar for variables bound by
-- a pattern match.
instance Out (Old.E2Ext LocArg Ty2) => Typeable (PreExp Old.E2Ext LocArg Ty2) where
  gRecoverType :: DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
env2 E2 LocArg Ty2
ex =
    case E2 LocArg Ty2
ex of
      VarE LocVar
v       -> Ty2 -> LocVar -> Map LocVar Ty2 -> Ty2
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault (String -> Ty2
forall a. HasCallStack => String -> a
error (String -> Ty2) -> String -> Ty2
forall a b. (a -> b) -> a -> b
$ String
"Cannot find type of variable " 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
" in " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Map LocVar Ty2 -> String
forall a. Show a => a -> String
show (Env2 Ty2 -> Map LocVar Ty2
forall a. Env2 a -> TyEnv a
vEnv Env2 (TyOf (E2 LocArg Ty2))
Env2 Ty2
env2)) LocVar
v (Env2 Ty2 -> Map LocVar Ty2
forall a. Env2 a -> TyEnv a
vEnv Env2 (TyOf (E2 LocArg Ty2))
Env2 Ty2
env2)
      LitE Int
_       -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ UrTy LocVar
forall loc. UrTy loc
IntTy
      CharE Char
_      -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ UrTy LocVar
forall loc. UrTy loc
CharTy
      FloatE{}     -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ UrTy LocVar
forall loc. UrTy loc
FloatTy
      LitSymE LocVar
_    -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ UrTy LocVar
forall loc. UrTy loc
SymTy
      AppE LocVar
v [LocArg]
locargs [E2 LocArg Ty2]
_ ->
                       let fnty :: ArrowTy2 Ty2
fnty  = Env2 Ty2 -> TyEnv (ArrowTy Ty2)
forall a. Env2 a -> TyEnv (ArrowTy a)
fEnv Env2 (TyOf (E2 LocArg Ty2))
Env2 Ty2
env2 Map LocVar (ArrowTy2 Ty2) -> LocVar -> ArrowTy2 Ty2
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# LocVar
v
                           outty :: Ty2
outty = ArrowTy2 Ty2 -> Ty2
forall ty2. ArrowTy2 ty2 -> ty2
Old.arrOut ArrowTy2 Ty2
fnty
                           mp :: Map LocVar LocVar
mp = [(LocVar, LocVar)] -> Map LocVar LocVar
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(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 (ArrowTy2 Ty2 -> [LocVar]
forall ty2. ArrowTy2 ty2 -> [LocVar]
Old.allLocVars ArrowTy2 Ty2
fnty) ((LocArg -> LocVar) -> [LocArg] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocVar
toLocVar [LocArg]
locargs)
                       in Map LocVar LocVar -> Ty2 -> Ty2
substLoc Map LocVar LocVar
mp Ty2
outty

      PrimAppE (DictInsertP Ty2
ty) ((VarE LocVar
v):[E2 LocArg Ty2]
_) -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ Maybe LocVar -> UrTy () -> UrTy LocVar
forall loc. Maybe LocVar -> UrTy () -> UrTy loc
SymDictTy (LocVar -> Maybe LocVar
forall a. a -> Maybe a
Just LocVar
v) (UrTy () -> UrTy LocVar) -> UrTy () -> UrTy LocVar
forall a b. (a -> b) -> a -> b
$ UrTy LocVar -> UrTy ()
forall a. UrTy a -> UrTy ()
stripTyLocs (Ty2 -> UrTy LocVar
unTy2 Ty2
ty)
      PrimAppE (DictEmptyP  Ty2
ty) ((VarE LocVar
v):[E2 LocArg Ty2]
_) -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ Maybe LocVar -> UrTy () -> UrTy LocVar
forall loc. Maybe LocVar -> UrTy () -> UrTy loc
SymDictTy (LocVar -> Maybe LocVar
forall a. a -> Maybe a
Just LocVar
v) (UrTy () -> UrTy LocVar) -> UrTy () -> UrTy LocVar
forall a b. (a -> b) -> a -> b
$ UrTy LocVar -> UrTy ()
forall a. UrTy a -> UrTy ()
stripTyLocs (Ty2 -> UrTy LocVar
unTy2 Ty2
ty)
      PrimAppE Prim Ty2
p [E2 LocArg Ty2]
_ -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ Prim (UrTy LocVar) -> UrTy LocVar
forall a. Prim (UrTy a) -> UrTy a
primRetTy ((Ty2 -> UrTy LocVar) -> Prim Ty2 -> Prim (UrTy LocVar)
forall a b. (a -> b) -> Prim a -> Prim b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Ty2 -> UrTy LocVar
unTy2 Prim Ty2
p)

      LetE (LocVar
v,[LocArg]
_,Ty2
t,E2 LocArg Ty2
_) E2 LocArg Ty2
e -> DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
ddfs (LocVar -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. LocVar -> a -> Env2 a -> Env2 a
extendVEnv LocVar
v Ty2
t Env2 (TyOf (E2 LocArg Ty2))
Env2 Ty2
env2) E2 LocArg Ty2
e
      IfE E2 LocArg Ty2
_ E2 LocArg Ty2
e E2 LocArg Ty2
_        -> DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
env2 E2 LocArg Ty2
e
      MkProdE [E2 LocArg Ty2]
es       -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ [UrTy LocVar] -> UrTy LocVar
forall loc. [UrTy loc] -> UrTy loc
ProdTy ([UrTy LocVar] -> UrTy LocVar) -> [UrTy LocVar] -> UrTy LocVar
forall a b. (a -> b) -> a -> b
$ (E2 LocArg Ty2 -> UrTy LocVar) -> [E2 LocArg Ty2] -> [UrTy LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map (Ty2 -> UrTy LocVar
unTy2 (Ty2 -> UrTy LocVar)
-> (E2 LocArg Ty2 -> Ty2) -> E2 LocArg Ty2 -> UrTy LocVar
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
env2) [E2 LocArg Ty2]
es
      DataConE LocArg
loc String
c [E2 LocArg Ty2]
_ -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ String -> LocVar -> UrTy LocVar
forall loc. String -> loc -> UrTy loc
PackedTy (DDefs Ty2 -> ShowS
forall a. Out a => DDefs a -> ShowS
getTyOfDataCon DDefs (TyOf (E2 LocArg Ty2))
DDefs Ty2
ddfs String
c) (LocArg -> LocVar
toLocVar LocArg
loc)
      TimeIt E2 LocArg Ty2
e Ty2
_ Bool
_     -> DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
env2 E2 LocArg Ty2
e
      MapE (LocVar, Ty2, E2 LocArg Ty2)
_ E2 LocArg Ty2
e         -> DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
env2 E2 LocArg Ty2
e
      FoldE (LocVar, Ty2, E2 LocArg Ty2)
_ (LocVar, Ty2, E2 LocArg Ty2)
_ E2 LocArg Ty2
e      -> DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
env2 E2 LocArg Ty2
e
      Ext E2Ext LocArg Ty2
ext          -> DDefs (TyOf (E2Ext LocArg Ty2))
-> Env2 (TyOf (E2Ext LocArg Ty2))
-> E2Ext LocArg Ty2
-> TyOf (E2Ext LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
DDefs (TyOf (E2Ext LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
Env2 (TyOf (E2Ext LocArg Ty2))
env2 E2Ext LocArg Ty2
ext
      ProjE Int
i E2 LocArg Ty2
e ->
        case Ty2 -> UrTy LocVar
unTy2 (Ty2 -> UrTy LocVar) -> Ty2 -> UrTy LocVar
forall a b. (a -> b) -> a -> b
$ DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
env2 E2 LocArg Ty2
e of
          (ProdTy [UrTy LocVar]
tys) -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ ([UrTy LocVar]
tys [UrTy LocVar] -> Int -> UrTy LocVar
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
          UrTy LocVar
oth -> String -> TyOf (E2 LocArg Ty2)
forall a. HasCallStack => String -> a
error(String -> TyOf (E2 LocArg Ty2)) -> String -> TyOf (E2 LocArg Ty2)
forall a b. (a -> b) -> a -> b
$ String
"typeExp: Cannot project fields from this type: "String -> ShowS
forall a. [a] -> [a] -> [a]
++UrTy LocVar -> String
forall a. Show a => a -> String
show UrTy LocVar
oth
                        String -> ShowS
forall a. [a] -> [a] -> [a]
++String
"\nExpression:\n  "String -> ShowS
forall a. [a] -> [a] -> [a]
++ E2 LocArg Ty2 -> String
forall a. Out a => a -> String
sdoc E2 LocArg Ty2
ex
                        String -> ShowS
forall a. [a] -> [a] -> [a]
++String
"\nEnvironment:\n  "String -> ShowS
forall a. [a] -> [a] -> [a]
++Map LocVar Ty2 -> String
forall a. Out a => a -> String
sdoc (Env2 Ty2 -> Map LocVar Ty2
forall a. Env2 a -> TyEnv a
vEnv Env2 (TyOf (E2 LocArg Ty2))
Env2 Ty2
env2)
      SpawnE LocVar
v [LocArg]
locargs [E2 LocArg Ty2]
_ ->
                         let fnty :: ArrowTy2 Ty2
fnty  = Env2 Ty2 -> TyEnv (ArrowTy Ty2)
forall a. Env2 a -> TyEnv (ArrowTy a)
fEnv Env2 (TyOf (E2 LocArg Ty2))
Env2 Ty2
env2 Map LocVar (ArrowTy2 Ty2) -> LocVar -> ArrowTy2 Ty2
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# LocVar
v
                             outty :: Ty2
outty = ArrowTy2 Ty2 -> Ty2
forall ty2. ArrowTy2 ty2 -> ty2
Old.arrOut ArrowTy2 Ty2
fnty
                             mp :: Map LocVar LocVar
mp = [(LocVar, LocVar)] -> Map LocVar LocVar
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(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 (ArrowTy2 Ty2 -> [LocVar]
forall ty2. ArrowTy2 ty2 -> [LocVar]
Old.allLocVars ArrowTy2 Ty2
fnty) ((LocArg -> LocVar) -> [LocArg] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocVar
toLocVar [LocArg]
locargs)
                         in Map LocVar LocVar -> Ty2 -> Ty2
substLoc Map LocVar LocVar
mp Ty2
outty
      E2 LocArg Ty2
SyncE -> UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$ UrTy LocVar
forall loc. UrTy loc
voidTy
      WithArenaE LocVar
_v E2 LocArg Ty2
e -> DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
env2 E2 LocArg Ty2
e
      CaseE E2 LocArg Ty2
_ [(String, [(LocVar, LocArg)], E2 LocArg Ty2)]
mp ->
        let (String
c,[(LocVar, LocArg)]
vlocargs,E2 LocArg Ty2
e) = [(String, [(LocVar, LocArg)], E2 LocArg Ty2)]
-> (String, [(LocVar, LocArg)], E2 LocArg Ty2)
forall a. HasCallStack => [a] -> a
head [(String, [(LocVar, LocArg)], E2 LocArg Ty2)]
mp
            ([LocVar]
vars,[LocArg]
locargs) = [(LocVar, LocArg)] -> ([LocVar], [LocArg])
forall a b. [(a, b)] -> ([a], [b])
unzip [(LocVar, LocArg)]
vlocargs
            locs :: [LocVar]
locs = (LocArg -> LocVar) -> [LocArg] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocVar
toLocVar [LocArg]
locargs

            env2' :: Env2 Ty2
env2' = HasCallStack =>
String -> DDefs Ty2 -> [LocVar] -> [LocVar] -> Env2 Ty2 -> Env2 Ty2
String -> DDefs Ty2 -> [LocVar] -> [LocVar] -> Env2 Ty2 -> Env2 Ty2
extendPatternMatchEnv String
c DDefs (TyOf (E2 LocArg Ty2))
DDefs Ty2
ddfs [LocVar]
vars [LocVar]
locs Env2 (TyOf (E2 LocArg Ty2))
Env2 Ty2
env2
        in DDefs (TyOf (E2 LocArg Ty2))
-> Env2 (TyOf (E2 LocArg Ty2))
-> E2 LocArg Ty2
-> TyOf (E2 LocArg Ty2)
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf (E2 LocArg Ty2))
ddfs Env2 (TyOf (E2 LocArg Ty2))
Env2 Ty2
env2' E2 LocArg Ty2
e

-------------------------------------------------------------------------------
-- Need to redefine the following because of the Ty2 newtype:

-- | Apply a location substitution to a type.
substLoc :: M.Map LocVar LocVar -> Ty2 -> Ty2
substLoc :: Map LocVar LocVar -> Ty2 -> Ty2
substLoc Map LocVar LocVar
mp Ty2
ty = UrTy LocVar -> Ty2
MkTy2 (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall a b. (a -> b) -> a -> b
$
  case Ty2 -> UrTy LocVar
unTy2 Ty2
ty of
   SymDictTy Maybe LocVar
v UrTy ()
te -> Maybe LocVar -> UrTy () -> UrTy LocVar
forall loc. Maybe LocVar -> UrTy () -> UrTy loc
SymDictTy Maybe LocVar
v UrTy ()
te -- (go te)
   ProdTy    [UrTy LocVar]
ts -> [UrTy LocVar] -> UrTy LocVar
forall loc. [UrTy loc] -> UrTy loc
ProdTy ((UrTy LocVar -> UrTy LocVar) -> [UrTy LocVar] -> [UrTy LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map (Ty2 -> UrTy LocVar
unTy2 (Ty2 -> UrTy LocVar)
-> (UrTy LocVar -> Ty2) -> UrTy LocVar -> UrTy LocVar
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ty2 -> Ty2
go (Ty2 -> Ty2) -> (UrTy LocVar -> Ty2) -> UrTy LocVar -> Ty2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UrTy LocVar -> Ty2
MkTy2) [UrTy LocVar]
ts)
   PackedTy String
k LocVar
l ->
       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
mp of
             Just LocVar
v  -> String -> LocVar -> UrTy LocVar
forall loc. String -> loc -> UrTy loc
PackedTy String
k LocVar
v
             Maybe LocVar
Nothing -> String -> LocVar -> UrTy LocVar
forall loc. String -> loc -> UrTy loc
PackedTy String
k LocVar
l
   UrTy LocVar
_ -> Ty2 -> UrTy LocVar
unTy2 Ty2
ty
  where go :: Ty2 -> Ty2
go = Map LocVar LocVar -> Ty2 -> Ty2
substLoc Map LocVar LocVar
mp

-- | List version of 'substLoc'.
substLocs :: M.Map LocVar LocVar -> [Ty2] -> [Ty2]
substLocs :: Map LocVar LocVar -> [Ty2] -> [Ty2]
substLocs Map LocVar LocVar
mp [Ty2]
tys = (Ty2 -> Ty2) -> [Ty2] -> [Ty2]
forall a b. (a -> b) -> [a] -> [b]
L.map (Map LocVar LocVar -> Ty2 -> Ty2
substLoc Map LocVar LocVar
mp) [Ty2]
tys

-- | Extend an environment for a pattern match. E.g.
--
--     data Foo = MkFoo Int Foo | ...
--
--     case foo1 of
--        MkFoo (i:loc1) (f:loc2) ->
--          new_env2 = extendPatternMatchEnv [loc1,loc2] old_env2
extendPatternMatchEnv :: HasCallStack => DataCon -> DDefs Ty2 -> [Var] -> [LocVar]
                      -> Env2 Ty2 -> Env2 Ty2
extendPatternMatchEnv :: HasCallStack =>
String -> DDefs Ty2 -> [LocVar] -> [LocVar] -> Env2 Ty2 -> Env2 Ty2
extendPatternMatchEnv String
dcon DDefs Ty2
ddefs [LocVar]
vars [LocVar]
locs Env2 Ty2
env2 =
  let tys :: [Ty2]
tys  = DDefs Ty2 -> String -> [Ty2]
forall a. Out a => DDefs a -> String -> [a]
lookupDataCon DDefs Ty2
ddefs String
dcon
      tys' :: [Ty2]
tys' = ((LocVar, Ty2) -> [Ty2] -> [Ty2])
-> [Ty2] -> [(LocVar, Ty2)] -> [Ty2]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
               (\(LocVar
loc,Ty2
ty) [Ty2]
acc ->
                  case Ty2 -> [LocVar]
locsInTy Ty2
ty of
                    []     -> Ty2
tyTy2 -> [Ty2] -> [Ty2]
forall a. a -> [a] -> [a]
:[Ty2]
acc
                    [LocVar
loc2] -> (Map LocVar LocVar -> Ty2 -> Ty2
substLoc (LocVar -> LocVar -> Map LocVar LocVar
forall k a. k -> a -> Map k a
M.singleton LocVar
loc2 LocVar
loc) Ty2
ty) Ty2 -> [Ty2] -> [Ty2]
forall a. a -> [a] -> [a]
: [Ty2]
acc
                    [LocVar]
_  -> String -> [Ty2]
forall a. HasCallStack => String -> a
error (String -> [Ty2]) -> String -> [Ty2]
forall a b. (a -> b) -> a -> b
$ String
"extendPatternMatchEnv': Found more than 1 location in type: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Ty2 -> String
forall a. Out a => a -> String
sdoc Ty2
ty)
               []
               ([LocVar] -> [Ty2] -> [(LocVar, Ty2)]
forall a b.
(Show a, Show b, HasCallStack) =>
[a] -> [b] -> [(a, b)]
fragileZip [LocVar]
locs [Ty2]
tys)
  in 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)] -> Map LocVar Ty2)
-> [(LocVar, Ty2)] -> Map LocVar Ty2
forall a b. (a -> b) -> a -> b
$ [LocVar] -> [Ty2] -> [(LocVar, Ty2)]
forall a b.
(Show a, Show b, HasCallStack) =>
[a] -> [b] -> [(a, b)]
fragileZip [LocVar]
vars [Ty2]
tys') Env2 Ty2
env2

-- | Collect all the locations mentioned in a type.
locsInTy :: Ty2 -> [LocVar]
locsInTy :: Ty2 -> [LocVar]
locsInTy Ty2
ty =
    case Ty2 -> UrTy LocVar
unTy2 Ty2
ty of
      PackedTy String
_ LocVar
lv -> [LocVar
lv]
      ProdTy [UrTy LocVar]
tys -> (UrTy LocVar -> [LocVar]) -> [UrTy LocVar] -> [LocVar]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Ty2 -> [LocVar]
locsInTy (Ty2 -> [LocVar])
-> (UrTy LocVar -> Ty2) -> UrTy LocVar -> [LocVar]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UrTy LocVar -> Ty2
MkTy2) [UrTy LocVar]
tys
      UrTy LocVar
_ -> []

-- Because L2 just adds a bit of metadata and enriched types, it is
-- possible to strip it back down to L1.
revertToL1 :: Prog2 -> Prog1
revertToL1 :: Prog2 -> Prog1
revertToL1 Prog{DDefs (TyOf (E2 LocArg Ty2))
ddefs :: DDefs (TyOf (E2 LocArg Ty2))
ddefs :: forall ex. Prog ex -> DDefs (TyOf ex)
ddefs,FunDefs (E2 LocArg Ty2)
fundefs :: FunDefs (E2 LocArg Ty2)
fundefs :: forall ex. Prog ex -> FunDefs ex
fundefs,Maybe (E2 LocArg Ty2, TyOf (E2 LocArg Ty2))
mainExp :: Maybe (E2 LocArg Ty2, TyOf (E2 LocArg Ty2))
mainExp :: forall ex. Prog ex -> Maybe (ex, TyOf ex)
mainExp} =
  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)
Map LocVar (DDef (UrTy ()))
ddefs' FunDefs Exp1
funefs' Maybe (Exp1, TyOf Exp1)
Maybe (Exp1, UrTy ())
mainExp'
  where
    ddefs' :: Map LocVar (DDef (UrTy ()))
ddefs'   = (DDef Ty2 -> DDef (UrTy ()))
-> DDefs Ty2 -> Map LocVar (DDef (UrTy ()))
forall a b k. (a -> b) -> Map k a -> Map k b
M.map DDef Ty2 -> DDef (UrTy ())
revertDDef DDefs (TyOf (E2 LocArg Ty2))
DDefs Ty2
ddefs
    funefs' :: FunDefs Exp1
funefs'  = (FunDef2 -> FunDef1) -> FunDefs (E2 LocArg Ty2) -> FunDefs Exp1
forall a b k. (a -> b) -> Map k a -> Map k b
M.map FunDef2 -> FunDef1
revertFunDef FunDefs (E2 LocArg Ty2)
fundefs
    mainExp' :: Maybe (Exp1, UrTy ())
mainExp' = case Maybe (E2 LocArg Ty2, TyOf (E2 LocArg Ty2))
mainExp of
                Maybe (E2 LocArg Ty2, TyOf (E2 LocArg Ty2))
Nothing -> Maybe (Exp1, UrTy ())
forall a. Maybe a
Nothing
                Just (E2 LocArg Ty2
e,TyOf (E2 LocArg Ty2)
ty) -> (Exp1, UrTy ()) -> Maybe (Exp1, UrTy ())
forall a. a -> Maybe a
Just (E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
e, UrTy LocVar -> UrTy ()
forall a. UrTy a -> UrTy ()
stripTyLocs (Ty2 -> UrTy LocVar
unTy2 TyOf (E2 LocArg Ty2)
Ty2
ty))

revertDDef :: DDef Ty2 -> DDef Ty1
revertDDef :: DDef Ty2 -> DDef (UrTy ())
revertDDef (DDef LocVar
tyargs [TyVar]
a [(String, [(Bool, Ty2)])]
b) =
  LocVar
-> [TyVar] -> [(String, [(Bool, UrTy ())])] -> DDef (UrTy ())
forall a. LocVar -> [TyVar] -> [(String, [(Bool, a)])] -> DDef a
DDef LocVar
tyargs [TyVar]
a
    (((String, [(Bool, UrTy ())]) -> Bool)
-> [(String, [(Bool, UrTy ())])] -> [(String, [(Bool, UrTy ())])]
forall a. (a -> Bool) -> [a] -> [a]
L.filter (\(String
dcon,[(Bool, UrTy ())]
_) -> Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ String -> Bool
isIndirectionTag String
dcon) ([(String, [(Bool, UrTy ())])] -> [(String, [(Bool, UrTy ())])])
-> [(String, [(Bool, UrTy ())])] -> [(String, [(Bool, UrTy ())])]
forall a b. (a -> b) -> a -> b
$
         ((String, [(Bool, Ty2)]) -> (String, [(Bool, UrTy ())]))
-> [(String, [(Bool, Ty2)])] -> [(String, [(Bool, UrTy ())])]
forall a b. (a -> b) -> [a] -> [b]
L.map (\(String
dcon,[(Bool, Ty2)]
tys) -> (String
dcon, ((Bool, Ty2) -> (Bool, UrTy ()))
-> [(Bool, Ty2)] -> [(Bool, UrTy ())]
forall a b. (a -> b) -> [a] -> [b]
L.map (\(Bool
x,Ty2
y) -> (Bool
x, UrTy LocVar -> UrTy ()
forall a. UrTy a -> UrTy ()
stripTyLocs (Ty2 -> UrTy LocVar
unTy2 Ty2
y))) [(Bool, Ty2)]
tys)) [(String, [(Bool, Ty2)])]
b)

revertFunDef :: FunDef2 -> FunDef1
revertFunDef :: FunDef2 -> FunDef1
revertFunDef FunDef{LocVar
funName :: LocVar
funName :: forall ex. FunDef ex -> LocVar
funName,[LocVar]
funArgs :: [LocVar]
funArgs :: forall ex. FunDef ex -> [LocVar]
funArgs,ArrowTy (TyOf (E2 LocArg Ty2))
funTy :: ArrowTy (TyOf (E2 LocArg Ty2))
funTy :: forall ex. FunDef ex -> ArrowTy (TyOf ex)
funTy,E2 LocArg Ty2
funBody :: E2 LocArg Ty2
funBody :: forall ex. FunDef ex -> ex
funBody,FunMeta
funMeta :: FunMeta
funMeta :: forall ex. FunDef ex -> FunMeta
funMeta} =
  FunDef { funName :: LocVar
funName = LocVar
funName
         , funArgs :: [LocVar]
funArgs = [LocVar]
funArgs
         , funTy :: ArrowTy (TyOf Exp1)
funTy   = ((Ty2 -> UrTy ()) -> [Ty2] -> [UrTy ()]
forall a b. (a -> b) -> [a] -> [b]
L.map (UrTy LocVar -> UrTy ()
forall a. UrTy a -> UrTy ()
stripTyLocs (UrTy LocVar -> UrTy ()) -> (Ty2 -> UrTy LocVar) -> Ty2 -> UrTy ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ty2 -> UrTy LocVar
unTy2) (ArrowTy2 Ty2 -> [Ty2]
forall ty2. ArrowTy2 ty2 -> [ty2]
Old.arrIns ArrowTy (TyOf (E2 LocArg Ty2))
ArrowTy2 Ty2
funTy), UrTy LocVar -> UrTy ()
forall a. UrTy a -> UrTy ()
stripTyLocs (Ty2 -> UrTy LocVar
unTy2 (ArrowTy2 Ty2 -> Ty2
forall ty2. ArrowTy2 ty2 -> ty2
Old.arrOut ArrowTy (TyOf (E2 LocArg Ty2))
ArrowTy2 Ty2
funTy)))
         , funBody :: Exp1
funBody = E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
funBody
         , funMeta :: FunMeta
funMeta = FunMeta
funMeta
         }

revertExp :: Exp2 -> Exp1
revertExp :: E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
ex =
  case E2 LocArg Ty2
ex of
    VarE LocVar
v    -> LocVar -> Exp1
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
v
    LitE Int
n    -> Int -> Exp1
forall (ext :: * -> * -> *) loc dec. Int -> PreExp ext loc dec
LitE Int
n
    CharE Char
n  -> Char -> Exp1
forall (ext :: * -> * -> *) loc dec. Char -> PreExp ext loc dec
CharE Char
n
    FloatE Double
n  -> Double -> Exp1
forall (ext :: * -> * -> *) loc dec. Double -> PreExp ext loc dec
FloatE Double
n
    LitSymE LocVar
v -> LocVar -> Exp1
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
LitSymE LocVar
v
    AppE LocVar
v [LocArg]
_ [E2 LocArg Ty2]
args   -> LocVar -> [()] -> [Exp1] -> Exp1
forall (ext :: * -> * -> *) loc dec.
LocVar -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE LocVar
v [] ((E2 LocArg Ty2 -> Exp1) -> [E2 LocArg Ty2] -> [Exp1]
forall a b. (a -> b) -> [a] -> [b]
L.map E2 LocArg Ty2 -> Exp1
revertExp [E2 LocArg Ty2]
args)
    PrimAppE Prim Ty2
p [E2 LocArg Ty2]
args -> Prim (UrTy ()) -> [Exp1] -> Exp1
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE (Prim Ty2 -> Prim (UrTy ())
revertPrim Prim Ty2
p) ([Exp1] -> Exp1) -> [Exp1] -> Exp1
forall a b. (a -> b) -> a -> b
$ (E2 LocArg Ty2 -> Exp1) -> [E2 LocArg Ty2] -> [Exp1]
forall a b. (a -> b) -> [a] -> [b]
L.map E2 LocArg Ty2 -> Exp1
revertExp [E2 LocArg Ty2]
args
    LetE (LocVar
v,[LocArg]
_, Ty2
ty, (Ext (Old.IndirectionE String
_ String
_ (LocArg, LocArg)
_ (LocArg, LocArg)
_ E2 LocArg Ty2
arg))) E2 LocArg Ty2
bod ->
      let PackedTy String
tycon LocVar
_ =  Ty2 -> UrTy LocVar
unTy2 Ty2
ty in
          (LocVar, [()], UrTy (), Exp1) -> Exp1 -> Exp1
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[],(UrTy LocVar -> UrTy ()
forall a. UrTy a -> UrTy ()
stripTyLocs (Ty2 -> UrTy LocVar
unTy2 Ty2
ty)), LocVar -> [()] -> [Exp1] -> Exp1
forall (ext :: * -> * -> *) loc dec.
LocVar -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE (String -> LocVar
mkCopyFunName String
tycon) [] [E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
arg]) (E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
bod)
    LetE (LocVar
v,[LocArg]
_,Ty2
ty,E2 LocArg Ty2
rhs) E2 LocArg Ty2
bod ->
      (LocVar, [()], UrTy (), Exp1) -> Exp1 -> Exp1
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (LocVar
v,[], UrTy LocVar -> UrTy ()
forall a. UrTy a -> UrTy ()
stripTyLocs (Ty2 -> UrTy LocVar
unTy2 Ty2
ty), E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
rhs) (E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
bod)
    IfE E2 LocArg Ty2
a E2 LocArg Ty2
b E2 LocArg Ty2
c  -> Exp1 -> Exp1 -> Exp1 -> Exp1
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> PreExp ext loc dec -> PreExp ext loc dec -> PreExp ext loc dec
IfE (E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
a) (E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
b) (E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
c)
    MkProdE [E2 LocArg Ty2]
ls -> [Exp1] -> Exp1
forall (ext :: * -> * -> *) loc dec.
[PreExp ext loc dec] -> PreExp ext loc dec
MkProdE ([Exp1] -> Exp1) -> [Exp1] -> Exp1
forall a b. (a -> b) -> a -> b
$ (E2 LocArg Ty2 -> Exp1) -> [E2 LocArg Ty2] -> [Exp1]
forall a b. (a -> b) -> [a] -> [b]
L.map E2 LocArg Ty2 -> Exp1
revertExp [E2 LocArg Ty2]
ls
    ProjE Int
i E2 LocArg Ty2
e  -> Int -> Exp1 -> Exp1
forall (ext :: * -> * -> *) loc dec.
Int -> PreExp ext loc dec -> PreExp ext loc dec
ProjE Int
i (E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
e)
    CaseE E2 LocArg Ty2
scrt [(String, [(LocVar, LocArg)], E2 LocArg Ty2)]
brs     -> Exp1 -> [(String, [(LocVar, ())], Exp1)] -> Exp1
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(String, [(LocVar, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE (E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
scrt) (((String, [(LocVar, LocArg)], E2 LocArg Ty2)
 -> (String, [(LocVar, ())], Exp1))
-> [(String, [(LocVar, LocArg)], E2 LocArg Ty2)]
-> [(String, [(LocVar, ())], Exp1)]
forall a b. (a -> b) -> [a] -> [b]
L.map (String, [(LocVar, LocArg)], E2 LocArg Ty2)
-> (String, [(LocVar, ())], Exp1)
docase [(String, [(LocVar, LocArg)], E2 LocArg Ty2)]
brs)
    DataConE LocArg
_ String
dcon [E2 LocArg Ty2]
ls -> () -> String -> [Exp1] -> Exp1
forall (ext :: * -> * -> *) loc dec.
loc -> String -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE () String
dcon ([Exp1] -> Exp1) -> [Exp1] -> Exp1
forall a b. (a -> b) -> a -> b
$ (E2 LocArg Ty2 -> Exp1) -> [E2 LocArg Ty2] -> [Exp1]
forall a b. (a -> b) -> [a] -> [b]
L.map E2 LocArg Ty2 -> Exp1
revertExp [E2 LocArg Ty2]
ls
    TimeIt E2 LocArg Ty2
e Ty2
ty Bool
b -> Exp1 -> UrTy () -> Bool -> Exp1
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> Bool -> PreExp ext loc dec
TimeIt (E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
e) (UrTy LocVar -> UrTy ()
forall a. UrTy a -> UrTy ()
stripTyLocs (Ty2 -> UrTy LocVar
unTy2 Ty2
ty)) Bool
b
    SpawnE LocVar
v [LocArg]
_ [E2 LocArg Ty2]
args -> LocVar -> [()] -> [Exp1] -> Exp1
forall (ext :: * -> * -> *) loc dec.
LocVar -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
SpawnE LocVar
v [] ((E2 LocArg Ty2 -> Exp1) -> [E2 LocArg Ty2] -> [Exp1]
forall a b. (a -> b) -> [a] -> [b]
L.map E2 LocArg Ty2 -> Exp1
revertExp [E2 LocArg Ty2]
args)
    E2 LocArg Ty2
SyncE -> Exp1
forall (ext :: * -> * -> *) loc dec. PreExp ext loc dec
SyncE
    WithArenaE LocVar
v E2 LocArg Ty2
e -> LocVar -> Exp1 -> Exp1
forall (ext :: * -> * -> *) loc dec.
LocVar -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE LocVar
v (E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
e)
    Ext E2Ext LocArg Ty2
ext ->
      case E2Ext LocArg Ty2
ext of
        Old.LetRegionE Region
_ RegionSize
_ Maybe RegionType
_ E2 LocArg Ty2
bod -> E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
bod
        Old.LetParRegionE Region
_ RegionSize
_ Maybe RegionType
_ E2 LocArg Ty2
bod -> E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
bod
        Old.LetLocE LocVar
_ LocExp
_ E2 LocArg Ty2
bod  -> E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
bod
        Old.TagCursor LocVar
a LocVar
_b -> E1Ext () (UrTy ()) -> Exp1
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (LocVar -> E1Ext () (UrTy ())
forall loc dec. LocVar -> E1Ext loc dec
L1.StartOfPkdCursor LocVar
a)
        Old.StartOfPkdCursor LocVar
v -> E1Ext () (UrTy ()) -> Exp1
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (LocVar -> E1Ext () (UrTy ())
forall loc dec. LocVar -> E1Ext loc dec
L1.StartOfPkdCursor LocVar
v)
        Old.RetE [LocArg]
_ LocVar
v -> LocVar -> Exp1
forall (ext :: * -> * -> *) loc dec. LocVar -> PreExp ext loc dec
VarE LocVar
v
        Old.AddFixed{} -> String -> Exp1
forall a. HasCallStack => String -> a
error String
"revertExp: TODO AddFixed."
        Old.FromEndE{} -> String -> Exp1
forall a. HasCallStack => String -> a
error String
"revertExp: TODO FromEndLE"
        Old.BoundsCheck{}   -> String -> Exp1
forall a. HasCallStack => String -> a
error String
"revertExp: TODO BoundsCheck"
        Old.IndirectionE{}  -> String -> Exp1
forall a. HasCallStack => String -> a
error String
"revertExp: TODO IndirectionE"
        E2Ext LocArg Ty2
Old.GetCilkWorkerNum-> Int -> Exp1
forall (ext :: * -> * -> *) loc dec. Int -> PreExp ext loc dec
LitE Int
0
        Old.LetAvail [LocVar]
_ E2 LocArg Ty2
bod  -> E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
bod
        Old.AllocateTagHere{} -> String -> Exp1
forall a. HasCallStack => String -> a
error String
"revertExp: TODO AddFixed."
        Old.AllocateScalarsHere{} -> String -> Exp1
forall a. HasCallStack => String -> a
error String
"revertExp: TODO AddFixed."
        Old.SSPush{} -> String -> Exp1
forall a. HasCallStack => String -> a
error String
"revertExp: TODO SSPush."
        Old.SSPop{} -> String -> Exp1
forall a. HasCallStack => String -> a
error String
"revertExp: TODO SSPop."
    MapE{}  -> String -> Exp1
forall a. HasCallStack => String -> a
error (String -> Exp1) -> String -> Exp1
forall a b. (a -> b) -> a -> b
$ String
"revertExp: TODO MapE"
    FoldE{} -> String -> Exp1
forall a. HasCallStack => String -> a
error (String -> Exp1) -> String -> Exp1
forall a b. (a -> b) -> a -> b
$ String
"revertExp: TODO FoldE"
  where
    -- Ugh .. this is bad. Can we remove the identity cases here ?
    -- TODO: Get rid of this (and L3.toL3Prim) soon.
    revertPrim :: Prim Ty2 -> Prim Ty1
    revertPrim :: Prim Ty2 -> Prim (UrTy ())
revertPrim Prim Ty2
pr = (Ty2 -> UrTy ()) -> Prim Ty2 -> Prim (UrTy ())
forall a b. (a -> b) -> Prim a -> Prim b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (UrTy LocVar -> UrTy ()
forall a. UrTy a -> UrTy ()
stripTyLocs (UrTy LocVar -> UrTy ()) -> (Ty2 -> UrTy LocVar) -> Ty2 -> UrTy ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ty2 -> UrTy LocVar
unTy2) Prim Ty2
pr

    docase :: (DataCon, [(Var,LocArg)], Exp2) -> (DataCon, [(Var,())], Exp1)
    docase :: (String, [(LocVar, LocArg)], E2 LocArg Ty2)
-> (String, [(LocVar, ())], Exp1)
docase (String
dcon,[(LocVar, LocArg)]
vlocargs,E2 LocArg Ty2
rhs) =
      let ([LocVar]
vars,[LocArg]
_) = [(LocVar, LocArg)] -> ([LocVar], [LocArg])
forall a b. [(a, b)] -> ([a], [b])
unzip [(LocVar, LocArg)]
vlocargs
      in (String
dcon, [LocVar] -> [()] -> [(LocVar, ())]
forall a b. [a] -> [b] -> [(a, b)]
zip [LocVar]
vars (() -> [()]
forall a. a -> [a]
repeat ()), E2 LocArg Ty2 -> Exp1
revertExp E2 LocArg Ty2
rhs)


-- | Build a dependency list which can be later converted to a graph
depList :: Exp2 -> [(Var, Var, [Var])]
-- The helper function, go, works with a map rather than list so that all
-- dependencies are properly grouped, without any duplicate keys. But we
-- convert it back to a list so that we can hand it off to 'graphFromEdges'.
-- Reversing the list makes it easy to peek at the return value of this AST later.
depList :: E2 LocArg Ty2 -> [(LocVar, LocVar, [LocVar])]
depList = ((LocVar, [LocVar]) -> (LocVar, LocVar, [LocVar]))
-> [(LocVar, [LocVar])] -> [(LocVar, LocVar, [LocVar])]
forall a b. (a -> b) -> [a] -> [b]
L.map (\(LocVar
a,[LocVar]
b) -> (LocVar
a,LocVar
a,[LocVar]
b)) ([(LocVar, [LocVar])] -> [(LocVar, LocVar, [LocVar])])
-> (E2 LocArg Ty2 -> [(LocVar, [LocVar])])
-> E2 LocArg Ty2
-> [(LocVar, LocVar, [LocVar])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map LocVar [LocVar] -> [(LocVar, [LocVar])]
forall k a. Map k a -> [(k, a)]
M.toList (Map LocVar [LocVar] -> [(LocVar, [LocVar])])
-> (E2 LocArg Ty2 -> Map LocVar [LocVar])
-> E2 LocArg Ty2
-> [(LocVar, [LocVar])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
forall k a. Map k a
M.empty
    where
      go :: M.Map Var [Var] -> Exp2 -> M.Map Var [Var]
      go :: Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc E2 LocArg Ty2
ex =
        case E2 LocArg Ty2
ex of
          VarE LocVar
v    -> ([LocVar] -> [LocVar] -> [LocVar])
-> LocVar -> [LocVar] -> Map LocVar [LocVar] -> Map LocVar [LocVar]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
(++) LocVar
v [LocVar
v] Map LocVar [LocVar]
acc
          LitE{}    -> Map LocVar [LocVar]
acc
          CharE{}  -> Map LocVar [LocVar]
acc
          FloatE{}  -> Map LocVar [LocVar]
acc
          LitSymE{} -> Map LocVar [LocVar]
acc
          AppE LocVar
_ [LocArg]
_ [E2 LocArg Ty2]
args   -> (Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar])
-> Map LocVar [LocVar] -> [E2 LocArg Ty2] -> Map LocVar [LocVar]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc [E2 LocArg Ty2]
args
          PrimAppE Prim Ty2
_ [E2 LocArg Ty2]
args -> (Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar])
-> Map LocVar [LocVar] -> [E2 LocArg Ty2] -> Map LocVar [LocVar]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc [E2 LocArg Ty2]
args
          LetE (LocVar
v,[LocArg]
_,Ty2
_,E2 LocArg Ty2
rhs) E2 LocArg Ty2
bod ->
            let acc_rhs :: Map LocVar [LocVar]
acc_rhs = Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc E2 LocArg Ty2
rhs
            in Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go (([LocVar] -> [LocVar] -> [LocVar])
-> LocVar -> [LocVar] -> Map LocVar [LocVar] -> Map LocVar [LocVar]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
(++) LocVar
v (Set LocVar -> [LocVar]
forall a. Set a -> [a]
S.toList (Set LocVar -> [LocVar]) -> Set LocVar -> [LocVar]
forall a b. (a -> b) -> a -> b
$ E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
rhs) Map LocVar [LocVar]
acc_rhs) E2 LocArg Ty2
bod
          IfE E2 LocArg Ty2
_ E2 LocArg Ty2
b E2 LocArg Ty2
c  -> Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go (Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc E2 LocArg Ty2
b) E2 LocArg Ty2
c
          MkProdE [E2 LocArg Ty2]
ls -> (Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar])
-> Map LocVar [LocVar] -> [E2 LocArg Ty2] -> Map LocVar [LocVar]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc [E2 LocArg Ty2]
ls
          ProjE Int
_ E2 LocArg Ty2
e  -> Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc E2 LocArg Ty2
e
          CaseE (VarE LocVar
v) [(String, [(LocVar, LocArg)], E2 LocArg Ty2)]
mp ->
            ((String, [(LocVar, LocArg)], E2 LocArg Ty2)
 -> Map LocVar [LocVar] -> Map LocVar [LocVar])
-> Map LocVar [LocVar]
-> [(String, [(LocVar, LocArg)], E2 LocArg 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 (\(String
_,[(LocVar, LocArg)]
vlocs,E2 LocArg Ty2
e) Map LocVar [LocVar]
acc' ->
                       let ([LocVar]
vars,[LocArg]
locs) = [(LocVar, LocArg)] -> ([LocVar], [LocArg])
forall a b. [(a, b)] -> ([a], [b])
unzip [(LocVar, LocArg)]
vlocs
                           acc'' :: Map LocVar [LocVar]
acc'' = (LocVar -> Map LocVar [LocVar] -> Map LocVar [LocVar])
-> Map 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
w Map LocVar [LocVar]
acc''' -> ([LocVar] -> [LocVar] -> [LocVar])
-> LocVar -> [LocVar] -> Map LocVar [LocVar] -> Map LocVar [LocVar]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
(++) LocVar
v [LocVar
w] Map LocVar [LocVar]
acc''')
                                           Map LocVar [LocVar]
acc'
                                           ([LocVar]
vars [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
++ ((LocArg -> LocVar) -> [LocArg] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocVar
toLocVar [LocArg]
locs))
                       in Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc'' E2 LocArg Ty2
e)
                    Map LocVar [LocVar]
acc
                    [(String, [(LocVar, LocArg)], E2 LocArg Ty2)]
mp
          CaseE E2 LocArg Ty2
_scrt [(String, [(LocVar, LocArg)], E2 LocArg Ty2)]
mp -> ((String, [(LocVar, LocArg)], E2 LocArg Ty2)
 -> Map LocVar [LocVar] -> Map LocVar [LocVar])
-> Map LocVar [LocVar]
-> [(String, [(LocVar, LocArg)], E2 LocArg 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 (\(String
_,[(LocVar, LocArg)]
_,E2 LocArg Ty2
e) Map LocVar [LocVar]
acc' -> Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc' E2 LocArg Ty2
e) Map LocVar [LocVar]
acc [(String, [(LocVar, LocArg)], E2 LocArg Ty2)]
mp
          DataConE LocArg
_ String
_ [E2 LocArg Ty2]
args -> (Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar])
-> Map LocVar [LocVar] -> [E2 LocArg Ty2] -> Map LocVar [LocVar]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc [E2 LocArg Ty2]
args
          TimeIt E2 LocArg Ty2
e Ty2
_ Bool
_ -> Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc E2 LocArg Ty2
e
          WithArenaE LocVar
_ E2 LocArg Ty2
e -> Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc E2 LocArg Ty2
e
          SpawnE LocVar
_ [LocArg]
_ [E2 LocArg Ty2]
ls  -> (Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar])
-> Map LocVar [LocVar] -> [E2 LocArg Ty2] -> Map LocVar [LocVar]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc [E2 LocArg Ty2]
ls
          E2 LocArg Ty2
SyncE          -> Map LocVar [LocVar]
acc
          MapE{}  -> Map LocVar [LocVar]
acc
          FoldE{} -> Map LocVar [LocVar]
acc
          Ext E2Ext LocArg Ty2
ext ->
            case E2Ext LocArg Ty2
ext of
              Old.LetRegionE Region
r RegionSize
_ Maybe RegionType
_ E2 LocArg Ty2
rhs ->
                Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go (([LocVar] -> [LocVar] -> [LocVar])
-> LocVar -> [LocVar] -> Map LocVar [LocVar] -> Map LocVar [LocVar]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
(++) (Region -> LocVar
Old.regionToVar Region
r) (Set LocVar -> [LocVar]
forall a. Set a -> [a]
S.toList (Set LocVar -> [LocVar]) -> Set LocVar -> [LocVar]
forall a b. (a -> b) -> a -> b
$ E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
rhs) Map LocVar [LocVar]
acc) E2 LocArg Ty2
rhs
              Old.LetParRegionE Region
r RegionSize
_ Maybe RegionType
_ E2 LocArg Ty2
rhs ->
                Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go (([LocVar] -> [LocVar] -> [LocVar])
-> LocVar -> [LocVar] -> Map LocVar [LocVar] -> Map LocVar [LocVar]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
(++) (Region -> LocVar
Old.regionToVar Region
r) (Set LocVar -> [LocVar]
forall a. Set a -> [a]
S.toList (Set LocVar -> [LocVar]) -> Set LocVar -> [LocVar]
forall a b. (a -> b) -> a -> b
$ E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
rhs) Map LocVar [LocVar]
acc) E2 LocArg Ty2
rhs
              Old.LetLocE LocVar
loc LocExp
phs E2 LocArg Ty2
rhs  ->
                Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go (([LocVar] -> [LocVar] -> [LocVar])
-> LocVar -> [LocVar] -> Map LocVar [LocVar] -> Map LocVar [LocVar]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
(++) LocVar
loc (LocExp -> [LocVar]
dep LocExp
phs [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
++ (Set LocVar -> [LocVar]
forall a. Set a -> [a]
S.toList (Set LocVar -> [LocVar]) -> Set LocVar -> [LocVar]
forall a b. (a -> b) -> a -> b
$ E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
rhs)) Map LocVar [LocVar]
acc) E2 LocArg Ty2
rhs
              Old.RetE{}         -> Map LocVar [LocVar]
acc
              Old.FromEndE{}     -> Map LocVar [LocVar]
acc
              Old.BoundsCheck{}  -> Map LocVar [LocVar]
acc
              Old.IndirectionE{} -> Map LocVar [LocVar]
acc
              Old.AddFixed LocVar
v Int
_   -> ([LocVar] -> [LocVar] -> [LocVar])
-> LocVar -> [LocVar] -> Map LocVar [LocVar] -> Map LocVar [LocVar]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
(++) LocVar
v [LocVar
v] Map LocVar [LocVar]
acc
              E2Ext LocArg Ty2
Old.GetCilkWorkerNum -> Map LocVar [LocVar]
acc
              Old.LetAvail [LocVar]
_ E2 LocArg Ty2
bod -> Map LocVar [LocVar] -> E2 LocArg Ty2 -> Map LocVar [LocVar]
go Map LocVar [LocVar]
acc E2 LocArg Ty2
bod
              Old.AllocateTagHere{} -> Map LocVar [LocVar]
acc
              Old.AllocateScalarsHere{} -> Map LocVar [LocVar]
acc
              Old.SSPush{} -> Map LocVar [LocVar]
acc
              Old.SSPop{} -> Map LocVar [LocVar]
acc
              Old.StartOfPkdCursor LocVar
cur -> ([LocVar] -> [LocVar] -> [LocVar])
-> LocVar -> [LocVar] -> Map LocVar [LocVar] -> Map LocVar [LocVar]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
(++) LocVar
cur [LocVar
cur] Map LocVar [LocVar]
acc
              Old.TagCursor LocVar
a LocVar
b -> ([LocVar] -> [LocVar] -> [LocVar])
-> LocVar -> [LocVar] -> Map LocVar [LocVar] -> Map LocVar [LocVar]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
(++) LocVar
b [LocVar
b] (([LocVar] -> [LocVar] -> [LocVar])
-> LocVar -> [LocVar] -> Map LocVar [LocVar] -> Map LocVar [LocVar]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
(++) LocVar
a [LocVar
a] Map LocVar [LocVar]
acc)

      dep :: Old.PreLocExp LocArg -> [Var]
      dep :: LocExp -> [LocVar]
dep LocExp
ex =
        case LocExp
ex of
          Old.StartOfRegionLE Region
r -> [Region -> LocVar
Old.regionToVar Region
r]
          Old.AfterConstantLE Int
_ LocArg
loc   -> [LocArg -> LocVar
toLocVar LocArg
loc]
          Old.AfterVariableLE LocVar
v LocArg
loc Bool
_ -> [LocVar
v,LocArg -> LocVar
toLocVar LocArg
loc]
          Old.InRegionLE Region
r  -> [Region -> LocVar
Old.regionToVar Region
r]
          Old.FromEndLE LocArg
loc -> [LocArg -> LocVar
toLocVar LocArg
loc]
          LocExp
Old.FreeLE -> []

-- gFreeVars ++ locations ++ region variables
allFreeVars :: Exp2 -> S.Set Var
allFreeVars :: E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
ex =
  case E2 LocArg Ty2
ex of
    AppE LocVar
_ [LocArg]
locs [E2 LocArg Ty2]
args -> [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList ((LocArg -> LocVar) -> [LocArg] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocVar
toLocVar [LocArg]
locs) Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` ([Set LocVar] -> Set LocVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((E2 LocArg Ty2 -> Set LocVar) -> [E2 LocArg Ty2] -> [Set LocVar]
forall a b. (a -> b) -> [a] -> [b]
map E2 LocArg Ty2 -> Set LocVar
allFreeVars [E2 LocArg Ty2]
args))
    PrimAppE Prim Ty2
_ [E2 LocArg Ty2]
args -> ([Set LocVar] -> Set LocVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((E2 LocArg Ty2 -> Set LocVar) -> [E2 LocArg Ty2] -> [Set LocVar]
forall a b. (a -> b) -> [a] -> [b]
map E2 LocArg Ty2 -> Set LocVar
allFreeVars [E2 LocArg Ty2]
args))
    LetE (LocVar
v,[LocArg]
locs,Ty2
_,E2 LocArg Ty2
rhs) E2 LocArg Ty2
bod -> ([LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList ((LocArg -> LocVar) -> [LocArg] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocVar
toLocVar [LocArg]
locs) Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` (E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
rhs) Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` (E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
bod))
                               Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` LocVar -> Set LocVar
forall a. a -> Set a
S.singleton LocVar
v
    IfE E2 LocArg Ty2
a E2 LocArg Ty2
b E2 LocArg Ty2
c -> E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
a Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
b Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
c
    MkProdE [E2 LocArg Ty2]
args -> ([Set LocVar] -> Set LocVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((E2 LocArg Ty2 -> Set LocVar) -> [E2 LocArg Ty2] -> [Set LocVar]
forall a b. (a -> b) -> [a] -> [b]
map E2 LocArg Ty2 -> Set LocVar
allFreeVars [E2 LocArg Ty2]
args))
    ProjE Int
_ E2 LocArg Ty2
bod -> E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
bod
    CaseE E2 LocArg Ty2
scrt [(String, [(LocVar, LocArg)], E2 LocArg Ty2)]
brs -> (E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
scrt) Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` ([Set LocVar] -> Set LocVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions (((String, [(LocVar, LocArg)], E2 LocArg Ty2) -> Set LocVar)
-> [(String, [(LocVar, LocArg)], E2 LocArg Ty2)] -> [Set LocVar]
forall a b. (a -> b) -> [a] -> [b]
map (\(String
_,[(LocVar, LocArg)]
vlocs,E2 LocArg Ty2
c) -> E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
c Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.difference`
                                                                                   [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList (((LocVar, LocArg) -> LocVar) -> [(LocVar, LocArg)] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
map (LocVar, LocArg) -> LocVar
forall a b. (a, b) -> a
fst [(LocVar, LocArg)]
vlocs) Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.difference`
                                                                                   [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList (((LocVar, LocArg) -> LocVar) -> [(LocVar, LocArg)] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
map (LocArg -> LocVar
toLocVar (LocArg -> LocVar)
-> ((LocVar, LocArg) -> LocArg) -> (LocVar, LocArg) -> LocVar
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (LocVar, LocArg) -> LocArg
forall a b. (a, b) -> b
snd) [(LocVar, LocArg)]
vlocs))
                                                                  [(String, [(LocVar, LocArg)], E2 LocArg Ty2)]
brs))
    DataConE LocArg
loc String
_ [E2 LocArg Ty2]
args -> LocVar -> Set LocVar
forall a. a -> Set a
S.singleton (LocArg -> LocVar
toLocVar LocArg
loc) Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` ([Set LocVar] -> Set LocVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((E2 LocArg Ty2 -> Set LocVar) -> [E2 LocArg Ty2] -> [Set LocVar]
forall a b. (a -> b) -> [a] -> [b]
map E2 LocArg Ty2 -> Set LocVar
allFreeVars [E2 LocArg Ty2]
args))
    TimeIt E2 LocArg Ty2
e Ty2
_ Bool
_ -> E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
e
    WithArenaE LocVar
_ E2 LocArg Ty2
e -> E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
e
    SpawnE LocVar
_ [LocArg]
locs [E2 LocArg Ty2]
args -> [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList ((LocArg -> LocVar) -> [LocArg] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocVar
toLocVar [LocArg]
locs) Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` ([Set LocVar] -> Set LocVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((E2 LocArg Ty2 -> Set LocVar) -> [E2 LocArg Ty2] -> [Set LocVar]
forall a b. (a -> b) -> [a] -> [b]
map E2 LocArg Ty2 -> Set LocVar
allFreeVars [E2 LocArg Ty2]
args))
    Ext E2Ext LocArg Ty2
ext ->
      case E2Ext LocArg Ty2
ext of
        Old.LetRegionE Region
r RegionSize
_ Maybe RegionType
_ E2 LocArg Ty2
bod -> LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => a -> Set a -> Set a
S.delete (Region -> LocVar
Old.regionToVar Region
r) (E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
bod)
        Old.LetParRegionE Region
r RegionSize
_ Maybe RegionType
_ E2 LocArg Ty2
bod -> LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => a -> Set a -> Set a
S.delete (Region -> LocVar
Old.regionToVar Region
r) (E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
bod)
        Old.LetLocE LocVar
loc LocExp
locexp E2 LocArg Ty2
bod -> LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => a -> Set a -> Set a
S.delete LocVar
loc (E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
bod Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` LocExp -> Set LocVar
forall a. FreeVars a => a -> Set LocVar
gFreeVars LocExp
locexp)
        Old.StartOfPkdCursor LocVar
v -> LocVar -> Set LocVar
forall a. a -> Set a
S.singleton LocVar
v
        Old.TagCursor LocVar
a LocVar
b-> [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar
a,LocVar
b]
        Old.RetE [LocArg]
locs LocVar
v     -> LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => a -> Set a -> Set a
S.insert LocVar
v ([LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList ((LocArg -> LocVar) -> [LocArg] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocVar
toLocVar [LocArg]
locs))
        Old.FromEndE LocArg
loc    -> LocVar -> Set LocVar
forall a. a -> Set a
S.singleton (LocArg -> LocVar
toLocVar LocArg
loc)
        Old.BoundsCheck Int
_ LocArg
reg LocArg
cur -> [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList ((LocArg -> LocVar) -> [LocArg] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocVar
toLocVar [LocArg
reg, LocArg
cur])
        Old.IndirectionE String
_ String
_ (LocArg
a,LocArg
b) (LocArg
c,LocArg
d) E2 LocArg Ty2
_ -> [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList ([LocVar] -> Set LocVar) -> [LocVar] -> Set LocVar
forall a b. (a -> b) -> a -> b
$ [LocArg -> LocVar
toLocVar LocArg
a, LocArg -> LocVar
toLocVar LocArg
b, LocArg -> LocVar
toLocVar LocArg
c, LocArg -> LocVar
toLocVar LocArg
d]
        Old.AddFixed LocVar
v Int
_    -> LocVar -> Set LocVar
forall a. a -> Set a
S.singleton LocVar
v
        E2Ext LocArg Ty2
Old.GetCilkWorkerNum-> Set LocVar
forall a. Set a
S.empty
        Old.LetAvail [LocVar]
vs E2 LocArg Ty2
bod -> [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar]
vs Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` E2 LocArg Ty2 -> Set LocVar
forall a. FreeVars a => a -> Set LocVar
gFreeVars E2 LocArg Ty2
bod
        Old.AllocateTagHere LocVar
loc String
_ -> LocVar -> Set LocVar
forall a. a -> Set a
S.singleton LocVar
loc
        Old.AllocateScalarsHere LocVar
loc -> LocVar -> Set LocVar
forall a. a -> Set a
S.singleton LocVar
loc
        Old.SSPush SSModality
_ LocVar
a LocVar
b String
_ -> [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar
a,LocVar
b]
        Old.SSPop SSModality
_ LocVar
a LocVar
b -> [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar
a,LocVar
b]
    E2 LocArg Ty2
_ -> E2 LocArg Ty2 -> Set LocVar
forall a. FreeVars a => a -> Set LocVar
gFreeVars E2 LocArg Ty2
ex

freeLocVars :: Exp2 -> [Var]
freeLocVars :: E2 LocArg Ty2 -> [LocVar]
freeLocVars E2 LocArg Ty2
ex = Set LocVar -> [LocVar]
forall a. Set a -> [a]
S.toList (Set LocVar -> [LocVar]) -> Set LocVar -> [LocVar]
forall a b. (a -> b) -> a -> b
$ (E2 LocArg Ty2 -> Set LocVar
allFreeVars E2 LocArg Ty2
ex) Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` (E2 LocArg Ty2 -> Set LocVar
forall a. FreeVars a => a -> Set LocVar
gFreeVars E2 LocArg Ty2
ex)