fix: hard line break issue in Format

Author: jrr6

src/Init/Data/Format/Basic.lean

@@ -142,17 +142,37 @@ private structure WorkItem where
   indent : Int
   activeTags : Nat
 
+/--
+A directive for how a given work group should be flattened.
+
+- `allow` indicates that the group is allowed to be flattened; its argument is `true` if
+  there is sufficient space for it to be flattened (and so it should be), or `false` if not.
+- `disallow` means that this group should not be flattened irrespective of space concerns.
+  This is used at levels of a `Format` outside of any flattening groups. It is necessary to track
+  this so that mechanisms like post-hard-line-break flattening can know whether to attempt to start
+  flattening.
+-/
+inductive FlattenDirective where
+  | allow (fits : Bool)
+  | disallow
+  deriving BEq
+
+/-- Whether the given directive indicates that flattening should occur. -/
+def FlattenDirective.shouldFlatten : FlattenDirective → Bool
+  | allow true => true
+  | _ => false
+
 private structure WorkGroup where
-  flatten : Bool
-  flb     : FlattenBehavior
-  items   : List WorkItem
+  fld   : FlattenDirective
+  flb   : FlattenBehavior
+  items : List WorkItem
 
 private partial def spaceUptoLine' : List WorkGroup → Nat → Nat → SpaceResult
   |   [],                         _,   _ => {}
   |   { items := [],    .. }::gs, col, w => spaceUptoLine' gs col w
   | g@{ items := i::is, .. }::gs, col, w =>
     merge w
-      (spaceUptoLine i.f g.flatten (w + col - i.indent) w)
+      (spaceUptoLine i.f g.fld.shouldFlatten (w + col - i.indent) w)
       (spaceUptoLine' ({ g with items := is }::gs) col)
 
 /-- A monad in which we can pretty-print `Format` objects. -/
@@ -169,11 +189,11 @@ open MonadPrettyFormat
 private def pushGroup (flb : FlattenBehavior) (items : List WorkItem) (gs : List WorkGroup) (w : Nat) [Monad m] [MonadPrettyFormat m] : m (List WorkGroup) := do
   let k  ← currColumn
   -- Flatten group if it + the remainder (gs) fits in the remaining space. For `fill`, measure only up to the next (ungrouped) line break.
-  let g  := { flatten := flb == FlattenBehavior.allOrNone, flb := flb, items := items : WorkGroup }
+  let g  := { fld := .allow (flb == FlattenBehavior.allOrNone), flb := flb, items := items : WorkGroup }
   let r  := spaceUptoLine' [g] k (w-k)
   let r' := merge (w-k) r (spaceUptoLine' gs k)
   -- Prevent flattening if any item contains a hard line break, except within `fill` if it is ungrouped (=> unflattened)
