feat: intra-line withPosition formatting

src/Init/Data/Format/Basic.lean

@@ -41,6 +41,19 @@ inductive Format where
   /-- A position where a newline may be inserted
   if the current group does not fit within the allotted column width. -/
   | line                : Format
+  /-- `align` tells the formatter to pad with spaces to the current indent,
+  or else add a newline if we are already past the indent. For example:
+  ```
+  nest 2 <| "." ++ align ++ "a" ++ line ++ "b"
+  ```
+  results in:
+  ```
+  . a
+    b
+  ```
+  If `force` is true, then it will pad to the indent even if it is in a flattened group.
+  -/
+  | align (force : Bool) : Format
   /-- A node containing a plain string. -/
   | text                : String → Format
   /-- `nest n f` tells the formatter that `f` is nested inside something with length `n`
@@ -70,6 +83,7 @@ namespace Format
 def isEmpty : Format → Bool
   | nil          => true
   | line         => false
+  | align _      => true
   | text msg     => msg == ""
   | nest _ f     => f.isEmpty
   | append f₁ f₂ => f₁.isEmpty && f₂.isEmpty
@@ -103,17 +117,23 @@ private structure SpaceResult where
     let r₂ := r₂ (w - r₁.space);
     { r₂ with space := r₁.space + r₂.space }
 
-private def spaceUptoLine : Format → Bool → Nat → SpaceResult
-  | nil,          _,       _ => {}
-  | line,         flatten, _ => if flatten then { space := 1 } else { foundLine := true }
-  | text s,       flatten, _ =>
-    let p := s.posOf '\n';
-    let off := s.offsetOfPos p;
+private def spaceUptoLine : Format → Bool → Int → Nat → SpaceResult
+  | nil,          _,       _, _ => {}
+  | line,         flatten, _, _ => if flatten then { space := 1 } else { foundLine := true }
+  | align force,  flatten, m, w =>
+    if flatten && !force then {}
+    else if w ≤ m then
+      { space := (m - w).toNat }
+    else
+      { foundLine := true }
+  | text s,       flatten, _, _ =>
+    let p := s.posOf '\n'
+    let off := s.offsetOfPos p
     { foundLine := p != s.endPos, foundFlattenedHardLine := flatten && p != s.endPos, space := off }
-  | append f₁ f₂, flatten, w => merge w (spaceUptoLine f₁ flatten w) (spaceUptoLine f₂ flatten)
-  | nest _ f,     flatten, w => spaceUptoLine f flatten w
-  | group f _,    _,       w => spaceUptoLine f true w
-  | tag _ f,      flatten, w => spaceUptoLine f flatten w
+  | append f₁ f₂, flatten, m, w => merge w (spaceUptoLine f₁ flatten m w) (spaceUptoLine f₂ flatten m)
+  | nest n f,     flatten, m, w => spaceUptoLine f flatten (m - n) w
+  | group f _,    _,       m, w => spaceUptoLine f true m w
+  | tag _ f,      flatten, m, w => spaceUptoLine f flatten m w
 
 private structure WorkItem where
   f : Format
@@ -125,10 +145,13 @@ private structure WorkGroup where
   flb     : FlattenBehavior
   items   : List WorkItem
 
-private partial def spaceUptoLine' : List WorkGroup → Nat → SpaceResult
-  |   [],                         _ => {}
-  |   { items := [],    .. }::gs, w => spaceUptoLine' gs w
-  | g@{ items := i::is, .. }::gs, w => merge w (spaceUptoLine i.f g.flatten w) (spaceUptoLine' ({ g with items := is }::gs))
+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' ({ g with items := is }::gs) col)
 
 /-- A monad in which we can pretty-print `Format` objects. -/
 class MonadPrettyFormat (m : Type → Type) where
@@ -145,8 +168,8 @@ private def pushGroup (flb : FlattenBehavior) (items : List WorkItem) (gs : List
   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 r  := spaceUptoLine' [g] (w-k)
-  let r' := merge (w-k) r (spaceUptoLine' gs);
+  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
 
@@ -199,13 +222,28 @@ private partial def be (w : Nat) [Monad m] [MonadPrettyFormat m] : List WorkGrou
           let gs'@(g'::_) ← pushGroup FlattenBehavior.fill is gs (w - " ".length)
             | panic "unreachable"
           if g'.flatten then
