3 files changed, 61 insertions, 8 deletions
diff --git a/src/Erebos/TextFormat.hs b/src/Erebos/TextFormat.hs
index 88fe0c2..6674ebc 100644
--- a/src/Erebos/TextFormat.hs
+++ b/src/Erebos/TextFormat.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE OverloadedStrings #-}
+
 module Erebos.TextFormat (
     FormattedText,
     plainText,
@@ -8,8 +10,11 @@ module Erebos.TextFormat (
     Color(..),
     setForegroundColor, setBackgroundColor,
 
+    endWithNewline,
+
     renderPlainText,
     formattedTextLength,
+    formattedTextHeight,
 ) where
 
 import Data.Text (Text)
@@ -35,14 +40,40 @@ setBackgroundColor :: Color -> TextStyle -> TextStyle
 setBackgroundColor color (CustomTextColor fg _) = CustomTextColor fg (Just color)
 
 
+endWithNewline :: FormattedText -> FormattedText
+endWithNewline = EndWithNewline
+
+
 renderPlainText :: FormattedText -> Text
 renderPlainText = \case
     PlainText text -> text
     ConcatenatedText ftexts -> mconcat $ map renderPlainText ftexts
     FormattedText _ ftext -> renderPlainText ftext
+    EndWithNewline ftext -> let res = renderPlainText ftext
+                             in case T.unsnoc res of
+                                    Just ( _, '\n') -> res
+                                    _               -> res <> "\n"
 
 formattedTextLength :: FormattedText -> Int
 formattedTextLength = \case
     PlainText text -> T.length text
     ConcatenatedText ftexts -> sum $ map formattedTextLength ftexts
     FormattedText _ ftext -> formattedTextLength ftext
+    EndWithNewline ftext -> formattedTextLength ftext
+
+formattedTextHeight :: FormattedText -> Int
+formattedTextHeight = countLines . collectParts
+  where
+    collectParts = \case
+        PlainText text -> [ text ]
+        ConcatenatedText ftexts -> concatMap collectParts ftexts
+        FormattedText _ ftext -> collectParts ftext
+        EndWithNewline ftext -> collectParts ftext
+    countLines (t : ts)
+        | T.null t = countLines ts
+        | otherwise = 1 + countLines (dropLine (t : ts))
+    countLines [] = 0
+    dropLine (t : ts)
+        | Just ( '\n', t' ) <- T.uncons (T.dropWhile (/= '\n') t) = t' : ts
+        | otherwise = dropLine ts
+    dropLine [] = []
diff --git a/src/Erebos/TextFormat/Ansi.hs b/src/Erebos/TextFormat/Ansi.hs
index c0ff978..faec0ad 100644
--- a/src/Erebos/TextFormat/Ansi.hs
+++ b/src/Erebos/TextFormat/Ansi.hs
@@ -8,6 +8,8 @@ module Erebos.TextFormat.Ansi (
 ) where
 
 import Control.Applicative
+import Control.Monad.State
+import Control.Monad.Writer
 
 import Data.String
 import Data.Text (Text)
@@ -20,19 +22,38 @@ newtype AnsiText = AnsiText { fromAnsiText :: Text }
     deriving (Eq, Ord, Semigroup, Monoid, IsString)
 
 
+data RenderState = RenderState
+    { rsEndedWithNewline :: Bool
+    }
+
+initialRenderState :: RenderState
+initialRenderState = RenderState
+    { rsEndedWithNewline = True
+    }
+
 renderAnsiText :: FormattedText -> AnsiText
-renderAnsiText = AnsiText . go ( Nothing, Nothing )
+renderAnsiText ft = AnsiText $ T.concat $ execWriter $ flip evalStateT initialRenderState $ go ( Nothing, Nothing ) ft
   where
+    go :: ( Maybe Color, Maybe Color ) -> FormattedText -> StateT RenderState (Writer [ Text ]) ()
     go cur@( cfg, cbg ) = \case
-        PlainText text -> text
-        ConcatenatedText ftexts -> mconcat $ map (go cur) ftexts
-        FormattedText (CustomTextColor fg bg) ftext -> mconcat
-            [ ansiColor fg bg
-            , go ( fg <|> cfg, bg <|> cbg ) ftext
-            , ansiColor
+        PlainText text -> do
+            tell [ text ]
+            case T.unsnoc text of
+                Just ( _, c ) -> modify (\s -> s { rsEndedWithNewline = c == '\n' })
+                Nothing -> return ()
+        ConcatenatedText ftexts -> mconcat <$> mapM (go cur) ftexts
+        FormattedText (CustomTextColor fg bg) ftext -> do
+            tell [ ansiColor fg bg ]
+            go ( fg <|> cfg, bg <|> cbg ) ftext
+            tell [ ansiColor
                 (if fg /= cfg then cfg <|> Just DefaultColor else Nothing)
                 (if bg /= cbg then cbg <|> Just DefaultColor else Nothing)
-            ]
+                ]
+        EndWithNewline ftext -> do
+            go cur ftext
+            gets rsEndedWithNewline >>= \case
+                True -> return ()
+                False -> tell [ "\n" ] >> modify (\s -> s { rsEndedWithNewline = True })
 
 
 ansiColor :: Maybe Color -> Maybe Color -> Text
diff --git a/src/Erebos/TextFormat/Types.hs b/src/Erebos/TextFormat/Types.hs
index a03bc71..a93026d 100644
--- a/src/Erebos/TextFormat/Types.hs
+++ b/src/Erebos/TextFormat/Types.hs
@@ -12,6 +12,7 @@ data FormattedText
     = PlainText Text
     | ConcatenatedText [ FormattedText ]
     | FormattedText TextStyle FormattedText
+    | EndWithNewline FormattedText
 
 instance IsString FormattedText where
     fromString = PlainText . fromString