{- © 2019 Serokell <hi@serokell.io>
 - © 2019 Lars Jellema <lars.jellema@gmail.com>
 -
 - SPDX-License-Identifier: MPL-2.0
 -}

{-# LANGUAGE DeriveFoldable, DeriveFunctor, FlexibleInstances,
             OverloadedStrings, StandaloneDeriving #-}

-- | This module implements a layer around the prettyprinter package, making it
-- easier to use.
module Nixfmt.Predoc
    ( text
    , sepBy
    , hcat
    , base
    , group
    , nest
    , softline'
    , line'
    , softline
    , line
    , hardspace
    , hardline
    , emptyline
    , newline
    , Doc
    , Pretty
    , pretty
    , fixup
    , layout
    , textWidth
    ) where

import Data.List (intersperse)
import Data.Text as Text (Text, concat, length, pack, replicate)

-- | Sequential Spacings are reduced to a single Spacing by taking the maximum.
-- This means that e.g. a Space followed by an Emptyline results in just an
-- Emptyline.
data Spacing
    = Softbreak
    | Break
    | Hardspace
    | Softspace
    | Space
    | Hardline
    | Emptyline
    | Newlines Int
    deriving (Int -> Spacing -> ShowS
[Spacing] -> ShowS
Spacing -> String
(Int -> Spacing -> ShowS)
-> (Spacing -> String) -> ([Spacing] -> ShowS) -> Show Spacing
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Spacing] -> ShowS
$cshowList :: [Spacing] -> ShowS
show :: Spacing -> String
$cshow :: Spacing -> String
showsPrec :: Int -> Spacing -> ShowS
$cshowsPrec :: Int -> Spacing -> ShowS
Show, Spacing -> Spacing -> Bool
(Spacing -> Spacing -> Bool)
-> (Spacing -> Spacing -> Bool) -> Eq Spacing
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Spacing -> Spacing -> Bool
$c/= :: Spacing -> Spacing -> Bool
== :: Spacing -> Spacing -> Bool
$c== :: Spacing -> Spacing -> Bool
Eq, Eq Spacing
Eq Spacing
-> (Spacing -> Spacing -> Ordering)
-> (Spacing -> Spacing -> Bool)
-> (Spacing -> Spacing -> Bool)
-> (Spacing -> Spacing -> Bool)
-> (Spacing -> Spacing -> Bool)
-> (Spacing -> Spacing -> Spacing)
-> (Spacing -> Spacing -> Spacing)
-> Ord Spacing
Spacing -> Spacing -> Bool
Spacing -> Spacing -> Ordering
Spacing -> Spacing -> Spacing
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Spacing -> Spacing -> Spacing
$cmin :: Spacing -> Spacing -> Spacing
max :: Spacing -> Spacing -> Spacing
$cmax :: Spacing -> Spacing -> Spacing
>= :: Spacing -> Spacing -> Bool
$c>= :: Spacing -> Spacing -> Bool
> :: Spacing -> Spacing -> Bool
$c> :: Spacing -> Spacing -> Bool
<= :: Spacing -> Spacing -> Bool
$c<= :: Spacing -> Spacing -> Bool
< :: Spacing -> Spacing -> Bool
$c< :: Spacing -> Spacing -> Bool
compare :: Spacing -> Spacing -> Ordering
$ccompare :: Spacing -> Spacing -> Ordering
$cp1Ord :: Eq Spacing
Ord)

data DocAnn
    -- | Node Group docs indicates either all or none of the Spaces and Breaks
    -- in docs should be converted to line breaks.
    = Group
    -- | Node (Nest n) doc indicates all line starts in doc should be indented
    -- by n more spaces than the surrounding Base.
    | Nest Int
    -- | Node Base doc sets the base indentation that Nests should be relative
    -- to to the indentation of the line where the Base starts.
    | Base
    deriving (Int -> DocAnn -> ShowS
[DocAnn] -> ShowS
DocAnn -> String
(Int -> DocAnn -> ShowS)
-> (DocAnn -> String) -> ([DocAnn] -> ShowS) -> Show DocAnn
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DocAnn] -> ShowS
$cshowList :: [DocAnn] -> ShowS
show :: DocAnn -> String
$cshow :: DocAnn -> String
showsPrec :: Int -> DocAnn -> ShowS
$cshowsPrec :: Int -> DocAnn -> ShowS
Show, DocAnn -> DocAnn -> Bool
(DocAnn -> DocAnn -> Bool)
-> (DocAnn -> DocAnn -> Bool) -> Eq DocAnn
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DocAnn -> DocAnn -> Bool
$c/= :: DocAnn -> DocAnn -> Bool
== :: DocAnn -> DocAnn -> Bool
$c== :: DocAnn -> DocAnn -> Bool
Eq)

