Lua filters: refactor lua module handling

The integration with Lua's package/module system is improved: A
pandoc-specific package searcher is prepended to the searchers in
`package.searchers`. The modules `pandoc` and `pandoc.mediabag` can now
be loaded via `require`.

4 files changed, 190 insertions, 56 deletions

diff --git a/src/Text/Pandoc/Lua.hs b/src/Text/Pandoc/Lua.hs
index 148e7a23d..1ca67dced 100644
--- a/src/Text/Pandoc/Lua.hs
+++ b/src/Text/Pandoc/Lua.hs
@@ -31,45 +31,46 @@ Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
 Pandoc lua utils.
 -}
-module Text.Pandoc.Lua (LuaException (..), pushPandocModule, runLuaFilter) where
+module Text.Pandoc.Lua
+  ( LuaException (..)
+  , LuaPackageParams (..)
+  , pushPandocModule
+  , runLuaFilter
+  , initLuaState
+  , luaPackageParams
+  ) where
 
 import Control.Monad (when, (>=>))
 import Control.Monad.Identity (Identity)
 import Control.Monad.Trans (MonadIO (..))
-import Data.IORef (IORef, newIORef, readIORef)
+import Data.IORef (newIORef, readIORef)
 import Foreign.Lua (FromLuaStack (peek), Lua, LuaException (..),
                     Status (OK), ToLuaStack (push))
-import Text.Pandoc.Class (CommonState, PandocIO, getCommonState, getMediaBag,
-                          setMediaBag)
+import Text.Pandoc.Class (PandocIO, getCommonState, getMediaBag, setMediaBag)
 import Text.Pandoc.Definition
+import Text.Pandoc.Lua.Packages (LuaPackageParams (..),
+                                 installPandocPackageSearcher)
 import Text.Pandoc.Lua.PandocModule (pushMediaBagModule, pushPandocModule)
 import Text.Pandoc.Lua.Filter (LuaFilter, walkMWithLuaFilter)
-import Text.Pandoc.MediaBag (MediaBag)
 import qualified Foreign.Lua as Lua
 import qualified Foreign.Lua.Module.Text as Lua
 
 runLuaFilter :: Maybe FilePath -> FilePath -> String
              -> Pandoc -> PandocIO (Either LuaException Pandoc)
 runLuaFilter datadir filterPath format pd = do
-  commonState <- getCommonState
-  mediaBag <- getMediaBag
-  mediaBagRef <- liftIO (newIORef mediaBag)
+  luaPkgParams <- luaPackageParams datadir
   res <- liftIO . Lua.runLuaEither $
-         runLuaFilter' commonState datadir filterPath format mediaBagRef pd
-  newMediaBag <- liftIO (readIORef mediaBagRef)
+         runLuaFilter' luaPkgParams filterPath format pd
+  newMediaBag <- liftIO (readIORef (luaPkgMediaBag luaPkgParams))
   setMediaBag newMediaBag
   return res
 
-runLuaFilter' :: CommonState
-              -> Maybe FilePath -> FilePath -> String -> IORef MediaBag
+runLuaFilter' :: LuaPackageParams
+              -> FilePath -> String
               -> Pandoc -> Lua Pandoc
-runLuaFilter' commonState datadir filterPath format mbRef pd = do
-  Lua.openlibs
-  Lua.preloadTextModule "text"
+runLuaFilter' luaPkgOpts filterPath format pd = do
+  initLuaState luaPkgOpts
   -- store module in global "pandoc"
-  pushPandocModule datadir
-  Lua.setglobal "pandoc"
-  addMediaBagModule
   registerFormat
   top <- Lua.gettop
   stat <- Lua.dofile filterPath
@@ -84,15 +85,33 @@ runLuaFilter' commonState datadir filterPath format mbRef pd = do
       luaFilters <- peek (-1)
       runAll luaFilters pd
  where
-  addMediaBagModule = do
-    Lua.getglobal "pandoc"
-    push "mediabag"
-    pushMediaBagModule commonState mbRef
-    Lua.rawset (-3)
   registerFormat = do
     push format
     Lua.setglobal "FORMAT"
 
