perf(iter): reduce number of table allocations

Packing and unpacking return values impairs performance considerably.
In an attempt to avoid creating tables as much as possible we can
instead pass return values between functions (which does not require
knowing the number of values a function might return). This makes the
code more complex, but improves benchmark numbers non-trivially.
This commit is contained in:
Gregory Anders 2023-04-28 11:37:53 -06:00
parent 5a0250c9a7
commit ef1801cc7c
2 changed files with 105 additions and 65 deletions

View File

@ -3080,7 +3080,7 @@ Iter:map({self}, {f}) *Iter:map()*
• {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.
stage. Nil return values are filtered from the output.
Return: ~
Iter

View File

@ -1,13 +1,14 @@
---@defgroup lua-iter
---
--- The \*vim.iter\* module provides a generic "iterator" interface over tables and iterator
--- functions.
--- The \*vim.iter\* module provides a generic "iterator" interface over tables
--- and iterator functions.
---
--- \*vim.iter()\* wraps its table or function argument into an \*Iter\* object with methods (such
--- as |Iter:filter()| and |Iter:map()|) that transform the underlying source data. These methods
--- can be chained together to create iterator "pipelines". 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:
--- \*vim.iter()\* wraps its table or function argument into an \*Iter\* object
--- with methods (such as |Iter:filter()| and |Iter:map()|) that transform the
--- underlying source data. These methods can be chained together to create
--- iterator "pipelines". 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
@ -47,8 +48,8 @@
--- -- true
--- </pre>
---
--- In addition to the |vim.iter()| function, the |vim.iter| module provides convenience functions
--- like |vim.iter.filter()| and |vim.iter.totable()|.
--- In addition to the |vim.iter()| function, the |vim.iter| module provides
--- convenience functions like |vim.iter.filter()| and |vim.iter.totable()|.
local M = {}
@ -61,9 +62,9 @@ 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)
---@field _table table Underlying table data
---@field _head number Index to the front of a table iterator
---@field _tail number Index to the end of a table iterator
local ListIter = {}
ListIter.__index = setmetatable(ListIter, Iter)
ListIter.__call = function(self)
@ -75,7 +76,7 @@ local packedmt = {}
---@private
local function unpack(t)
if getmetatable(t) == packedmt then
if type(t) == 'table' and getmetatable(t) == packedmt then
return _G.unpack(t, 1, t.n)
end
return t
@ -92,13 +93,47 @@ end
---@private
local function sanitize(t)
if getmetatable(t) == packedmt then
if type(t) == 'table' and getmetatable(t) == packedmt then
-- Remove length tag
t.n = nil
end
return t
end
--- Determine if the current iterator stage should continue.
---
--- If any arguments are passed to this function, then return those arguments
--- and stop the current iterator stage. Otherwise, return true to signal that
--- the current stage should continue.
---
---@param ... any Function arguments.
---@return boolean True if the iterator stage should continue, false otherwise
---@return any Function arguments.
---@private
local function continue(...)
if select('#', ...) > 0 then
return false, ...
end
return true
end
--- If no input arguments are given return false, indicating the current
--- iterator stage should stop. Otherwise, apply the arguments to the function
--- f. If that function returns no values, the current iterator stage continues.
--- Otherwise, those values are returned.
---
---@param f function Function to call with the given arguments
---@param ... any Arguments to apply to f
---@return boolean True if the iterator pipeline should continue, false otherwise
---@return any Return values of f
---@private
local function apply(f, ...)
if select('#', ...) > 0 then
return continue(f(...))
end
return false
end
--- Add a filter step to the iterator pipeline.
---
--- Example:
@ -106,33 +141,16 @@ end
--- 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.
---@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(...)
return self:map(function(...)
if f(...) then
return ...
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)
end
---@private
@ -165,31 +183,52 @@ end
--- -- { 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
---@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
--- 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
-- Implementation note: the reader may be forgiven for observing that this
-- function appears excessively convoluted. The problem to solve is that each
-- stage of the iterator pipeline can return any number of values, and the
-- number of values could even change per iteration. And the return values
-- must be checked to determine if the pipeline has ended, so we cannot
-- naively forward them along to the next stage.
--
-- A simple approach is to pack all of the return values into a table, check
-- for nil, then unpack the table for the next stage. However, packing and
-- unpacking tables is quite slow. There is no other way in Lua to handle an
-- unknown number of function return values than to simply forward those
-- values along to another function. Hence the intricate function passing you
-- see here.
local next = self.next
self.next = function(this)
local cont, result
repeat
cont, result = fn(next(this))
until not cont
return unpack(result)
--- Drain values from the upstream iterator source until a value can be
--- returned.
---
--- This is a recursive function. The base case is when the first argument is
--- false, which indicates that the rest of the arguments should be returned
--- as the values for the current iteration stage.
---
---@param cont boolean If true, the current iterator stage should continue to
--- pull values from its upstream pipeline stage.
--- Otherwise, this stage is complete and returns the
--- values passed.
---@param ... any Values to return if cont is false.
---@return any
---@private
local function fn(cont, ...)
if cont then
return fn(apply(f, next(self)))
end
return ...
end
self.next = function()
return fn(apply(f, next(self)))
end
return self
end
@ -211,17 +250,18 @@ 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 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.
---@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
if select('#', ...) > 0 then
f(...)
return true
end