diff --git a/src/Tidy.purs b/src/Tidy.purs
index 79d288c..e9b7f22 100644
--- a/src/Tidy.purs
+++ b/src/Tidy.purs
@@ -32,7 +32,7 @@ import Dodo as Dodo
 import Partial.Unsafe (unsafeCrashWith)
 import PureScript.CST.Errors (RecoveredError(..))
 import PureScript.CST.Types (Binder(..), ClassFundep(..), ClassHead, Comment(..), DataCtor(..), DataHead, DataMembers(..), Declaration(..), Delimited, DelimitedNonEmpty, DoStatement(..), Export(..), Expr(..), FixityOp(..), Foreign(..), Guarded(..), GuardedExpr(..), Ident, IfThenElse, Import(..), ImportDecl(..), Instance(..), InstanceBinding(..), InstanceHead, Label, Labeled(..), LetBinding(..), LineFeed, Module(..), ModuleBody(..), ModuleHeader(..), Name(..), OneOrDelimited(..), Operator, PatternGuard(..), Proper, QualifiedName(..), RecordLabeled(..), RecordUpdate(..), Row(..), Separated(..), SourceStyle(..), SourceToken, Token(..), Type(..), TypeVarBinding(..), ValueBindingFields, Where(..), Wrapped(..))
-import Tidy.Doc (FormatDoc, align, alignCurrentColumn, anchor, blockComment, break, flattenMax, flexDoubleBreak, flexGroup, flexSoftBreak, flexSpaceBreak, forceMinSourceBreaks, fromDoc, indent, joinWith, joinWithMap, leadingLineComment, locally, softBreak, softSpace, sourceBreak, space, spaceBreak, text, trailingLineComment)
+import Tidy.Doc (FormatDoc, align, alignCurrentColumn, anchor, break, flattenMax, flexDoubleBreak, flexGroup, flexSoftBreak, flexSpaceBreak, forceMinSourceBreaks, fromDoc, indent, joinWith, joinWithMap, leadingBlockComment, leadingLineComment, locally, softBreak, softSpace, sourceBreak, space, spaceBreak, text, trailingBlockComment, trailingLineComment)
 import Tidy.Doc (FormatDoc, toDoc) as Exports
 import Tidy.Hang (HangingDoc, HangingOp(..), hang, hangApp, hangBreak, hangOps, hangWithIndent)
 import Tidy.Hang as Hang
@@ -125,13 +125,19 @@ type Format f e a = FormatOptions e a -> f -> FormatDoc a
 type FormatHanging f e a = FormatOptions e a -> f -> HangingDoc a
 type FormatSpace a = FormatDoc a -> FormatDoc a -> FormatDoc a
 
-formatComment :: forall l a. (String -> FormatDoc a) -> Comment l -> FormatDoc a -> FormatDoc a
-formatComment lineComment com next = case com of
+formatComment
+  :: forall l a
+   . (String -> FormatDoc a -> FormatDoc a)
+  -> (String -> FormatDoc a -> FormatDoc a)
+  -> Comment l
+  -> FormatDoc a
+  -> FormatDoc a
+formatComment lineComment blockComment com next = case com of
   Comment str
     | SCU.take 2 str == "--" ->
-        lineComment str `break` next
+        lineComment str next
     | otherwise ->
-        blockComment str `space` next
+        blockComment str next
   Line _ n ->
     sourceBreak n next
   Space _ ->
@@ -140,8 +146,8 @@ formatComment lineComment com next = case com of
 formatWithComments :: forall a. Array (Comment LineFeed) -> Array (Comment Void) -> FormatDoc a -> FormatDoc a
 formatWithComments leading trailing doc =
   foldr
-    (formatComment leadingLineComment)
-    (doc `space` foldr (formatComment trailingLineComment) mempty trailing)
+    (formatComment leadingLineComment leadingBlockComment)
+    (doc <> foldr (formatComment trailingLineComment trailingBlockComment) mempty trailing)
     leading
 
 formatToken :: forall a r. { unicode :: UnicodeOption | r } -> SourceToken -> FormatDoc a
@@ -155,8 +161,14 @@ formatToken conf tok = formatWithComments tok.leadingComments tok.trailingCommen
 
 formatRawString :: forall a. String -> FormatDoc a
 formatRawString = splitLines >>> Array.uncons >>> foldMap \{ head, tail } ->
-  text head `break` locally (_ { indent = 0, indentSpaces = "" }) do
-    joinWithMap break text tail
+  if Array.null tail then
+    text head
+  else
+    fromDoc $ Dodo.lines
+      [ Dodo.text head
+      , Dodo.locally (_ { indent = 0, indentSpaces = "" }) do
+          Array.intercalate Dodo.break $ Dodo.text <$> tail
+      ]
 
 formatString :: forall a. String -> FormatDoc a
 formatString = splitStringEscapeLines >>> Array.uncons >>> foldMap \{ head, tail } ->
