{-|

With the new location calculus and `inferLocations`, this pass should've been
redundant. But not quite.. We still need it to reorder location variables bound
in case expressions.

For example, after running the add1 program through the pipeline till `RouteEnds`,
this is what the `Node` case looks like:

    ("Node",
     [("x9", "l10"),("y11", "l12")],
        ...)

To "unpack" these fields, `Cursorize` just binds `x9` to `l10` and `y11` to `l12`.
But, the cursor `l12` is not bound (or known) until we call `(add1 x9)` and
get the end_of_read cursor it returns. This happens later in the program.
Thus, `y11` refers to an unbound cursor `l12` here.
FindWitnesses fixes this by moving the `let y11 = l12` binding to its proper place.

Another strategy would be to actually handle this properly in Cursorize.

-}

module Gibbon.Passes.FindWitnesses
  (findWitnesses) where


import Data.Graph
import qualified Data.Map as Map
import qualified Data.Set as Set

import Gibbon.Common
-- import Gibbon.L3.Syntax
import Gibbon.L2.Syntax hiding (mapMExprs)

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

data DelayedBind = DelayVar (Var,[LocVar], Ty2, Exp2)
                 | DelayLoc (LocVar, LocExp)
  deriving (DelayedBind -> DelayedBind -> Bool
(DelayedBind -> DelayedBind -> Bool)
-> (DelayedBind -> DelayedBind -> Bool) -> Eq DelayedBind
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DelayedBind -> DelayedBind -> Bool
== :: DelayedBind -> DelayedBind -> Bool
$c/= :: DelayedBind -> DelayedBind -> Bool
/= :: DelayedBind -> DelayedBind -> Bool
Eq, Int -> DelayedBind -> ShowS
[DelayedBind] -> ShowS
DelayedBind -> String
(Int -> DelayedBind -> ShowS)
-> (DelayedBind -> String)
-> ([DelayedBind] -> ShowS)
-> Show DelayedBind
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DelayedBind -> ShowS
showsPrec :: Int -> DelayedBind -> ShowS
$cshow :: DelayedBind -> String
show :: DelayedBind -> String
$cshowList :: [DelayedBind] -> ShowS
showList :: [DelayedBind] -> ShowS
Show, Eq DelayedBind
Eq DelayedBind
-> (DelayedBind -> DelayedBind -> Ordering)
-> (DelayedBind -> DelayedBind -> Bool)
-> (DelayedBind -> DelayedBind -> Bool)
-> (DelayedBind -> DelayedBind -> Bool)
-> (DelayedBind -> DelayedBind -> Bool)
-> (DelayedBind -> DelayedBind -> DelayedBind)
-> (DelayedBind -> DelayedBind -> DelayedBind)
-> Ord DelayedBind
DelayedBind -> DelayedBind -> Bool
DelayedBind -> DelayedBind -> Ordering
DelayedBind -> DelayedBind -> DelayedBind
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 :: DelayedBind -> DelayedBind -> Ordering
compare :: DelayedBind -> DelayedBind -> Ordering
$c< :: DelayedBind -> DelayedBind -> Bool
< :: DelayedBind -> DelayedBind -> Bool
$c<= :: DelayedBind -> DelayedBind -> Bool
<= :: DelayedBind -> DelayedBind -> Bool
$c> :: DelayedBind -> DelayedBind -> Bool
> :: DelayedBind -> DelayedBind -> Bool
$c>= :: DelayedBind -> DelayedBind -> Bool
>= :: DelayedBind -> DelayedBind -> Bool
$cmax :: DelayedBind -> DelayedBind -> DelayedBind
max :: DelayedBind -> DelayedBind -> DelayedBind
$cmin :: DelayedBind -> DelayedBind -> DelayedBind
min :: DelayedBind -> DelayedBind -> DelayedBind
Ord)

bigNumber :: Int
bigNumber :: Int
bigNumber = Int
10 -- limit number of loops

