Rl.hs

{-# LANGUAGE TemplateHaskell #-}

module Rl where

import Control.Monad.Trans.State
import Control.Lens
import Text.Read (readMaybe)
import qualified Data.Map as M
import Control.Monad.IO.Class
import Control.Applicative (liftA2)

type VM a = StateT World IO a

data World = World
    { _stack :: [Int]
    , _dict :: M.Map String (VM ())
    }

makeLenses ''World

pop :: VM Int
pop = get >>= \w -> case w ^. stack of
    []     -> fail "stack underflow"
    (x:xs) -> do
        put $ stack .~ xs $ w
        return x

push :: Int -> VM ()
push x = modify (stack %~ (x:))

run :: [String] -> VM ()
run []     = return ()
run (w:ws) = do
    d <- gets (^. dict)
    case readMaybe w of
        Just x -> push x
        Nothing -> word w
    run ws

word :: String -> VM ()
word w = gets (^. dict) >>= \d -> case M.lookup w d of
    Nothing -> fail $ "undefined word \"" ++ w ++ "\""
    Just x  -> x

world :: World
world = World
    { _stack = []
    , _dict  = library
    }

library :: M.Map String (VM ())
library = M.fromList
    [ "pick" ~> do
        (x:xs) <- gets (^. stack)
        modify (stack .~ (xs !! x):xs)
    , "roll" ~> do
        x <- pop
        (xs, y:ys) <- fmap (splitAt x) $ gets (^. stack)
        modify (stack .~ (y:xs ++ ys))
    , "dup"  ~> go "0 pick"
    , "2dup" ~> go "dup dup"
    , "swap" ~> go "1 roll"
    , "over" ~> go "1 pick"
    , "rot"  ~> go "2 roll"
    , "tuck" ~> go "1 roll 1 pick"
    , "+"    ~> op (+)
    , "-"    ~> op (-)
    , "*"    ~> op (*)
    , "/"    ~> op div
    , "."    ~> (pop >>= liftIO . print)
    ] where op f = liftA2 (flip f) pop pop >>= push
            go = run . words
            (~>) = (,)