Macro.hs

module Zwirn.Language.Macro where

import Control.Monad.State (StateT, lift, modify, runStateT)
import qualified Data.Map as Map
import qualified Data.Text as T
import Zwirn.Language.Location (Located (..), RealSrcLoc (..), SrcLoc (..))
import Zwirn.Language.Pretty (ppterm)
import Zwirn.Language.Syntax (LocTerm, Term (..))

type MacroMap = Map.Map T.Text LocTerm

-- represents the code edit that replaces the text at editPos with editText
data CodeEdit = CodeEdit
  { editPos :: RealSrcLoc,
    editText :: T.Text
  }
  deriving (Eq, Show)

-- a simple monad that accumulates code edits
type MacroMonad = StateT [CodeEdit] Maybe

defaultMacroMap :: MacroMap
defaultMacroMap = Map.fromList [("num", Located (SrcLoc (RealSrcLoc "internal" 0 0 0 1)) (TNum "1"))]

runMacros :: LocTerm -> MacroMap -> Maybe (LocTerm, [CodeEdit])
runMacros t mmap = runStateT (substituteMacros mmap t) []

substituteMacros :: MacroMap -> LocTerm -> MacroMonad LocTerm
substituteMacros _ t@(Located _ (TNum _)) = return t
substituteMacros _ t@(Located _ (TText _)) = return t
substituteMacros _ t@(Located _ (TVar _)) = return t
substituteMacros ms (Located (SrcLoc oldPos@(RealSrcLoc doc stln stch _ _)) (TMacro name)) = do
  (Located p m) <- lift $ Map.lookup name ms
  let replace = ppterm m
  case p of
    NoLoc -> lift Nothing
    SrcLoc (RealSrcLoc _ stln' _ enln' ench') -> do
      let newPos = RealSrcLoc doc stln stch (stln + (enln' - stln')) ench'
          edit = CodeEdit oldPos replace
      modify (edit :)
      return $ Located (SrcLoc newPos) m
substituteMacros _ t@(Located _ TRest) = return t
substituteMacros ms (Located p (TBracket t)) = do
  s <- substituteMacros ms t
  return $ Located p (TBracket s)
substituteMacros ms (Located p (TRepeat t i)) = do
  s <- substituteMacros ms t
  return $ Located p (TRepeat s i)
substituteMacros ms (Located p (TLambda x t)) = do
  s <- substituteMacros ms t
  return $ Located p (TLambda x s)
substituteMacros ms (Located p (TSectionL t x)) = do
  s <- substituteMacros ms t
  return $ Located p (TSectionL s x)
substituteMacros ms (Located p (TSectionR x t)) = do
  s <- substituteMacros ms t
  return $ Located p (TSectionR x s)
substituteMacros ms (Located p (TSeq ts)) = do
  ss <- mapM (substituteMacros ms) ts
  return $ Located p (TSeq ss)
substituteMacros ms (Located p (TStack ts)) = do
  ss <- mapM (substituteMacros ms) ts
  return $ Located p (TStack ss)
substituteMacros ms (Located p (TAlt ts)) = do
  ss <- mapM (substituteMacros ms) ts
  return $ Located p (TAlt ss)
substituteMacros ms (Located p (TChoice i ts)) = do
  ss <- mapM (substituteMacros ms) ts
  return $ Located p (TChoice i ss)
substituteMacros ms (Located p (TPoly t1 t2)) = do
  s1 <- substituteMacros ms t1
  s2 <- substituteMacros ms t2
  return $ Located p (TPoly s1 s2)
substituteMacros ms (Located p (TApp t1 t2)) = do
  s1 <- substituteMacros ms t1
  s2 <- substituteMacros ms t2
  return $ Located p (TApp s1 s2)
substituteMacros ms (Located p (TInfix t1 n t2)) = do
  s1 <- substituteMacros ms t1
  s2 <- substituteMacros ms t2
  return $ Located p (TInfix s1 n s2)
substituteMacros ms (Located p (TIfThenElse t1 t2 Nothing)) = do
  s1 <- substituteMacros ms t1
  s2 <- substituteMacros ms t2
  return $ Located p (TIfThenElse s1 s2 Nothing)
substituteMacros ms (Located p (TEnum k t1 t2)) = do
  s1 <- substituteMacros ms t1
  s2 <- substituteMacros ms t2
  return $ Located p (TEnum k s1 s2)
substituteMacros ms (Located p (TEnumThen k t1 t2 t3)) = do
  s1 <- substituteMacros ms t1
  s2 <- substituteMacros ms t2
  s3 <- substituteMacros ms t3
  return $ Located p (TEnumThen k s1 s2 s3)
substituteMacros ms (Located p (TIfThenElse t1 t2 (Just t3))) = do
  s1 <- substituteMacros ms t1
  s2 <- substituteMacros ms t2
  s3 <- substituteMacros ms t3
  return $ Located p (TIfThenElse s1 s2 (Just s3))
substituteMacros _ _ = lift Nothing