-- | Single document element. Documents are modeled as lists of these elements
-- in order to make concatenation simple.
data DocE
    = Text Text
    | Spacing Spacing
    | Node DocAnn Doc
    deriving (Int -> DocE -> ShowS
[DocE] -> ShowS
DocE -> String
(Int -> DocE -> ShowS)
-> (DocE -> String) -> ([DocE] -> ShowS) -> Show DocE
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DocE] -> ShowS
$cshowList :: [DocE] -> ShowS
show :: DocE -> String
$cshow :: DocE -> String
showsPrec :: Int -> DocE -> ShowS
$cshowsPrec :: Int -> DocE -> ShowS
Show, DocE -> DocE -> Bool
(DocE -> DocE -> Bool) -> (DocE -> DocE -> Bool) -> Eq DocE
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DocE -> DocE -> Bool
$c/= :: DocE -> DocE -> Bool
== :: DocE -> DocE -> Bool
$c== :: DocE -> DocE -> Bool
Eq)

type Doc = [DocE]

class Pretty a where
    pretty :: a -> Doc

instance Pretty Text where
    pretty :: Text -> [DocE]
pretty = DocE -> [DocE]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DocE -> [DocE]) -> (Text -> DocE) -> Text -> [DocE]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> DocE
Text

instance Pretty String where
    pretty :: String -> [DocE]
pretty = DocE -> [DocE]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DocE -> [DocE]) -> (String -> DocE) -> String -> [DocE]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> DocE
Text (Text -> DocE) -> (String -> Text) -> String -> DocE
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
pack

instance Pretty Doc where
    pretty :: [DocE] -> [DocE]
pretty = [DocE] -> [DocE]
forall a. a -> a
id

instance Pretty a => Pretty (Maybe a) where
    pretty :: Maybe a -> [DocE]
pretty Maybe a
Nothing  = [DocE]
forall a. Monoid a => a
mempty
    pretty (Just a
x) = a -> [DocE]
forall a. Pretty a => a -> [DocE]
pretty a
x

text :: Text -> Doc
text :: Text -> [DocE]
text Text
"" = []
text Text
t  = [Text -> DocE
Text Text
t]

group :: Pretty a => a -> Doc
group :: a -> [DocE]
group = DocE -> [DocE]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DocE -> [DocE]) -> (a -> DocE) -> a -> [DocE]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DocAnn -> [DocE] -> DocE
Node DocAnn
Group ([DocE] -> DocE) -> (a -> [DocE]) -> a -> DocE
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> [DocE]
forall a. Pretty a => a -> [DocE]
pretty

-- | @nest n doc@ sets the indentation for lines in @doc@ to @n@ more than the
-- indentation of the part before it. This is based on the actual indentation of
-- the line, rather than the indentation it should have used: If multiple
-- indentation levels start on the same line, only the last indentation level
-- will be applied on the next line. This prevents unnecessary nesting.
nest :: Int -> Doc -> Doc
nest :: Int -> [DocE] -> [DocE]
nest Int
level = DocE -> [DocE]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DocE -> [DocE]) -> ([DocE] -> DocE) -> [DocE] -> [DocE]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DocAnn -> [DocE] -> DocE
Node (Int -> DocAnn
Nest Int
level)

base :: Doc -> Doc
base :: [DocE] -> [DocE]
base = DocE -> [DocE]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DocE -> [DocE]) -> ([DocE] -> DocE) -> [DocE] -> [DocE]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DocAnn -> [DocE] -> DocE
Node DocAnn
Base

softline' :: Doc
softline' :: [DocE]
softline' = [Spacing -> DocE
Spacing Spacing
Softbreak]

line' :: Doc
line' :: [DocE]
line' = [Spacing -> DocE
Spacing Spacing
Break]

