From c341d8432f0ca9d4e89ecec997297f87e386b118 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 18 Jul 2022 22:14:41 -0400
Subject: [PATCH] feat: remove leading spaces from docstrings

---
 src/Init/Data/String/Extra.lean | 17 ++++++++++++++
 src/Lean/DocString.lean         | 41 +++++++++++++++++++++++++++++++--
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/src/Init/Data/String/Extra.lean b/src/Init/Data/String/Extra.lean
index 18709a010a..d8a19e7a58 100644
--- a/src/Init/Data/String/Extra.lean
+++ b/src/Init/Data/String/Extra.lean
@@ -66,4 +66,21 @@ theorem Iterator.sizeOf_next_lt_of_atEnd (i : String.Iterator) (h : ¬ i.atEnd =
 
 macro_rules | `(tactic| decreasing_trivial) => `(tactic| apply String.Iterator.sizeOf_next_lt_of_atEnd; assumption)
 
+namespace Iterator
+
+@[specialize] def find (it : Iterator) (p : Char → Bool) : Iterator :=
+  if it.atEnd then it
+  else if p it.curr then it
+  else find it.next p
+
+@[specialize] def foldUntil (it : Iterator) (init : α) (f : α → Char → Option α) : α × Iterator :=
+  if it.atEnd then 
+    (init, it)
+  else if let some a := f init it.curr then 
+    foldUntil it.next a f
+  else
+    (init, it)
+
+end Iterator
+
 end String
diff --git a/src/Lean/DocString.lean b/src/Lean/DocString.lean
index 8b2fdf97b9..9b233b2da1 100644
--- a/src/Lean/DocString.lean
+++ b/src/Lean/DocString.lean
@@ -11,11 +11,48 @@ namespace Lean
 private builtin_initialize builtinDocStrings : IO.Ref (NameMap String) ← IO.mkRef {}
 private builtin_initialize docStringExt : MapDeclarationExtension String ← mkMapDeclarationExtension `docstring
 
+private def findLeadingSpacesSize (s : String) : Nat :=     
+  let it := s.iter
+  let it := it.find (· == '\n') |>.next
+  let (min, it) := it.foldUntil 0 fun num c => if c == ' ' || c == '\t' then some (num + 1) else none
+  findNextLine it min 
+where 
+  consumeSpaces (it : String.Iterator) (curr min : Nat) : Nat :=
+    if it.atEnd then min
+    else if it.curr == ' ' || it.curr == '\t' then consumeSpaces it.next (curr + 1) min
+    else findNextLine it.next (Nat.min curr min)
+  findNextLine (it : String.Iterator) (min : Nat) : Nat :=
+    if it.atEnd then min
+    else if it.curr == '\n' then consumeSpaces it.next 0 min 
+    else findNextLine it.next min 
+
+private def removeNumLeadingSpaces (n : Nat) (s : String) : String :=
+  consumeSpaces n s.iter ""
+where
+  consumeSpaces (n : Nat) (it : String.Iterator) (r : String) : String :=
+     match n with
+     | 0 => saveLine it r
+     | n+1 =>
+       if it.atEnd then r 
+       else if it.curr == ' ' || it.curr == '\t' then consumeSpaces n it.next r 
+       else saveLine it r
+  saveLine (it : String.Iterator) (r : String) : String :=
+    if it.atEnd then r
+    else if it.curr == '\n' then consumeSpaces n it.next (r.push '\n')
+    else saveLine it.next (r.push it.curr)  
+termination_by 
+  consumeSpaces n it r => (it, 1) 
+  saveLine it r => (it, 0)
+
+private def removeLeadingSpaces (s : String) : String := 
+  let n := findLeadingSpacesSize s
+  if n == 0 then s else removeNumLeadingSpaces n s
+
 def addBuiltinDocString (declName : Name) (docString : String) : IO Unit :=
-  builtinDocStrings.modify (·.insert declName docString)
+  builtinDocStrings.modify (·.insert declName (removeLeadingSpaces docString))
 
 def addDocString [MonadEnv m] (declName : Name) (docString : String) : m Unit :=
-  modifyEnv fun env => docStringExt.insert env declName docString
+  modifyEnv fun env => docStringExt.insert env declName (removeLeadingSpaces docString)
 
 def addDocString' [Monad m] [MonadEnv m] (declName : Name) (docString? : Option String) : m Unit :=
   match docString? with