+luaPackageParams :: Maybe FilePath -> PandocIO LuaPackageParams
+luaPackageParams datadir = do
+  commonState <- getCommonState
+  mbRef <- liftIO . newIORef =<< getMediaBag
+  return LuaPackageParams
+    { luaPkgCommonState = commonState
+    , luaPkgDataDir = datadir
+    , luaPkgMediaBag = mbRef
+    }
+
+-- Initialize the lua state with all required values
+initLuaState :: LuaPackageParams -> Lua ()
+initLuaState luaPkgParams@(LuaPackageParams commonState datadir mbRef) = do
+  Lua.openlibs
+  Lua.preloadTextModule "text"
+  installPandocPackageSearcher luaPkgParams
+  pushPandocModule datadir
+  -- add MediaBag module
+  push "mediabag"
+  pushMediaBagModule commonState mbRef
+  Lua.rawset (-3)
+  Lua.setglobal "pandoc"
+  return ()
 
 pushGlobalFilter :: Lua ()
 pushGlobalFilter = do
diff --git a/src/Text/Pandoc/Lua/Packages.hs b/src/Text/Pandoc/Lua/Packages.hs
new file mode 100644
index 000000000..ede7beccd
--- /dev/null
+++ b/src/Text/Pandoc/Lua/Packages.hs
@@ -0,0 +1,109 @@
+{-
+Copyright © 2017 Albert Krewinkel <tarleb+pandoc@moltkeplatz.de>
+
+This program is free software; you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation; either version 2 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program; if not, write to the Free Software
+Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
+-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE FlexibleContexts #-}
+{- |
+   Module      : Text.Pandoc.Lua.Packages
+   Copyright   : Copyright © 2017 Albert Krewinkel
+   License     : GNU GPL, version 2 or above
+
+   Maintainer  : Albert Krewinkel <tarleb+pandoc@moltkeplatz.de>
+   Stability   : alpha
+
+Pandoc module for lua.
+-}
+module Text.Pandoc.Lua.Packages
+  ( LuaPackageParams (..)
+  , installPandocPackageSearcher
+  ) where
+
+import Control.Monad (forM_)
+import Data.ByteString.Char8 (unpack)
+import Data.IORef (IORef)
+import Foreign.Lua (Lua, NumResults, liftIO)
+import Text.Pandoc.Class (CommonState, readDataFile, runIO, setUserDataDir)
+import Text.Pandoc.MediaBag (MediaBag)
+import Text.Pandoc.Lua.PandocModule (pushPandocModule, pushMediaBagModule)
+import Text.Pandoc.Lua.Util (dostring')
+
+import qualified Foreign.Lua as Lua
+
+-- | Parameters used to create lua packages/modules.
+data LuaPackageParams = LuaPackageParams
+  { luaPkgCommonState :: CommonState
+  , luaPkgDataDir :: Maybe FilePath
+  , luaPkgMediaBag :: IORef MediaBag
+  }
+
+-- | Insert pandoc's package loader as the first loader, making it the default.
+installPandocPackageSearcher :: LuaPackageParams -> Lua ()
+installPandocPackageSearcher luaPkgParams = do
+  Lua.getglobal' "package.searchers"
+  shiftArray
+  Lua.pushHaskellFunction (pandocPackageSearcher luaPkgParams)
+  Lua.wrapHaskellFunction
+  Lua.rawseti (-2) 1
+  Lua.pop 1           -- remove 'package.searchers' from stack
+ where
+  shiftArray = forM_ [4, 3, 2, 1] $ \i -> do
+    Lua.rawgeti (-1) i
+    Lua.rawseti (-2) (i + 1)
+
+-- | Load a pandoc module.
+pandocPackageSearcher :: LuaPackageParams -> String -> Lua NumResults
+pandocPackageSearcher luaPkgParams pkgName =
+  case pkgName of
+    "pandoc"          -> let datadir = luaPkgDataDir luaPkgParams
+                         in pushWrappedHsFun (pushPandocModule datadir)
+    "pandoc.mediabag" -> let st    = luaPkgCommonState luaPkgParams
+                             mbRef = luaPkgMediaBag luaPkgParams
+                         in pushWrappedHsFun (pushMediaBagModule st mbRef)
+    _ -> searchPureLuaLoader
+ where
+  pushWrappedHsFun f = do
+    Lua.pushHaskellFunction f
+    Lua.wrapHaskellFunction
+    return 1
+  searchPureLuaLoader = do
+    let filename = pkgName ++ ".lua"
+    modScript <- liftIO (dataDirScript (luaPkgDataDir luaPkgParams) filename)
+    case modScript of
+      Just script -> pushWrappedHsFun (loadStringAsPackage pkgName script)
+      Nothing -> do
+        Lua.push ("no file '" ++ filename ++ "' in pandoc's datadir")
+        return 1
+
+loadStringAsPackage :: String -> String -> Lua NumResults
+loadStringAsPackage pkgName script = do
+  status <- dostring' script
+  if status == Lua.OK
+    then return (1 :: NumResults)
+    else do
+      msg <- Lua.peek (-1) <* Lua.pop 1
+      Lua.push ("Error while loading ``" ++ pkgName ++ "`.\n" ++ msg)
+      Lua.lerror
+      return (2 :: NumResults)
+
+-- | Get the string representation of the pandoc module
+dataDirScript :: Maybe FilePath -> FilePath -> IO (Maybe String)
+dataDirScript datadir moduleFile = do
+  res <- runIO $ setUserDataDir datadir >> readDataFile moduleFile
+  return $ case res of
+    Left _ -> Nothing
+    Right s -> Just (unpack s)
+
diff --git a/src/Text/Pandoc/Lua/PandocModule.hs b/src/Text/Pandoc/Lua/PandocModule.hs
index 4df01f019..744edfe82 100644
--- a/src/Text/Pandoc/Lua/PandocModule.hs
+++ b/src/Text/Pandoc/Lua/PandocModule.hs
@@ -33,21 +33,21 @@ module Text.Pandoc.Lua.PandocModule
   , pushMediaBagModule
   ) where
 
-import Control.Monad (unless, zipWithM_)
-import Data.ByteString.Char8 (unpack)
+import Control.Monad (zipWithM_)
 import Data.Default (Default (..))
 import Data.Digest.Pure.SHA (sha1, showDigest)
-import Data.IORef
+import Data.IORef (IORef, modifyIORef', readIORef)
 import Data.Maybe (fromMaybe)
 import Data.Text (pack)
 import Foreign.Lua (ToLuaStack, FromLuaStack, Lua, NumResults, liftIO)
 import Foreign.Lua.FunctionCalling (ToHaskellFunction)
 import System.Exit (ExitCode (..))
 import Text.Pandoc.Class (CommonState (..), fetchItem, putCommonState,
-                          readDataFile, runIO, runIOorExplode, setMediaBag,
-                          setUserDataDir)
-import Text.Pandoc.Lua.StackInstances ()
+                          runIO, runIOorExplode, setMediaBag)
 import Text.Pandoc.Definition (Block, Inline)
+import Text.Pandoc.Lua.Filter (walkInlines, walkBlocks, LuaFilter)
+import Text.Pandoc.Lua.StackInstances ()
+import Text.Pandoc.Lua.Util (loadScriptFromDataDir)
 import Text.Pandoc.Walk (Walkable)
 import Text.Pandoc.MIME (MimeType)
 import Text.Pandoc.Options (ReaderOptions (readerExtensions))
@@ -57,43 +57,18 @@ import Text.Pandoc.Readers (Reader (..), getReader)
 import qualified Data.ByteString.Lazy as BL
 import qualified Foreign.Lua as Lua
 import qualified Text.Pandoc.MediaBag as MB
-import Text.Pandoc.Lua.Filter (walkInlines, walkBlocks, LuaFilter)
 
 -- | Push the "pandoc" on the lua stack. Requires the `list` module to be
 -- loaded.
 pushPandocModule :: Maybe FilePath -> Lua ()
 pushPandocModule datadir = do
-  loadListModule datadir
-  script <- liftIO (moduleScript datadir "pandoc.lua")
-  status <- Lua.loadstring script
-  unless (status /= Lua.OK) $ Lua.call 0 1
+  loadScriptFromDataDir datadir "pandoc.lua"
   addFunction "_pipe" pipeFn
   addFunction "_read" readDoc
   addFunction "sha1" sha1HashFn
   addFunction "walk_block" walkBlock
   addFunction "walk_inline" walkInline
 
--- | Get the string representation of the pandoc module
-moduleScript :: Maybe FilePath -> FilePath -> IO String
-moduleScript datadir moduleFile = unpack <$>
-  runIOorExplode (setUserDataDir datadir >> readDataFile moduleFile)
-
--- Loads pandoc's list module without assigning it to a variable.
-pushListModule :: Maybe FilePath -> Lua ()
-pushListModule datadir = do
-  script <- liftIO (moduleScript datadir "List.lua")
-  status <- Lua.loadstring script
-  if status == Lua.OK
-    then Lua.call 0 1
-    else Lua.throwTopMessageAsError' ("Error while loading module `list`\n" ++)
-
-loadListModule :: Maybe FilePath -> Lua ()
-loadListModule datadir = do
-  Lua.getglobal' "package.loaded"
-  pushListModule datadir
-  Lua.setfield (-2) "pandoc.List"
-  Lua.pop 1
-
 walkElement :: (ToLuaStack a, Walkable [Inline] a, Walkable [Block] a)
             => a -> LuaFilter -> Lua NumResults
 walkElement x f = do
diff --git a/src/Text/Pandoc/Lua/Util.hs b/src/Text/Pandoc/Lua/Util.hs
index 7960c0670..5803e62dc 100644
--- a/src/Text/Pandoc/Lua/Util.hs
+++ b/src/Text/Pandoc/Lua/Util.hs
@@ -38,11 +38,18 @@ module Text.Pandoc.Lua.Util
   , PushViaCall
   , pushViaCall
   , pushViaConstructor
+  , loadScriptFromDataDir
+  , dostring'
   ) where
 
