Compiler.hs

{-# language OverloadedStrings #-}

module Compiler (genTop) where

import Common hiding (Lvl)
import StringBuilder
import qualified Zonk   as Z

import Control.Monad.State.Strict
import Data.String
import Data.Void
import Lens.Micro.Platform
import Prelude hiding (const, tail)
import qualified Data.Set as S
import System.Directory

import Paths_dtt_rtcg


-- Closure conversion
--------------------------------------------------------------------------------

type ZTm = Z.Tm Void

data Top
  = TLet Name Tm Top
  | TBind Name Tm Top
  | TSeq Tm Top
  | TBody Tm
  | TClosure Name [Name] Name Tm Top -- name, env, arg, body
  deriving Show

data Tm
  = Var Name
  | Let Name Tm Tm
  | LiftedLam Name Name [Name] -- function name, arg name, env application
  | Lam Name Tm
  | App Tm Tm
  | Erased String
  | Quote Tm
  | Splice Tm (Maybe String)
  | Return Tm
  | Bind Name Tm Tm
  | Seq Tm Tm
  | New Tm
  | Write Tm Tm
  | Read Tm
  | Log String
  | ReadNat
  | PrintNat Tm
  | Rec [(Name, Tm)]
  | Proj Tm Name
  | NatLit Integer
  | Suc Tm
  | NatElim Tm Tm Tm
  deriving Show

type TopClosures = [(Name, [Name], Name, Tm)]

data S = S {
    sFreeVars :: S.Set Name
  , sNextId   :: Int
  , sClosures :: TopClosures
  } deriving Show

makeFields ''S

type Env     = (?env     :: [(Bool, Name)])  -- is top-level, name
type Mode    = (?mode    :: Maybe Int)
type TopName = (?topName :: String)

mangle :: Name -> Name
mangle x = '$' : concatMap (\case
  '\'' -> "__"
  c    -> [c]) x

freshenName :: Env => Name -> Name
freshenName x = go (mangle x) where
  go :: Env => Name -> Name
  go x | any ((==x).(\(_, x) -> x)) ?env =
         go $ x ++ show (length ?env)
       | otherwise = x

fresh :: Name -> (Env => Name -> a) -> Env => a
fresh x act = let x' = freshenName x in
              let ?env = (False, x') : ?env in
              act x'

-- create fresh name, run action, delete bound name from freevars of the result
bind :: Name -> (Env => Name -> State S a) -> Env => State S a
bind x act = fresh x \x -> do
  a <- act x
  s <- get
  freeVars %= S.delete x
  pure a

-- desugar Open to iterated Let
ccOpen :: Env => Mode => TopName => [Name] -> ZTm -> ZTm -> State S Tm
ccOpen xs t u = do
  t <- cconv t
  let tmp = "$tmp" ++ show (length ?env)
  let go :: Env => Mode => TopName => [Name] -> State S Tm
      go []     = cconv u
      go (x:xs) = bind x \x' -> Let x' (Proj (Var tmp) x) <$> go xs
  Let tmp t <$> go xs

cconv :: Env => Mode => TopName => ZTm -> State S Tm
cconv = \case
  Z.Var x -> do
    (top, x) <- pure $! ?env !! coerce x
    when (not top) $ freeVars %= S.insert x
    pure $ Var x

  Z.Let x t u -> do
    t <- cconv t
    bind x \x -> Let x t <$> cconv u

  Z.Lam x t -> case ?mode of
    Nothing -> fresh x \x -> do
      old_fvs    <- use freeVars
      freeVars   .= S.empty
      t          <- cconv t
      captureSet <- S.delete x <$> use freeVars
      capture    <- pure $ S.toList captureSet
      clId       <- nextId <<%= (+1)
      clName     <- pure $! ?topName++show clId++"_"
      closures   %= ((clName, capture, x, t):)
      freeVars   .= S.union old_fvs captureSet
      pure $ LiftedLam clName x capture
    Just{} ->
      bind x \x -> Lam x <$> cconv t

  Z.App t u       -> App <$> cconv t <*> cconv u
  Z.Erased s      -> pure $ Erased s
  Z.Quote t       -> case ?mode of
                       Nothing -> do
                         let ?mode = Just 1
                         Quote <$> cconv t
                       Just s -> do
                         let ?mode = Just $! s + 1
                         Quote <$> cconv t
  Z.Splice t p    -> do let ?mode = case ?mode of
                              Nothing -> Nothing
                              Just 0  -> Just 0
                              Just s  -> Just $! s - 1
                        Splice <$> cconv t <*> pure p

  Z.Return t      -> Return <$> cconv t
  Z.Bind x t u    -> do {t <- cconv t; bind x \x -> Bind x t <$> cconv u}
  Z.Seq t u       -> Seq <$> cconv t <*> cconv u
  Z.New t         -> New <$> cconv t
  Z.Write t u     -> Write <$> cconv t <*> cconv u
  Z.Read t        -> Read <$> cconv t
  Z.NatLit n      -> pure $ NatLit n
  Z.Suc t         -> Suc <$> cconv t
  Z.NatElim s z n -> NatElim <$> cconv s <*> cconv z <*> cconv n
  Z.Rec ts        -> Rec <$> traverse (traverse cconv) ts
  Z.Proj t x      -> Proj <$> cconv t <*> pure x
  Z.Open xs t u   -> ccOpen xs t u
  Z.PrintNat t    -> PrintNat <$> cconv t
  Z.ReadNat       -> pure ReadNat
  Z.Log s         -> pure $ Log s


cconv0 :: Env => Mode => Name -> ZTm -> (Tm, TopClosures)
cconv0 x t =
  let ?topName = x in
  case runState (cconv t) (S mempty 0 []) of
    (t, S _ _ cs) -> (t, cs)

addClosures :: TopClosures -> Top -> Top
addClosures cs t =
  foldl' (\acc (x, env, arg, t) -> TClosure x env arg t acc) t cs

-- desugar Open to iterated Let
ccOpenTop :: Env => Mode => [Name] -> ZTm -> ZTm -> Top
ccOpenTop xs t u =
  let tmp = "$tmp" ++ show (length ?env) in
  let (t', cs) = cconv0 tmp t in

  let go :: Env => Mode => [Name] -> Top
      go []     = cconvTop u
      go (x:xs) = fresh x \x' -> TLet x' (Proj (Var tmp) x) (go xs) in

  addClosures cs $ TLet tmp t' $ go xs

cconvTop :: Env => Mode => ZTm -> Top
cconvTop = \case
  Z.Let (freshenName -> x) t u ->
    let (t', cs) = cconv0 x t in
    let ?env     = (True, x) : ?env in
    addClosures cs $ TLet x t' (cconvTop u)
  Z.Bind (freshenName -> x) t u ->
    let (t', cs) = cconv0 x t in
    let ?env     = (True, x) : ?env in
    addClosures cs $ TBind x t' (cconvTop u)
  Z.Seq t u ->
    let (t', cs) = cconv0 "$cl" t in
    addClosures cs $ TSeq t' (cconvTop u)
  Z.Open xs t u ->
    ccOpenTop xs t u
  t ->
    case cconv0 "$cl" t of
      (t', cs) -> addClosures cs (TBody t')

runCConv :: ZTm -> Top
runCConv t = let ?env = []; ?mode = Nothing in cconvTop t


-- Code generation
--------------------------------------------------------------------------------

type Stage   = (?stage :: Int)
data IsTail' = Tail | NonTail deriving Show
type IsTail  = (?isTail :: IsTail')
type Cxt     = (?cxt :: [(Name, Bool)]) -- (Name, is closed)
                                        -- We need to box if we reference a closed value from open code

stage :: Int -> (Stage => a) -> a
stage s act = let ?stage = s in act

tail :: (IsTail => Out) -> Out
tail act = let ?isTail = Tail in act

nonTail :: (IsTail => Out) -> Out
nonTail act = let ?isTail = NonTail in act

jLet :: Name -> Bool -> (IsTail => Cxt => Out) -> (IsTail => Cxt => Out) -> (IsTail => Cxt => Out)
jLet x closed t u = case ?isTail of
  Tail -> let u' = (let ?cxt = (x, closed): ?cxt in tail u) in
          "const " <> str x <> " = " <> indent (nonTail t) <> ";" <> newl <> u'
  _    -> let u' = (let ?cxt = (x, closed): ?cxt in nonTail u) in
            "((" <> str x <> ") => " <> parens u' <> ")(" <> nonTail t <> ")"

jSeq :: IsTail => Cxt => (IsTail => Out) -> (IsTail => Out) -> Out
jSeq t u = case ?isTail of
  Tail -> indent (nonTail t) <> ";" <> newl <> tail u
  _    -> "(_) => " <> parens (nonTail u) <> ")(" <> nonTail t <> ")"

jTuple :: [Out] -> Out
jTuple xs = "(" <> go xs <> ")" where
  go = \case
    []     -> mempty
    [x]    -> x
    (x:xs) -> x <> ", " <> go xs

jReturn :: (IsTail => Out) -> (IsTail => Out)
jReturn t = case ?isTail of
  Tail -> "return " <> nonTail t
  _    -> t

jLam :: Cxt => [Name] -> Bool -> (IsTail => Cxt => Out) -> (IsTail => Cxt => Out)
jLam xs closed t =
  let t' = (let ?cxt = map (,closed) xs ++ ?cxt in tail t) in
  jReturn $ jTuple (map str xs) <> " => {" <> t' <> "}"

jLamExp :: [Name] -> Bool -> (IsTail => Cxt => Out) -> (IsTail => Cxt => Out)
jLamExp xs closed t =
  let t' = (let ?cxt = map (,closed) xs ++ ?cxt in nonTail t) in
  jReturn $ jTuple (map str xs) <> " => " <> t'

cApp :: (IsTail => Out) -> (IsTail => Out) -> (IsTail => Out)
cApp t u = jReturn $ parens t <> "._1(" <> u <> ")"

parens :: Out -> Out
parens t = "(" <> t <> ")"

jApp :: (IsTail => Out) -> (IsTail => [Out]) -> (IsTail => Out)
jApp t args = jReturn $ t <> jTuple args

cRun :: (IsTail => Out) -> (IsTail => Out)
cRun t = jReturn $ t <> "()"

cProj :: (IsTail => Out) -> Name -> (IsTail => Out)
cProj t x = jReturn $ t <> "." <> str x

jClosure :: [Name] -> Name -> Bool -> (Cxt => IsTail => Out) -> (Cxt => IsTail => Out)
jClosure []  x closed t = jLam [x] closed t
jClosure env x closed t = jLamExp env closed $ jLam [x] closed t

jAppClosure :: (IsTail => Out) -> (IsTail => [Out]) -> (IsTail => Out)
jAppClosure t args = jReturn $ case args of
  []   -> t
  args -> t <> jTuple args

closeVar :: Name -> Name
closeVar x = x ++ "c"

openVar :: Name -> Name
openVar x = x ++ "o"

cRec :: Cxt => [(Name, Tm)] -> Out
cRec []          = mempty
cRec [(x, t)]    = str x <> ": " <> nonTail (ceval t)
cRec ((x, t):ts) = str x <> ": " <> nonTail (ceval t) <> ", " <> cRec ts

spliceLoc :: Maybe String -> Out
spliceLoc = \case
  Nothing -> "undefined"
  Just loc -> "[" <> go (lines loc) <> "]" where
    go []     = ""
    go [l]    = strLit l
    go (l:ls) = strLit l <> ", " <> go ls

execTop :: Cxt => Top -> Out
execTop t = tail $ go t where
  go :: IsTail => Cxt => Top -> Out
  go = \case
    TLet x t u  ->
      jLet x True (ceval t) (newl <> go u)
    TBind x t u ->
      jLet x True (exec t) (newl <> go u)
    TSeq t u    ->
      jSeq (exec t) (newl <> go u)
    TClosure x env arg body t ->
      jLet (closeVar x) True (jClosure env arg True $ ceval body) $
      jLet (openVar x)  True (jClosure env arg False $ stage 0 $ oeval body) $
      go t

    -- finalize
    TBody t ->
      "const main_ = () => {" <> exec t <> "};" <> newl <>
      "console.log('RESULT:');console.log(util_.inspect(main_(), false, null))" <> newl

exec :: IsTail => Cxt => Tm -> Out
exec = \case
  Var x         -> cRun (str x)
  Let x t u     -> jLet x True (ceval t) (exec u)
  LiftedLam{}   -> impossible
  Lam{}         -> impossible
  App t u       -> cRun (ceval t `cApp` ceval u)
  Erased{}      -> impossible
  Quote{}       -> impossible
  Splice t loc  -> cRun (jApp "codegenClosed_" [ceval t, spliceLoc loc])
  Return t      -> jReturn $ ceval t
  Bind x t u    -> jLet x True (exec t) (exec u)
  Seq t u       -> jSeq (exec t) (exec u)
  New t         -> jReturn $ "{_1 : " <> ceval t <> "}"
  Write t u     -> nonTail $ ceval t <> "._1 = " <> ceval u
  Read t        -> jReturn $ ceval t <> "._1"
  Log s         -> jApp "log_" [strLit s]
  ReadNat       -> jApp "readNat_" []
  PrintNat t    -> jApp "printNat_" [ceval t]
  Rec{}         -> impossible
  Proj t x      -> cRun (cProj (ceval t) x)
  Suc{}         -> impossible
  NatLit{}      -> impossible
  NatElim s z n -> cRun (jApp "cNatElim_" [ceval s, ceval z, ceval n])

oevalVar :: Cxt => IsTail => Name -> Out
oevalVar x = case lookup x ?cxt of
  Nothing    -> impossible
  Just True  -> jApp "CSP_" [str x, strLit x]
  Just False -> jReturn (str x)

oRec :: Cxt => Stage => [(Name, Tm)] -> Out
oRec []         = mempty
oRec [(x,t)]    = "[" <> strLit x <> ", " <> nonTail (oeval t) <> "]"
oRec ((x,t):ts) = "[" <> strLit x <> ", " <> nonTail (oeval t) <> "], " <> oRec ts

ceval :: IsTail => Cxt => Tm -> Out
ceval = \case
  Var x             -> jReturn (str x)
  Let x t u         -> jLet x True (ceval t) (ceval u)
  LiftedLam f x env -> jReturn $ nonTail $
                          "{ _1 : " <> jAppClosure (str (closeVar f)) (map str env) <>
                          ", _2 : " <> jAppClosure (str (openVar f))  (map oevalVar env) <> "}"
  Lam{}           -> impossible
  App t u         -> cApp (ceval t) (ceval u)
  Erased s        -> jReturn "undefined"
  Quote t         -> stage 1 $ oeval t
  Splice t loc    -> jApp "codegenClosed_" [ceval t, spliceLoc loc]
  t@Return{}      -> jLam [] True $ exec t
  t@Bind{}        -> jLam [] True $ exec t
  t@Seq{}         -> jLam [] True $ exec t
  t@Write{}       -> jLam [] True $ exec t
  t@Read{}        -> jLam [] True $ exec t
  t@New{}         -> jLam [] True $ exec t
  t@Log{}         -> jLam [] True $ exec t
  t@ReadNat{}     -> jLam [] True $ exec t
  t@PrintNat{}    -> jLam [] True $ exec t
  Rec ts          -> jReturn $ "{" <> cRec ts <> "}"
  Proj t x        -> cProj (ceval t) x
  NatLit n        -> jReturn $ str (show n)
  Suc t           -> jApp "cSuc_" [ceval t]
  NatElim s z n   -> jApp "cNatElim_" [ceval s, ceval z, ceval n]

oeval :: IsTail => Cxt => Stage => Tm -> Out
oeval = \case
  Var x             -> oevalVar x
  Let x t u         -> case ?stage of
                         0 -> jLet x False (oeval t) (oeval u)
                         _ -> jApp "Let_" [strLit x, oeval t, jLam [x] False (oeval u)]
  LiftedLam f x env -> jApp "Lam_" [strLit x, jAppClosure (str (openVar f)) (map str env)]
  Lam x t           -> jApp "Lam_" [strLit x, jLam [x] False (oeval t)]
  App t u           -> case ?stage of
                         0 -> jApp "app_" [oeval t, oeval u]
                         _ -> jApp "App_" [oeval t, oeval u]
  Erased s          -> jReturn "CSP_undefined_"
  Quote t           -> jApp "quote_" [stage (?stage + 1) (oeval t)]
  Splice t loc      -> case ?stage of
                         0 -> jApp "codegenOpen_" [oeval t, spliceLoc loc]
                         _ -> stage (?stage - 1) $ jApp "splice_" [oeval t]
  Return t          -> jApp "Return_" [oeval t]
  Bind x t u        -> jApp "Bind_" [strLit x, oeval t, jLam [x] False (oeval u)]
  Seq t u           -> jApp "Seq_" [oeval t, oeval u]
  New t             -> jApp "New_" [oeval t]
  Write t u         -> jApp "Write_" [oeval t, oeval u]
  Read t            -> jApp "Read_" [oeval t]
  Log s             -> jApp "Log_" [strLit s]
  ReadNat           -> jReturn "ReadNat_"
  PrintNat t        -> jApp "PrintNat_" [oeval t]
  Rec ts            -> jApp "Rec_" [ "new Map([" <> oRec ts <> "])" ]
  Proj t x          -> case ?stage of
                         0 -> jApp "proj_" [oeval t, strLit x]
                         _ -> jApp "Proj_" [oeval t, strLit x]
  NatLit n          -> jApp "CSP_" [str (show n), strLit (show n)]

  Suc t             -> case ?stage of
                         0 -> jApp "suc_" [oeval t]
                         _ -> jApp "Suc_" [oeval t]
  NatElim s z n     -> case ?stage of
                         0 -> jApp "natElim_" [oeval s, oeval z, oeval n]
                         _ -> jApp "NatElim_" [oeval s, oeval z, oeval n]

genTop :: Z.Tm Void -> IO Out
genTop t = do
  src <- do
    exec_path <- getDataFileName "rts.js"
    doesFileExist exec_path >>= \case
      True -> readFile exec_path
      _    -> readFile "rts.js"
  let ?cxt = []
  return $! str src <> newl <> newl <> execTop (runCConv t)