From 71122696a1dd8725a5203081560c8a69fe7f7c74 Mon Sep 17 00:00:00 2001
From: Kim Morrison <scott.morrison@gmail.com>
Date: Tue, 22 Oct 2024 10:35:32 +1100
Subject: [PATCH] feat: rename Array.shrink to take, and relate to List.take
 (#5796)

---
 src/Init/Data/Array/Basic.lean                |  5 ++-
 src/Init/Data/Array/Lemmas.lean               | 31 +++++++++++++++++++
 src/Lean/Compiler/LCNF/PullLetDecls.lean      |  2 +-
 src/Lean/Elab/ParseImportsFast.lean           |  2 +-
 src/Lean/Meta/Tactic/LinearArith/Solver.lean  |  6 ++--
 src/Lean/Meta/WHNF.lean                       |  4 +--
 src/Lean/Parser/Basic.lean                    |  8 ++---
 src/Lean/Parser/Types.lean                    | 20 +++++++-----
 .../Delaborator/TopDownAnalyze.lean           |  2 +-
 src/Lean/PrettyPrinter/Formatter.lean         |  2 +-
 src/lake/Lake/Util/Log.lean                   |  2 +-
 tests/playground/parser/parser.lean           | 24 +++++++-------
 12 files changed, 73 insertions(+), 35 deletions(-)

diff --git a/src/Init/Data/Array/Basic.lean b/src/Init/Data/Array/Basic.lean
index d057e5e87d..88117d0654 100644
--- a/src/Init/Data/Array/Basic.lean
+++ b/src/Init/Data/Array/Basic.lean
@@ -241,12 +241,15 @@ def swapAt! (a : Array α) (i : Nat) (v : α) : α × Array α :=
     have : Inhabited (α × Array α) := ⟨(v, a)⟩
     panic! ("index " ++ toString i ++ " out of bounds")
 
-def shrink (a : Array α) (n : Nat) : Array α :=
+/-- `take a n` returns the first `n` elements of `a`. -/
+def take (a : Array α) (n : Nat) : Array α :=
   let rec loop
     | 0,   a => a
     | n+1, a => loop n a.pop
   loop (a.size - n) a
 
+@[deprecated take (since := "2024-10-22")] abbrev shrink := @take
+
 @[inline]
 unsafe def modifyMUnsafe [Monad m] (a : Array α) (i : Nat) (f : α → m α) : m (Array α) := do
   if h : i < a.size then
diff --git a/src/Init/Data/Array/Lemmas.lean b/src/Init/Data/Array/Lemmas.lean
index 0a60920f7c..4a156c9adf 100644
--- a/src/Init/Data/Array/Lemmas.lean
+++ b/src/Init/Data/Array/Lemmas.lean
@@ -671,6 +671,33 @@ theorem getElem_range {n : Nat} {x : Nat} (h : x < (Array.range n).size) : (Arra
         true_and, Nat.not_lt] at h
       rw [List.getElem?_eq_none_iff.2 ‹_›, List.getElem?_eq_none_iff.2 (a.toList.length_reverse ▸ ‹_›)]
 
+/-! ### take -/
+
+@[simp] theorem size_take_loop (a : Array α) (n : Nat) : (take.loop n a).size = a.size - n := by
+  induction n generalizing a with
+  | zero => simp [take.loop]
+  | succ n ih =>
+    simp [take.loop, ih]
+    omega
+
+@[simp] theorem getElem_take_loop (a : Array α) (n : Nat) (i : Nat) (h : i < (take.loop n a).size) :
+    (take.loop n a)[i] = a[i]'(by simp at h; omega) := by
+  induction n generalizing a i with
+  | zero => simp [take.loop]
+  | succ n ih =>
+    simp [take.loop, ih]
+
+@[simp] theorem size_take (a : Array α) (n : Nat) : (a.take n).size = min n a.size  := by
+  simp [take]
+  omega
+
+@[simp] theorem getElem_take (a : Array α) (n : Nat) (i : Nat) (h : i < (a.take n).size) :
+    (a.take n)[i] = a[i]'(by simp at h; omega) := by
+  simp [take]
+
+@[simp] theorem toList_take (a : Array α) (n : Nat) : (a.take n).toList = a.toList.take n := by
+  apply List.ext_getElem <;> simp
+
 /-! ### foldl / foldr -/
 
 @[simp] theorem foldlM_loop_empty [Monad m] (f : β → α → m β) (init : β) (i j : Nat) :
@@ -1339,6 +1366,10 @@ Our goal is to have `simp` "pull `List.toArray` outwards" as much as possible.
   apply ext'
   simp
 
+@[simp] theorem take_toArray (l : List α) (n : Nat) : l.toArray.take n = (l.take n).toArray := by
+  apply ext'
+  simp
+
 @[simp] theorem mapM_toArray [Monad m] [LawfulMonad m] (f : α → m β) (l : List α) :
     l.toArray.mapM f = List.toArray <$> l.mapM f := by
   simp only [← mapM'_eq_mapM, mapM_eq_foldlM]
diff --git a/src/Lean/Compiler/LCNF/PullLetDecls.lean b/src/Lean/Compiler/LCNF/PullLetDecls.lean
index 31496f2ae2..81a5acca3b 100644
--- a/src/Lean/Compiler/LCNF/PullLetDecls.lean
+++ b/src/Lean/Compiler/LCNF/PullLetDecls.lean
@@ -46,7 +46,7 @@ partial def withCheckpoint (x : PullM Code) : PullM Code := do
     else
       return c
   let (c, keep) := go toPullSizeSaved (← read).included |>.run #[]
-  modify fun s => { s with toPull := s.toPull.shrink toPullSizeSaved ++ keep }
+  modify fun s => { s with toPull := s.toPull.take toPullSizeSaved ++ keep }
   return c
 
 def attachToPull (c : Code) : PullM Code := do
diff --git a/src/Lean/Elab/ParseImportsFast.lean b/src/Lean/Elab/ParseImportsFast.lean
index aafa94ef64..59e480e3a7 100644
--- a/src/Lean/Elab/ParseImportsFast.lean
+++ b/src/Lean/Elab/ParseImportsFast.lean
@@ -182,7 +182,7 @@ partial def moduleIdent (runtimeOnly : Bool) : Parser := fun input s =>
   let s := p input s
   match s.error? with
   | none => many p input s
-  | some _ => { pos, error? := none, imports := s.imports.shrink size }
+  | some _ => { pos, error? := none, imports := s.imports.take size }
 
 @[inline] partial def preludeOpt (k : String) : Parser :=
   keywordCore k (fun _ s => s.pushModule `Init false) (fun _ s => s)
diff --git a/src/Lean/Meta/Tactic/LinearArith/Solver.lean b/src/Lean/Meta/Tactic/LinearArith/Solver.lean
index 8a560658d7..9d00100491 100644
--- a/src/Lean/Meta/Tactic/LinearArith/Solver.lean
+++ b/src/Lean/Meta/Tactic/LinearArith/Solver.lean
@@ -36,8 +36,8 @@ abbrev Assignment.get? (a : Assignment) (x : Var) : Option Rat :=
 abbrev Assignment.push (a : Assignment) (v : Rat) : Assignment :=
   { a with val := a.val.push v }
 
-abbrev Assignment.shrink (a : Assignment) (newSize : Nat) : Assignment :=
-  { a with val := a.val.shrink newSize }
+abbrev Assignment.take (a : Assignment) (newSize : Nat) : Assignment :=
+  { a with val := a.val.take newSize }
 
 structure Poly where
   val : Array (Int × Var)
@@ -242,7 +242,7 @@ def resolve (s : State) (cl : Cnstr) (cu : Cnstr) : Sum Result State :=
     let maxVarIdx := c.lhs.getMaxVar.id
     match s with -- Hack: we avoid { s with ... } to make sure we get a destructive update
     | { lowers, uppers, int, assignment, } =>
-      let assignment := assignment.shrink maxVarIdx
+      let assignment := assignment.take maxVarIdx
       if c.lhs.getMaxVarCoeff < 0 then
         let lowers := lowers.modify maxVarIdx (·.push c)
         Sum.inr { lowers, uppers, int, assignment }
diff --git a/src/Lean/Meta/WHNF.lean b/src/Lean/Meta/WHNF.lean
index c7b5ce4e36..55db7b40d4 100644
--- a/src/Lean/Meta/WHNF.lean
+++ b/src/Lean/Meta/WHNF.lean
@@ -84,7 +84,7 @@ private def mkNullaryCtor (type : Expr) (nparams : Nat) : MetaM (Option Expr) :=
   let .const d lvls := type.getAppFn
     | return none
   let (some ctor) ← getFirstCtor d | pure none
-  return mkAppN (mkConst ctor lvls) (type.getAppArgs.shrink nparams)
+  return mkAppN (mkConst ctor lvls) (type.getAppArgs.take nparams)
 
 private def getRecRuleFor (recVal : RecursorVal) (major : Expr) : Option RecursorRule :=
   match major.getAppFn with
@@ -152,7 +152,7 @@ private def toCtorWhenStructure (inductName : Name) (major : Expr) : MetaM Expr
       else
         let some ctorName ← getFirstCtor d | pure major
         let ctorInfo ← getConstInfoCtor ctorName
-        let params := majorType.getAppArgs.shrink ctorInfo.numParams
+        let params := majorType.getAppArgs.take ctorInfo.numParams
         let mut result := mkAppN (mkConst ctorName us) params
         for i in [:ctorInfo.numFields] do
           result := mkApp result (← mkProjFn ctorInfo us params i major)
diff --git a/src/Lean/Parser/Basic.lean b/src/Lean/Parser/Basic.lean
index b47e4cde1c..8dd33413de 100644
--- a/src/Lean/Parser/Basic.lean
+++ b/src/Lean/Parser/Basic.lean
@@ -1305,7 +1305,7 @@ namespace ParserState
 
 def keepTop (s : SyntaxStack) (startStackSize : Nat) : SyntaxStack :=
   let node  := s.back
-  s.shrink startStackSize |>.push node
+  s.take startStackSize |>.push node
 
 def keepNewError (s : ParserState) (oldStackSize : Nat) : ParserState :=
   match s with
@@ -1314,13 +1314,13 @@ def keepNewError (s : ParserState) (oldStackSize : Nat) : ParserState :=
 def keepPrevError (s : ParserState) (oldStackSize : Nat) (oldStopPos : String.Pos) (oldError : Option Error) (oldLhsPrec : Nat) : ParserState :=
   match s with
   | ⟨stack, _, _, cache, _, errs⟩ =>
-    ⟨stack.shrink oldStackSize, oldLhsPrec, oldStopPos, cache, oldError, errs⟩
+    ⟨stack.take oldStackSize, oldLhsPrec, oldStopPos, cache, oldError, errs⟩
 
 def mergeErrors (s : ParserState) (oldStackSize : Nat) (oldError : Error) : ParserState :=
   match s with
   | ⟨stack, lhsPrec, pos, cache, some err, errs⟩ =>
     let newError := if oldError == err then err else oldError.merge err
-    ⟨stack.shrink oldStackSize, lhsPrec, pos, cache, some newError, errs⟩
+    ⟨stack.take oldStackSize, lhsPrec, pos, cache, some newError, errs⟩
   | other                         => other
 
 def keepLatest (s : ParserState) (startStackSize : Nat) : ParserState :=
@@ -1363,7 +1363,7 @@ def runLongestMatchParser (left? : Option Syntax) (startLhsPrec : Nat) (p : Pars
     s -- success or error with the expected number of nodes
   else if s.hasError then
     -- error with an unexpected number of nodes.
-    s.shrinkStack startSize |>.pushSyntax Syntax.missing
+    s.takeStack startSize |>.pushSyntax Syntax.missing
   else
     -- parser succeeded with incorrect number of nodes
     invalidLongestMatchParser s
diff --git a/src/Lean/Parser/Types.lean b/src/Lean/Parser/Types.lean
index a410f406bf..9f6d8b3606 100644
--- a/src/Lean/Parser/Types.lean
+++ b/src/Lean/Parser/Types.lean
@@ -158,8 +158,10 @@ def size (stack : SyntaxStack) : Nat :=
 def isEmpty (stack : SyntaxStack) : Bool :=
   stack.size == 0
 
-def shrink (stack : SyntaxStack) (n : Nat) : SyntaxStack :=
-  { stack with raw := stack.raw.shrink (stack.drop + n) }
+def take (stack : SyntaxStack) (n : Nat) : SyntaxStack :=
+  { stack with raw := stack.raw.take (stack.drop + n) }
+
+@[deprecated take (since := "2024-10-22")] abbrev shrink := @take
 
 def push (stack : SyntaxStack) (a : Syntax) : SyntaxStack :=
   { stack with raw := stack.raw.push a }
@@ -212,7 +214,7 @@ def stackSize (s : ParserState) : Nat :=
   s.stxStack.size
 
 def restore (s : ParserState) (iniStackSz : Nat) (iniPos : String.Pos) : ParserState :=
-  { s with stxStack := s.stxStack.shrink iniStackSz, errorMsg := none, pos := iniPos }
+  { s with stxStack := s.stxStack.take iniStackSz, errorMsg := none, pos := iniPos }
 
 def setPos (s : ParserState) (pos : String.Pos) : ParserState :=
   { s with pos := pos }
@@ -226,8 +228,10 @@ def pushSyntax (s : ParserState) (n : Syntax) : ParserState :=
 def popSyntax (s : ParserState) : ParserState :=
   { s with stxStack := s.stxStack.pop }
 
-def shrinkStack (s : ParserState) (iniStackSz : Nat) : ParserState :=
-  { s with stxStack := s.stxStack.shrink iniStackSz }
+def takeStack (s : ParserState) (iniStackSz : Nat) : ParserState :=
+  { s with stxStack := s.stxStack.take iniStackSz }
+
+@[deprecated takeStack (since := "2024-10-22")] abbrev shrinkStack := @takeStack
 
 def next (s : ParserState) (input : String) (pos : String.Pos) : ParserState :=
   { s with pos := input.next pos }
@@ -250,7 +254,7 @@ def mkNode (s : ParserState) (k : SyntaxNodeKind) (iniStackSz : Nat) : ParserSta
       ⟨stack, lhsPrec, pos, cache, err, recovered⟩
     else
       let newNode := Syntax.node SourceInfo.none k (stack.extract iniStackSz stack.size)
-      let stack   := stack.shrink iniStackSz
+      let stack   := stack.take iniStackSz
       let stack   := stack.push newNode
       ⟨stack, lhsPrec, pos, cache, err, recovered⟩
 
@@ -258,7 +262,7 @@ def mkTrailingNode (s : ParserState) (k : SyntaxNodeKind) (iniStackSz : Nat) : P
   match s with
   | ⟨stack, lhsPrec, pos, cache, err, errs⟩ =>
     let newNode := Syntax.node SourceInfo.none k (stack.extract (iniStackSz - 1) stack.size)
-    let stack   := stack.shrink (iniStackSz - 1)
+    let stack   := stack.take (iniStackSz - 1)
     let stack   := stack.push newNode
     ⟨stack, lhsPrec, pos, cache, err, errs⟩
 
@@ -283,7 +287,7 @@ def mkEOIError (s : ParserState) (expected : List String := []) : ParserState :=
 def mkErrorsAt (s : ParserState) (ex : List String) (pos : String.Pos) (initStackSz? : Option Nat := none) : ParserState := Id.run do
   let mut s := s.setPos pos
   if let some sz := initStackSz? then
-    s := s.shrinkStack sz
+    s := s.takeStack sz
   s := s.setError { expected := ex }
   s.pushSyntax .missing
 
diff --git a/src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean b/src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean
index 3d9e050755..1e749cd211 100644
--- a/src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean
+++ b/src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean
@@ -398,7 +398,7 @@ mutual
       let fType ← replaceLPsWithVars (← inferType f)
       let (mvars, bInfos, resultType) ← forallMetaBoundedTelescope fType args.size
       let rest := args.extract mvars.size args.size
-      let args := args.shrink mvars.size
+      let args := args.take mvars.size
 
       -- Unify with the expected type
       if (← read).knowsType then tryUnify (← inferType (mkAppN f args)) resultType
diff --git a/src/Lean/PrettyPrinter/Formatter.lean b/src/Lean/PrettyPrinter/Formatter.lean
index 12129b2641..8b9324fe3d 100644
--- a/src/Lean/PrettyPrinter/Formatter.lean
+++ b/src/Lean/PrettyPrinter/Formatter.lean
@@ -144,7 +144,7 @@ def fold (fn : Array Format → Format) (x : FormatterM Unit) : FormatterM Unit
   x
   let stack ← getStack
   let f := fn $ stack.extract sp stack.size
-  setStack $ (stack.shrink sp).push f
+  setStack $ (stack.take sp).push f
 
 /-- Execute `x` and concatenate generated Format objects. -/
 def concat (x : FormatterM Unit) : FormatterM Unit := do
diff --git a/src/lake/Lake/Util/Log.lean b/src/lake/Lake/Util/Log.lean
index c106fe8172..dd34583cb9 100644
--- a/src/lake/Lake/Util/Log.lean
+++ b/src/lake/Lake/Util/Log.lean
@@ -292,7 +292,7 @@ instance : Append Log := ⟨Log.append⟩
 
 /-- Removes log entries after `pos` (inclusive). -/
 @[inline] def dropFrom (log : Log) (pos : Log.Pos) : Log :=
-  .mk <| log.entries.shrink pos.val
+  .mk <| log.entries.take pos.val
 
 /-- Takes log entries before `pos` (exclusive). -/
 @[inline] def takeFrom (log : Log) (pos : Log.Pos) : Log :=
diff --git a/tests/playground/parser/parser.lean b/tests/playground/parser/parser.lean
index 555c3756a3..50ac4e0a61 100644
--- a/tests/playground/parser/parser.lean
+++ b/tests/playground/parser/parser.lean
@@ -64,7 +64,7 @@ d.errorMsg != none
 d.stxStack.size
 
 def ParserData.restore (d : ParserData) (iniStackSz : Nat) (iniPos : Nat) : ParserData :=
-{ stxStack := d.stxStack.shrink iniStackSz, errorMsg := none, pos := iniPos, .. d}
+{ stxStack := d.stxStack.take iniStackSz, errorMsg := none, pos := iniPos, .. d}
 
 def ParserData.setPos (d : ParserData) (pos : Nat) : ParserData :=
 { pos := pos, .. d }
@@ -75,8 +75,8 @@ def ParserData.setCache (d : ParserData) (cache : ParserCache) : ParserData :=
 def ParserData.pushSyntax (d : ParserData) (n : Syntax) : ParserData :=
 { stxStack := d.stxStack.push n, .. d }
 
-def ParserData.shrinkStack (d : ParserData) (iniStackSz : Nat) : ParserData :=
-{ stxStack := d.stxStack.shrink iniStackSz, .. d }
+def ParserData.takeStack (d : ParserData) (iniStackSz : Nat) : ParserData :=
+{ stxStack := d.stxStack.take iniStackSz, .. d }
 
 def ParserData.next (d : ParserData) (s : String) (pos : Nat) : ParserData :=
 { pos := s.next pos, .. d }
@@ -114,7 +114,7 @@ match d with
     d
   else
     let newNode := Syntax.node k (stack.extract iniStackSz stack.size) [] in
-    let stack   := stack.shrink iniStackSz in
+    let stack   := stack.take iniStackSz in
     let stack   := stack.push newNode in
     ⟨stack, pos, cache, err⟩
 
@@ -144,7 +144,7 @@ match d with
   let iniSz  := d.stackSize in
   let iniPos := d.pos in
   match p s d with
-  | ⟨stack, _, cache, some msg⟩ := ⟨stack.shrink iniSz, iniPos, cache, some msg⟩
+  | ⟨stack, _, cache, some msg⟩ := ⟨stack.take iniSz, iniPos, cache, some msg⟩
   | other                       := other
 
 @[noinline] def noFirstTokenInfo (info : ParserInfo) : ParserInfo :=
@@ -516,15 +516,15 @@ partial def identFnAux (startPos : Nat) (tk : Option TokenConfig) : Name → Par
 
 def ParserData.keepNewError (d : ParserData) (oldStackSize : Nat) : ParserData :=
 match d with
-| ⟨stack, pos, cache, err⟩ := ⟨stack.shrink oldStackSize, pos, cache, err⟩
+| ⟨stack, pos, cache, err⟩ := ⟨stack.take oldStackSize, pos, cache, err⟩
 
 def ParserData.keepPrevError (d : ParserData) (oldStackSize : Nat) (oldStopPos : String.Pos) (oldError : Option String) : ParserData :=
 match d with
-| ⟨stack, _, cache, _⟩ := ⟨stack.shrink oldStackSize, oldStopPos, cache, oldError⟩
+| ⟨stack, _, cache, _⟩ := ⟨stack.take oldStackSize, oldStopPos, cache, oldError⟩
 
 def ParserData.mergeErrors (d : ParserData) (oldStackSize : Nat) (oldError : String) : ParserData :=
 match d with
-| ⟨stack, pos, cache, some err⟩ := ⟨stack.shrink oldStackSize, pos, cache, some (err ++ "; " ++ oldError)⟩
+| ⟨stack, pos, cache, some err⟩ := ⟨stack.take oldStackSize, pos, cache, some (err ++ "; " ++ oldError)⟩
 | other                         := other
 
 def ParserData.mkLongestNodeAlt (d : ParserData) (startSize : Nat) : ParserData :=
@@ -535,14 +535,14 @@ match d with
   else
     -- parser created more than one node, combine them into a single node
     let node := Syntax.node nullKind (stack.extract startSize stack.size) [] in
-    let stack := stack.shrink startSize in
+    let stack := stack.take startSize in
     ⟨stack.push node, pos, cache, none⟩
 
 def ParserData.keepLatest (d : ParserData) (startStackSize : Nat) : ParserData :=
 match d with
 | ⟨stack, pos, cache, _⟩ :=
   let node  := stack.back in
-  let stack := stack.shrink startStackSize in
+  let stack := stack.take startStackSize in
   let stack := stack.push node in
   ⟨stack, pos, cache, none⟩
 
@@ -591,7 +591,7 @@ def longestMatchFn₂ (p q : ParserFn) : ParserFn :=
 let startSize := d.stackSize in
 let startPos  := d.pos in
 let d         := p s d in
-let d         := if d.hasError then d.shrinkStack startSize else d.mkLongestNodeAlt startSize in
+let d         := if d.hasError then d.takeStack startSize else d.mkLongestNodeAlt startSize in
 let d         := longestMatchStep startSize startPos q s d in
 longestMatchMkResult startSize d
 
@@ -603,7 +603,7 @@ def longestMatchFn : List ParserFn → ParserFn
   let startPos  := d.pos in
   let d         := p s d in
   if d.hasError then
-    let d := d.shrinkStack startSize in
+    let d := d.takeStack startSize in
     longestMatchFnAux startSize startPos ps s d
   else
     let d := d.mkLongestNodeAlt startSize in