feat: align take/drop/extract across List/Array/Vector (leanprover#6860)

This PR makes `take`/`drop`/`extract` available for each of
`List`/`Array`/`Vector`. The simp normal forms differ, however: in
`List`, we simplify `extract` to `take+drop`, while in `Array` and
`Vector` we simplify `take` and `drop` to `extract`. We also provide
`Array/Vector.shrink`, which simplifies to `take`, but is implemented by
repeatedly popping. Verification lemmas for `Array/Vector.extract` to
follow in a subsequent PR.

Author: kim-em

src/Init/Data/Array/Basic.lean

@@ -270,14 +270,22 @@ def swapAt! (a : Array α) (i : Nat) (v : α) : α × Array α :=
     have : Inhabited (α × Array α) := ⟨(v, a)⟩
     panic! ("index " ++ toString i ++ " out of bounds")
 
-/-- `take a n` returns the first `n` elements of `a`. -/
-def take (a : Array α) (n : Nat) : Array α :=
+/-- `shrink a n` returns the first `n` elements of `a`, implemented by repeatedly popping the last element. -/
+def shrink (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
+/-- `take a n` returns the first `n` elements of `a`, implemented by copying the first `n` elements. -/
+abbrev take (a : Array α) (n : Nat) : Array α := extract a 0 n
+
+@[simp] theorem take_eq_extract (a : Array α) (n : Nat) : a.take n = a.extract 0 n := rfl
+
+/-- `drop a n` removes the first `n` elements of `a`, implemented by copying the remaining elements. -/
+abbrev drop (a : Array α) (n : Nat) : Array α := extract a n a.size
+
+@[simp] theorem drop_eq_extract (a : Array α) (n : Nat) : a.drop n = a.extract n a.size := rfl
 
 @[inline]
 unsafe def modifyMUnsafe [Monad m] (a : Array α) (i : Nat) (f : α → m α) : m (Array α) := do

src/Init/Data/Array/Lemmas.lean

@@ -2565,8 +2565,14 @@ theorem getElem?_extract {as : Array α} {start stop : Nat} :
     · omega
   · rfl
 
+@[congr] theorem extract_congr {as bs : Array α}
+    (w : as = bs) (h : start = start') (h' : stop = stop') :
+    as.extract start stop = bs.extract start' stop' := by
+  subst w h h'
+  rfl
+
 @[simp] theorem toList_extract (as : Array α) (start stop : Nat) :
-    (as.extract start stop).toList = (as.toList.drop start).take (stop - start) := by
+    (as.extract start stop).toList = as.toList.extract start stop := by
   apply List.ext_getElem
   · simp only [length_toList, size_extract, List.length_take, List.length_drop]
     omega
@@ -2595,7 +2601,7 @@ theorem extract_empty_of_size_le_start (as : Array α) {start stop : Nat} (h : a
   extract_empty_of_size_le_start _ (Nat.zero_le _)
 
 @[simp] theorem _root_.List.extract_toArray (l : List α) (start stop : Nat) :
-    l.toArray.extract start stop = ((l.drop start).take (stop - start)).toArray := by
+    l.toArray.extract start stop = (l.extract start stop).toArray := by
   apply ext'
   simp
 
@@ -3363,36 +3369,36 @@ theorem size_eq_length_toList (as : Array α) : as.size = as.toList.length := rf
 theorem getElem_range {n : Nat} {x : Nat} (h : x < (Array.range n).size) : (Array.range n)[x] = x := by
   simp [← getElem_toList]
 
+/-! ### shrink -/
 
-
-
-/-! ### take -/
-
-@[simp] theorem size_take_loop (a : Array α) (n : Nat) : (take.loop n a).size = a.size - n := by
+@[simp] theorem size_shrink_loop (a : Array α) (n : Nat) : (shrink.loop n a).size = a.size - n := by
   induction n generalizing a with
-  | zero => simp [take.loop]
+  | zero => simp [shrink.loop]
   | succ n ih =>
-    simp [take.loop, ih]
+    simp [shrink.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
+@[simp] theorem getElem_shrink_loop (a : Array α) (n : Nat) (i : Nat) (h : i < (shrink.loop n a).size) :
+    (shrink.loop n a)[i] = a[i]'(by simp at h; omega) := by
   induction n generalizing a i with
-  | zero => simp [take.loop]
+  | zero => simp [shrink.loop]
   | succ n ih =>
-    simp [take.loop, ih]
+    simp [shrink.loop, ih]
 
-@[simp] theorem size_take (a : Array α) (n : Nat) : (a.take n).size = min n a.size  := by
-  simp [take]
+@[simp] theorem size_shrink (a : Array α) (n : Nat) : (a.shrink n).size = min n a.size  := by
+  simp [shrink]
   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 getElem_shrink (a : Array α) (n : Nat) (i : Nat) (h : i < (a.shrink n).size) :
+    (a.shrink n)[i] = a[i]'(by simp at h; omega) := by
+  simp [shrink]
 
-@[simp] theorem toList_take (a : Array α) (n : Nat) : (a.take n).toList = a.toList.take n := by
+@[simp] theorem toList_shrink (a : Array α) (n : Nat) : (a.shrink n).toList = a.toList.take n := by
   apply List.ext_getElem <;> simp
 
+@[simp] theorem shrink_eq_take (a : Array α) (n : Nat) : a.shrink n = a.take n := by
+  ext <;> simp
+
 /-! ### forIn -/
 
 @[simp] theorem forIn_toList [Monad m] (as : Array α) (b : β) (f : α → β → m (ForInStep β)) :

src/Init/Data/List/Basic.lean

@@ -823,6 +823,17 @@ theorem drop_eq_nil_of_le {as : List α} {i : Nat} (h : as.length ≤ i) : as.dr
   | _::_,  0   => simp at h
   | _::as, i+1 => simp only [length_cons] at h; exact @drop_eq_nil_of_le as i (Nat.le_of_succ_le_succ h)
 
+/-! ### extract -/
+
+/-- `extract l start stop` returns the slice of `l` from indices `start` to `stop` (exclusive). -/
+-- This is only an abbreviation for the operation in terms of `drop` and `take`.
+-- We do not prove properties of extract itself.
+abbrev extract (l : List α) (start : Nat := 0) (stop : Nat := l.length) : List α :=
+  (l.drop start).take (stop - start)
+
+@[simp] theorem extract_eq_drop_take (l : List α) (start stop : Nat) :
+    l.extract start stop = (l.drop start).take (stop - start) := rfl
+
 /-! ### takeWhile -/
 
 /--

src/Init/Data/Vector/Basic.lean

@@ -163,9 +163,36 @@ instance : HAppend (Vector α n) (Vector α m) (Vector α (n + m)) where
 Extracts the slice of a vector from indices `start` to `stop` (exclusive). If `start ≥ stop`, the
 result is empty. If `stop` is greater than the size of the vector, the size is used instead.
 -/
-@[inline] def extract (v : Vector α n) (start stop : Nat) : Vector α (min stop n - start) :=
+@[inline] def extract (v : Vector α n) (start : Nat := 0) (stop : Nat := n) : Vector α (min stop n - start) :=
   ⟨v.toArray.extract start stop, by simp⟩
 
+/--
+Extract the first `m` elements of a vector. If `m` is greater than or equal to the size of the
+vector then the vector is returned unchanged.
+-/
+@[inline] def take (v : Vector α n) (m : Nat) : Vector α (min m n) :=
+  ⟨v.toArray.take m, by simp⟩
+
+@[simp] theorem take_eq_extract (v : Vector α n) (m : Nat) : v.take m = v.extract 0 m := rfl
+
+/--
+Deletes the first `m` elements of a vector. If `m` is greater than or equal to the size of the
+vector then the empty vector is returned.
+-/
+@[inline] def drop (v : Vector α n) (m : Nat) : Vector α (n - m) :=
+  ⟨v.toArray.drop m, by simp⟩
+
+@[simp] theorem drop_eq_cast_extract (v : Vector α n) (m : Nat) :
+    v.drop m = (v.extract m n).cast (by simp) := by
+  simp [drop, extract, Vector.cast]
+
+/-- Shrinks a vector to the first `m` elements, by repeatedly popping the last element. -/
+@[inline] def shrink (v : Vector α n) (m : Nat) : Vector α (min m n) :=
+  ⟨v.toArray.shrink m, by simp⟩
+
+@[simp] theorem shrink_eq_take (v : Vector α n) (m : Nat) : v.shrink m = v.take m := by
+  simp [shrink, take]
+
 /-- Maps elements of a vector using the function `f`. -/
 @[inline] def map (f : α → β) (v : Vector α n) : Vector β n :=
   ⟨v.toArray.map f, by simp⟩
@@ -291,20 +318,6 @@ This will perform the update destructively provided that the vector has a refere
 /-- The vector `#v[0,1,2,...,n-1]`. -/
 @[inline] def range (n : Nat) : Vector Nat n := ⟨Array.range n, by simp⟩
 
-/--
-Extract the first `m` elements of a vector. If `m` is greater than or equal to the size of the
-vector then the vector is returned unchanged.
--/
-@[inline] def take (v : Vector α n) (m : Nat) : Vector α (min m n) :=
-  ⟨v.toArray.take m, by simp⟩
-
-/--
-Deletes the first `m` elements of a vector. If `m` is greater than or equal to the size of the
-vector then the empty vector is returned.
--/
-@[inline] def drop (v : Vector α n) (m : Nat) : Vector α (n - m) :=
-  ⟨v.toArray.extract m v.size, by simp⟩
-
 /--
 Compares two vectors of the same size using a given boolean relation `r`. `isEqv v w r` returns
 `true` if and only if `r v[i] w[i]` is true for all indices `i`.

src/Init/Prelude.lean

@@ -2706,7 +2706,7 @@ protected def Array.appendCore {α : Type u}  (as : Array α) (bs : Array α) :
   If `start` is greater or equal to `stop`, the result is empty.
   If `stop` is greater than the length of `as`, the length is used instead. -/
 -- NOTE: used in the quotation elaborator output
-def Array.extract (as : Array α) (start stop : Nat) : Array α :=
+def Array.extract (as : Array α) (start : Nat := 0) (stop : Nat := as.size) : Array α :=
   let rec loop (i : Nat) (j : Nat) (bs : Array α) : Array α :=
     dite (LT.lt j as.size)
       (fun hlt =>

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.take toPullSizeSaved ++ keep }
+  modify fun s => { s with toPull := s.toPull.shrink toPullSizeSaved ++ keep }
   return c
 
 def attachToPull (c : Code) : PullM Code := do

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.take size }
+  | some _ => { pos, error? := none, imports := s.imports.shrink size }
 
 @[inline] partial def preludeOpt (k : String) : Parser :=
   keywordCore k (fun _ s => s.pushModule `Init false) (fun _ s => s)

src/Lean/Meta/Tactic/LinearArith/Solver.lean

@@ -37,8 +37,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.take (a : Assignment) (newSize : Nat) : Assignment :=
-  { a with val := a.val.take newSize }
+abbrev Assignment.shrink (a : Assignment) (newSize : Nat) : Assignment :=
+  { a with val := a.val.shrink newSize }
 
 structure Poly where
   val : Array (Int × Var)
@@ -243,7 +243,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.take maxVarIdx
+      let assignment := assignment.shrink maxVarIdx
       if c.lhs.getMaxVarCoeff < 0 then
         let lowers := lowers.modify maxVarIdx (·.push c)
         Sum.inr { lowers, uppers, int, assignment }

src/Lean/Meta/WHNF.lean

@@ -112,7 +112,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.take nparams)
+  return mkAppN (mkConst ctor lvls) (type.getAppArgs.shrink nparams)
 
 private def getRecRuleFor (recVal : RecursorVal) (major : Expr) : Option RecursorRule :=
   match major.getAppFn with
@@ -180,7 +180,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.take ctorInfo.numParams
+        let params := majorType.getAppArgs.shrink 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)

src/Lean/Parser/Basic.lean

@@ -1332,7 +1332,7 @@ namespace ParserState
 
 def keepTop (s : SyntaxStack) (startStackSize : Nat) : SyntaxStack :=
   let node  := s.back
-  s.take startStackSize |>.push node
+  s.shrink startStackSize |>.push node
 
 def keepNewError (s : ParserState) (oldStackSize : Nat) : ParserState :=
   match s with
@@ -1341,13 +1341,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.take oldStackSize, oldLhsPrec, oldStopPos, cache, oldError, errs⟩
+    ⟨stack.shrink 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.take oldStackSize, lhsPrec, pos, cache, some newError, errs⟩
+    ⟨stack.shrink oldStackSize, lhsPrec, pos, cache, some newError, errs⟩
   | other                         => other
 
 def keepLatest (s : ParserState) (startStackSize : Nat) : ParserState :=
@@ -1390,7 +1390,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.takeStack startSize |>.pushSyntax Syntax.missing
+    s.shrinkStack startSize |>.pushSyntax Syntax.missing
   else
     -- parser succeeded with incorrect number of nodes
     invalidLongestMatchParser s

src/Lean/Parser/Types.lean

@@ -158,10 +158,8 @@ def size (stack : SyntaxStack) : Nat :=
 def isEmpty (stack : SyntaxStack) : Bool :=
   stack.size == 0
 
-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 shrink (stack : SyntaxStack) (n : Nat) : SyntaxStack :=
+  { stack with raw := stack.raw.shrink (stack.drop + n) }
 
 def push (stack : SyntaxStack) (a : Syntax) : SyntaxStack :=
   { stack with raw := stack.raw.push a }
@@ -214,7 +212,7 @@ def stackSize (s : ParserState) : Nat :=
   s.stxStack.size
 
 def restore (s : ParserState) (iniStackSz : Nat) (iniPos : String.Pos) : ParserState :=
-  { s with stxStack := s.stxStack.take iniStackSz, errorMsg := none, pos := iniPos }
+  { s with stxStack := s.stxStack.shrink iniStackSz, errorMsg := none, pos := iniPos }
 
 def setPos (s : ParserState) (pos : String.Pos) : ParserState :=
   { s with pos := pos }
@@ -228,10 +226,8 @@ def pushSyntax (s : ParserState) (n : Syntax) : ParserState :=
 def popSyntax (s : ParserState) : ParserState :=
   { s with stxStack := s.stxStack.pop }
 
-def takeStack (s : ParserState) (iniStackSz : Nat) : ParserState :=
-  { s with stxStack := s.stxStack.take iniStackSz }
-
-@[deprecated takeStack (since := "2024-10-22")] abbrev shrinkStack := @takeStack
+def shrinkStack (s : ParserState) (iniStackSz : Nat) : ParserState :=
+  { s with stxStack := s.stxStack.shrink iniStackSz }
 
 def next (s : ParserState) (input : String) (pos : String.Pos) : ParserState :=
   { s with pos := input.next pos }
@@ -254,15 +250,15 @@ 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.take iniStackSz
+      let stack   := stack.shrink iniStackSz
       let stack   := stack.push newNode
       ⟨stack, lhsPrec, pos, cache, err, recovered⟩
 
 def mkTrailingNode (s : ParserState) (k : SyntaxNodeKind) (iniStackSz : Nat) : ParserState :=
   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.take (iniStackSz - 1)
+    let stack   := stack.shrink (iniStackSz - 1)
     let stack   := stack.push newNode
     ⟨stack, lhsPrec, pos, cache, err, errs⟩
 
@@ -287,7 +283,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.takeStack sz
+    s := s.shrinkStack sz
   s := s.setError { expected := ex }
   s.pushSyntax .missing
 

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.take mvars.size
+      let args := args.shrink mvars.size
 
       -- Unify with the expected type
       if (← read).knowsType then tryUnify (← inferType (mkAppN f args)) resultType

src/Lean/PrettyPrinter/Formatter.lean

@@ -146,7 +146,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.take sp).push f
+  setStack $ (stack.shrink sp).push f
 
 /-- Execute `x` and concatenate generated Format objects. -/
 def concat (x : FormatterM Unit) : FormatterM Unit := do

src/lake/Lake/Util/Log.lean

@@ -321,7 +321,7 @@ instance : Append Log := ⟨Log.append⟩
 
 /-- Removes log entries after `pos` (inclusive). -/
 @[inline] def dropFrom (log : Log) (pos : Log.Pos) : Log :=
-  .mk <| log.entries.take pos.val
+  .mk <| log.entries.shrink pos.val
 
 /-- Takes log entries before `pos` (exclusive). -/
 @[inline] def takeFrom (log : Log) (pos : Log.Pos) : Log :=