{-# LANGUAGE LambdaCase #-}

-- | Pretty-printer API over 'Doc' document trees.
--
-- Smart constructors, combinators, layout algorithms, and rendering.
module Data.Fmt.Tree (
    -- * Smart constructors
    fail_,
    emptyDoc,
    leaf,
    hardline,
    line,
    line',
    flatAlt,
    nest,
    union,
    annotate,
    column,
    nesting,

    -- * Combinators
    flatten,
    group,
    align,

    -- * Line breaks
    softline,
    softline',

    -- * Separators
    (<+>),
    concatWith,
    hsep,
    vsep,
    fillSep,
    sep,

    -- * Concatenation
    hcat,
    vcat,
    fillCat,
    cat,

    -- * Indentation
    hang,
    indent,

    -- * Enclosure
    surround,
    encloseSep,
    list,
    tupled,
    punctuate,

    -- * Filling
    width,
    fill,
    fillBreak,

    -- * Annotations
    reAnnotate,
    unAnnotate,
    alterAnnotations,

    -- * Optimized group
    FlattenResult (..),
    changesUponFlattening,
    group',

    -- * Fusion
    fuse,

    -- * Trailing whitespace
    removeTrailingWhitespace,

    -- * Tokens
    Token (..),

    -- * Layout
    PageWidth (..),
    LayoutOptions (..),
    defaultLayoutOptions,
    layoutPretty,
    layoutSmart,
    layoutCompact,

    -- * Streaming layout
    layoutStream,
    renderStream,
    prettyStream,

    -- * Rendering
    render,
    pretty,
) where

import Data.Fmt.Cons (Cons (..))
import Data.Fmt.Fixed (Mu, Nu (..), fold, foldWithContext, hoistMu, unwrap, wrap)
import Data.Fmt.Functor (Doc (..), Tree)
import Data.String (IsString (..))

---------------------------------------------------------------------
-- Smart constructors
---------------------------------------------------------------------

-- | Layout failure.
{-# INLINE fail_ #-}
fail_ :: Tree m ann
fail_ :: forall m ann. Tree m ann
fail_ = Doc m ann (Mu (Doc m ann)) -> Mu (Doc m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap Doc m ann (Mu (Doc m ann))
forall m ann r. Doc m ann r
Fail

-- | Empty document.
{-# INLINE emptyDoc #-}
emptyDoc :: Tree m ann
emptyDoc :: forall m ann. Tree m ann
emptyDoc = Doc m ann (Mu (Doc m ann)) -> Mu (Doc m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap Doc m ann (Mu (Doc m ann))
forall m ann r. Doc m ann r
Empty

-- | Literal content with display width.
{-# INLINE leaf #-}
leaf :: Int -> m -> Tree m ann
leaf :: forall m ann. Int -> m -> Tree m ann
leaf Int
n m
m = Doc m ann (Mu (Doc m ann)) -> Mu (Doc m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Int -> m -> Doc m ann (Mu (Doc m ann))
forall m ann r. Int -> m -> Doc m ann r
Leaf Int
n m
m)

-- | Hard line break. Cannot be flattened.
{-# INLINE hardline #-}
hardline :: Tree m ann
hardline :: forall m ann. Tree m ann
hardline = Doc m ann (Mu (Doc m ann)) -> Mu (Doc m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap Doc m ann (Mu (Doc m ann))
forall m ann r. Doc m ann r
Line

-- | Line break, or space when flattened.
{-# INLINE line #-}
line :: IsString m => Tree m ann
line :: forall m ann. IsString m => Tree m ann
line = Tree m ann -> Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann -> Tree m ann
flatAlt Tree m ann
forall m ann. Tree m ann
hardline (Int -> m -> Tree m ann
forall m ann. Int -> m -> Tree m ann
leaf Int
1 (String -> m
forall a. IsString a => String -> a
fromString String
" "))

-- | Line break, or empty when flattened.
{-# INLINE line' #-}
line' :: Tree m ann
line' :: forall m ann. Tree m ann
line' = Tree m ann -> Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann -> Tree m ann
flatAlt Tree m ann
forall m ann. Tree m ann
hardline Tree m ann
forall m ann. Tree m ann
emptyDoc

-- | @flatAlt default flat@: use @default@ normally,
-- @flat@ when flattened by 'group'.
{-# INLINE flatAlt #-}
flatAlt :: Tree m ann -> Tree m ann -> Tree m ann
flatAlt :: forall m ann. Tree m ann -> Tree m ann -> Tree m ann
flatAlt Tree m ann
x Tree m ann
y = Doc m ann (Tree m ann) -> Tree m ann
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Tree m ann -> Tree m ann -> Doc m ann (Tree m ann)
forall m ann r. r -> r -> Doc m ann r
FlatAlt Tree m ann
x Tree m ann
y)

-- | Increase nesting by @i@.
{-# INLINE nest #-}
nest :: Int -> Tree m ann -> Tree m ann
nest :: forall m ann. Int -> Tree m ann -> Tree m ann
nest Int
i = Doc m ann (Mu (Doc m ann)) -> Mu (Doc m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Doc m ann (Mu (Doc m ann)) -> Mu (Doc m ann))
-> (Mu (Doc m ann) -> Doc m ann (Mu (Doc m ann)))
-> Mu (Doc m ann)
-> Mu (Doc m ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Mu (Doc m ann) -> Doc m ann (Mu (Doc m ann))
forall m ann r. Int -> r -> Doc m ann r
Nest Int
i

-- | Layout alternatives. Invariant: first argument is the
-- flattened form of the second.
{-# INLINE union #-}
union :: Tree m ann -> Tree m ann -> Tree m ann
union :: forall m ann. Tree m ann -> Tree m ann -> Tree m ann
union Tree m ann
x Tree m ann
y = Doc m ann (Tree m ann) -> Tree m ann
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Tree m ann -> Tree m ann -> Doc m ann (Tree m ann)
forall m ann r. r -> r -> Doc m ann r
Union Tree m ann
x Tree m ann
y)

-- | Attach an annotation.
{-# INLINE annotate #-}
annotate :: ann -> Tree m ann -> Tree m ann
annotate :: forall ann m. ann -> Tree m ann -> Tree m ann
annotate ann
a = Doc m ann (Mu (Doc m ann)) -> Mu (Doc m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Doc m ann (Mu (Doc m ann)) -> Mu (Doc m ann))
-> (Mu (Doc m ann) -> Doc m ann (Mu (Doc m ann)))
-> Mu (Doc m ann)
-> Mu (Doc m ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ann -> Mu (Doc m ann) -> Doc m ann (Mu (Doc m ann))
forall m ann r. ann -> r -> Doc m ann r
Ann ann
a

-- | React to the current column position.
{-# INLINE column #-}
column :: (Int -> Tree m ann) -> Tree m ann
column :: forall m ann. (Int -> Tree m ann) -> Tree m ann
column = Doc m ann (Mu (Doc m ann)) -> Mu (Doc m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Doc m ann (Mu (Doc m ann)) -> Mu (Doc m ann))
-> ((Int -> Mu (Doc m ann)) -> Doc m ann (Mu (Doc m ann)))
-> (Int -> Mu (Doc m ann))
-> Mu (Doc m ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Mu (Doc m ann)) -> Doc m ann (Mu (Doc m ann))
forall m ann r. (Int -> r) -> Doc m ann r
Column

-- | React to the current nesting level.
{-# INLINE nesting #-}
nesting :: (Int -> Tree m ann) -> Tree m ann
nesting :: forall m ann. (Int -> Tree m ann) -> Tree m ann
nesting = Doc m ann (Mu (Doc m ann)) -> Mu (Doc m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Doc m ann (Mu (Doc m ann)) -> Mu (Doc m ann))
-> ((Int -> Mu (Doc m ann)) -> Doc m ann (Mu (Doc m ann)))
-> (Int -> Mu (Doc m ann))
-> Mu (Doc m ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Mu (Doc m ann)) -> Doc m ann (Mu (Doc m ann))
forall m ann r. (Int -> r) -> Doc m ann r
Nesting

---------------------------------------------------------------------
-- Combinators
---------------------------------------------------------------------

-- | Replace 'FlatAlt' with its flat branch, 'Line' with 'Fail'.
--
-- Implemented as a fold — children are flattened first,
-- so 'Column'/'Nesting' functions automatically produce
-- flattened subtrees.
flatten :: Tree m ann -> Tree m ann
flatten :: forall m ann. Tree m ann -> Tree m ann
flatten = Algebra (Doc m ann) (Tree m ann) -> Tree m ann -> Tree m ann
forall (f :: * -> *) a. Algebra f a -> Mu f -> a
fold (Algebra (Doc m ann) (Tree m ann) -> Tree m ann -> Tree m ann)
-> Algebra (Doc m ann) (Tree m ann) -> Tree m ann -> Tree m ann
forall a b. (a -> b) -> a -> b
$ \case
    FlatAlt Tree m ann
_ Tree m ann
y -> Tree m ann
y     -- use the flat alternative
    Doc m ann (Tree m ann)
Line -> Tree m ann
forall m ann. Tree m ann
fail_        -- hardline can't be flattened
    Union Tree m ann
a Tree m ann
_ -> Tree m ann
a       -- already the flatter branch
    Doc m ann (Tree m ann)
other -> Algebra (Doc m ann) (Tree m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap Doc m ann (Tree m ann)
other  -- re-wrap with flattened children

-- | Try to lay out the document on a single line.
-- Falls back to the original if flattening fails.
--
-- @group x = union (flatten x) x@
group :: Tree m ann -> Tree m ann
group :: forall m ann. Tree m ann -> Tree m ann
group Tree m ann
x = Tree m ann -> Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann -> Tree m ann
union (Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann
flatten Tree m ann
x) Tree m ann
x

-- | Lay out relative to the current column rather than
-- the current nesting level.
--
-- @align d = column (\\k -> nesting (\\i -> nest (k - i) d))@
align :: Tree m ann -> Tree m ann
align :: forall m ann. Tree m ann -> Tree m ann
align Tree m ann
d = (Int -> Tree m ann) -> Tree m ann
forall m ann. (Int -> Tree m ann) -> Tree m ann
column ((Int -> Tree m ann) -> Tree m ann)
-> (Int -> Tree m ann) -> Tree m ann
forall a b. (a -> b) -> a -> b
$ \Int
k -> (Int -> Tree m ann) -> Tree m ann
forall m ann. (Int -> Tree m ann) -> Tree m ann
nesting ((Int -> Tree m ann) -> Tree m ann)
-> (Int -> Tree m ann) -> Tree m ann
forall a b. (a -> b) -> a -> b
$ \Int
i -> Int -> Tree m ann -> Tree m ann
forall m ann. Int -> Tree m ann -> Tree m ann
nest (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) Tree m ann
d

---------------------------------------------------------------------
-- Line breaks
---------------------------------------------------------------------

-- | A line break that behaves like a space when flattened by 'group'.
--
-- @softline = group line@
softline :: IsString m => Tree m ann
softline :: forall m ann. IsString m => Tree m ann
softline = Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann
group Tree m ann
forall m ann. IsString m => Tree m ann
line

-- | A line break that vanishes when flattened by 'group'.
--
-- @softline' = group line'@
softline' :: Tree m ann
softline' :: forall m ann. Tree m ann
softline' = Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann
group Tree m ann
forall m ann. Tree m ann
line'

---------------------------------------------------------------------
-- Separators
---------------------------------------------------------------------

-- | Concatenate with a space in between.
infixr 6 <+>
(<+>) :: IsString m => Tree m ann -> Tree m ann -> Tree m ann
Tree m ann
x <+> :: forall m ann. IsString m => Tree m ann -> Tree m ann -> Tree m ann
<+> Tree m ann
y = Tree m ann
x Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Int -> m -> Tree m ann
forall m ann. Int -> m -> Tree m ann
leaf Int
1 (String -> m
forall a. IsString a => String -> a
fromString String
" ") Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
y

-- | Concatenate documents using a binary operator.
concatWith :: (Tree m ann -> Tree m ann -> Tree m ann) -> [Tree m ann] -> Tree m ann
concatWith :: forall m ann.
(Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
concatWith Tree m ann -> Tree m ann -> Tree m ann
_ [] = Tree m ann
forall m ann. Tree m ann
emptyDoc
concatWith Tree m ann -> Tree m ann -> Tree m ann
f (Tree m ann
x : [Tree m ann]
xs) = (Tree m ann -> Tree m ann -> Tree m ann)
-> Tree m ann -> [Tree m ann] -> Tree m ann
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Tree m ann -> Tree m ann -> Tree m ann
f Tree m ann
x [Tree m ann]
xs

-- | Concatenate with spaces.
hsep :: IsString m => [Tree m ann] -> Tree m ann
hsep :: forall m ann. IsString m => [Tree m ann] -> Tree m ann
hsep = (Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
forall m ann.
(Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
concatWith Tree m ann -> Tree m ann -> Tree m ann
forall m ann. IsString m => Tree m ann -> Tree m ann -> Tree m ann
(<+>)

-- | Concatenate with 'line' separators.
vsep :: IsString m => [Tree m ann] -> Tree m ann
vsep :: forall m ann. IsString m => [Tree m ann] -> Tree m ann
vsep = (Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
forall m ann.
(Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
concatWith (\Tree m ann
x Tree m ann
y -> Tree m ann
x Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
forall m ann. IsString m => Tree m ann
line Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
y)

-- | Concatenate with 'softline' separators.
fillSep :: IsString m => [Tree m ann] -> Tree m ann
fillSep :: forall m ann. IsString m => [Tree m ann] -> Tree m ann
fillSep = (Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
forall m ann.
(Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
concatWith (\Tree m ann
x Tree m ann
y -> Tree m ann
x Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
forall m ann. IsString m => Tree m ann
softline Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
y)

-- | 'vsep' that tries to fit on one line ('group').
sep :: IsString m => [Tree m ann] -> Tree m ann
sep :: forall m ann. IsString m => [Tree m ann] -> Tree m ann
sep = Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann
group (Tree m ann -> Tree m ann)
-> ([Tree m ann] -> Tree m ann) -> [Tree m ann] -> Tree m ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Tree m ann] -> Tree m ann
forall m ann. IsString m => [Tree m ann] -> Tree m ann
vsep

---------------------------------------------------------------------
-- Concatenation
---------------------------------------------------------------------

-- | Concatenate without separators.
hcat :: [Tree m ann] -> Tree m ann
hcat :: forall m ann. [Tree m ann] -> Tree m ann
hcat = (Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
forall m ann.
(Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
concatWith Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
(<>)

-- | Concatenate with 'line'' separators (line or empty).
vcat :: [Tree m ann] -> Tree m ann
vcat :: forall m ann. [Tree m ann] -> Tree m ann
vcat = (Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
forall m ann.
(Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
concatWith (\Tree m ann
x Tree m ann
y -> Tree m ann
x Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
forall m ann. Tree m ann
line' Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
y)

-- | Concatenate with 'softline'' separators.
fillCat :: [Tree m ann] -> Tree m ann
fillCat :: forall m ann. [Tree m ann] -> Tree m ann
fillCat = (Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
forall m ann.
(Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> Tree m ann
concatWith (\Tree m ann
x Tree m ann
y -> Tree m ann
x Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
forall m ann. Tree m ann
softline' Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
y)

-- | 'vcat' that tries to fit on one line ('group').
cat :: [Tree m ann] -> Tree m ann
cat :: forall m ann. [Tree m ann] -> Tree m ann
cat = Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann
group (Tree m ann -> Tree m ann)
-> ([Tree m ann] -> Tree m ann) -> [Tree m ann] -> Tree m ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Tree m ann] -> Tree m ann
forall m ann. [Tree m ann] -> Tree m ann
vcat

---------------------------------------------------------------------
-- Indentation
---------------------------------------------------------------------

-- | @hang i doc = align (nest i doc)@
hang :: Int -> Tree m ann -> Tree m ann
hang :: forall m ann. Int -> Tree m ann -> Tree m ann
hang Int
i Tree m ann
d = Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann
align (Int -> Tree m ann -> Tree m ann
forall m ann. Int -> Tree m ann -> Tree m ann
nest Int
i Tree m ann
d)

-- | @indent i doc@ inserts @i@ spaces then aligns.
indent :: IsString m => Int -> Tree m ann -> Tree m ann
indent :: forall m ann. IsString m => Int -> Tree m ann -> Tree m ann
indent Int
i Tree m ann
d = Int -> Tree m ann -> Tree m ann
forall m ann. Int -> Tree m ann -> Tree m ann
hang Int
i (Int -> m -> Tree m ann
forall m ann. Int -> m -> Tree m ann
leaf Int
i (String -> m
forall a. IsString a => String -> a
fromString (Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
i Char
' ')) Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
d)

---------------------------------------------------------------------
-- Enclosure
---------------------------------------------------------------------

-- | @surround mid left right = left \<\> mid \<\> right@
surround :: Tree m ann -> Tree m ann -> Tree m ann -> Tree m ann
surround :: forall m ann. Tree m ann -> Tree m ann -> Tree m ann -> Tree m ann
surround Tree m ann
x Tree m ann
l Tree m ann
r = Tree m ann
l Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
x Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
r

-- | Enclose a list with separators.
--
-- @encloseSep lbrace rbrace comma [a, b, c] = lbrace \<\> a \<\> comma \<\> b \<\> comma \<\> c \<\> rbrace@
--
-- When the content fits, renders on one line. Otherwise, each
-- element gets its own line, aligned.
encloseSep :: IsString m => Tree m ann -> Tree m ann -> Tree m ann -> [Tree m ann] -> Tree m ann
encloseSep :: forall m ann.
IsString m =>
Tree m ann
-> Tree m ann -> Tree m ann -> [Tree m ann] -> Tree m ann
encloseSep Tree m ann
l Tree m ann
r Tree m ann
_ [] = Tree m ann
l Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
r
encloseSep Tree m ann
l Tree m ann
r Tree m ann
s [Tree m ann]
ds = Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann
group (Tree m ann -> Tree m ann) -> Tree m ann -> Tree m ann
forall a b. (a -> b) -> a -> b
$
    Tree m ann
l Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> [Tree m ann] -> Tree m ann
forall m ann. [Tree m ann] -> Tree m ann
hcat ((Tree m ann -> Tree m ann -> Tree m ann)
-> [Tree m ann] -> [Tree m ann] -> [Tree m ann]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
(<>) (Tree m ann
forall m ann. Tree m ann
emptyDoc Tree m ann -> [Tree m ann] -> [Tree m ann]
forall a. a -> [a] -> [a]
: Tree m ann -> [Tree m ann]
forall a. a -> [a]
repeat (Tree m ann
s Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
forall m ann. Tree m ann
line')) [Tree m ann]
ds) Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
r

-- | @list = encloseSep \"[\" \"]\" \", \"@
list :: IsString m => [Tree m ann] -> Tree m ann
list :: forall m ann. IsString m => [Tree m ann] -> Tree m ann
list = Tree m ann
-> Tree m ann -> Tree m ann -> [Tree m ann] -> Tree m ann
forall m ann.
IsString m =>
Tree m ann
-> Tree m ann -> Tree m ann -> [Tree m ann] -> Tree m ann
encloseSep (Int -> m -> Tree m ann
forall m ann. Int -> m -> Tree m ann
leaf Int
1 (String -> m
forall a. IsString a => String -> a
fromString String
"["))
                  (Int -> m -> Tree m ann
forall m ann. Int -> m -> Tree m ann
leaf Int
1 (String -> m
forall a. IsString a => String -> a
fromString String
"]"))
                  (Int -> m -> Tree m ann
forall m ann. Int -> m -> Tree m ann
leaf Int
2 (String -> m
forall a. IsString a => String -> a
fromString String
", "))

-- | @tupled = encloseSep \"(\" \")\" \", \"@
tupled :: IsString m => [Tree m ann] -> Tree m ann
tupled :: forall m ann. IsString m => [Tree m ann] -> Tree m ann
tupled = Tree m ann
-> Tree m ann -> Tree m ann -> [Tree m ann] -> Tree m ann
forall m ann.
IsString m =>
Tree m ann
-> Tree m ann -> Tree m ann -> [Tree m ann] -> Tree m ann
encloseSep (Int -> m -> Tree m ann
forall m ann. Int -> m -> Tree m ann
leaf Int
1 (String -> m
forall a. IsString a => String -> a
fromString String
"("))
                    (Int -> m -> Tree m ann
forall m ann. Int -> m -> Tree m ann
leaf Int
1 (String -> m
forall a. IsString a => String -> a
fromString String
")"))
                    (Int -> m -> Tree m ann
forall m ann. Int -> m -> Tree m ann
leaf Int
2 (String -> m
forall a. IsString a => String -> a
fromString String
", "))

-- | Append a separator to all but the last element.
punctuate :: Tree m ann -> [Tree m ann] -> [Tree m ann]
punctuate :: forall m ann. Tree m ann -> [Tree m ann] -> [Tree m ann]
punctuate Tree m ann
_ [] = []
punctuate Tree m ann
_ [Tree m ann
d] = [Tree m ann
d]
punctuate Tree m ann
s (Tree m ann
d : [Tree m ann]
ds) = (Tree m ann
d Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
s) Tree m ann -> [Tree m ann] -> [Tree m ann]
forall a. a -> [a] -> [a]
: Tree m ann -> [Tree m ann] -> [Tree m ann]
forall m ann. Tree m ann -> [Tree m ann] -> [Tree m ann]
punctuate Tree m ann
s [Tree m ann]
ds

---------------------------------------------------------------------
-- Filling
---------------------------------------------------------------------

-- | @width doc f@ renders @doc@ then passes its rendered width to @f@.
width :: Tree m ann -> (Int -> Tree m ann) -> Tree m ann
width :: forall m ann. Tree m ann -> (Int -> Tree m ann) -> Tree m ann
width Tree m ann
d Int -> Tree m ann
f = (Int -> Tree m ann) -> Tree m ann
forall m ann. (Int -> Tree m ann) -> Tree m ann
column ((Int -> Tree m ann) -> Tree m ann)
-> (Int -> Tree m ann) -> Tree m ann
forall a b. (a -> b) -> a -> b
$ \Int
start -> Tree m ann
d Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> (Int -> Tree m ann) -> Tree m ann
forall m ann. (Int -> Tree m ann) -> Tree m ann
column (\Int
end -> Int -> Tree m ann
f (Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start))

-- | @fill n doc@ pads @doc@ to width @n@ with spaces.
fill :: IsString m => Int -> Tree m ann -> Tree m ann
fill :: forall m ann. IsString m => Int -> Tree m ann -> Tree m ann
fill Int
n Tree m ann
d = Tree m ann -> (Int -> Tree m ann) -> Tree m ann
forall m ann. Tree m ann -> (Int -> Tree m ann) -> Tree m ann
width Tree m ann
d ((Int -> Tree m ann) -> Tree m ann)
-> (Int -> Tree m ann) -> Tree m ann
forall a b. (a -> b) -> a -> b
$ \Int
w ->
    if Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n
        then Tree m ann
forall m ann. Tree m ann
emptyDoc
        else Int -> m -> Tree m ann
forall m ann. Int -> m -> Tree m ann
leaf (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
w) (String -> m
forall a. IsString a => String -> a
fromString (Int -> Char -> String
forall a. Int -> a -> [a]
replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
w) Char
' '))

-- | @fillBreak n doc@ pads or breaks after @doc@ if it exceeds @n@.
fillBreak :: IsString m => Int -> Tree m ann -> Tree m ann
fillBreak :: forall m ann. IsString m => Int -> Tree m ann -> Tree m ann
fillBreak Int
n Tree m ann
d = Tree m ann -> (Int -> Tree m ann) -> Tree m ann
forall m ann. Tree m ann -> (Int -> Tree m ann) -> Tree m ann
width Tree m ann
d ((Int -> Tree m ann) -> Tree m ann)
-> (Int -> Tree m ann) -> Tree m ann
forall a b. (a -> b) -> a -> b
$ \Int
w ->
    if Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n
        then Int -> Tree m ann -> Tree m ann
forall m ann. Int -> Tree m ann -> Tree m ann
nest Int
n Tree m ann
forall m ann. Tree m ann
line'
        else Int -> m -> Tree m ann
forall m ann. Int -> m -> Tree m ann
leaf (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
w) (String -> m
forall a. IsString a => String -> a
fromString (Int -> Char -> String
forall a. Int -> a -> [a]
replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
w) Char
' '))

---------------------------------------------------------------------
-- Annotations
---------------------------------------------------------------------

-- | Map over annotations.
reAnnotate :: (ann -> ann') -> Tree m ann -> Tree m ann'
reAnnotate :: forall ann ann' m. (ann -> ann') -> Tree m ann -> Tree m ann'
reAnnotate ann -> ann'
f = (forall a. Doc m ann a -> Doc m ann' a)
-> Mu (Doc m ann) -> Mu (Doc m ann')
forall (f :: * -> *) (g :: * -> *).
(forall a. f a -> g a) -> Mu f -> Mu g
hoistMu Doc m ann a -> Doc m ann' a
forall a. Doc m ann a -> Doc m ann' a
forall {m} {r}. Doc m ann r -> Doc m ann' r
go
  where
    go :: Doc m ann r -> Doc m ann' r
go Doc m ann r
Fail = Doc m ann' r
forall m ann r. Doc m ann r
Fail
    go Doc m ann r
Empty = Doc m ann' r
forall m ann r. Doc m ann r
Empty
    go (Leaf Int
n m
m) = Int -> m -> Doc m ann' r
forall m ann r. Int -> m -> Doc m ann r
Leaf Int
n m
m
    go (Cat r
a r
b) = r -> r -> Doc m ann' r
forall m ann r. r -> r -> Doc m ann r
Cat r
a r
b
    go Doc m ann r
Line = Doc m ann' r
forall m ann r. Doc m ann r
Line
    go (FlatAlt r
a r
b) = r -> r -> Doc m ann' r
forall m ann r. r -> r -> Doc m ann r
FlatAlt r
a r
b
    go (Nest Int
i r
a) = Int -> r -> Doc m ann' r
forall m ann r. Int -> r -> Doc m ann r
Nest Int
i r
a
    go (Union r
a r
b) = r -> r -> Doc m ann' r
forall m ann r. r -> r -> Doc m ann r
Union r
a r
b
    go (Ann ann
a r
x) = ann' -> r -> Doc m ann' r
forall m ann r. ann -> r -> Doc m ann r
Ann (ann -> ann'
f ann
a) r
x
    go (Column Int -> r
k) = (Int -> r) -> Doc m ann' r
forall m ann r. (Int -> r) -> Doc m ann r
Column Int -> r
k
    go (Nesting Int -> r
k) = (Int -> r) -> Doc m ann' r
forall m ann r. (Int -> r) -> Doc m ann r
Nesting Int -> r
k

-- | Remove all annotations.
unAnnotate :: Tree m ann -> Tree m ann'
unAnnotate :: forall m ann ann'. Tree m ann -> Tree m ann'
unAnnotate = (ann -> [ann']) -> Tree m ann -> Tree m ann'
forall ann ann' m. (ann -> [ann']) -> Tree m ann -> Tree m ann'
alterAnnotations ([ann'] -> ann -> [ann']
forall a b. a -> b -> a
const [])

-- | Alter annotations, potentially adding or removing layers.
alterAnnotations :: (ann -> [ann']) -> Tree m ann -> Tree m ann'
alterAnnotations :: forall ann ann' m. (ann -> [ann']) -> Tree m ann -> Tree m ann'
alterAnnotations ann -> [ann']
f = Algebra (Doc m ann) (Tree m ann') -> Mu (Doc m ann) -> Tree m ann'
forall (f :: * -> *) a. Algebra f a -> Mu f -> a
fold (Algebra (Doc m ann) (Tree m ann')
 -> Mu (Doc m ann) -> Tree m ann')
-> Algebra (Doc m ann) (Tree m ann')
-> Mu (Doc m ann)
-> Tree m ann'
forall a b. (a -> b) -> a -> b
$ \case
    Ann ann
a Tree m ann'
x -> (ann' -> Tree m ann' -> Tree m ann')
-> Tree m ann' -> [ann'] -> Tree m ann'
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ann'
a' Tree m ann'
d -> Doc m ann' (Tree m ann') -> Tree m ann'
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (ann' -> Tree m ann' -> Doc m ann' (Tree m ann')
forall m ann r. ann -> r -> Doc m ann r
Ann ann'
a' Tree m ann'
d)) Tree m ann'
x (ann -> [ann']
f ann
a)
    Doc m ann (Tree m ann')
Fail -> Doc m ann' (Tree m ann') -> Tree m ann'
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap Doc m ann' (Tree m ann')
forall m ann r. Doc m ann r
Fail
    Doc m ann (Tree m ann')
Empty -> Doc m ann' (Tree m ann') -> Tree m ann'
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap Doc m ann' (Tree m ann')
forall m ann r. Doc m ann r
Empty
    Leaf Int
n m
m -> Doc m ann' (Tree m ann') -> Tree m ann'
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Int -> m -> Doc m ann' (Tree m ann')
forall m ann r. Int -> m -> Doc m ann r
Leaf Int
n m
m)
    Cat Tree m ann'
a Tree m ann'
b -> Doc m ann' (Tree m ann') -> Tree m ann'
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Tree m ann' -> Tree m ann' -> Doc m ann' (Tree m ann')
forall m ann r. r -> r -> Doc m ann r
Cat Tree m ann'
a Tree m ann'
b)
    Doc m ann (Tree m ann')
Line -> Doc m ann' (Tree m ann') -> Tree m ann'
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap Doc m ann' (Tree m ann')
forall m ann r. Doc m ann r
Line
    FlatAlt Tree m ann'
a Tree m ann'
b -> Doc m ann' (Tree m ann') -> Tree m ann'
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Tree m ann' -> Tree m ann' -> Doc m ann' (Tree m ann')
forall m ann r. r -> r -> Doc m ann r
FlatAlt Tree m ann'
a Tree m ann'
b)
    Nest Int
i Tree m ann'
a -> Doc m ann' (Tree m ann') -> Tree m ann'
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Int -> Tree m ann' -> Doc m ann' (Tree m ann')
forall m ann r. Int -> r -> Doc m ann r
Nest Int
i Tree m ann'
a)
    Union Tree m ann'
a Tree m ann'
b -> Doc m ann' (Tree m ann') -> Tree m ann'
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Tree m ann' -> Tree m ann' -> Doc m ann' (Tree m ann')
forall m ann r. r -> r -> Doc m ann r
Union Tree m ann'
a Tree m ann'
b)
    Column Int -> Tree m ann'
k -> Doc m ann' (Tree m ann') -> Tree m ann'
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap ((Int -> Tree m ann') -> Doc m ann' (Tree m ann')
forall m ann r. (Int -> r) -> Doc m ann r
Column Int -> Tree m ann'
k)
    Nesting Int -> Tree m ann'
k -> Doc m ann' (Tree m ann') -> Tree m ann'
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap ((Int -> Tree m ann') -> Doc m ann' (Tree m ann')
forall m ann r. (Int -> r) -> Doc m ann r
Nesting Int -> Tree m ann'
k)

---------------------------------------------------------------------
-- Optimized group
---------------------------------------------------------------------

-- | Result of checking whether flattening changes a document.
data FlattenResult a
    = Flattened a   -- ^ Flattening produces a different document
    | AlreadyFlat   -- ^ Document is already flat (no FlatAlt/Line)
    | NeverFlat      -- ^ Document can never be flattened (bare Line)
    deriving (Int -> FlattenResult a -> ShowS
[FlattenResult a] -> ShowS
FlattenResult a -> String
(Int -> FlattenResult a -> ShowS)
-> (FlattenResult a -> String)
-> ([FlattenResult a] -> ShowS)
-> Show (FlattenResult a)
forall a. Show a => Int -> FlattenResult a -> ShowS
forall a. Show a => [FlattenResult a] -> ShowS
forall a. Show a => FlattenResult a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> FlattenResult a -> ShowS
showsPrec :: Int -> FlattenResult a -> ShowS
$cshow :: forall a. Show a => FlattenResult a -> String
show :: FlattenResult a -> String
$cshowList :: forall a. Show a => [FlattenResult a] -> ShowS
showList :: [FlattenResult a] -> ShowS
Show, FlattenResult a -> FlattenResult a -> Bool
(FlattenResult a -> FlattenResult a -> Bool)
-> (FlattenResult a -> FlattenResult a -> Bool)
-> Eq (FlattenResult a)
forall a. Eq a => FlattenResult a -> FlattenResult a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => FlattenResult a -> FlattenResult a -> Bool
== :: FlattenResult a -> FlattenResult a -> Bool
$c/= :: forall a. Eq a => FlattenResult a -> FlattenResult a -> Bool
/= :: FlattenResult a -> FlattenResult a -> Bool
Eq)

instance Functor FlattenResult where
    fmap :: forall a b. (a -> b) -> FlattenResult a -> FlattenResult b
fmap a -> b
f (Flattened a
a) = b -> FlattenResult b
forall a. a -> FlattenResult a
Flattened (a -> b
f a
a)
    fmap a -> b
_ FlattenResult a
AlreadyFlat = FlattenResult b
forall a. FlattenResult a
AlreadyFlat
    fmap a -> b
_ FlattenResult a
NeverFlat = FlattenResult b
forall a. FlattenResult a
NeverFlat

-- | Check whether flattening changes a document, and if so,
-- produce the flattened version.
--
-- This is the key optimization for 'group'': by checking first,
-- we avoid creating unnecessary 'Union' nodes.
--
-- Uses direct recursion via 'unwrap' (matching prettyprinter's
-- approach). A zygomorphism formulation is possible but less
-- readable.
changesUponFlattening :: Tree m ann -> FlattenResult (Tree m ann)
changesUponFlattening :: forall m ann. Tree m ann -> FlattenResult (Tree m ann)
changesUponFlattening Tree m ann
t = case Tree m ann -> Doc m ann (Tree m ann)
forall (f :: * -> *). Functor f => Mu f -> f (Mu f)
unwrap Tree m ann
t of
    Doc m ann (Tree m ann)
Fail -> FlattenResult (Tree m ann)
forall a. FlattenResult a
NeverFlat
    Doc m ann (Tree m ann)
Empty -> FlattenResult (Tree m ann)
forall a. FlattenResult a
AlreadyFlat
    Leaf Int
_ m
_ -> FlattenResult (Tree m ann)
forall a. FlattenResult a
AlreadyFlat
    Doc m ann (Tree m ann)
Line -> FlattenResult (Tree m ann)
forall a. FlattenResult a
NeverFlat
    FlatAlt Tree m ann
_ Tree m ann
y -> Tree m ann -> FlattenResult (Tree m ann)
forall a. a -> FlattenResult a
Flattened (Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann
flatten Tree m ann
y)
    Cat Tree m ann
x Tree m ann
y -> case (Tree m ann -> FlattenResult (Tree m ann)
forall m ann. Tree m ann -> FlattenResult (Tree m ann)
changesUponFlattening Tree m ann
x, Tree m ann -> FlattenResult (Tree m ann)
forall m ann. Tree m ann -> FlattenResult (Tree m ann)
changesUponFlattening Tree m ann
y) of
        (Flattened Tree m ann
x', Flattened Tree m ann
y') -> Tree m ann -> FlattenResult (Tree m ann)
forall a. a -> FlattenResult a
Flattened (Tree m ann
x' Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
y')
        (Flattened Tree m ann
x', FlattenResult (Tree m ann)
AlreadyFlat) -> Tree m ann -> FlattenResult (Tree m ann)
forall a. a -> FlattenResult a
Flattened (Tree m ann
x' Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
y)
        (FlattenResult (Tree m ann)
AlreadyFlat, Flattened Tree m ann
y') -> Tree m ann -> FlattenResult (Tree m ann)
forall a. a -> FlattenResult a
Flattened (Tree m ann
x Tree m ann -> Tree m ann -> Tree m ann
forall a. Semigroup a => a -> a -> a
<> Tree m ann
y')
        (FlattenResult (Tree m ann)
AlreadyFlat, FlattenResult (Tree m ann)
AlreadyFlat) -> FlattenResult (Tree m ann)
forall a. FlattenResult a
AlreadyFlat
        (FlattenResult (Tree m ann)
NeverFlat, FlattenResult (Tree m ann)
_) -> FlattenResult (Tree m ann)
forall a. FlattenResult a
NeverFlat
        (FlattenResult (Tree m ann)
_, FlattenResult (Tree m ann)
NeverFlat) -> FlattenResult (Tree m ann)
forall a. FlattenResult a
NeverFlat
    Nest Int
i Tree m ann
x -> (Tree m ann -> Tree m ann)
-> FlattenResult (Tree m ann) -> FlattenResult (Tree m ann)
forall a b. (a -> b) -> FlattenResult a -> FlattenResult b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> Tree m ann -> Tree m ann
forall m ann. Int -> Tree m ann -> Tree m ann
nest Int
i) (Tree m ann -> FlattenResult (Tree m ann)
forall m ann. Tree m ann -> FlattenResult (Tree m ann)
changesUponFlattening Tree m ann
x)
    Union Tree m ann
x Tree m ann
_ -> case Tree m ann -> FlattenResult (Tree m ann)
forall m ann. Tree m ann -> FlattenResult (Tree m ann)
changesUponFlattening Tree m ann
x of
        Flattened Tree m ann
x' -> Tree m ann -> FlattenResult (Tree m ann)
forall a. a -> FlattenResult a
Flattened Tree m ann
x'
        FlattenResult (Tree m ann)
AlreadyFlat -> FlattenResult (Tree m ann)
forall a. FlattenResult a
AlreadyFlat
        FlattenResult (Tree m ann)
NeverFlat -> FlattenResult (Tree m ann)
forall a. FlattenResult a
NeverFlat
    Ann ann
a Tree m ann
x -> (Tree m ann -> Tree m ann)
-> FlattenResult (Tree m ann) -> FlattenResult (Tree m ann)
forall a b. (a -> b) -> FlattenResult a -> FlattenResult b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ann -> Tree m ann -> Tree m ann
forall ann m. ann -> Tree m ann -> Tree m ann
annotate ann
a) (Tree m ann -> FlattenResult (Tree m ann)
forall m ann. Tree m ann -> FlattenResult (Tree m ann)
changesUponFlattening Tree m ann
x)
    Column Int -> Tree m ann
_ -> Tree m ann -> FlattenResult (Tree m ann)
forall a. a -> FlattenResult a
Flattened (Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann
flatten Tree m ann
t)
    Nesting Int -> Tree m ann
_ -> Tree m ann -> FlattenResult (Tree m ann)
forall a. a -> FlattenResult a
Flattened (Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann
flatten Tree m ann
t)

-- | Optimized 'group': avoids creating 'Union' when the document
-- is already flat or can never be flattened.
--
-- @pretty opts (group' doc) = pretty opts (group doc)@
group' :: Tree m ann -> Tree m ann
group' :: forall m ann. Tree m ann -> Tree m ann
group' Tree m ann
x = case Tree m ann -> FlattenResult (Tree m ann)
forall m ann. Tree m ann -> FlattenResult (Tree m ann)
changesUponFlattening Tree m ann
x of
    Flattened Tree m ann
x' -> Tree m ann -> Tree m ann -> Tree m ann
forall m ann. Tree m ann -> Tree m ann -> Tree m ann
union Tree m ann
x' Tree m ann
x
    FlattenResult (Tree m ann)
AlreadyFlat -> Tree m ann
x
    FlattenResult (Tree m ann)
NeverFlat -> Tree m ann
x

---------------------------------------------------------------------
-- Fusion
---------------------------------------------------------------------

-- | Merge adjacent 'Leaf' nodes and collapse nested 'Nest'.
--
-- This is a semantic no-op — the rendered output is identical —
-- but reduces tree size for more efficient layout.
--
-- Operates as a fold (catamorphism). Does not recurse into
-- 'Column'/'Nesting' functions.
--
-- @pretty opts (fuse doc) = pretty opts doc@
fuse :: Semigroup m => Tree m ann -> Tree m ann
fuse :: forall m ann. Semigroup m => Tree m ann -> Tree m ann
fuse = Algebra (Doc m ann) (Tree m ann) -> Tree m ann -> Tree m ann
forall (f :: * -> *) a. Algebra f a -> Mu f -> a
fold (Algebra (Doc m ann) (Tree m ann) -> Tree m ann -> Tree m ann)
-> Algebra (Doc m ann) (Tree m ann) -> Tree m ann -> Tree m ann
forall a b. (a -> b) -> a -> b
$ \case
    -- Merge adjacent leaves: Cat (Leaf n1 m1) (Leaf n2 m2) → Leaf (n1+n2) (m1<>m2)
    Cat Tree m ann
x Tree m ann
y -> case (Tree m ann -> Doc m ann (Tree m ann)
forall (f :: * -> *). Functor f => Mu f -> f (Mu f)
unwrap Tree m ann
x, Tree m ann -> Doc m ann (Tree m ann)
forall (f :: * -> *). Functor f => Mu f -> f (Mu f)
unwrap Tree m ann
y) of
        (Leaf Int
n1 m
m1, Leaf Int
n2 m
m2) -> Algebra (Doc m ann) (Tree m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Int -> m -> Doc m ann (Tree m ann)
forall m ann r. Int -> m -> Doc m ann r
Leaf (Int
n1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n2) (m
m1 m -> m -> m
forall a. Semigroup a => a -> a -> a
<> m
m2))
        (Doc m ann (Tree m ann), Doc m ann (Tree m ann))
_ -> Algebra (Doc m ann) (Tree m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Tree m ann -> Tree m ann -> Doc m ann (Tree m ann)
forall m ann r. r -> r -> Doc m ann r
Cat Tree m ann
x Tree m ann
y)
    -- Collapse nested Nest: Nest i (Nest j x) → Nest (i+j) x
    Nest Int
i Tree m ann
x -> case Tree m ann -> Doc m ann (Tree m ann)
forall (f :: * -> *). Functor f => Mu f -> f (Mu f)
unwrap Tree m ann
x of
        Nest Int
j Tree m ann
inner -> Algebra (Doc m ann) (Tree m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Int -> Tree m ann -> Doc m ann (Tree m ann)
forall m ann r. Int -> r -> Doc m ann r
Nest (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) Tree m ann
inner)
        Doc m ann (Tree m ann)
_ -> Algebra (Doc m ann) (Tree m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap (Int -> Tree m ann -> Doc m ann (Tree m ann)
forall m ann r. Int -> r -> Doc m ann r
Nest Int
i Tree m ann
x)
    Doc m ann (Tree m ann)
other -> Algebra (Doc m ann) (Tree m ann)
forall (f :: * -> *). Functor f => f (Mu f) -> Mu f
wrap Doc m ann (Tree m ann)
other

---------------------------------------------------------------------
-- Trailing whitespace
---------------------------------------------------------------------

-- | Remove trailing whitespace from rendered output.
--
-- Drops spaces at the end of each line. Applied as a
-- post-processing step on the rendered string.
removeTrailingWhitespace :: String -> String
removeTrailingWhitespace :: ShowS
removeTrailingWhitespace = [String] -> String
unlines ([String] -> String) -> (String -> [String]) -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS -> [String] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map ShowS
stripEnd ([String] -> [String])
-> (String -> [String]) -> String -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String]
lines
  where
    stripEnd :: ShowS
stripEnd = ShowS
forall a. [a] -> [a]
reverse ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
' ') ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
forall a. [a] -> [a]
reverse

---------------------------------------------------------------------
-- Tokens
---------------------------------------------------------------------

-- | A single rendered token. The output of the layout algorithm.
data Token m ann
    = TLeaf !Int !m     -- ^ Content with display width
    | TLine !Int        -- ^ Newline followed by @n@ indentation spaces
    | TAnnPush ann      -- ^ Begin annotation scope
    | TAnnPop           -- ^ End annotation scope
    deriving (Int -> Token m ann -> ShowS
[Token m ann] -> ShowS
Token m ann -> String
(Int -> Token m ann -> ShowS)
-> (Token m ann -> String)
-> ([Token m ann] -> ShowS)
-> Show (Token m ann)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall m ann. (Show m, Show ann) => Int -> Token m ann -> ShowS
forall m ann. (Show m, Show ann) => [Token m ann] -> ShowS
forall m ann. (Show m, Show ann) => Token m ann -> String
$cshowsPrec :: forall m ann. (Show m, Show ann) => Int -> Token m ann -> ShowS
showsPrec :: Int -> Token m ann -> ShowS
$cshow :: forall m ann. (Show m, Show ann) => Token m ann -> String
show :: Token m ann -> String
$cshowList :: forall m ann. (Show m, Show ann) => [Token m ann] -> ShowS
showList :: [Token m ann] -> ShowS
Show, Token m ann -> Token m ann -> Bool
(Token m ann -> Token m ann -> Bool)
-> (Token m ann -> Token m ann -> Bool) -> Eq (Token m ann)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall m ann. (Eq m, Eq ann) => Token m ann -> Token m ann -> Bool
$c== :: forall m ann. (Eq m, Eq ann) => Token m ann -> Token m ann -> Bool
== :: Token m ann -> Token m ann -> Bool
$c/= :: forall m ann. (Eq m, Eq ann) => Token m ann -> Token m ann -> Bool
/= :: Token m ann -> Token m ann -> Bool
Eq)

---------------------------------------------------------------------
-- Layout
---------------------------------------------------------------------

-- | Page width configuration.
data PageWidth
    = AvailablePerLine !Int !Double
    -- ^ @AvailablePerLine maxColumns ribbonFraction@
    | Unbounded
    deriving (Int -> PageWidth -> ShowS
[PageWidth] -> ShowS
PageWidth -> String
(Int -> PageWidth -> ShowS)
-> (PageWidth -> String)
-> ([PageWidth] -> ShowS)
-> Show PageWidth
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PageWidth -> ShowS
showsPrec :: Int -> PageWidth -> ShowS
$cshow :: PageWidth -> String
show :: PageWidth -> String
$cshowList :: [PageWidth] -> ShowS
showList :: [PageWidth] -> ShowS
Show, PageWidth -> PageWidth -> Bool
(PageWidth -> PageWidth -> Bool)
-> (PageWidth -> PageWidth -> Bool) -> Eq PageWidth
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PageWidth -> PageWidth -> Bool
== :: PageWidth -> PageWidth -> Bool
$c/= :: PageWidth -> PageWidth -> Bool
/= :: PageWidth -> PageWidth -> Bool
Eq)

-- | Layout options.
newtype LayoutOptions = LayoutOptions
    { LayoutOptions -> PageWidth
layoutPageWidth :: PageWidth
    }
    deriving (Int -> LayoutOptions -> ShowS
[LayoutOptions] -> ShowS
LayoutOptions -> String
(Int -> LayoutOptions -> ShowS)
-> (LayoutOptions -> String)
-> ([LayoutOptions] -> ShowS)
-> Show LayoutOptions
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LayoutOptions -> ShowS
showsPrec :: Int -> LayoutOptions -> ShowS
$cshow :: LayoutOptions -> String
show :: LayoutOptions -> String
$cshowList :: [LayoutOptions] -> ShowS
showList :: [LayoutOptions] -> ShowS
Show, LayoutOptions -> LayoutOptions -> Bool
(LayoutOptions -> LayoutOptions -> Bool)
-> (LayoutOptions -> LayoutOptions -> Bool) -> Eq LayoutOptions
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LayoutOptions -> LayoutOptions -> Bool
== :: LayoutOptions -> LayoutOptions -> Bool
$c/= :: LayoutOptions -> LayoutOptions -> Bool
/= :: LayoutOptions -> LayoutOptions -> Bool
Eq)

-- | 80 columns, ribbon fraction 1.0.
defaultLayoutOptions :: LayoutOptions
defaultLayoutOptions :: LayoutOptions
defaultLayoutOptions = PageWidth -> LayoutOptions
LayoutOptions (Int -> Double -> PageWidth
AvailablePerLine Int
80 Double
1.0)

-- | Commands for the layout pipeline.
data Cmd m ann
    = CDoc !Int (Tree m ann)   -- ^ Process document at nesting level
    | CPopAnn                  -- ^ Emit TAnnPop

-- | Wadler/Leijen layout with one-line lookahead.
--
-- At 'Union' nodes, tries the flattened (first) branch. If the
-- remainder of the line fits within the page width, uses it.
-- Otherwise falls back to the second (default) branch.
layoutPretty :: LayoutOptions -> Tree m ann -> [Token m ann]
layoutPretty :: forall m ann. LayoutOptions -> Tree m ann -> [Token m ann]
layoutPretty LayoutOptions
opts Tree m ann
doc = case Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
0 Int
0 [Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
0 Tree m ann
doc] of
    Maybe [Token m ann]
Nothing -> []  -- should not happen for well-formed docs
    Just [Token m ann]
tokens -> [Token m ann]
tokens
  where
    (Int
maxWidth, Int
ribbonWidth) = case LayoutOptions -> PageWidth
layoutPageWidth LayoutOptions
opts of
        AvailablePerLine Int
w Double
r ->
            let rw :: Int
rw = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
w (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
w)))
             in (Int
w, Int
rw)
        PageWidth
Unbounded -> (Int
forall a. Bounded a => a
maxBound, Int
forall a. Bounded a => a
maxBound)

    -- | Available width for fitting, accounting for ribbon.
    -- The ribbon limits how far past the nesting level we go.
    availableWidth :: Int -> Int -> Int
    availableWidth :: Int -> Int -> Int
availableWidth Int
nl Int
cc = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
maxWidth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cc) (Int
ribbonWidth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nl)

    best :: Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
    best :: forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
_ Int
_ [] = [Token m ann] -> Maybe [Token m ann]
forall a. a -> Maybe a
Just []
    best Int
nl Int
cc (Cmd m ann
CPopAnn : [Cmd m ann]
rest) = (Token m ann
forall m ann. Token m ann
TAnnPop Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
:) ([Token m ann] -> [Token m ann])
-> Maybe [Token m ann] -> Maybe [Token m ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc [Cmd m ann]
rest
    best Int
nl Int
cc (CDoc Int
i Tree m ann
d : [Cmd m ann]
rest) = case Tree m ann -> Doc m ann (Tree m ann)
forall (f :: * -> *). Functor f => Mu f -> f (Mu f)
unwrap Tree m ann
d of
        Doc m ann (Tree m ann)
Fail -> Maybe [Token m ann]
forall a. Maybe a
Nothing
        Doc m ann (Tree m ann)
Empty -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc [Cmd m ann]
rest
        Leaf Int
len m
m -> (Int -> m -> Token m ann
forall m ann. Int -> m -> Token m ann
TLeaf Int
len m
m Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
:) ([Token m ann] -> [Token m ann])
-> Maybe [Token m ann] -> Maybe [Token m ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl (Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) [Cmd m ann]
rest
        Cat Tree m ann
x Tree m ann
y -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
y Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Doc m ann (Tree m ann)
Line -> (Int -> Token m ann
forall m ann. Int -> Token m ann
TLine Int
i Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
:) ([Token m ann] -> [Token m ann])
-> Maybe [Token m ann] -> Maybe [Token m ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
i Int
i [Cmd m ann]
rest
        FlatAlt Tree m ann
x Tree m ann
_ -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Nest Int
j Tree m ann
x -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Union Tree m ann
x Tree m ann
y ->
            case Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest) of
                Just [Token m ann]
flatTokens | Int -> [Token m ann] -> Bool
forall m ann. Int -> [Token m ann] -> Bool
fits (Int -> Int -> Int
availableWidth Int
nl Int
cc) [Token m ann]
flatTokens -> [Token m ann] -> Maybe [Token m ann]
forall a. a -> Maybe a
Just [Token m ann]
flatTokens
                Maybe [Token m ann]
_ -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
y Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Ann ann
a Tree m ann
x -> (ann -> Token m ann
forall m ann. ann -> Token m ann
TAnnPush ann
a Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
:) ([Token m ann] -> [Token m ann])
-> Maybe [Token m ann] -> Maybe [Token m ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: Cmd m ann
forall m ann. Cmd m ann
CPopAnn Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Column Int -> Tree m ann
f -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i (Int -> Tree m ann
f Int
cc) Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Nesting Int -> Tree m ann
f -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i (Int -> Tree m ann
f Int
i) Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)

-- | One-line lookahead: does the content fit in @w@ characters
-- before the next line break?
fits :: Int -> [Token m ann] -> Bool
fits :: forall m ann. Int -> [Token m ann] -> Bool
fits Int
w [Token m ann]
_ | Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = Bool
False
fits Int
_ [] = Bool
True
fits Int
_ (TLine Int
_ : [Token m ann]
_) = Bool
True
fits Int
w (TLeaf Int
len m
_ : [Token m ann]
rest) = Int -> [Token m ann] -> Bool
forall m ann. Int -> [Token m ann] -> Bool
fits (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len) [Token m ann]
rest
fits Int
w (TAnnPush ann
_ : [Token m ann]
rest) = Int -> [Token m ann] -> Bool
forall m ann. Int -> [Token m ann] -> Bool
fits Int
w [Token m ann]
rest
fits Int
w (Token m ann
TAnnPop : [Token m ann]
rest) = Int -> [Token m ann] -> Bool
forall m ann. Int -> [Token m ann] -> Bool
fits Int
w [Token m ann]
rest

-- | Smart layout with multi-line lookahead.
--
-- Like 'layoutPretty', but the fitting predicate checks beyond
-- the first line break — it continues checking until it finds a
-- line that starts at the same or shallower indentation level.
-- This prevents surprises where content looks like it fits but
-- subsequent lines overflow.
--
-- Use this for deeply nested structures where 'layoutPretty'
-- makes suboptimal choices.
layoutSmart :: LayoutOptions -> Tree m ann -> [Token m ann]
layoutSmart :: forall m ann. LayoutOptions -> Tree m ann -> [Token m ann]
layoutSmart LayoutOptions
opts Tree m ann
doc = case Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
0 Int
0 [Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
0 Tree m ann
doc] of
    Maybe [Token m ann]
Nothing -> []
    Just [Token m ann]
tokens -> [Token m ann]
tokens
  where
    (Int
maxWidth, Int
ribbonWidth) = case LayoutOptions -> PageWidth
layoutPageWidth LayoutOptions
opts of
        AvailablePerLine Int
w Double
r ->
            let rw :: Int
rw = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
w (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
w)))
             in (Int
w, Int
rw)
        PageWidth
Unbounded -> (Int
forall a. Bounded a => a
maxBound, Int
forall a. Bounded a => a
maxBound)

    availableWidth :: Int -> Int -> Int
    availableWidth :: Int -> Int -> Int
availableWidth Int
nl Int
cc = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
maxWidth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cc) (Int
ribbonWidth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nl)

    best :: Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
    best :: forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
_ Int
_ [] = [Token m ann] -> Maybe [Token m ann]
forall a. a -> Maybe a
Just []
    best Int
nl Int
cc (Cmd m ann
CPopAnn : [Cmd m ann]
rest) = (Token m ann
forall m ann. Token m ann
TAnnPop Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
:) ([Token m ann] -> [Token m ann])
-> Maybe [Token m ann] -> Maybe [Token m ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc [Cmd m ann]
rest
    best Int
nl Int
cc (CDoc Int
i Tree m ann
d : [Cmd m ann]
rest) = case Tree m ann -> Doc m ann (Tree m ann)
forall (f :: * -> *). Functor f => Mu f -> f (Mu f)
unwrap Tree m ann
d of
        Doc m ann (Tree m ann)
Fail -> Maybe [Token m ann]
forall a. Maybe a
Nothing
        Doc m ann (Tree m ann)
Empty -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc [Cmd m ann]
rest
        Leaf Int
len m
m -> (Int -> m -> Token m ann
forall m ann. Int -> m -> Token m ann
TLeaf Int
len m
m Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
:) ([Token m ann] -> [Token m ann])
-> Maybe [Token m ann] -> Maybe [Token m ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl (Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) [Cmd m ann]
rest
        Cat Tree m ann
x Tree m ann
y -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
y Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Doc m ann (Tree m ann)
Line -> (Int -> Token m ann
forall m ann. Int -> Token m ann
TLine Int
i Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
:) ([Token m ann] -> [Token m ann])
-> Maybe [Token m ann] -> Maybe [Token m ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
i Int
i [Cmd m ann]
rest
        FlatAlt Tree m ann
x Tree m ann
_ -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Nest Int
j Tree m ann
x -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Union Tree m ann
x Tree m ann
y ->
            case Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest) of
                Just [Token m ann]
flatTokens | Int -> Int -> [Token m ann] -> Bool
forall m ann. Int -> Int -> [Token m ann] -> Bool
fitsSmart Int
nl (Int -> Int -> Int
availableWidth Int
nl Int
cc) [Token m ann]
flatTokens -> [Token m ann] -> Maybe [Token m ann]
forall a. a -> Maybe a
Just [Token m ann]
flatTokens
                Maybe [Token m ann]
_ -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
y Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Ann ann
a Tree m ann
x -> (ann -> Token m ann
forall m ann. ann -> Token m ann
TAnnPush ann
a Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
:) ([Token m ann] -> [Token m ann])
-> Maybe [Token m ann] -> Maybe [Token m ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: Cmd m ann
forall m ann. Cmd m ann
CPopAnn Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Column Int -> Tree m ann
f -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i (Int -> Tree m ann
f Int
cc) Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Nesting Int -> Tree m ann
f -> Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> Int -> [Cmd m ann] -> Maybe [Token m ann]
best Int
nl Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i (Int -> Tree m ann
f Int
i) Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)

-- | Multi-line lookahead: checks fitting beyond the first line break.
--
-- After a line break, continues checking if the content on the
-- next line fits — but only if the next line is more deeply
-- indented than the current nesting level. Stops when it finds
-- a line at the same or shallower indentation.
fitsSmart :: Int -> Int -> [Token m ann] -> Bool
fitsSmart :: forall m ann. Int -> Int -> [Token m ann] -> Bool
fitsSmart Int
_ Int
w [Token m ann]
_ | Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = Bool
False
fitsSmart Int
_ Int
_ [] = Bool
True
fitsSmart Int
nl Int
_ (TLine Int
i : [Token m ann]
rest)
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
nl = Int -> Int -> [Token m ann] -> Bool
forall m ann. Int -> Int -> [Token m ann] -> Bool
fitsSmart Int
nl (Int
i) [Token m ann]
rest  -- deeper: keep checking
    | Bool
otherwise = Bool
True                 -- same or shallower: done
fitsSmart Int
nl Int
w (TLeaf Int
len m
_ : [Token m ann]
rest) = Int -> Int -> [Token m ann] -> Bool
forall m ann. Int -> Int -> [Token m ann] -> Bool
fitsSmart Int
nl (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len) [Token m ann]
rest
fitsSmart Int
nl Int
w (TAnnPush ann
_ : [Token m ann]
rest) = Int -> Int -> [Token m ann] -> Bool
forall m ann. Int -> Int -> [Token m ann] -> Bool
fitsSmart Int
nl Int
w [Token m ann]
rest
fitsSmart Int
nl Int
w (Token m ann
TAnnPop : [Token m ann]
rest) = Int -> Int -> [Token m ann] -> Bool
forall m ann. Int -> Int -> [Token m ann] -> Bool
fitsSmart Int
nl Int
w [Token m ann]
rest

-- | Compact layout: no width-sensitivity, always breaks.
--
-- Every 'FlatAlt' uses its default, every 'Union' uses the
-- narrow branch. 'Column' and 'Nesting' receive 0.
layoutCompact :: Tree m ann -> [Token m ann]
layoutCompact :: forall m ann. Tree m ann -> [Token m ann]
layoutCompact Tree m ann
doc = Int -> [Tree m ann] -> [Token m ann]
forall {m} {ann}. Int -> [Mu (Doc m ann)] -> [Token m ann]
go Int
0 [Tree m ann
doc]
  where
    go :: Int -> [Mu (Doc m ann)] -> [Token m ann]
go Int
_ [] = []
    go Int
cc (Mu (Doc m ann)
d : [Mu (Doc m ann)]
rest) = case Mu (Doc m ann) -> Doc m ann (Mu (Doc m ann))
forall (f :: * -> *). Functor f => Mu f -> f (Mu f)
unwrap Mu (Doc m ann)
d of
        Doc m ann (Mu (Doc m ann))
Fail -> []
        Doc m ann (Mu (Doc m ann))
Empty -> Int -> [Mu (Doc m ann)] -> [Token m ann]
go Int
cc [Mu (Doc m ann)]
rest
        Leaf Int
len m
m -> Int -> m -> Token m ann
forall m ann. Int -> m -> Token m ann
TLeaf Int
len m
m Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
: Int -> [Mu (Doc m ann)] -> [Token m ann]
go (Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) [Mu (Doc m ann)]
rest
        Cat Mu (Doc m ann)
x Mu (Doc m ann)
y -> Int -> [Mu (Doc m ann)] -> [Token m ann]
go Int
cc (Mu (Doc m ann)
x Mu (Doc m ann) -> [Mu (Doc m ann)] -> [Mu (Doc m ann)]
forall a. a -> [a] -> [a]
: Mu (Doc m ann)
y Mu (Doc m ann) -> [Mu (Doc m ann)] -> [Mu (Doc m ann)]
forall a. a -> [a] -> [a]
: [Mu (Doc m ann)]
rest)
        Doc m ann (Mu (Doc m ann))
Line -> Int -> Token m ann
forall m ann. Int -> Token m ann
TLine Int
0 Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
: Int -> [Mu (Doc m ann)] -> [Token m ann]
go Int
0 [Mu (Doc m ann)]
rest
        FlatAlt Mu (Doc m ann)
x Mu (Doc m ann)
_ -> Int -> [Mu (Doc m ann)] -> [Token m ann]
go Int
cc (Mu (Doc m ann)
x Mu (Doc m ann) -> [Mu (Doc m ann)] -> [Mu (Doc m ann)]
forall a. a -> [a] -> [a]
: [Mu (Doc m ann)]
rest)
        Nest Int
_ Mu (Doc m ann)
x -> Int -> [Mu (Doc m ann)] -> [Token m ann]
go Int
cc (Mu (Doc m ann)
x Mu (Doc m ann) -> [Mu (Doc m ann)] -> [Mu (Doc m ann)]
forall a. a -> [a] -> [a]
: [Mu (Doc m ann)]
rest)  -- ignore nesting in compact
        Union Mu (Doc m ann)
_ Mu (Doc m ann)
y -> Int -> [Mu (Doc m ann)] -> [Token m ann]
go Int
cc (Mu (Doc m ann)
y Mu (Doc m ann) -> [Mu (Doc m ann)] -> [Mu (Doc m ann)]
forall a. a -> [a] -> [a]
: [Mu (Doc m ann)]
rest)  -- always narrow
        Ann ann
a Mu (Doc m ann)
x -> ann -> Token m ann
forall m ann. ann -> Token m ann
TAnnPush ann
a Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
: Int -> [Mu (Doc m ann)] -> [Token m ann]
go Int
cc (Mu (Doc m ann)
x Mu (Doc m ann) -> [Mu (Doc m ann)] -> [Mu (Doc m ann)]
forall a. a -> [a] -> [a]
: [Mu (Doc m ann)]
rest) -- no pop tracking in compact
        Column Int -> Mu (Doc m ann)
f -> Int -> [Mu (Doc m ann)] -> [Token m ann]
go Int
cc (Int -> Mu (Doc m ann)
f Int
cc Mu (Doc m ann) -> [Mu (Doc m ann)] -> [Mu (Doc m ann)]
forall a. a -> [a] -> [a]
: [Mu (Doc m ann)]
rest)
        Nesting Int -> Mu (Doc m ann)
f -> Int -> [Mu (Doc m ann)] -> [Token m ann]
go Int
cc (Int -> Mu (Doc m ann)
f Int
0 Mu (Doc m ann) -> [Mu (Doc m ann)] -> [Mu (Doc m ann)]
forall a. a -> [a] -> [a]
: [Mu (Doc m ann)]
rest)

---------------------------------------------------------------------
-- Streaming layout
---------------------------------------------------------------------

-- | Wadler/Leijen layout producing a 'Nu' stream.
--
-- Unlike 'layoutPretty' which materializes a @[Token]@ list,
-- 'layoutStream' produces a @Nu (Cons (Token m ann))@ — a seed
-- + step function that generates tokens lazily on demand.
-- Construction is O(1); tokens are computed as they are consumed.
layoutStream :: LayoutOptions -> Tree m ann -> Nu (Cons (Token m ann))
layoutStream :: forall m ann.
LayoutOptions -> Tree m ann -> Nu (Cons (Token m ann))
layoutStream LayoutOptions
opts Tree m ann
doc = ((Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann]))
-> (Int, [Cmd m ann]) -> Nu (Cons (Token m ann))
forall (f :: * -> *) a. (a -> f a) -> a -> Nu f
Nu (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall m ann.
(Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
step (Int
0, [Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
0 Tree m ann
doc])
  where
    pageWidth :: Int
pageWidth = case LayoutOptions -> PageWidth
layoutPageWidth LayoutOptions
opts of
        AvailablePerLine Int
w Double
_ -> Int
w
        PageWidth
Unbounded -> Int
forall a. Bounded a => a
maxBound

    step :: (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
    step :: forall m ann.
(Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
step (Int
_, []) = Cons (Token m ann) (Int, [Cmd m ann])
forall a b. Cons a b
Nil
    step (Int
cc, Cmd m ann
CPopAnn : [Cmd m ann]
rest) = Token m ann
-> (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall a b. a -> b -> Cons a b
Cons (Token m ann
forall m ann. Token m ann
TAnnPop) (Int
cc, [Cmd m ann]
rest)
    step (Int
cc, CDoc Int
i Tree m ann
d : [Cmd m ann]
rest) = case Tree m ann -> Doc m ann (Tree m ann)
forall (f :: * -> *). Functor f => Mu f -> f (Mu f)
unwrap Tree m ann
d of
        Doc m ann (Tree m ann)
Fail -> Cons (Token m ann) (Int, [Cmd m ann])
forall a b. Cons a b
Nil
        Doc m ann (Tree m ann)
Empty -> (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall m ann.
(Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
step (Int
cc, [Cmd m ann]
rest)
        Leaf Int
len m
m -> Token m ann
-> (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall a b. a -> b -> Cons a b
Cons (Int -> m -> Token m ann
forall m ann. Int -> m -> Token m ann
TLeaf Int
len m
m) (Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len, [Cmd m ann]
rest)
        Cat Tree m ann
x Tree m ann
y -> (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall m ann.
(Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
step (Int
cc, Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
y Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Doc m ann (Tree m ann)
Line -> Token m ann
-> (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall a b. a -> b -> Cons a b
Cons (Int -> Token m ann
forall m ann. Int -> Token m ann
TLine Int
i) (Int
i, [Cmd m ann]
rest)
        FlatAlt Tree m ann
x Tree m ann
_ -> (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall m ann.
(Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
step (Int
cc, Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Nest Int
j Tree m ann
x -> (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall m ann.
(Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
step (Int
cc, Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Union Tree m ann
x Tree m ann
y ->
            -- Speculative evaluation: try the flat branch via list-based
            -- layout. If it fits, commit to the flat branch; otherwise
            -- continue streaming with the narrow branch.
            case Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
cc (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest) of
                Just [Token m ann]
flatTokens | Int -> [Token m ann] -> Bool
forall m ann. Int -> [Token m ann] -> Bool
fits (Int
pageWidth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cc) [Token m ann]
flatTokens ->
                    (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall m ann.
(Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
step (Int
cc, Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
                Maybe [Token m ann]
_ -> (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall m ann.
(Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
step (Int
cc, Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
y Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Ann ann
a Tree m ann
x -> Token m ann
-> (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall a b. a -> b -> Cons a b
Cons (ann -> Token m ann
forall m ann. ann -> Token m ann
TAnnPush ann
a) (Int
cc, Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: Cmd m ann
forall m ann. Cmd m ann
CPopAnn Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Column Int -> Tree m ann
f -> (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall m ann.
(Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
step (Int
cc, Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i (Int -> Tree m ann
f Int
cc) Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)
        Nesting Int -> Tree m ann
f -> (Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
forall m ann.
(Int, [Cmd m ann]) -> Cons (Token m ann) (Int, [Cmd m ann])
step (Int
cc, Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i (Int -> Tree m ann
f Int
i) Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest)

    -- Reuse the list-based layout for Union speculative evaluation
    bestList :: Int -> [Cmd m ann] -> Maybe [Token m ann]
    bestList :: forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
_ [] = [Token m ann] -> Maybe [Token m ann]
forall a. a -> Maybe a
Just []
    bestList Int
cc' (Cmd m ann
CPopAnn : [Cmd m ann]
rest') = (Token m ann
forall m ann. Token m ann
TAnnPop Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
:) ([Token m ann] -> [Token m ann])
-> Maybe [Token m ann] -> Maybe [Token m ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
cc' [Cmd m ann]
rest'
    bestList Int
cc' (CDoc Int
i' Tree m ann
d' : [Cmd m ann]
rest') = case Tree m ann -> Doc m ann (Tree m ann)
forall (f :: * -> *). Functor f => Mu f -> f (Mu f)
unwrap Tree m ann
d' of
        Doc m ann (Tree m ann)
Fail -> Maybe [Token m ann]
forall a. Maybe a
Nothing
        Doc m ann (Tree m ann)
Empty -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
cc' [Cmd m ann]
rest'
        Leaf Int
len m
m -> (Int -> m -> Token m ann
forall m ann. Int -> m -> Token m ann
TLeaf Int
len m
m Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
:) ([Token m ann] -> [Token m ann])
-> Maybe [Token m ann] -> Maybe [Token m ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList (Int
cc' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) [Cmd m ann]
rest'
        Cat Tree m ann
x Tree m ann
y -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
cc' (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i' Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i' Tree m ann
y Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest')
        Doc m ann (Tree m ann)
Line -> (Int -> Token m ann
forall m ann. Int -> Token m ann
TLine Int
i' Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
:) ([Token m ann] -> [Token m ann])
-> Maybe [Token m ann] -> Maybe [Token m ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
i' [Cmd m ann]
rest'
        FlatAlt Tree m ann
x Tree m ann
_ -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
cc' (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i' Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest')
        Nest Int
j Tree m ann
x -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
cc' (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc (Int
i' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest')
        Union Tree m ann
x Tree m ann
y -> case Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
cc' (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i' Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest') of
            Just [Token m ann]
flatTokens | Int -> [Token m ann] -> Bool
forall m ann. Int -> [Token m ann] -> Bool
fits (Int
pageWidth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cc') [Token m ann]
flatTokens -> [Token m ann] -> Maybe [Token m ann]
forall a. a -> Maybe a
Just [Token m ann]
flatTokens
            Maybe [Token m ann]
_ -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
cc' (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i' Tree m ann
y Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest')
        Ann ann
a Tree m ann
x -> (ann -> Token m ann
forall m ann. ann -> Token m ann
TAnnPush ann
a Token m ann -> [Token m ann] -> [Token m ann]
forall a. a -> [a] -> [a]
:) ([Token m ann] -> [Token m ann])
-> Maybe [Token m ann] -> Maybe [Token m ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
cc' (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i' Tree m ann
x Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: Cmd m ann
forall m ann. Cmd m ann
CPopAnn Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest')
        Column Int -> Tree m ann
f -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
cc' (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i' (Int -> Tree m ann
f Int
cc') Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest')
        Nesting Int -> Tree m ann
f -> Int -> [Cmd m ann] -> Maybe [Token m ann]
forall m ann. Int -> [Cmd m ann] -> Maybe [Token m ann]
bestList Int
cc' (Int -> Tree m ann -> Cmd m ann
forall m ann. Int -> Tree m ann -> Cmd m ann
CDoc Int
i' (Int -> Tree m ann
f Int
i') Cmd m ann -> [Cmd m ann] -> [Cmd m ann]
forall a. a -> [a] -> [a]
: [Cmd m ann]
rest')

    tokenWidth :: Token m ann -> Int
tokenWidth (TLeaf Int
len m
_) = Int
len
    tokenWidth Token m ann
_ = Int
0

-- | Render a 'Nu' token stream to the output monoid.
renderStream :: (Monoid m, IsString m) => Nu (Cons (Token m ann)) -> m
renderStream :: forall m ann.
(Monoid m, IsString m) =>
Nu (Cons (Token m ann)) -> m
renderStream (Nu a -> Cons (Token m ann) a
step a
seed) = a -> m
go a
seed
  where
    go :: a -> m
go a
s = case a -> Cons (Token m ann) a
step a
s of
        Cons (Token m ann) a
Nil -> m
forall a. Monoid a => a
mempty
        Cons (TLeaf Int
_ m
m) a
s' -> m
m m -> m -> m
forall a. Semigroup a => a -> a -> a
<> a -> m
go a
s'
        Cons (TLine Int
i) a
s' -> String -> m
forall a. IsString a => String -> a
fromString (Char
'\n' Char -> ShowS
forall a. a -> [a] -> [a]
: Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
i Char
' ') m -> m -> m
forall a. Semigroup a => a -> a -> a
<> a -> m
go a
s'
        Cons (TAnnPush ann
_) a
s' -> a -> m
go a
s'
        Cons (Token m ann
TAnnPop) a
s' -> a -> m
go a
s'

-- | Streaming layout + render.
--
-- @prettyStream opts = renderStream . layoutStream opts@
prettyStream :: (Monoid m, IsString m) => LayoutOptions -> Tree m ann -> m
prettyStream :: forall m ann.
(Monoid m, IsString m) =>
LayoutOptions -> Tree m ann -> m
prettyStream LayoutOptions
opts = Nu (Cons (Token m ann)) -> m
forall m ann.
(Monoid m, IsString m) =>
Nu (Cons (Token m ann)) -> m
renderStream (Nu (Cons (Token m ann)) -> m)
-> (Tree m ann -> Nu (Cons (Token m ann))) -> Tree m ann -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LayoutOptions -> Tree m ann -> Nu (Cons (Token m ann))
forall m ann.
LayoutOptions -> Tree m ann -> Nu (Cons (Token m ann))
layoutStream LayoutOptions
opts

---------------------------------------------------------------------
-- Rendering
---------------------------------------------------------------------

-- | Render a token stream to the output monoid, discarding annotations.
render :: (Monoid m, IsString m) => [Token m ann] -> m
render :: forall m ann. (Monoid m, IsString m) => [Token m ann] -> m
render = [Token m ann] -> m
forall m ann. (Monoid m, IsString m) => [Token m ann] -> m
go
  where
    go :: [Token a ann] -> a
go [] = a
forall a. Monoid a => a
mempty
    go (TLeaf Int
_ a
m : [Token a ann]
rest) = a
m a -> a -> a
forall a. Semigroup a => a -> a -> a
<> [Token a ann] -> a
go [Token a ann]
rest
    go (TLine Int
i : [Token a ann]
rest) = String -> a
forall a. IsString a => String -> a
fromString (Char
'\n' Char -> ShowS
forall a. a -> [a] -> [a]
: Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
i Char
' ') a -> a -> a
forall a. Semigroup a => a -> a -> a
<> [Token a ann] -> a
go [Token a ann]
rest
    go (TAnnPush ann
_ : [Token a ann]
rest) = [Token a ann] -> a
go [Token a ann]
rest
    go (Token a ann
TAnnPop : [Token a ann]
rest) = [Token a ann] -> a
go [Token a ann]
rest

-- | Lay out and render a document.
--
-- @pretty opts = render . layoutPretty opts@
pretty :: (Monoid m, IsString m) => LayoutOptions -> Tree m ann -> m
pretty :: forall m ann.
(Monoid m, IsString m) =>
LayoutOptions -> Tree m ann -> m
pretty LayoutOptions
opts = [Token m ann] -> m
forall m ann. (Monoid m, IsString m) => [Token m ann] -> m
render ([Token m ann] -> m)
-> (Tree m ann -> [Token m ann]) -> Tree m ann -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LayoutOptions -> Tree m ann -> [Token m ann]
forall m ann. LayoutOptions -> Tree m ann -> [Token m ann]
layoutPretty LayoutOptions
opts