{-# LANGUAGE OverloadedLists #-}
module Gyehoek.CPS.Stackify
  ( stackifyExp
  , stackifyProgram
  , module Gyehoek.CPS.Syntax
  ) where

import Gyehoek.CPS.Syntax
import Gyehoek.Stack.Syntax qualified as Stk
import Data.Sequence (Seq)
import Data.Sequence qualified as Seq
import Gyehoek.GenSym
import Effectful.Writer.Static.Shared
import Data.Foldable
import qualified Data.HashMap.Strict as H
import Data.List (elemIndex)
import Gyehoek.Prelude


type Stackify = Writer Stk.Program

runStackify :: Eff (Stackify : es) a -> Eff es (a, Stk.Program)
runStackify = runWriter

live :: Free a => Env -> a -> List Name
live g e = free' e & filter \x ->
  x `H.member` g.bound
  && not (x `elem` g.contStack)

data BlockBuilder
  = Code (List Stk.Instr) BlockBuilder
  | Tail Stk.Tail
  deriving (Show, Generic)

buildBlock :: BlockBuilder -> Stk.Block
buildBlock = go [] where
  go acc (Code xs bb) = go (acc ++ xs) bb
  go acc (Tail t) = Stk.MkBlock acc t

emitRoutine :: Stackify :> es => Stk.Routine -> Eff es ()
emitRoutine rt = tell [rt]

stackify
  :: (GenSym :> es, Stackify :> es)
  => Env -> Exp -> Eff es BlockBuilder

stackify g (ExpLetRec [(f, kap@(AbsKappa' xs m))] e) = do
  let vs = (f, Stk.ValLabel f) : (bindReg <$> xs)
  let ls = live g kap
  m' <- stackify (g & #bound .~ H.fromList (vs ++ (bindReg <$> ls))) m
  emitRoutine $
    Stk.MkRoutine f xs . buildBlock $
      Code [Stk.Pop x | x <- ls] m'
  let g' = g & #bound . at f ?~ Stk.ValLabel f
             & #liveness . at f ?~ ls
  stackify g' e

stackify g (ExpLetRec [(f, AbsLambda' xs k m)] e) = do
  let vs = (k:xs) <&> \x -> (x, Stk.ValReg x)
  m' <- stackify (g & #bound .~ H.fromList vs
                    & #contStack %~ (k:)) m
  emitRoutine $ Stk.MkRoutine f xs (buildBlock m')
  stackify g e

stackify g (ExpIf c t f) = do
  let c' = stackifyVal g c
  t' <- buildBlock <$> stackify g t
  f' <- buildBlock <$> stackify g f
  pure . Tail $ Stk.If c' t' f'

stackify g (ExpApply f xs ktail) = pure $
    Code [ Stk.PushCont k ] $
    Code [ Stk.Push (Stk.ValReg l) | l <- ls ] $
    Tail (Stk.TailCall (stackifyVal g f) (stackifyVal g <$> xs))
  where
    k = var g ktail
    ls = fold $ (k ^? #ValImm . #ImmLabel)
      >>= \klbl -> g ^. #liveness . at klbl

stackify g (ExpContinue k xs) =
  -- return continuations require popping the stack. how do we know
  -- when a continuation is a return continuation? is this a correct
  -- test?
  case elemIndex k g.contStack of
    Nothing -> pure . Tail $ Stk.TailCall (Stk.ValLabel k) xs'
    Just j -> do
      ktail <- gensym' @Name $ k ^. _Wrapped'
      pure $
        Code (replicate j $ Stk.PopCont "_") $
        Code [Stk.PopCont ktail] $
        Tail (Stk.TailCall (Stk.ValReg ktail) xs')
  where xs' = stackifyVal g <$> xs

stackify g (ExpPrim p (MkKappa [x] e)) = do
  e' <- stackify (g & #bound . at x ?~ Stk.ValReg x) e
  pure $
    Code [ Stk.Prim x (stackifyVal g <$> p) ] $
    e'

stackify _ e = error [i|unimplemented exp: #{e}|]

stackifyVal :: Env -> Val -> Stk.Val
stackifyVal g = \case
  ValImm imm -> Stk.ValImm imm
  ValVar v -> var g v
  v -> error [i|unimplemented val: #{v}|]

var :: Env -> Name -> Stk.Val
var g v = case g ^. #bound . at v of
  Just x -> x
  Nothing -> Stk.ValLabel v

bindReg :: Name -> (Name, Stk.Val)
bindReg x = (x, Stk.ValReg x)



data Env = MkEnv
  { bound :: HashMap Name Stk.Val
  -- | for each locally-bound continuation @k@, @liveness@ has an
  -- entry @(k,ls)@ where @ls@ is the sequence of registers @k@
  -- expects to find saved on the stack.
  , liveness :: HashMap Name (List Name)
  , contStack :: List Name
  }
  deriving (Show, Generic)

emptyEnv :: Env
emptyEnv = MkEnv mempty mempty ["halt"]



stackifyExp :: GenSym :> es => Name -> Exp -> Eff es Stk.Program
stackifyExp lbl e = do
  (code,p) <- runStackify $ stackify emptyEnv e
  pure $ p <> [ Stk.MkRoutine lbl [] (buildBlock code) ]

stackifyProgram :: GenSym :> es => Program -> Eff es Stk.Program
stackifyProgram (MkProgram e) = stackifyExp "main" e



fac :: Program
fac = [cps|
  (letrec ((fac (λ (n ktail)
                  (prim (zero? n)
                        (κ (x0)
                           (if x0
                               (continue ktail 1)
                               (prim (- n 1)
                                 (κ (x1)
                                    (letrec ((fac-k0
                                              (κ (x2)
                                                 (prim (* n x2)
                                                   (κ (x3)
                                                     (continue ktail x3))))))
                                      (fac x1 fac-k0))))))))))
    (fac 6 halt))
|]