perf: scan the keys on the first BoundedSet lookup or change after a read

docs/how-to/store-a-large-set-or-dict.md

@@ -10,10 +10,11 @@ of keys costs milliseconds each way. Declare the field as
 [`BoundedSet`][sharedbox.BoundedSet] instead, and a read copies the stored
 bytes without building any element, while assigning a `BoundedSet` to the
 field copies its bytes back in one step. An `in` test encodes only the key
-it looks for and compares it with the stored bytes. Up to 8 elements it
-compares every element. Above that, the first test after a read hashes
-every element into a table the value keeps in its own memory, not in the
-box, and every later test looks up only the one key.
+it looks for and compares it with the stored bytes. The first 64 calls to
+`in`, `add`, `discard` or `remove` after a read compare the key with every
+element. In a set of more than 8 elements, the next one hashes every
+element once into a table the value keeps in its own memory, not in the
+box, and from then on each test or change looks up only the one key.
 
 Use it for a large set that you read or write whole and test for a few
 keys at a time. Going through every element decodes each one, and so does

src/sharedbox/_native/keys.cpp

@@ -79,7 +79,7 @@ KeyPlan plan_of(const detail::type_table &tree, const detail::type_ref &key) {
     plan.size = key.size;
     plan.mask.assign(key.size, std::byte{0xff});
     add_checks(tree, key, 0, no_guard, plan);
-    if (plan.checks.empty())
+    if (std::ranges::all_of(plan.mask, [](std::byte b) { return b == std::byte{0xff}; }))
         plan.mask.clear();
     return plan;
 }
@@ -194,6 +194,67 @@ std::uint32_t KeyedStore::find(const std::byte *key) {
     }
 }
 