@@ -188,7 +200,7 @@ formatModule conf (Module { header: ModuleHeader header, body: ModuleBody body }
         ImportWrapSource ->
           locally (_ { pageWidth = top, ribbonRatio = 1.0 }) imports
     , forceMinSourceBreaks 2 $ formatTopLevelGroups conf body.decls
-    , foldr (formatComment leadingLineComment) mempty body.trailingComments
+    , foldr (formatComment leadingLineComment leadingBlockComment) mempty body.trailingComments
     ]
   where
   imports =
@@ -1066,7 +1078,7 @@ formatParens format conf (Wrapped { open, value, close }) =
 
 formatParensBlock :: forall e a b. Format b e a -> Format (Wrapped b) e a
 formatParensBlock format conf (Wrapped { open, value, close }) =
-  commentedFlexGroup open softSpace conf do
+  flexGroup $ formatToken conf open `softSpace` do
     align 2 (anchor (format conf value))
       `softBreak` formatToken conf close
 
@@ -1109,7 +1121,8 @@ formatList :: forall e a b. FormatSpace a -> FormatSpace a -> Int -> FormatGroup
 formatList openSpace closeSpace alignment grouped format conf { open, head, tail, close } =
   case grouped of
     Grouped ->
-      commentedFlexGroup open openSpace conf listElems
+      flexGroup $ formatToken conf open
+        `openSpace` listElems
     NotGrouped ->
       formatToken conf open
         `openSpace` listElems
@@ -1154,14 +1167,6 @@ toQualifiedOperatorTree precMap opNs =
 overLabel :: forall a b c. (a -> b) -> Labeled a c -> Labeled b c
 overLabel k (Labeled lbl) = Labeled lbl { label = k lbl.label }
 
--- TODO: Properly fix this grouping issue by buffering comments in FormatDoc.
--- This would let us group the semantic elements of the doc and not the comments.
-commentedFlexGroup :: forall e a. SourceToken -> FormatSpace a -> Format (FormatDoc a) e a
-commentedFlexGroup tok spc conf doc =
-  formatWithComments tok.leadingComments [] $ flexGroup do
-    formatToken conf (tok { leadingComments = [] })
-      `spc` doc
-
 data DeclGroup
   = DeclGroupValueSignature Ident
   | DeclGroupValue Ident
diff --git a/src/Tidy/Doc.purs b/src/Tidy/Doc.purs
index 7042f08..e92fb26 100644
--- a/src/Tidy/Doc.purs
+++ b/src/Tidy/Doc.purs
@@ -1,14 +1,16 @@
 module Tidy.Doc
   ( FormatDoc(..)
+  , LeadingComment(..)
+  , TrailingComment(..)
   , ForceBreak(..)
   , text
   , leadingLineComment
   , trailingLineComment
-  , blockComment
+  , leadingBlockComment
+  , trailingBlockComment
   , anchor
   , flatten
   , flattenMax
-  , forceBreak
   , indent
   , align
   , alignCurrentColumn
@@ -28,6 +30,8 @@ module Tidy.Doc
   , fromDoc
   , toDoc
   , mapDoc
+  , breakDoc
+  , breaks
   , joinWithMap
   , joinWith
   ) where
@@ -36,7 +40,7 @@ import Prelude
 
 import Control.Alternative (guard)
 import Data.Array as Array
-import Data.Foldable (class Foldable, foldMap, foldl, intercalate)
+import Data.Foldable (class Foldable, foldl, intercalate)
 import Data.Maybe (Maybe(..), fromMaybe)
 import Data.Monoid (power)
 import Data.String as String
@@ -47,15 +51,95 @@ import Dodo as Dodo
 import Dodo.Internal (LocalOptions)
 import Tidy.Util (splitLines)
 
-data ForceBreak
-  = ForceNone
-  | ForceSpace
-  | ForceBreak
+data ForceBreak = ForceNone | ForceSpace | ForceBreak
 
 derive instance eqForceBreak :: Eq ForceBreak
 derive instance ordForceBreak :: Ord ForceBreak
 
