Lua: support extensions in custom writers

Custom writers can define the extensions that they support via the
global `writer_extensions`. The variable's value must be a table with
all supported extensions as keys, and their default status as values.

E.g., the below specifies that the writer support the extensions `smart`
and `sourcepos`, but only the `smart` extension is enabled by default:

    writer_extensions = {
      smart = true,
      sourcepos = false,
    }

9 files changed, 119 insertions, 49 deletions

diff --git a/pandoc-lua-engine/pandoc-lua-engine.cabal b/pandoc-lua-engine/pandoc-lua-engine.cabal
index a122b935c..3c632026b 100644
--- a/pandoc-lua-engine/pandoc-lua-engine.cabal
+++ b/pandoc-lua-engine/pandoc-lua-engine.cabal
@@ -22,6 +22,7 @@ extra-source-files:  README.md
                    , test/bytestring.bin
                    , test/bytestring.lua
                    , test/bytestring-reader.lua
+                   , test/extensions.lua
                    , test/lua/*.lua
                    , test/lua/module/*.lua
                    , test/lua/module/partial.test
@@ -66,7 +67,7 @@ library
                      , Text.Pandoc.Lua.Init
                      , Text.Pandoc.Lua.Marshal.CommonState
                      , Text.Pandoc.Lua.Marshal.Context
-                     , Text.Pandoc.Lua.Marshal.Extensions
+                     , Text.Pandoc.Lua.Marshal.Format
                      , Text.Pandoc.Lua.Marshal.PandocError
                      , Text.Pandoc.Lua.Marshal.ReaderOptions
                      , Text.Pandoc.Lua.Marshal.Reference
diff --git a/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/Extensions.hs b/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/Extensions.hs
deleted file mode 100644
index 99e5bc442..000000000
--- a/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/Extensions.hs
+++ /dev/null
@@ -1,32 +0,0 @@
-{-# OPTIONS_GHC -fno-warn-orphans #-}
-{- |
-   Module      : Text.Pandoc.Lua.Marshaling.Extensions
-   Copyright   : © 2022 Albert Krewinkel
-   License     : GPL-2.0-or-later
-   Maintainer  : Albert Krewinkel <tarleb+pandoc@moltkeplatz.de>
-
-Marshaling functions and instance for 'Extensions'.
--}
-module Text.Pandoc.Lua.Marshal.Extensions
-  ( peekExtensions
-  , pushExtensions
-  ) where
-
-import HsLua
-import Text.Pandoc.Extensions (Extensions)
-
--- | Retrieves an 'Extensions' set from the Lua stack.
-peekExtensions :: LuaError e => Peeker e Extensions
-peekExtensions = peekViaJSON
-{-# INLINE peekExtensions #-}
-
--- | Pushes a set of 'Extensions' to the top of the Lua stack.
-pushExtensions :: LuaError e => Pusher e Extensions
-pushExtensions = pushViaJSON
-{-# INLINE pushExtensions #-}
-
-instance Peekable Extensions where
-  safepeek = peekExtensions
-
-instance Pushable Extensions where
-  push = pushExtensions
diff --git a/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/Format.hs b/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/Format.hs
new file mode 100644
index 000000000..a71aeb952
--- /dev/null
+++ b/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/Format.hs
@@ -0,0 +1,57 @@
+{-# LANGUAGE LambdaCase           #-}
+{-# LANGUAGE OverloadedStrings    #-}
+{-# OPTIONS_GHC -fno-warn-orphans #-}
+{- |
+   Module      : Text.Pandoc.Lua.Marshaling.Format
+   Copyright   : © 2022 Albert Krewinkel
+   License     : GPL-2.0-or-later
+   Maintainer  : Albert Krewinkel <tarleb+pandoc@moltkeplatz.de>
+
+Marshaling functions and instance for format related types, including
+'Extensions' and 'ExtensionConfig'.
+-}
+module Text.Pandoc.Lua.Marshal.Format
+  ( peekExtensions
+  , pushExtensions
+  , peekExtensionsConfig
+  ) where
+
+import HsLua
+import Text.Pandoc.Extensions (Extension, Extensions, extensionsFromList, readExtension)
+import Text.Pandoc.Format (ExtensionsConfig (..))
+import qualified HsLua.Core.Utf8 as UTF8
+
+-- | Retrieves an 'Extensions' set from the Lua stack.
+peekExtension :: LuaError e => Peeker e Extension
+peekExtension idx = do
+  extString <- peekString idx
+  case readExtension extString of
+    Just ext -> return ext
+    Nothing  -> failPeek . UTF8.fromString $
+                "Unknown extension: " <> extString
+{-# INLINE peekExtension #-}
+
+-- | Retrieves an 'Extensions' set from the Lua stack.
+peekExtensions :: LuaError e => Peeker e Extensions
+peekExtensions = peekViaJSON
+{-# INLINE peekExtensions #-}
+
+-- | Pushes a set of 'Extensions' to the top of the Lua stack.
+pushExtensions :: LuaError e => Pusher e Extensions
+pushExtensions = pushViaJSON
+{-# INLINE pushExtensions #-}
+
+instance Peekable Extensions where
+  safepeek = peekExtensions
+
+instance Pushable Extensions where
+  push = pushExtensions
+
+-- | Retrieves an 'ExtensionsConfig' value from the Lua stack.
+peekExtensionsConfig :: LuaError e => Peeker e ExtensionsConfig
+peekExtensionsConfig idx = do
+  exts <- peekKeyValuePairs peekExtension peekBool idx
+  return $ ExtensionsConfig
+    { extsDefault   = extensionsFromList . map fst $ filter snd exts
+    , extsSupported = extensionsFromList . map fst $ exts
+    }
diff --git a/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/ReaderOptions.hs b/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/ReaderOptions.hs
index 0d23e0fb7..8a02a6d7e 100644
--- a/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/ReaderOptions.hs
+++ b/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/ReaderOptions.hs
@@ -22,7 +22,7 @@ module Text.Pandoc.Lua.Marshal.ReaderOptions
 
 import Data.Default (def)
 import HsLua as Lua
-import Text.Pandoc.Lua.Marshal.Extensions (peekExtensions, pushExtensions)
+import Text.Pandoc.Lua.Marshal.Format (peekExtensions, pushExtensions)
 import Text.Pandoc.Lua.Marshal.List (pushPandocList)
 import Text.Pandoc.Options (ReaderOptions (..))
 
diff --git a/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/WriterOptions.hs b/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/WriterOptions.hs
index bbd878907..a96c3d4b8 100644
--- a/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/WriterOptions.hs
+++ b/pandoc-lua-engine/src/Text/Pandoc/Lua/Marshal/WriterOptions.hs
@@ -21,7 +21,7 @@ module Text.Pandoc.Lua.Marshal.WriterOptions
 import Control.Applicative (optional)
 import Data.Default (def)
 import HsLua as Lua
-import Text.Pandoc.Lua.Marshal.Extensions (peekExtensions, pushExtensions)
+import Text.Pandoc.Lua.Marshal.Format (peekExtensions, pushExtensions)
 import Text.Pandoc.Lua.Marshal.List (pushPandocList)
 import Text.Pandoc.Lua.Marshal.Template (peekTemplate, pushTemplate)
 import Text.Pandoc.Options (WriterOptions (..))
diff --git a/pandoc-lua-engine/src/Text/Pandoc/Lua/Module/Format.hs b/pandoc-lua-engine/src/Text/Pandoc/Lua/Module/Format.hs
index 7eb66b3ca..d9d413e09 100644
--- a/pandoc-lua-engine/src/Text/Pandoc/Lua/Module/Format.hs
+++ b/pandoc-lua-engine/src/Text/Pandoc/Lua/Module/Format.hs
@@ -13,9 +13,8 @@ module Text.Pandoc.Lua.Module.Format
 
 import HsLua
 import Text.Pandoc.Error (PandocError)
-import Text.Pandoc.Extensions
-  ( getAllExtensions, getDefaultExtensions )
-import Text.Pandoc.Lua.Marshal.Extensions (pushExtensions)
+import Text.Pandoc.Extensions (getAllExtensions, getDefaultExtensions)
+import Text.Pandoc.Lua.Marshal.Format (pushExtensions)
 import Text.Pandoc.Lua.PandocLua ()
 
 import qualified Data.Text as T
diff --git a/pandoc-lua-engine/src/Text/Pandoc/Lua/Writer.hs b/pandoc-lua-engine/src/Text/Pandoc/Lua/Writer.hs
index f2d75b905..a9a044fe6 100644
--- a/pandoc-lua-engine/src/Text/Pandoc/Lua/Writer.hs
+++ b/pandoc-lua-engine/src/Text/Pandoc/Lua/Writer.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE LambdaCase          #-}
 {-# LANGUAGE OverloadedStrings   #-}
+{-# LANGUAGE TupleSections       #-}
 {-# LANGUAGE TypeApplications    #-}
 {- |
    Module      : Text.Pandoc.Lua.Writer
@@ -25,14 +26,16 @@ import HsLua.Core.Run (newGCManagedState, withGCManagedState)
 import Control.Monad.IO.Class (MonadIO)
 import Text.Pandoc.Class (PandocMonad, findFileWithDataFallback)
 import Text.Pandoc.Error (PandocError)
+import Text.Pandoc.Format (ExtensionsConfig (..))
 import Text.Pandoc.Lua.Global (Global (..), setGlobals)
 import Text.Pandoc.Lua.Init (runLuaWith)
+import Text.Pandoc.Lua.Marshal.Format (peekExtensionsConfig)
 import Text.Pandoc.Writers (Writer (..))
 import qualified Text.Pandoc.Lua.Writer.Classic as Classic
 
 -- | Convert Pandoc to custom markup.
 writeCustom :: (PandocMonad m, MonadIO m)
-            => FilePath -> m (Writer m)
+            => FilePath -> m (Writer m, ExtensionsConfig)
 writeCustom luaFile = do
   luaState <- liftIO newGCManagedState
   luaFile' <- fromMaybe luaFile <$> findFileWithDataFallback "writers" luaFile
@@ -55,19 +58,23 @@ writeCustom luaFile = do
 
     let writerField = "PANDOC Writer function"
 
+    extsConf <- rawgetglobal "writer_extensions" >>= \case
+      TypeNil   -> pure $ ExtensionsConfig mempty mempty
+      _         -> forcePeek $ peekExtensionsConfig top `lastly` pop 1
+
     rawgetglobal "Writer" >>= \case
       TypeNil -> rawgetglobal "ByteStringWriter" >>= \case
         TypeNil -> do
           -- Neither `Writer` nor `BinaryWriter` are defined. Try to
           -- use the file as a classic writer.
           pop 1  -- remove nil
-          return . TextWriter $ \opts doc ->
+          pure $ (,extsConf) . TextWriter $ \opts doc ->
             liftIO $ withGCManagedState luaState $ do
               Classic.runCustom @PandocError opts doc
         _ -> do
           -- Binary writer. Writer function is on top of the stack.
           setfield registryindex writerField
-          return . ByteStringWriter $ \opts doc ->
+          pure $ (,extsConf) . ByteStringWriter $ \opts doc ->
             -- Call writer with document and writer options as arguments.
             liftIO $ withGCManagedState luaState $ do
               getfield registryindex writerField
@@ -78,7 +85,7 @@ writeCustom luaFile = do
       _ -> do
         -- New-type text writer. Writer function is on top of the stack.
         setfield registryindex writerField
-        return . TextWriter $ \opts doc ->
+        pure $ (,extsConf) . TextWriter $ \opts doc ->
           liftIO $ withGCManagedState luaState $ do
             getfield registryindex writerField
             push doc
diff --git a/pandoc-lua-engine/test/Tests/Lua/Writer.hs b/pandoc-lua-engine/test/Tests/Lua/Writer.hs
index 18d4e700d..8b6e82816 100644
--- a/pandoc-lua-engine/test/Tests/Lua/Writer.hs
+++ b/pandoc-lua-engine/test/Tests/Lua/Writer.hs
@@ -1,12 +1,10 @@
 {-# LANGUAGE LambdaCase        #-}
+{-# LANGUAGE OverloadedStrings #-}
 {- |
 Module      : Tests.Lua.Writer
 Copyright   : © 2019-2022 Albert Krewinkel
 License     : GNU GPL, version 2 or above
-
 Maintainer  : Albert Krewinkel <albert@zeitkraut.de>
-Stability   : alpha
-Portability : portable
 
 Tests for custom Lua writers.
 -}
@@ -14,13 +12,19 @@ module Tests.Lua.Writer (tests) where
 
 import Data.Default (Default (def))
 import Text.Pandoc.Class (runIOorExplode, readFileStrict)
+import Text.Pandoc.Extensions (Extension (..))
+import Text.Pandoc.Format (ExtensionsDiff (..), FlavoredFormat (..),
+                           applyExtensionsDiff)
 import Text.Pandoc.Lua (writeCustom)
+import Text.Pandoc.Options (WriterOptions (..))
 import Text.Pandoc.Readers (readNative)
 import Text.Pandoc.Writers (Writer (ByteStringWriter, TextWriter))
 import Test.Tasty (TestTree)
 import Test.Tasty.Golden (goldenVsString)
+import Test.Tasty.HUnit (testCase, (@?=))
 
 import qualified Data.ByteString.Lazy as BL
+import qualified Text.Pandoc.Builder as B
 import qualified Text.Pandoc.UTF8 as UTF8
 
 tests :: [TestTree]
@@ -31,7 +35,7 @@ tests =
         source <- UTF8.toText <$> readFileStrict "testsuite.native"
         doc <- readNative def source
         txt <- writeCustom "sample.lua" >>= \case
-          TextWriter f -> f def doc
+          (TextWriter f, _) -> f def doc
           _            -> error "Expected a text writer"
         pure $ BL.fromStrict (UTF8.fromText txt))
 
@@ -41,7 +45,7 @@ tests =
         source <- UTF8.toText <$> readFileStrict "tables.native"
         doc <- readNative def source
         txt <- writeCustom "sample.lua" >>= \case
-          TextWriter f -> f def doc
+          (TextWriter f, _) -> f def doc
           _            -> error "Expected a text writer"
         pure $ BL.fromStrict (UTF8.fromText txt))
 
@@ -49,7 +53,29 @@ tests =
     "bytestring.bin"
     (runIOorExplode $ do
         txt <- writeCustom "bytestring.lua" >>= \case
-          ByteStringWriter f -> f def mempty
-          _                  -> error "Expected a bytestring writer"
+          (ByteStringWriter f, _) -> f def mempty
+          _                       -> error "Expected a bytestring writer"
         pure txt)
+
+  , testCase "preset extensions" $ do
+      let ediff = ExtensionsDiff{extsToEnable = [], extsToDisable = []}
+      let format = FlavoredFormat "extensions.lua" ediff
+      result <- runIOorExplode $ writeCustom "extensions.lua" >>= \case
+          (TextWriter write, extsConf) -> do
+            exts <- applyExtensionsDiff extsConf format
+            write def{writerExtensions = exts} (B.doc mempty)
+          _                        -> error "Expected a text writer"
+      result @?= "smart extension is enabled;\ncitations extension is disabled\n"
+  , testCase "modified extensions" $ do
+      let ediff = ExtensionsDiff
+            { extsToEnable = [Ext_citations]
+            , extsToDisable = []
+            }
+      let format = FlavoredFormat "extensions.lua" ediff
+      result <- runIOorExplode $ writeCustom "extensions.lua" >>= \case
+          (TextWriter write, extsConf) -> do
+            exts <- applyExtensionsDiff extsConf format
+            write def{writerExtensions = exts} (B.doc mempty)
+          _                        -> error "Expected a text writer"
+      result @?= "smart extension is enabled;\ncitations extension is enabled\n"
   ]
diff --git a/pandoc-lua-engine/test/extensions.lua b/pandoc-lua-engine/test/extensions.lua
new file mode 100644
index 000000000..cea9a45a1
--- /dev/null
+++ b/pandoc-lua-engine/test/extensions.lua
@@ -0,0 +1,12 @@
+function Writer (doc, opts)
+  local output = 'smart extension is %s;\ncitations extension is %s\n'
+  local status = function (ext)
+    return opts.extensions:includes(ext) and 'enabled' or 'disabled'
+  end
+  return output:format(status('smart'), status('citations'))
+end
+
+writer_extensions = {
+  smart = true,
+  citations = false,
+}