{-# LANGUAGE ViewPatterns #-}
module Gyehoek.Stack.VM
  ( VM(..)
  , Env(..)
  , eval
  , trace
  , module Gyehoek.Stack.Syntax
  ) where

import Gyehoek.Stack.Syntax
import Data.List (List)
import GHC.Generics (Generic)
import Control.Lens
import Data.HashMap.Strict (HashMap)
import Data.Text (Text)
import qualified Data.HashMap.Strict as H
import Data.String.Interpolate (i)
import Gyehoek.Scheme.Syntax (Sexp(..))
import Debug.Pretty.Simple (pTraceShowIdForceColor)
import qualified Data.List.NonEmpty as NE
import Data.Functor (($>))
import Data.List (unfoldr)


data VM = MkVM
  { stack :: List Obj
  , kstack :: List Name
  , code :: List Instr
  , registers :: HashMap Name Obj
  , stdout :: Text
  , result :: Maybe (List Obj)
  }
  deriving (Show, Generic)

data Env = MkEnv
  { blocks :: HashMap Name Block
  }
  deriving (Show, Generic)

step :: Env -> VM -> VM
step e vm = case vm ^. #code of
  c:cs -> stepI e (vm & #code .~ cs) c
  _ -> error "halt never called"

stepI :: Env -> VM -> Instr -> VM

stepI e vm (Push v) = vm & #stack %~ (evalVal e vm v :)

stepI e vm (PushCont k) = vm & #kstack %~ (k:)

stepI e vm (Prim r p) = case evalVal e vm <$> p of
  PrimZeroP x -> case x of
    ObjImm (ImmInt n) -> ret . ObjImm . ImmBool $ n == 0
    _ -> error [i|bad arg to zero?: #{x}|]
  PrimAdd x y -> arith_binop (+) x y
  PrimMul x y -> arith_binop (*) x y
  PrimSub x y -> arith_binop (-) x y
  PrimDiv x y -> arith_binop div x y
  x -> error [i|unimplemented prim: #{p}|]
  where
    ret v = vm & #registers . at r ?~ v
    arith_binop op (ObjImm (ImmInt x)) (ObjImm (ImmInt y)) =
      ret $ ObjImm (ImmInt (op x y))
    arith_binop _ x y = error [i|bad arith: #{x}, #{y}|]

stepI e vm (Pop r) = case vm ^. #stack of
  [] -> error "empty stack"
  (x:xs) -> vm & #registers . at r ?~ x
               & #stack .~ xs

stepI e vm (PopCont r) = case vm ^. #kstack of
  [] -> error "empty stack"
  (x:xs) -> vm & #registers . at r ?~ ObjImm (ImmLabel x)
               & #kstack .~ xs

stepI e vm (Call v xs) =
  case evalToLabel e vm v of
    "halt" -> vm & #result ?~ fmap (evalVal e vm) xs
    l -> vm & #code .~ b.code
            & #registers .~ fmap (evalVal e vm) (H.fromList $ b.params `zip` xs)
      where
        b = case e ^. #blocks . at l of
          Just x -> x
          Nothing -> error [i|undefined label: #{l}|]

stepI e vm (If c t f) =
  case evalVal e vm c of
    ObjImm (ImmBool False) -> vm & #code .~ f
    _ -> vm & #code .~ t

stepI e vm ins = error [i|unimplemented instruction: #{ins}|]

evalToLabel e vm v =
  case evalVal e vm v of
    ObjImm (ImmLabel x) -> x
    x -> error [i|not a label: #{x}|]

evalVal :: Env -> VM -> Val -> Obj
evalVal e vm = \case
  ValImm imm -> ObjImm imm
  ValReg r -> case vm ^. #registers . at r of
    Just x -> x
    Nothing -> error [i|undefined register: #{r}|]

initialVM :: VM
initialVM = MkVM
  { stack = []
  , kstack = ["halt"]
  , code = [Call (ValImm $ ImmLabel "main") []]
  , registers = mempty
  , stdout = ""
  , result = Nothing
  }

initialEnv :: Program -> Env
initialEnv (MkProgram bs) = MkEnv
  { blocks = bs & foldMap \b -> H.singleton b.label b
  }

loop :: (a -> Either b a) -> a -> b
loop f a = case f a of
  Right a' -> loop f a'
  Left b -> b

eval :: Program -> List Obj
eval p = initialVM & loop \vm -> case vm ^. #result of
  Nothing -> Right $ step (initialEnv p) vm
  Just rs -> Left rs

trace :: Program -> List VM
trace p = initialVM & unfoldr \vm ->
    case vm.result of
      Just _ -> Nothing
      Nothing -> Just (vm, step e vm)
  where e = initialEnv p