softline :: Doc
softline :: [DocE]
softline = [Spacing -> DocE
Spacing Spacing
Softspace]

line :: Doc
line :: [DocE]
line = [Spacing -> DocE
Spacing Spacing
Space]

hardspace :: Doc
hardspace :: [DocE]
hardspace = [Spacing -> DocE
Spacing Spacing
Hardspace]

hardline :: Doc
hardline :: [DocE]
hardline = [Spacing -> DocE
Spacing Spacing
Hardline]

emptyline :: Doc
emptyline :: [DocE]
emptyline = [Spacing -> DocE
Spacing Spacing
Emptyline]

newline :: Doc
newline :: [DocE]
newline = [Spacing -> DocE
Spacing (Int -> Spacing
Newlines Int
1)]

sepBy :: Pretty a => Doc -> [a] -> Doc
sepBy :: [DocE] -> [a] -> [DocE]
sepBy [DocE]
separator = [[DocE]] -> [DocE]
forall a. Monoid a => [a] -> a
mconcat ([[DocE]] -> [DocE]) -> ([a] -> [[DocE]]) -> [a] -> [DocE]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [DocE] -> [[DocE]] -> [[DocE]]
forall a. a -> [a] -> [a]
intersperse [DocE]
separator ([[DocE]] -> [[DocE]]) -> ([a] -> [[DocE]]) -> [a] -> [[DocE]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> [DocE]) -> [a] -> [[DocE]]
forall a b. (a -> b) -> [a] -> [b]
map a -> [DocE]
forall a. Pretty a => a -> [DocE]
pretty

-- | Concatenate documents horizontally without spacing.
hcat :: Pretty a => [a] -> Doc
hcat :: [a] -> [DocE]
hcat = [[DocE]] -> [DocE]
forall a. Monoid a => [a] -> a
mconcat ([[DocE]] -> [DocE]) -> ([a] -> [[DocE]]) -> [a] -> [DocE]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> [DocE]) -> [a] -> [[DocE]]
forall a b. (a -> b) -> [a] -> [b]
map a -> [DocE]
forall a. Pretty a => a -> [DocE]
pretty

isSpacing :: DocE -> Bool
isSpacing :: DocE -> Bool
isSpacing (Spacing Spacing
_) = Bool
True
isSpacing DocE
_           = Bool
False

