{-# LANGUAGE DeriveFunctor, GeneralizedNewtypeDeriving #-}

module Compiler(compile) where

import Control.Monad.State.Strict
import Control.Monad.Writer.Strict
import Data.List

import AST


-- Execution starts in the middle of the tape
tapeLength :: Offset
tapeLength = 60000


newtype LineWriter a = LineWriter {unLineWriter :: WriterT [String] (State Int) a}
  deriving (Functor, Applicative, Monad, MonadWriter [String], MonadState Int)

runLineWriter :: LineWriter a -> String
runLineWriter = (++ "\n") . intercalate "\n" . flip evalState 1 . execWriterT . unLineWriter

emit0 :: String -> LineWriter ()
emit0 = tell . pure

emit :: String -> LineWriter ()
emit = tell . pure . ('\t' :)

genId :: LineWriter Int
genId = state $ \i -> (i, i + 1)


compile :: Program -> String
compile (Program inss) = runLineWriter $ do
    emit0 prologue
    mapM_ compileIns inss
    emit0 epilogue

cursorReg :: String
cursorReg = "rbx"

cursorOffset :: Offset -> String
cursorOffset 0 = cursorReg
cursorOffset off = cursorReg ++ " + " ++ show off

compileIns :: Instruction -> LineWriter ()
compileIns (IAdd v off) = emit $ "add byte [" ++ cursorOffset off ++ "], " ++ show v
compileIns (ISet v off) = emit $ "mov byte [" ++ cursorOffset off ++ "], " ++ show v
compileIns (ICopy _ _ 0) = return ()
compileIns (ICopy from to m) = do
    emit $ "mov al, [" ++ cursorOffset from ++ "]"
    case m of
        1 -> return ()
        _ | Just p <- isTwoPower m -> emit $ "shl al, " ++ show p
          | otherwise -> do
                             emit $ "mov cl, " ++ show m
                             emit "mul byte cl"
    emit $ "add [" ++ cursorOffset to ++ "], al"
compileIns (ISlide off) = emit $ "add " ++ cursorReg ++ ", " ++ show off
compileIns (ILoop inss off) = do
    loopid <- genId
    emit $ "cmp byte [" ++ cursorOffset off ++ "], 0"
    emit $ "jz .Laf_" ++ show loopid
    emit0 $ ".Lbd_" ++ show loopid ++ ":"
    mapM_ compileIns inss
    emit $ "cmp byte [" ++ cursorOffset off ++ "], 0"
    emit $ "jnz .Lbd_" ++ show loopid
    emit0 $ ".Laf_" ++ show loopid ++ ":"
compileIns (IInput off) = do
    emit "call _getchar"
    emit $ "mov [" ++ cursorOffset off ++ "], al"
compileIns (IOutput off) = do
    emit "xor edi, edi"
    emit $ "mov dil, [" ++ cursorOffset off ++ "]"
    emit "call _putchar"
compileIns IStart = return ()

isTwoPower :: Byte -> Maybe Int
isTwoPower v = findIndex (==v) (take 8 $ iterate (* 2) 1)


prologue :: String
prologue =
         "global _main"
    ++ "\nextern _calloc, _free, _putchar, _getchar, _write"
    ++ "\ndefault rel"
    ++ "\nsection .text"

    ++ "\n_main:"
    ++ "\n\tpush rbp"
    ++ "\n\tmov rbp, rsp"
    ++ "\n\tpush rbx"
    ++ "\n\tsub rsp, 8"
    ++ "\n\tmov edi, " ++ show tapeLength
    ++ "\n\tmov esi, 1"
    ++ "\n\tcall _calloc"
    ++ "\n\ttest rax, rax"
    ++ "\n\tjz .allocation_failure"
    ++ "\n\tlea rbx, [rax + " ++ show (tapeLength `div` 2) ++ "]  ; rbx = cursor"
    ++ "\n\tmov [rsp], rax  ; [rsp] = start of tape buffer"
    ++ "\n"

epilogue :: String
epilogue =
       "\n\tmov rdi, [rsp]"
    ++ "\n\tcall _free"
    ++ "\n\txor eax, eax"
    ++ "\n.return:"
    ++ "\n\tadd rsp, 8"
    ++ "\n\tpop rbx"
    ++ "\n\tmov rsp, rbp"
    ++ "\n\tpop rbp"
    ++ "\n\tret"
    ++ "\n.allocation_failure:"
    ++ "\n\tsub rsp, 32"
    ++ "\n\tmov rax, 0x697461636f6c6c41  ; \"Allocati\""
    ++ "\n\tmov rbx, 0x756c696166206e6f  ; \"on failu\""
    ++ "\n\tmov ecx, 0x000a6572          ; \"re\\n\\0\""
    ++ "\n\tmov [rsp], rax"
    ++ "\n\tmov [rsp+8], rbx"
    ++ "\n\tmov [rsp+16], ecx"
    ++ "\n\tmov edi, 2"
    ++ "\n\tmov rsi, rsp"
    ++ "\n\tmov edx, 20"
    ++ "\n\tcall _write"
    ++ "\n\tadd rsp, 32"
    ++ "\n\tmov eax, 1"
    ++ "\n\tjmp .return"