From d0797a025be2dbf8430c4cdb42cab6c7d957fb44 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sun, 11 Feb 2024 12:56:43 +0530
Subject: [PATCH] Add dedicated tests for find_either_of_two

---
 kitty/simd-string.c   | 32 +++++++++++++++++++++++++++++
 kitty_tests/parser.py | 48 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/kitty/simd-string.c b/kitty/simd-string.c
index 851909e1d..e7d03cc9e 100644
--- a/kitty/simd-string.c
+++ b/kitty/simd-string.c
@@ -103,10 +103,42 @@ test_utf8_decode_to_sentinel(PyObject *self UNUSED, PyObject *args) {
     utf8_decoder_free(&d);
     return Py_BuildValue("OOi", found_sentinel ? Py_True : Py_False, ans, p);
 }
+
+static PyObject*
+test_find_either_of_two_bytes(PyObject *self UNUSED, PyObject *args) {
+    RAII_PY_BUFFER(buf);
+    int which_function = 0, align_offset = 0;
+    const uint8_t*(*func)(const uint8_t*, const size_t sz, const uint8_t, const uint8_t) = find_either_of_two_bytes;
+    unsigned char a, b;
+    if (!PyArg_ParseTuple(args, "s*BB|ii", &buf, &a, &b, &which_function, &align_offset)) return NULL;
+    switch (which_function) {
+        case 1:
+            func = find_either_of_two_bytes_scalar; break;
+        case 2:
+            func = find_either_of_two_bytes_128; break;
+        case 3:
+            func = find_either_of_two_bytes_256; break;
+    }
+    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);
+    const uint8_t *ans = func(p, buf.len, a, b);
+    free(abuf);
+    if (ans == NULL) return PyLong_FromLong(-1);
+    unsigned long long n = ans - p;
+    return PyLong_FromUnsignedLongLong(n);
+}
+
 // }}}
 
 static PyMethodDef module_methods[] = {
     METHODB(test_utf8_decode_to_sentinel, METH_VARARGS),
+    METHODB(test_find_either_of_two_bytes, METH_VARARGS),
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
 
diff --git a/kitty_tests/parser.py b/kitty_tests/parser.py
index e55503fd6..b0473989a 100644
--- a/kitty_tests/parser.py
+++ b/kitty_tests/parser.py
@@ -5,7 +5,16 @@ from base64 import standard_b64encode
 from binascii import hexlify
 from functools import partial
 
-from kitty.fast_data_types import CURSOR_BLOCK, VT_PARSER_BUFFER_SIZE, base64_decode, base64_encode, has_avx2, has_sse4_2, test_utf8_decode_to_sentinel
+from kitty.fast_data_types import (
+    CURSOR_BLOCK,
+    VT_PARSER_BUFFER_SIZE,
+    base64_decode,
+    base64_encode,
+    has_avx2,
+    has_sse4_2,
+    test_find_either_of_two_bytes,
+    test_utf8_decode_to_sentinel,
+)
 from kitty.notify import NotificationCommand, handle_notification_cmd, notification_activated, reset_registry
 
 from . import BaseTest, parse_bytes
@@ -280,6 +289,43 @@ class TestParser(BaseTest):
             pb(b'"\xf0\x9f\x98"', '"\ufffd"')
             pb(b'"\xef\x93\x94\x95"', '"\uf4d4\ufffd"')
 
+    def test_find_either_of_two_bytes(self):
+
+        def test(buf, a, b, align_offset=0):
+            a, b = ord(a), ord(b)
+            sizes = [0]
+            if has_sse4_2:
+                sizes.append(2)
+            if has_avx2:
+                sizes.append(3)
+            expected = test_find_either_of_two_bytes(buf, a, b, 1)
+
+            for sz in sizes:
+                actual = test_find_either_of_two_bytes(buf, a, b, sz, align_offset)
+                self.ae(expected, actual, f'Failed for: {buf!r} at {sz=} and {align_offset=}')
+
+        q = 'abc'
+        for off in range(32):
+            test(q, '<', '>', off)
+            test(q, ' ', 'b', off)
+            test(q, '<', 'a', off)
+            test(q, '<', 'b', off)
+            test(q, 'c', '>', off)
+
+        def tests(buf, a, b):
+            for sz in (0, 16, 32, 64, 79):
+                buf = (' ' * sz) + buf
+                for align_offset in range(32):
+                    test(buf, a, b, align_offset)
+        tests("", '<', '>')
+        tests("a", '\0', '\0')
+        tests("a", '<', '>')
+        tests("dsdfsfa", '1', 'a')
+        tests("xa", 'a', 'a')
+        tests("bbb", 'a', '1')
+        tests("bba", 'a', '<')
+        tests("baa", '>', 'a')
+
     def test_esc_codes(self):
         s = self.create_screen()
         pb = partial(self.parse_bytes_dump, s)