{- | A minimal Wadler/Leijen-style document combinator and width-aware renderer.
A 'Doc' describes a layout abstractly; 'render' chooses where soft breaks become
newlines to fit a target width. 'Group' lays a region flat when it fits.
-}
module DataFrame.Internal.Pretty (
    Doc,
    text,
    line,
    hardline,
    nest,
    group,
    (<+>),
    hcat,
    punctuate,
    parens,
    parensWhenBroken,
    defaultWidth,
    render,
) where

data Doc
    = Empty
    | Text String
    | Line
    | Cat Doc Doc
    | Nest Int Doc
    | Group Doc
    | Hard
    | Alt Doc Doc

instance Semigroup Doc where
    <> :: Doc -> Doc -> Doc
(<>) = Doc -> Doc -> Doc
Cat

instance Monoid Doc where
    mempty :: Doc
mempty = Doc
Empty

-- | A literal chunk of text. Must not contain newlines (use 'line'/'hardline').
text :: String -> Doc
text :: String -> Doc
text = String -> Doc
Text

{- | A soft break: a single space when its enclosing 'group' fits the width,
otherwise a newline + current indentation.
-}
line :: Doc
line :: Doc
line = Doc
Line

-- | A hard break that never flattens; any enclosing 'group' is forced to break.
hardline :: Doc
hardline :: Doc
hardline = Doc
Hard

-- | Add @k@ spaces to the indentation applied at line breaks inside @d@.
nest :: Int -> Doc -> Doc
nest :: Int -> Doc -> Doc
nest = Int -> Doc -> Doc
Nest

-- | Lay the document out flat if it fits the remaining width, broken otherwise.
group :: Doc -> Doc
group :: Doc -> Doc
group = Doc -> Doc
Group

-- | Concatenate two documents separated by a single space.
(<+>) :: Doc -> Doc -> Doc
Doc
x <+> :: Doc -> Doc -> Doc
<+> Doc
y = Doc
x Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
Text String
" " Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
y

infixr 6 <+>

