content: 
  module Interpreter (execTop, readBackClosed) where
  
  import Common hiding (Lvl)
  import Zonk   hiding (Tm)
  
  import qualified Zonk   as Z
  import qualified Common as C
  import Pretty
  
  import Data.IORef
  
  --------------------------------------------------------------------------------
  
  type Tm    = Z.Tm Closed
  type Lvl   = (?lvl   :: C.Lvl)
  type Env a = (?env   :: [(Name, a)])
  type Stage = (?stage :: Int)
  type CEnv  = Env Closed
  type OEnv  = Env Open
  
  data Closed
    = CLam Name (NoShow (Closed -> Closed)) (NoShow (C.Lvl -> Open -> Open))
    | CAction (NoShow (IO Closed))
    | CRef (NoShow (IORef Closed))
    | CQuote Open
    | CErased String
    | CRec [(Name, Closed)]
    | CNat Integer
    deriving Show
  
  data Open
    = OVar C.Lvl
    | OLet Name Open (NoShow (C.Lvl -> Open))
    | OLam Name (NoShow (C.Lvl -> Open -> Open))
    | OApp Open Open
    | OErased String
    | OQuote Open
    | OSplice Open
    | OReturn Open
    | OBind Name Open (NoShow (C.Lvl -> Open))
    | OSeq Open Open
    | ONew Open
    | OWrite Open Open
    | ORead Open
    | OClosed (Maybe Name) Closed
    | OSuc Open
    | ONatElim Open Open Open
    | OProj Open Name
    | ORec [(Name, Open)]
    | OOpen [Name] Open (NoShow (C.Lvl -> Open))
    | OReadNat
    | OPrintNat Open
    | OLog String
    deriving Show
  
  def :: Name -> v -> (Env v => a) -> (Env v => a)
  def x t act = let ?env = (x, t): ?env in act
  
  stage :: Int -> (Stage => a) -> a
  stage s act = let ?stage = s in act
  
  lvl :: C.Lvl -> (Lvl => a) -> a
  lvl l act = let ?lvl = l in act
  
  env :: [(Name, v)] -> (Env v => a) -> a
  env e act = let ?env = e in act
  
  fresh :: (Lvl => Open -> a) -> Lvl => a
  fresh act = let v = OVar ?lvl in let ?lvl = ?lvl + 1 in act v
  
  freshes :: [Name] -> (Lvl => a) -> Lvl => a
  freshes xs act = let ?lvl = ?lvl + coerce (length xs) in seq ?lvl act
  
  closeEnv :: (OEnv => a) -> CEnv => a
  closeEnv act = let ?env = map (\(x, v) -> (x, OClosed (Just x) v)) ?env in act
  
  idEnv :: C.Lvl -> [(Name, Open)]
  idEnv l = map (\x -> ("lvl"++show x, OVar x)) [l-1,l-2..0]
  
  --------------------------------------------------------------------------------
  
  eRun :: Closed -> IO Closed
  eRun = \case
    CAction t -> coerce t
    _         -> impossible
  
  cApp :: Closed -> Closed -> Closed
  cApp t u = case t of
    CLam x f g -> coerce f u
    t          -> error $ show t
  
  cProj :: Closed -> Name -> Closed
  cProj t x = case t of
    CRec ts -> fromJust $ lookup x ts
    _       -> impossible
  
  cSplice :: Closed -> Maybe String -> Closed
  cSplice t loc = case t of
    CQuote t -> env [] $ lvl 0 $ ceval $ traceGen t loc
    t        -> error $ show t
  
  cSuc :: Closed -> Closed
  cSuc = \case
    CNat n -> CNat (n + 1)
    _      -> impossible
  
  cNatElim :: Closed -> Closed -> Closed -> Closed
  cNatElim s z n =
    let go 0 = z
        go n = s `cApp` CNat (n - 1) `cApp` go (n - 1)
    in case n of
      CNat n -> go n
      _      -> impossible
  
  cOpen :: [Name] -> Closed -> (CEnv => a) -> CEnv => a
  cOpen xs t u = let ?env = foldl' (\env x -> ((x,) $! cProj t x) : env) ?env xs
                 in seq ?env u
  
  oNatElim :: Lvl => Open -> Open -> Open -> Open
  oNatElim s z n =
    let go 0 = z
        go n = let m = n - 1 in
               s `oApp` OClosed Nothing (CNat m) `oApp` go m
    in case n of
      OClosed _ (CNat n) -> go n
      OClosed _ _        -> impossible
      OSuc n             -> s `oApp` n `oApp` oNatElim s z n
      n                  -> ONatElim s z n
  
  oQuote :: Open -> Open
  oQuote = \case
    OSplice t -> t
    t         -> OQuote t
  
  exec :: CEnv => Tm -> IO Closed
  exec = \case
    Var x         -> eRun (snd (?env !! coerce x))
    Let x t u     -> def x (ceval t) (exec u)
    Lam x t       -> impossible
    App t u       -> eRun $ cApp (ceval t) (ceval u)
    Erased{}      -> impossible
    Quote t       -> impossible
    Splice t pos  -> eRun $ cSplice (ceval t) pos
    Return t      -> pure $! ceval t
    Bind x t u    -> do {t <- exec t; def x t $ exec u}
    Seq t u       -> exec t >> exec u
    New t         -> CRef <$!> (coerce (newIORef $! ceval (coerce t)))
    Write t u     -> case ceval t of
                       CRef r -> CErased "tt" <$ (writeIORef (coerce r) $! ceval u)
                       _      -> impossible
    Read t        -> case ceval t of
                       CRef r -> readIORef (coerce r)
                       _      -> impossible
    NatLit n      -> impossible
    Suc t         -> impossible
    NatElim s z n -> eRun $ cNatElim (ceval s) (ceval z) (ceval n)
    Rec ts        -> impossible
    Proj t x      -> eRun (cProj (ceval t) x)
    CSP x t       -> eRun t
    Open xs t u   -> cOpen xs (ceval t) (exec u)
    ReadNat       -> do putStr "> "
                        n <- read <$> getLine
                        if n < 0 then error "negative integer"
                                 else pure $ CNat n
    PrintNat t    -> case ceval t of
                       CNat n -> print n >> pure (CRec [])
                       _      -> impossible
    Log s         -> putStrLn s >> pure (CRec [])
  
  ceval :: CEnv => Tm -> Closed
  ceval = \case
    Var x         -> snd (?env !! coerce x)
    Let x t u     -> def x (ceval t) (ceval u)
    Lam x t       -> CLam x (coerce (\v -> def x v $ ceval t))
                            (coerce (\l v -> closeEnv $ def x v $ lvl l $ stage 0 $ oeval t))
    App t u       -> cApp (ceval t) (ceval u)
    Erased s      -> CErased s
    Quote t       -> CQuote (closeEnv $ lvl 0 $ stage 1 $ oeval t)
    Splice t pos  -> cSplice (ceval t) pos
    t@Return{}    -> CAction (coerce (exec t))
    t@Bind{}      -> CAction (coerce (exec t))
    t@Seq{}       -> CAction (coerce (exec t))
    t@New{}       -> CAction (coerce (exec t))
    t@Write{}     -> CAction (coerce (exec t))
    t@Read{}      -> CAction (coerce (exec t))
    t@ReadNat     -> CAction (coerce (exec t))
    t@PrintNat{}  -> CAction (coerce (exec t))
    t@Log{}       -> CAction (coerce (exec t))
    NatLit n      -> CNat n
    Suc t         -> cSuc (ceval t)
    NatElim s z n -> cNatElim (ceval s) (ceval z) (ceval n)
    Rec ts        -> CRec (fmap (fmap ceval) ts)
    Proj t x      -> cProj (ceval t) x
    CSP _ t       -> t
    Open xs t u   -> cOpen xs (ceval t) (ceval u)
  
  
  oApp :: Lvl => Open -> Open -> Open
  oApp t u = case (t, u) of
    (OClosed x (CLam _ f _), OClosed x' u) -> OClosed Nothing (coerce f u)
    (OClosed x (CLam _ _ f), u           ) -> coerce f ?lvl u
    (OLam _ f              , u           ) -> coerce f ?lvl u
    (t                     , u           ) -> OApp t u
  
  oSplice :: Open -> Open
  oSplice = \case
    OClosed _ (CQuote t) -> t
    OQuote t             -> t
    t                    -> OSplice t
  
  oProj :: Open -> Name -> Open
  oProj t x = case t of
    OClosed _ (CRec ts) -> OClosed Nothing (fromJust $ lookup x ts)
    ORec ts             -> fromJust $ lookup x ts
    t                   -> OProj t x
  
  oOpen :: [Name] -> Open -> (OEnv => a) -> OEnv => a
  oOpen xs t u =
    let ?env = foldl' (\env x -> (x, oProj t x):env) ?env xs
    in seq ?env u
  
  oSuc :: Open -> Open
  oSuc = \case
    OClosed _ (CNat n) -> OClosed Nothing (CNat (n + 1))
    t                  -> OSuc t
  
  defs :: [Name] -> C.Lvl -> (OEnv => a) -> (OEnv => a)
  defs xs l act =
    let ?env = snd $ foldl' (\(l, env) x -> (l+1, (x, OVar l):env)) (l-coerce (length xs), ?env) xs in
    seq ?env act
  
  oeval :: OEnv => Lvl => Stage => Tm -> Open
  oeval = \case
    Var x      -> snd (?env !! coerce x)
    Lam x t    -> OLam x (coerce \l v -> lvl l $ def x v $ oeval t)
    Erased s   -> OErased s
    Quote t    -> oQuote $ stage (?stage + 1) $ oeval t
    Return t   -> OReturn (oeval t)
    Bind x t u -> OBind x (oeval t) (NoShow \l -> lvl l $ def x (OVar (l-1)) $ oeval u)
    Seq t u    -> OSeq (oeval t) (oeval u)
    New t      -> ONew (oeval t)
    Write t u  -> OWrite (oeval t) (oeval u)
    Read t     -> ORead (oeval t)
    ReadNat    -> OReadNat
    PrintNat t -> OPrintNat (oeval t)
    Log s      -> OLog s
    CSP x t    -> OClosed x t
    Rec ts     -> ORec (fmap (fmap oeval) ts)
    NatLit n   -> OClosed Nothing (CNat n)
    t -> case ?stage of
      0 -> case t of
        Let x t u     -> def x (oeval t) (oeval u)
        Open xs t u   -> oOpen xs (oeval t) (oeval u)
        App t u       -> oApp (oeval t) (oeval u)
        Splice t loc  -> case oeval t of
                           OClosed _ (CQuote t) -> OClosed Nothing $ lvl 0 $ env [] $ ceval $ traceGen t loc
                           OQuote t             -> env (idEnv ?lvl) $ oeval $ traceGen t loc
                           t                    -> OSplice t
        NatElim s z n -> oNatElim (oeval s) (oeval z) (oeval n)
        Proj t x      -> oProj (oeval t) x
        Suc t         -> oSuc (oeval t)
      _ -> case t of
        Let x t u     -> OLet x (oeval t) (NoShow \l -> lvl l $ def x (OVar (l-1)) $ oeval u)
        Open xs t u   -> OOpen xs (oeval t) (NoShow \l -> lvl l $ defs xs l $ oeval u)
        App t u       -> OApp (oeval t) (oeval u)
        Splice t pos  -> oSplice $ stage (?stage - 1) $ oeval t
        NatElim s z n -> ONatElim (oeval s) (oeval z) (oeval n)
        Proj t x      -> OProj (oeval t) x
        Suc t         -> OSuc (oeval t)
  
  traceGen :: Lvl => Open -> Maybe String -> Tm
  traceGen t loc =
    let t' = gen t
        freevars = map (\l -> "@" ++ show (?lvl - l - 1)) [0.. ?lvl - 1]
        displayCode  = prettyTm False 0 0 freevars (Z.unzonk t') ""
        -- displayCode = show t' ++ " | " ++ (show $ Z.unzonk t')
    in case loc of
      Nothing ->
        trace ("CODE GENERATED: \n\n" ++ displayCode ++ "\n") t'
      Just loc ->
        trace ("CODE GENERATED AT:\n" ++ loc ++ "\nCODE:\n\n" ++ displayCode ++ "\n") t'
  
  gen :: Lvl => Open -> Tm
  gen = \case
    OVar x             -> Var (coerce (?lvl - x - 1))
    OLet x t u         -> Let x (gen t) $ fresh \_ -> gen (coerce u ?lvl)
    OLam x t           -> Lam x $ fresh \v -> gen (coerce t ?lvl v)
    OApp t u           -> App (gen t) (gen u)
    OErased s          -> Erased s
    OQuote t           -> Quote (gen t)
    OSplice t          -> Splice (gen t) Nothing
    OReturn t          -> Return (gen t)
    OBind x t u        -> Bind x (gen t) $ fresh \_ -> gen (coerce u ?lvl)
    OSeq t u           -> Seq (gen t) (gen u)
    ONew t             -> New (gen t)
    OWrite t u         -> Write (gen t) (gen u)
    ORead t            -> Read (gen t)
    OSuc t             -> Suc (gen t)
    ONatElim s z n     -> NatElim (gen s) (gen z) (gen n)
    ORec ts            -> Rec (fmap (fmap gen) ts)
    OProj t x          -> Proj (gen t) x
    OClosed x (CNat n) -> NatLit n -- inline closed numeral into code
    OClosed x t        -> CSP x t
    OOpen xs t u       -> Open xs (gen t) $ freshes xs $ gen (coerce u ?lvl)
    OPrintNat t        -> PrintNat (gen t)
    OReadNat           -> ReadNat
    OLog s             -> Log s
  
  -- Only for pretty printing purposes
  readBackClosed :: Closed -> Tm
  readBackClosed t = let ?lvl = 0 in case t of
    CLam x _ t -> Lam x $ fresh \v -> gen (coerce t ?lvl v)
    CAction _  -> Erased "Action"
    CRef _     -> Erased "Ref"
    CQuote t   -> Quote (gen t)
    CErased s  -> Erased s
    CRec ts    -> Rec (fmap (fmap readBackClosed) ts)
    CNat n     -> NatLit n
  
  execTop :: Tm -> IO Closed
  execTop t = env [] $ exec t
  