{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{- |
   Module      : Text.Pandoc.Lua.Util
   Copyright   : © 2012–2020 John MacFarlane,
                 © 2017-2020 Albert Krewinkel
   License     : GNU GPL, version 2 or above

   Maintainer  : Albert Krewinkel <tarleb+pandoc@moltkeplatz.de>
   Stability   : alpha

Lua utility functions.
-}
module Text.Pandoc.Lua.Util
  ( getTag
  , rawField
  , addField
  , addFunction
  , addValue
  , pushViaConstructor
  , loadScriptFromDataDir
  , defineHowTo
  , throwTopMessageAsError'
  , callWithTraceback
  , dofileWithTraceback
  ) where

import Control.Monad (unless, when)
import Foreign.Lua ( Lua, NumArgs, NumResults, Peekable, Pushable, StackIndex
                   , Status, ToHaskellFunction )
import Text.Pandoc.Class.PandocIO (runIOorExplode)
import Text.Pandoc.Class.PandocMonad (readDataFile, setUserDataDir)
import qualified Foreign.Lua as Lua
import qualified Text.Pandoc.UTF8 as UTF8
import Data.Text (Text)

-- | Get value behind key from table at given index.
rawField :: Peekable a => StackIndex -> String -> Lua a
rawField :: StackIndex -> String -> Lua a
rawField idx :: StackIndex
idx key :: String
key = do
  StackIndex
absidx <- StackIndex -> Lua StackIndex
Lua.absindex StackIndex
idx
  String -> Lua ()
forall a. Pushable a => a -> Lua ()
Lua.push String
key
  StackIndex -> Lua ()
Lua.rawget StackIndex
absidx
  Lua a
forall a. Peekable a => Lua a
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 :: String -> a -> Lua ()
addField = String -> a -> Lua ()
forall a b. (Pushable a, Pushable b) => a -> b -> Lua ()
addValue

-- | Add a key-value pair to the table at the top of the stack.
addValue :: (Pushable a, Pushable b) => a -> b -> Lua ()
addValue :: a -> b -> Lua ()
addValue key :: a
key value :: b
value = do
  a -> Lua ()
forall a. Pushable a => a -> Lua ()
Lua.push a
key
  b -> Lua ()
forall a. Pushable a => a -> Lua ()
Lua.push b
value
  StackIndex -> Lua ()
Lua.rawset (CInt -> StackIndex
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 :: String -> a -> Lua ()
addFunction name :: String
name fn :: a
fn = do
  String -> Lua ()
forall a. Pushable a => a -> Lua ()
Lua.push String
name
  a -> Lua ()
forall a. ToHaskellFunction a => a -> Lua ()
Lua.pushHaskellFunction a
fn
  StackIndex -> Lua ()
Lua.rawset (-3)

-- | Helper class for pushing a single value to the stack via a lua function.
-- See @pushViaCall@.
class PushViaCall a where
  pushViaCall' :: String -> Lua () -> NumArgs -> a

instance PushViaCall (Lua ()) where
  pushViaCall' :: String -> Lua () -> NumArgs -> Lua ()
pushViaCall' fn :: String
fn pushArgs :: Lua ()
pushArgs num :: NumArgs
num = do
    String -> Lua ()
forall a. Pushable a => a -> Lua ()
Lua.push String
fn
    StackIndex -> Lua ()
Lua.rawget StackIndex
Lua.registryindex
    Lua ()
pushArgs
    NumArgs -> NumResults -> Lua ()
Lua.call NumArgs
num 1

instance (Pushable a, PushViaCall b) => PushViaCall (a -> b) where
  pushViaCall' :: String -> Lua () -> NumArgs -> a -> b
pushViaCall' fn :: String
fn pushArgs :: Lua ()
pushArgs num :: NumArgs
num x :: a
x =
    String -> Lua () -> NumArgs -> b
forall a. PushViaCall a => String -> Lua () -> NumArgs -> a
pushViaCall' String
fn (Lua ()
pushArgs Lua () -> Lua () -> Lua ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> a -> Lua ()
forall a. Pushable a => a -> Lua ()
Lua.push a
x) (NumArgs
num NumArgs -> NumArgs -> NumArgs
forall a. Num a => a -> a -> a
+ 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
-- a single value.
pushViaCall :: PushViaCall a => String -> a
pushViaCall :: String -> a
pushViaCall fn :: String
fn = String -> Lua () -> NumArgs -> a
forall a. PushViaCall a => String -> Lua () -> NumArgs -> a
pushViaCall' String
fn (() -> Lua ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()) 0

-- | Call a pandoc element constructor within lua, passing all given arguments.
pushViaConstructor :: PushViaCall a => String -> a
pushViaConstructor :: String -> a
pushViaConstructor pandocFn :: String
pandocFn = String -> a
forall a. PushViaCall a => String -> a
pushViaCall ("pandoc." String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
pandocFn)

-- | Load a file from pandoc's data directory.
loadScriptFromDataDir :: Maybe FilePath -> FilePath -> Lua ()
loadScriptFromDataDir :: Maybe String -> String -> Lua ()
loadScriptFromDataDir datadir :: Maybe String
datadir scriptFile :: String
scriptFile = do
  ByteString
script <- IO ByteString -> Lua ByteString
forall (m :: * -> *) a. MonadIO m => IO a -> m a
Lua.liftIO (IO ByteString -> Lua ByteString)
-> (PandocIO ByteString -> IO ByteString)
-> PandocIO ByteString
-> Lua ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PandocIO ByteString -> IO ByteString
forall a. PandocIO a -> IO a
runIOorExplode (PandocIO ByteString -> Lua ByteString)
-> PandocIO ByteString -> Lua ByteString
forall a b. (a -> b) -> a -> b
$
            Maybe String -> PandocIO ()
forall (m :: * -> *). PandocMonad m => Maybe String -> m ()
setUserDataDir Maybe String
datadir PandocIO () -> PandocIO ByteString -> PandocIO ByteString
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> String -> PandocIO ByteString
forall (m :: * -> *). PandocMonad m => String -> m ByteString
readDataFile String
scriptFile
  Status
status <- ByteString -> Lua Status
Lua.dostring ByteString
script
  Bool -> Lua () -> Lua ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Status
status Status -> Status -> Bool
forall a. Eq a => a -> a -> Bool
/= Status
Lua.OK) (Lua () -> Lua ()) -> Lua () -> Lua ()
forall a b. (a -> b) -> a -> b
$
    (String -> String) -> Lua ()
forall a. (String -> String) -> Lua a
throwTopMessageAsError' (("Couldn't load '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
scriptFile String -> String -> String
forall a. [a] -> [a] -> [a]
++ "'.\n") String -> String -> String
forall a. [a] -> [a] -> [a]
++)

