LaTeXParser.hs

{-# OPTIONS_GHC -fno-warn-tabs #-}
{-# LANGUAGE OverloadedStrings, RecordWildCards, ViewPatterns, TupleSections #-}

module LaTeXParser (parseString,
	Token(Token), Context(..), defaultContext, Signature(..), Macros(..), Environment(..), Command(..), ParseResult(ParseResult),
	defaultMacros,
	nullCmd, storeCmd, codeEnv, normalCmd,
	storeEnv) where

import LaTeXBase (LaTeXUnit(..), LaTeX, TeXArg, ArgKind(..), MathType(..), concatRaws)
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Char (isAlphaNum, isSpace, isAlpha)
import Data.Maybe (fromJust)
import Control.Arrow (first)
import Data.Map (Map)
import qualified Data.Map as Map
import Prelude hiding ((++), (.))
import Util ((.), (++), getDigit, stripInfix)

newtype Token = Token { tokenChars :: String }
	deriving (Eq, Show)

data Environment = Environment (Context -> [Token] -> ParseResult)
data Command = Command { runCommand :: Context -> String {- ws -} -> [Token] -> ParseResult }

data Macros = Macros
	{ commands :: Map Text Command
	, environments :: Map Text Environment
	, counters :: Map Text Int }

newCommand :: Bool {- overwrite -} -> (Text, Command) -> Macros -> Macros
newCommand True (name, cmd) Macros{..} = Macros{commands = Map.insert name cmd commands, ..}
newCommand False (name, cmd) Macros{..} = Macros{commands = Map.insertWith (\_ y -> y) name cmd commands, ..}

instance Semigroup Macros where
	x <> y = Macros
		(commands x ++ commands y)
		(environments x ++ environments y)
		(counters x ++ counters y)

instance Monoid Macros where
	mempty = Macros mempty mempty mempty

data ParseResult = ParseResult
	{ content :: LaTeX
	, newMacros :: Macros
	, remainder :: [Token] }

data Signature = Signature
	{ nrFixArgs :: Int
	, defaultArg :: Maybe [Token] }
	deriving Show

data Context = Context
	{ commentsEnabled :: Bool
	, parsingOptArg :: Bool
	, macros :: Macros }

prependContent :: LaTeX -> ParseResult -> ParseResult
prependContent t p = p{content = t ++ content p}

combineMacros :: Bool {- left biased -} -> Macros -> Macros -> Macros
combineMacros b x y = if b then x ++ y else y ++ x

addMacros :: Bool {- overwrite -} -> Macros -> ParseResult -> ParseResult
addMacros b m p = p{newMacros = combineMacros b m (newMacros p)}

defaultEnvs :: [(Text, Environment)]
defaultEnvs = [outputblockEnv]

codeEnv :: Text -> Signature -> (Text, Environment)
codeEnv name sig = (name, Environment f)
	where
		f :: Context -> [Token] -> ParseResult
		f ctx toks = ParseResult [env] mempty rest'
			where
				(arguments, rest) = parseArgs sig toks
				Just (code, rest') = stripInfix [Token "\\end", Token "{", Token (Text.unpack name), Token "}"] rest
				env = TeXEnv (Text.unpack name) (map ((FixArg, ) . fullParse ctx) arguments) (parseCode name ctx code)

outputblockEnv :: (Text, Environment)
outputblockEnv = ("outputblock", Environment f)
	where
		f :: Context -> [Token] -> ParseResult
		f ctx toks = ParseResult [env] mempty rest
			where
				Just (content, rest) = stripInfix [Token "\\end", Token "{", Token "outputblock", Token "}"] toks
				env = TeXEnv "outputblock" [] (parseOutputBlock ctx content)

parseOutputBlock :: Context -> [Token] -> LaTeX
parseOutputBlock c = concatRaws . go
	where
		go :: [Token] -> LaTeX
		go [] = []
		go (Token "@" : rest) = fullParse c cmd ++ go rest'
			where (cmd, Token "@" : rest') = break (== Token "@") rest
		go s = TeXRaw (Text.pack $ concatMap tokenChars code) : go rest
			where (code, rest) = break (== Token "@") s

storeEnv :: String -> Signature -> (Text, Environment)
storeEnv name sig = (Text.pack name, Environment act)
	where
		act :: Context -> [Token] -> ParseResult
		act ctx toks = ParseResult [env] mempty afterend
			where
				(arguments, rest) = parseArgs sig toks
				ParseResult body _ afterend = parse ctx rest
				env = TeXEnv name (map ((FixArg, ) . fullParse ctx) arguments) (concatRaws body)
							-- todo: not all fixargs

defaultCmds :: [(Text, Command)]
defaultCmds =
	[ ("newcommand", newCommandCommand)
	, ("renewcommand", newCommandCommand)
	, ("DeclareMathOperator", declareMathOperator)
	, ("newcolumntype", newColumnTypeCommand)
	, ("newenvironment", newEnvCommand)
	, ("lstnewenvironment", newEnvCommand)
	, ("raisebox", raiseBoxCommand)
	, ("let", Command $ \ctx _ws rest -> parse ctx (drop 2 rest))
	, beginCommand
	, endCommand
	, oldDefCommand
	]

oldDefCommand :: (Text, Command)
oldDefCommand = ("def", Command pars)
	where
		pars ctx@Context{..} _ws rest
			| (Token ('\\' : name) : rest') <- rest
			, Just (body, rest'') <- balanced ('{', '}') rest' =
				let
					m = Macros (Map.fromList [defCmd (Text.pack name) (Signature 0 Nothing) body]) mempty mempty
					ParseResult p mm r = parse ctx{macros=macros++m} rest''
				in
					ParseResult p (m ++ mm) r
			| otherwise = parse ctx $ snd $ fromJust $ balanced ('{', '}') $ dropWhile (/= Token "{") rest

endCommand :: (Text, Command)
endCommand = ("end", Command $ \c _ws rest ->
	let Just (_, rest') = parseFixArg c rest in ParseResult mempty mempty rest')

beginCommand :: (Text, Command)
beginCommand = ("begin", normalCmd $ Command pars)
	where
		pars c@Context{..} _ws rest
			| Just (Environment f) <- Map.lookup (envname) (environments macros) = f c rest'
			| otherwise = error $ "undefined env: " ++ Text.unpack envname
			where
				Just (arg, rest') = parseFixArg c rest
				[TeXRaw envname] = concatRaws arg

raiseBoxCommand :: Command
raiseBoxCommand = normalCmd $ Command $ \c@Context{..} _ws rest ->
	let
		Just (a0, rest') = balanced ('{', '}') rest
		(a1, rest'') = case parseOptArg rest' of
			Nothing -> (Nothing, rest')
			Just (x, y) -> (Just x, y)
		Just (a2, rest''') = balanced ('{', '}') rest''
		args = [(FixArg, fullParse c a0)]
			++ case a1 of
				Nothing -> []
				Just x -> [(OptArg, fullParse c x)]
			++ [(FixArg, fullParse c a2)]
	in
		ParseResult [TeXComm "raisebox" "" args] mempty rest'''

newCommandCommand :: Command
newCommandCommand = normalCmd $ Command $ \Context{..} _ws (Token "{" : Token ('\\' : name) : Token "}" : rest) ->
	let
		(sig, rest') = parseSignature rest
		Just (body, rest'') = balanced ('{', '}') rest'
		newMacros = newCommand True (defCmd (Text.pack name) sig body) mempty
	in
		ParseResult [] newMacros rest''

declareMathOperator :: Command
declareMathOperator = normalCmd $ Command $ \Context{..} _ws (Token "{" : Token ('\\' : name) : Token "}" : rest) ->
	let
		Just (body, rest') = balanced ('{', '}') rest
		newBody = [Token "\\operatorname", Token "{"] ++ body ++ [Token "}"]
		newMacros = newCommand True (defCmd (Text.pack name) (Signature 0 Nothing) newBody) mempty
	in
		ParseResult [] newMacros rest'

newColumnTypeCommand :: Command
newColumnTypeCommand = normalCmd $ Command $ \Context{..} _ws (Token "{" : Token _ : Token "}" : rest) ->
	let
		(_, rest') = parseSignature rest
		Just (_, rest'') = balanced ('{', '}') rest'
	in
		ParseResult [] mempty rest''

defaultMacros :: Macros
defaultMacros = Macros (Map.fromList defaultCmds) (Map.fromList defaultEnvs) mempty

defaultContext :: Context
defaultContext = Context
	{ commentsEnabled = True
	, parsingOptArg = False
	, macros = defaultMacros }

rmLine :: [Token] -> [Token]
rmLine s = case dropWhile (/= Token "\n") s of
	Token "\n" : x -> x
	x -> x

parseOptArg :: [Token] -> Maybe ([Token], [Token])
parseOptArg = balanced ('[', ']')

parseOptArgs :: [Token] -> ([[Token]], [Token])
parseOptArgs s
	| Just (r, s') <- parseOptArg s = first (r:) (parseOptArgs s')
	| otherwise = ([], s)

parseFixArg :: Context -> [Token] -> Maybe (LaTeX, [Token])
parseFixArg ctx (Token [c] : more) | isSpace c = parseFixArg ctx more
parseFixArg ctx (Token "{" : more) =
	let ParseResult t _macros s = parse ctx more in Just (t, s)
parseFixArg _ _ = Nothing

parseSignature :: [Token] -> (Signature, [Token])
parseSignature t = case optArgs of
	[] -> (Signature 0 Nothing, t')
	[[Token a]] -> (Signature (read a) Nothing, t')
	[[Token a], deflt] -> (Signature (read a) (Just deflt), t')
	_ -> error "unrecognized signature"
	where (optArgs, t') = parseOptArgs t

balanced :: (Char, Char) -> [Token] -> Maybe ([Token], [Token])
balanced (open, close) (dropWhile (all isSpace . tokenChars) -> (Token [o] : s))
	| o == open = Just $ go 0 s
	where
		go :: Int -> [Token] -> ([Token], [Token])
		go 0 [] = ([], [])
		go 0 (Token [c] : x) | c == close = ([], x)
		go n (Token "}" : x) = first (Token "}" :) (go (n-1) x)
		go n (Token "{" : x) = first (Token "{" :) (go (n+1) x)
		go n (x:y) = first (x :) (go n y)
		go n x = error $ "\n\nbalanced: " ++ show (n, x)
balanced oc (dropWhile (all isSpace. tokenChars) -> (Token "%" : x)) = balanced oc (dropWhile (/= Token "\n") x)
balanced _ _ = Nothing

balanced_body :: Context -> String -> [Token] -> ([Token], [Token])
balanced_body ctx end = go 0
	where
		go :: Int -> [Token] -> ([Token], [Token])
		go 0 [] = ([], [])
		go 0 (Token "\\end" : Token "{" : e : Token "}" : x) | fullParse ctx [e] == [TeXRaw $ Text.pack end] = ([], x)
		go n (Token "}" : x) = first (Token "}" :) (go (n-1) x)
		go n (Token "{" : x) = first (Token "{" :) (go (n+1) x)
		go n (x:y) = first (x :) (go n y)
		go n s = error $ "\n\nbalanced: " ++ show (n, s)

parseArgs :: Signature -> [Token] -> ([[Token]], [Token])
parseArgs Signature{..} s = case defaultArg of
	Nothing -> n_balanced ('{', '}') nrFixArgs s
	Just dfl -> case parseOptArg s of
		Nothing ->
			first (dfl :) (n_balanced ('{', '}') (nrFixArgs - 1) s)
		Just (optArg, s') ->
			first (optArg :) (n_balanced ('{', '}') (nrFixArgs - 1) s')

parseArgs2 :: Context -> Signature -> [Token] -> ([TeXArg], [Token])
parseArgs2 c Signature{..} s
	| defaultArg == Nothing = first (map fa) (n_balanced ('{', '}') nrFixArgs s)
	| Just (optArg, s') <- parseOptArg s =
			first (\a -> (OptArg, fullParse c optArg) : map fa a)
			(n_balanced ('{', '}') (nrFixArgs - 1) s')
	| otherwise = first (map fa) (n_balanced ('{', '}') (nrFixArgs - 1) s)
	where
		fa = (FixArg, ) . fullParse c

-- todo: clean up parseArgs/parseArgs2 above

n_balanced :: (Char, Char) -> Int -> [Token] -> ([[Token]], [Token])
n_balanced oc n s
	| n > 0, Just (x, s') <- balanced oc s = first (x:) $ n_balanced oc (n-1) s'
	| otherwise = ([], s)

newEnvCommand :: Command
newEnvCommand = normalCmd $ Command $ \Context{..} _ws (Token "{" : (span (/= Token "}") -> (name, Token "}" : rest))) ->
	let
		nameStr = concatMap tokenChars name
		(sig, rest') = parseSignature rest
		Just (begin, rest'') = balanced ('{', '}') rest'
		Just (end, rest''') = balanced ('{', '}') rest''
		pa :: Context -> [Token] -> ParseResult
		pa c' toks = ParseResult replaced mempty toks''
			where
				replaced = fullParse c' $ replArgs args begin ++ body ++ end
				(args, toks') = parseArgs sig toks
				(body, toks'') = balanced_body c' nameStr toks'
		m = Macros mempty (Map.singleton (Text.pack nameStr) (Environment pa)) mempty
	in
		ParseResult [] m rest'''

parseString :: Context -> String -> (LaTeX, Macros, [Token])
parseString c s = (concatRaws x, y, z)
	where ParseResult x y z = parse c (tokenize s)

literal :: String
literal = " @_{}&,%-#/~>!$;:^"

breakComment :: [Token] -> ([Token], [Token])
breakComment x@(Token "\n" : _) = ([], x)
breakComment (Token ('\\' : cmd) : xs)
	| (c, r@(_:_)) <- span (/= '\n') cmd = ([Token ('\\':c)], Token r : xs)
breakComment (Token "%" : Token "\n" : x) = first ((Token "%" :) . (Token "\n" :)) (breakComment x)
breakComment (x : xs) = first (x:) (breakComment xs)
breakComment [] = ([], [])

data LiteralKind = StringLiteral | CharLiteral

parseCode :: Text -> Context -> [Token] -> LaTeX
parseCode envname c = concatRaws . go Nothing
	where
		go :: Maybe LiteralKind -> [Token] -> LaTeX
		go _ [] = []
		go b (Token "@" : rest) = fullParse c cmd ++ go b rest'
			where (cmd, Token "@" : rest') = break (== Token "@") rest
		go (Just StringLiteral) (Token "\"" : rest) = TeXRaw "\"" : go Nothing rest
		go (Just CharLiteral) (Token "'" : rest) = TeXRaw "'" : go Nothing rest
		go Nothing (Token "\"" : rest) = TeXRaw "\"" : (go (Just StringLiteral) lit ++ go Nothing rest')
			where (lit, rest') = stringLiteral rest
		go Nothing (Token "'" : rest)
		       | envname == "codeblockdigitsep" = TeXRaw "'" : go Nothing rest
		       | otherwise = TeXRaw "'" : (go (Just CharLiteral) lit ++ go Nothing rest')
		       where (lit, rest') = charLiteral rest
		go Nothing (Token "/" : Token "/" : (breakComment -> (comment, rest')))
			= TeXComm "comment" "" [(FixArg, TeXRaw "//" : noncode comment)] : go Nothing rest'
		go Nothing (Token "/" : Token "*" : rest)
		    | Just (comment, rest') <- stripInfix [Token "*", Token "/"] rest
		    = TeXComm "comment" "" [(FixArg, [TeXRaw "/*"] ++ noncode comment ++ [TeXRaw "*/"])] : go Nothing rest'
		go b (Token "/" : rest) = TeXRaw "/" : go b rest
		go b s = TeXRaw (Text.pack $ concatMap tokenChars code) : go b rest
			where
			  breakToks = [Token "@", Token "/"] ++ 
			    case b of
			      Nothing -> [Token "\"", Token "'"]
			      Just StringLiteral -> [Token "\""]
			      Just CharLiteral -> [Token "'"]
			  (code, rest) = break (`elem` breakToks) s
		noncode :: [Token] -> LaTeX
		noncode toks =
		  fullParse c nc ++ case more of
		    [] -> []
		    Token "@" : (break (== Token "@") -> (code, _ : rest)) ->
		        TeXComm "tcode" "" [(FixArg, fullParse c code)] : noncode rest
		    _ -> error "no"
		    where (nc, more) = span (/= Token "@") toks
		stringLiteral :: [Token] -> ([Token], [Token])
		stringLiteral (Token "\\" : Token "\"" : x) = first (Token "\\\"" :) (stringLiteral x)
		stringLiteral (Token "\\" : Token "\\" : x) = first (Token "\\\\" :) (stringLiteral x)
		stringLiteral (Token "\"" : x) = ([Token "\""], x)
		stringLiteral (y : x) = first (y :) (stringLiteral x)
		stringLiteral [] = ([], [])
		charLiteral :: [Token] -> ([Token], [Token])
		charLiteral (Token "\\" : Token "'" : x) = first (Token "\\'" :) (charLiteral x)
		charLiteral (Token "\\" : Token "\\" : x) = first (Token "\\\\" :) (charLiteral x)
		charLiteral (Token "'" : x) = ([Token "'"], x)
		charLiteral (y : x) = first (y :) (charLiteral x)
		charLiteral [] = ([], [])

isCommandChar :: Char -> Bool
isCommandChar c = isAlpha c || c == '*'

tokenize :: String -> [Token]
tokenize "" = []
tokenize ('\\':'v':'e':'r':'b': delim : (break (== delim) -> (arg, _ : rest))) =
	Token ("\\verb:" ++ arg) : tokenize rest
tokenize ('\\' : (span isCommandChar -> (cmd@(_:_), (span isSpace -> (ws, rest)))))
	= Token ('\\' : cmd ++ ws) : tokenize rest
tokenize ('\\' : c : rest) = Token ['\\', c] : tokenize rest
tokenize x@((isAlpha -> True): _) = let (a, b) = span isAlphaNum x in Token a : tokenize b
tokenize (x:y) = Token [x] : tokenize y

	-- \verb is handled in tokenize so that the 'balanced' function doesn't
	-- get confused by \verb|{|

	-- Notice how the whitespace following a command like \bla is included in the Token
	-- This lets the parser include it in the TeXComm/TeXCommS's command field, so that
	-- the whitespace is not lost when serializing back to text when sending to MathJax.

replArgs :: [[Token]] -> [Token] -> [Token]
replArgs args = go
	where
		go [] = []
		go (Token "%" : (span (/= Token "\n") -> (x, y))) = Token "%" : x ++ go y
		go (Token "#" : Token "#" : y) = Token "#" : go y
		go (Token "#" : Token [getDigit -> Just i] : y)
			| length args >= i = (args !! (i-1)) ++ go y
			| otherwise = error $ "need more args than " ++ show args ++ " to replace in " ++ show (concatMap tokenChars y)
		go (x:y) = x : go y

nullCmd :: Text -> Signature -> (Text, Command)
nullCmd name sig = defCmd name sig []

storeCmd :: String -> Signature -> (Text, Command)
storeCmd name sig = (Text.pack name, normalCmd $ Command pars)
	where
		pars context ws tokens = ParseResult [TeXComm name ws args] mempty rest
			where (args, rest) = parseArgs2 context sig tokens

defCmd :: Text -> Signature -> [Token] -> (Text, Command)
defCmd name sig body = (name, normalCmd $ Command pars)
	where
		pars context _ws tokens = ParseResult (fullParse context $ replArgs args body) mempty rest
			where (args, rest) = parseArgs sig tokens

normalCmd :: Command -> Command
normalCmd (Command f) = Command $ \ctx ws toks ->
	let ParseResult content newMacros rest = f ctx ws toks
	in addMacros False newMacros (prependContent content (parse ctx{macros=macros ctx ++ newMacros} rest))

consumeMath :: [Token] -> ([Token], [Token])
consumeMath = f 0
    where
        f :: Integer -> [Token] -> ([Token], [Token])
        f 0 (Token "$" : rest) = ([], rest)
        f depth (Token "{" : rest) = first (Token "{" :) (f (depth + 1) rest)
        f depth (Token "}" : rest) = first (Token "}" :) (f (depth - 1) rest)
        f depth (tok : rest) = first (tok :) (f depth rest)
        f _ [] = error "unexpected end of math"

parse :: Context -> [Token] -> ParseResult
parse c (Token "$" : (consumeMath -> (math, rest))) =
	prependContent [TeXMath Dollar (fullParse c math)] (parse c rest)
parse c (Token "\\[" : (span (/= Token "\\]") -> (math, Token "\\]" : rest))) =
	prependContent [TeXMath Square (fullParse c math)] (parse c rest)
parse c (Token "]" : x)
	| parsingOptArg c = ParseResult mempty mempty x
parse _ (Token "}" : x) = ParseResult mempty mempty x
parse c (Token "{" : x) = prependContent [TeXBraces y] $ parse c rest
	where ParseResult y _ rest = parse c x
parse c (Token "%" : x)
	| commentsEnabled c = parse c (rmLine x)
parse _ [] = ParseResult mempty mempty mempty
parse c (Token "\\\\" : x) = prependContent [TeXLineBreak] (parse c x)
parse c (Token ['\\', ch] : x)
	| ch `elem` literal = prependContent [TeXComm [ch] "" []] (parse c x)
parse c (Token ('\\':'v':'e':'r':'b':':':arg) : rest) =
	prependContent [TeXComm "verb" "" [(FixArg, [TeXRaw $ Text.pack arg])]] (parse c rest)
parse c (Token "\\rSec" : Token [getDigit -> Just i] : s)
		= prependContent [TeXComm "rSec" "" args] $ parse c s''
	where
		Just (a, s') = parseOptArg s
		Just (b, s'') = parseFixArg c s'
		args = [(FixArg, [TeXRaw $ Text.pack $ show i]), (FixArg, fullParse c a), (FixArg, b)]
parse c@Context{..} (Token ('\\' : (span (not . isSpace) -> (nos, w))) : rest)
	| Just f <- Map.lookup (Text.pack cmd) (commands macros) = runCommand f c ws rest
	| otherwise = error $
		"\n\nundefined command: " ++ show cmd ++ " at: " ++ take 50 (concatMap tokenChars rest)
	where (cmd, ws) | nos == "", (x : xx) <- w = ([x], xx)
	                | otherwise = (nos, w)
parse ctx (Token c : rest)
	| all isAlphaNum c
		= prependContent [TeXRaw $ Text.pack c] $ parse ctx rest
parse ctx (Token [c] : rest)
		= prependContent [TeXRaw $ Text.pack [c]] $ parse ctx rest
parse _ s = error $ "parse: unexpected: " ++ take 100 (concatMap tokenChars s)

fullParse :: Context -> [Token] -> LaTeX
fullParse c t
	| all isSpace (concatMap tokenChars remainder) = concatRaws content
	| otherwise = error $ "could not fully parse: "
		++ concatMap tokenChars t
		++ "\n\nremainder: "
		++ concatMap tokenChars remainder
	where ParseResult{..} = parse c t