diff options
| -rw-r--r-- | runtime/doc/lua.txt | 511 | ||||
| -rw-r--r-- | runtime/doc/news.txt | 3 | ||||
| -rw-r--r-- | runtime/lua/vim/iter.lua | 836 | ||||
| -rw-r--r-- | runtime/lua/vim/shared.lua | 105 | ||||
| -rwxr-xr-x | scripts/gen_vimdoc.py | 4 | ||||
| -rw-r--r-- | test/functional/lua/vim_spec.lua | 354 |
6 files changed, 1812 insertions, 1 deletions
diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt index 6fdf3775f6..60067f1bc3 100644 --- a/runtime/doc/lua.txt +++ b/runtime/doc/lua.txt @@ -1653,6 +1653,26 @@ endswith({s}, {suffix}) *vim.endswith()* Return: ~ (boolean) `true` if `suffix` is a suffix of `s` +filter({f}, {src}, {...}) *vim.filter()* + Filter a table or iterator. + + This is a convenience function that performs: >lua + + vim.iter(src):filter(f):totable() +< + + Parameters: ~ + • {f} function(...):bool Filter function. Accepts the current + iterator or table values as arguments and returns true if those + values should be kept in the final table + • {src} table|function Table or iterator function to filter + + Return: ~ + (table) + + See also: ~ + • |Iter:filter()| + gsplit({s}, {sep}, {opts}) *vim.gsplit()* Splits a string at each instance of a separator. @@ -1698,6 +1718,64 @@ is_callable({f}) *vim.is_callable()* Return: ~ (boolean) `true` if `f` is callable, else `false` +iter({src}, {...}) *vim.iter()* + Create an Iter |lua-iter| object from a table or iterator. + + The input value can be a table or a function iterator (see |luaref-in|). + + This function wraps the input value into an interface which allows + chaining multiple pipeline stages in an efficient manner. Each pipeline + stage receives as input the output values from the prior stage. The values + used in the first stage of the pipeline depend on the type passed to this + function: + + • List tables pass only the value of each element + • Non-list tables pass both the key and value of each element + • Function iterators pass all of the values returned by their respective + function + + Examples: >lua + + local it = vim.iter({ 1, 2, 3, 4, 5 }) + it:map(function(v) + return v * 3 + end) + it:rev() + it:skip(2) + it:totable() + -- { 9, 6, 3 } + + vim.iter(ipairs({ 1, 2, 3, 4, 5 })):map(function(i, v) + if i > 2 then return v end + end):totable() + -- { 3, 4, 5 } + + local it = vim.iter(vim.gsplit('1,2,3,4,5', ',')) + it:map(function(s) return tonumber(s) end) + for i, d in it:enumerate() do + print(string.format("Column %d is %d", i, d)) + end + -- Column 1 is 1 + -- Column 2 is 2 + -- Column 3 is 3 + -- Column 4 is 4 + -- Column 5 is 5 + + vim.iter({ a = 1, b = 2, c = 3, z = 26 }):any(function(k, v) + return k == 'z' + end) + -- true +< + + Parameters: ~ + • {src} table|function Table or iterator. + + Return: ~ + Iter |lua-iter| + + See also: ~ + • |lua-iter| + list_contains({t}, {value}) *vim.list_contains()* Checks if a list-like table (integer keys without gaps) contains `value`. @@ -1740,6 +1818,26 @@ list_slice({list}, {start}, {finish}) *vim.list_slice()* Return: ~ (list) Copy of table sliced from start to finish (inclusive) +map({f}, {src}, {...}) *vim.map()* + Map and filter a table or iterator. + + This is a convenience function that performs: >lua + + vim.iter(src):map(f):totable() +< + + Parameters: ~ + • {f} function(...):?any Map function. Accepts the current iterator + or table values as arguments and returns one or more new + values. Nil values are removed from the final table. + • {src} table|function Table or iterator function to filter + + Return: ~ + (table) + + See also: ~ + • |Iter:map()| + pesc({s}) *vim.pesc()* Escapes magic chars in |lua-patterns|. @@ -2001,6 +2099,20 @@ tbl_values({t}) *vim.tbl_values()* Return: ~ (list) List of values +totable({f}, {...}) *vim.totable()* + Collect an iterator into a table. + + This is a convenience function that performs: >lua + + vim.iter(f):totable() +< + + Parameters: ~ + • {f} (function) Iterator function + + Return: ~ + (table) + trim({s}) *vim.trim()* Trim whitespace (Lua pattern "%s") from both sides of a string. @@ -2817,4 +2929,403 @@ range({spec}) *vim.version.range()* See also: ~ • # https://github.com/npm/node-semver#ranges + +============================================================================== +Lua module: iter *lua-iter* + +Iter:all({self}, {pred}) *Iter:all()* + Return true if all of the items in the iterator match the given predicate. + + Parameters: ~ + • {pred} function(...):bool Predicate function. Takes all values + returned from the previous stage in the pipeline as arguments + and returns true if the predicate matches. + +Iter:any({self}, {pred}) *Iter:any()* + Return true if any of the items in the iterator match the given predicate. + + Parameters: ~ + • {pred} function(...):bool Predicate function. Takes all values + returned from the previous stage in the pipeline as arguments + and returns true if the predicate matches. + +Iter:each({self}, {f}) *Iter:each()* + Call a function once for each item in the pipeline. + + This is used for functions which have side effects. To modify the values + in the iterator, use |Iter:map()|. + + This function drains the iterator. + + Parameters: ~ + • {f} function(...) Function to execute for each item in the pipeline. + Takes all of the values returned by the previous stage in the + pipeline as arguments. + +Iter:enumerate({self}) *Iter:enumerate()* + Add an iterator stage that returns the current iterator count as well as + the iterator value. + + For list tables, prefer >lua + + vim.iter(ipairs(t)) +< + + over + + >lua + + vim.iter(t):enumerate() +< + + as the former is faster. + + Example: >lua + + local it = vim.iter(vim.gsplit('abc', '')):enumerate() + it:next() + -- 1 'a' + it:next() + -- 2 'b' + it:next() + -- 3 'c' +< + + Return: ~ + Iter + +Iter:filter({self}, {f}) *Iter:filter()* + Add a filter step to the iterator pipeline. + + Example: >lua + + local bufs = vim.iter(vim.api.nvim_list_bufs()):filter(vim.api.nvim_buf_is_loaded) +< + + Parameters: ~ + • {f} function(...):bool Takes all values returned from the previous + stage in the pipeline and returns false or nil if the current + iterator element should be removed. + + Return: ~ + Iter + +Iter:find({self}, {f}) *Iter:find()* + Find the first value in the iterator that satisfies the given predicate. + + Advances the iterator. Returns nil and drains the iterator if no value is + found. + + Examples: >lua + + local it = vim.iter({ 3, 6, 9, 12 }) + it:find(12) + -- 12 + + local it = vim.iter({ 3, 6, 9, 12 }) + it:find(20) + -- nil + + local it = vim.iter({ 3, 6, 9, 12 }) + it:find(function(v) return v % 4 == 0 end) + -- 12 +< + + Return: ~ + any + +Iter:fold({self}, {init}, {f}) *Iter:fold()* + Fold an iterator or table into a single value. + + Parameters: ~ + • {init} any Initial value of the accumulator. + • {f} function(acc:any, ...):A Accumulation function. + + Return: ~ + any + +Iter:last({self}) *Iter:last()* + Return the last item in the iterator. + + Drains the iterator. + + Example: >lua + + local it = vim.iter(vim.gsplit('abcdefg', '')) + it:last() + -- 'g' + + local it = vim.iter({ 3, 6, 9, 12, 15 }) + it:last() + -- 15 +< + + Return: ~ + any + +Iter:map({self}, {f}) *Iter:map()* + Add a map step to the iterator pipeline. + + If the map function returns nil, the value is filtered from the iterator. + + Example: >lua + + local it = vim.iter({ 1, 2, 3, 4 }):map(function(v) + if v % 2 == 0 then + return v * 3 + end + end) + it:totable() + -- { 6, 12 } +< + + Parameters: ~ + • {f} function(...):any Mapping function. Takes all values returned + from the previous stage in the pipeline as arguments and returns + one or more new values, which are used in the next pipeline + stage. Nil return values returned are filtered from the output. + + Return: ~ + Iter + +Iter:next({self}) *Iter:next()* + Return the next value from the iterator. + + Example: >lua + + local it = vim.iter(string.gmatch('1 2 3', 'd+')):map(tonumber) + it:next() + -- 1 + it:next() + -- 2 + it:next() + -- 3 +< + + Return: ~ + any + +Iter:nextback({self}) *Iter:nextback()* + Return the next value from the end of the iterator. + + Only supported for iterators on list-like tables. + + Example: >lua + + local it = vim.iter({1, 2, 3, 4}) + it:nextback() + -- 4 + it:nextback() + -- 3 +< + + Return: ~ + any + +Iter:nth({self}, {n}) *Iter:nth()* + Return the nth value in the iterator. + + This function advances the iterator. + + Example: >lua + + local it = vim.iter({ 3, 6, 9, 12 }) + it:nth(2) + -- 6 + it:nth(2) + -- 12 +< + + Parameters: ~ + • {n} (number) The index of the value to return. + + Return: ~ + any + +Iter:nthback({self}, {n}) *Iter:nthback()* + Return the nth value from the end of the iterator. + + This function advances the iterator. + + Only supported for iterators on list-like tables. + + Example: >lua + + local it = vim.iter({ 3, 6, 9, 12 }) + it:nthback(2) + -- 9 + it:nthback(2) + -- 3 +< + + Parameters: ~ + • {n} (number) The index of the value to return. + + Return: ~ + any + +Iter:peek({self}) *Iter:peek()* + Peek at the next value in the iterator without consuming it. + + Only supported for iterators on list-like tables. + + Example: >lua + + local it = vim.iter({ 3, 6, 9, 12 }) + it:peek() + -- 3 + it:peek() + -- 3 + it:next() + -- 3 +< + + Return: ~ + any + +Iter:peekback({self}) *Iter:peekback()* + Return the next value from the end of the iterator without consuming it. + + Only supported for iterators on list-like tables. + + Example: >lua + + local it = vim.iter({1, 2, 3, 4}) + it:peekback() + -- 4 + it:peekback() + -- 4 + it:nextback() + -- 4 +< + + Return: ~ + any + +Iter:rev({self}) *Iter:rev()* + Reverse an iterator. + + Only supported for iterators on list-like tables. + + Example: >lua + + local it = vim.iter({ 3, 6, 9, 12 }):rev() + it:totable() + -- { 12, 9, 6, 3 } +< + + Return: ~ + Iter + +Iter:rfind({self}, {f}) *Iter:rfind()* + Find the first value in the iterator that satisfies the given predicate, + starting from the end. + + Advances the iterator. Returns nil and drains the iterator if no value is + found. + + Only supported for iterators on list-like tables. + + Examples: >lua + + local it = vim.iter({ 1, 2, 3, 2, 1 }):enumerate() + it:rfind(1) + -- 5 1 + it:rfind(1) + -- 1 1 +< + + Return: ~ + any + + See also: ~ + • Iter.find + +Iter:skip({self}, {n}) *Iter:skip()* + Skip values in the iterator. + + Example: >lua + + local it = vim.iter({ 3, 6, 9, 12 }):skip(2) + it:next() + -- 9 +< + + Parameters: ~ + • {n} (number) Number of values to skip. + + Return: ~ + Iter + +Iter:skipback({self}, {n}) *Iter:skipback()* + Skip values in the iterator starting from the end. + + Only supported for iterators on list-like tables. + + Example: >lua + + local it = vim.iter({ 1, 2, 3, 4, 5 }):skipback(2) + it:next() + -- 1 + it:nextback() + -- 3 +< + + Parameters: ~ + • {n} (number) Number of values to skip. + + Return: ~ + Iter + +Iter:slice({self}, {first}, {last}) *Iter:slice()* + Slice an iterator, changing its start and end positions. + + This is equivalent to :skip(first - 1):skipback(len - last + 1) + + Only supported for iterators on list-like tables. + + Parameters: ~ + • {first} (number) + • {last} (number) + + Return: ~ + Iter + +Iter:totable({self}) *Iter:totable()* + Collect the iterator into a table. + + The resulting table depends on the initial source in the iterator + pipeline. List-like tables and function iterators will be collected into a + list-like table. If multiple values are returned from the final stage in + the iterator pipeline, each value will be included in a table. If a + map-like table was used as the initial source, then a map-like table is + returned. + + Examples: >lua + + vim.iter(string.gmatch('100 20 50', 'd+')):map(tonumber):totable() + -- { 100, 20, 50 } + + vim.iter({ 1, 2, 3 }):map(function(v) return v, 2 * v end):totable() + -- { { 1, 2 }, { 2, 4 }, { 3, 6 } } + + vim.iter({ a = 1, b = 2, c = 3 }):filter(function(k, v) return v % 2 ~= 0 end):totable() + -- { a = 1, c = 3 } +< + + Return: ~ + (table) + +new({src}, {...}) *new()* + Create a new Iter object from a table or iterator. + + Parameters: ~ + • {src} table|function Table or iterator to drain values from + + Return: ~ + Iter + +next() *next()* + TODO: Documentation + vim:tw=78:ts=8:sw=4:sts=4:et:ft=help:norl: diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt index dc541fcf37..bbf63a5ee0 100644 --- a/runtime/doc/news.txt +++ b/runtime/doc/news.txt @@ -35,7 +35,8 @@ ADDED FEATURES *news-added* The following new APIs or features were added. -• ... +• |vim.iter()| provides a generic iterator interface for tables and Lua +iterators |luaref-in|. ============================================================================== CHANGED FEATURES *news-changed* diff --git a/runtime/lua/vim/iter.lua b/runtime/lua/vim/iter.lua new file mode 100644 index 0000000000..b73c03ba9a --- /dev/null +++ b/runtime/lua/vim/iter.lua @@ -0,0 +1,836 @@ +--- Iterator implementation. + +---@class Iter +local Iter = {} +Iter.__index = Iter +Iter.__call = function(self) + return self:next() +end + +--- Special case implementations for iterators on list tables. +---@class ListIter : Iter +---@field _table table Underlying table data (table iterators only) +---@field _head number Index to the front of a table iterator (table iterators only) +---@field _tail number Index to the end of a table iterator (table iterators only) +local ListIter = {} +ListIter.__index = setmetatable(ListIter, Iter) +ListIter.__call = function(self) + return self:next() +end + +--- Special case implementations for iterators on non-list tables. +---@class TableIter : Iter +local TableIter = {} +TableIter.__index = setmetatable(TableIter, Iter) +TableIter.__call = function(self) + return self:next() +end + +---@private +local function unpack(t) + if type(t) == 'table' then + return _G.unpack(t) + end + return t +end + +---@private +local function pack(...) + if select('#', ...) > 1 then + return { ... } + end + return ... +end + +--- Add a filter step to the iterator pipeline. +--- +--- Example: +--- <pre>lua +--- local bufs = vim.iter(vim.api.nvim_list_bufs()):filter(vim.api.nvim_buf_is_loaded) +--- </pre> +--- +---@param f function(...):bool Takes all values returned from the previous stage in the pipeline and +--- returns false or nil if the current iterator element should be +--- removed. +---@return Iter +function Iter.filter(self, f) + ---@private + local function fn(...) + local result = nil + if select(1, ...) ~= nil then + if not f(...) then + return true, nil + else + result = pack(...) + end + end + return false, result + end + + local next = self.next + self.next = function(this) + local cont, result + repeat + cont, result = fn(next(this)) + until not cont + return unpack(result) + end + return self +end + +---@private +function ListIter.filter(self, f) + local inc = self._head < self._tail and 1 or -1 + local n = self._head + for i = self._head, self._tail - inc, inc do + local v = self._table[i] + if f(unpack(v)) then + self._table[n] = v + n = n + inc + end + end + self._tail = n + return self +end + +--- Add a map step to the iterator pipeline. +--- +--- If the map function returns nil, the value is filtered from the iterator. +--- +--- Example: +--- <pre>lua +--- local it = vim.iter({ 1, 2, 3, 4 }):map(function(v) +--- if v % 2 == 0 then +--- return v * 3 +--- end +--- end) +--- it:totable() +--- -- { 6, 12 } +--- </pre> +--- +---@param f function(...):any Mapping function. Takes all values returned from the previous stage +--- in the pipeline as arguments and returns one or more new values, +--- which are used in the next pipeline stage. Nil return values returned +--- are filtered from the output. +---@return Iter +function Iter.map(self, f) + ---@private + local function fn(...) + local result = nil + if select(1, ...) ~= nil then + result = pack(f(...)) + if result == nil then + return true, nil + end + end + return false, result + end + + local next = self.next + self.next = function(this) + local cont, result + repeat + cont, result = fn(next(this)) + until not cont + return unpack(result) + end + return self +end + +---@private +function ListIter.map(self, f) + local inc = self._head < self._tail and 1 or -1 + local n = self._head + for i = self._head, self._tail - inc, inc do + local v = pack(f(unpack(self._table[i]))) + if v ~= nil then + self._table[n] = v + n = n + inc + end + end + self._tail = n + return self +end + +--- Call a function once for each item in the pipeline. +--- +--- This is used for functions which have side effects. To modify the values in the iterator, use +--- |Iter:map()|. +--- +--- This function drains the iterator. +--- +---@param f function(...) Function to execute for each item in the pipeline. Takes all of the +--- values returned by the previous stage in the pipeline as arguments. +function Iter.each(self, f) + ---@private + local function fn(...) + if select(1, ...) ~= nil then + f(...) + return true + end + end + while fn(self:next()) do + end +end + +---@private +function ListIter.each(self, f) + local inc = self._head < self._tail and 1 or -1 + for i = self._head, self._tail - inc, inc do + f(unpack(self._table[i])) + end + self._head = self._tail +end + +--- Collect the iterator into a table. +--- +--- The resulting table depends on the initial source in the iterator pipeline. List-like tables +--- and function iterators will be collected into a list-like table. If multiple values are returned +--- from the final stage in the iterator pipeline, each value will be included in a table. If a +--- map-like table was used as the initial source, then a map-like table is returned. +--- +--- Examples: +--- <pre>lua +--- vim.iter(string.gmatch('100 20 50', '%d+')):map(tonumber):totable() +--- -- { 100, 20, 50 } +--- +--- vim.iter({ 1, 2, 3 }):map(function(v) return v, 2 * v end):totable() +--- -- { { 1, 2 }, { 2, 4 }, { 3, 6 } } +--- +--- vim.iter({ a = 1, b = 2, c = 3 }):filter(function(k, v) return v % 2 ~= 0 end):totable() +--- -- { a = 1, c = 3 } +--- </pre> +--- +---@return table +function Iter.totable(self) + local t = {} + + while true do + local args = pack(self:next()) + if args == nil then + break + end + t[#t + 1] = args + end + return t +end + +---@private +function ListIter.totable(self) + if self._head == 1 and self._tail == #self._table + 1 and self.next == ListIter.next then + return self._table + end + + return Iter.totable(self) +end + +---@private +function TableIter.totable(self) + local t = {} + for k, v in self do + t[k] = v + end + return t +end + +--- Fold an iterator or table into a single value. +--- +---@generic A +--- +---@param init A Initial value of the accumulator. +---@param f function(acc:A, ...):A Accumulation function. +---@return A +function Iter.fold(self, init, f) + local acc = init + + --- Use a closure to handle var args returned from iterator + ---@private + local function fn(...) + if select(1, ...) ~= nil then + acc = f(acc, ...) + return true + end + end + + while fn(self:next()) do + end + return acc +end + +---@private +function ListIter.fold(self, init, f) + local acc = init + local inc = self._head < self._tail and 1 or -1 + for i = self._head, self._tail - inc, inc do + acc = f(acc, unpack(self._table[i])) + end + return acc +end + +--- Return the next value from the iterator. +--- +--- Example: +--- <pre>lua +--- +--- local it = vim.iter(string.gmatch('1 2 3', '%d+')):map(tonumber) +--- it:next() +--- -- 1 +--- it:next() +--- -- 2 +--- it:next() +--- -- 3 +--- +--- </pre> +--- +---@return any +function Iter.next(self) -- luacheck: no unused args + -- This function is provided by the source iterator in Iter.new. This definition exists only for + -- the docstring +end + +---@private +function ListIter.next(self) + if self._head ~= self._tail then + local v = self._table[self._head] + local inc = self._head < self._tail and 1 or -1 + self._head = self._head + inc + return unpack(v) + end +end + +--- Reverse an iterator. +--- +--- Only supported for iterators on list-like tables. +--- +--- Example: +--- <pre>lua +--- +--- local it = vim.iter({ 3, 6, 9, 12 }):rev() +--- it:totable() +--- -- { 12, 9, 6, 3 } +--- +--- </pre> +--- +---@return Iter +function Iter.rev(self) + error('rev() requires a list-like table') + return self +end + +---@private +function ListIter.rev(self) + local inc = self._head < self._tail and 1 or -1 + self._head, self._tail = self._tail - inc, self._head - inc + return self +end + +--- Peek at the next value in the iterator without consuming it. +--- +--- Only supported for iterators on list-like tables. +--- +--- Example: +--- <pre>lua +--- +--- local it = vim.iter({ 3, 6, 9, 12 }) +--- it:peek() +--- -- 3 +--- it:peek() +--- -- 3 +--- it:next() +--- -- 3 +--- +--- </pre> +--- +---@return any +function Iter.peek(self) -- luacheck: no unused args + error('peek() requires a list-like table') +end + +---@private +function ListIter.peek(self) + if self._head ~= self._tail then + return self._table[self._head] + end +end + +--- Find the first value in the iterator that satisfies the given predicate. +--- +--- Advances the iterator. Returns nil and drains the iterator if no value is found. +--- +--- Examples: +--- <pre>lua +--- +--- local it = vim.iter({ 3, 6, 9, 12 }) +--- it:find(12) +--- -- 12 +--- +--- local it = vim.iter({ 3, 6, 9, 12 }) +--- it:find(20) +--- -- nil +--- +--- local it = vim.iter({ 3, 6, 9, 12 }) +--- it:find(function(v) return v % 4 == 0 end) +--- -- 12 +--- +--- </pre> +--- +---@return any +function Iter.find(self, f) + if type(f) ~= 'function' then + local val = f + f = function(v) + return v == val + end + end + + local result = nil + + --- Use a closure to handle var args returned from iterator + ---@private + local function fn(...) + if select(1, ...) ~= nil then + if f(...) then + result = pack(...) + else + return true + end + end + end + + while fn(self:next()) do + end + return unpack(result) +end + +--- Find the first value in the iterator that satisfies the given predicate, starting from the end. +--- +--- Advances the iterator. Returns nil and drains the iterator if no value is found. +--- +--- Only supported for iterators on list-like tables. +--- +--- Examples: +--- <pre>lua +--- +--- local it = vim.iter({ 1, 2, 3, 2, 1 }):enumerate() +--- it:rfind(1) +--- -- 5 1 +--- it:rfind(1) +--- -- 1 1 +--- +--- </pre> +--- +---@see Iter.find +--- +---@return any +function Iter.rfind(self, f) -- luacheck: no unused args + error('rfind() requires a list-like table') +end + +---@private +function ListIter.rfind(self, f) -- luacheck: no unused args + if type(f) ~= 'function' then + local val = f + f = function(v) + return v == val + end + end + + local inc = self._head < self._tail and 1 or -1 + for i = self._tail - inc, self._head, -inc do + local v = self._table[i] + if f(unpack(v)) then + self._tail = i + return unpack(v) + end + end + self._head = self._tail +end + +--- Return the next value from the end of the iterator. +--- +--- Only supported for iterators on list-like tables. +--- +--- Example: +--- <pre>lua +--- local it = vim.iter({1, 2, 3, 4}) +--- it:nextback() +--- -- 4 +--- it:nextback() +--- -- 3 +--- </pre> +--- +---@return any +function Iter.nextback(self) -- luacheck: no unused args + error('nextback() requires a list-like table') +end + +function ListIter.nextback(self) + if self._head ~= self._tail then + local inc = self._head < self._tail and 1 or -1 + self._tail = self._tail - inc + return self._table[self._tail] + end +end + +--- Return the next value from the end of the iterator without consuming it. +--- +--- Only supported for iterators on list-like tables. +--- +--- Example: +--- <pre>lua +--- local it = vim.iter({1, 2, 3, 4}) +--- it:peekback() +--- -- 4 +--- it:peekback() +--- -- 4 +--- it:nextback() +--- -- 4 +--- </pre> +--- +---@return any +function Iter.peekback(self) -- luacheck: no unused args + error('peekback() requires a list-like table') +end + +function ListIter.peekback(self) + if self._head ~= self._tail then + local inc = self._head < self._tail and 1 or -1 + return self._table[self._tail - inc] + end +end + +--- Skip values in the iterator. +--- +--- Example: +--- <pre>lua +--- +--- local it = vim.iter({ 3, 6, 9, 12 }):skip(2) +--- it:next() +--- -- 9 +--- +--- </pre> +--- +---@param n number Number of values to skip. +---@return Iter +function Iter.skip(self, n) + for _ = 1, n do + local _ = self:next() + end + return self +end + +---@private +function ListIter.skip(self, n) + local inc = self._head < self._tail and n or -n + self._head = self._head + inc + if (inc > 0 and self._head > self._tail) or (inc < 0 and self._head < self._tail) then + self._head = self._tail + end + return self +end + +--- Skip values in the iterator starting from the end. +--- +--- Only supported for iterators on list-like tables. +--- +--- Example: +--- <pre>lua +--- local it = vim.iter({ 1, 2, 3, 4, 5 }):skipback(2) +--- it:next() +--- -- 1 +--- it:nextback() +--- -- 3 +--- </pre> +--- +---@param n number Number of values to skip. +---@return Iter +function Iter.skipback(self, n) -- luacheck: no unused args + error('skipback() requires a list-like table') + return self +end + +---@private +function ListIter.skipback(self, n) + local inc = self._head < self._tail and n or -n + self._tail = self._tail - inc + if (inc > 0 and self._head > self._tail) or (inc < 0 and self._head < self._tail) then + self._head = self._tail + end + return self +end + +--- Return the nth value in the iterator. +--- +--- This function advances the iterator. +--- +--- Example: +--- <pre>lua +--- +--- local it = vim.iter({ 3, 6, 9, 12 }) +--- it:nth(2) +--- -- 6 +--- it:nth(2) +--- -- 12 +--- +--- </pre> +--- +---@param n number The index of the value to return. +---@return any +function Iter.nth(self, n) + if n > 0 then + return self:skip(n - 1):next() + end +end + +--- Return the nth value from the end of the iterator. +--- +--- This function advances the iterator. +--- +--- Only supported for iterators on list-like tables. +--- +--- Example: +--- <pre>lua +--- +--- local it = vim.iter({ 3, 6, 9, 12 }) +--- it:nthback(2) +--- -- 9 +--- it:nthback(2) +--- -- 3 +--- +--- </pre> +--- +---@param n number The index of the value to return. +---@return any +function Iter.nthback(self, n) + if n > 0 then + return self:skipback(n - 1):nextback() + end +end + +--- Slice an iterator, changing its start and end positions. +--- +--- This is equivalent to :skip(first - 1):skipback(len - last + 1) +--- +--- Only supported for iterators on list-like tables. +--- +---@param first number +---@param last number +---@return Iter +function Iter.slice(self, first, last) -- luacheck: no unused args + return self:skip(math.max(0, first - 1)):skipback(math.max(0, self._tail - last - 1)) +end + +--- Return true if any of the items in the iterator match the given predicate. +--- +---@param pred function(...):bool Predicate function. Takes all values returned from the previous +--- stage in the pipeline as arguments and returns true if the +--- predicate matches. +function Iter.any(self, pred) + local any = false + + --- Use a closure to handle var args returned from iterator + ---@private + local function fn(...) + if select(1, ...) ~= nil then + if pred(...) then + any = true + else + return true + end + end + end + + while fn(self:next()) do + end + return any +end + +--- Return true if all of the items in the iterator match the given predicate. +--- +---@param pred function(...):bool Predicate function. Takes all values returned from the previous +--- stage in the pipeline as arguments and returns true if the +--- predicate matches. +function Iter.all(self, pred) + local all = true + + ---@private + local function fn(...) + if select(1, ...) ~= nil then + if not pred(...) then + all = false + else + return true + end + end + end + + while fn(self:next()) do + end + return all +end + +--- Return the last item in the iterator. +--- +--- Drains the iterator. +--- +--- Example: +--- <pre>lua +--- +--- local it = vim.iter(vim.gsplit('abcdefg', '')) +--- it:last() +--- -- 'g' +--- +--- local it = vim.iter({ 3, 6, 9, 12, 15 }) +--- it:last() +--- -- 15 +--- +--- </pre> +--- +---@return any +function Iter.last(self) + local last = self:next() + local cur = self:next() + while cur do + last = cur + cur = self:next() + end + return last +end + +---@private +function ListIter.last(self) + local inc = self._head < self._tail and 1 or -1 + local v = self._table[self._tail - inc] + self._head = self._tail + return v +end + +--- Add an iterator stage that returns the current iterator count as well as the iterator value. +--- +--- For list tables, prefer +--- <pre>lua +--- vim.iter(ipairs(t)) +--- </pre> +--- +--- over +--- +--- <pre>lua +--- vim.iter(t):enumerate() +--- </pre> +--- +--- as the former is faster. +--- +--- Example: +--- <pre>lua +--- +--- local it = vim.iter(vim.gsplit('abc', '')):enumerate() +--- it:next() +--- -- 1 'a' +--- it:next() +--- -- 2 'b' +--- it:next() +--- -- 3 'c' +--- +--- </pre> +--- +---@return Iter +function Iter.enumerate(self) + local i = 0 + return self:map(function(...) + i = i + 1 + return i, ... + end) +end + +---@private +function ListIter.enumerate(self) + local inc = self._head < self._tail and 1 or -1 + for i = self._head, self._tail - inc, inc do + local v = self._table[i] + self._table[i] = { i, v } + end + return self +end + +--- Create a new Iter object from a table or iterator. +--- +---@param src table|function Table or iterator to drain values from +---@return Iter +function Iter.new(src, ...) + local it = {} + if type(src) == 'table' then + local t = {} + + -- Check if source table can be treated like a list (indices are consecutive integers + -- starting from 1) + local count = 0 + for _ in pairs(src) do + count = count + 1 + local v = src[count] + if v == nil then + return TableIter.new(src) + end + t[count] = v + end + return ListIter.new(t) + end + + if type(src) == 'function' then + local s, var = ... + + --- Use a closure to handle var args returned from iterator + ---@private + local function fn(...) + if select(1, ...) ~= nil then + var = select(1, ...) + return ... + end + end + + function it.next() + return fn(src(s, var)) + end + + setmetatable(it, Iter) + else + error('src must be a table or function') + end + return it +end + +--- Create a new ListIter +--- +---@param t table List-like table. Caller guarantees that this table is a valid list. +---@return Iter +---@private +function ListIter.new(t) + local it = {} + it._table = t + it._head = 1 + it._tail = #t + 1 + setmetatable(it, ListIter) + return it +end + +--- Create a new TableIter +--- +---@param t table Table to iterate over. For list-like tables, use ListIter.new instead. +---@return Iter +---@private +function TableIter.new(t) + local it = {} + + local index = nil + function it.next() + local k, v = next(t, index) + if k ~= nil then + index = k + return k, v + end + end + + setmetatable(it, TableIter) + return it +end + +return Iter diff --git a/runtime/lua/vim/shared.lua b/runtime/lua/vim/shared.lua index f700f4a6b3..1a96ef9bc4 100644 --- a/runtime/lua/vim/shared.lua +++ b/runtime/lua/vim/shared.lua @@ -884,4 +884,109 @@ function vim.defaulttable(create) }) end +--- Create an Iter |lua-iter| object from a table or iterator. +--- +--- The input value can be a table or a function iterator (see |luaref-in|). +--- +--- This function wraps the input value into an interface which allows chaining +--- multiple pipeline stages in an efficient manner. Each pipeline stage +--- receives as input the output values from the prior stage. The values used in +--- the first stage of the pipeline depend on the type passed to this function: +--- +--- - List tables pass only the value of each element +--- - Non-list tables pass both the key and value of each element +--- - Function iterators pass all of the values returned by their respective +--- function +--- +--- Examples: +--- <pre>lua +--- local it = vim.iter({ 1, 2, 3, 4, 5 }) +--- it:map(function(v) +--- return v * 3 +--- end) +--- it:rev() +--- it:skip(2) +--- it:totable() +--- -- { 9, 6, 3 } +--- +--- vim.iter(ipairs({ 1, 2, 3, 4, 5 })):map(function(i, v) +--- if i > 2 then return v end +--- end):totable() +--- -- { 3, 4, 5 } +--- +--- local it = vim.iter(vim.gsplit('1,2,3,4,5', ',')) +--- it:map(function(s) return tonumber(s) end) +--- for i, d in it:enumerate() do +--- print(string.format("Column %d is %d", i, d)) +--- end +--- -- Column 1 is 1 +--- -- Column 2 is 2 +--- -- Column 3 is 3 +--- -- Column 4 is 4 +--- -- Column 5 is 5 +--- +--- vim.iter({ a = 1, b = 2, c = 3, z = 26 }):any(function(k, v) +--- return k == 'z' +--- end) +--- -- true +--- </pre> +--- +---@see |lua-iter| +--- +---@param src table|function Table or iterator. +---@return Iter @|lua-iter| +function vim.iter(src, ...) + local Iter = require('vim.iter') + return Iter.new(src, ...) +end + +--- Collect an iterator into a table. +--- +--- This is a convenience function that performs: +--- <pre>lua +--- vim.iter(f):totable() +--- </pre> +--- +---@param f function Iterator function +---@return table +function vim.totable(f, ...) + return vim.iter(f, ...):totable() +end + +--- Filter a table or iterator. +--- +--- This is a convenience function that performs: +--- <pre>lua +--- vim.iter(src):filter(f):totable() +--- </pre> +--- +---@see |Iter:filter()| +--- +---@param f function(...):bool Filter function. Accepts the current iterator or table values as +--- arguments and returns true if those values should be kept in the +--- final table +---@param src table|function Table or iterator function to filter +---@return table +function vim.filter(f, src, ...) + return vim.iter(src, ...):filter(f):totable() +end + +--- Map and filter a table or iterator. +--- +--- This is a convenience function that performs: +--- <pre>lua +--- vim.iter(src):map(f):totable() +--- </pre> +--- +---@see |Iter:map()| +--- +---@param f function(...):?any Map function. Accepts the current iterator or table values as +--- arguments and returns one or more new values. Nil values are removed +--- from the final table. +---@param src table|function Table or iterator function to filter +---@return table +function vim.map(f, src, ...) + return vim.iter(src, ...):map(f):totable() +end + return vim diff --git a/scripts/gen_vimdoc.py b/scripts/gen_vimdoc.py index 8860bdd5c6..52d03d9746 100755 --- a/scripts/gen_vimdoc.py +++ b/scripts/gen_vimdoc.py @@ -154,8 +154,10 @@ CONFIG = { 'fs.lua', 'secure.lua', 'version.lua', + 'iter.lua', ], 'files': [ + 'runtime/lua/vim/iter.lua', 'runtime/lua/vim/_editor.lua', 'runtime/lua/vim/shared.lua', 'runtime/lua/vim/loader.lua', @@ -185,6 +187,8 @@ CONFIG = { 'fn_helptag_fmt': lambda fstem, name: ( f'*vim.{name}()*' if fstem.lower() == '_editor' + else f'*{name}()*' + if fstem in ('iter.lua') else f'*{fstem}.{name}()*'), 'module_override': { # `shared` functions are exposed on the `vim` module. diff --git a/test/functional/lua/vim_spec.lua b/test/functional/lua/vim_spec.lua index 1ee1a13fd5..e5caf6f6f7 100644 --- a/test/functional/lua/vim_spec.lua +++ b/test/functional/lua/vim_spec.lua @@ -3029,6 +3029,360 @@ describe('lua stdlib', function() eq(false, if_nil(d, c)) eq(NIL, if_nil(a)) end) + + describe('vim.iter', function() + it('filter()', function() + local function odd(v) + return v % 2 ~= 0 + end + + local t = { 1, 2, 3, 4, 5 } + eq({ 1, 3, 5 }, vim.iter(t):filter(odd):totable()) + eq({ 2, 4 }, vim.iter(t):filter(function(v) return not odd(v) end):totable()) + eq({}, vim.iter(t):filter(function(v) if v > 5 then return v end end):totable()) + + do + local it = vim.iter(ipairs(t)) + it:filter(function(i, v) return i > 1 and v < 5 end) + it:map(function(_, v) return v * 2 end) + eq({ 4, 6, 8 }, it:totable()) + end + + local it = vim.iter(string.gmatch('the quick brown fox', '%w+')) + eq({'the', 'fox'}, it:filter(function(s) return #s <= 3 end):totable()) + end) + + it('map()', function() + local t = { 1, 2, 3, 4, 5 } + eq( + { 2, 4, 6, 8, 10 }, + vim + .iter(t) + :map(function(v) + return 2 * v + end) + :totable() + ) + + local it = vim.gsplit( + [[ + Line 1 + Line 2 + Line 3 + Line 4 + ]], + '\n' + ) + + eq( + { 'Lion 2', 'Lion 4' }, + vim + .iter(it) + :map(function(s) + local lnum = s:match('(%d+)') + if lnum and tonumber(lnum) % 2 == 0 then + return vim.trim(s:gsub('Line', 'Lion')) + end + end) + :totable() + ) + end) + + it('for loops', function() + local t = {1, 2, 3, 4, 5} + local acc = 0 + for v in vim.iter(t):map(function(v) return v * 3 end) do + acc = acc + v + end + eq(45, acc) + end) + + it('totable()', function() + do + local it = vim.iter({1, 2, 3}):map(function(v) return v, v*v end) + eq({{1, 1}, {2, 4}, {3, 9}}, it:totable()) + end + + do + local it = vim.iter(string.gmatch('1,4,lol,17,blah,2,9,3', '%d+')):map(tonumber) + eq({1, 4, 17, 2, 9, 3}, it:totable()) + end + end) + + it('next()', function() + local it = vim.iter({1, 2, 3}):map(function(v) return 2 * v end) + eq(2, it:next()) + eq(4, it:next()) + eq(6, it:next()) + eq(nil, it:next()) + end) + + it('rev()', function() + eq({3, 2, 1}, vim.iter({1, 2, 3}):rev():totable()) + + local it = vim.iter(string.gmatch("abc", "%w")) + matches('rev%(%) requires a list%-like table', pcall_err(it.rev, it)) + end) + + it('skip()', function() + do + local t = {4, 3, 2, 1} + eq(t, vim.iter(t):skip(0):totable()) + eq({3, 2, 1}, vim.iter(t):skip(1):totable()) + eq({2, 1}, vim.iter(t):skip(2):totable()) + eq({1}, vim.iter(t):skip(#t - 1):totable()) + eq({}, vim.iter(t):skip(#t):totable()) + eq({}, vim.iter(t):skip(#t + 1):totable()) + end + + do + local function skip(n) + return vim.iter(vim.gsplit('a|b|c|d', '|')):skip(n):totable() + end + eq({'a', 'b', 'c', 'd'}, skip(0)) + eq({'b', 'c', 'd'}, skip(1)) + eq({'c', 'd'}, skip(2)) + eq({'d'}, skip(3)) + eq({}, skip(4)) + eq({}, skip(5)) + end + end) + + it('skipback()', function() + do + local t = {4, 3, 2, 1} + eq(t, vim.iter(t):skipback(0):totable()) + eq({4, 3, 2}, vim.iter(t):skipback(1):totable()) + eq({4, 3}, vim.iter(t):skipback(2):totable()) + eq({4}, vim.iter(t):skipback(#t - 1):totable()) + eq({}, vim.iter(t):skipback(#t):totable()) + eq({}, vim.iter(t):skipback(#t + 1):totable()) + end + + local it = vim.iter(vim.gsplit('a|b|c|d', '|')) + matches('skipback%(%) requires a list%-like table', pcall_err(it.skipback, it, 0)) + end) + + it('slice()', function() + local t = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10} + eq({3, 4, 5, 6, 7}, vim.iter(t):slice(3, 7):totable()) + eq({}, vim.iter(t):slice(6, 5):totable()) + eq({}, vim.iter(t):slice(0, 0):totable()) + eq({1}, vim.iter(t):slice(1, 1):totable()) + eq({1, 2}, vim.iter(t):slice(1, 2):totable()) + eq({10}, vim.iter(t):slice(10, 10):totable()) + eq({8, 9, 10}, vim.iter(t):slice(8, 11):totable()) + end) + + it('nth()', function() + do + local t = {4, 3, 2, 1} + eq(nil, vim.iter(t):nth(0)) + eq(4, vim.iter(t):nth(1)) + eq(3, vim.iter(t):nth(2)) + eq(2, vim.iter(t):nth(3)) + eq(1, vim.iter(t):nth(4)) + eq(nil, vim.iter(t):nth(5)) + end + + do + local function nth(n) + return vim.iter(vim.gsplit('a|b|c|d', '|')):nth(n) + end + eq(nil, nth(0)) + eq('a', nth(1)) + eq('b', nth(2)) + eq('c', nth(3)) + eq('d', nth(4)) + eq(nil, nth(5)) + end + end) + + it('nthback()', function() + do + local t = {4, 3, 2, 1} + eq(nil, vim.iter(t):nthback(0)) + eq(1, vim.iter(t):nthback(1)) + eq(2, vim.iter(t):nthback(2)) + eq(3, vim.iter(t):nthback(3)) + eq(4, vim.iter(t):nthback(4)) + eq(nil, vim.iter(t):nthback(5)) + end + + local it = vim.iter(vim.gsplit('a|b|c|d', '|')) + matches('skipback%(%) requires a list%-like table', pcall_err(it.nthback, it, 1)) + end) + + it('any()', function() + local function odd(v) + return v % 2 ~= 0 + end + + do + local t = { 4, 8, 9, 10 } + eq(true, vim.iter(t):any(odd)) + end + + do + local t = { 4, 8, 10 } + eq(false, vim.iter(t):any(odd)) + end + + do + eq(true, vim.iter(vim.gsplit('a|b|c|d', '|')):any(function(s) return s == 'd' end)) + eq(false, vim.iter(vim.gsplit('a|b|c|d', '|')):any(function(s) return s == 'e' end)) + end + end) + + it('all()', function() + local function odd(v) + return v % 2 ~= 0 + end + + do + local t = { 3, 5, 7, 9 } + eq(true, vim.iter(t):all(odd)) + end + + do + local t = { 3, 5, 7, 10 } + eq(false, vim.iter(t):all(odd)) + end + + do + eq(true, vim.iter(vim.gsplit('a|a|a|a', '|')):all(function(s) return s == 'a' end)) + eq(false, vim.iter(vim.gsplit('a|a|a|b', '|')):all(function(s) return s == 'a' end)) + end + end) + + it('last()', function() + local s = 'abcdefghijklmnopqrstuvwxyz' + eq('z', vim.iter(vim.split(s, '')):last()) + eq('z', vim.iter(vim.gsplit(s, '')):last()) + end) + + it('enumerate()', function() + local it = vim.iter(vim.gsplit('abc', '')):enumerate() + eq({1, 'a'}, {it:next()}) + eq({2, 'b'}, {it:next()}) + eq({3, 'c'}, {it:next()}) + eq({}, {it:next()}) + end) + + it('peek()', function() + do + local it = vim.iter({ 3, 6, 9, 12 }) + eq(3, it:peek()) + eq(3, it:peek()) + eq(3, it:next()) + end + + do + local it = vim.iter(vim.gsplit('hi', '')) + matches('peek%(%) requires a list%-like table', pcall_err(it.peek, it)) + end + end) + + it('find()', function() + local t = {3, 6, 9, 12} + eq(12, vim.iter(t):find(12)) + eq(nil, vim.iter(t):find(15)) + eq(12, vim.iter(t):find(function(v) return v % 4 == 0 end)) + + do + local it = vim.iter(t) + local pred = function(v) return v % 3 == 0 end + eq(3, it:find(pred)) + eq(6, it:find(pred)) + eq(9, it:find(pred)) + eq(12, it:find(pred)) + eq(nil, it:find(pred)) + end + + do + local it = vim.iter(vim.gsplit('AbCdE', '')) + local pred = function(s) return s:match('[A-Z]') end + eq('A', it:find(pred)) + eq('C', it:find(pred)) + eq('E', it:find(pred)) + eq(nil, it:find(pred)) + end + end) + + it('rfind()', function() + local t = {1, 2, 3, 2, 1} + do + local it = vim.iter(t) + eq(1, it:rfind(1)) + eq(1, it:rfind(1)) + eq(nil, it:rfind(1)) + end + + do + local it = vim.iter(t):enumerate() + local pred = function(i) return i % 2 ~= 0 end + eq({5, 1}, {it:rfind(pred)}) + eq({3, 3}, {it:rfind(pred)}) + eq({1, 1}, {it:rfind(pred)}) + eq(nil, it:rfind(pred)) + end + + do + local it = vim.iter(vim.gsplit('AbCdE', '')) + matches('rfind%(%) requires a list%-like table', pcall_err(it.rfind, it, 'E')) + end + end) + + it('nextback()', function() + do + local it = vim.iter({ 1, 2, 3, 4 }) + eq(4, it:nextback()) + eq(3, it:nextback()) + eq(2, it:nextback()) + eq(1, it:nextback()) + eq(nil, it:nextback()) + eq(nil, it:nextback()) + end + + do + local it = vim.iter(vim.gsplit('hi', '')) + matches('nextback%(%) requires a list%-like table', pcall_err(it.nextback, it)) + end + end) + + it('peekback()', function() + do + local it = vim.iter({ 1, 2, 3, 4 }) + eq(4, it:peekback()) + eq(4, it:peekback()) + eq(4, it:peekback()) + end + + do + local it = vim.iter(vim.gsplit('hi', '')) + matches('peekback%(%) requires a list%-like table', pcall_err(it.peekback, it)) + end + end) + + it('fold()', function() + local t = {1, 2, 3, 4, 5} + eq(115, vim.iter(t):fold(100, function(acc, v) return acc + v end)) + eq({5, 4, 3, 2, 1}, vim.iter(t):fold({}, function(acc, v) + table.insert(acc, 1, v) + return acc + end)) + end) + + it('handles map-like tables', function() + local t = { a = 1, b = 2, c = 3 } + local it = vim.iter(t):map(function(k, v) + if v % 2 ~= 0 then + return k:upper(), v * 2 + end + end) + eq({ A = 2, C = 6 }, it:totable()) + end) + end) end) describe('lua: builtin modules', function() |