-- | 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
-- @idx@ and on its metatable, also ignoring any @__index@ value on the
-- metatable.
getTag :: StackIndex -> Lua String
getTag :: StackIndex -> Lua String
getTag idx :: StackIndex
idx = do
  -- push metatable or just the table
  StackIndex -> Lua Bool
Lua.getmetatable StackIndex
idx Lua Bool -> (Bool -> Lua ()) -> Lua ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \hasMT :: Bool
hasMT -> Bool -> Lua () -> Lua ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
hasMT (StackIndex -> Lua ()
Lua.pushvalue StackIndex
idx)
  Text -> Lua ()
forall a. Pushable a => a -> Lua ()
Lua.push ("tag" :: Text)
  StackIndex -> Lua ()
Lua.rawget (CInt -> StackIndex
Lua.nthFromTop 2)
  StackIndex -> Lua (Maybe ByteString)
Lua.tostring StackIndex
Lua.stackTop Lua (Maybe ByteString) -> Lua () -> Lua (Maybe ByteString)
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* StackIndex -> Lua ()
Lua.pop 2 Lua (Maybe ByteString)
-> (Maybe ByteString -> Lua String) -> Lua String
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Nothing -> String -> Lua String
forall a. String -> Lua a
Lua.throwException "untagged value"
    Just x :: ByteString
x -> String -> Lua String
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> String
UTF8.toString ByteString
x)