-  return { g with flatten := !r.foundFlattenedHardLine && r'.space <= w-k }::gs
+  return { g with fld := .allow (!r.foundFlattenedHardLine && r'.space <= w-k) }::gs
 
 private partial def be (w : Nat) [Monad m] [MonadPrettyFormat m] : List WorkGroup → m Unit
   | []                           => pure ()
@@ -200,11 +220,15 @@ private partial def be (w : Nat) [Monad m] [MonadPrettyFormat m] : List WorkGrou
         pushNewline i.indent.toNat
         let is := { i with f := text (s.extract (s.next p) s.endPos) }::is
         -- after a hard line break, re-evaluate whether to flatten the remaining group
-        pushGroup g.flb is gs w >>= be w
+        -- note that we shouldn't start flattening after a hard break outside a group
+        if g.fld == .disallow then
+          be w (gs' is)
+        else
+          pushGroup g.flb is gs w >>= be w
     | line =>
       match g.flb with
       | FlattenBehavior.allOrNone =>
-        if g.flatten then
+        if g.fld.shouldFlatten then
           -- flatten line = text " "
           pushOutput " "
           endTags i.activeTags
@@ -220,10 +244,10 @@ private partial def be (w : Nat) [Monad m] [MonadPrettyFormat m] : List WorkGrou
           endTags i.activeTags
           pushGroup FlattenBehavior.fill is gs w >>= be w
         -- if preceding fill item fit in a single line, try to fit next one too
-        if g.flatten then
+        if g.fld.shouldFlatten then
           let gs'@(g'::_) ← pushGroup FlattenBehavior.fill is gs (w - " ".length)
             | panic "unreachable"
-          if g'.flatten then
+          if g'.fld.shouldFlatten then
             pushOutput " "
             endTags i.activeTags
             be w gs'  -- TODO: use `return`
@@ -232,7 +256,7 @@ private partial def be (w : Nat) [Monad m] [MonadPrettyFormat m] : List WorkGrou
         else
           breakHere
     | align force =>
-      if g.flatten && !force then
+      if g.fld.shouldFlatten && !force then
         -- flatten (align false) = nil
         endTags i.activeTags
         be w (gs' is)
@@ -247,7 +271,7 @@ private partial def be (w : Nat) [Monad m] [MonadPrettyFormat m] : List WorkGrou
           endTags i.activeTags
           be w (gs' is)
     | group f flb =>
-      if g.flatten then
+      if g.fld.shouldFlatten then
         -- flatten (group f) = flatten f
         be w (gs' ({ i with f }::is))
       else
@@ -256,7 +280,7 @@ private partial def be (w : Nat) [Monad m] [MonadPrettyFormat m] : List WorkGrou
 /-- Render the given `f : Format` with a line width of `w`.
 `indent` is the starting amount to indent each line by. -/
 def prettyM (f : Format) (w : Nat) (indent : Nat := 0) [Monad m] [MonadPrettyFormat m] : m Unit :=
-  be w [{ flb := FlattenBehavior.allOrNone, flatten := false, items := [{ f := f, indent, activeTags := 0 }]}]
+  be w [{ flb := FlattenBehavior.allOrNone, fld := .disallow, items := [{ f := f, indent, activeTags := 0 }]}]
 
 /-- Create a format `l ++ f ++ r` with a flatten group.
 FlattenBehaviour is `allOrNone`; for `fill` use `bracketFill`. -/

tests/lean/1079.lean.expected.out

@@ -11,4 +11,7 @@
       ?a = ?a
     with
       Ordering.eq = Ordering.lt
-[Meta.Tactic.simp.rewrite] imp_self:10000: False → False ==> True
+[Meta.Tactic.simp.rewrite] imp_self:10000:
+      False → False
+    ==>
+      True

tests/lean/973b.lean.expected.out

@@ -1,5 +1,6 @@
 973b.lean:5:8-5:10: warning: declaration uses 'sorry'
 [Meta.Tactic.simp.discharge] ex discharge ❌️
       ?p x
-[Meta.Tactic.simp.discharge] ex discharge ❌️ ?p (f x)
+[Meta.Tactic.simp.discharge] ex discharge ❌️
+      ?p (f x)
 973b.lean:9:8-9:11: warning: declaration uses 'sorry'

tests/lean/discrTreeIota.lean.expected.out

@@ -2,4 +2,7 @@
       default
     ==>
       PUnit.unit
-[Meta.Tactic.simp.rewrite] eq_self:1000: PUnit.unit = x ==> True
+[Meta.Tactic.simp.rewrite] eq_self:1000:
+      PUnit.unit = x
+    ==>
+      True

tests/lean/exceptionalTrace.lean.expected.out

@@ -5,4 +5,5 @@ use `set_option diagnostics true` to get diagnostic information
     rt + 1
   [Elab.step] expected type: <not-available>, term
       binop% HAdd.hAdd✝ rt 1
-    [Elab.step] expected type: <not-available>, term rt
+    [Elab.step] expected type: <not-available>, term
+        rt

tests/lean/run/1234.lean

@@ -28,7 +28,10 @@ trace: [Meta.Tactic.simp.rewrite] h₁:1000:
       if True then ⟨v, ⋯⟩ else ⟨v - 1, ⋯⟩
     ==>
       ⟨v, ⋯⟩
-[Meta.Tactic.simp.rewrite] eq_self:1000: ⟨v, ⋯⟩ = ⟨v, ⋯⟩ ==> True
+[Meta.Tactic.simp.rewrite] eq_self:1000:
+      ⟨v, ⋯⟩ = ⟨v, ⋯⟩
+    ==>
+      True
 -/
 #guard_msgs in
 example (h₁: k ≤ v - 1) (h₂: 0 < v):
@@ -65,7 +68,10 @@ trace: [Meta.Tactic.simp.rewrite] h₁:1000:
       if True then ⟨v, ⋯⟩ else ⟨v - 1, ⋯⟩
     ==>
       ⟨v, ⋯⟩
-[Meta.Tactic.simp.rewrite] eq_self:1000: ⟨v, ⋯⟩ = ⟨v, ⋯⟩ ==> True
+[Meta.Tactic.simp.rewrite] eq_self:1000:
+      ⟨v, ⋯⟩ = ⟨v, ⋯⟩
+    ==>
+      True
 -/
 #guard_msgs in
 example (h₁: k ≤ v - 1) (h₂: 0 < v):
@@ -100,7 +106,10 @@ trace: [Meta.Tactic.simp.rewrite] h₁:1000:
       if True then ⟨v, ⋯⟩ else ⟨v - 1, ⋯⟩
     ==>
       ⟨v, ⋯⟩
-[Meta.Tactic.simp.rewrite] eq_self:1000: ⟨v, ⋯⟩ = ⟨v, ⋯⟩ ==> True
+[Meta.Tactic.simp.rewrite] eq_self:1000:
+      ⟨v, ⋯⟩ = ⟨v, ⋯⟩
+    ==>
+      True
 -/
 #guard_msgs in
 example (h₁: k ≤ v - 1) (h₂: 0 < v):

tests/lean/run/1380.lean

@@ -21,7 +21,10 @@ trace: [Meta.Tactic.simp.unify] eq_self:1000, failed to unify
         v₁ < v₂
       ==>
         True
-[Meta.Tactic.simp.rewrite] Nat.ne_of_gt:1000: v₂ = v₁ ==> False
+[Meta.Tactic.simp.rewrite] Nat.ne_of_gt:1000:
+      v₂ = v₁
+    ==>
+      False
 -/
 #guard_msgs in
 set_option trace.Meta.Tactic.simp true in

tests/lean/run/1815.lean

@@ -13,7 +13,10 @@ trace: [Meta.Tactic.simp.rewrite] mul_comm:1000:perm, perm rejected Left a ==> d
     ==>
       a * default
 [Meta.Tactic.simp.rewrite] mul_comm:1000:perm, perm rejected a * default ==> default * a
-[Meta.Tactic.simp.rewrite] eq_self:1000: Left a = a * default ==> True
+[Meta.Tactic.simp.rewrite] eq_self:1000:
+      Left a = a * default
+    ==>
+      True
 -/
 #guard_msgs in
 example (a : α) : Left a = Right a := by

tests/lean/run/3257.lean

@@ -20,7 +20,10 @@ trace: [Meta.Tactic.simp.discharge] bar discharge ✅️
         T
       ==>
         True
-[Meta.Tactic.simp.rewrite] bar:1000: U ==> True
+[Meta.Tactic.simp.rewrite] bar:1000:
+      U
+    ==>
+      True
 -/
 #guard_msgs in
 example : U := by

tests/lean/run/790.lean

@@ -20,7 +20,10 @@ trace: [Meta.Tactic.simp.rewrite] differential_of_linear:1000:
       differential f x dx
     ==>
       f dx
-[Meta.Tactic.simp.rewrite] eq_self:1000: f dx = f dx ==> True
+[Meta.Tactic.simp.rewrite] eq_self:1000:
+      f dx = f dx
+    ==>
+      True
 -/
 #guard_msgs in
 example {Y : Type} [Vec Y] (f : Nat → Y) (x dx : Nat)