Files
gyehoek-hs/src/Gyehoek/Stack/VM.hs
T
2026-08-30 02:12:16 -06:00

497 lines
14 KiB
Haskell

{-# LANGUAGE ViewPatterns, MultilineStrings #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DeriveAnyClass #-}
module Gyehoek.Stack.VM
( VM(..)
, Env(..)
, eval
, trace
, module Gyehoek.Stack.Syntax
, writeObj
, traceEval
) where
import Gyehoek.Stack.Syntax
import Control.Lens
import qualified Data.HashMap.Strict as H
import Data.List (unfoldr, intersperse, compareLength)
import Gyehoek.Prelude
import qualified Data.List.NonEmpty as NE
import Lucid
import Data.Foldable (traverse_)
import qualified Gyehoek.Sexp as S
import Gyehoek.Jalmot
import Text.Pretty.Simple (pStringNoColor, pShowNoColor)
import Effectful.State.Static.Local (runState, evalState, get)
import Data.Traversable
import Control.Applicative (Alternative(..))
import Gyehoek.Sexp.Print (htmlData)
import Control.DeepSeq (deepseq, ($!!))
import Gyehoek.Sexp.Print (htmlData, htmlDatum)
import Control.DeepSeq (deepseq, ($!!))
import Data.String (fromString)
import Data.Monoid (First)
import GHC.Stack (popCallStack)
import Data.Maybe (fromMaybe)
-- | non-essential information maintained only to aide in debugging.
data DebugVM = MkDebugVM
{ activeRoutine :: Label
}
deriving (Show, Generic)
newtype Frame = MkFrame { locals :: List Obj }
deriving stock (Show, Generic)
-- affine
returnAddress :: Traversal' Frame Obj
returnAddress = #locals . _last
-- affine
activeProcedure :: Traversal' Frame Obj
activeProcedure = #locals . _init . _last
newtype Stack = MkStack { frames :: NonEmpty Frame }
deriving stock (Show, Generic)
data VM = MkVM
{ stack :: Stack
, code :: List Instr
, tail :: Tail
, registers :: HashMap Reg Obj
, stdout :: Text
, result :: Maybe (List Obj)
, debug :: DebugVM
}
deriving (Show, Generic)
type instance Index Frame = Int
type instance IxValue Frame = Obj
instance Ixed Frame where
ix j = wrappedIso . ix j
instance Cons Frame Frame Obj Obj where
_Cons = prism'
(\(x,MkFrame xs) -> MkFrame (x:xs))
\case
MkFrame (x:xs) -> Just (x, MkFrame xs)
MkFrame [] -> Nothing
instance Each Frame Frame Obj Obj where each = wrappedIso . each
instance Each Stack Stack Frame Frame where each = wrappedIso . each
pushes :: Foldable f => f Obj -> Frame -> Frame
pushes = flip $ foldr cons
_NonEmpty :: Iso (NonEmpty a) (NonEmpty b) (a, List a) (b, List b)
_NonEmpty = iso
(\(x:|xs) -> (x,xs))
(\(x,xs) -> x:|xs)
pushFrame :: Frame -> Stack -> Stack
pushFrame f (MkStack xs) = MkStack $ NE.cons f xs
activeFrame :: Lens' VM Frame
activeFrame = #stack . #frames . _NonEmpty . _1
data Env = MkEnv
{ labels :: HashMap Label Routine
}
deriving (Show, Generic)
step :: Jalmot :> es => Env -> VM -> Eff es VM
step g vm = case vm ^. #code of
c:cs -> stepI g (vm & #code .~ cs) c
[] -> stepT g vm vm.tail
vmerror :: (HasCallStack, Jalmot :> es) => Text -> Eff es a
vmerror = throwError . VMError
stepI :: Jalmot :> es => Env -> VM -> Instr -> Eff es VM
stepI e vm (Load r j) = do
x <- expectOf [i|object at index #{j}|] (activeFrame . ix j) vm
pure $ vm & #registers . at r ?~ x
stepI e vm (Push v) = traverseOf activeFrame push vm
where push xs = cons <$> evalVal e vm v <*> pure xs
stepI g vm (Prim p) = stepP g vm p
stepI e vm (Pop r) = case vm ^? activeFrame . _Cons of
Nothing -> vmerror "empty stack"
Just (x,xs) -> pure $ vm & #registers . at r ?~ x
& activeFrame .~ xs
stepI e vm ins = vmerror [i|unimplemented instruction: #{ins}|]
stepT :: Jalmot :> es => Env -> VM -> Tail -> Eff es VM
stepT g vm tc@(Call nargs) = do
(args,f,ret,frm) <- parseCall nargs (vm ^. activeFrame)
& expectOf [i|bad call: #{show tc}|] _Just
rt <- getRoutine g f
let newFrame = MkFrame $ args ++ [f,ret]
pure $ vm
& jumpToRoutine rt
& activeFrame .~ frm
& #stack %~ pushFrame newFrame
-- it is not essential we clear the registers, but it'll
-- make bugs more obvious.
& #registers .~ mempty
stepT g vm tc@(Return nret) = do
(xs,_) <- splitAtExact nret (vm ^. activeFrame . #locals)
& expectOf [i|bad return: #{show tc}|] _Just
expectOf [i|no return addr|] (activeFrame . returnAddress) vm >>= \case
ObjLabel "halt" -> pure $ vm & #result ?~ xs
ra -> do
rt <- getRoutine g ra
vm & traverseOf #stack (fmap snd . popFrame)
& mapped . activeFrame %~ pushes xs
& mapped %~ jumpToRoutine rt
-- it is not essential we clear the registers, but it'll make
-- bugs more obvious.
& mapped . #registers .~ mempty
stepT g vm tc@(TailCall nargs) = do
(args,f,ra) <- parseTailCall nargs (vm ^. activeFrame)
& expectOf [i|bad call: #{show tc}|] _Just
case f of
ObjLabel "halt" -> pure $ vm & #result ?~ args
_ -> do
rt <- getRoutine g f
let newFrame = MkFrame $ args ++ [f, ra]
pure $ vm
& jumpToRoutine rt
-- replace the active frame; don't push a new one.
& activeFrame .~ newFrame
-- it is not essential we clear the registers, but it'll make
-- bugs more obvious.
& #registers .~ mempty
stepT g vm (If c t f) = do
branch <- evalVal g vm c <&> \case
ObjImm (ImmBool False) -> f
_ -> t
pure $ jumpToBlock branch vm
stepP :: Jalmot :> es => Env -> VM -> Prim Val -> Eff es VM
stepP g vm p = traverse (evalVal g vm) p >>= \case
PrimZeroP x -> case x of
ObjImm (ImmInt n) -> ret1 . ObjImm . ImmBool $ n == 0
_ -> vmerror [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
PrimMakeClosure f env ->
case f of
ObjImm (ImmLabel l) -> ret1 . ObjHob $ HobClosure l env
_ -> vmerror [i|expected label, got #{f}|]
PrimEnvCode env ->
case env of
ObjHob (HobClosure l _) -> ret1 . ObjImm . ImmLabel $ l
_ -> vmerror [i|expected closure, got #{env}|]
PrimEnv -> do
x <- vm & expectOf "expected closure" (activeFrame . activeProcedure)
ret1 x
PrimEnvRef n -> do
(label,env) <- vm & expectOf "expected closure"
(activeFrame . activeProcedure . #_ObjHob . #_HobClosure)
x <- env & expectOf "expected upval" (ix n)
ret1 x
PrimCons x y -> ret1 $ ObjHob $ HobPair x y
PrimCar x -> case x of
ObjHob (HobPair car _) -> ret1 car
_ -> vmerror [i|expected pair, got ${x}|]
PrimCdr x -> case x of
ObjHob (HobPair _ cdr) -> ret1 cdr
_ -> vmerror [i|expected pair, got ${x}|]
x -> vmerror [i|unimplemented prim: #{p}|]
where
ret vs = pure $ vm & activeFrame . #locals <>:~ vs
ret1 v = ret [v]
arith_binop op (ObjImm (ImmInt x)) (ObjImm (ImmInt y)) =
ret1 $ ObjImm (ImmInt (op x y))
arith_binop _ x y = vmerror [i|bad arith: #{x}, #{y}|]
popFrame :: (HasCallStack, Jalmot :> es) => Stack -> Eff es (Frame, Stack)
popFrame stk = case stk ^. #frames . to NE.uncons of
(_, Nothing) -> vmerror "no frame to pop"
(f, Just fs) -> pure (f, stk & #frames .~ fs)
jumpToBlock :: Block -> VM -> VM
jumpToBlock b vm = vm
& #code .~ b.code
& #tail .~ b.tail
jumpToRoutine :: Routine -> VM -> VM
jumpToRoutine rt vm = vm
& jumpToBlock rt.start
& #debug . #activeRoutine .~ rt.label
getLabel :: Obj -> Maybe Label
getLabel = \case
ObjHob (HobClosure {label}) -> Just label
ObjImm (ImmLabel label) -> Just label
x -> Nothing
getRoutine :: (HasCallStack, Jalmot :> es) => Env -> Obj -> Eff es Routine
getRoutine g f = do
l <- getLabel f & expectOf [i|no label for #{f}|] _Just
case g ^. #labels . at l of
Just rt -> pure rt
Nothing -> vmerror [i|undefined label #{l}|]
expectOf
:: (HasCallStack, Jalmot :> es)
=> Text -> Getting (First a) s a -> s -> Eff es a
expectOf msg l = maybe (vmerror msg) pure . preview l
evalToLabel :: Jalmot :> es => Env -> VM -> Val -> Eff es Label
evalToLabel e vm v =
evalVal e vm v >>= \case
ObjImm (ImmLabel x) -> pure x
x -> vmerror [i|not a label: #{x}|]
evalVal :: Jalmot :> es => Env -> VM -> Val -> Eff es Obj
evalVal e vm = \case
ValImm imm -> pure $ ObjImm imm
ValReg r -> case vm ^. #registers . at r of
Just x -> pure x
Nothing -> vmerror [i|undefined register: #{r}|]
splitAtExact :: Int -> List a -> Maybe (List a, List a)
splitAtExact n xs = case compareLength xs n of
(EQ;GT) -> Just $ splitAt n xs
LT -> Nothing
takeExact :: Int -> List a -> Maybe (List a)
takeExact n xs = case compareLength xs n of
(EQ;GT) -> Just $ take n xs
LT -> Nothing
parseCall :: Int -> Frame -> Maybe (List Obj, Obj, Obj, Frame)
parseCall nargs frm = do
(xs,ys) <- splitAtExact (nargs+2) (frm ^. #locals)
let (xs',[f,ret]) = splitAt nargs xs
pure (xs',f,ret,MkFrame ys)
parseTailCall :: Int -> Frame -> Maybe (List Obj, Obj, Obj)
parseTailCall nargs frm = do
(xs,_) <- splitAtExact (nargs+1) (frm ^. #locals)
let (xs',f) = xs ^?! _Snoc
pure (xs',f,frm ^?! returnAddress)
initialVM :: VM
initialVM = MkVM
{ stack = MkStack . NE.singleton . MkFrame $
[ ObjLabel "start"
, ObjLabel "<nowhere at all>"
, ObjLabel "halt"
]
, tail = TailCall 0
, code = []
, registers = mempty
, stdout = ""
, result = Nothing
, debug = MkDebugVM
{ activeRoutine = "<nowhere>"
}
}
initialEnv :: Program -> Env
initialEnv p = MkEnv
{ labels = p.routines
}
loop :: (a -> Either b a) -> a -> b
loop f a = case f a of
Right a' -> loop f a'
Left b -> b
loopM :: Monad m => (a -> m (Either b a)) -> a -> m b
loopM f a = f a >>= \case
Right a' -> loopM f a'
Left b -> pure b
eval :: Jalmot :> es => Program -> Eff es (List Obj)
eval p = initialVM & loopM \vm -> case vm ^. #result of
Nothing -> Right <$> step (initialEnv p) vm
Just rs -> pure . Left $ rs
data Trace
= Step { vm :: VM, next :: Trace }
| StepToSuccess { vm :: VM, result :: List Obj }
| StepToFailure { vm :: VM, err :: AJalmotCS }
deriving (Show)
trace :: Program -> Trace
trace p = go (initialEnv p) initialVM
where
go g vm =
case vm.result of
Just rs -> StepToSuccess vm rs
Nothing ->
case runPureEff . runJalmotCS $ step g vm of
Left err -> StepToFailure vm err
Right vm' -> Step vm (go g vm')
writeObj :: Obj -> Text
writeObj = runJalmotUnsafe . S.encodeWith' S.datumIso
traceEval :: IOE :> es => Program -> Eff es ()
traceEval p = do
let t = trace p
liftIO . renderToFile "trace.html" . ppDoc p $ t
ppDoc :: Program -> Trace -> Html ()
ppDoc p t =
html_ do
head_ do
title_ "stackify trace"
style_ """
pre {
max-width: 95vw;
overflow: scroll;
}
table {
max-width: 95vw;
}
tbody > tr:nth-of-type(even) {
background-color: rgb(237 238 242);
}
.loc {
font-size: 0.8rem;
}
.syn-builtin, .syn-macro {
color: purple;
font-style: italic;
font-weight: bold;
}
.syn-constant {
color: olive;
}
.syn-procedure {
color: teal;
}
td pre {
display: inline
}
.syn-paren-0 { color: maroon; }
.syn-paren-1 { color: olive; }
.syn-paren-2 { color: green; }
.syn-paren-3 { color: navy; }
.syn-paren-4 { color: purple; }
.stack-frame
{ display: inline-flex
; flex-direction: row
; column-gap: 0.5em
}
"""
body_ do
details_ do
summary_ "stack code"
pre_ $ code_ do
htmlData . runJalmotUnsafe . S.toData S.dataIso $ p
ppTrace t
ppTrace :: Trace -> Html ()
ppTrace trace =
table_ do
thead_ $ tr_ do
traverse_ (th_ [scope_ "col"])
["routine","next instruction","stack frame"]
tbody_ do
go trace
where
go :: Trace -> Html ()
go = \case
Step vm next -> ppVM vm >> go next
StepToSuccess vm rs -> do
tr_ [class_ "trace-result"] do
td_ do
details_ do
summary_ "result"
pre_ do
code_ . toHtml . pShowNoColor $ vm
td_ [colspan_ "2"] do
sequence_ . intersperse " | " $ code_ . ppDatum <$> rs
StepToFailure vm err -> do
ppVM vm
tr_ [class_ "trace-failure"] do
td_ [colspan_ "3"] do
details_ do
summary_ "error"
pre_ do
samp_ do
fromString $ displayException err
ppVM :: VM -> Html ()
ppVM vm = do
tr_ do
td_ do
details_ do
summary_ do
var_ [class_ "loc"] do
vm ^. #debug . #activeRoutine . to ppDatum
pre_ do
code_ . toHtml . pShowNoColor $ vm
td_ do
code_ curi
td_ do
ppStack vm.stack
where
curi = vm ^?! failing (#code . _head . to ppDatum) (#tail . to ppDatum)
ppStack :: Stack -> Html ()
ppStack stk = do
span_ [class_ "stack"] do
stk ^.. each
& fmap ppFrame
& intersperse " | "
& sequence_
ppFrame :: Frame -> Html ()
ppFrame frm = do
span_ [class_ "stack-frame"] do
sequence_ $ frm ^.. #locals . each . to ppDatum
ppData :: S.DataIso a => a -> Html ()
ppData = htmlData . runJalmotUnsafe . S.toData S.dataIso
ppDatum :: S.DatumIso a => a -> Html ()
ppDatum = htmlDatum . runJalmotUnsafe . S.toDatum S.datumIso
fac (n :: Int) = [stkP|
(define $start
(push! $fac)
(push! #{n})
(tail-call 1))
(define $fac
(load %n 0)
(prim %x0 (zero? %n))
(if %x0
(then (push! 1)
(return 1))
(else (prim %x1 (- %n 1))
(push! $fac-c0)
(push! $fac)
(push! %x1)
(call 1))))
(define $fac-c0
(pop! %x2)
(pop! %n)
(prim %x3 (* %n %x2))
(push! %x3)
(return 1))
|]