1 files changed, 85 insertions, 101 deletions

diff --git a/src/Text/Pandoc/Lua/Util.hs b/src/Text/Pandoc/Lua/Util.hs
index ea9ec2554..77b27b88e 100644
--- a/src/Text/Pandoc/Lua/Util.hs
+++ b/src/Text/Pandoc/Lua/Util.hs
@@ -1,4 +1,3 @@
-{-# LANGUAGE NoImplicitPrelude #-}
 {-
 Copyright © 2012-2018 John MacFarlane <jgm@berkeley.edu>
             2017-2018 Albert Krewinkel <tarleb+pandoc@moltkeplatz.de>
@@ -18,6 +17,8 @@ along with this program; if not, write to the Free Software
 Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 -}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE NoImplicitPrelude #-}
 {- |
    Module      : Text.Pandoc.Lua.Util
    Copyright   : © 2012–2018 John MacFarlane,
@@ -31,101 +32,53 @@ Lua utility functions.
 -}
 module Text.Pandoc.Lua.Util
   ( getTag
-  , getTable
-  , addValue
+  , rawField
+  , addField
   , addFunction
-  , getRawInt
-  , setRawInt
-  , addRawInt
-  , typeCheck
-  , raiseError
-  , popValue
-  , PushViaCall
-  , pushViaCall
+  , addValue
   , pushViaConstructor
   , loadScriptFromDataDir
-  , dostring'
+  , defineHowTo
+  , throwTopMessageAsError'
+  , callWithTraceback
+  , dofileWithTraceback
   ) where
 
 import Prelude
-import Control.Monad (when)
-import Control.Monad.Catch (finally)
-import Data.ByteString.Char8 (unpack)
-import Foreign.Lua (FromLuaStack (..), NumResults, Lua, NumArgs, StackIndex,
-                    ToLuaStack (..), ToHaskellFunction)
-import Foreign.Lua.Api (Status, call, pop, rawget, rawgeti, rawset, rawseti)
+import Control.Monad (unless, when)
+import Foreign.Lua ( Lua, NumArgs, NumResults, Peekable, Pushable, StackIndex
+                   , Status, ToHaskellFunction )
 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.
-adjustIndexBy :: StackIndex -> StackIndex -> StackIndex
-adjustIndexBy idx n =
-  if idx < 0
-  then idx - n
-  else idx
+import qualified Text.Pandoc.UTF8 as UTF8
 
 -- | Get value behind key from table at given index.
-getTable :: (ToLuaStack a, FromLuaStack b) => StackIndex -> a -> Lua b
-getTable idx key = do
-  push key
-  rawget (idx `adjustIndexBy` 1)
-  popValue
+rawField :: Peekable a => StackIndex -> String -> Lua a
+rawField idx key = do
+  absidx <- Lua.absindex idx
+  Lua.push key
+  Lua.rawget absidx
+  Lua.popValue
+
+-- | Add a value to the table at the top of the stack at a string-index.
+addField :: Pushable a => String -> a -> Lua ()
+addField = addValue
 
 -- | Add a key-value pair to the table at the top of the stack.
-addValue :: (ToLuaStack a, ToLuaStack b) => a -> b -> Lua ()
+addValue :: (Pushable a, Pushable b) => a -> b -> Lua ()
 addValue key value = do
-  push key
-  push value
-  rawset (-3)
+  Lua.push key
+  Lua.push value
+  Lua.rawset (Lua.nthFromTop 3)
 
 -- | Add a function to the table at the top of the stack, using the given name.
 addFunction :: ToHaskellFunction a => String -> a -> Lua ()
 addFunction name fn = do
   Lua.push name
   Lua.pushHaskellFunction fn
-  Lua.wrapHaskellFunction
   Lua.rawset (-3)
 
--- | Get value behind key from table at given index.
-getRawInt :: FromLuaStack a => StackIndex -> Int -> Lua a
-getRawInt idx key = do
-  rawgeti idx key
-  popValue
-
--- | Set numeric key/value in table at the given index
-setRawInt :: ToLuaStack a => StackIndex -> Int -> a -> Lua ()
-setRawInt idx key value = do
-  push value
-  rawseti (idx `adjustIndexBy` 1) key
-
--- | Set numeric key/value in table at the top of the stack.
-addRawInt :: ToLuaStack a => Int -> a -> Lua ()
-addRawInt = setRawInt (-1)
-
-typeCheck :: StackIndex -> Lua.Type -> Lua ()
-typeCheck idx expected = do
-  actual <- Lua.ltype idx
-  when (actual /= expected) $ do
-    expName <- Lua.typename expected
-    actName <- Lua.typename actual
-    Lua.throwLuaError $ "expected " ++ expName ++ " but got " ++ actName ++ "."
-
-raiseError :: ToLuaStack a => a -> Lua NumResults
-raiseError e = do
-  Lua.push e
-  fromIntegral <$> Lua.lerror
-
--- | Get, then pop the value at the top of the stack.
-popValue :: FromLuaStack a => Lua a
-popValue = do
-  resOrError <- Lua.peekEither (-1)
-  pop 1
-  case resOrError of
-    Left err -> Lua.throwLuaError err
-    Right x -> return x
-
 -- | Helper class for pushing a single value to the stack via a lua function.
 -- See @pushViaCall@.
 class PushViaCall a where
@@ -136,11 +89,11 @@ instance PushViaCall (Lua ()) where
     Lua.push fn
     Lua.rawget Lua.registryindex
     pushArgs
-    call num 1
+    Lua.call num 1
 
-instance (ToLuaStack a, PushViaCall b) => PushViaCall (a -> b) where
+instance (Pushable a, PushViaCall b) => PushViaCall (a -> b) where
   pushViaCall' fn pushArgs num x =
-    pushViaCall' fn (pushArgs *> push x) (num + 1)
+    pushViaCall' fn (pushArgs *> Lua.push x) (num + 1)
 
 -- | Push an value to the stack via a lua function. The lua function is called
 -- with all arguments that are passed to this function and is expected to return
@@ -155,26 +108,11 @@ 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 $
+  script <- 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, …), then the function's finalizer, and the full program,
--- will hang.
-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
+  status <- Lua.dostring script
+  when (status /= Lua.OK) $
+    throwTopMessageAsError' (("Couldn't load '" ++ scriptFile ++ "'.\n") ++)
 
 -- | Get the tag of a value. This is an optimized and specialized version of
 -- @Lua.getfield idx "tag"@. It only checks for the field on the table at index
