{-# LANGUAGE OverloadedLists #-}
module Gyehoek.CPS.Stackify
  ( 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, nub)
import Data.Text qualified as T
import Gyehoek.Prelude
import Debug.Pretty.Simple
import qualified Gyehoek.Sexp as S


type Stackify = Writer Stk.Program

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

live :: Free a => Env -> a -> List Name
-- TODO: free' should return an OSet lol
live g e = nub (free' e) & filter \x ->
  x `elem` 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, AbsKappa kap)] e) = do
  kap' <- stackifyKappa g kap
  emitRoutine . Stk.MkRoutine (MkLabel f) . buildBlock $ kap'
  stackify g e

stackify g (ExpLetRec [(f, AbsLambda lam)] e) = do
  lam' <- stackifyLambda g (MkLabel f) lam
  emitRoutine lam'
  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) = do
  pure $
    Code [ Stk.Push $ stackifyVal g (ValVar ktail)
         , Stk.Push $ stackifyVal g f
         ] $
    Code (pushArgs g xs) $
    Tail (Stk.Call (length xs))

-- assume that `k` is the continuation on top of the stack lol.
stackify g e@(ExpContinue k xs)
  | isn't (#_ValVar . only g.tail) k = pure $
    Code [ Stk.Push (stackifyVal g k) ] $
    Code (pushArgs g xs) $
    Tail $ Stk.TailCall (length xs)
  | otherwise = pure $
    Code (pushArgs g xs) $
    Tail (Stk.Return (length xs))

stackify g (ExpPrim (PrimCallCC withcc) cc) = do
  cc' <- stackifyKappa g cc
  cc_l <- gensym' @Name "cc"
  reified_cc_l <- gensym' @Name "reified-cc"
  emitRoutine . Stk.MkRoutine (MkLabel cc_l) . buildBlock $ cc'
  stackify g $
    ExpPrim PrimCaptureCC $
    MkKappa [reified_cc_l] $
    ExpApply withcc [ValVar reified_cc_l] cc_l

stackify g (ExpPrim p kap) = do
  kap' <- stackifyKappa g kap
  pure $ Code [ Stk.Prim (stackifyVal g <$> p) ] kap'

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

loadArgs :: List Name -> List Stk.Instr
loadArgs = imapOf itraversed \n x -> Stk.Load (MkReg x) n

pushArgs :: Env -> List Val -> List Stk.Instr
pushArgs g args = [ Stk.Push $ stackifyVal g x | x <- reverse args ]

-- affine
_ValName :: Traversal' Val Name
_ValName = failing #_ValVar (#_ValImm . #_ImmLabel . #_MkLabel)

stackifyKappa
  :: (Stackify :> es, GenSym :> es)
  => Env -> Kappa
  -> Eff es BlockBuilder
stackifyKappa g (MkKappa xs m) = do
  let g' = g & #bound <>:~ xs
  Code (loadArgs g'.bound)
    <$> stackify g' m

stackifyLambda
  :: (Stackify :> es, GenSym :> es)
  => Env -> Label -> Lambda
  -> Eff es Stk.Routine
stackifyLambda g name (MkLambda xs k m) = do
  m' <- stackify (g & #bound .~ xs & #tail .~ k) m
  pure $
    Stk.MkRoutine name . buildBlock $
    Code (loadArgs xs) $
    Code [Stk.Load (MkReg k) (length xs + 1)] m'

stackifyVal :: Env -> Val -> Stk.Val
stackifyVal g = \case
  ValImm imm -> Stk.ValImm imm
  ValVar v -> case regOf g v of
    Just r -> Stk.ValReg r
    Nothing -> Stk.ValLabel (MkLabel v)
  v -> error [i|unimplemented val: #{v}|]

regOf :: Env -> Name -> Maybe Reg
regOf g x
  | x `elem` g.bound || x == g.tail = Just . MkReg $ x
  | otherwise                       = Nothing



data Env = MkEnv
  -- | `bound` tracks the stack lifetime of bound variables.
  { bound :: List Name
  -- | 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 Label (List Name)
  , tail :: Name
  }
  deriving (Show, Generic)

emptyEnv :: Env
emptyEnv = MkEnv
  { bound = mempty
  , liveness = mempty
  , tail = "halt"
  }



stackifyProgram :: GenSym :> es => Program -> Eff es Stk.Program
stackifyProgram (MkProgram lam) = do
  let g = emptyEnv
  (_,p) <- runStackify $ emitRoutine =<< stackifyLambda g "start" lam
  pure p

letfn :: Program
letfn = [cps|
(λ (start-ktail0)
  (letrec ((lambda-body1
            (λ (x lambda-tail2)
              (prim (* x x) (κ (r3) (continue lambda-tail2 r3))))))
    (letrec ((let-body6
              (κ (square)
                (letrec ((r4 (κ (x5) (continue start-ktail0 x5))))
                  (square 4 r4)))))
      (continue let-body6 lambda-body1))))
|]

blah :: Program
blah = [cps|
(λ (ktail0)
          (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)))
|]