Script serializer

Main.hs

module Main where

import qualified Data.Map as M
import Text.Printf (printf)
import Control.Monad (ap)

data Script a = Script { genScript :: Int -> (Int, [Operation], a) }

instance Functor Script where
  fmap f (Script gen) = Script $ fmap f . gen

instance Applicative Script where
  pure = return
  (<*>) = ap

instance Monad Script where
  return x = Script $ \regs -> (regs, [], x)
  s >>= f = Script genRet where
    genRet regs = (regs'', ops1 ++ ops2, val2) where
      (regs', ops1, val1) = genScript s regs
      (regs'', ops2, val2) = genScript (f val1) regs'

data Operation
  = OGetArg  Int String
  | OReturn  Int
  | OIntPure Int Integer
  | OIntAdd  Int Int Int
  | OIntMul  Int Int Int

instance Show Operation where
  show (OGetArg  o s  ) = printf "%%%d <- arg '%s'" o s
  show (OReturn  i    ) = printf "return %%%d" i
  show (OIntPure o x  ) = printf "%%%d <- $%d" o x
  show (OIntAdd  o a b) = printf "%%%d <- %%%d + %%%d" o a b
  show (OIntMul  o a b) = printf "%%%d <- %%%d * %%%d" o a b

class SType a where
  var :: a -> Script a

data SInt
  = SPure Integer
  | SRef  Int
  | SAdd  SInt SInt
  | SMul  SInt SInt

instance SType SInt where
  var (SPure  x) = Script $ \regs -> (regs + 1, [OIntPure regs x], SRef regs)
  var (SRef   x) = pure $ SRef x
  var (SAdd a b) = do
    a' <- var a
    b' <- var b
    case (a', b') of
      (SRef a'', SRef b'') -> Script $ \regs -> (regs + 1, [OIntAdd regs a'' b''], SRef regs)
      _ -> undefined
  var (SMul a b) = do
    a' <- var a
    b' <- var b
    case (a', b') of
      (SRef a'', SRef b'') -> Script $ \regs -> (regs + 1, [OIntMul regs a'' b''], SRef regs)
      _ -> undefined

getArg :: String -> Script SInt
getArg arg = Script $ \regs -> (regs + 1, [OGetArg regs arg], SRef regs)

data SerializedScript = SerializedScript { regs :: Int, ops :: [Operation] }

instance Show SerializedScript where
  show (SerializedScript regs ops) =
    printf "; registers: %d\n%s" regs (concat $ map (\s -> show s ++ "\n") ops)

serializeScript :: Script SInt -> SerializedScript
serializeScript s = SerializedScript regs ops' where
  (regs, ops, SRef v) = genScript (s >>= var) 0
  ops' = ops ++ [OReturn v]

script :: Script SInt
script = do
  a <- getArg "foo"
  b <- getArg "bar"
  -- c <- var $ a * 2 + b + 1 -- If using RebindableSyntax.
  c <- var $ SAdd (SAdd (SMul a (SPure 2)) b) (SPure 1)
  pure $ SAdd c c

main :: IO ()
main = do
  let ser = serializeScript script
  -- runScript ser (M.fromList [("foo", 20), ("bar", 1)])
  print ser