-data FormatDoc a = FormatEmpty | FormatDoc ForceBreak Int Boolean (Doc a) ForceBreak
+newtype LeadingComment a = LeadingComment
+  { doc :: Doc a
+  , left :: ForceBreak
+  , lines :: Int
+  , multiline :: Boolean
+  , right :: ForceBreak
+  }
+
+newtype TrailingComment a = TrailingComment
+  { doc :: Doc a
+  , left :: ForceBreak
+  , multiline :: Boolean
+  , right :: ForceBreak
+  }
+
+instance Semigroup (LeadingComment a) where
+  append (LeadingComment c1) (LeadingComment c2)
+    | Dodo.isEmpty c1.doc =
+        LeadingComment c2
+          { left = max c1.left c2.left
+          , lines = c1.lines + c2.lines
+          }
+    | Dodo.isEmpty c2.doc =
+        LeadingComment c1
+          { doc = c1.doc <> breaks ForceNone c2.lines
+          , multiline = c1.multiline || c2.lines > 0
+          , right =
+              if c2.lines > 0 then
+                ForceNone
+              else
+                max c1.right c2.right
+          }
+    | otherwise = do
+        let br = max c1.right c2.left
+        if c2.lines > 0 || br == ForceBreak then
+          LeadingComment c1
+            { doc = c1.doc <> breaks ForceBreak c2.lines <> c2.doc
+            , multiline = true
+            , right = c2.right
+            }
+        else
+          LeadingComment c1
+            { doc = c1.doc <> breakDoc br c2.doc
+            , multiline = c1.multiline || c2.multiline
+            , right = c2.right
+            }
+
+instance Monoid (LeadingComment a) where
+  mempty = LeadingComment
+    { doc: mempty
+    , left: ForceNone
+    , lines: 0
+    , multiline: false
+    , right: ForceNone
+    }
+
+instance Semigroup (TrailingComment a) where
+  append (TrailingComment c1) (TrailingComment c2)
+    | Dodo.isEmpty c1.doc =
+        TrailingComment c2 { left = max c1.left c2.left }
+    | Dodo.isEmpty c2.doc =
+        TrailingComment c1 { right = max c1.right c2.right }
+    | otherwise =
+        TrailingComment c1
+          { doc = c1.doc <> breakDoc (max c1.right c2.left) c2.doc
+          , multiline = c1.multiline || c2.multiline
+          , right = c2.right
+          }
+
+instance Monoid (TrailingComment a) where
+  mempty = TrailingComment
+    { doc: mempty
+    , left: ForceNone
+    , multiline: false
+    , right: ForceNone
+    }
+
+newtype FormatDoc a = FormatDoc
+  { doc :: Doc a
+  , isEmpty :: Boolean
+  , leading :: LeadingComment a
+  , multiline :: Boolean
+  , trailing :: TrailingComment a
+  }
 
 type FormatDocOperator a = FormatDoc a -> FormatDoc a -> FormatDoc a
 
@@ -63,23 +147,82 @@ instance semigroupFormatDoc :: Semigroup (FormatDoc a) where
   append = joinDoc (force identity)
 
 instance monoidFormatDoc :: Monoid (FormatDoc a) where
-  mempty = FormatEmpty
+  mempty = FormatDoc
+    { doc: mempty
+    , leading: mempty
+    , isEmpty: true
+    , multiline: false
+    , trailing: mempty
+    }
 
 fromDoc :: forall a. Doc a -> FormatDoc a
-fromDoc doc = FormatDoc ForceNone 0 false doc ForceNone
+fromDoc doc
+  | Dodo.isEmpty doc = mempty
+  | otherwise = FormatDoc
+      { doc
+      , leading: mempty
+      , isEmpty: false
+      , multiline: false
+      , trailing: mempty
+      }
 
 text :: forall a. String -> FormatDoc a
-text str = FormatDoc ForceNone 0 false (Dodo.text str) ForceNone
+text = fromDoc <<< Dodo.text
 
-leadingLineComment :: forall a. String -> FormatDoc a
-leadingLineComment str = FormatDoc ForceBreak 0 false (Dodo.text str) ForceBreak
-
-trailingLineComment :: forall a. String -> FormatDoc a
-trailingLineComment str = FormatDoc ForceSpace 0 false (Dodo.text str) ForceBreak
-
-blockComment :: forall a. String -> FormatDoc a
-blockComment = splitLines >>> Array.uncons >>> foldMap \{ head, tail } -> do
-  let
+leadingLineComment :: forall a. String -> FormatDoc a -> FormatDoc a
+leadingLineComment str (FormatDoc doc) = FormatDoc doc
+  { leading = comm <> doc.leading
+  , isEmpty = false
+  }
+  where
+  comm = LeadingComment
+    { doc: Dodo.text str
+    , left: ForceBreak
+    , lines: 0
+    , multiline: false
+    , right: ForceBreak
+    }
+
+trailingLineComment :: forall a. String -> FormatDoc a -> FormatDoc a
+trailingLineComment str (FormatDoc doc) = FormatDoc doc
+  { trailing = comm <> doc.trailing
+  , isEmpty = false
+  }
+  where
+  comm = TrailingComment
+    { doc: Dodo.text str
+    , left: ForceSpace
+    , multiline: false
+    , right: ForceBreak
+    }
+
+formatBlockComment :: forall a. String -> Tuple Boolean (Doc a)
+formatBlockComment = splitLines >>> Array.uncons >>> case _ of
+  Nothing ->
+    Tuple false mempty
+  Just { head, tail } ->
+    case prefixSpaces of
+      Nothing ->
+        Tuple false (Dodo.text head)
+      Just ind ->
+        Tuple true $ Dodo.withPosition \pos -> do
+          let newIndent = if ind < pos.indent then 0 else ind
+          let spaces = power " " newIndent
+          let tailDocs = map (\str -> Dodo.text $ fromMaybe str $ String.stripPrefix (String.Pattern spaces) str) tail
+          Dodo.lines
+            [ Dodo.text head
+            , Dodo.locally
+                ( \prev ->
+                    if newIndent < prev.indent then
+                      prev
+                        { indentSpaces = spaces
+                        , indent = newIndent
+                        }
+                    else prev
+                )
+                (intercalate Dodo.break tailDocs)
+            ]
+    where
     prefixSpaces =
       tail
         # Array.mapMaybe
@@ -89,87 +232,93 @@ blockComment = splitLines >>> Array.uncons >>> foldMap \{ head, tail } -> do
             )
         # Array.sort
         # Array.head