+import Control.Monad (when)
+import Data.ByteString.Char8 (unpack)
 import Foreign.Lua (FromLuaStack (..), Lua, NumArgs, StackIndex,
                     ToLuaStack (..), getglobal')
-import Foreign.Lua.Api (call, pop, rawget, rawgeti, rawset, rawseti)
+import Foreign.Lua.Api (Status, call, pop, rawget, rawgeti, rawset, rawseti)
+import Text.Pandoc.Class (readDataFile, runIOorExplode, setUserDataDir)
+
+import qualified Foreign.Lua as Lua
 
 -- | Adjust the stack index, assuming that @n@ new elements have been pushed on
 -- the stack.
@@ -107,3 +114,27 @@ pushViaCall fn = pushViaCall' fn (return ()) 0
 -- | Call a pandoc element constructor within lua, passing all given arguments.
 pushViaConstructor :: PushViaCall a => String -> a
 pushViaConstructor pandocFn = pushViaCall ("pandoc." ++ pandocFn)
+
+-- | Load a file from pandoc's data directory.
+loadScriptFromDataDir :: Maybe FilePath -> FilePath -> Lua ()
+loadScriptFromDataDir datadir scriptFile = do
+  script <- fmap unpack . Lua.liftIO . runIOorExplode $
+            setUserDataDir datadir >> readDataFile scriptFile
+  status <- dostring' script
+  when (status /= Lua.OK) .
+    Lua.throwTopMessageAsError' $ \msg ->
+      "Couldn't load '" ++ scriptFile ++ "'.\n" ++ msg
+
+-- | Load a string and immediately perform a full garbage collection. This is
+-- important to keep the program from hanging: If the program contained a call
+-- to @require@, the a new loader function was created which then become
+-- garbage. If that function is collected at an inopportune times, i.e. when the
+-- Lua API is called via a function that doesn't allow calling back into Haskell
+-- (getraw, setraw, …). The function's finalizer, and the full program, hangs
+-- when that happens.
+dostring' :: String -> Lua Status
+dostring' script = do
+  loadRes <- Lua.loadstring script
+  if loadRes == Lua.OK
+    then Lua.pcall 0 1 Nothing <* Lua.gc Lua.GCCOLLECT 0
+    else return loadRes