spanEnd :: (a -> Bool) -> [a] -> ([a], [a])
spanEnd :: (a -> Bool) -> [a] -> ([a], [a])
spanEnd a -> Bool
p = ([a] -> [a]) -> ([a], [a]) -> ([a], [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [a] -> [a]
forall a. [a] -> [a]
reverse (([a], [a]) -> ([a], [a]))
-> ([a] -> ([a], [a])) -> [a] -> ([a], [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Bool) -> [a] -> ([a], [a])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span a -> Bool
p ([a] -> ([a], [a])) -> ([a] -> [a]) -> [a] -> ([a], [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. [a] -> [a]
reverse

-- | Fix up a Doc in multiple stages:
-- - First, all spacings are moved out of Groups and Nests and empty Groups and
--   Nests are removed.
-- - Now, all consecutive Spacings are ensured to be in the same list, so each
--   sequence of Spacings can be merged into a single one.
-- - Finally, Spacings right before a Nest should be moved inside in order to
--   get the right indentation.
fixup :: Doc -> Doc
fixup :: [DocE] -> [DocE]
fixup = [DocE] -> [DocE]
moveLinesIn ([DocE] -> [DocE]) -> ([DocE] -> [DocE]) -> [DocE] -> [DocE]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [DocE] -> [DocE]
mergeLines ([DocE] -> [DocE]) -> ([DocE] -> [DocE]) -> [DocE] -> [DocE]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DocE -> [DocE]) -> [DocE] -> [DocE]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DocE -> [DocE]
moveLinesOut

moveLinesOut :: DocE -> Doc
moveLinesOut :: DocE -> [DocE]
moveLinesOut (Node DocAnn
ann [DocE]
xs) =
    let movedOut :: [DocE]
movedOut     = (DocE -> [DocE]) -> [DocE] -> [DocE]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DocE -> [DocE]
moveLinesOut [DocE]
xs
        ([DocE]
pre, [DocE]
rest)  = (DocE -> Bool) -> [DocE] -> ([DocE], [DocE])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span DocE -> Bool
isSpacing [DocE]
movedOut
        ([DocE]
post, [DocE]
body) = (DocE -> Bool) -> [DocE] -> ([DocE], [DocE])
forall a. (a -> Bool) -> [a] -> ([a], [a])
spanEnd DocE -> Bool
isSpacing [DocE]
rest
    in case [DocE]
body of
            [] -> [DocE]
pre [DocE] -> [DocE] -> [DocE]
forall a. [a] -> [a] -> [a]
++ [DocE]
post
            [DocE]
_  -> [DocE]
pre [DocE] -> [DocE] -> [DocE]
forall a. [a] -> [a] -> [a]
++ (DocAnn -> [DocE] -> DocE
Node DocAnn
ann [DocE]
body DocE -> [DocE] -> [DocE]
forall a. a -> [a] -> [a]
: [DocE]
post)

moveLinesOut DocE
x = [DocE
x]

mergeSpacings :: Spacing -> Spacing -> Spacing
mergeSpacings :: Spacing -> Spacing -> Spacing
mergeSpacings Spacing
x Spacing
y | Spacing
x Spacing -> Spacing -> Bool
forall a. Ord a => a -> a -> Bool
> Spacing
y               = Spacing -> Spacing -> Spacing
mergeSpacings Spacing
y Spacing
x
mergeSpacings Spacing
Break        Spacing
Softspace    = Spacing
Space
mergeSpacings Spacing
Break        Spacing
Hardspace    = Spacing
Space
mergeSpacings Spacing
Softbreak    Spacing
Hardspace    = Spacing
Softspace
mergeSpacings (Newlines Int
x) (Newlines Int
y) = Int -> Spacing
Newlines (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y)
mergeSpacings Spacing
Emptyline    (Newlines Int
x) = Int -> Spacing
Newlines (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
mergeSpacings Spacing
Hardspace    (Newlines Int
x) = Int -> Spacing
Newlines Int
x
mergeSpacings Spacing
_            (Newlines Int
x) = Int -> Spacing
Newlines (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
mergeSpacings Spacing
_            Spacing
y            = Spacing
y

mergeLines :: Doc -> Doc
mergeLines :: [DocE] -> [DocE]
mergeLines []                           = []
mergeLines (Spacing Spacing
a : Spacing Spacing
b : [DocE]
xs) = [DocE] -> [DocE]
mergeLines ([DocE] -> [DocE]) -> [DocE] -> [DocE]
forall a b. (a -> b) -> a -> b
$ Spacing -> DocE
Spacing (Spacing -> Spacing -> Spacing
mergeSpacings Spacing
a Spacing
b) DocE -> [DocE] -> [DocE]
forall a. a -> [a] -> [a]
: [DocE]
xs
mergeLines (Text Text
a : Text Text
b : [DocE]
xs)       = [DocE] -> [DocE]
mergeLines ([DocE] -> [DocE]) -> [DocE] -> [DocE]
forall a b. (a -> b) -> a -> b
$ Text -> DocE
Text (Text
a Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
b) DocE -> [DocE] -> [DocE]
forall a. a -> [a] -> [a]
: [DocE]
xs
mergeLines (Node DocAnn
ann [DocE]
xs : [DocE]
ys)           = DocAnn -> [DocE] -> DocE
Node DocAnn
ann ([DocE] -> [DocE]
mergeLines [DocE]
xs) DocE -> [DocE] -> [DocE]
forall a. a -> [a] -> [a]
: [DocE] -> [DocE]
mergeLines [DocE]
ys
mergeLines (DocE
x : [DocE]
xs)                     = DocE
x DocE -> [DocE] -> [DocE]
forall a. a -> [a] -> [a]
: [DocE] -> [DocE]
mergeLines [DocE]
xs

moveLinesIn :: Doc -> Doc
moveLinesIn :: [DocE] -> [DocE]
moveLinesIn [] = []
moveLinesIn (Spacing Spacing
l : Node (Nest Int
level) [DocE]
xs : [DocE]
ys) =
    DocAnn -> [DocE] -> DocE
Node (Int -> DocAnn
Nest Int
level) (Spacing -> DocE
Spacing Spacing
l DocE -> [DocE] -> [DocE]
forall a. a -> [a] -> [a]
: [DocE] -> [DocE]
moveLinesIn [DocE]
xs) DocE -> [DocE] -> [DocE]
forall a. a -> [a] -> [a]
: [DocE] -> [DocE]
moveLinesIn [DocE]
ys

moveLinesIn (Node DocAnn
ann [DocE]
xs : [DocE]
ys) =
    DocAnn -> [DocE] -> DocE
Node DocAnn
ann ([DocE] -> [DocE]
moveLinesIn [DocE]
xs) DocE -> [DocE] -> [DocE]
forall a. a -> [a] -> [a]
: [DocE] -> [DocE]
moveLinesIn [DocE]
ys

moveLinesIn (DocE
x : [DocE]
xs) = DocE
x DocE -> [DocE] -> [DocE]
forall a. a -> [a] -> [a]
: [DocE] -> [DocE]
moveLinesIn [DocE]
xs

layout :: Pretty a => Int -> a -> Text
layout :: Int -> a -> Text
layout Int
w = Int -> [DocE] -> Text
layoutGreedy Int
w ([DocE] -> Text) -> (a -> [DocE]) -> a -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [DocE] -> [DocE]
fixup ([DocE] -> [DocE]) -> (a -> [DocE]) -> a -> [DocE]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> [DocE]
forall a. Pretty a => a -> [DocE]
pretty

-- 1. Move and merge Spacings.
-- 2. Convert Softlines to Grouped Lines and Hardspaces to Texts.
-- 3. For each Text or Group, try to fit as much as possible on a line
-- 4. For each Group, if it fits on a single line, render it that way.
-- 5. If not, convert lines to hardlines and unwrap the group

-- | To support i18n, this function needs to be patched.
textWidth :: Text -> Int
textWidth :: Text -> Int
textWidth = Text -> Int
Text.length

-- | Attempt to fit a list of documents in a single line of a specific width.
fits :: Int -> Doc -> Maybe Text
fits :: Int -> [DocE] -> Maybe Text
fits Int
c [DocE]
_ | Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = Maybe Text
forall a. Maybe a
Nothing
fits Int
_ [] = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
""
fits Int
c (DocE
x:[DocE]
xs) = case DocE
x of
    Text Text
t               -> (Text
tText -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>) (Text -> Text) -> Maybe Text -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [DocE] -> Maybe Text
fits (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Text -> Int
textWidth Text
t) [DocE]
xs
    Spacing Spacing
Softbreak    -> Int -> [DocE] -> Maybe Text
fits Int
c [DocE]
xs
    Spacing Spacing
Break        -> Int -> [DocE] -> Maybe Text
fits Int
c [DocE]
xs
    Spacing Spacing
Softspace    -> (Text
" "Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>) (Text -> Text) -> Maybe Text -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [DocE] -> Maybe Text
fits (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [DocE]
xs
    Spacing Spacing
Space        -> (Text
" "Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>) (Text -> Text) -> Maybe Text -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [DocE] -> Maybe Text
fits (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [DocE]
xs
    Spacing Spacing
Hardspace    -> (Text
" "Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>) (Text -> Text) -> Maybe Text -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [DocE] -> Maybe Text
fits (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [DocE]
xs
    Spacing Spacing
Hardline     -> Maybe Text
forall a. Maybe a
Nothing
    Spacing Spacing
Emptyline    -> Maybe Text
forall a. Maybe a
Nothing
    Spacing (Newlines Int
_) -> Maybe Text
forall a. Maybe a
Nothing
    Node DocAnn
_ [DocE]
ys            -> Int -> [DocE] -> Maybe Text
fits Int
c ([DocE] -> Maybe Text) -> [DocE] -> Maybe Text
forall a b. (a -> b) -> a -> b
$ [DocE]
ys [DocE] -> [DocE] -> [DocE]
forall a. [a] -> [a] -> [a]
++ [DocE]
xs

-- | Find the width of the first line in a list of documents, using target
-- width 0, which always forces line breaks when possible.
firstLineWidth :: Doc -> Int
firstLineWidth :: [DocE] -> Int
firstLineWidth []                       = Int
0
firstLineWidth (Text Text
t : [DocE]
xs)            = Text -> Int
textWidth Text
t Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [DocE] -> Int
firstLineWidth [DocE]
xs
firstLineWidth (Spacing Spacing
Hardspace : [DocE]
xs) = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [DocE] -> Int
firstLineWidth [DocE]
xs
firstLineWidth (Spacing Spacing
_ : [DocE]
_)          = Int
0
firstLineWidth (Node DocAnn
_ [DocE]
xs : [DocE]
ys)         = [DocE] -> Int
firstLineWidth ([DocE]
xs [DocE] -> [DocE] -> [DocE]
forall a. [a] -> [a] -> [a]
++ [DocE]
ys)

-- | Check if the first line in a list of documents fits a target width given
-- a maximum width, without breaking up groups.
firstLineFits :: Int -> Int -> Doc -> Bool
firstLineFits :: Int -> Int -> [DocE] -> Bool
firstLineFits Int
targetWidth Int
maxWidth [DocE]
docs = Int -> [DocE] -> Bool
go Int
maxWidth [DocE]
docs
    where go :: Int -> [DocE] -> Bool
go Int
c [DocE]
_ | Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0                = Bool
False
          go Int
c []                       = Int
maxWidth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
targetWidth
          go Int
c (Text Text
t : [DocE]
xs)            = Int -> [DocE] -> Bool
go (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Text -> Int
textWidth Text
t) [DocE]
xs
          go Int
c (Spacing Spacing
Hardspace : [DocE]
xs) = Int -> [DocE] -> Bool
go (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [DocE]
xs
          go Int
c (Spacing Spacing
_ : [DocE]
_)          = Int
maxWidth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
targetWidth
          go Int
c (Node DocAnn
Group [DocE]
ys : [DocE]
xs)     =
              case Int -> [DocE] -> Maybe Text
fits (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- [DocE] -> Int
firstLineWidth [DocE]
xs) [DocE]
ys of
                   Maybe Text
Nothing -> Int -> [DocE] -> Bool
go Int
c ([DocE]
ys [DocE] -> [DocE] -> [DocE]
forall a. [a] -> [a] -> [a]
++ [DocE]
xs)
                   Just Text
t  -> Int -> [DocE] -> Bool
go (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Text -> Int
textWidth Text
t) [DocE]
xs

          go Int
c (Node DocAnn
_ [DocE]
ys : [DocE]
xs)         = Int -> [DocE] -> Bool
go Int
c ([DocE]
ys [DocE] -> [DocE] -> [DocE]
forall a. [a] -> [a] -> [a]
++ [DocE]
xs)

-- |
data Chunk = Chunk Int DocE

indent :: Int -> Int -> Text
indent :: Int -> Int -> Text
indent Int
n Int
i = Int -> Text -> Text
Text.replicate Int
n Text
"\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text -> Text
Text.replicate Int
i Text
" "

unChunk :: Chunk -> DocE
unChunk :: Chunk -> DocE
unChunk (Chunk Int
_ DocE
doc) = DocE
doc

-- tw   Target Width
-- cc   Current Column
-- ci   Current Indentation
-- ti   Target Indentation
layoutGreedy :: Int -> Doc -> Text
layoutGreedy :: Int -> [DocE] -> Text
layoutGreedy Int
tw [DocE]
doc = [Text] -> Text
Text.concat ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$ Int -> Int -> [Chunk] -> [Text]
go Int
0 Int
0 [Int -> DocE -> Chunk
Chunk Int
0 (DocE -> Chunk) -> DocE -> Chunk
forall a b. (a -> b) -> a -> b
$ DocAnn -> [DocE] -> DocE
Node DocAnn
Group [DocE]
doc]
    where go :: Int -> Int -> [Chunk] -> [Text]
go Int
_ Int
_ [] = []
          go Int
cc Int
ci (Chunk Int
ti DocE
x : [Chunk]
xs) = case DocE
x of
            Text Text
t               -> Text
t Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> Int -> [Chunk] -> [Text]
go (Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Text -> Int
textWidth Text
t) Int
ci [Chunk]
xs

            Spacing Spacing
Break        -> Int -> Int -> Text
indent Int
1 Int
ti Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> Int -> [Chunk] -> [Text]
go Int
ti Int
ti [Chunk]
xs
            Spacing Spacing
Space        -> Int -> Int -> Text
indent Int
1 Int
ti Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> Int -> [Chunk] -> [Text]
go Int
ti Int
ti [Chunk]
xs
            Spacing Spacing
Hardspace    -> Text
" "         Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> Int -> [Chunk] -> [Text]
go (Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
ci [Chunk]
xs
            Spacing Spacing
Hardline     -> Int -> Int -> Text
indent Int
1 Int
ti Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> Int -> [Chunk] -> [Text]
go Int
ti Int
ti [Chunk]
xs
            Spacing Spacing
Emptyline    -> Int -> Int -> Text
indent Int
2 Int
ti Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> Int -> [Chunk] -> [Text]
go Int
ti Int
ti [Chunk]
xs
            Spacing (Newlines Int
n) -> Int -> Int -> Text
indent Int
n Int
ti Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> Int -> [Chunk] -> [Text]
go Int
ti Int
ti [Chunk]
xs

            Spacing Spacing
Softbreak
              | Int -> Int -> [DocE] -> Bool
firstLineFits (Int
tw Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cc) (Int
tw Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ti) ((Chunk -> DocE) -> [Chunk] -> [DocE]
forall a b. (a -> b) -> [a] -> [b]
map Chunk -> DocE
unChunk [Chunk]
xs)
                                 -> Int -> Int -> [Chunk] -> [Text]
go Int
cc Int
ci [Chunk]
xs
              | Bool
otherwise        -> Int -> Int -> Text
indent Int
1 Int
ti Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> Int -> [Chunk] -> [Text]
go Int
ti Int
ti [Chunk]
xs

            Spacing Spacing
Softspace
              | Int -> Int -> [DocE] -> Bool
firstLineFits (Int
tw Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
tw Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ti) ((Chunk -> DocE) -> [Chunk] -> [DocE]
forall a b. (a -> b) -> [a] -> [b]
map Chunk -> DocE
unChunk [Chunk]
xs)
                                 -> Text
" " Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> Int -> [Chunk] -> [Text]
go (Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
ci [Chunk]
xs
              | Bool
otherwise        -> Int -> Int -> Text
indent Int
1 Int
ti Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> Int -> [Chunk] -> [Text]
go Int
ti Int
ti [Chunk]
xs

            Node (Nest Int
l) [DocE]
ys     -> Int -> Int -> [Chunk] -> [Text]
go Int
cc Int
ci ([Chunk] -> [Text]) -> [Chunk] -> [Text]
forall a b. (a -> b) -> a -> b
$ (DocE -> Chunk) -> [DocE] -> [Chunk]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> DocE -> Chunk
Chunk (Int
ti Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l)) [DocE]
ys [Chunk] -> [Chunk] -> [Chunk]
forall a. [a] -> [a] -> [a]
++ [Chunk]
xs
            Node DocAnn
Base [DocE]
ys         -> Int -> Int -> [Chunk] -> [Text]
go Int
cc Int
ci ([Chunk] -> [Text]) -> [Chunk] -> [Text]
forall a b. (a -> b) -> a -> b
$ (DocE -> Chunk) -> [DocE] -> [Chunk]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> DocE -> Chunk
Chunk Int
ci) [DocE]
ys [Chunk] -> [Chunk] -> [Chunk]
forall a. [a] -> [a] -> [a]
++ [Chunk]
xs
            Node DocAnn
Group [DocE]
ys        ->
                case Int -> [DocE] -> Maybe Text
fits (Int
tw Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
- [DocE] -> Int
firstLineWidth ((Chunk -> DocE) -> [Chunk] -> [DocE]
forall a b. (a -> b) -> [a] -> [b]
map Chunk -> DocE
unChunk [Chunk]
xs)) [DocE]
ys of
                     Maybe Text
Nothing     -> Int -> Int -> [Chunk] -> [Text]
go Int
cc Int
ci ([Chunk] -> [Text]) -> [Chunk] -> [Text]
forall a b. (a -> b) -> a -> b
$ (DocE -> Chunk) -> [DocE] -> [Chunk]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> DocE -> Chunk
Chunk Int
ti) [DocE]
ys [Chunk] -> [Chunk] -> [Chunk]
forall a. [a] -> [a] -> [a]
++ [Chunk]
xs
                     Just Text
t      -> Text
t Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> Int -> [Chunk] -> [Text]
go (Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Text -> Int
textWidth Text
t) Int
ci [Chunk]
xs