-  case prefixSpaces of
-    Nothing ->
-      FormatDoc ForceSpace 0 false (Dodo.text head) ForceSpace
-    Just ind ->
-      FormatDoc ForceSpace 0 true commentDoc ForceSpace
-      where
-      commentDoc = Dodo.withPosition \pos -> do
-        let newIndent = if ind < pos.indent then 0 else ind
-        let spaces = power " " newIndent
-        let tailDocs = map (\str -> Dodo.text $ fromMaybe str $ String.stripPrefix (String.Pattern spaces) str) tail
-        Dodo.lines
-          [ Dodo.text head
-          , Dodo.locally
-              ( \prev ->
-                  if newIndent < prev.indent then
-                    prev
-                      { indentSpaces = spaces
-                      , indent = newIndent
-                      }
-                  else prev
-              )
-              (intercalate Dodo.break tailDocs)
-          ]
+
+leadingBlockComment :: forall a. String -> FormatDoc a -> FormatDoc a
+leadingBlockComment str (FormatDoc doc) = FormatDoc doc
+  { leading = comm <> doc.leading
+  , isEmpty = false
+  }
+  where
+  Tuple multi commDoc =
+    formatBlockComment str
+
+  comm = LeadingComment
+    { doc: commDoc
+    , left: ForceSpace
+    , lines: 0
+    , multiline: multi
+    , right: ForceSpace
+    }
+
+trailingBlockComment :: forall a. String -> FormatDoc a -> FormatDoc a
+trailingBlockComment str (FormatDoc doc) = FormatDoc doc
+  { trailing = comm <> doc.trailing
+  , isEmpty = false
+  }
+  where
+  Tuple multi commDoc =
+    formatBlockComment str
+
+  comm = TrailingComment
+    { doc: commDoc
+    , left: ForceSpace
+    , multiline: multi
+    , right: ForceSpace
+    }
 
 anchor :: forall a. FormatDoc a -> FormatDoc a
-anchor = case _ of
-  FormatEmpty ->
-    FormatEmpty
-  FormatDoc fl n m doc fr ->
-    FormatDoc fl 0 (if n > 0 then true else m) doc fr
+anchor (FormatDoc doc) = case doc.leading of
+  LeadingComment comm | comm.lines > 0 ->
+    FormatDoc doc
+      { leading = LeadingComment comm { lines = 0 }
+      , multiline = true
+      }
+  _ ->
+    FormatDoc doc
 
 flatten :: forall a. FormatDoc a -> FormatDoc a
 flatten = flattenMax 0
 
 flattenMax :: forall a. Int -> FormatDoc a -> FormatDoc a