-- | Modify the message at the top of the stack before throwing it as an
-- Exception.
throwTopMessageAsError' :: (String -> String) -> Lua a
throwTopMessageAsError' :: (String -> String) -> Lua a
throwTopMessageAsError' modifier :: String -> String
modifier = do
  ByteString
msg <- StackIndex -> Lua ByteString
Lua.tostring' StackIndex
Lua.stackTop
  StackIndex -> Lua ()
Lua.pop 2 -- remove error and error string pushed by tostring'
  String -> Lua a
forall a. String -> Lua a
Lua.throwException (String -> String
modifier (ByteString -> String
UTF8.toString ByteString
msg))

-- | Mark the context of a Lua computation for better error reporting.
defineHowTo :: String -> Lua a -> Lua a
defineHowTo :: String -> Lua a -> Lua a
defineHowTo ctx :: String
ctx = (String -> String) -> Lua a -> Lua a
forall a. (String -> String) -> Lua a -> Lua a
Lua.withExceptionMessage (("Could not " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
ctx String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ": ") String -> String -> String
forall a. Semigroup a => a -> a -> a
<>)

-- | Like @'Lua.pcall'@, but uses a predefined error handler which adds a
-- traceback on error.
pcallWithTraceback :: NumArgs -> NumResults -> Lua Status
pcallWithTraceback :: NumArgs -> NumResults -> Lua Status
pcallWithTraceback nargs :: NumArgs
nargs nresults :: NumResults
nresults = do
  let traceback' :: Lua NumResults
      traceback' :: Lua NumResults
traceback' = do
        State
l <- Lua State
Lua.state
        ByteString
msg <- StackIndex -> Lua ByteString
Lua.tostring' (CInt -> StackIndex
Lua.nthFromBottom 1)
        State -> Maybe String -> Int -> Lua ()
Lua.traceback State
l (String -> Maybe String
forall a. a -> Maybe a
Just (ByteString -> String
UTF8.toString ByteString
msg)) 2
        NumResults -> Lua NumResults
forall (m :: * -> *) a. Monad m => a -> m a
return 1
  StackIndex
tracebackIdx <- StackIndex -> Lua StackIndex
Lua.absindex (CInt -> StackIndex
Lua.nthFromTop (NumArgs -> CInt
Lua.fromNumArgs NumArgs
nargs CInt -> CInt -> CInt
forall a. Num a => a -> a -> a
+ 1))
  Lua NumResults -> Lua ()
forall a. ToHaskellFunction a => a -> Lua ()
Lua.pushHaskellFunction Lua NumResults
traceback'
  StackIndex -> Lua ()
Lua.insert StackIndex
tracebackIdx
  Status
result <- NumArgs -> NumResults -> Maybe StackIndex -> Lua Status
Lua.pcall NumArgs
nargs NumResults
nresults (StackIndex -> Maybe StackIndex
forall a. a -> Maybe a
Just StackIndex
tracebackIdx)
  StackIndex -> Lua ()
Lua.remove StackIndex
tracebackIdx
  Status -> Lua Status
forall (m :: * -> *) a. Monad m => a -> m a
return Status
result

-- | Like @'Lua.call'@, but adds a traceback to the error message (if any).
callWithTraceback :: NumArgs -> NumResults -> Lua ()
callWithTraceback :: NumArgs -> NumResults -> Lua ()
callWithTraceback nargs :: NumArgs
nargs nresults :: NumResults
nresults = do
  Status
result <- NumArgs -> NumResults -> Lua Status
pcallWithTraceback NumArgs
nargs NumResults
nresults
  Bool -> Lua () -> Lua ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Status
result Status -> Status -> Bool
forall a. Eq a => a -> a -> Bool
/= Status
Lua.OK) Lua ()
forall a. Lua a
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 :: String -> Lua Status
dofileWithTraceback fp :: String
fp = do
  Status
loadRes <- String -> Lua Status
Lua.loadfile String
fp
  case Status
loadRes of
    Lua.OK -> NumArgs -> NumResults -> Lua Status
pcallWithTraceback 0 NumResults
Lua.multret
    _ -> Status -> Lua Status
forall (m :: * -> *) a. Monad m => a -> m a
return Status
loadRes