-- | This pass must find witnesses if they exist in the lexical
-- environment, and it must *reorder* let bindings to bring start/end
-- witnesses into scope.
--
-- Phase Ordering: This must run after flatten.
findWitnesses :: Prog2 -> PassM Prog2
findWitnesses :: Prog2 -> PassM Prog2
findWitnesses p :: Prog2
p@Prog{FunDefs Exp2
fundefs :: FunDefs Exp2
fundefs :: forall ex. Prog ex -> FunDefs ex
fundefs} = (Env2 Ty2 -> Set Var -> Exp2 -> PassM Exp2) -> Prog2 -> PassM Prog2
forall (m :: * -> *).
Monad m =>
(Env2 Ty2 -> Set Var -> Exp2 -> m Exp2) -> Prog2 -> m Prog2
mapMExprs Env2 Ty2 -> Set Var -> Exp2 -> PassM Exp2
forall {m :: * -> *} {a}.
Monad m =>
Env2 a -> Set Var -> Exp2 -> m Exp2
fn Prog2
p
 where
  fn :: Env2 a -> Set Var -> Exp2 -> m Exp2
fn Env2{TyEnv a
vEnv :: TyEnv a
vEnv :: forall a. Env2 a -> TyEnv a
vEnv,TyEnv (ArrowTy a)
fEnv :: TyEnv (ArrowTy a)
fEnv :: forall a. Env2 a -> TyEnv (ArrowTy a)
fEnv} Set Var
boundlocs Exp2
ex = Exp2 -> m Exp2
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Set Var -> Exp2 -> Int -> Exp2
forall {t}. (Eq t, Num t) => Set Var -> Exp2 -> t -> Exp2
goFix (TyEnv a -> Set Var
forall k a. Map k a -> Set k
Map.keysSet TyEnv a
vEnv Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` TyEnv (ArrowTy a) -> Set Var
forall k a. Map k a -> Set k
Map.keysSet TyEnv (ArrowTy a)
fEnv
                                                  Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Set Var
boundlocs
                                                  )
                                            Exp2
ex Int
bigNumber)
  goFix :: Set Var -> Exp2 -> t -> Exp2
goFix Set Var
_    Exp2
ex t
0 = String -> Exp2
forall a. HasCallStack => String -> a
error (String -> Exp2) -> String -> Exp2
forall a b. (a -> b) -> a -> b
$ String
"timeout in findWitness on " String -> ShowS
forall a. [a] -> [a] -> [a]
++ (Exp2 -> String
forall a. Show a => a -> String
show Exp2
ex)
  goFix Set Var
bound0 Exp2
ex0 t
n = let ex1 :: Exp2
ex1 = Set Var -> Map Var DelayedBind -> Exp2 -> Exp2
goE Set Var
bound0 Map Var DelayedBind
forall k a. Map k a
Map.empty Exp2
ex0
                           ex2 :: Exp2
ex2 = Set Var -> Map Var DelayedBind -> Exp2 -> Exp2
goE Set Var
bound0 Map Var DelayedBind
forall k a. Map k a
Map.empty Exp2
ex1
                       in if Exp2
ex1 Exp2 -> Exp2 -> Bool
forall a. Eq a => a -> a -> Bool
== Exp2
ex2 then Exp2
ex2
                          else Set Var -> Exp2 -> t -> Exp2
goFix Set Var
bound0 Exp2
ex2 (t
n t -> t -> t
forall a. Num a => a -> a -> a
- t
1)

  docase :: Set Var
-> Map Var DelayedBind
-> (String, [(Var, Var)], Exp2)
-> (String, [(Var, Var)], Exp2)
docase Set Var
bound Map Var DelayedBind
mp (String
k,[(Var, Var)]
vs,Exp2
e) =
    let ([Var]
vars,[Var]
locs) = [(Var, Var)] -> ([Var], [Var])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Var, Var)]
vs
        bound' :: Set Var
bound' = [Var] -> Set Var
forall a. Ord a => [a] -> Set a
Set.fromList ([Var]
vars [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ [Var]
locs) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Set Var
bound
    in (String
k,[(Var, Var)]
vs,Set Var -> Map Var DelayedBind -> Exp2 -> Exp2
goE Set Var
bound' Map Var DelayedBind
mp Exp2
e)

  goE :: Set.Set Var -> Map.Map Var DelayedBind -> Exp2 -> Exp2
  goE :: Set Var -> Map Var DelayedBind -> Exp2 -> Exp2
goE Set Var
bound Map Var DelayedBind
mp Exp2
ex =
    let go :: Map Var DelayedBind -> Exp2 -> Exp2
go      = Set Var -> Map Var DelayedBind -> Exp2 -> Exp2
goE Set Var
bound -- Shorthand.
        goClear :: Exp2 -> Exp2
goClear = Set Var -> Map Var DelayedBind -> Exp2 -> Exp2
goE (Set Var
bound Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Map Var DelayedBind -> Set Var
forall k a. Map k a -> Set k
Map.keysSet Map Var DelayedBind
mp) Map Var DelayedBind
forall k a. Map k a
Map.empty
        -- shorthand for applying (L p)
        handle' :: Exp2 -> Exp2
handle' Exp2
e = Set Var -> FunDefs Exp2 -> Map Var DelayedBind -> Exp2 -> Exp2
handle Set Var
bound FunDefs Exp2
fundefs Map Var DelayedBind
mp Exp2
e
    in
      case Exp2
ex of
        LetE (Var
v,[Var]
locs,Ty2
t, (TimeIt Exp2
e Ty2
ty Bool
b)) Exp2
bod ->
            Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ (Var, [Var], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
v,[Var]
locs,Ty2
t, Exp2 -> Ty2 -> Bool -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> Bool -> PreExp ext loc dec
TimeIt (Map Var DelayedBind -> Exp2 -> Exp2
go Map Var DelayedBind
forall k a. Map k a
Map.empty Exp2
e) Ty2
ty Bool
b)
                      (Set Var -> Map Var DelayedBind -> Exp2 -> Exp2
goE (Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
Set.insert Var
v (Set Var
bound Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Map Var DelayedBind -> Set Var
forall k a. Map k a -> Set k
Map.keysSet Map Var DelayedBind
mp)) Map Var DelayedBind
forall k a. Map k a
Map.empty Exp2
bod)

        Ext E2Ext Var Ty2
ext ->
          case E2Ext Var Ty2
ext of
            LetLocE Var
loc LocExp
locexp Exp2
bod ->
              let freelocs :: Set Var
freelocs = LocExp -> Set Var
forall a. FreeVars a => a -> Set Var
gFreeVars LocExp
locexp Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` Set Var
bound
                  chk :: Bool
chk = Set Var -> Bool
forall a. Set a -> Bool
Set.null Set Var
freelocs
              in if Bool
chk
                 -- dbgTraceIt (if loc == "loc_17052" then (sdoc (loc, locexp, freelocs, chk)) else "")
                 then E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> LocExp -> Exp2 -> E2Ext Var Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc LocExp
locexp (Exp2 -> E2Ext Var Ty2) -> Exp2 -> E2Ext Var Ty2
forall a b. (a -> b) -> a -> b
$ Set Var -> Map Var DelayedBind -> Exp2 -> Exp2
goE (Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
Set.insert Var
loc Set Var
bound) Map Var DelayedBind
mp Exp2
bod
                 else
                   case LocExp
locexp of
                     AfterVariableLE Var
v Var
loc2 Bool
b ->
                       (Map Var DelayedBind -> Exp2 -> Exp2
go (Var -> DelayedBind -> Map Var DelayedBind -> Map Var DelayedBind
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Var
loc ((Var, LocExp) -> DelayedBind
DelayLoc (Var
loc, (Var -> Var -> Bool -> LocExp
forall loc. Var -> loc -> Bool -> PreLocExp loc
AfterVariableLE Var
v Var
loc2 Bool
b))) Map Var DelayedBind
mp) Exp2
bod)
                     AfterConstantLE Int
i Var
loc2 ->
                       Map Var DelayedBind -> Exp2 -> Exp2
go (Var -> DelayedBind -> Map Var DelayedBind -> Map Var DelayedBind
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Var
loc ((Var, LocExp) -> DelayedBind
DelayLoc (Var
loc, (Int -> Var -> LocExp
forall loc. Int -> loc -> PreLocExp loc
AfterConstantLE Int
i Var
loc2))) Map Var DelayedBind
mp) Exp2
bod
                     LocExp
_ -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> LocExp -> Exp2 -> E2Ext Var Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc LocExp
locexp (Exp2 -> E2Ext Var Ty2) -> Exp2 -> E2Ext Var Ty2
forall a b. (a -> b) -> a -> b
$ Set Var -> Map Var DelayedBind -> Exp2 -> Exp2
goE (Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
Set.insert Var
loc Set Var
bound) Map Var DelayedBind
mp Exp2
bod
            LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Region -> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var Ty2
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
r RegionSize
sz Maybe RegionType
ty (Exp2 -> E2Ext Var Ty2) -> Exp2 -> E2Ext Var Ty2
forall a b. (a -> b) -> a -> b
$ Map Var DelayedBind -> Exp2 -> Exp2
go Map Var DelayedBind
mp Exp2
bod
            LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Region -> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var Ty2
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty (Exp2 -> E2Ext Var Ty2) -> Exp2 -> E2Ext Var Ty2
forall a b. (a -> b) -> a -> b
$ Map Var DelayedBind -> Exp2 -> Exp2
go Map Var DelayedBind
mp Exp2
bod
            E2Ext Var Ty2
_ -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Exp2
ex

        LetE (Var
v,[Var]
locs,Ty2
t,Exp2
rhs) Exp2
bod ->
          let rhs' :: Exp2
rhs' = Map Var DelayedBind -> Exp2 -> Exp2
go Map Var DelayedBind
forall k a. Map k a
Map.empty Exp2
rhs -- recur on rhs, but flatten makes these pretty boring.
              freelocs :: Set Var
freelocs = Exp2 -> Set Var
ex_freeVars Exp2
rhs' Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` Set Var
bound
              chk :: Bool
chk = Set Var -> Bool
forall a. Set a -> Bool
Set.null Set Var
freelocs
          in if Bool
chk
             then (Var, [Var], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
v,[Var]
locs,Ty2
t,Exp2
rhs') (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Set Var -> Map Var DelayedBind -> Exp2 -> Exp2
goE (Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
Set.insert Var
v Set Var
bound) Map Var DelayedBind
mp (Set Var -> FunDefs Exp2 -> Map Var DelayedBind -> Exp2 -> Exp2
handle (Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
Set.insert Var
v Set Var
bound) FunDefs Exp2
fundefs Map Var DelayedBind
mp Exp2
bod)
             else Map Var DelayedBind -> Exp2 -> Exp2
go (Var -> DelayedBind -> Map Var DelayedBind -> Map Var DelayedBind
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Var
v ((Var, [Var], Ty2, Exp2) -> DelayedBind
DelayVar (Var
v,[Var]
locs,Ty2
t,Exp2
rhs')) Map Var DelayedBind
mp) Exp2
bod

        VarE Var
v         -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> Exp2
forall (ext :: * -> * -> *) loc dec. Var -> PreExp ext loc dec
VarE Var
v
        LitE Int
n         -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Int -> Exp2
forall (ext :: * -> * -> *) loc dec. Int -> PreExp ext loc dec
LitE Int
n
        CharE Char
c        -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Char -> Exp2
forall (ext :: * -> * -> *) loc dec. Char -> PreExp ext loc dec
CharE Char
c
        FloatE Double
n       -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Double -> Exp2
forall (ext :: * -> * -> *) loc dec. Double -> PreExp ext loc dec
FloatE Double
n
        LitSymE Var
v      -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> Exp2
forall (ext :: * -> * -> *) loc dec. Var -> PreExp ext loc dec
LitSymE Var
v
        AppE Var
v [Var]
locs [Exp2]
ls -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> [Var] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
v [Var]
locs ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
goClear [Exp2]
ls)

        SpawnE Var
v [Var]
locs [Exp2]
ls -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> [Var] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
SpawnE Var
v [Var]
locs ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
goClear [Exp2]
ls)
        Exp2
SyncE            -> Exp2
forall (ext :: * -> * -> *) loc dec. PreExp ext loc dec
SyncE

        PrimAppE Prim Ty2
pr [Exp2]
ls -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Prim Ty2 -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim Ty2
pr ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
goClear [Exp2]
ls)
        ProjE Int
i Exp2
e      -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Int -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Int -> PreExp ext loc dec -> PreExp ext loc dec
ProjE Int
i (Exp2 -> Exp2
goClear Exp2
e)

        -- It's ok that we don't go deeper on scrutinees, subexpressions
        -- of tuples, and so on, because flatten should have done
        -- let-lifting out of these contexts:
        MkProdE [Exp2]
ls     -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
[PreExp ext loc dec] -> PreExp ext loc dec
MkProdE ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
goClear [Exp2]
ls)
        DataConE Var
