-- | This module provides some utility functions to help with aligning pretty
-- printed values by column.
module Codec.CBOR.Cuddle.Pretty.Columnar (
  CellAlign (..),
  Row (..),
  Cell (..),
  cellL,
  cellR,
  emptyCell,
  isEmptyCell,
  Columnar (..),
  prettyColumnar,
  columnarListing,
  columnarSepBy,
  singletonRow,
) where

import Codec.CBOR.Cuddle.Pretty.Utils (fillLeft, fillRight, renderedLen)
import Prettyprinter (Doc, Pretty (..), line', vcat, (<+>))

data CellAlign
  = LeftAlign
  | RightAlign

data Cell ann = Cell
  { forall ann. Cell ann -> Doc ann
cellDoc :: Doc ann
  , forall ann. Cell ann -> CellAlign
cellAlign :: CellAlign
  }

-- | Creates a cell by pretty printing the input value and then left-aligning
-- the resulting `Doc` within the table.
cellL :: Pretty a => a -> Cell ann
cellL :: forall a ann. Pretty a => a -> Cell ann
cellL = (Doc ann -> CellAlign -> Cell ann
forall ann. Doc ann -> CellAlign -> Cell ann
`Cell` CellAlign
LeftAlign) (Doc ann -> Cell ann) -> (a -> Doc ann) -> a -> Cell ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty

-- | Creates a cell by pretty printing the input value and then right-aligning
-- the resulting `Doc` within the table.
cellR :: Pretty a => a -> Cell ann
cellR :: forall a ann. Pretty a => a -> Cell ann
cellR = (Doc ann -> CellAlign -> Cell ann
forall ann. Doc ann -> CellAlign -> Cell ann
`Cell` CellAlign
RightAlign) (Doc ann -> Cell ann) -> (a -> Doc ann) -> a -> Cell ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty

-- | A cell that takes up a cell but has no content.
emptyCell :: Cell ann
emptyCell :: forall ann. Cell ann
emptyCell = Doc ann -> CellAlign -> Cell ann
forall ann. Doc ann -> CellAlign -> Cell ann
Cell Doc ann
forall a. Monoid a => a
mempty CellAlign
LeftAlign

-- | Checks whether the cell contains a `Doc` with a rendered width of zero.
isEmptyCell :: Cell ann -> Bool
isEmptyCell :: forall ann. Cell ann -> Bool
isEmptyCell (Cell Doc ann
d CellAlign
_) = Doc ann -> Int
forall ann. Doc ann -> Int
renderedLen Doc ann
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0

-- | A row within the table.
newtype Row ann = Row {forall ann. Row ann -> [Cell ann]
rowCells :: [Cell ann]}

-- | Adds a cell at the beginning of the row.
prependCell :: Cell ann -> Row ann -> Row ann
prependCell :: forall ann. Cell ann -> Row ann -> Row ann
prependCell Cell ann
c (Row [Cell ann]
cs) = [Cell ann] -> Row ann
forall ann. [Cell ann] -> Row ann
Row ([Cell ann] -> Row ann) -> [Cell ann] -> Row ann
forall a b. (a -> b) -> a -> b
$ Cell ann
c Cell ann -> [Cell ann] -> [Cell ann]
forall a. a -> [a] -> [a]
: [Cell ann]
cs

-- | A row with a single left-aligned document
singletonRow :: Doc ann -> Row ann
singletonRow :: forall ann. Doc ann -> Row ann
singletonRow Doc ann
x = [Cell ann] -> Row ann
forall ann. [Cell ann] -> Row ann
Row [Doc ann -> CellAlign -> Cell ann
forall ann. Doc ann -> CellAlign -> Cell ann
Cell Doc ann
x CellAlign
LeftAlign]

-- | `Columnar` is a two-dimensional table of `Doc`s. When rendered, the cells
-- within each row will be aligned with the cells of every other row in the
-- same column.
newtype Columnar ann = Columnar {forall ann. Columnar ann -> [Row ann]
colRows :: [Row ann]}

prettyRow :: [Int] -> [Cell ann] -> Doc ann
prettyRow :: forall ann. [Int] -> [Cell ann] -> Doc ann
prettyRow = [Int] -> [Cell ann] -> Doc ann
forall ann. [Int] -> [Cell ann] -> Doc ann
prettyRow'
  where
    prettyRow' :: [Int] -> [Cell ann] -> Doc ann
prettyRow' [] (Cell Doc ann
c CellAlign
_ : [Cell ann]
cs) = Doc ann
c Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Int] -> [Cell ann] -> Doc ann
prettyRow' [] [Cell ann]
cs
    prettyRow' [Int]