-flattenMax n' = case _ of
-  FormatEmpty ->
-    FormatEmpty
-  FormatDoc fl n m doc fr ->
-    FormatDoc fl (min n' n) m doc fr
-
-forceBreak :: forall a. FormatDoc a -> FormatDoc a
-forceBreak = case _ of
-  FormatEmpty ->
-    FormatEmpty
-  FormatDoc _ n m doc fr ->
-    FormatDoc ForceBreak n m doc fr
+flattenMax n (FormatDoc doc) = case doc.leading of
+  LeadingComment comm ->
+    FormatDoc doc
+      { leading = LeadingComment comm { lines = min comm.lines n }
+      }
 
 flexGroup :: forall a. FormatDoc a -> FormatDoc a
-flexGroup = case _ of
-  FormatEmpty ->
-    FormatEmpty
-  a@(FormatDoc fl n m doc fr)
-    | not m -> FormatDoc fl n false (Dodo.flexGroup doc) fr
-    | otherwise -> a
+flexGroup (FormatDoc doc)
+  | doc.multiline = FormatDoc doc
+  | otherwise = FormatDoc doc { doc = Dodo.flexGroup doc.doc }
 
 indent :: forall a. FormatDoc a -> FormatDoc a
-indent = mapDoc Dodo.indent
+indent = mapDocs Dodo.indent
 
 align :: forall a. Int -> FormatDoc a -> FormatDoc a
-align = mapDoc <<< Dodo.align
+align = mapDocs <<< Dodo.align
 
 alignCurrentColumn :: forall a. FormatDoc a -> FormatDoc a
-alignCurrentColumn = mapDoc Dodo.alignCurrentColumn
+alignCurrentColumn = mapDocs Dodo.alignCurrentColumn
 
 locally :: forall a. (LocalOptions -> LocalOptions) -> FormatDoc a -> FormatDoc a
-locally = mapDoc <<< Dodo.locally
+locally k (FormatDoc doc) = FormatDoc doc { doc = Dodo.locally k doc.doc }
 
 sourceBreak :: forall a. Int -> FormatDoc a -> FormatDoc a
-sourceBreak m = case _ of
-  FormatEmpty
-    | m <= 0 -> FormatEmpty
-    | otherwise -> FormatDoc ForceNone m false mempty ForceNone
-  FormatDoc fl n multi doc fr ->
-    FormatDoc fl (m + n) multi doc fr
+sourceBreak n (FormatDoc doc) = do
+  let LeadingComment comm = doc.leading
+  FormatDoc doc
+    { isEmpty = false
+    , leading = LeadingComment comm { lines = comm.lines + n }
+    }
 
 forceMinSourceBreaks :: forall a. Int -> FormatDoc a -> FormatDoc a
-forceMinSourceBreaks m = case _ of
-  FormatEmpty -> FormatEmpty
-  FormatDoc fl n multi doc fr ->
-    FormatDoc fl (max m n) multi doc fr
+forceMinSourceBreaks n (FormatDoc doc)
+  | doc.isEmpty = FormatDoc doc
+  | otherwise = do
+      let LeadingComment comm = doc.leading
+      FormatDoc doc
+        { leading = LeadingComment comm { lines = max comm.lines n }
+        }
 
 space :: forall a. FormatDocOperator a
 space = joinDoc (force (append Dodo.space))
@@ -242,32 +391,69 @@ softSpace = joinDoc \f m doc -> case f of
 -- | right associativity. You will always get double breaks when used
 -- | with left associativity.
 flexDoubleBreak :: forall a. FormatDocOperator a
-flexDoubleBreak = case _, _ of
-  FormatEmpty, b -> b
-  a, FormatEmpty -> a
-  FormatDoc fl1 n1 m1 doc1 _, FormatDoc _ n2 _ doc2 fr2
-    | n2 >= 2 || m1 ->
-        FormatDoc fl1 n1 true (doc1 <> Dodo.break <> Dodo.break <> doc2) fr2
-    | otherwise ->
-        FormatDoc fl1 n1 true (Dodo.flexSelect doc1 mempty Dodo.break <> Dodo.break <> doc2) fr2
+flexDoubleBreak (FormatDoc doc1) (FormatDoc doc2)
+  | doc1.isEmpty = FormatDoc doc2
+  | doc2.isEmpty = FormatDoc doc1
+  | otherwise = do
+      let TrailingComment comm1 = doc1.trailing
+      let LeadingComment comm2 = doc2.leading
+      let docLeft = doc1.doc <> breakDoc comm1.left comm1.doc
+      let docRight = comm2.doc <> breakDoc comm2.right doc2.doc
+      if comm2.lines >= 2 || doc1.multiline then
+        FormatDoc doc1
+          { doc = docLeft <> Dodo.break <> Dodo.break <> docRight
+          , multiline = true
+          , trailing = doc2.trailing
+          }
+      else
+        FormatDoc doc1
+          { doc = Dodo.flexSelect docLeft mempty Dodo.break <> Dodo.break <> docRight
+          , multiline = true
+          , trailing = doc2.trailing
+          }
 
 isEmpty :: forall a. FormatDoc a -> Boolean
-isEmpty = case _ of
-  FormatEmpty -> true
-  _ -> false
+isEmpty (FormatDoc doc) = doc.isEmpty
+
+breakDoc :: forall a. ForceBreak -> Doc a -> Doc a
+breakDoc br doc
+  | Dodo.isEmpty doc = doc
+  | otherwise = case br of
+      ForceBreak -> Dodo.break <> doc
+      ForceSpace -> Dodo.space <> doc
+      ForceNone -> doc
+
+breaks :: forall a. ForceBreak -> Int -> Doc a
+breaks fl n
+  | n >= 2 = Dodo.break <> Dodo.break
+  | n == 1 = Dodo.break
+  | otherwise = case fl of
+      ForceBreak -> Dodo.break
+      ForceSpace -> Dodo.space
+      ForceNone -> mempty
 
 joinDoc :: forall a. (ForceBreak -> Boolean -> Doc a -> Tuple Boolean (Doc a)) -> FormatDocOperator a
-joinDoc spc = case _, _ of
-  FormatEmpty, b -> b
-  a, FormatEmpty -> a
-  FormatDoc fl1 n1 m1 doc1 fr1, FormatDoc fl2 n2 m2 doc2 fr2
-    | n2 == 1 ->
-        FormatDoc fl1 n1 true (doc1 <> Dodo.break <> doc2) fr2
-    | n2 >= 2 ->
-        FormatDoc fl1 n1 true (doc1 <> Dodo.break <> Dodo.break <> doc2) fr2
-    | otherwise -> do
-        let (Tuple m3 doc3) = spc (max fr1 fl2) m2 doc2
-        FormatDoc fl1 n1 (m1 || m3) (doc1 <> doc3) fr2
+joinDoc spaceFn (FormatDoc doc1) (FormatDoc doc2)
+  | doc1.isEmpty = FormatDoc doc2
+  | doc2.isEmpty = FormatDoc doc1
+  | otherwise = do
+      let TrailingComment comm1 = doc1.trailing
+      let LeadingComment comm2 = doc2.leading
+      let docLeft = doc1.doc <> breakDoc comm1.left comm1.doc
+      let docRight = comm2.doc <> breakDoc comm2.right doc2.doc
+      if comm2.lines > 0 then
+        FormatDoc doc1
+          { doc = docLeft <> breaks ForceBreak comm2.lines <> docRight
+          , multiline = true
+          , trailing = doc2.trailing
+          }
+      else do
+        let Tuple m3 doc3 = spaceFn (max comm1.right comm2.left) (comm2.multiline || doc2.multiline) docRight
+        FormatDoc doc1
+          { doc = docLeft <> doc3
+          , multiline = comm1.multiline || doc1.multiline || m3
+          , trailing = doc2.trailing
+          }
 
 force :: forall a. (Doc a -> Doc a) -> (ForceBreak -> Boolean -> Doc a -> Tuple Boolean (Doc a))
 force k f m doc = case f of
@@ -276,17 +462,32 @@ force k f m doc = case f of
   _ ->
     Tuple m (k doc)
 
-mapDoc :: forall a b. (Doc a -> Doc b) -> FormatDoc a -> FormatDoc b
-mapDoc k = case _ of
-  FormatEmpty ->
-    FormatEmpty
-  FormatDoc fl n m doc fr ->
-    FormatDoc fl n m (k doc) fr
+mapDoc :: forall a. (Doc a -> Doc a) -> FormatDoc a -> FormatDoc a
+mapDoc k (FormatDoc doc)
+  | doc.isEmpty = FormatDoc doc
+  | otherwise = FormatDoc doc { doc = k doc.doc }
+
+mapDocs :: forall a. (Doc a -> Doc a) -> FormatDoc a -> FormatDoc a
+mapDocs k (FormatDoc doc)
+  | doc.isEmpty = FormatDoc doc
+  | otherwise = do
+      let LeadingComment comm1 = doc.leading
+      let TrailingComment comm2 = doc.trailing
+      FormatDoc doc
+        { doc = k doc.doc
+        , leading = LeadingComment comm1 { doc = k comm1.doc }
+        , trailing = TrailingComment comm2 { doc = k comm2.doc }
+        }
 
 toDoc :: forall a. FormatDoc a -> Doc a
-toDoc = case _ of
-  FormatEmpty -> mempty
-  FormatDoc _ _ _ doc _ -> doc
+toDoc (FormatDoc doc)
+  | doc.isEmpty = mempty
+  | otherwise = do
+      let LeadingComment comm1 = doc.leading
+      let TrailingComment comm2 = doc.trailing
+      comm1.doc
+        <> breakDoc comm1.right doc.doc
+        <> breakDoc comm2.left comm2.doc
 
 joinWithMap
   :: forall f a b
diff --git a/src/Tidy/Hang.purs b/src/Tidy/Hang.purs
index 855279d..82837f7 100644
--- a/src/Tidy/Hang.purs
+++ b/src/Tidy/Hang.purs
@@ -21,8 +21,7 @@ import Data.Maybe (maybe)
 import Data.Tuple (Tuple(..), fst, snd)
 import Dodo (Doc)
 import Dodo as Dodo
-import Partial.Unsafe (unsafeCrashWith)
-import Tidy.Doc (ForceBreak(..), FormatDoc(..), align, break, flatten, flexGroup, forceBreak, indent)
+import Tidy.Doc (ForceBreak(..), FormatDoc(..), LeadingComment(..), TrailingComment(..), align, break, breakDoc, flatten, flexGroup, forceMinSourceBreaks, indent)
 
 data HangingDoc a
   = HangBreak (FormatDoc a)
@@ -188,7 +187,7 @@ toFormatDoc = fst <<< goInit
         { init, last } = NonEmptyArray.unsnoc tail
         next = Array.foldr goInitApp (goLastApp last) init
         this = fst $ goInit $ case head of
-          HangApp _ _ _ -> overHangHead forceBreak head
+          HangApp _ _ _ -> overHangHead (forceMinSourceBreaks 1) head
           _ -> head
         docIndent = indMulti head ind (fst next)
         docGroup = flexSelectJoin this (fst next) docIndent
@@ -204,43 +203,123 @@ toFormatDoc = fst <<< goInit
         docBreak = flexSelectJoin (prevInd this) (prevAlgn docIndent) (prevInd docIndent)
       Tuple docGroup docBreak
 
-  flexSelect = case _, _, _ of
-    FormatDoc fl1 n1 m1 doc1 fr1, FormatDoc fl2 n2 m2 doc2 fr2, FormatDoc fl3 n3 m3 doc3 fr3 -> do
-      let
-        doc2' = breaks (max fr1 fl2) n2 <> doc2
-        doc3' = breaks (max fr1 fl3) n3 <> doc3
-      FormatDoc fl1 n1 (m1 || (m2 && m3))
-        (Dodo.flexSelect doc1 doc2' doc3')
-        (max fr2 fr3)
-    _, _, _ ->
-      unsafeCrashWith "flexSelect/FormatEmpty"
-
-  flexSelectJoin = case _, _, _ of
-    FormatDoc fl1 n1 m1 doc1 fr1, FormatDoc fl2 n2 m2 doc2 fr2, FormatDoc fl3 n3 m3 doc3 fr3 -> do
-      let
-        doc2' = breaks (max fr1 fl2) n2 <> doc2
-        doc3' = breaks (max fr1 fl3) n3 <> doc3
-        br = if fl1 == ForceBreak || n1 > 0 then Dodo.break else Dodo.spaceBreak
-      FormatDoc ForceNone 0 (m1 || (m2 && m3))
-        (Dodo.flexSelect (br <> doc1) doc2' doc3')
-        (max fr2 fr3)
-    _, _, _ ->
-      unsafeCrashWith "flexSelect/FormatEmpty"
-
-  docJoin = case _ of
-    FormatEmpty -> FormatEmpty
-    fdoc@(FormatDoc fl n m doc fr) ->
-      if fl == ForceBreak || n > 0 then fdoc
-      else if m then FormatDoc ForceBreak n m doc fr
-      else FormatDoc ForceNone 0 false (Dodo.spaceBreak <> doc) fr
+  flexSelect (FormatDoc doc1) (FormatDoc doc2) (FormatDoc doc3) = do
+    let
+      TrailingComment comm1r = doc1.trailing
+
+      doc1' =
+        doc1.doc
+          <> breakDoc comm1r.left comm1r.doc
+
+      LeadingComment comm2l = doc2.leading
+      TrailingComment comm2r = doc2.trailing
+
+      doc2' =
+        breaks (max comm1r.right comm2l.left) comm2l.lines
+          <> comm2l.doc
+          <> breakDoc comm2l.right doc2.doc
+          <> breakDoc comm2r.left comm2r.doc
+
+      LeadingComment comm3l = doc3.leading
+      TrailingComment comm3r = doc3.trailing
+
+      doc3' =
+        breaks (max comm1r.right comm3l.left) comm3l.lines
+          <> comm3l.doc
+          <> breakDoc comm2l.right doc3.doc
+          <> breakDoc comm3r.left comm3r.doc
+
+      m1 = doc1.multiline || comm1r.multiline
+      m2 = comm2l.multiline || doc2.multiline || comm2r.multiline
+      m3 = comm3l.multiline || doc3.multiline || comm3r.multiline
+
+    FormatDoc
+      { doc: Dodo.flexSelect doc1' doc2' doc3'
+      , leading: doc1.leading
+      , isEmpty: false
+      , multiline: m1 || (m2 && m3)
+      , trailing:
+          TrailingComment
+            { doc: mempty
+            , left: ForceNone
+            , multiline: false
+            , right: max comm2r.right comm3r.right
+            }
+      }
+
+  flexSelectJoin (FormatDoc doc1) (FormatDoc doc2) (FormatDoc doc3) = do
+    let
+      LeadingComment comm1l = doc1.leading
+      TrailingComment comm1r = doc1.trailing
+
+      break
+        | comm1l.left == ForceBreak || comm1l.lines > 0 = Dodo.break
+        | otherwise = Dodo.spaceBreak
+
+      doc1' =
+        break
+          <> comm1l.doc
+          <> breakDoc comm1l.right doc1.doc
+          <> breakDoc comm1r.left comm1r.doc
+
+      LeadingComment comm2l = doc2.leading
+      TrailingComment comm2r = doc2.trailing
+
+      doc2' =
+        breaks (max comm1r.right comm2l.left) comm2l.lines
+          <> comm2l.doc
+          <> breakDoc comm2l.right doc2.doc
+          <> breakDoc comm2r.left comm2r.doc
+
+      LeadingComment comm3l = doc3.leading
+      TrailingComment comm3r = doc3.trailing
+
+      doc3' =
+        breaks (max comm1r.right comm3l.left) comm3l.lines
+          <> comm3l.doc
+          <> breakDoc comm2l.right doc3.doc
+          <> breakDoc comm3r.left comm3r.doc
+
+      m1 = comm1l.multiline || doc1.multiline || comm1r.multiline
+      m2 = comm2l.multiline || doc2.multiline || comm2r.multiline
+      m3 = comm3l.multiline || doc3.multiline || comm3r.multiline
+
+    FormatDoc
+      { doc: Dodo.flexSelect doc1' doc2' doc3'
+      , leading: mempty
+      , isEmpty: false
+      , multiline: m1 || (m2 && m3)
+      , trailing:
+          TrailingComment
+            { doc: mempty
+            , left: ForceNone
+            , multiline: false
+            , right: max comm2r.right comm3r.right
+            }
+      }
+
+  docJoin fdoc@(FormatDoc doc)
+    | doc.isEmpty = fdoc
+    | otherwise = do
+        let LeadingComment comm = doc.leading
+        if comm.left == ForceBreak || comm.lines > 0 then
+          fdoc
+        else if comm.multiline || doc.multiline then
+          forceMinSourceBreaks 1 fdoc
+        else
+          FormatDoc doc
+            { doc = Dodo.spaceBreak <> comm.doc <> breakDoc comm.right doc.doc
+            , leading = mempty
+            }
 
   realignOp op doc = case op, hangHead doc of
-    FormatDoc fl1 n1 _ _ fr1, FormatDoc fl2 n2 m2 _ _
-      | fl1 /= ForceBreak && n1 == 0 && fr1 /= ForceBreak && fl2 /= ForceBreak && n2 > 0 ->
-          realignOp (forceBreak op) (overHangHead flatten doc)
+    FormatDoc doc1@{ leading: LeadingComment comm1, trailing: TrailingComment comm2 },
+    FormatDoc { leading: LeadingComment comm3 }
+      | comm1.left /= ForceBreak && comm1.lines == 0 && comm2.right /= ForceBreak && comm3.left /= ForceBreak && comm3.lines > 0 ->
+          realignOp (forceMinSourceBreaks 1 op) (overHangHead flatten doc)
       | HangBreak _ <- doc
-      , fr1 /= ForceBreak && fl2 /= ForceBreak && n2 == 0 && m2 ->
-          Tuple op (overHangHead forceBreak doc)
+      , comm2.right /= ForceBreak && comm3.left /= ForceBreak && comm3.lines == 0 && (comm3.multiline || doc1.multiline) ->
+          Tuple op (overHangHead (forceMinSourceBreaks 1) doc)
     _, _ ->
       Tuple op doc
 
diff --git a/test/snapshots/DeclarationBreaks.output b/test/snapshots/DeclarationBreaks.output
index 828741c..efa19d7 100644
--- a/test/snapshots/DeclarationBreaks.output
+++ b/test/snapshots/DeclarationBreaks.output
@@ -8,8 +8,10 @@ type Ok = String
 type Ok2 :: Type
 type Ok2 = String
 
+-- Comment
 type Ok3 = String
 type Ok4 = Int
+type OK5 = Int
 
 class Foo a where
   foo :: a
@@ -62,8 +64,10 @@ type Ok = String
 type Ok2 :: Type
 type Ok2 = String
 
+-- Comment
 type Ok3 = String
 type Ok4 = Int
+type OK5 = Int
 
 class Foo a where
   foo :: a
diff --git a/test/snapshots/DeclarationBreaks.purs b/test/snapshots/DeclarationBreaks.purs
index f4b1b3a..1304193 100644
--- a/test/snapshots/DeclarationBreaks.purs
+++ b/test/snapshots/DeclarationBreaks.purs
@@ -5,8 +5,10 @@ import Bar
 type Ok = String
 type Ok2 :: Type
 type Ok2 = String
+-- Comment
 type Ok3 = String
 type Ok4 = Int
+type OK5 = Int
 class Foo a where foo :: a
 instance Foo Int where foo = 12
 else instance Foo String where foo = "foo"
diff --git a/test/snapshots/MultiCase.output b/test/snapshots/MultiCase.output
index f4508b8..98fe728 100644
--- a/test/snapshots/MultiCase.output
+++ b/test/snapshots/MultiCase.output
@@ -11,3 +11,8 @@ test = case _, _ of
     , c
     } ->
     a <> b <> c
