Allow lua filters to return lists of elements

Closes: #3918

3 files changed, 93 insertions, 61 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
 
diff --git a/src/Text/Pandoc/Lua/StackInstances.hs b/src/Text/Pandoc/Lua/StackInstances.hs
index 15a7cdd84..73b04e50f 100644
--- a/src/Text/Pandoc/Lua/StackInstances.hs
+++ b/src/Text/Pandoc/Lua/StackInstances.hs
@@ -35,14 +35,15 @@ module Text.Pandoc.Lua.StackInstances () where
 import Control.Applicative ((<|>))
 import Foreign.Lua (Lua, LuaInteger, LuaNumber, Type (..), FromLuaStack (peek),
                     ToLuaStack (push), StackIndex, throwLuaError, tryLua)
-import Foreign.Lua.Api (getmetatable, ltype, newtable, pop, rawget, rawlen)
 import Text.Pandoc.Definition
 import Text.Pandoc.Lua.Util (addValue, adjustIndexBy, getTable, pushViaConstructor)
 import Text.Pandoc.Shared (safeRead)
 
+import qualified Foreign.Lua as Lua
+
 instance ToLuaStack Pandoc where
   push (Pandoc meta blocks) = do
-    newtable
+    Lua.newtable
     addValue "blocks" blocks
     addValue "meta"   meta
 
@@ -156,7 +157,7 @@ peekMetaValue idx = do
   -- Get the contents of an AST element.
   let elementContent :: FromLuaStack a => Lua a
       elementContent = peek idx
-  luatype <- ltype idx
+  luatype <- Lua.ltype idx
   case luatype of
     TypeBoolean -> MetaBool <$> peek idx
     TypeString  -> MetaString <$> peek idx
@@ -172,13 +173,13 @@ peekMetaValue idx = do
         Right t             -> throwLuaError ("Unknown meta tag: " ++ t)
         Left _ -> do
           -- no meta value tag given, try to guess.
-          len <- rawlen idx
+          len <- Lua.rawlen idx
           if len <= 0
             then MetaMap <$> peek idx
             else  (MetaInlines <$> peek idx)
                   <|> (MetaBlocks <$> peek idx)
                   <|> (MetaList <$> peek idx)
-    _        -> throwLuaError ("could not get meta value")
+    _        -> throwLuaError "could not get meta value"
 
 -- | Push an block element to the top of the lua stack.
 pushBlock :: Block -> Lua ()
@@ -284,16 +285,15 @@ peekInline idx = do
 
 getTag :: StackIndex -> Lua String
 getTag idx = do
-  hasMT <- getmetatable idx
-  if hasMT
-    then do
-      push "tag"
-      rawget (-2)
-      peek (-1) <* pop 2
-    else do
-      push "tag"
-      rawget (idx `adjustIndexBy` 1)
-      peek (-1) <* pop 1
+  top <- Lua.gettop
+  hasMT <- Lua.getmetatable idx
+  push "tag"
+  if hasMT then Lua.rawget (-2) else Lua.rawget (idx `adjustIndexBy` 1)
+  r <- tryLua (peek (-1))
+  Lua.settop top
+  case r of
+    Left (Lua.LuaException err) -> throwLuaError err
+    Right res -> return res
 
 withAttr :: (Attr -> a -> b) -> (LuaAttr, a) -> b
 withAttr f (attributes, x) = f (fromLuaAttr attributes) x
diff --git a/test/lua/undiv.lua b/test/lua/undiv.lua
new file mode 100644
index 000000000..1cbb6d30e
--- /dev/null
+++ b/test/lua/undiv.lua
@@ -0,0 +1,3 @@
+function Div(el)
+  return el.content
+end