{-# 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)

stackify
  :: (GenSym :> es, Stackify :> es)
  => Env -> Exp -> Eff es (Seq Stk.Instr)

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
  tell [Stk.MkBlock f xs $
        [Stk.Pop x | x <- ls] <> toList 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
  tell [Stk.MkBlock f xs . toList $ m']
  stackify g e

stackify g (ExpIf c t f) = do
  t' <- stackify g t
  f' <- stackify g f
  pure [ Stk.If (stackifyVal g c) (toList t') (toList f') ]

stackify g (ExpApply f xs ktail) = do
  pure $
    [ Stk.PushCont k ]
    <> fromList [ Stk.Push (Stk.ValReg l) | l <- ls ]
    <> [ Stk.Call (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 [ Stk.Call (Stk.ValLabel k) xs' ]
    Just j -> do
      ktail <- gensym' $ k ^. _Wrapped'
      pure $
       Seq.replicate j (Stk.PopCont "_")
       <> [ Stk.PopCont ktail
          , Stk.Call (Stk.ValReg ktail) (stackifyVal g <$> 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 $ [ 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.MkProgram [ Stk.MkBlock lbl [] (toList 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))
|]