module Gibbon.Passes.DirectL3
    (directL3) where

import qualified Data.List as L
import qualified Data.Map as M

import           Gibbon.Common
import           Gibbon.L1.Syntax
import           Gibbon.L3.Syntax
import Data.Bifunctor


-- | Directly convert the source program to L3. Used in the pointer mode
--
directL3 :: Prog1 -> PassM Prog3
directL3 :: Prog1 -> PassM Prog3
directL3 prg :: Prog1
prg@(Prog DDefs (TyOf Exp1)
ddfs FunDefs Exp1
fndefs Maybe (Exp1, TyOf Exp1)
mnExp) = do
    let mnExp' :: Maybe (Exp3, Ty1)
mnExp' = case Maybe (Exp1, TyOf Exp1)
mnExp of
                   Maybe (Exp1, TyOf Exp1)
Nothing -> Maybe (Exp3, Ty1)
forall a. Maybe a
Nothing
                   Just (Exp1
ex,TyOf Exp1
ty) -> (Exp3, Ty1) -> Maybe (Exp3, Ty1)
forall a. a -> Maybe a
Just (Env2 Ty1 -> Exp1 -> Exp3
go Env2 (TyOf Exp1)
Env2 Ty1
init_fun_env Exp1
ex, TyOf Exp1
Ty1
ty)
        fndefs' :: Map Var FunDef3
fndefs' = (FunDef1 -> FunDef3) -> FunDefs Exp1 -> Map Var FunDef3
forall a b k. (a -> b) -> Map k a -> Map k b
M.map FunDef1 -> FunDef3
fd FunDefs Exp1
fndefs
        ddfs' :: Map Var DDef3
ddfs' = (DDef3 -> DDef3) -> Map Var DDef3 -> Map Var DDef3
forall a b. (a -> b) -> Map Var a -> Map Var b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap DDef3 -> DDef3
goDDef DDefs (TyOf Exp1)
Map Var DDef3
ddfs
    Prog3 -> PassM Prog3
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (DDefs (TyOf Exp3)
-> Map Var FunDef3 -> Maybe (Exp3, TyOf Exp3) -> Prog3
forall ex.
DDefs (TyOf ex) -> FunDefs ex -> Maybe (ex, TyOf ex) -> Prog ex
Prog DDefs (TyOf Exp3)
Map Var DDef3
ddfs' Map Var FunDef3
fndefs' Maybe (Exp3, TyOf Exp3)
Maybe (Exp3, Ty1)
mnExp')
  where
    init_fun_env :: Env2 (TyOf Exp1)
init_fun_env = Prog1 -> Env2 (TyOf Exp1)
forall a. Prog a -> Env2 (TyOf a)
progToEnv Prog1
prg

    goDDef :: DDef1 -> DDef3 
    goDDef :: DDef3 -> DDef3
goDDef ddf :: DDef3
ddf@DDef{[(DataCon, [(IsBoxed, Ty1)])]
dataCons :: [(DataCon, [(IsBoxed, Ty1)])]
dataCons :: forall a. DDef a -> [(DataCon, [(IsBoxed, a)])]
dataCons} = 
      DDef3
ddf{dataCons :: [(DataCon, [(IsBoxed, Ty1)])]
dataCons  = ((DataCon, [(IsBoxed, Ty1)]) -> (DataCon, [(IsBoxed, Ty1)]))
-> [(DataCon, [(IsBoxed, Ty1)])] -> [(DataCon, [(IsBoxed, Ty1)])]
forall a b. (a -> b) -> [a] -> [b]
map(([(IsBoxed, Ty1)] -> [(IsBoxed, Ty1)])
-> (DataCon, [(IsBoxed, Ty1)]) -> (DataCon, [(IsBoxed, Ty1)])
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (((IsBoxed, Ty1) -> (IsBoxed, Ty1))
-> [(IsBoxed, Ty1)] -> [(IsBoxed, Ty1)]
forall a b. (a -> b) -> [a] -> [b]
map((Ty1 -> Ty1) -> (IsBoxed, Ty1) -> (IsBoxed, Ty1)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Ty1 -> Ty1
goTy))) [(DataCon, [(IsBoxed, Ty1)])]
dataCons}

    fd :: FunDef1 -> FunDef3
    fd :: FunDef1 -> FunDef3
fd FunDef{Var
funName :: Var
funName :: forall ex. FunDef ex -> Var
funName,[Var]
funArgs :: [Var]
funArgs :: forall ex. FunDef ex -> [Var]
funArgs,ArrowTy (TyOf Exp1)
funTy :: ArrowTy (TyOf Exp1)
funTy :: forall ex. FunDef ex -> ArrowTy (TyOf ex)
funTy,Exp1
funBody :: Exp1
funBody :: forall ex. FunDef ex -> ex
funBody,FunMeta
funMeta :: FunMeta
funMeta :: forall ex. FunDef ex -> FunMeta
funMeta} =
        let env2 :: Env2 Ty1
env2 = Map Var Ty1 -> Env2 Ty1 -> Env2 Ty1
forall a. Map Var a -> Env2 a -> Env2 a
extendsVEnv ([(Var, Ty1)] -> Map Var Ty1
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Var, Ty1)] -> Map Var Ty1) -> [(Var, Ty1)] -> Map Var Ty1
forall a b. (a -> b) -> a -> b
$ [Var] -> [Ty1] -> [(Var, Ty1)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Var]
funArgs (([Ty1], Ty1) -> [Ty1]
forall a b. (a, b) -> a
fst ([Ty1], Ty1)
ArrowTy (TyOf Exp1)
funTy)) Env2 (TyOf Exp1)
Env2 Ty1
init_fun_env in
        FunDef { funName :: Var
funName = Var
funName
               , funTy :: ArrowTy (TyOf Exp3)
funTy   = ((Ty1 -> Ty1) -> [Ty1] -> [Ty1]
forall a b. (a -> b) -> [a] -> [b]
map Ty1 -> Ty1
goTy ([Ty1] -> [Ty1]) -> [Ty1] -> [Ty1]
forall a b. (a -> b) -> a -> b
$ ([Ty1], Ty1) -> [Ty1]
forall a b. (a, b) -> a
fst ([Ty1], Ty1)
ArrowTy (TyOf Exp1)
funTy, Ty1 -> Ty1
goTy (Ty1 -> Ty1) -> Ty1 -> Ty1
forall a b. (a -> b) -> a -> b
$ ([Ty1], Ty1) -> Ty1
forall a b. (a, b) -> b
snd ([Ty1], Ty1)
ArrowTy (TyOf Exp1)
funTy)
               , funArgs :: [Var]
funArgs = [Var]
funArgs
               , funBody :: Exp3
funBody = Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
funBody
               , funMeta :: FunMeta
funMeta = FunMeta
funMeta
               }

    go :: Env2 Ty1 -> Exp1 -> Exp3
    go :: Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
