+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+module Compiler where
+
+import Language
+import Serialise
+
+import qualified Data.Map as DM
+import Control.Monad.Writer
+import Control.Monad.State
+import Control.Monad.ST
+import Debug.Trace
+import Data.Array
+import Data.Array.ST
+import Data.Function
+
+newtype Compiler a = Compiler { unCompiler :: StateT CS (WriterT [Instr] (Either String)) a }
+ deriving
+ ( Functor
+ , Applicative
+ , Monad
+ , MonadWriter [Instr]
+ , MonadState CS
+ )
+instance MonadFail Compiler where fail s = Compiler $ lift $ lift $ Left s
+data CS = CS
+ { fresh :: [Int]
+ , functions :: DM.Map Int [Instr]
+ }
+
+runCompiler :: Compiler a -> Either String [Instr]
+runCompiler c = execWriterT
+ $ evalStateT (unCompiler (c >> instr [Halt] >> writeFunctions))
+ $ CS {fresh=[0..], functions=DM.empty}
+ where
+ writeFunctions :: Compiler ()
+ writeFunctions = gets (DM.elems . functions) >>= tell . concat
+
+instr :: [Instr] -> Compiler a
+instr i = tell i >> pure undefined
+
+freshLabel :: Compiler Int
+freshLabel = get >>= \cs->put (cs { fresh=tail (fresh cs) }) >> pure (head $ fresh cs)
+
+binop :: Instr -> Compiler a1 -> Compiler a2 -> Compiler b
+binop i l r = l >> r >> instr [i]
+
+unop :: Instr -> Compiler a -> Compiler b
+unop i l = l >> instr [i]
+
+instance Expression Compiler where
+ lit v = instr $ map Push $ serialise v []
+ (+.) = binop Add
+ (-.) = binop Sub
+ (/.) = binop Div
+ (*.) = binop Mul
+ (^.) = binop Pow
+-- (^.) l r = freshLabel >>= \lblstart->freshLabel >>= \lblend->
+-- l >> r >> instr -- pow (x, y) {
+-- [ Str 1
+-- , Str 0
+-- , Push 1 -- res = 1
+-- , Lbl lblstart -- while
+-- , Ldr 1 -- (y == 0)
+-- , Push 0 --
+-- , Neq --
+-- , Brf lblend --
+-- , Ldr 0 -- res *= x
+-- , Mul --
+-- , Ldr 1 -- y -= 1
+-- , Push 1 --
+-- , Sub --
+-- , Str 1 --
+-- , Bra lblstart --
+-- , Lbl lblend --
+-- ]
+ neg = unop Neg
+ (&.) = binop And
+ (|.) = binop Or
+ not = unop Not
+ (==.) = binop Eq
+ (/=.) = binop Neq
+ (<.) = binop Le
+ (>.) = binop Ge
+ (<=.) = binop Leq
+ (>=.) = binop Geq
+ if' p t e = freshLabel >>= \elselabel-> freshLabel >>= \endiflabel->
+ p >> instr [Brf elselabel] >>
+ t >> instr [Bra endiflabel, Lbl elselabel] >>
+ e >> instr [Lbl endiflabel]
+
+instance Function () Compiler where
+ fun def = Main $
+ freshLabel >>= \funlabel->
+ let g :- m = def (\()->instr [Jsr funlabel])
+ in liftFunction funlabel 0 (g ()) >> unmain m
+
+instance Function (Compiler a) Compiler where
+ fun def = Main $
+ freshLabel >>= \funlabel->
+ let g :- m = def (\a->a >> instr [Jsr funlabel])
+ in liftFunction funlabel 1 (g (instr [Arg 0])) >> unmain m
+
+instance Function (Compiler a, Compiler b) Compiler where
+ fun def = Main $
+ freshLabel >>= \funlabel->
+ let g :- m = def (\(a, b)->a >> b >> instr [Jsr funlabel])
+ in liftFunction funlabel 2 (g (instr [Arg 1], instr [Arg 0])) >> unmain m
+
+instance Function (Compiler a, Compiler b, Compiler c) Compiler where
+ fun def = Main $
+ freshLabel >>= \funlabel->
+ let g :- m = def (\(a, b, c)->a >> b >> c >> instr [Jsr funlabel])
+ in liftFunction funlabel 3 (g (instr [Arg 2], instr [Arg 1], instr [Arg 0])) >> unmain m
+
+liftFunction :: Int -> Int -> Compiler a -> Compiler ()
+liftFunction lbl nargs body = do
+ is <- snd <$> censor (\_->[]) (listen body)
+ let instructions = Lbl lbl : is ++ [Ret nargs]
+ modify (\s->s { functions=DM.insert lbl instructions $ functions s })
+
+data Instr
+ = Push Int | Pop Int | Dup | Roll Int Int
+ | Add | Sub | Mul | Div | Neg | Pow
+ | And | Or | Not
+ | Eq | Neq | Le | Ge | Leq | Geq
+ | Lbl Int | Bra Int | Brf Int
+ | Str Int | Ldr Int
+ | Sth Int | Ldh Int
+ | Jsr Int | Ret Int | Arg Int
+ | Halt
+ deriving Show
+
+data Registers = Registers
+ { pc :: Int
+ , hp :: Int
+ , sp :: Int
+ , mp :: Int
+ , gp :: DM.Map Int Int
+ }
+ deriving Show
+
+interpret :: Int -> [Instr] -> Array Int Int
+interpret memsize prog = runSTArray $ do
+ program <- newListArray (0, length prog) prog
+ mem <- newArray (0, memsize-1) 0
+ int program mem (Registers {pc=0, sp=memsize-1, mp=0, hp=0, gp=DM.empty})
+ where
+ pushh :: STArray s Int Int -> Int -> Registers -> ST s Registers
+ pushh memory value reg = do
+ writeArray memory (hp reg) value
+ pure (reg { hp = hp reg + 1} )
+
+ loadh :: STArray s Int Int -> Int -> Registers -> ST s Registers
+ loadh memory hptr registers = readArray memory hptr >>= flip (push memory) registers
+
+ push :: STArray s Int Int -> Int -> Registers -> ST s Registers
+ push memory value reg = do
+ writeArray memory (sp reg) value
+ pure (reg { sp = sp reg - 1} )
+
+ pop :: STArray s Int Int -> Registers -> ST s (Registers, Int)
+ pop memory reg = do
+ v <- readArray memory (sp reg + 1)
+ pure (reg { sp = sp reg + 1}, v)
+
+ popn :: STArray s Int Int -> Int -> Registers -> ST s (Registers, [Int])
+ popn _ 0 reg = pure (reg, [])
+ popn memory n reg = do
+ (reg', v) <- pop memory reg
+ (reg'', vs) <- popn memory (n - 1) reg'
+ pure (reg'', v:vs)
+
+ bop :: (Int -> Int -> Int) -> STArray s Int Int -> Registers -> ST s Registers
+ bop op memory reg = do
+ (reg1, r) <- pop memory reg
+ uop (flip op r) memory reg1
+
+ uop :: (Int -> Int) -> STArray s Int Int -> Registers -> ST s Registers
+ uop op memory reg = do
+ (reg1, r) <- pop memory reg
+ push memory (op r) reg1
+
+ int :: STArray s Int Instr -> STArray s Int Int -> Registers -> ST s (STArray s Int Int)
+ int program memory registers = do
+ instruction <- readArray program $ pc registers
+ stack <- getElems memory
+ let reg = registers { pc = pc registers + 1 }
+ case trace ("Interpret: " ++ show instruction ++ " with registers: " ++ show registers ++ " and stack: " ++ show stack) instruction of
+-- case instruction of
+ Str r -> do
+ (reg', v) <- pop memory reg
+ int program memory $ reg' { gp = DM.insert r v (gp reg')}
+ Ldr r -> push memory (DM.findWithDefault 0 r $ gp reg) reg >>= int program memory
+ Roll 0 _ -> int program memory reg
+ Roll 1 _ -> int program memory reg
+ Roll _ 0 -> int program memory reg
+ Roll depth num -> do
+ (reg', vs) <- popn memory depth reg
+ foldM (flip $ push memory) reg' (roll num [] $ reverse vs) >>= int program memory
+ where
+ roll 0 acc vs = vs ++ reverse acc
+ roll n acc [] = roll n [] $ reverse acc
+ roll n acc (v:vs) = roll (n-1) (v:acc) vs
+ Pop n -> popn memory n reg >>= int program memory . fst
+ Push v -> push memory v reg >>= int program memory
+ Dup -> pop memory reg >>= \(r', v)->push memory v r' >>= push memory v >>= int program memory
+ Add -> bop (+) memory reg >>= int program memory
+ Sub -> bop (-) memory reg >>= int program memory
+ Mul -> bop (*) memory reg >>= int program memory
+ Div -> bop div memory reg >>= int program memory
+ Neg -> uop negate memory reg >>= int program memory
+ Pow -> bop (^) memory reg >>= int program memory
+ And -> bop ((b2i .) . on (&&) i2b) memory reg >>= int program memory
+ Or -> bop ((b2i .) . on (||) i2b) memory reg >>= int program memory
+ Not -> uop (b2i . Prelude.not . i2b) memory reg >>= int program memory
+ Eq -> bop ((b2i .) . (==)) memory reg >>= int program memory
+ Neq -> bop ((b2i .) . (/=)) memory reg >>= int program memory
+ Le -> bop ((b2i .) . (<)) memory reg >>= int program memory
+ Ge -> bop ((b2i .) . (>)) memory reg >>= int program memory
+ Leq -> bop ((b2i .) . (<=)) memory reg >>= int program memory
+ Geq -> bop ((b2i .) . (>=)) memory reg >>= int program memory
+ Lbl _ -> int program memory reg
+ Bra l -> branch l program reg >>= int program memory
+ Brf l -> do
+ (reg', v) <- pop memory reg
+ reg'' <- if i2b v then pure reg' else branch l program reg'
+ int program memory reg''
+ Sth n ->
+ popn memory n reg
+ >>= uncurry (foldM $ flip $ pushh memory)
+ >>= push memory (hp reg + n - 1)
+ >>= int program memory
+ Ldh n -> pop memory reg >>= \(reg', hptr)->loadh memory (hptr - n - 1) reg'
+ >>= int program memory
+ Jsr i -> push memory (pc reg) reg
+ >>= push memory (mp reg)
+ >>= branch i program
+ >>= \r->int program memory (r { mp = sp r})
+ Ret n -> do
+ (reg1, rval) <- pop memory reg
+ (reg2, omp) <- pop memory reg1
+ (reg3, ra) <- pop memory reg2
+ (reg4, _) <- popn memory n reg3
+ reg5 <- push memory rval reg4
+ int program memory $ reg5 { pc=ra, mp=omp }
+ Arg n -> do
+ v <- readArray memory (mp reg + 3 + n)
+ push memory v reg >>= int program memory
+ Halt -> pure memory
+
+ branch :: Int -> STArray s Int Instr -> Registers -> ST s Registers
+ branch label program reg = case pc reg of
+ -1 -> getBounds program >>= \(_, m)->branch label program $ reg { pc = m - 1}
+ _ -> readArray program (pc reg) >>= \case
+ Lbl l | label == l -> pure $ reg
+ _ -> branch label program $ reg { pc = pc reg - 1 }
+
+b2i :: Bool -> Int
+b2i True = 1
+b2i False = 0
+
+i2b :: Int -> Bool
+i2b 0 = False
+i2b _ = True