ANSI writer: Apply row spans in tables (#11294)

The ANSI writer is now able to keep track of row spans and apply them in rows.

[API change] T.P.Writers.Shared:
Add functions `insertCurrentSpansAtColumn`, `takePreviousSpansAtColumn`
and `decrementTrailingRowSpans` for applying and keeping track of
RowSpans over multiple rows.

Fixes: #10149

3 files changed, 86 insertions, 32 deletions

diff --git a/src/Text/Pandoc/Writers/ANSI.hs b/src/Text/Pandoc/Writers/ANSI.hs
index cc731aea1..e20713f61 100644
--- a/src/Text/Pandoc/Writers/ANSI.hs
+++ b/src/Text/Pandoc/Writers/ANSI.hs
@@ -15,6 +15,7 @@ module Text.Pandoc.Writers.ANSI ( writeANSI ) where
 import Control.Monad.State.Strict ( StateT, gets, modify, evalStateT )
 import Control.Monad (foldM)
 import Data.List (intersperse)
+import qualified Data.Map as M
 import Data.Maybe (fromMaybe)
 import Data.Text (Text)
 import Text.DocLayout ((<+>), ($$), ($+$))
@@ -108,9 +109,7 @@ titleBlock :: Maybe Int -> Context Text -> D.Doc Text
 titleBlock width meta =
   if null most
      then D.empty
-     else (case width of
-             Just w -> D.cblock w
-             Nothing -> id) $ most $+$ hr
+     else (maybe id D.cblock width) $ most $+$ hr
   where
     title = D.bold (fromMaybe D.empty $ getField "title" meta)
     subtitle = fromMaybe D.empty $ getField "subtitle" meta
@@ -193,7 +192,6 @@ blockToANSI opts (BlockQuote blocks) = do
   contents <- withFewerColumns 2 $ blockListToANSI opts blocks
   return ( D.prefixed "│ " contents $$ D.blankline)
 
--- TODO: Row spans don't work
 blockToANSI opts (Table _ (Caption _ caption) colSpecs (TableHead _ thead) tbody (TableFoot _ tfoot)) = do
   let captionInlines = blocksToInlines caption
   captionMarkup <-
@@ -215,22 +213,40 @@ blockToANSI opts (Table _ (Caption _ caption) colSpecs (TableHead _ thead) tbody
       maxWidth k = claimWidth k
   let widths = map maxWidth inWidths
   let decor = [D.hsep $ map rule widths]
-  head' <- mapM (goRow widths . unRow) thead
-  body' <- mapM (goRow widths . unRow) (unBodies tbody)
-  foot' <- mapM (goRow widths . unRow) tfoot
+  head' <- (makeRows widths . map unRow) thead
+  body' <- (makeRows widths . map unRow) (tableBodiesToRows tbody)
+  foot' <- (makeRows widths . map unRow) tfoot
   modify $ \s -> s{stInTable = wasTable}
   return $ D.vcat (head' <> decor <> body' <> decor <> foot') $+$ captionMarkup
   where
     unRow (Row _ cs) = cs
-    unBody (TableBody _ _ hd bd) = hd <> bd
-    unBodies = concatMap unBody
-    goRow ws cs = do
-      (d, _) <- foldM goCell ([], ws) cs
-      return $ D.hcat $ intersperse (D.vfill " ") $ reverse d
-    goCell (r, ws) (Cell _ aln _ (ColSpan cspan) inner) = do
+    makeRows ws rows = do
+      (docs, _) <- foldM (goRow ws) ([], M.empty) rows
+      return $ reverse docs
+    goRow _ (r, spans) [] =
+      -- Empty rows are not displayed but previous row spans still apply for them.
+      let spans' = M.map decrementPreviousRowSpans spans
+      in  return (r, spans')
+    goRow ws (r, spans) cs = do
+      (d, (nextPos, spans'), _) <- foldM goCell ([], (0, spans), ws) cs
+      let spans'' = decrementTrailingRowSpans nextPos spans' -- Handle previous row spans next to the end of the current row
+      return (D.hcat (intersperse (D.vfill " ") $ reverse d):r, spans'')
+    goCell (r, (colPos, spans), ws) cell@(Cell _ aln (RowSpan rspan) (ColSpan cspan) inner)
+      | Just (ColSpan previousColSpan, spans') <- takePreviousSpansAtColumn colPos spans = do
+          (r', nextPos, ws') <- makeCell r colPos ws AlignDefault previousColSpan []
+          goCell (r', (nextPos, spans'), ws') cell
+      | otherwise = do
+          (r', nextPos, ws') <- makeCell r colPos ws aln cspan inner
+          let spans' = insertCurrentSpansAtColumn colPos spans (RowSpan rspan) (ColSpan cspan)
+          return (r', (nextPos, spans'), ws')
+    decrementPreviousRowSpans spans@(RowSpan rspan, cspan) =
+      if rspan >= 1
+        then (RowSpan rspan - 1, cspan)
+        else spans
+    makeCell r colPos ws aln cspan inner = do
       let (ws', render) = next ws aln cspan
       innerDoc <- blockListToANSI opts inner
-      return ((render innerDoc):r, ws')
+      return ((render innerDoc):r, colPos + cspan, ws')
     tcell AlignLeft    = D.lblock
     tcell AlignRight   = D.rblock
     tcell AlignCenter  = D.cblock
diff --git a/src/Text/Pandoc/Writers/AsciiDoc.hs b/src/Text/Pandoc/Writers/AsciiDoc.hs
index d0150bb37..aca0a4241 100644
--- a/src/Text/Pandoc/Writers/AsciiDoc.hs
+++ b/src/Text/Pandoc/Writers/AsciiDoc.hs
@@ -818,25 +818,19 @@ adjustFooters rows = adjustFooters' [] (0, length rows) M.empty rows
 
         -- Apply row spans from a previous row that are next to the end of the
         -- current row's cells to keep track of the correct column position.
-        previousRowSpans'' = M.mapWithKey (applyTrailingPreviousRowSpans nextColumnPosition) previousRowSpans'
+        previousRowSpans'' = decrementTrailingRowSpans nextColumnPosition previousRowSpans'
     in  (previousRowSpans'', Row attr cells'', catMaybes columnIndices)
 
-  applyTrailingPreviousRowSpans nextColumnPosition columnPosition previousRowSpan@(RowSpan rowSpan, ColSpan colSpan) =
-    if columnPosition >= nextColumnPosition && rowSpan >= 1
-      then (RowSpan rowSpan - 1, ColSpan colSpan)
-      else previousRowSpan
-
 -- | Adjust footer cell for the fact that AsciiDoc only supports a single footer row.
 --
 -- Collects cells whose RowSpan would reach to the last footer row and applies
 -- them as empty cells to that last footer row.
 adjustFooterCell :: (Int, Int) -> (Int, M.Map Int (RowSpan, ColSpan)) -> Cell -> ((Int, M.Map Int (RowSpan, ColSpan)), (Cell, Maybe (Int, Cell)))
 adjustFooterCell rowInfo@(rowIndex, footerLength) (columnPosition, previousSpans) cell@(Cell _ _ (RowSpan rowSpan) (ColSpan colSpan) _)
-  | Just previous@(RowSpan previousRowSpan, ColSpan previousColSpan) <- M.lookup columnPosition previousSpans
-  , previousRowSpan >= 1 =
+  | Just (ColSpan previousColSpan, previousSpans') <- takePreviousSpansAtColumn columnPosition previousSpans =
       -- Apply row span from a previous row that occupies this column to keep
       -- track of the correct column position.
-      adjustFooterCell rowInfo (columnPosition + previousColSpan, updatePreviousRowSpan previous) cell
+      adjustFooterCell rowInfo (columnPosition + previousColSpan, previousSpans') cell
   | rowSpan > 1 && rowIndex + rowSpan >= footerLength =
       -- Adjust row span that would reach all the way to the last footer row and
       -- keep track of that to apply it to the last footer row.
@@ -844,15 +838,7 @@ adjustFooterCell rowInfo@(rowIndex, footerLength) (columnPosition, previousSpans
   | otherwise = ((nextColumnPosition, previousRowSpans'), (cell, Nothing))
  where
   -- Keep track of this cell's RowSpan for the rows following it.
-  previousRowSpans' = if rowSpan > 1
-                      then M.insert columnPosition (RowSpan rowSpan - 1, ColSpan colSpan) previousSpans -- Minus its own row.
-                      else previousSpans
-
-  updatePreviousRowSpan (RowSpan previousRowSpan, previousColSpan) =
-    if previousRowSpan > 1
-      then M.insert columnPosition (RowSpan previousRowSpan - 1, previousColSpan) previousSpans
-      else M.delete columnPosition previousSpans
-
+  previousRowSpans' = insertCurrentSpansAtColumn columnPosition previousSpans (RowSpan rowSpan) (ColSpan colSpan)
   nextColumnPosition = columnPosition + colSpan
   emptyCellWithColSpan = Cell nullAttr AlignDefault (RowSpan 1) (ColSpan colSpan) []
 
diff --git a/src/Text/Pandoc/Writers/Shared.hs b/src/Text/Pandoc/Writers/Shared.hs
index 2fdb8de5d..c95ba06d7 100644
--- a/src/Text/Pandoc/Writers/Shared.hs
+++ b/src/Text/Pandoc/Writers/Shared.hs
@@ -52,6 +52,9 @@ module Text.Pandoc.Writers.Shared (
                      , delimited
                      , allRowsEmpty
                      , tableBodiesToRows
+                     , insertCurrentSpansAtColumn
+                     , takePreviousSpansAtColumn
+                     , decrementTrailingRowSpans
                      )
 where
 import Safe (lastMay, maximumMay)
@@ -873,3 +876,52 @@ tableBodiesToRows :: [TableBody] -> [Row]
 tableBodiesToRows = concatMap tableBodyToRows
   where
     tableBodyToRows (TableBody _ _ headerRows bodyRows) = headerRows ++ bodyRows
+
+-- | Insert the current span information of a table cell to keep track of it in
+-- subsequent rows.
+--
+-- If 'RowSpan' @> 1@, the current span information will be inserted. Otherwise
+-- the previous span information will be left unchanged.
+--
+-- Use 'takePreviousSpansAtColumn' to take previous span information at
+-- subsequent rows. Use 'decrementTrailingRowSpans' to handle previous trailing
+-- spans at the end of a row.
+--
+-- For writers that need to manually apply the 'RowSpan' of cells over multiple
+-- rows or otherwise have to keep track of it.
+insertCurrentSpansAtColumn :: Int -> M.Map Int (RowSpan, ColSpan) -> RowSpan -> ColSpan -> M.Map Int (RowSpan, ColSpan)
+insertCurrentSpansAtColumn columnPosition previousSpans (RowSpan rowSpan) colSpan =
+  if (rowSpan > 1)
+    then M.insert columnPosition (RowSpan rowSpan - 1, colSpan) previousSpans -- Minus its own row.
+    else previousSpans
+
+-- | Take previous span information at a column position that was added with
+-- 'insertCurrentSpansAtColumn' if available.
+--
+-- If the previous 'RowSpan' @>= 1@, this will return 'Just' the previous
+-- 'ColSpan' and an adjusted span information where that 'RowSpan' is either
+-- decremented or deleted if it would fall to 0. Otherwise this will return
+-- 'Nothing'.
+takePreviousSpansAtColumn :: Int -> M.Map Int (RowSpan, ColSpan) -> Maybe (ColSpan, M.Map Int (RowSpan, ColSpan))
+takePreviousSpansAtColumn columnPosition previousSpans
+  | Just previous@(RowSpan previousRowSpan, previousColSpan) <- M.lookup columnPosition previousSpans
+  , previousRowSpan >= 1 = Just (previousColSpan, decrementPreviousRowSpans previous)
+  | otherwise = Nothing
+ where
+  decrementPreviousRowSpans (RowSpan previousRowSpan, previousColSpan) =
+    if previousRowSpan > 1
+      then M.insert columnPosition (RowSpan previousRowSpan - 1, previousColSpan) previousSpans
+      else M.delete columnPosition previousSpans
+
+-- | Decrement all previously tracked trailing 'RowSpan' elements at or after a
+-- column position.
+--
+-- For handling previous row spans that are next to the end of a row's cells
+-- that were previously added with 'insertCurrentSpansAtColumn'.
+decrementTrailingRowSpans :: Int -> M.Map Int (RowSpan, ColSpan) -> M.Map Int (RowSpan, ColSpan)
+decrementTrailingRowSpans columnPosition = M.mapWithKey decrementTrailing
+  where
+    decrementTrailing previousColumnPosition previousSpan@(RowSpan rowSpan, colSpan) =
+      if previousColumnPosition >= columnPosition && rowSpan >= 1
+        then (RowSpan rowSpan - 1, colSpan)
+        else previousSpan