Macro AST plus application

gistfile1.hs

{-# LANGUAGE OverloadedStrings #-}

import Control.Monad (guard)
import Data.ByteString.Char8 ()
import qualified Data.ByteString as B
import qualified Data.Map as M

type MacroMap = M.Map B.ByteString Macro

type ArgMap = M.Map B.ByteString B.ByteString

data MacroChunk =
    StringChunk B.ByteString
  | VariableChunk B.ByteString
  | CallChunk B.ByteString [MacroChunk]
  deriving (Show, Eq)

data Macro = Macro {
  macroName   :: B.ByteString,
  macroArgs   :: [B.ByteString],
  macroChunks :: [MacroChunk]
}
  deriving (Show, Eq)

callMacro :: MacroMap -> Macro -> [B.ByteString] -> Maybe B.ByteString
callMacro macros macro args =
  B.concat `fmap` (applyMacro macros macro args >>= mapM strings)

applyMacro :: MacroMap -> Macro -> [B.ByteString] -> Maybe [MacroChunk]
applyMacro macros (Macro _ mArgs chunks) args =
  concat `fmap` (mapM (substArgs argMap) chunks >>= mapM (substCalls macros))
  where
    argMap = M.fromList $ zip mArgs args

substArgs :: ArgMap -> MacroChunk -> Maybe MacroChunk
substArgs argMap (VariableChunk name) =
  StringChunk `fmap` M.lookup name argMap
substArgs argMap (CallChunk name chunks) =
  CallChunk name `fmap` mapM (substArgs argMap) chunks
substArgs _      other = return other

substCalls :: MacroMap -> MacroChunk -> Maybe [MacroChunk]
substCalls _   ch@(StringChunk _) = return [ch]
substCalls macros (CallChunk name args) = do
  macro <- M.lookup name macros
  guard $ length (macroArgs macro) == length args
  strArgs <- mapM strings args
  applyMacro macros macro strArgs
substCalls _      _ =
 -- Variable chunks should've been substituted
  Nothing

strings :: MacroChunk -> Maybe B.ByteString
strings (StringChunk s) = return s
strings _               = Nothing


-- Example macros

testMacro :: Macro
testMacro = Macro "test" ["a", "b"] [StringChunk "foo", VariableChunk "b"]

testMacro2 :: Macro
testMacro2 = Macro "test2" ["a"] [StringChunk "bar", CallChunk "test" [StringChunk "10", VariableChunk "a"]]

testMacros :: MacroMap
testMacros = M.fromList [("test", testMacro), ("test2", testMacro2)]