Lua: traverse nested blocks and inlines in correct order

Traversal methods are updated to use the new Walk module such that
sequences with nested Inline (or Block) elements are traversed in the
order in which they appear in the linearized document.

Fixes: #5667

2 files changed, 24 insertions, 11 deletions

diff --git a/src/Text/Pandoc/Lua/Filter.hs b/src/Text/Pandoc/Lua/Filter.hs
index e8958347d..9efc2b9ae 100644
--- a/src/Text/Pandoc/Lua/Filter.hs
+++ b/src/Text/Pandoc/Lua/Filter.hs
@@ -20,6 +20,7 @@ module Text.Pandoc.Lua.Filter ( LuaFilterFunction
                               , walkBlocks
                               , blockElementNames
                               , inlineElementNames
+                              , module Text.Pandoc.Lua.Walk
                               ) where
 import Prelude
 import Control.Monad (mplus, (>=>))
@@ -31,7 +32,8 @@ import Data.Map (Map)
 import Foreign.Lua (Lua, Peekable, Pushable)
 import Text.Pandoc.Definition
 import Text.Pandoc.Lua.Marshaling ()
-import Text.Pandoc.Walk (walkM, Walkable)
+import Text.Pandoc.Lua.Walk (SingletonsList (..))
+import Text.Pandoc.Walk (Walkable (walkM))
 
 import qualified Data.Map.Strict as Map
 import qualified Foreign.Lua as Lua
@@ -115,6 +117,12 @@ tryFilter (LuaFilter fnMap) x =
     Just fn -> runFilterFunction fn x *> elementOrList x
     Nothing -> return [x]
 
+-- | Apply filter on a sequence of AST elements.
+runOnSequence :: (Data a, Peekable a, Pushable a)
+              => LuaFilter -> SingletonsList a -> Lua (SingletonsList a)
+runOnSequence lf (SingletonsList xs) =
+  SingletonsList <$> mconcatMapM (tryFilter lf) xs
+
 -- | Push a value to the stack via a lua filter function. The filter function is
 -- 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
@@ -135,16 +143,20 @@ mconcatMapM f = fmap mconcat . mapM f
 hasOneOf :: LuaFilter -> [String] -> Bool
 hasOneOf (LuaFilter fnMap) = any (\k -> Map.member k fnMap)
 
-walkInlines :: Walkable [Inline] a => LuaFilter -> a -> Lua a
-walkInlines f =
-  if f `hasOneOf` inlineElementNames
-     then walkM (mconcatMapM (tryFilter f :: Inline -> Lua [Inline]))
+walkInlines :: Walkable (SingletonsList Inline) a => LuaFilter -> a -> Lua a
+walkInlines lf =
+  let f :: SingletonsList Inline -> Lua (SingletonsList Inline)
+      f = runOnSequence lf
+  in if lf `hasOneOf` inlineElementNames
+     then walkM f
      else return
 
-walkBlocks :: Walkable [Block] a => LuaFilter -> a -> Lua a
-walkBlocks f =
-  if f `hasOneOf` blockElementNames
-     then walkM (mconcatMapM (tryFilter f :: Block -> Lua [Block]))
+walkBlocks :: Walkable (SingletonsList Block) a => LuaFilter -> a -> Lua a
+walkBlocks lf =
+  let f :: SingletonsList Block -> Lua (SingletonsList Block)
+      f = runOnSequence lf
+  in if lf `hasOneOf` blockElementNames
+     then walkM f
      else return
 
 walkMeta :: LuaFilter -> Pandoc -> Lua Pandoc
diff --git a/src/Text/Pandoc/Lua/Module/Pandoc.hs b/src/Text/Pandoc/Lua/Module/Pandoc.hs
index 09892db49..8950c4b7f 100644
--- a/src/Text/Pandoc/Lua/Module/Pandoc.hs
+++ b/src/Text/Pandoc/Lua/Module/Pandoc.hs
@@ -23,7 +23,7 @@ import Foreign.Lua (Lua, NumResults, Optional, Peekable, Pushable)
 import System.Exit (ExitCode (..))
 import Text.Pandoc.Class (runIO)
 import Text.Pandoc.Definition (Block, Inline)
-import Text.Pandoc.Lua.Filter (walkInlines, walkBlocks, LuaFilter)
+import Text.Pandoc.Lua.Filter (walkInlines, walkBlocks, LuaFilter, SingletonsList (..))
 import Text.Pandoc.Lua.Marshaling ()
 import Text.Pandoc.Walk (Walkable)
 import Text.Pandoc.Options (ReaderOptions (readerExtensions))
@@ -46,7 +46,8 @@ pushModule datadir = do
   LuaUtil.addFunction "walk_inline" walkInline
   return 1
 
-walkElement :: (Walkable [Inline] a, Walkable [Block] a)
+walkElement :: (Walkable (SingletonsList Inline) a,
+                Walkable (SingletonsList Block) a)
             => a -> LuaFilter -> Lua a
 walkElement x f = walkInlines f x >>= walkBlocks f