_ [] = Doc ann
forall a. Monoid a => a
mempty
    prettyRow' [Int]
_ [Cell Doc ann
c CellAlign
LeftAlign] = Doc ann
c -- Do not add white space to the last cell
    prettyRow' (Int
0 : [Int]
ws) (Cell ann
_ : [Cell ann]
cs) = [Int] -> [Cell ann] -> Doc ann
prettyRow' [Int]
ws [Cell ann]
cs -- Skip empty columns
    prettyRow' (Int
w : [Int]
ws) (Cell Doc ann
c CellAlign
alignment : [Cell ann]
cs) =
      let
        align' :: Int -> Doc ann -> Doc ann
align' = case CellAlign
alignment of
          CellAlign
LeftAlign -> Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
fillRight
          CellAlign
RightAlign -> Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
fillLeft
       in
        Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
align' Int
w Doc ann
c Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Int] -> [Cell ann] -> Doc ann
prettyRow' [Int]
ws [Cell ann]
cs

-- | Pretty print the `Columnar` as a table.
prettyColumnar :: forall ann. Columnar ann -> Doc ann
prettyColumnar :: forall ann. Columnar ann -> Doc ann
prettyColumnar (Columnar [Row ann]
rows) = [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
vcat ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ [Int] -> [Cell ann] -> Doc ann
forall ann. [Int] -> [Cell ann] -> Doc ann
prettyRow [Int]
columnWidths ([Cell ann] -> Doc ann)
-> (Row ann -> [Cell ann]) -> Row ann -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Row ann -> [Cell ann]
forall ann. Row ann -> [Cell ann]
rowCells (Row ann -> Doc ann) -> [Row ann] -> [Doc ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Row ann]
rows
  where
    columnWidths :: [Int]
columnWidths =
      (Row ann -> [Int] -> [Int]) -> [Int] -> [Row ann] -> [Int]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Ord a => a -> a -> a
max ([Int] -> [Int] -> [Int])
-> (Row ann -> [Int]) -> Row ann -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Cell ann -> Int) -> [Cell ann] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Doc ann -> Int
forall ann. Doc ann -> Int
renderedLen (Doc ann -> Int) -> (Cell ann -> Doc ann) -> Cell ann -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Cell ann -> Doc ann
forall ann. Cell ann -> Doc ann
cellDoc) ([Cell ann] -> [Int])
-> (Row ann -> [Cell ann]) -> Row ann -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Row ann -> [Cell ann]
forall ann. Row ann -> [Cell ann]
rowCells) (Int -> [Int]
forall a. a -> [a]
repeat Int
0) [Row ann]
rows

-- | Pretty prints the `Columnar` so that the rows are separated by by the
-- separator `Doc` provided as the third argument and then everything is
-- enclosed within the left and right brackets provided as the first and second
-- argument accordingly. The brackets will be aligned with the separators in the
-- first column, e.g.
-- ```
-- [ foo
-- , bar
-- ]
-- ```
columnarListing :: Doc ann -> Doc ann -> Doc ann -> Columnar ann -> Doc ann
columnarListing :: forall ann.
Doc ann -> Doc ann -> Doc ann -> Columnar ann -> Doc ann
columnarListing Doc ann
lEnc Doc ann
rEnc Doc ann
_ (Columnar []) = Doc ann
lEnc Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
rEnc
columnarListing Doc ann
lEnc Doc ann
rEnc Doc ann
s (Columnar (Row ann
row : [Row ann]
rows)) =
  Columnar ann -> Doc ann