@@ -182,8 +120,54 @@ dostring' script = do
 -- metatable.
 getTag :: StackIndex -> Lua String
 getTag idx = do
-  top <- Lua.gettop
-  hasMT <- Lua.getmetatable idx
-  push "tag"
-  if hasMT then Lua.rawget (-2) else Lua.rawget (idx `adjustIndexBy` 1)
-  peek Lua.stackTop `finally` Lua.settop top
+  -- push metatable or just the table
+  Lua.getmetatable idx >>= \hasMT -> unless hasMT (Lua.pushvalue idx)
+  Lua.push "tag"
+  Lua.rawget (Lua.nthFromTop 2)
+  Lua.tostring Lua.stackTop <* Lua.pop 2 >>= \case
+    Nothing -> Lua.throwException "untagged value"
+    Just x -> return (UTF8.toString x)
+
+-- | Modify the message at the top of the stack before throwing it as an
+-- Exception.
+throwTopMessageAsError' :: (String -> String) -> Lua a
+throwTopMessageAsError' modifier = do
+  msg <- Lua.tostring' Lua.stackTop
+  Lua.pop 2 -- remove error and error string pushed by tostring'
+  Lua.throwException (modifier (UTF8.toString msg))
+
+-- | Mark the context of a Lua computation for better error reporting.
+defineHowTo :: String -> Lua a -> Lua a
+defineHowTo ctx = Lua.withExceptionMessage (("Could not " <> ctx <> ": ") <>)
+
+-- | Like @'Lua.pcall'@, but uses a predefined error handler which adds a
+-- traceback on error.
+pcallWithTraceback :: NumArgs -> NumResults -> Lua Status
+pcallWithTraceback nargs nresults = do
+  let traceback' :: Lua NumResults
+      traceback' = do
+        l <- Lua.state
+        msg <- Lua.tostring' (Lua.nthFromBottom 1)
+        Lua.traceback l (Just (UTF8.toString msg)) 2
+        return 1
+  tracebackIdx <- Lua.absindex (Lua.nthFromTop (Lua.fromNumArgs nargs + 1))
+  Lua.pushHaskellFunction traceback'
+  Lua.insert tracebackIdx
+  result <- Lua.pcall nargs nresults (Just tracebackIdx)
+  Lua.remove tracebackIdx
+  return result
+
+-- | Like @'Lua.call'@, but adds a traceback to the error message (if any).
+callWithTraceback :: NumArgs -> NumResults -> Lua ()
+callWithTraceback nargs nresults = do
+  result <- pcallWithTraceback nargs nresults
+  when (result /= Lua.OK) Lua.throwTopMessage
+
+-- | Run the given string as a Lua program, while also adding a traceback to the
+-- error message if an error occurs.
+dofileWithTraceback :: FilePath -> Lua Status
+dofileWithTraceback fp = do
+  loadRes <- Lua.loadfile fp
+  case loadRes of
+    Lua.OK -> pcallWithTraceback 0 Lua.multret
+    _ -> return loadRes