hcat :: [Doc] -> Doc
hcat :: [Doc] -> Doc
hcat = [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat

-- | Append @sep@ after every element but the last.
punctuate :: Doc -> [Doc] -> [Doc]
punctuate :: Doc -> [Doc] -> [Doc]
punctuate Doc
_ [] = []
punctuate Doc
_ [Doc
d] = [Doc
d]
punctuate Doc
sep (Doc
d : [Doc]
ds) = (Doc
d Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
sep) Doc -> [Doc] -> [Doc]
forall a. a -> [a] -> [a]
: Doc -> [Doc] -> [Doc]
punctuate Doc
sep [Doc]
ds

parens :: Doc -> Doc
parens :: Doc -> Doc
parens Doc
d = String -> Doc
Text String
"(" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
d Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
Text String
")"

{- | Render @d@ bare when it fits flat on the current line, wrapped in parens when
it must break across lines. Keeps operator grouping unambiguous once a
sub-expression wraps, without parenthesis noise on one-line expressions.
-}
parensWhenBroken :: Doc -> Doc
parensWhenBroken :: Doc -> Doc
parensWhenBroken Doc
d = Doc -> Doc
Group (Doc -> Doc -> Doc
Alt Doc
d (Doc -> Doc
parens Doc
d))

defaultWidth :: Int
defaultWidth :: Int
defaultWidth = Int
80

data Mode = Flat | Break

-- | Render a document, breaking soft lines so output fits @width@ columns.
render :: Int -> Doc -> String
render :: Int -> Doc -> String
render Int
width Doc
doc = Int -> [(Int, Mode, Doc)] -> String
layout Int
0 [(Int
0, Mode
Break, Doc
doc)]
  where
    layout :: Int -> [(Int, Mode, Doc)] -> String
    layout :: Int -> [(Int, Mode, Doc)] -> String
layout Int
_ [] = String
""
    layout Int
col ((Int
i, Mode
m, Doc
d) : [(Int, Mode, Doc)]
rest) = case Doc
d of
        Doc
Empty -> Int -> [(Int, Mode, Doc)] -> String
layout Int
col [(Int, Mode, Doc)]
rest
        Text String
s -> String
s String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> [(Int, Mode, Doc)] -> String
layout (Int
col Int -> Int -> Int
forall a. Num a => a -> a -> a
+ String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
s) [(Int, Mode, Doc)]
rest
        Cat Doc
x Doc
y -> Int -> [(Int, Mode, Doc)] -> String
layout Int
col ((Int
i, Mode
m, Doc
x) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: (Int
i, Mode
m, Doc
y) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: [(Int, Mode, Doc)]
rest)
        Nest Int
j Doc
x -> Int -> [(Int, Mode, Doc)] -> String
layout Int
col ((Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j, Mode
m, Doc
x) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: [(Int, Mode, Doc)]
rest)
        Doc
Line -> case Mode
m of
            Mode
Flat -> Char
' ' Char -> String -> String
forall a. a -> [a] -> [a]
: Int -> [(Int, Mode, Doc)] -> String
layout (Int
col Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [(Int, Mode, Doc)]
rest
            Mode
Break -> Char
'\n' Char -> String -> String
forall a. a -> [a] -> [a]
: Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
i Char
' ' String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> [(Int, Mode, Doc)] -> String
layout Int
i [(Int, Mode, Doc)]
rest
        Doc
Hard -> Char
'\n' Char -> String -> String
forall a. a -> [a] -> [a]
: Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
i Char
' ' String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> [(Int, Mode, Doc)] -> String
layout Int
i [(Int, Mode, Doc)]
rest
        Group Doc
x ->
            if Int -> [(Int, Mode, Doc)] -> Bool
fits (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
col) ((Int
i, Mode
Flat, Doc
x) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: [(Int, Mode, Doc)]
rest)
                then Int -> [(Int, Mode, Doc)] -> String
layout Int
col ((Int
i, Mode
Flat, Doc
x) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: [(Int, Mode, Doc)]
rest)
                else Int -> [(Int, Mode, Doc)] -> String
layout Int
col ((Int
i, Mode
Break, Doc
x) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: [(Int, Mode, Doc)]
rest)
        Alt Doc
flat Doc
broken -> case Mode
m of
            Mode
Flat -> Int -> [(Int, Mode, Doc)] -> String
layout Int
col ((Int
i, Mode
Flat, Doc
flat) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: [(Int, Mode, Doc)]
rest)
            Mode
Break -> Int -> [(Int, Mode, Doc)] -> String
layout Int
col ((Int
i, Mode
Break, Doc
broken) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: [(Int, Mode, Doc)]
rest)

    fits :: Int -> [(Int, Mode, Doc)] -> Bool
    fits :: Int -> [(Int, Mode, Doc)] -> Bool
fits Int
w [(Int, Mode, Doc)]
_ | Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = Bool
False
    fits Int
_ [] = Bool
True
    fits Int
w ((Int
i, Mode
m, Doc
d) : [(Int, Mode, Doc)]
rest) = case Doc
d of
        Doc
Empty -> Int -> [(Int, Mode, Doc)] -> Bool
fits Int
w [(Int, Mode, Doc)]
rest
        Text String
s -> Int -> [(Int, Mode, Doc)] -> Bool
fits (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
s) [(Int, Mode, Doc)]
rest
        Cat Doc
x Doc
y -> Int -> [(Int, Mode, Doc)] -> Bool
fits Int
w ((Int
i, Mode
m, Doc
x) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: (Int
i, Mode
m, Doc
y) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: [(Int, Mode, Doc)]
rest)
        Nest Int
j Doc
x -> Int -> [(Int, Mode, Doc)] -> Bool
fits Int
w ((Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j, Mode
m, Doc
x) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: [(Int, Mode, Doc)]
rest)
        Doc
Line -> case Mode
m of
            Mode
Flat -> Int -> [(Int, Mode, Doc)] -> Bool
fits (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [(Int, Mode, Doc)]
rest
            Mode
Break -> Bool
True
        Doc
Hard -> case Mode
m of
            Mode
Flat -> Bool
False
            Mode
Break -> Bool
True
        Group Doc
x -> Int -> [(Int, Mode, Doc)] -> Bool
fits Int
w ((Int
i, Mode
Flat, Doc
x) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: [(Int, Mode, Doc)]
rest)
        Alt Doc
flat Doc
broken -> case Mode
m of
            Mode
Flat -> Int -> [(Int, Mode, Doc)] -> Bool
fits Int
w ((Int
i, Mode
Flat, Doc
flat) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: [(Int, Mode, Doc)]
rest)
            Mode
Break -> Int -> [(Int, Mode, Doc)] -> Bool
fits Int
w ((Int
i, Mode
Break, Doc
broken) (Int, Mode, Doc) -> [(Int, Mode, Doc)] -> [(Int, Mode, Doc)]
forall a. a -> [a] -> [a]
: [(Int, Mode, Doc)]
rest)