feat(iter): add `Iter.take` (#26525)

2 files changed, 54 insertions, 0 deletions

diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
index f7f722bc0e..7e0ad5f4c3 100644
--- a/runtime/doc/lua.txt
+++ b/runtime/doc/lua.txt
@@ -3639,6 +3639,25 @@ Iter:slice({first}, {last})                                     *Iter:slice()*
     Return: ~
         Iter
 
+Iter:take({n})                                                   *Iter:take()*
+    Transforms an iterator to yield only the first n values.
+
+    Example: >lua
+        local it = vim.iter({ 1, 2, 3, 4 }):take(2)
+        it:next()
+        -- 1
+        it:next()
+        -- 2
+        it:next()
+        -- nil
+<
+
+    Parameters: ~
+      • {n}  (integer)
+
+    Return: ~
+        Iter
+
 Iter:totable()                                                *Iter:totable()*
     Collect the iterator into a table.
 
diff --git a/runtime/lua/vim/iter.lua b/runtime/lua/vim/iter.lua
index e9c2b66bf2..8e602c406a 100644
--- a/runtime/lua/vim/iter.lua
+++ b/runtime/lua/vim/iter.lua
@@ -592,6 +592,41 @@ function ListIter.rfind(self, f) -- luacheck: no unused args
   self._head = self._tail
 end
 
+--- Transforms an iterator to yield only the first n values.
+---
+--- Example:
+---
+--- ```lua
+--- local it = vim.iter({ 1, 2, 3, 4 }):take(2)
+--- it:next()
+--- -- 1
+--- it:next()
+--- -- 2
+--- it:next()
+--- -- nil
+--- ```
+---
+---@param n integer
+---@return Iter
+function Iter.take(self, n)
+  local next = self.next
+  local i = 0
+  self.next = function()
+    if i < n then
+      i = i + 1
+      return next(self)
+    end
+  end
+  return self
+end
+
+---@private
+function ListIter.take(self, n)
+  local inc = self._head < self._tail and 1 or -1
+  self._tail = math.min(self._tail, self._head + n * inc)
+  return self
+end
+
 --- "Pops" a value from a |list-iterator| (gets the last value and decrements the tail).
 ---
 --- Example: