Make XOR64 test also test alignment issues

kitty/disk-cache.c

@@ -5,7 +5,6 @@
  * Distributed under terms of the GPL3 license.
  */
 
-#define EXTRA_INIT if (PyModule_AddFunctions(module, module_methods) != 0) return false;
 #define MAX_KEY_SIZE 16u
 
 #include "disk-cache.h"
@@ -687,18 +686,6 @@ PYWRAP(ensure_state) {
     Py_RETURN_NONE;
 }
 
-PYWRAP(xor_data64) {
-    (void) self;
-    const char *key, *data;
-    Py_ssize_t keylen, data_sz;
-    PA("s#s#", &key, &keylen, &data, &data_sz);
-    if (keylen != 64) { PyErr_SetString(PyExc_TypeError, "key must be 64 bytes long"); return NULL; }
-    PyObject *ans = PyBytes_FromStringAndSize(data, data_sz);
-    if (ans == NULL) return NULL;
-    xor_data64((const uint8_t*)key, (uint8_t*)PyBytes_AS_STRING(ans), data_sz);
-    return ans;
-}
-
 PYWRAP(read_from_cache_file) {
     Py_ssize_t pos = 0, sz = -1;
     PA("|nn", &pos, &sz);
@@ -836,10 +823,5 @@ PyTypeObject DiskCache_Type = {
     .tp_new = new_diskcache_object,
 };
 
-static PyMethodDef module_methods[] = {
-    MW(xor_data64, METH_VARARGS),
-    {NULL, NULL, 0, NULL}        /* Sentinel */
-};
-
 INIT_TYPE(DiskCache)
 PyObject* create_disk_cache(void) { return new_diskcache_object(&DiskCache_Type, NULL, NULL); }

kitty/simd-string.c

@@ -144,11 +144,49 @@ test_find_either_of_two_bytes(PyObject *self UNUSED, PyObject *args) {
     return PyLong_FromUnsignedLongLong(n);
 }
 
+static PyObject*
+test_xor64(PyObject *self UNUSED, PyObject *args) {
+    RAII_PY_BUFFER(buf);
+    RAII_PY_BUFFER(key);
+    int which_function = 0, align_offset = 0;
+    void (*func)(const uint8_t key[64], uint8_t* data, const size_t data_sz) = xor_data64;
+    if (!PyArg_ParseTuple(args, "s*s*BB|ii", &key, &buf, &which_function, &align_offset)) return NULL;
+    switch (which_function) {
+        case 1:
+            func = xor_data64_scalar; break;
+        case 2:
+            func = xor_data64_128; break;
+        case 3:
+            func = xor_data64_256; break;
+        case 0: break;
+        default:
+            PyErr_SetString(PyExc_ValueError, "Unknown which_function");
+            return NULL;
+    }
+    uint8_t *abuf;
+    if (posix_memalign((void**)&abuf, 64, 256 + buf.len) != 0) {
+        return PyErr_NoMemory();
+    }
+    uint8_t *p = abuf;
+    memset(p, '<', 64 + align_offset); p += 64 + align_offset;
+    memcpy(p, buf.buf, buf.len);
+    memset(p + buf.len, '>', 64);
+    func(key.buf, p, buf.len);
+    PyObject *ans = NULL;
+    for (int i = 0; i < 64 + align_offset; i++) if (abuf[i] != '<') { PyErr_SetString(PyExc_SystemError, "xor wrote before start of data region"); }
+    for (int i = 0; i < 64; i++) if (p[i + buf.len] != '>') { PyErr_SetString(PyExc_SystemError, "xor wrote after end of data region"); }
+    if (!PyErr_Occurred()) ans = PyBytes_FromStringAndSize((const char*)p, buf.len);
+    free(abuf);
+    return ans;
+}
+
+
 // }}}
 
 static PyMethodDef module_methods[] = {
     METHODB(test_utf8_decode_to_sentinel, METH_VARARGS),
     METHODB(test_find_either_of_two_bytes, METH_VARARGS),
+    METHODB(test_xor64, METH_VARARGS),
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
 

kitty_tests/graphics.py

@@ -10,9 +10,8 @@ import zlib
 from contextlib import suppress
 from dataclasses import dataclass
 from io import BytesIO
-from itertools import cycle
 
-from kitty.fast_data_types import base64_decode, base64_encode, load_png_data, shm_unlink, shm_write, xor_data64
+from kitty.fast_data_types import base64_decode, base64_encode, has_avx2, has_sse4_2, load_png_data, shm_unlink, shm_write, test_xor64
 
 from . import BaseTest, parse_bytes
 
@@ -184,17 +183,28 @@ def make_send_command(screen):
 class TestGraphics(BaseTest):
 
     def test_xor_data(self):
+        base_data = b'\x01' * 64
+        key = b'\x02' * 64
+        sizes = []
+        if has_sse4_2:
+            sizes.append(2)
+        if has_avx2:
+            sizes.append(3)
+        sizes.append(0)
 
-        def xor(skey, data):
-            ckey = cycle(bytearray(skey))
-            return bytes(bytearray(k ^ d for k, d in zip(ckey, bytearray(data))))
+        def t(key, data, align_offset=0):
+            expected = test_xor64(key, data, 1, 0)
+            for which_function in sizes:
+                actual = test_xor64(key, data, which_function, align_offset)
+                self.ae(expected, actual, f'{align_offset=} {len(data)=}')
 
-        base_data = os.urandom(61)
-        key = os.urandom(64)
-        for base in (b'', base_data, base_data * 3):
+        t(key, b'')
+
+        for base in (b'abc', base_data):
             for extra in range(len(base_data)):
-                data = base + base_data[:extra]
-                self.assertEqual(xor(key, data), xor_data64(key, data))
+                for align_offset in range(64):
+                    data = base + base_data[:extra]
+                    t(key, data, align_offset)
 
     def test_disk_cache(self):
         s = self.create_screen()