zorth/zorth.hs

{-# LANGUAGE LambdaCase #-}
import qualified Control.Applicative as Ap
import Control.Applicative ((<|>), Alternative, many, some)
import Data.Char (isDigit, digitToInt)
import Control.Monad 
import Text.Printf (printf)
newtype Parser a = Parser {runParser :: String -> [(a, String)]}

instance Functor Parser where
    fmap f p = Parser $ \inp -> [ (f x,xs) | (x,xs) <- runParser p inp ]

instance Applicative Parser where
    pure v = Parser $ \inp -> [(v,inp)]
    d <*> e = d >>= (<$> e)

instance Monad Parser where
    return = pure
    p >>= f = Parser $ \inp -> concat [runParser (f v) out | (v,out) <- runParser p inp]

instance Alternative Parser where
    empty = Parser $ const []
    p <|> q = Parser $ \inp -> runParser p inp ++ runParser q inp
    many p = do { x <- p; xs <- many p; return $ x:xs } <|> return []
    some p = do
        x <- p
        xs <- many p
        return $ x:xs

instance MonadPlus Parser

instance MonadFail Parser where
    fail _ = mzero

(<++) :: Parser a -> Parser a -> Parser a
p <++ q = Parser $ \inp -> let r1 = runParser p inp in
    if null r1 then runParser q inp else r1

get :: Parser Char 
get = Parser $ \case
    [] -> [] 
    x:xs -> [(x,xs)]


look :: Parser Char 
look = Parser $ \case
    [] -> []
    x:xs -> [(x,x:xs)]

sat :: (Char -> Bool) -> Parser Char 
sat p = do
    c <- get
    if p c then return c else mzero

char :: Char -> Parser Char
char c = sat (== c)

string :: String -> Parser String
string "" = return "" 
string (x:xs) = do
    char x
    string xs
    return (x:xs)

eof :: Parser () 
eof = Parser $ \inp -> if null inp then [((),inp)] else [] 


data ZorthAST = ZorthASTInteger Int
              | ZorthASTWord String
              deriving Show
skipNonsenseSymbols :: Parser [Char] 
skipNonsenseSymbols = many (char ' ' <|> char '\n')

digit :: Parser Int
digit = digitToInt <$> sat isDigit 

option :: a -> Parser a -> Parser a
option x p = p <|> return x

pZorthUnsignedInteger :: Parser ZorthAST
pZorthUnsignedInteger = ZorthASTInteger .  foldr (\n t -> n + t*10) 0 . reverse <$> some digit

pZorthSignedInteger :: Parser ZorthAST
pZorthSignedInteger = do
    char '-'
    (ZorthASTInteger i) <- pZorthUnsignedInteger
    return $ ZorthASTInteger $ negate i

pZorthInteger :: Parser ZorthAST
pZorthInteger = do
    skipNonsenseSymbols
    i <- pZorthSignedInteger <|> pZorthUnsignedInteger 
    eof <|> ((sat (\x -> x== ' ' || x == '\n')) >> return ())
    return i

pZorthWord :: Parser ZorthAST
pZorthWord = do
    skipNonsenseSymbols
    ZorthASTWord <$> some (sat (\x -> x /= ' ' && x /= '\n'))

pZorth :: Parser [ZorthAST]
pZorth = some (pZorthWord <|> pZorthInteger)

parseZorth :: String -> [ZorthAST]
parseZorth = fst . head . runParser pZorth

handleSymbol :: ZorthAST -> IO ()
handleSymbol (ZorthASTInteger i) = void $ printf "    push %d\n" i
handleSymbol (ZorthASTWord "+") = do
    putStrLn "    pop rbx"
    putStrLn "    pop rax"
    putStrLn "    add rax, rbx"
    putStrLn "    push rax"
handleSymbol (ZorthASTWord "-") = do
    putStrLn "    pop rbx"
    putStrLn "    pop rax"
    putStrLn "    sub rax, rbx"
    putStrLn "    push rax"
handleSymbol (ZorthASTWord "ret") = do
    putStrLn "    mov rax,60"
    putStrLn "    mov rdi,0"
    putStrLn "    syscall"
handleSymbol (ZorthASTWord "dup") = do
    putStrLn "    pop rax"
    putStrLn "    push rax"
    putStrLn "    push rax"
handleSymbol (ZorthASTWord "swap") = do
    putStrLn "    pop rax"
    putStrLn "    pop rbx"
    putStrLn "    push rax"
    putStrLn "    push rbx"
handleSymbol (ZorthASTWord "drop") = do
    putStrLn "    add rsp, 8"
handleSymbol (ZorthASTWord ".") = do
    putStrLn "    pop rbx"
    putStrLn "    call print_number"

compileZorth :: [ZorthAST] -> IO ()
compileZorth [] = return ()
compileZorth xs = do
    putStrLn "global _start"
    putStrLn "print_number:"
    putStrLn "    push rbp"
    putStrLn "    mov rbp, rsp"
    putStrLn "    sub rsp, 128 "
    putStrLn "    mov rdx, -1 "
    putStrLn "    jmp l2"
    putStrLn "    add rsp, 28"
    putStrLn "    pop rbp"
    putStrLn "    ret"
    putStrLn "l1:"
    putStrLn "    dec  rdx"
    putStrLn "    imul rax, rbx, 1717986919"
    putStrLn "    shr  rax, 34"
    putStrLn "    imul rcx, rax, 10"
    putStrLn "    sub  rbx, rcx"
    putStrLn "    add  rbx, '0'"
    putStrLn "    mov  qword[rbp+8*rdx], rbx"
    putStrLn ""
    putStrLn "    mov rbx, rax"
    putStrLn "l2:"
    putStrLn "    cmp rbx, 0"
    putStrLn "    jne l1"
    putStrLn "    mov qword[rbp-8], `\\n` "
    putStrLn "    mov rax, 1"
    putStrLn "    mov rdi, 1"
    putStrLn "    lea rsi, [rbp+8*rdx]"
    putStrLn "    neg rdx"
    putStrLn "    imul rdx, 8"
    putStrLn "    syscall"
    putStrLn "    leave"
    putStrLn "    ret"
    putStrLn "_start:"
    foldr ((>>) . handleSymbol) (return ()) xs

main :: IO ()
main = do
    contents <- getContents
    compileZorth $ parseZorth contents