1 files changed, 75 insertions, 46 deletions

diff --git a/src/Text/Pandoc/Lua.hs b/src/Text/Pandoc/Lua.hs
index d6e5def4a..477076191 100644
--- a/src/Text/Pandoc/Lua.hs
+++ b/src/Text/Pandoc/Lua.hs
@@ -1,8 +1,9 @@
+{-# LANGUAGE CPP                   #-}
 {-# LANGUAGE FlexibleContexts      #-}
 {-# LANGUAGE FlexibleInstances     #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
-{-# LANGUAGE DeriveDataTypeable    #-}
 {-# LANGUAGE ScopedTypeVariables   #-}
+{-# OPTIONS_GHC -fno-warn-orphans #-}
 {-
 Copyright © 2017 Albert Krewinkel <tarleb+pandoc@moltkeplatz.de>
 
@@ -33,6 +34,7 @@ Pandoc lua utils.
 module Text.Pandoc.Lua (LuaException (..), pushPandocModule, runLuaFilter) where
 
 import Control.Monad (mplus, unless, when, (>=>))
+import Control.Monad.Identity (Identity)
 import Control.Monad.Trans (MonadIO (..))
 import Data.Data (DataType, Data, toConstr, showConstr, dataTypeOf,
                   dataTypeConstrs, dataTypeName, tyconUQname)
@@ -40,10 +42,10 @@ import Data.Foldable (foldrM)
 import Data.Map (Map)
 import Data.Maybe (isJust)
 import Foreign.Lua (Lua, FromLuaStack (peek), LuaException (..), StackIndex,
-                    Status(OK), ToLuaStack (push))
+                    Status (OK), ToLuaStack (push))
 import Text.Pandoc.Definition
 import Text.Pandoc.Lua.PandocModule (pushPandocModule)
-import Text.Pandoc.Walk (Walkable (walkM))
+import Text.Pandoc.Walk (walkM)
 
 import qualified Data.Map as Map
 import qualified Foreign.Lua as Lua
@@ -56,7 +58,7 @@ runLuaFilter datadir filterPath args pd = liftIO . Lua.runLua $ do
   pushPandocModule datadir
   Lua.setglobal "pandoc"
   top <- Lua.gettop
-  stat<- Lua.dofile filterPath
+  stat <- Lua.dofile filterPath
   if stat /= OK
     then do
       luaErrMsg <- peek (-1) <* Lua.pop 1
@@ -64,7 +66,7 @@ runLuaFilter datadir filterPath args pd = liftIO . Lua.runLua $ do
     else do
       newtop <- Lua.gettop
       -- Use the implicitly defined global filter if nothing was returned
-      when (newtop - top < 1) $ pushGlobalFilter
+      when (newtop - top < 1) pushGlobalFilter
       luaFilters <- peek (-1)
       push args
       Lua.setglobal "PandocParameters"
@@ -81,27 +83,36 @@ runAll :: [LuaFilter] -> Pandoc -> Lua Pandoc
 runAll = foldr ((>=>) . walkMWithLuaFilter) return
 
 walkMWithLuaFilter :: LuaFilter -> Pandoc -> Lua Pandoc
-walkMWithLuaFilter (LuaFilter fnMap) = walkLua
-  where
-    walkLua :: Pandoc -> Lua Pandoc
-    walkLua =
-      (if hasOneOf inlineFilterNames
-        then walkM (tryFilter fnMap :: Inline -> Lua Inline)
-        else return)
-      >=>
-      (if hasOneOf blockFilterNames
-        then walkM (tryFilter fnMap :: Block -> Lua Block)
-        else return)
-      >=>
-      (case Map.lookup "Meta" fnMap of
-          Just fn -> walkM (\(Pandoc meta blocks) -> do
-                               meta' <- runFilterFunction fn meta
-                               return $ Pandoc meta' blocks)
-          Nothing -> return)
-      >=>
-      (case foldl mplus Nothing (map (`Map.lookup` fnMap) pandocFilterNames) of
-          Just fn -> runFilterFunction fn :: Pandoc -> Lua Pandoc
-          Nothing -> return)
+walkMWithLuaFilter (LuaFilter fnMap) =
+  walkInlines >=> walkBlocks >=> walkMeta >=> walkPandoc
+ where
+    walkInlines :: Pandoc -> Lua Pandoc
+    walkInlines =
+      if hasOneOf inlineFilterNames
+      then walkM (mconcatMapM (tryFilter fnMap :: Inline -> Lua [Inline]))
+      else return
+
+    walkBlocks :: Pandoc -> Lua Pandoc
+    walkBlocks =
+      if hasOneOf blockFilterNames
+      then walkM (mconcatMapM (tryFilter fnMap :: Block -> Lua [Block]))
+      else return
+
+    walkMeta :: Pandoc -> Lua Pandoc
+    walkMeta =
+      case Map.lookup "Meta" fnMap of
+        Just fn -> walkM (\(Pandoc meta blocks) -> do
+                             meta' <- runFilterFunction fn meta *> singleElement meta
+                             return $ Pandoc meta' blocks)
+        Nothing -> return
+
+    walkPandoc :: Pandoc -> Lua Pandoc
+    walkPandoc =
+      case foldl mplus Nothing (map (`Map.lookup` fnMap) pandocFilterNames) of
+        Just fn -> \x -> runFilterFunction fn x *> singleElement x
+        Nothing -> return
+
+    mconcatMapM f = fmap mconcat . mapM f
     hasOneOf = any (\k -> isJust (Map.lookup k fnMap))
 
 constructorsFor :: DataType -> [String]
@@ -124,14 +135,15 @@ newtype LuaFilter = LuaFilter FunctionMap
 newtype LuaFilterFunction = LuaFilterFunction { functionIndex :: Int }
 
 -- | Try running a filter for the given element
-tryFilter :: (Data a, FromLuaStack a, ToLuaStack a) => FunctionMap -> a -> Lua a
+tryFilter :: (Data a, FromLuaStack a, ToLuaStack a)
+          => FunctionMap -> a -> Lua [a]
 tryFilter fnMap x =
   let filterFnName = showConstr (toConstr x)
       catchAllName = tyconUQname $ dataTypeName (dataTypeOf x)
   in
   case Map.lookup filterFnName fnMap `mplus` Map.lookup catchAllName fnMap of
-    Just fn -> runFilterFunction fn x
-    Nothing -> return x
+    Just fn -> runFilterFunction fn x *> elementOrList x
+    Nothing -> return [x]
 
 instance FromLuaStack LuaFilter where
   peek idx =
@@ -151,28 +163,42 @@ instance FromLuaStack LuaFilter where
 -- called with given element as argument and is expected to return an element.
 -- Alternatively, the function can return nothing or nil, in which case the
 -- element is left unchanged.
-runFilterFunction :: (FromLuaStack a, ToLuaStack a)
-                  => LuaFilterFunction -> a -> Lua a
+runFilterFunction :: ToLuaStack a => LuaFilterFunction -> a -> Lua ()
 runFilterFunction lf x = do
   pushFilterFunction lf
   push x
   z <- Lua.pcall 1 1 Nothing
-  if z /= OK
-    then do
-      msg <- peek (-1)
-      let prefix = "Error while running filter function: "
-      Lua.throwLuaError $ prefix ++ msg
+  when (z /= OK) $ do
+    msg <- Lua.peek (-1) <* Lua.pop 1
+    let prefix = "Error while running filter function: "
+    Lua.throwLuaError $ prefix ++ msg
+
+elementOrList :: FromLuaStack a => a -> Lua [a]
+elementOrList x = do
+  let topOfStack = Lua.StackIndex (-1)
+  elementUnchanged <- Lua.isnil topOfStack
+  if elementUnchanged
+    then [x] <$ Lua.pop 1
+    else do
+       mbres <- Lua.peekEither topOfStack
+       case mbres of
+         Right res -> [res] <$ Lua.pop 1
+         Left _  -> Lua.toList topOfStack <* Lua.pop 1
+
+singleElement :: FromLuaStack a => a -> Lua a
+singleElement x = do
+  elementUnchanged <- Lua.isnil (-1)
+  if elementUnchanged
+    then x <$ Lua.pop 1
     else do
-      noExplicitFilter <- Lua.isnil (-1)
-      if noExplicitFilter
-        then  Lua.pop 1 *> return x
-        else do
-          mbres <- Lua.peekEither (-1)
-          case mbres of
-            Left err -> Lua.throwLuaError
-                        ("Error while trying to get a filter's return "
-                         ++ "value from lua stack.\n" ++ err)
-            Right res -> res <$ Lua.pop 1
+    mbres <- Lua.peekEither (-1)
+    case mbres of
+      Right res -> res <$ Lua.pop 1
+      Left err  -> do
+        Lua.pop 1
+        Lua.throwLuaError $
+          "Error while trying to get a filter's return " ++
+          "value from lua stack.\n" ++ err
 
 -- | Push the filter function to the top of the stack.
 pushFilterFunction :: LuaFilterFunction -> Lua ()
@@ -188,6 +214,9 @@ registerFilterFunction idx = do
   refIdx <- Lua.ref Lua.registryindex
   return $ LuaFilterFunction refIdx
 
+instance (FromLuaStack a) => FromLuaStack (Identity a) where
+  peek = fmap return . peek
+
 instance ToLuaStack LuaFilterFunction where
   push = pushFilterFunction