-            pushOutput " ";
+            pushOutput " "
             endTags i.activeTags
             be w gs'  -- TODO: use `return`
           else
             breakHere
         else
           breakHere
+    | align force =>
+      if g.flatten && !force then
+        -- flatten (align false) = nil
+        endTags i.activeTags
+        be w (gs' is)
+      else
+        let k ← currColumn
+        if k ≤ i.indent then
+          pushOutput ("".pushn ' ' (i.indent - k).toNat)
+          endTags i.activeTags
+          be w (gs' is)
+        else
+          pushNewline i.indent.toNat
+          endTags i.activeTags
+          be w (gs' is)
     | group f flb =>
       if g.flatten then
         -- flatten (group f) = flatten f

src/Init/NotationExtra.lean

@@ -339,13 +339,9 @@ section
 open Lean.Parser.Tactic
 syntax cdotTk := patternIgnore("·" <|> ".")
 /-- `· tac` focuses on the main goal and tries to solve it using `tac`, or else fails. -/
-syntax cdotTk ppHardSpace many1Indent(tactic ";"? ppLine) : tactic
+syntax cdotTk ppSpace sepBy1IndentSemicolon(tactic) : tactic
 macro_rules
-  | `(tactic| $cdot:cdotTk $[$tacs $[;%$sc]?]*) => do
-    let tacs ← tacs.zip sc |>.mapM fun
-      | (tac, none)    => pure tac
-      | (tac, some sc) => `(tactic| ($tac; with_annotate_state $sc skip))
-    `(tactic| { with_annotate_state $cdot skip; $[$tacs]* })
+  | `(tactic| $cdot:cdotTk $tacs*) => `(tactic| {%$cdot $tacs* }%$cdot)
 end
 
 /--

src/Lean/Data/Format.lean

@@ -41,7 +41,7 @@ open Std
 export Std
   (Format ToFormat Format.nest Format.nil Format.joinSep Format.line
    Format.sbracket Format.bracket Format.group Format.tag Format.pretty
-   Format.fill Format.paren Format.join)
+   Format.fill Format.paren Format.join Format.align)
 export ToFormat (format)
 
 instance : ToFormat Name where

src/Lean/Parser/Extra.lean

@@ -73,18 +73,17 @@ open PrettyPrinter Syntax.MonadTraverser Formatter in
 @[combinator_formatter sepByIndent]
 def sepByIndent.formatter (p : Formatter) (_sep : String) (pSep : Formatter) : Formatter := do
   let stx ← getCur
-  let hasNewlineSep := stx.getArgs.mapIdx (fun ⟨i, _⟩ n => i % 2 == 1 && n.matchesNull 0) |>.any id
+  let hasNewlineSep := stx.getArgs.mapIdx (fun ⟨i, _⟩ n =>
+    i % 2 == 1 && n.matchesNull 0 && i != stx.getArgs.size - 1) |>.any id
   visitArgs do
     for i in (List.range stx.getArgs.size).reverse do
       if i % 2 == 0 then p else pSep <|>
         -- If the final separator is a newline, skip it.
         ((if i == stx.getArgs.size - 1 then pure () else pushWhitespace "\n") *> goLeft)
-  -- If there is any newline separator, then we need to force a newline at the
-  -- start so that `withPosition` will pick up the right column.
+  -- If there is any newline separator, then we add an `align` at the start
+  -- so that `withPosition` will pick up the right column.
   if hasNewlineSep then
-    pushWhitespace "\n"
-    -- HACK: allow formatter to put initial brace on previous line in structure instances
-    modify ({ · with mustBeGrouped := false })
+    pushAlign (force := true)
 
 @[combinator_formatter sepBy1Indent] def sepBy1Indent.formatter := sepByIndent.formatter
 

src/Lean/PrettyPrinter/Formatter.lean

@@ -125,6 +125,10 @@ def pushWhitespace (f : Format) : FormatterM Unit := do
 def pushLine : FormatterM Unit :=
   pushWhitespace Format.line
 
+def pushAlign (force : Bool) : FormatterM Unit :=
+  -- don't reset leadWord because .align may introduce zero space
+  push (.align force)
+
 /-- Execute `x` at the right-most child of the current node, if any, then advance to the left. -/
 def visitArgs (x : FormatterM Unit) : FormatterM Unit := do
   let stx ← getCur

tests/lean/formatTerm.lean

@@ -1,7 +1,7 @@
 import Lean
-open Lean
+open Lean PrettyPrinter
 
-def fmt (stx : CoreM Syntax) : CoreM Format := stx >>= PrettyPrinter.formatTerm
+def fmt (stx : CoreM Syntax) : CoreM Format := stx >>= formatTerm
 
 #eval fmt `(if c then do t else e)
 #eval fmt `(if c then do t; t else e)
@@ -10,3 +10,45 @@ def fmt (stx : CoreM Syntax) : CoreM Format := stx >>= PrettyPrinter.formatTerm
 #eval fmt `(if c then do t else if c then do t else do e) -- FIXME: make this cascade better?
 #eval fmt `(do if c then t else e)
 #eval fmt `(do if c then t else if c then t else e)
+
+#eval fmt `(def foo := by
+  · skip; skip
+  · skip; skip
+    skip
+    (skip; skip)
+    (skip; skip
+     try skip; skip
+     try skip
+         skip
+     skip))
+
+#eval fmt `(by
+  try
+  skip
+  skip)
+
+set_option format.indent 3 in
+#eval fmt `(by
+  try
+  skip
+  skip)
+set_option format.indent 4 in
+#eval fmt `(by
+  try
+  skip
+  skip)
+set_option format.indent 4 in
+#eval fmt `(by
+  try
+    skip
+    skip
+  skip)
+
+#eval fmt `({
+  foo := bar
+  bar := foo + bar
+})
+
+#eval fmt `(let x := { foo := bar
+                       bar := foo + bar }
+            x)

