1 files changed, 85 insertions, 49 deletions
diff --git a/src/Text/Pandoc/Lua/Init.hs b/src/Text/Pandoc/Lua/Init.hs
index baa6f0295..835da1fc9 100644
--- a/src/Text/Pandoc/Lua/Init.hs
+++ b/src/Text/Pandoc/Lua/Init.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE LambdaCase        #-}
+{-# LANGUAGE OverloadedStrings #-}
 {- |
    Module      : Text.Pandoc.Lua
    Copyright   : Copyright © 2017-2021 Albert Krewinkel
@@ -12,25 +14,24 @@ module Text.Pandoc.Lua.Init
   ( runLua
   ) where
 
-import Control.Monad (when)
-import Control.Monad.Catch (try)
+import Control.Monad (forM, forM_, when)
+import Control.Monad.Catch (throwM, try)
 import Control.Monad.Trans (MonadIO (..))
-import Data.Data (Data, dataTypeConstrs, dataTypeOf, showConstr)
-import Foreign.Lua (Lua)
+import Data.Maybe (catMaybes)
+import HsLua as Lua hiding (status, try)
 import GHC.IO.Encoding (getForeignEncoding, setForeignEncoding, utf8)
-import Text.Pandoc.Class.PandocMonad (readDataFile)
-import Text.Pandoc.Class.PandocIO (PandocIO)
-import Text.Pandoc.Error (PandocError)
+import Text.Pandoc.Class.PandocMonad (PandocMonad, readDataFile)
+import Text.Pandoc.Error (PandocError (PandocLuaError))
 import Text.Pandoc.Lua.Packages (installPandocPackageSearcher)
 import Text.Pandoc.Lua.PandocLua (PandocLua, liftPandocLua, runPandocLua)
-import Text.Pandoc.Lua.Util (throwTopMessageAsError')
-import qualified Foreign.Lua as Lua
-import qualified Text.Pandoc.Definition as Pandoc
+import qualified Data.Text as T
+import qualified Lua.LPeg as LPeg
 import qualified Text.Pandoc.Lua.Module.Pandoc as ModulePandoc
 
 -- | Run the lua interpreter, using pandoc's default way of environment
 -- initialization.
-runLua :: Lua a -> PandocIO (Either PandocError a)
+runLua :: (PandocMonad m, MonadIO m)
+       => LuaE PandocError a -> m (Either PandocError a)
 runLua luaOp = do
   enc <- liftIO $ getForeignEncoding <* setForeignEncoding utf8
   res <- runPandocLua . try $ do
@@ -39,12 +40,27 @@ runLua luaOp = do
   liftIO $ setForeignEncoding enc
   return res
 
+-- | Modules that are loaded at startup and assigned to fields in the
+-- pandoc module.
+loadedModules :: [(Name, Name)]
+loadedModules =
+  [ ("pandoc.List", "List")
+  , ("pandoc.mediabag", "mediabag")
+  , ("pandoc.path", "path")
+  , ("pandoc.system", "system")
+  , ("pandoc.types", "types")
+  , ("pandoc.utils", "utils")
+  , ("text", "text")
+  ]
+
 -- | Initialize the lua state with all required values
 initLuaState :: PandocLua ()
 initLuaState = do
   liftPandocLua Lua.openlibs
   installPandocPackageSearcher
   initPandocModule
+  installLpegSearcher
+  setGlobalModules
   loadInitScript "init.lua"
  where
   initPandocModule :: PandocLua ()
@@ -53,12 +69,16 @@ initLuaState = do
     ModulePandoc.pushModule
     -- register as loaded module
     liftPandocLua $ do
-      Lua.pushvalue Lua.stackTop
-      Lua.getfield Lua.registryindex Lua.loadedTableRegistryField
-      Lua.setfield (Lua.nthFromTop 2) "pandoc"
-      Lua.pop 1
-    -- copy constructors into registry
-    putConstructorsInRegistry
+      Lua.getfield Lua.registryindex Lua.loaded
+      Lua.pushvalue (Lua.nth 2)
+      Lua.setfield (Lua.nth 2) "pandoc"
+      Lua.pop 1  -- remove LOADED table
+    -- load modules and add them to the `pandoc` module table.
+    liftPandocLua $ forM_ loadedModules $ \(pkgname, fieldname) -> do
+      Lua.getglobal "require"
+      Lua.pushName pkgname
+      Lua.call 1 1
+      Lua.setfield (nth 2) fieldname
     -- assign module to global variable
     liftPandocLua $ Lua.setglobal "pandoc"
 
@@ -66,38 +86,54 @@ initLuaState = do
   loadInitScript scriptFile = do
     script <- readDataFile scriptFile
     status <- liftPandocLua $ Lua.dostring script
-    when (status /= Lua.OK) . liftPandocLua $
-      throwTopMessageAsError'
-      (("Couldn't load '" ++ scriptFile ++ "'.\n") ++)
+    when (status /= Lua.OK) . liftPandocLua $ do
+      err <- popException
+      let prefix = "Couldn't load '" <> T.pack scriptFile <> "':\n"
+      throwM . PandocLuaError . (prefix <>) $ case err of
+        PandocLuaError msg -> msg
+        _                  -> T.pack $ show err
 
+  setGlobalModules :: PandocLua ()
+  setGlobalModules = liftPandocLua $ do
+    let globalModules =
+          [ ("lpeg", LPeg.luaopen_lpeg_ptr)  -- must be loaded first
+          , ("re", LPeg.luaopen_re_ptr)      -- re depends on lpeg
+          ]
+    loadedBuiltInModules <- fmap catMaybes . forM globalModules $
+      \(pkgname, luaopen) -> do
+        Lua.pushcfunction luaopen
+        usedBuiltIn <- Lua.pcall 0 1 Nothing >>= \case
+          OK -> do               -- all good, loading succeeded
+            -- register as loaded module so later modules can rely on this
+            Lua.getfield Lua.registryindex Lua.loaded
+            Lua.pushvalue (Lua.nth 2)
+            Lua.setfield (Lua.nth 2) pkgname
+            Lua.pop 1  -- pop _LOADED
+            return True
+          _  -> do               -- built-in library failed, load system lib
+            Lua.pop 1  -- ignore error message
+            -- Try loading via the normal package loading mechanism.
+            Lua.getglobal "require"
+            Lua.pushName pkgname
+            Lua.call 1 1  -- Throws an exception if loading failed again!
+            return False
 
--- | AST elements are marshaled via normal constructor functions in the
--- @pandoc@ module. However, accessing Lua globals from Haskell is
--- expensive (due to error handling). Accessing the Lua registry is much
--- cheaper, which is why the constructor functions are copied into the
--- Lua registry and called from there.
---
--- This function expects the @pandoc@ module to be at the top of the
--- stack.
-putConstructorsInRegistry :: PandocLua ()
-putConstructorsInRegistry = liftPandocLua $ do
-  constrsToReg $ Pandoc.Pandoc mempty mempty
-  constrsToReg $ Pandoc.Str mempty
-  constrsToReg $ Pandoc.Para mempty
-  constrsToReg $ Pandoc.Meta mempty
-  constrsToReg $ Pandoc.MetaList mempty
-  constrsToReg $ Pandoc.Citation mempty mempty mempty Pandoc.AuthorInText 0 0
-  putInReg "Attr"  -- used for Attr type alias
-  putInReg "ListAttributes"  -- used for ListAttributes type alias
-  putInReg "List"  -- pandoc.List
-  putInReg "SimpleTable"  -- helper for backward-compatible table handling
- where
-  constrsToReg :: Data a => a -> Lua ()
-  constrsToReg = mapM_ (putInReg . showConstr) . dataTypeConstrs . dataTypeOf
+        -- Module on top of stack. Register as global
+        Lua.setglobal pkgname
+        return $ if usedBuiltIn then Just pkgname else Nothing
+
+    -- Remove module entry from _LOADED table in registry if we used a
+    -- built-in library. This ensures that later calls to @require@ will
+    -- prefer the shared library, if any.
+    forM_ loadedBuiltInModules $ \pkgname -> do
+      Lua.getfield Lua.registryindex Lua.loaded
+      Lua.pushnil
+      Lua.setfield (Lua.nth 2) pkgname
+      Lua.pop 1  -- registry
 
-  putInReg :: String -> Lua ()
-  putInReg name = do
-    Lua.push ("pandoc." ++ name) -- name in registry
-    Lua.push name -- in pandoc module
-    Lua.rawget (Lua.nthFromTop 3)
-    Lua.rawset Lua.registryindex
+  installLpegSearcher :: PandocLua ()
+  installLpegSearcher = liftPandocLua $ do
+    Lua.getglobal' "package.searchers"
+    Lua.pushHaskellFunction $ Lua.state >>= liftIO . LPeg.lpeg_searcher
+    Lua.rawseti (Lua.nth 2) . (+1) . fromIntegral =<< Lua.rawlen (Lua.nth 2)
+    Lua.pop 1  -- remove 'package.searchers' from stack