+
+test = case _, _ of
+  -- Comment
+  Foo a, Bar b ->
+    a <> b
diff --git a/test/snapshots/MultiCase.purs b/test/snapshots/MultiCase.purs
index 9482b3f..0ad0c85 100644
--- a/test/snapshots/MultiCase.purs
+++ b/test/snapshots/MultiCase.purs
@@ -9,3 +9,8 @@ test = case _, _ of
              , c
              } ->
     a <> b <> c
+
+test = case _, _ of
+  -- Comment
+  Foo a, Bar b ->
+    a <> b
diff --git a/test/snapshots/MultilineStringLiterals.output b/test/snapshots/MultilineStringLiterals.output
index f7c913f..96ae694 100644
--- a/test/snapshots/MultilineStringLiterals.output
+++ b/test/snapshots/MultilineStringLiterals.output
@@ -1,7 +1,20 @@
 module MultilineStringLiterals where
 
+test = """ """
+
+test =
+  """
+this is a string
+"""
+
 test =
   """
+this is a string
+
+this is a string
+
+
+
 this is a string
 """
 
diff --git a/test/snapshots/MultilineStringLiterals.purs b/test/snapshots/MultilineStringLiterals.purs
index ed89c7b..0e68d9c 100644
--- a/test/snapshots/MultilineStringLiterals.purs
+++ b/test/snapshots/MultilineStringLiterals.purs
@@ -1,6 +1,18 @@
 module MultilineStringLiterals where
 
+test = """ """
+
+test = """
+this is a string
+"""
+
 test = """
+this is a string
+
+this is a string
+
+
+
 this is a string
 """
 
diff --git a/test/snapshots/TrailingLineComments.output b/test/snapshots/TrailingLineComments.output
new file mode 100644
index 0000000..b61009e
--- /dev/null
+++ b/test/snapshots/TrailingLineComments.output
@@ -0,0 +1,20 @@
+module TrailingLineComments where
+
+test -- ok
+  = test
+
+test = test -- ok
+
+test = test {- a -} {- b -} -- ok
+
+-- @format --width=10
+module TrailingLineComments where
+
+test -- ok
+  = test
+
+test =
+  test -- ok
+
+test =
+  test {- a -} {- b -} -- ok
diff --git a/test/snapshots/TrailingLineComments.purs b/test/snapshots/TrailingLineComments.purs
new file mode 100644
index 0000000..5255db3
--- /dev/null
+++ b/test/snapshots/TrailingLineComments.purs
@@ -0,0 +1,9 @@
+-- @format --width=10
+module TrailingLineComments where
+
+test -- ok
+  = test
+
+test = test -- ok
+
+test = test {- a -} {- b -} -- ok