+std::uint32_t KeyedStore::scan(const std::byte *key, std::uint32_t end) const {
+    if (plan_->mask.empty() && key_size_ == 8 && stride_ == 8) {
+        std::uint64_t wanted;
+        std::memcpy(&wanted, key, sizeof wanted);
+        const std::byte *first = slot(0);
+        constexpr std::uint32_t block = 32;
+        while (end > 0) {
+            const std::uint32_t start = end > block ? end - block : 0;
+            // Bit 63 of (x - 1) & ~x is set only when x is 0. It takes a subtraction and bit operations, which
+            // SSE2 has for 64-bit lanes where it has no 64-bit compare, so GCC and MSVC turn the loop into SIMD
+            // instructions; only a block holding key is searched one slot at a time.
+            std::uint64_t any = 0;
+            for (std::uint32_t i = start; i < end; ++i) {
+                std::uint64_t stored;
+                std::memcpy(&stored, first + 8 * std::size_t{i}, sizeof stored);
+                const std::uint64_t x = stored ^ wanted;
+                any |= (x - 1) & ~x;
+            }
+            if (any >> 63)
+                for (std::uint32_t i = end; i-- > start;)
+                    if (std::memcmp(first + 8 * std::size_t{i}, key, 8) == 0)
+                        return i;
+            end = start;
+        }
+        return no_slot;
+    }
+    for (std::uint32_t at = end; at-- > 0;)
+        if (same(slot(at), key))
+            return at;
+    return no_slot;
+}
+
+std::uint32_t KeyedStore::lookup(const std::byte *key) {
+    if (checked_ || scans_ == scans_before_index) {
+        ensure_checked(0);
+        return find(key);
+    }
+    codec();
+    // The first scan checks every slot, as ensure_checked would, so a refused key raises whichever slot holds
+    // it; the count goes up only after, so a refused value raises again on its next lookup.
+    if (scans_ == 0 && !plan_->checks.empty())
+        for (std::uint32_t at = 0; at < count_; ++at)
+            if (canonical(slot(at)) != nullptr)
+                corrupt(label_);
+    ++scans_;
+    return scan(key, count_);
+}
+
+bool KeyedStore::remove_key(const std::byte *key) {
+    const bool scanning = !checked_ && scans_ < scans_before_index;
+    std::uint32_t at = lookup(key);
+    if (at == no_slot)
+        return false;
+    remove_moving_last(at);
+    // scan found the last copy, and the slot moved into at was after it, so earlier slots hold any other.
+    if (scanning)
+        while ((at = scan(key, at)) != no_slot)
+            remove_moving_last(at);
+    return true;
+}
+
 void KeyedStore::allocate(std::uint32_t keys) {
     const std::uint32_t size = std::bit_ceil(std::max<std::uint32_t>(2 * keys, 2));
     index_.table.reset(new std::uint64_t[size]);

src/sharedbox/_native/keys.hpp

@@ -21,6 +21,9 @@ inline constexpr std::uint32_t no_guard = static_cast<std::uint32_t>(-1);
 
 /// The most keys a lookup scans; above it, the lookup probes the index.
 inline constexpr std::uint32_t scan_limit = 8;
+/// How many lookups a value scans before checking its keys and building the index: a scan costs well under
+/// a hundredth of building the index, so a value read for a few tests or changes never builds one.
+inline constexpr std::uint32_t scans_before_index = 64;
 /// No slot; a capacity is at most 2**20, so no slot index reaches it.
 inline constexpr std::uint32_t no_slot = 0xffffffff;
 
@@ -35,7 +38,7 @@ struct KeyCheck {
 /// How the bytes of a key are checked and compared.
 struct KeyPlan {
     /// The bytes that count when keys are compared: 0xff except each datetime and time fold bit. Empty when
-    /// there are no checks, so keys compare whole.
+    /// no bit is masked, so keys compare whole.
     std::vector<std::byte> mask;
     /// Empty for a key of int, str or enum, whose stored bytes are compared as they are.
     std::vector<KeyCheck> checks;
@@ -71,6 +74,13 @@ class KeyedStore : public SlotStore {
     bool encode_key(nb::handle value, std::byte *out, bool quiet) const;
     /// The slot holding key, or no_slot.
     std::uint32_t find(const std::byte *key);
+    /// The slot of a set holding key, or no_slot. Until scans_before_index lookups on keys not yet checked,
+    /// a lookup compares key with every slot through scan, dropping no duplicate and building no index; the
+    /// first of them checks each slot as ensure_checked does. Later ones run ensure_checked and find.
+    std::uint32_t lookup(const std::byte *key);
+    /// Removes key from a set and returns whether it held it. While lookups scan keys that may hold a
+    /// duplicate, it removes every slot holding key.
+    bool remove_key(const std::byte *key);
     /// Appends a slot holding encoded and indexes it; the caller has checked the capacity.
     void append_slot(const std::byte *encoded);
     /// Moves the last slot into slot at, which it removes.
@@ -85,6 +95,8 @@ class KeyedStore : public SlotStore {
     KeyIndex index_;
     /// Whether the stored keys have passed ensure_checked; a value that starts empty has.
     bool checked_ = false;
+    /// How many lookups have scanned the keys while they were not yet checked.
+    std::uint32_t scans_ = 0;
     const KeyPlan *plan_ = nullptr;
     std::uint32_t key_size_ = 0;
 
@@ -95,6 +107,8 @@ class KeyedStore : public SlotStore {
     /// The hash of key, masked through scratch, key_size_ bytes, when the plan has checks.
     std::uint64_t hash_of(const std::byte *key, std::byte *scratch) const;
     bool same(const std::byte *a, const std::byte *b) const;
+    /// The last slot before end holding key, or no_slot, comparing key with each slot.
+    std::uint32_t scan(const std::byte *key, std::uint32_t end) const;
     /// Builds the index over every slot, at a power of two of entries at least twice the count.
     void build();
     /// Replaces the index with an empty one of a power of two of entries at least twice keys.

src/sharedbox/_native/set.cpp

@@ -30,18 +30,19 @@ bool NativeSet::contains(nb::handle value) {
     if (!encode_key(value, key.data(), true))
         return false;
     nb::ft_lock_guard guard(mutex_);
-    ensure_checked(0);
-    return find(key.data()) != no_slot;
+    return lookup(key.data()) != no_slot;
 }
 
 void NativeSet::add(nb::handle value) {
     SlotBuffer key(stride_);
     encode_key(value, key.data(), false);
     {
         nb::ft_lock_guard guard(mutex_);
-        ensure_checked(0);
-        if (find(key.data()) != no_slot)
+        if (lookup(key.data()) != no_slot)
             return;
+        // Duplicates another program stored may fill the capacity; dropping them may make room.
+        if (count_ == capacity_)
+            ensure_checked(0);
         if (count_ < capacity_)
             return append_slot(key.data());
     }
@@ -53,18 +54,15 @@ void NativeSet::discard(nb::handle value) {
     if (!encode_key(value, key.data(), true))
         return;
     nb::ft_lock_guard guard(mutex_);
-    ensure_checked(0);
-    if (const std::uint32_t at = find(key.data()); at != no_slot)
-        remove_moving_last(at);
+    remove_key(key.data());
 }
 
 void NativeSet::remove(nb::handle value) {
     SlotBuffer key(stride_);
     if (encode_key(value, key.data(), true)) {
         nb::ft_lock_guard guard(mutex_);
-        ensure_checked(0);
-        if (const std::uint32_t at = find(key.data()); at != no_slot)
-            return remove_moving_last(at);
+        if (remove_key(key.data()))
+            return;
     }
     // In a tuple, so a tuple value is the exception's one argument, as for set.remove.
     PyErr_SetObject(PyExc_KeyError, nb::make_tuple(value).ptr());

src/sharedbox/benchmarks/ops.py

@@ -429,26 +429,62 @@ def resources(kind: str) -> tuple[Any, ...]:
     ("read set16", "SharedBox", TYPED, ["typed.s16"]),
     ("in set16", "SharedBox", TYPED, ["8 in typed.s16"]),
     ("iterate set16", "SharedBox", TYPED, ["for _ in typed.s16: pass"]),
+    (
+        "change set16",
+        "SharedBox",
+        TYPED,
+        ["s = typed.s16", "s.discard(0)", "s.add(0)", "typed.s16 = s"],
+    ),
     ("write set1k", "SharedBox", TYPED, ["typed.s1k = S1K"]),
     ("read set1k", "SharedBox", TYPED, ["typed.s1k"]),
     ("in set1k", "SharedBox", TYPED, ["512 in typed.s1k"]),
     ("iterate set1k", "SharedBox", TYPED, ["for _ in typed.s1k: pass"]),
+    (
+        "change set1k",
+        "SharedBox",
+        TYPED,
+        ["s = typed.s1k", "s.discard(0)", "s.add(0)", "typed.s1k = s"],
+    ),
     ("write set64k", "SharedBox", TYPED, ["typed.s64k = S64K"]),
     ("read set64k", "SharedBox", TYPED, ["typed.s64k"]),
     ("in set64k", "SharedBox", TYPED, ["32768 in typed.s64k"]),
     ("iterate set64k", "SharedBox", TYPED, ["for _ in typed.s64k: pass"]),
+    (
+        "change set64k",
+        "SharedBox",
+        TYPED,
+        ["s = typed.s64k", "s.discard(0)", "s.add(0)", "typed.s64k = s"],
+    ),
     ("write set16", "BoundedSet", TYPED, ["typed.bs16 = BS16"]),
     ("read set16", "BoundedSet", TYPED, ["typed.bs16"]),
     ("in set16", "BoundedSet", TYPED, ["8 in typed.bs16"]),
     ("iterate set16", "BoundedSet", TYPED, ["for _ in typed.bs16: pass"]),
+    (
+        "change set16",
+        "BoundedSet",
+        TYPED,
+        ["s = typed.bs16", "s.discard(0)", "s.add(0)", "typed.bs16 = s"],
+    ),
     ("write set1k", "BoundedSet", TYPED, ["typed.bs1k = BS1K"]),
     ("read set1k", "BoundedSet", TYPED, ["typed.bs1k"]),
     ("in set1k", "BoundedSet", TYPED, ["512 in typed.bs1k"]),
     ("iterate set1k", "BoundedSet", TYPED, ["for _ in typed.bs1k: pass"]),
+    (
+        "change set1k",
+        "BoundedSet",
+        TYPED,
+        ["s = typed.bs1k", "s.discard(0)", "s.add(0)", "typed.bs1k = s"],
+    ),
     ("write set64k", "BoundedSet", TYPED, ["typed.bs64k = BS64K"]),
     ("read set64k", "BoundedSet", TYPED, ["typed.bs64k"]),
     ("in set64k", "BoundedSet", TYPED, ["32768 in typed.bs64k"]),
     ("iterate set64k", "BoundedSet", TYPED, ["for _ in typed.bs64k: pass"]),
+    (
+        "change set64k",
+        "BoundedSet",
+        TYPED,
+        ["s = typed.bs64k", "s.discard(0)", "s.add(0)", "typed.bs64k = s"],
+    ),
     ("write dict16", "SharedBox", TYPED, ["typed.d16 = D16"]),
     ("read dict16", "SharedBox", TYPED, ["typed.d16"]),
     ("in dict16", "SharedBox", TYPED, ["8 in typed.d16"]),

tests/forged_open.py

@@ -88,6 +88,9 @@ def forged_open(edits: list[tuple[tuple[int, int], int, int]]) -> None:
         except LockTimeoutError:
             # A forged seq or writer_pid looks like a writer that never finishes.
             pass
+        except TypeError as error:
+            # A forged kind above 5 needs a class's Types, which a bare Segment has not.
+            assert "kind above 5" in str(error), error
         finally:
             other.close()
     finally:
@@ -99,9 +102,15 @@ def forged_open(edits: list[tuple[tuple[int, int], int, int]]) -> None:
 def main(examples: int) -> None:
     # A field count lowered to one the mapping still fits passes every geometry check.
     pinned = example([((64, 64), 20, 1)])
+    # Field 2's kind byte set to 10.
+    kind = example([((128, 40), 383, 10)])
     given(EDITS)(
         pinned(
-            settings(max_examples=examples, deadline=None, database=None)(forged_open)
+            kind(
+                settings(max_examples=examples, deadline=None, database=None)(
+                    forged_open
+                )
+            )
         )
     )()
 

tests/test_bounded_set.py

@@ -426,10 +426,10 @@ class Blue(enum.Enum):
     A = 1
 
 
-def keys_field(key: object) -> Any:
+def keys_field(key: object, capacity: int = 16) -> Any:
     """Return the layout of a class `Catalogue` with one field `keys`, a BoundedSet of `key`."""
     return build_layout(
-        make(Annotated[BoundedSet[key], Capacity(16)], "Catalogue", "keys")  # type: ignore[valid-type]
+        make(Annotated[BoundedSet[key], Capacity(capacity)], "Catalogue", "keys")  # type: ignore[valid-type]
     )
 
 
@@ -440,7 +440,7 @@ def empty_value(key: object) -> Any:
 
 def forged(key: object, slots: list[bytes]) -> Any:
     """Return a BoundedSet of `key` read from stored slots `slots`, as another program could write them."""
-    types = keys_field(key).types
+    types = keys_field(key, max(16, len(slots))).types
     head = bytearray(types.decode(0, bytes(64))._stored())
     head[:4] = len(slots).to_bytes(4, "little")
     return types.decode(0, bytes(head) + b"".join(slots))
@@ -561,11 +561,11 @@ def pack_float(x: float) -> bytes:
 def test_forged_duplicates_and_nan_are_found_when_a_key_is_looked_up(
     size: int,
 ) -> None:
-    """Drop a forged duplicate, refuse a forged nan and rewrite a forged -0.0 on the first lookup, by scan and by index."""
+    """Refuse a forged nan and rewrite a forged -0.0 on the first lookup, and drop a forged duplicate when a change checks the keys, by scan and by index."""
     rest = list(range(3, size))
     ints = forged(int, [pack_int(k) for k in [1, 2, 1, *rest]])
-    assert len(ints) == size
     assert 1 in ints
+    ints.update(())
     assert len(ints) == size - 1
     assert list(ints) == [1, 2, *rest]
     others = [float(k) for k in rest]
@@ -577,6 +577,43 @@ def test_forged_duplicates_and_nan_are_found_when_a_key_is_looked_up(
     assert math.copysign(1, next(iter(zero))) == 1
 
 
+@pytest.mark.parametrize("size", [3, 40])
+@pytest.mark.parametrize(("key", "pack"), [(int, pack_int), (float, pack_float)])
+def test_a_first_change_after_a_read_handles_forged_copies(
+    size: int, key: type, pack: collections.abc.Callable[[Any], bytes]
+) -> None:
+    """Remove every forged copy of a key on the first discard or remove after a read, add no key it holds, and add one when forged copies fill the capacity."""
+    keys = [1, 2, *range(3, size), 1]
+    discarded = forged(key, [pack(k) for k in keys])
+    discarded.discard(1)
+    assert 1 not in discarded
+    assert len(discarded) == size - 2
+    removed = forged(key, [pack(k) for k in keys])
+    removed.remove(1)
+    assert 1 not in removed
+    with pytest.raises(KeyError):
+        removed.remove(1)
+    added = forged(key, [pack(k) for k in keys])
+    added.add(2)
+    assert len(added) == size
+    # 16 slots fill a capacity of 16, but 1 is stored twice, so a 16th key fits.
+    full = forged(key, [pack(k) for k in [1, *range(1, 16)]])
+    full.add(key(100))
+    assert key(100) in full
+
+
+def test_lookups_find_every_key_before_and_after_the_index_is_built() -> None:
+    """Find every stored key and no other through many lookups and changes on one read, scanning first and then through the index."""
+    value = forged(int, [pack_int(k) for k in range(40)])
+    for k in range(100):
+        assert (k in value) == (k < 40)
+        value.discard(1000 + k)
+    value.discard(5)
+    value.add(5)
+    value.remove(7)
+    assert sorted(value) == [k for k in range(40) if k != 7]
+
+
 @pytest.mark.parametrize("size", [3, 12])
 def test_a_pop_checks_forged_keys_first(size: int) -> None:
     """Drop a forged duplicate before a pop takes the last slot, and refuse a forged nan, by scan and by index."""