forall ann. Columnar ann -> Doc ann
prettyColumnar
    ( [Row ann] -> Columnar ann
forall ann. [Row ann] -> Columnar ann
Columnar ([Row ann] -> Columnar ann) -> [Row ann] -> Columnar ann
forall a b. (a -> b) -> a -> b
$
        Cell ann -> Row ann -> Row ann
forall ann. Cell ann -> Row ann -> Row ann
prependCell (Doc ann -> CellAlign -> Cell ann
forall ann. Doc ann -> CellAlign -> Cell ann
Cell Doc ann
lEnc CellAlign
LeftAlign) Row ann
row
          Row ann -> [Row ann] -> [Row ann]
forall a. a -> [a] -> [a]
: (Cell ann -> Row ann -> Row ann
forall ann. Cell ann -> Row ann -> Row ann
prependCell (Doc ann -> CellAlign -> Cell ann
forall ann. Doc ann -> CellAlign -> Cell ann
Cell Doc ann
s CellAlign
LeftAlign) (Row ann -> Row ann) -> [Row ann] -> [Row ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Row ann]
rows)
    )
    Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
line'
    Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
rEnc

-- | Pretty prints the `Columnar` so that every line after the first has a
-- separator prepended to it. This can be useful when you want to align the rows,
-- but the separator would cause all the other rows after the first to be shifted
-- right by one. The way you use this is you reduce the indentation on the
-- following lines by the width of the separator.
-- ```
-- foo = x
--     , y
--     , z
-- ```
columnarSepBy :: Doc ann -> Columnar ann -> Doc ann
columnarSepBy :: forall ann. Doc ann -> Columnar ann -> Doc ann
columnarSepBy Doc ann
_ (Columnar []) = Doc ann
forall a. Monoid a => a
mempty
columnarSepBy Doc ann
s (Columnar rows :: [Row ann]
rows@(Row [Cell ann]
x : [Row ann]
xs)) =
  [Int] -> [Cell ann] -> Doc ann
forall ann. [Int] -> [Cell ann] -> Doc ann
prettyRow [Int]
columnWidths [Cell ann]
x Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
line' Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Columnar ann -> Doc ann
forall ann. Columnar ann -> Doc ann
prettyColumnar ([Row ann] -> Columnar ann
forall ann. [Row ann] -> Columnar ann
Columnar ([Row ann] -> Columnar ann) -> [Row ann] -> Columnar ann
forall a b. (a -> b) -> a -> b
$ Row ann -> Row ann
prependRow (Row ann -> Row ann) -> [Row ann] -> [Row ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Row ann]
xs)
  where
    prependRow :: Row ann -> Row ann
prependRow (Row (Cell Doc ann
c CellAlign
al : [Cell ann]
cs)) = [Cell ann] -> Row ann
forall ann. [Cell ann] -> Row ann
Row ([Cell ann] -> Row ann) -> [Cell ann] -> Row ann
forall a b. (a -> b) -> a -> b
$ Doc ann -> CellAlign -> Cell ann
forall ann. Doc ann -> CellAlign -> Cell ann
Cell (Doc ann
s Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
c) CellAlign
al Cell ann -> [Cell ann] -> [Cell ann]
forall a. a -> [a] -> [a]
: [Cell ann]
cs
    prependRow (Row []) = [Cell ann] -> Row ann
forall ann. [Cell ann] -> Row ann
Row []
    columnWidths :: [Int]
columnWidths =
      (Row ann -> [Int] -> [Int]) -> [Int] -> [Row ann] -> [Int]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Ord a => a -> a -> a
max ([Int] -> [Int] -> [Int])
-> (Row ann -> [Int]) -> Row ann -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Cell ann -> Int) -> [Cell ann] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Doc ann -> Int
forall ann. Doc ann -> Int
renderedLen (Doc ann -> Int) -> (Cell ann -> Doc ann) -> Cell ann -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Cell ann -> Doc ann
forall ann. Cell ann -> Doc ann
cellDoc) ([Cell ann] -> [Int])
-> (Row ann -> [Cell ann]) -> Row ann -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Row ann -> [Cell ann]
forall ann. Row ann -> [Cell ann]
rowCells) (Int -> [Int]
forall a. a -> [a]
repeat Int
0) [Row ann]
rows