ex =
      case Exp1
ex of
        VarE Var
v    -> Var -> Exp3
forall (ext :: * -> * -> *) loc dec. Var -> PreExp ext loc dec
VarE Var
v
        LitE Int
n    -> Int -> Exp3
forall (ext :: * -> * -> *) loc dec. Int -> PreExp ext loc dec
LitE Int
n
        CharE Char
c   -> Char -> Exp3
forall (ext :: * -> * -> *) loc dec. Char -> PreExp ext loc dec
CharE Char
c
        FloatE Double
n  -> Double -> Exp3
forall (ext :: * -> * -> *) loc dec. Double -> PreExp ext loc dec
FloatE Double
n
        LitSymE Var
v -> Var -> Exp3
forall (ext :: * -> * -> *) loc dec. Var -> PreExp ext loc dec
LitSymE Var
v
        AppE Var
v [()]
locs [Exp1]
ls   -> Var -> [()] -> [Exp3] -> Exp3
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
v [()]
locs ([Exp3] -> Exp3) -> [Exp3] -> Exp3
forall a b. (a -> b) -> a -> b
$ (Exp1 -> Exp3) -> [Exp1] -> [Exp3]
forall a b. (a -> b) -> [a] -> [b]
map (Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2) [Exp1]
ls
        PrimAppE Prim Ty1
pr [Exp1]
args -> Prim Ty1 -> [Exp3] -> Exp3
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim Ty1
pr ([Exp3] -> Exp3) -> [Exp3] -> Exp3
forall a b. (a -> b) -> a -> b
$ (Exp1 -> Exp3) -> [Exp1] -> [Exp3]
forall a b. (a -> b) -> [a] -> [b]
L.map (Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2) [Exp1]
args
        LetE (Var
v,[()]
locs,Ty1
ty,ProjE Int
i Exp1
arg) Exp1
bod ->
            (Var, [()], Ty1, Exp3) -> Exp3 -> Exp3
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
v, [()]
locs, Ty1 -> Ty1
goTy Ty1
ty, Int -> Exp3 -> Exp3
forall (ext :: * -> * -> *) loc dec.
Int -> PreExp ext loc dec -> PreExp ext loc dec
ProjE Int
i (Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
arg)) (Exp3 -> Exp3) -> Exp3 -> Exp3
forall a b. (a -> b) -> a -> b
$
            Env2 Ty1 -> Exp1 -> Exp3
