Speed up parsing of file transfer OSC codes

Avoid roundtrips between unicode and bytes objects

kittens/transfer/rsync.c

@@ -76,8 +76,6 @@ begin_create_signature(PyObject *self UNUSED, PyObject *args) {
     if (!job) { PyErr_SetString(PyExc_TypeError, "Not a job capsule"); return NULL; } \
 
 
-#define FREE_BUFFER_AFTER_FUNCTION __attribute__((cleanup(PyBuffer_Release)))
-
 static PyObject*
 iter_job(PyObject *self UNUSED, PyObject *args) {
     FREE_BUFFER_AFTER_FUNCTION Py_buffer input_buf = {0};
@@ -173,6 +171,61 @@ begin_patch(PyObject *self UNUSED, PyObject *callback) {
     return job_capsule;
 }
 
+static bool
+call_ftc_callback(PyObject *callback, char *src, Py_ssize_t key_start, Py_ssize_t key_length, Py_ssize_t val_start, Py_ssize_t val_length, PyObject *has_semicolons) {
+    DECREF_AFTER_FUNCTION PyObject *k = PyMemoryView_FromMemory(src + key_start, key_length, PyBUF_READ);
+    if (!k) return false;
+    DECREF_AFTER_FUNCTION PyObject *v = PyMemoryView_FromMemory(src + val_start, val_length, PyBUF_READ);
+    if (!v) return false;
+    DECREF_AFTER_FUNCTION PyObject *ret = PyObject_CallFunctionObjArgs(callback, k, v, has_semicolons, NULL);
+    return ret != NULL;
+}
+
+static PyObject*
+decode_utf8_buffer(PyObject *self UNUSED, PyObject *args) {
+    FREE_BUFFER_AFTER_FUNCTION Py_buffer buf = {0};
+    if (!PyArg_ParseTuple(args, "s*", &buf)) return NULL;
+    return PyUnicode_FromStringAndSize(buf.buf, buf.len);
+}
+
+static PyObject*
+parse_ftc(PyObject *self UNUSED, PyObject *args) {
+    FREE_BUFFER_AFTER_FUNCTION Py_buffer buf = {0};
+    PyObject *callback;
+    size_t i = 0, key_start = 0, key_length = 0, val_start = 0, val_length = 0;
+    if (!PyArg_ParseTuple(args, "s*O", &buf, &callback)) return NULL;
+    char *src = buf.buf;
+    size_t sz = buf.len;
+    if (!PyCallable_Check(callback)) { PyErr_SetString(PyExc_TypeError, "callback must be callable"); return NULL; }
+    PyObject *has_semicolons = Py_False;
+    for (i = 0; i < sz; i++) {
+        char ch = src[i];
+        if (key_length == 0) {
+            if (ch == '=') {
+                key_length = i - key_start;
+                val_start = i + 1;
+                has_semicolons = Py_False;
+            }
+        } else {
+            if (ch == ';') {
+                if (i + 1 < sz && src[i + 1] == ';') {
+                    has_semicolons = Py_True;
+                    i++;
+                } else {
+                    val_length = i - val_start;
+                    if (!call_ftc_callback(callback, src, key_start, key_length, val_start, val_length, has_semicolons)) return NULL;
+                    key_length = 0; key_start = i + 1; val_start = 0;
+                }
+            }
+        }
+    }
+    if (key_length && val_start) {
+        val_length = sz - val_start;
+        if (!call_ftc_callback(callback, src, key_start, key_length, val_start, val_length, has_semicolons)) return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef module_methods[] = {
     {"begin_create_signature", begin_create_signature, METH_VARARGS, ""},
     {"begin_load_signature", begin_load_signature, METH_NOARGS, ""},
@@ -180,6 +233,8 @@ static PyMethodDef module_methods[] = {
     {"begin_patch", begin_patch, METH_O, ""},
     {"begin_create_delta", begin_create_delta, METH_VARARGS, ""},
     {"iter_job", iter_job, METH_VARARGS, ""},
+    {"parse_ftc", parse_ftc, METH_VARARGS, ""},
+    {"decode_utf8_buffer", decode_utf8_buffer, METH_VARARGS, ""},
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
 

kittens/transfer/rsync.pyi

@@ -1,4 +1,4 @@
-from typing import Callable, Tuple
+from typing import Callable, Tuple, Union
 
 IO_BUFFER_SIZE: int
 
@@ -37,3 +37,11 @@ def begin_patch(callback: Callable[[memoryview, int], int]) -> JobCapsule:
 
 def iter_job(job_capsule: JobCapsule, input_data: bytes, output_buf: bytearray) -> Tuple[bool, int, int]:
     pass
+
+
+def parse_ftc(src: Union[str, bytes, memoryview], callback: Callable[[memoryview, memoryview, bool], None]) -> None:
+    pass
+
+
+def decode_utf8_buffer(src: Union[str, bytes, memoryview]) -> str:
+    pass

kittens/tui/loop.py

@@ -64,6 +64,7 @@ class Debug:
 
 
 debug = Debug()
+ftc_code = str(FILE_TRANSFER_CODE)
 
 
 class TermManager:
@@ -310,18 +311,25 @@ class Loop:
         pass
 
     def _on_osc(self, osc: str) -> None:
-        parts = osc.split(';', 1)
-        if len(parts) == 2 and parts[0].isdigit():
-            code = int(parts[0])
-            rest = parts[1]
-            if code == 52:
-                where, rest = rest.split(';', 1)
-                from_primary = 'p' in where
+        idx = osc.find(';')
+        if idx <= 0:
+            return
+        q = osc[:idx]
+        if q == '52':
+            widx = osc.find(';', idx + 1)
+            if widx < idx:
+                from_primary = osc.find('p', idx + 1) > -1
+                payload = ''
+            else:
                 from base64 import standard_b64decode
-                self.handler.on_clipboard_response(standard_b64decode(rest).decode('utf-8'), from_primary)
-            elif code == FILE_TRANSFER_CODE:
-                from kitty.file_transmission import FileTransmissionCommand
-                self.handler.on_file_transfer_response(FileTransmissionCommand.deserialize(rest))
+                from_primary = osc.find('p', idx+1, widx) > -1
+                data = memoryview(osc.encode('ascii'))
+                payload = standard_b64decode(data[widx+1:]).decode('utf-8')
+            self.handler.on_clipboard_response(payload, from_primary)
+        elif q == ftc_code:
+            from kitty.file_transmission import FileTransmissionCommand
+            data = memoryview(osc.encode('ascii'))
+            self.handler.on_file_transfer_response(FileTransmissionCommand.deserialize(data[idx+1:]))
 
     def _on_apc(self, apc: str) -> None:
         if apc.startswith('G'):

kitty/data-types.h

@@ -40,6 +40,9 @@ void log_error(const char *fmt, ...) __attribute__ ((format (printf, 1, 2)));
 #define fatal(...) { log_error(__VA_ARGS__); exit(EXIT_FAILURE); }
 static inline void cleanup_free(void *p) { free(*(void**)p); }
 #define FREE_AFTER_FUNCTION __attribute__((cleanup(cleanup_free)))
+static inline void cleanup_clear(void *p) { Py_XDECREF((PyObject *)p); }
+#define DECREF_AFTER_FUNCTION __attribute__((cleanup(cleanup_clear)))
+#define FREE_BUFFER_AFTER_FUNCTION __attribute__((cleanup(PyBuffer_Release)))
 
 typedef unsigned long long id_type;
 typedef uint32_t char_type;

kitty/file_transmission.py

@@ -195,30 +195,34 @@ class FileTransmissionCommand:
         return ''.join(map(as_unicode, self.get_serialized_fields(prefix_with_osc_code)))
 
     @classmethod
-    def deserialize(cls, data: str) -> 'FileTransmissionCommand':
+    def deserialize(cls, data: Union[str, bytes, memoryview]) -> 'FileTransmissionCommand':
         ans = FileTransmissionCommand()
-        parts = (x.replace('\0', ';').partition('=')[::2] for x in data.replace(';;', '\0').split(';'))
-        if not hasattr(cls, 'fmap'):
-            setattr(cls, 'fmap', {k.name: k for k in fields(cls)})
-        fmap: Dict[str, Field] = getattr(cls, 'fmap')
+        fmap: Dict[bytes, Field] = getattr(cls, 'fmap', None)
+        if not fmap:
+            fmap = {k.name.encode('ascii'): k for k in fields(cls)}
+            setattr(cls, 'fmap', fmap)
+        from kittens.transfer.rsync import parse_ftc, decode_utf8_buffer
 
-        for k, v in parts:
-            field = fmap.get(k)
+        def handle_item(key: memoryview, val: memoryview, has_semicolons: bool) -> None:
+            field = fmap.get(key)
             if field is None:
-                continue
+                return
             if issubclass(field.type, Enum):
-                setattr(ans, field.name, field.type[v])
+                setattr(ans, field.name, field.type[decode_utf8_buffer(val)])
             elif field.type is bytes:
-                setattr(ans, field.name, standard_b64decode(v))
+                setattr(ans, field.name, standard_b64decode(val))
             elif field.type is int:
-                setattr(ans, field.name, int(v))
+                setattr(ans, field.name, int(val))
             elif field.type is str:
                 if field.metadata.get('base64'):
-                    sval = standard_b64decode(v).decode('utf-8')
+                    sval = standard_b64decode(val).decode('utf-8')
                 else:
-                    sval = v
+                    if has_semicolons:
+                        val = memoryview(bytes(val).replace(b';;', b';'))
+                    sval = decode_utf8_buffer(val)
                 setattr(ans, field.name, sanitize_control_codes(sval))
 
+        parse_ftc(data, handle_item)
         if ans.action is Action.invalid:
             raise ValueError('No valid action specified in file transmission command')
 

kitty_tests/file_transmission.py

@@ -14,6 +14,7 @@ from kittens.transfer.librsync import (
     LoadSignature, PatchFile, delta_for_file, signature_of_file
 )
 from kittens.transfer.main import parse_transfer_args
+from kittens.transfer.rsync import decode_utf8_buffer, parse_ftc
 from kittens.transfer.send import files_for_send
 from kittens.transfer.utils import set_paths
 from kitty.file_transmission import (
@@ -261,6 +262,25 @@ class TestFileTransmission(BaseTest):
         self.ae(os.stat(dest + 'd2').st_mtime_ns, 29)
         self.assertFalse(ft.active_receives)
 
+    def test_parse_ftc(self):
+        def t(raw, *expected):
+            a = []
+
+            def c(k, v, has_semicolons):
+                a.append(decode_utf8_buffer(k))
+                if has_semicolons:
+                    v = bytes(v).replace(b';;', b';')
+                a.append(decode_utf8_buffer(v))
+
+            parse_ftc(raw, c)
+            self.ae(tuple(a), expected)
+
+        t('a=b', 'a', 'b')
+        t('a=b;', 'a', 'b')
+        t('a1=b1;c=d;;', 'a1', 'b1', 'c', 'd;')
+        t('a1=b1;c=d;;e', 'a1', 'b1', 'c', 'd;e')
+        t('a1=b1;c=d;;;1=1', 'a1', 'b1', 'c', 'd;', '1', '1')
+
     def test_path_mapping(self):
         opts = parse_transfer_args([])[0]
         b = Path(os.path.join(self.tdir, 'b'))