loc String
k [Exp2]
ls -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> String -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
loc -> String -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE Var
loc String
k ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
goClear [Exp2]
ls)
        TimeIt Exp2
e Ty2
t Bool
b   -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> 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 -> Exp2
goClear Exp2
e) Ty2
t Bool
b -- prevent pushing work into timeit

        -- FIXME: give CaseE a treatment like IfE:
        CaseE Exp2
scrt [(String, [(Var, Var)], Exp2)]
ls  -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Exp2 -> [(String, [(Var, Var)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(String, [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE Exp2
scrt (((String, [(Var, Var)], Exp2) -> (String, [(Var, Var)], Exp2))
-> [(String, [(Var, Var)], Exp2)] -> [(String, [(Var, Var)], Exp2)]
forall a b. (a -> b) -> [a] -> [b]
map (Set Var
-> Map Var DelayedBind
-> (String, [(Var, Var)], Exp2)
-> (String, [(Var, Var)], Exp2)
docase Set Var
bound Map Var DelayedBind
mp) [(String, [(Var, Var)], Exp2)]
ls)

        IfE Exp2
a Exp2
b Exp2
c ->
            -- If we have "succeeded" in accumulating all the bindings
            -- we need, then we can discharge their topological sort
            -- here, without duplicating bindings:
            if Set Var -> Map Var DelayedBind -> Bool
closed Set Var
bound Map Var DelayedBind
mp
            then Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> 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 -> Exp2
goClear Exp2
b) (Exp2 -> Exp2
goClear Exp2
c)
            else Exp2 -> Exp2 -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> PreExp ext loc dec -> PreExp ext loc dec -> PreExp ext loc dec
IfE (Map Var DelayedBind -> Exp2 -> Exp2
go Map Var DelayedBind
mp Exp2
a) -- Otherwise we duplicate...
                     (Map Var DelayedBind -> Exp2 -> Exp2
go Map Var DelayedBind
mp Exp2
b)
                     (Map Var DelayedBind -> Exp2 -> Exp2
go Map Var DelayedBind
mp Exp2
c)
        -- Like MkProdE
        MapE  (Var
v,Ty2
t,Exp2
rhs) Exp2
bod -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ (Var, Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
MapE (Var
v,Ty2
t,Exp2
rhs) (Exp2 -> Exp2
goClear Exp2
bod)
        FoldE (Var
v1,Ty2
t1,Exp2
r1) (Var
v2,Ty2
t2,Exp2
r2) Exp2
bod -> Exp2 -> Exp2
handle' (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ (Var, Ty2, Exp2) -> (Var, Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, dec, PreExp ext loc dec)
-> (Var, dec, PreExp ext loc dec)
-> PreExp ext loc dec
-> PreExp ext loc dec
FoldE (Var
v1,Ty2
t1,Exp2
r1) (Var
v2,Ty2
t2,Exp2
r2) (Exp2 -> Exp2
goClear Exp2
bod)
        WithArenaE{} -> String -> Exp2
forall a. HasCallStack => String -> a
error String
"findWitnesses: WithArenaE not handled."


-- TODO: this needs to preserve any bindings that have TimeIt forms (hasTimeIt).
-- OR we can only match a certain pattern like (Let (_,_,TimeIt _ _) _)
handle :: Set.Set Var -> FunDefs2 -> Map.Map Var DelayedBind -> Exp2 -> Exp2
handle :: Set Var -> FunDefs Exp2 -> Map Var DelayedBind -> Exp2 -> Exp2
handle Set Var
bound FunDefs Exp2
fundefs Map Var DelayedBind
mp Exp2
expr =
    Int -> String -> Exp2 -> Exp2
forall a. Int -> String -> a -> a
dbgTrace Int
6 (String
" [findWitnesses] building lets using vars "String -> ShowS
forall a. [a] -> [a] -> [a]
++[Var] -> String
forall a. Show a => a -> String
show [Var]
vsString -> ShowS
forall a. [a] -> [a] -> [a]
++String
" for expr: "String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> ShowS
forall a. Int -> [a] -> [a]
take Int
80 (Exp2 -> String
forall a. Show a => a -> String
show Exp2
expr)) (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$
    -- dbgTraceIt (if vars /= [] then "binding: " ++ sdoc vars else "")
    Map Var DelayedBind -> [Var] -> Exp2 -> Exp2
buildLets Map Var DelayedBind
mp [Var]
vars Exp2
expr
    where freeInBind :: Var -> [Var]
freeInBind Var
v = case Var -> Map Var DelayedBind -> Maybe DelayedBind
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Var -> Var
view Var
v) Map Var DelayedBind
mp of
                           Maybe DelayedBind
Nothing -> []
                           Just (DelayVar (Var
_v,[Var]
_locs,Ty2
_t,Exp2
e)) -> Set Var -> [Var]
forall a. Set a -> [a]
Set.toList (Set Var -> [Var]) -> Set Var -> [Var]
forall a b. (a -> b) -> a -> b
$ (Exp2 -> Set Var
ex_freeVars Exp2
e) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` (FunDefs Exp2 -> Set Var
forall k a. Map k a -> Set k
Map.keysSet FunDefs Exp2
fundefs)
                           Just (DelayLoc (Var
_loc, LocExp
locexp)) -> Set Var -> [Var]
forall a. Set a -> [a]
Set.toList (Set Var -> [Var]) -> Set Var -> [Var]
forall a b. (a -> b) -> a -> b
$ (LocExp -> Set Var
forall a. FreeVars a => a -> Set Var
gFreeVars LocExp
locexp) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` (FunDefs Exp2 -> Set Var
forall k a. Map k a -> Set k
Map.keysSet FunDefs Exp2
fundefs)

          (Graph
g,Int -> (Var, Var, [Var])
vf,Var -> Maybe Int
_) = [(Var, Var, [Var])]
-> (Graph, Int -> (Var, Var, [Var]), Var -> Maybe Int)
forall key node.
Ord key =>
[(node, key, [key])]
-> (Graph, Int -> (node, key, [key]), key -> Maybe Int)
graphFromEdges ([(Var, Var, [Var])]
 -> (Graph, Int -> (Var, Var, [Var]), Var -> Maybe Int))
-> [(Var, Var, [Var])]
-> (Graph, Int -> (Var, Var, [Var]), Var -> Maybe Int)
forall a b. (a -> b) -> a -> b
$ [Var] -> [Var] -> [[Var]] -> [(Var, Var, [Var])]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Var]
vs [Var]
vs ([[Var]] -> [(Var, Var, [Var])]) -> [[Var]] -> [(Var, Var, [Var])]
forall a b. (a -> b) -> a -> b
$ (Var -> [Var]) -> [Var] -> [[Var]]
forall a b. (a -> b) -> [a] -> [b]
map Var -> [Var]
freeInBind [Var]
vs
          vars :: [Var]
vars = [Var] -> [Var]
forall a. [a] -> [a]
reverse ([Var] -> [Var]) -> [Var] -> [Var]
forall a b. (a -> b) -> a -> b
$ ((Var, Var, [Var]) -> Var) -> [(Var, Var, [Var])] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (\(Var
x,Var
_,[Var]
_) -> Var
x) ([(Var, Var, [Var])] -> [Var]) -> [(Var, Var, [Var])] -> [Var]
forall a b. (a -> b) -> a -> b
$ (Int -> (Var, Var, [Var])) -> [Int] -> [(Var, Var, [Var])]
forall a b. (a -> b) -> [a] -> [b]
map Int -> (Var, Var, [Var])
vf ([Int] -> [(Var, Var, [Var])]) -> [Int] -> [(Var, Var, [Var])]
forall a b. (a -> b) -> a -> b
$ Graph -> [Int]
topSort Graph
g
          vs :: [Var]
vs = Map Var DelayedBind -> [Var]
forall k a. Map k a -> [k]
Map.keys (Map Var DelayedBind -> [Var]) -> Map Var DelayedBind -> [Var]
forall a b. (a -> b) -> a -> b
$ (Var -> DelayedBind -> Bool)
-> Map Var DelayedBind -> Map Var DelayedBind
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (\Var
k DelayedBind
_v -> Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Var
k Set Var
bound) Map Var DelayedBind
mp




-- withWitnesses :: [LocVar] -> [LocVar]
-- withWitnesses ls = concatMap f ls
--     where f v = if isWitnessVar v
--                 then [v]
--                 else [v,toWitnessVar v] -- maybe?

-- From the point of view of this pass, we "see through" witness markerS:
view :: Var -> Var
view :: Var -> Var
view Var
v = Var
v  -- RRN: actually, coming up with a good policy here is problematic.

-- view v | isWitnessVar v = let Just v' = fromWitnessVar v in v'
--        | otherwise      = v


buildLets :: Map.Map Var DelayedBind -> [Var] -> Exp2-> Exp2
buildLets :: Map Var DelayedBind -> [Var] -> Exp2 -> Exp2
buildLets Map Var DelayedBind
_mp [] Exp2
bod = Exp2
bod
buildLets Map Var DelayedBind
mp (Var
v:[Var]
vs) Exp2
bod =
    case Var -> Map Var DelayedBind -> Maybe DelayedBind
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Var -> Var
view Var
v) Map Var DelayedBind
mp of
      Maybe DelayedBind
Nothing -> Map Var DelayedBind -> [Var] -> Exp2 -> Exp2
buildLets Map Var DelayedBind
mp [Var]
vs Exp2
bod
      Just (DelayVar (Var, [Var], Ty2, Exp2)
bnd) -> (Var, [Var], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var, [Var], Ty2, Exp2)
bnd (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Map Var DelayedBind -> [Var] -> Exp2 -> Exp2
buildLets Map Var DelayedBind
mp [Var]
vs Exp2
bod
      Just (DelayLoc (Var
loc, LocExp
bnd)) -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> LocExp -> Exp2 -> E2Ext Var Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc LocExp
bnd (Map Var DelayedBind -> [Var] -> Exp2 -> Exp2
buildLets Map Var DelayedBind
mp [Var]
vs Exp2
bod)


-- | Are all the free variables currently bound (transitively) in the
-- environment?
closed :: Set.Set Var -> Map.Map Var DelayedBind -> Bool
closed :: Set Var -> Map Var DelayedBind -> Bool
closed Set Var
bound Map Var DelayedBind
mp = Set Var -> Bool
forall a. Set a -> Bool
Set.null (Set Var
allBound Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` Set Var
allUsed)
  where
   allBound :: Set Var
allBound = Set Var
bound Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Map Var DelayedBind -> Set Var
forall k a. Map k a -> Set k
Map.keysSet Map Var DelayedBind
mp
   -- allUsed = Set.unions [ gFreeVars rhs | (_,_,_,rhs) <- Map.elems mp ]
   allUsed :: Set Var
allUsed = [Set Var] -> Set Var
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ([Set Var] -> Set Var) -> [Set Var] -> Set Var
forall a b. (a -> b) -> a -> b
$ (DelayedBind -> Set Var) -> [DelayedBind] -> [Set Var]
forall a b. (a -> b) -> [a] -> [b]
map (\DelayedBind
db -> case DelayedBind
db of
                                  DelayVar (Var
_,[Var]
_,Ty2
_,Exp2
rhs) -> Exp2 -> Set Var
ex_freeVars Exp2
rhs
                                  DelayLoc (Var
_,LocExp
locexp)  -> LocExp -> Set Var
forall a. FreeVars a => a -> Set Var
gFreeVars LocExp
locexp)
                          (Map Var DelayedBind -> [DelayedBind]
forall k a. Map k a -> [a]
Map.elems Map Var DelayedBind
mp)

mapMExprs :: Monad m => (Env2 Ty2 -> Set.Set LocVar -> Exp2 -> m Exp2) -> Prog2 -> m Prog2
mapMExprs :: forall (m :: * -> *).
Monad m =>
(Env2 Ty2 -> Set Var -> Exp2 -> m Exp2) -> Prog2 -> m Prog2
mapMExprs Env2 Ty2 -> Set Var -> Exp2 -> m Exp2
fn (Prog DDefs (TyOf Exp2)
ddfs FunDefs Exp2
fundefs Maybe (Exp2, TyOf Exp2)
mainExp) =
  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)
ddfs (FunDefs Exp2 -> Maybe (Exp2, Ty2) -> Prog2)
-> m (FunDefs Exp2) -> m (Maybe (Exp2, Ty2) -> Prog2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
    ((FunDef Exp2 -> m (FunDef Exp2))
-> FunDefs Exp2 -> m (FunDefs 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) -> Map Var a -> m (Map Var b)
mapM (\f :: FunDef Exp2
f@FunDef{[Var]
funArgs :: [Var]
funArgs :: forall ex. FunDef ex -> [Var]
funArgs,ArrowTy (TyOf Exp2)
funTy :: ArrowTy (TyOf Exp2)
funTy :: forall ex. FunDef ex -> ArrowTy (TyOf ex)
funTy,Exp2
funBody :: Exp2
funBody :: forall ex. FunDef ex -> ex
funBody} ->
              let env :: Env2 Ty2
env = TyEnv Ty2 -> TyEnv (ArrowTy Ty2) -> Env2 Ty2
forall a. TyEnv a -> TyEnv (ArrowTy a) -> Env2 a
Env2 ([(Var, Ty2)] -> TyEnv Ty2
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Var, Ty2)] -> TyEnv Ty2) -> [(Var, Ty2)] -> TyEnv Ty2
forall a b. (a -> b) -> a -> b
$ [Var] -> [Ty2] -> [(Var, Ty2)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Var]
funArgs (ArrowTy Ty2 -> [Ty2]
forall ty. FunctionTy ty => ArrowTy ty -> [ty]
inTys ArrowTy (TyOf Exp2)
ArrowTy Ty2
funTy)) TyEnv (ArrowTy Ty2)
Map Var (ArrowTy2 Ty2)
funEnv
                  boundlocs :: Set Var
boundlocs = [Var] -> Set Var
forall a. Ord a => [a] -> Set a
Set.fromList (ArrowTy2 Ty2 -> [Var]
forall ty2. ArrowTy2 ty2 -> [Var]
allLocVars ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`Set.union`
                              [Var] -> Set Var
forall a. Ord a => [a] -> Set a
Set.fromList [Var]
funArgs
              in do
                Exp2
bod' <- Env2 Ty2 -> Set Var -> Exp2 -> m Exp2
fn Env2 Ty2
env Set Var
boundlocs Exp2
funBody
                FunDef Exp2 -> m (FunDef Exp2)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (FunDef Exp2 -> m (FunDef Exp2)) -> FunDef Exp2 -> m (FunDef Exp2)
forall a b. (a -> b) -> a -> b
$ FunDef Exp2
f { funBody :: Exp2
funBody =  Exp2
bod' })
     FunDefs Exp2
fundefs)
    m (Maybe (Exp2, Ty2) -> Prog2) -> m (Maybe (Exp2, Ty2)) -> m Prog2
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
    (((Exp2, Ty2) -> m (Exp2, Ty2))
-> Maybe (Exp2, Ty2) -> m (Maybe (Exp2, Ty2))
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) -> Maybe a -> m (Maybe b)
mapM (\ (Exp2
e,Ty2
t) -> (,Ty2
t) (Exp2 -> (Exp2, Ty2)) -> m Exp2 -> m (Exp2, Ty2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env2 Ty2 -> Set Var -> Exp2 -> m Exp2
fn (TyEnv Ty2 -> TyEnv (ArrowTy Ty2) -> Env2 Ty2
forall a. TyEnv a -> TyEnv (ArrowTy a) -> Env2 a
Env2 TyEnv Ty2
forall k a. Map k a
Map.empty TyEnv (ArrowTy Ty2)
Map Var (ArrowTy2 Ty2)
funEnv) Set Var
forall a. Set a
Set.empty Exp2
e) Maybe (Exp2, TyOf Exp2)
Maybe (Exp2, Ty2)
mainExp)
  where funEnv :: Map Var (ArrowTy2 Ty2)
funEnv = (FunDef Exp2 -> ArrowTy2 Ty2)
-> FunDefs Exp2 -> Map Var (ArrowTy2 Ty2)
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map FunDef Exp2 -> ArrowTy (TyOf Exp2)
FunDef Exp2 -> ArrowTy2 Ty2
forall ex. FunDef ex -> ArrowTy (TyOf ex)
funTy FunDefs Exp2
fundefs

ex_freeVars :: Exp2 -> Set.Set Var
ex_freeVars :: Exp2 -> Set Var
ex_freeVars = Exp2 -> Set Var
allFreeVars
-- ex_freeVars = gFreeVars