From d4f1e673aeb80884107aa3f66e8c69ecba187c96 Mon Sep 17 00:00:00 2001
From: Kevin Modzelewski <kevmod@gmail.com>
Date: Tue, 7 Jan 2025 20:06:33 -0500
Subject: [PATCH] Microoptimize the walk() function

I noticed that it takes pynvim about 4ms to attach to an nvim instance for me,
and 3ms of that is due to the single line:
  metadata = walk(decode_if_bytes, metadata)

This commit reduces the walk() time down to 1.5ms, which brings the total
attach time down to 2.5ms. This is helpful for me because in my use case I end
up connecting to all of the currently-running nvim processes and this starts to
take a noticeable amount of time. Unfortunately parallelization does not help
here due to the nature of the slowness.

walk() is expensive because it does a very large amount of pure-python
manipulation, so this commit is just some tweaks to reduce the overheads:

- *args and **kw make the function call slow, and we can avoid needing them by
  pre-packing the args into fn via functools.partial
- The comprehensions can be written to directly construct the objects rather
  than create a generator which is passed to a constructor
- The typechecking is microoptimized by calling type() once and unrolling the
  `type_ in [list, tuple]` check

I did notice that in my setup the metadata contains no byte objects, so the
entire call is a noop. I'm not sure if that is something that could be relied
on or detected, which could be an even bigger speedup.
---
 pynvim/api/common.py  | 20 ++++++++++++--------
 pynvim/api/nvim.py    |  2 +-
 pynvim/plugin/host.py |  2 +-
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/pynvim/api/common.py b/pynvim/api/common.py
index 0e109a10..6a291e66 100644
--- a/pynvim/api/common.py
+++ b/pynvim/api/common.py
@@ -241,11 +241,15 @@ def decode_if_bytes(obj: T, mode: TDecodeMode = True) -> Union[T, str]:
     return obj
 
 
-def walk(fn: Callable[..., Any], obj: Any, *args: Any, **kwargs: Any) -> Any:
-    """Recursively walk an object graph applying `fn`/`args` to objects."""
-    if type(obj) in [list, tuple]:
-        return list(walk(fn, o, *args) for o in obj)
-    if type(obj) is dict:
-        return dict((walk(fn, k, *args), walk(fn, v, *args)) for k, v in
-                    obj.items())
-    return fn(obj, *args, **kwargs)
+def walk(fn: Callable[[Any], Any], obj: Any) -> Any:
+    """Recursively walk an object graph applying `fn` to objects."""
+
+    # Note: this function is very hot, so it is worth being careful
+    # about performance.
+    type_ = type(obj)
+
+    if type_ is list or type_ is tuple:
+        return [walk(fn, o) for o in obj]
+    if type_ is dict:
+        return {walk(fn, k): walk(fn, v) for k, v in obj.items()}
+    return fn(obj)
diff --git a/pynvim/api/nvim.py b/pynvim/api/nvim.py
index f0c33fdc..a9cc3d44 100644
--- a/pynvim/api/nvim.py
+++ b/pynvim/api/nvim.py
@@ -209,7 +209,7 @@ def request(self, name: str, *args: Any, **kwargs: Any) -> Any:
         decode = kwargs.pop('decode', self._decode)
         args = walk(self._to_nvim, args)
         res = self._session.request(name, *args, **kwargs)
-        return walk(self._from_nvim, res, decode=decode)
+        return walk(partial(self._from_nvim, decode=decode), res)
 
     def next_message(self) -> Any:
         """Block until a message(request or notification) is available.
diff --git a/pynvim/plugin/host.py b/pynvim/plugin/host.py
index 4a5a209b..b2f8bced 100644
--- a/pynvim/plugin/host.py
+++ b/pynvim/plugin/host.py
@@ -112,7 +112,7 @@ def _wrap_delayed_function(self, cls, delayed_handlers, name, sync,
 
     def _wrap_function(self, fn, sync, decode, nvim_bind, name, *args):
         if decode:
-            args = walk(decode_if_bytes, args, decode)
+            args = walk(partial(decode_if_bytes, mode=decode), args)
         if nvim_bind is not None:
             args.insert(0, nvim_bind)
         try: