feat(benchmark): add util to benchmark functions (#228)

* feat(benchmark): init

* update interface and output

* make it work with a single function

* apply conni patch and update interface

co-authored-by: Conni2461 <[email protected]>

* enhance

co-authored-by: Conni2461 <[email protected]>

Co-authored-by: Conni2461 <[email protected]>

Author: kkharji

lua/plenary/benchmark/init.lua

@@ -0,0 +1,126 @@
+local B = {}
+local stat = require "plenary.benchmark.stat"
+
+local get_stats = function(results)
+  local ret = {}
+
+  ret.max, ret.min = stat.maxmin(results)
+  ret.mean = stat.mean(results)
+  ret.median = stat.median(results)
+  ret.std = stat.std_dev(results)
+
+  return ret
+end
+
+local get_output = function(index, res, runs)
+  -- divine with a sutable one / 1e3, 1e6, 1e9
+  local time_types = { "ns", "μs", "ms" }
+
+  local get_leading = function(time)
+    time = math.floor(time)
+    local count = 0
+    repeat
+      time = math.floor(time / 10)
+      count = count + 1
+    until time <= 0
+    return count
+  end
+
+  local get_best_fmt = function(time)
+    for _, v in ipairs(time_types) do
+      if math.abs(time) < 1000.0 then
+        return string.format("%s%3.1f %s", string.rep(" ", 3 - get_leading(time)), time, v)
+      end
+      time = time / 1000.0
+    end
+    return string.format("%.1f %s", time, "s")
+  end
+
+  return string.format(
+    "Benchmark #%d: '%s'\n  Time(mean ± σ):    %s ± %s\n  Range(min … max):  %s … %s  %d runs\n",
+    index,
+    res.name,
+    get_best_fmt(res.stats.mean),
+    get_best_fmt(res.stats.std),
+    get_best_fmt(res.stats.min),
+    get_best_fmt(res.stats.max),
+    runs
+  )
+end
+
+local get_summary = function(res)
+  if #res == 1 then
+    return ""
+  end
+
+  local fastest_mean = math.huge
+  local fastest_index = 1
+  for i, benchmark in ipairs(res) do
+    if benchmark.stats.mean < fastest_mean then
+      fastest_mean = benchmark.stats.mean
+      fastest_index = i
+    end
+  end
+
+  if fastest_mean == math.huge then
+    return ""
+  end
+
+  local output = {}
+  local fastest = res[fastest_index].stats
+  for i, benchmark in ipairs(res) do
+    if i ~= fastest_index then
+      local result = benchmark.stats
+      local ratio = result.mean / fastest.mean
+
+      -- // https://en.wikipedia.org/wiki/Propagation_of_uncertainty#Example_formulas
+      -- // Covariance asssumed to be 0, i.e. variables are assumed to be independent
+      local ratio_std = ratio
+        * math.sqrt(math.pow(result.std / result.mean, 2) + math.pow(fastest.std / fastest.mean, 2))
+
+      table.insert(output, string.format("  %.1f ± %.1f times faster than '%s'\n", ratio, ratio_std, benchmark.name))
+    end
+  end
+
+  return string.format("Summary\n  '%s' ran\n%s", res[fastest_index].name, table.concat(output, ""))
+end
+
+---@class benchmark_run_opts
+---@field warmup number @number of initial runs before starting to track time.
+---@field runs number @number of runs to make
+---@field fun table<array<string, function>> @functions to execute
+
+---Benchmark a function
+---@param name string @benchmark name
+---@param opts benchmark_run_opts
+local bench = function(name, opts)
+  vim.validate {
+    opts = { opts, "table" },
+    fun = { opts.fun, "table" },
+  }
+  opts.warmup = vim.F.if_nil(opts.warmup, 3)
+  opts.runs = vim.F.if_nil(opts.runs, 5)
+
+  opts.fun = type(opts.fun) == "function" and { opts.fun } or opts.fun
+  local output = { string.format("Benchmark Group: '%s' -----------------------\n", name) }
+  local res = {}
+  for i, fun in ipairs(opts.fun) do
+    res[i] = { name = fun[1], results = {} }
+    for _ = 1, opts.warmup do
+      fun[2]()
+    end
+    for j = 1, opts.runs do
+      local start = vim.loop.hrtime()
+      fun[2]()
+      res[i].results[j] = vim.loop.hrtime() - start
+    end
+    res[i].stats = get_stats(res[i].results)
+    table.insert(output, get_output(i, res[i], opts.runs))
+  end
+
+  print(string.format("%s\n%s", table.concat(output, ""), get_summary(res)))
+
+  return res
+end
+
+return bench

lua/plenary/benchmark/stat.lua

@@ -0,0 +1,86 @@
+local stat = {}
+
+---Calculate mean
+---@param t number[] @double
+---@return number @double
+stat.mean = function(t)
+  local sum = 0
+  local count = 0
+
+  for _, v in pairs(t) do
+    if type(v) == "number" then
+      sum = sum + v
+      count = count + 1
+    end
+  end
+
+  return (sum / count)
+end
+
+-- Get the median of a table.
+---@param t number[]
+---@return number
+stat.median = function(t)
+  local temp = {}
+
+  -- deep copy table so that when we sort it, the original is unchanged
+  -- also weed out any non numbers
+  for _, v in pairs(t) do
+    if type(v) == "number" then
+      table.insert(temp, v)
+    end
+  end
+
+  table.sort(temp)
+
+  -- If we have an even number of table elements or odd.
+  if math.fmod(#temp, 2) == 0 then
+    -- return mean value of middle two elements
+    return (temp[#temp / 2] + temp[(#temp / 2) + 1]) / 2
+  else
+    -- return middle element
+    return temp[math.ceil(#temp / 2)]
+  end
+end
+
+--- Get the standard deviation of a table
+---@param t number[]
+stat.std_dev = function(t)
+  local m, vm, result
+  local sum = 0
+  local count = 0
+
+  m = stat.mean(t)
+
+  for _, v in pairs(t) do
+    if type(v) == "number" then
+      vm = v - m
+      sum = sum + (vm * vm)
+      count = count + 1
+    end
+  end
+
+  result = math.sqrt(sum / (count - 1))
+
+  return result
+end
+
+---Get the max and min for a table
+---@param t number[]
+---@return number
+---@return number
+stat.maxmin = function(t)
+  local max = -math.huge
+  local min = math.huge
+
+  for _, v in pairs(t) do
+    if type(v) == "number" then
+      max = math.max(max, v)
+      min = math.min(min, v)
+    end
+  end
+
+  return max, min
+end
+
+return stat

lua/plenary/test_harness.lua

@@ -61,6 +61,10 @@ function harness.test_directory(directory, opts)
     if res.border_win_id then
       vim.api.nvim_win_set_option(res.border_win_id, "winhl", "Normal:Normal")
     end
+
+    if res.bufnr then
+      vim.api.nvim_buf_set_option(res.bufnr, "filetype", "PlenaryTestPopup")
+    end
     vim.cmd "mode"
   end