go (Var -> Ty1 -> Env2 Ty1 -> Env2 Ty1
forall a. Var -> a -> Env2 a -> Env2 a
extendVEnv Var
v Ty1
ty Env2 Ty1
env2) Exp1
bod
        LetE (Var
v,[()]
locs,Ty1
ty,Exp1
rhs) Exp1
bod -> (Var, [()], Ty1, Exp3) -> Exp3 -> Exp3
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
v, [()]
locs, Ty1 -> Ty1
goTy Ty1
ty, Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
rhs) (Exp3 -> Exp3) -> Exp3 -> Exp3
forall a b. (a -> b) -> a -> b
$
                                      Env2 Ty1 -> Exp1 -> Exp3
go (Var -> Ty1 -> Env2 Ty1 -> Env2 Ty1
forall a. Var -> a -> Env2 a -> Env2 a
extendVEnv Var
v Ty1
ty Env2 Ty1
env2) Exp1
bod
        IfE Exp1
a Exp1
b Exp1
c   -> Exp3 -> Exp3 -> Exp3 -> Exp3
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> PreExp ext loc dec -> PreExp ext loc dec -> PreExp ext loc dec
IfE (Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
a) (Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
b) (Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
c)
        MkProdE [Exp1]
ls  -> [Exp3] -> Exp3
forall (ext :: * -> * -> *) loc dec.
[PreExp ext loc dec] -> PreExp ext loc dec
MkProdE ([Exp3] -> Exp3) -> [Exp3] -> Exp3
forall a b. (a -> b) -> a -> b
$ (Exp1 -> Exp3) -> [Exp1] -> [Exp3]
forall a b. (a -> b) -> [a] -> [b]
L.map (Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2) [Exp1]
ls
        ProjE Int
i Exp1
arg ->
            let ty :: TyOf Exp1
ty = DDefs (TyOf Exp1) -> Env2 (TyOf Exp1) -> Exp1 -> TyOf Exp1
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf Exp1)
ddfs Env2 (TyOf Exp1)
Env2 Ty1
env2 Exp1
ex
                rhs' :: Exp3
rhs' = Int -> Exp3 -> Exp3
forall (ext :: * -> * -> *) loc dec.
Int -> PreExp ext loc dec -> PreExp ext loc dec
ProjE Int
i (Exp3 -> Exp3) -> Exp3 -> Exp3
forall a b. (a -> b) -> a -> b
$ Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
arg
            in (Var, [()], Ty1, Exp3) -> Exp3 -> Exp3
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
"prjtmp", [], TyOf Exp1
Ty1
ty, Exp3
rhs') (Var -> Exp3
forall (ext :: * -> * -> *) loc dec. Var -> PreExp ext loc dec
VarE Var
"prjtmp")
        CaseE Exp1
