nvim_cmp_buffer_mem_usage_benchmark.py

#!/usr/bin/env python3
# type: ignore

import matplotlib.pyplot as plt
import pynvim

CHUNK_SIZES = [
  1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75, 80, 85, 90,
  95, 100, 150, 200, 250, 300, 350, 400, 450, 500, 550, 600, 650, 700, 750, 800, 850, 900, 950,
  1000, 1250, 1500, 1750, 2000, 2250, 2500, 2750, 3000, 3250, 3500, 3750, 4000, 4250, 4500, 4750,
  5000, 5250, 5500, 5750, 6000, 6250, 6500, 6750, 7000, 7250, 7500, 7750, 8000, 8250, 8500, 8750,
  9000, 9250, 9500, 9750, 10000, 15000, 20000, 25000, 30000, 35000, 40000, 45000, 50000, 55000,
  60000, 65000, 70000, 75000, 80000, 85000, 90000, 95000, 100000, 150000, 200000
]

results = []
vim_version_str = None
for i, chunk_size in enumerate(CHUNK_SIZES):
  print(f"{i+1}/{len(CHUNK_SIZES)} {chunk_size}")
  with pynvim.attach(
    'child',
    argv=[
      "/usr/bin/env", "nvim", "--embed", "--headless", "-u", "NONE", "--cmd",
      "set rtp+=~/.config/nvim/plugged/cmp-buffer", "-n", "-i", "NONE",
      "/home/dmitmel/crosscode/src.js"
    ]
  ) as nvim:
    if vim_version_str is None:
      vim_version_str = "NVIM v{0.major}.{0.minor}.{0.patch} {1}".format(
        nvim.version, nvim.exec_lua("return jit and jit.version or _VERSION")
      )
    result = nvim.exec_lua(
      """
      local indexing_chunk_size = ...
      local buffer = require('cmp_buffer.buffer')
      local start_time = vim.loop.hrtime()
      local start_mem = vim.loop.resident_set_memory()
      local bufnr = vim.api.nvim_get_current_buf()
      local buf = buffer.new(bufnr, {
        keyword_length = 3,
        keyword_pattern = [[\k\+]],
        indexing_interval = -1,
        indexing_chunk_size = indexing_chunk_size,
      })
      buf:index()
      local end_mem = vim.loop.resident_set_memory()
      local end_time = vim.loop.hrtime()
      return {
        chunk_size = indexing_chunk_size,
        time = (end_time - start_time) / 1e9,
        mem = (end_mem - start_mem) / 1024 / 1024
      }
      """, chunk_size
    )
    print(result)
    results.append(result)

fig, (ax1, ax2) = plt.subplots(
  2,
  sharex=True,
  subplot_kw={
    "xscale": "log",
    "xlabel": "chunk_size",
  },
)
fig.suptitle(vim_version_str)
ax1.set_ylabel("time")
ax2.set_ylabel("mem")
ax1.plot([x["chunk_size"] for x in results], [x["time"] for x in results])
ax2.plot([x["chunk_size"] for x in results], [x["mem"] for x in results])
plt.show()