tests/lean/formatTerm.lean.expected.out

@@ -32,3 +32,33 @@ do
     t.1
   else
     e.1
+def foo.1 := by 
+  · skip; skip
+  · skip; skip
+    skip
+    (skip; skip)
+    ( skip; skip
+      try skip; skip
+      try 
+        skip
+        skip
+      skip)
+by
+  try 
+    skip
+    skip
+by
+   try 
+      skip
+      skip
+by try  skip
+        skip
+by  try skip
+        skip
+    skip
+{ foo.1 := bar.1
+  bar.1 := foo.1 + bar.1 }
+let x.1 :=
+  { foo.1 := bar.1
+    bar.1 := foo.1 + bar.1 }
+x.1

tests/lean/sepByIndentQuot.lean.expected.out

@@ -1,6 +1,6 @@
 frobnicate✝ { x := x, x := x, x := x, x := x }
-frobnicate✝ { 
-  x := x
-  x := x
-  x := x
-  x := x }
+frobnicate✝
+  { x := x
+    x := x
+    x := x
+    x := x }

tests/lean/tacUnsolvedGoalsErrors.lean.expected.out

@@ -32,13 +32,13 @@ p q r : Prop
 h2 : p → q
 h✝ : q
 ⊢ q
-tacUnsolvedGoalsErrors.lean:26:2-27:8: error: unsolved goals
+tacUnsolvedGoalsErrors.lean:26:2-26:3: error: unsolved goals
 case inl
 p q r : Prop
 h2 : p → q
 h✝ : p
 ⊢ q
-tacUnsolvedGoalsErrors.lean:28:2-29:8: error: unsolved goals
+tacUnsolvedGoalsErrors.lean:28:2-28:3: error: unsolved goals
 case inr
 p q r : Prop
 h2 : p → q

tests/lean/traceTacticSteps.lean.expected.out

@@ -2,7 +2,7 @@
     True
   [Elab.step.result] True
 [Elab.step] expected type: True, term
-    by
+    by 
       skip
       trivial
   [Elab.step.result] ?m
@@ -15,7 +15,7 @@
 [Elab.step] skip
 [Elab.step] trivial
 [Elab.step] (apply And.intro✝) <;> trivial
-[Elab.step] focus
+[Elab.step] focus 
       apply And.intro✝
       with_annotate_state"<;>" skip
       all_goals trivial