scrt [(DataCon, [(Var, ())], Exp1)]
ls -> Exp3 -> [(DataCon, [(Var, ())], Exp3)] -> Exp3
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(DataCon, [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE (Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
scrt) ([(DataCon, [(Var, ())], Exp3)] -> Exp3)
-> [(DataCon, [(Var, ())], Exp3)] -> Exp3
forall a b. (a -> b) -> a -> b
$
                           ((DataCon, [(Var, ())], Exp1) -> (DataCon, [(Var, ())], Exp3))
-> [(DataCon, [(Var, ())], Exp1)] -> [(DataCon, [(Var, ())], Exp3)]
forall a b. (a -> b) -> [a] -> [b]
L.map (\(DataCon
dcon,[(Var, ())]
vs,Exp1
rhs) -> (DataCon
dcon,[(Var, ())]
vs,Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
rhs)) [(DataCon, [(Var, ())], Exp1)]
ls
        DataConE ()
loc DataCon
dcon [Exp1]
args -> () -> DataCon -> [Exp3] -> Exp3
forall (ext :: * -> * -> *) loc dec.
loc -> DataCon -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE ()
loc DataCon
dcon ([Exp3] -> Exp3) -> [Exp3] -> Exp3
forall a b. (a -> b) -> a -> b
$ (Exp1 -> Exp3) -> [Exp1] -> [Exp3]
forall a b. (a -> b) -> [a] -> [b]
L.map (Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2) [Exp1]
args
        TimeIt Exp1
arg Ty1
ty IsBoxed
b -> Exp3 -> Ty1 -> IsBoxed -> Exp3
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> IsBoxed -> PreExp ext loc dec
TimeIt (Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
arg) Ty1
ty IsBoxed
b
        WithArenaE Var
a Exp1
e  -> Var -> Exp3 -> Exp3
forall (ext :: * -> * -> *) loc dec.
Var -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE Var
a (Exp3 -> Exp3) -> Exp3 -> Exp3
forall a b. (a -> b) -> a -> b
$ Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
e
        SpawnE Var
fn [()]
locs [Exp1]
ls -> Var -> [()] -> [Exp3] -> Exp3
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
SpawnE Var
fn [()]
locs ([Exp3] -> Exp3) -> [Exp3] -> Exp3
forall a b. (a -> b) -> a -> b
$ (Exp1 -> Exp3) -> [Exp1] -> [Exp3]
forall a b. (a -> b) -> [a] -> [b]
map (Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2) [Exp1]
ls
        Exp1
SyncE -> Exp3
forall (ext :: * -> * -> *) loc dec. PreExp ext loc dec
SyncE
        Ext (BenchE Var
fn [()]
_locs [Exp1]
args IsBoxed
b) ->
          let fn_ty :: ArrowTy Ty1
fn_ty  = Var -> Env2 Ty1 -> ArrowTy Ty1
forall a. Out (ArrowTy a) => Var -> Env2 a -> ArrowTy a
lookupFEnv Var
fn Env2 Ty1
env2
              ret_ty :: Ty1
ret_ty = ([Ty1], Ty1) -> Ty1
forall a b. (a, b) -> b
snd ([Ty1], Ty1)
ArrowTy Ty1
fn_ty
              ex' :: Exp1
ex'    = Exp1 -> Ty1 -> IsBoxed -> Exp1
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> IsBoxed -> PreExp ext loc dec
TimeIt (Var -> [()] -> [Exp1] -> Exp1
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
fn [] [Exp1]
args) Ty1
ret_ty IsBoxed
b
          in Env2 Ty1 -> Exp1 -> Exp3
go Env2 Ty1
env2 Exp1
ex'
        Ext (AddFixed{}) -> DataCon -> Exp3
forall a. HasCallStack => DataCon -> a
error DataCon
"directL3: AddFixed not handled."
        Ext (StartOfPkdCursor{}) -> DataCon -> Exp3
forall a. HasCallStack => DataCon -> a
error DataCon
"directL3: StartOfPkdCursor not handled."
        MapE{}  -> DataCon -> Exp3
forall a. HasCallStack => DataCon -> a
error DataCon
"directL3: todo MapE"
        FoldE{} -> DataCon -> Exp3
forall a. HasCallStack => DataCon -> a
error DataCon
"directL3: todo FoldE"

    goTy :: Ty1 -> Ty3
    goTy :: Ty1 -> Ty1
goTy Ty1
ty =
      case Ty1
ty of
        Ty1
IntTy  -> Ty1
forall loc. UrTy loc
IntTy
        Ty1
CharTy -> Ty1
forall loc. UrTy loc
CharTy
        Ty1
FloatTy-> Ty1
forall loc. UrTy loc
FloatTy
        Ty1
SymTy -> Ty1
forall loc. UrTy loc
SymTy
        Ty1
BoolTy -> Ty1
forall loc. UrTy loc
BoolTy
        ProdTy [Ty1]
tys -> [Ty1] -> Ty1
forall loc. [UrTy loc] -> UrTy loc
ProdTy ([Ty1] -> Ty1) -> [Ty1] -> Ty1
forall a b. (a -> b) -> a -> b
$ (Ty1 -> Ty1) -> [Ty1] -> [Ty1]
forall a b. (a -> b) -> [a] -> [b]
map Ty1 -> Ty1
goTy [Ty1]
tys
        SymDictTy Maybe Var
mv Ty1
_ty -> Maybe Var -> Ty1 -> Ty1
forall loc. Maybe Var -> Ty1 -> UrTy loc
SymDictTy Maybe Var
mv Ty1
forall loc. UrTy loc
CursorTy
        PDictTy Ty1
k Ty1
v -> Ty1 -> Ty1 -> Ty1
forall loc. UrTy loc -> UrTy loc -> UrTy loc
PDictTy (Ty1 -> Ty1
goTy Ty1
k) (Ty1 -> Ty1
goTy Ty1
v)
        PackedTy DataCon
_ ()
_ -> Ty1
forall loc. UrTy loc
CursorTy
        Ty1
ArenaTy -> Ty1
forall loc. UrTy loc
ArenaTy
        VectorTy Ty1
t -> Ty1 -> Ty1
forall loc. UrTy loc -> UrTy loc
VectorTy (Ty1 -> Ty1
goTy Ty1
t)
        ListTy Ty1
t -> Ty1 -> Ty1
forall loc. UrTy loc -> UrTy loc
ListTy (Ty1 -> Ty1
goTy Ty1
t)
        Ty1
PtrTy -> Ty1
forall loc. UrTy loc
PtrTy
        Ty1
CursorTy -> Ty1
forall loc. UrTy loc
CursorTy
        Ty1
SymSetTy -> Ty1
forall loc. UrTy loc
SymSetTy
        Ty1
SymHashTy -> Ty1
forall loc. UrTy loc
SymHashTy
        Ty1
IntHashTy -> Ty1
forall loc. UrTy loc
IntHashTy