diff --git a/probables/cuckoo/countingcuckoo.py b/probables/cuckoo/countingcuckoo.py index f3df32c..023efee 100644 --- a/probables/cuckoo/countingcuckoo.py +++ b/probables/cuckoo/countingcuckoo.py @@ -245,10 +245,13 @@ def _insert_fingerprint_alt( # either move everything around or hit the maximum number of swaps idx = random.choice([idx_1, idx_2]) prv_bin = CountingCuckooBin(fingerprint, 1) + orig_bin = prv_bin + swaps = [] for _ in range(self.max_swaps): # select one element to be swapped out... swap_elm = random.randint(0, self.bucket_size - 1) swap_finger = self.buckets[idx][swap_elm] + swaps.append((idx, swap_elm, swap_finger)) prv_bin, self.buckets[idx][swap_elm] = swap_finger, prv_bin # now find another place to put this fingerprint @@ -261,8 +264,13 @@ def _insert_fingerprint_alt( self.__unique_elements += 1 return None + # unable to place the new fingerprint; undo the swaps so that the + # failed insertion leaves the filter unchanged + for idx, swap_elm, elm in reversed(swaps): + self.buckets[idx][swap_elm] = elm + # if we got here we have an error... we might need to know what is left - return prv_bin + return orig_bin def _check_if_present(self, idx_1: int, idx_2: int, fingerprint: int) -> int | None: """wrapper for checking if fingerprint is already inserted""" @@ -304,16 +312,30 @@ def _parse_buckets(self, d: ByteString) -> None: def _expand_logic(self, extra_fingerprint: Union["CountingCuckooBin", None]) -> None: """the logic to acutally expand the cuckoo filter""" - # get all the fingerprints - fingerprints = self._setup_expand(extra_fingerprint) - self.__unique_elements = 0 # this needs to be reset! - - for elm in fingerprints: - idx_1, idx_2 = self._indicies_from_fingerprint(elm.finger) - res = self._insert_fingerprint_alt(elm.finger, idx_1, idx_2, elm.count) - if res is not None: # again, this *shouldn't* happen - msg = "The CountingCuckooFilter failed to expand" - raise CuckooFilterFullError(msg) + # snapshot the current state so that a failed expansion leaves the + # filter unchanged + old_capacity = self._cuckoo_capacity + old_buckets = self._buckets + old_inserted = self._inserted_elements + old_unique = self.__unique_elements + + try: + # get all the fingerprints + fingerprints = self._setup_expand(extra_fingerprint) + self.__unique_elements = 0 # this needs to be reset! + + for elm in fingerprints: + idx_1, idx_2 = self._indicies_from_fingerprint(elm.finger) + res = self._insert_fingerprint_alt(elm.finger, idx_1, idx_2, elm.count) + if res is not None: # again, this *shouldn't* happen + msg = "The CountingCuckooFilter failed to expand" + raise CuckooFilterFullError(msg) + except Exception: + self._cuckoo_capacity = old_capacity + self._buckets = old_buckets + self._inserted_elements = old_inserted + self.__unique_elements = old_unique + raise def __insert_element(self, fingerprint, idx, count=1) -> bool: """insert an element""" diff --git a/probables/cuckoo/cuckoo.py b/probables/cuckoo/cuckoo.py index 0625b57..b83c171 100644 --- a/probables/cuckoo/cuckoo.py +++ b/probables/cuckoo/cuckoo.py @@ -371,12 +371,15 @@ def _insert_fingerprint(self, fingerprint, idx_1, idx_2): # and move things around to the other index, if possible, until we # either move everything around or hit the maximum number of swaps idx = random.choice([idx_1, idx_2]) + orig_fingerprint = fingerprint + swaps = [] for _ in range(self.max_swaps): # select one element to be swapped out... swap_elm = random.randint(0, self.bucket_size - 1) swb = self.buckets[idx][swap_elm] + swaps.append((idx, swap_elm, swb)) fingerprint, self.buckets[idx][swap_elm] = swb, fingerprint # now find another place to put this fingerprint @@ -388,8 +391,13 @@ def _insert_fingerprint(self, fingerprint, idx_1, idx_2): self._inserted_elements += 1 return None + # unable to place the new fingerprint; undo the swaps so that the + # failed insertion leaves the filter unchanged + for idx, swap_elm, elm in reversed(swaps): + self.buckets[idx][swap_elm] = elm + # if we got here we have an error... we might need to know what is left - return fingerprint + return orig_fingerprint def _load(self, file: Path | str | IOBase | mmap | bytes) -> None: """load a cuckoo filter from file""" @@ -454,15 +462,27 @@ def __insert_element(self, fingerprint, idx) -> bool: def _expand_logic(self, extra_fingerprint): """the logic to acutally expand the cuckoo filter""" - # get all the fingerprints - fingerprints = self._setup_expand(extra_fingerprint) - - for finger in fingerprints: - idx_1, idx_2 = self._indicies_from_fingerprint(finger) - res = self._insert_fingerprint(finger, idx_1, idx_2) - if res is not None: # again, this *shouldn't* happen - msg = "The CuckooFilter failed to expand" - raise CuckooFilterFullError(msg) + # snapshot the current state so that a failed expansion leaves the + # filter unchanged + old_capacity = self._cuckoo_capacity + old_buckets = self._buckets + old_inserted = self._inserted_elements + + try: + # get all the fingerprints + fingerprints = self._setup_expand(extra_fingerprint) + + for finger in fingerprints: + idx_1, idx_2 = self._indicies_from_fingerprint(finger) + res = self._insert_fingerprint(finger, idx_1, idx_2) + if res is not None: # again, this *shouldn't* happen + msg = "The CuckooFilter failed to expand" + raise CuckooFilterFullError(msg) + except Exception: + self._cuckoo_capacity = old_capacity + self._buckets = old_buckets + self._inserted_elements = old_inserted + raise def _setup_expand(self, extra_fingerprint): """setup this thing""" diff --git a/tests/countingcuckoo_test.py b/tests/countingcuckoo_test.py index da4e4ae..c9b4d3b 100755 --- a/tests/countingcuckoo_test.py +++ b/tests/countingcuckoo_test.py @@ -3,6 +3,7 @@ import hashlib import os +import random import sys import unittest from pathlib import Path @@ -127,6 +128,33 @@ def test_c_cuckoo_full_msg(self): else: self.assertEqual(True, False) + def test_c_cuckoo_filter_full_preserves_elements(self): + """test that a failed insertion does not remove previously added elements""" + self.addCleanup(random.setstate, random.getstate()) + random.seed(0) + cko = CountingCuckooFilter(capacity=4, bucket_size=1, max_swaps=1, auto_expand=False) + added = ["k54", "k24", "k48", "k56"] + for key in added: + cko.add(key) + self.assertRaises(CuckooFilterFullError, cko.add, "k36") + self.assertEqual(len(added), cko.elements_added) + for key in added: + self.assertTrue(cko.check(key)) + + def test_c_cuckoo_filter_failed_expand_preserves_elements(self): + """test that a failed expansion does not remove previously added elements""" + self.addCleanup(random.setstate, random.getstate()) + random.seed(15) + cko = CountingCuckooFilter(capacity=8, bucket_size=1, max_swaps=2) + added = ["k{}".format(i) for i in range(18)] + for key in added: + cko.add(key) + self.assertEqual(32, cko.capacity) # earlier expansions succeeded + self.assertRaises(CuckooFilterFullError, cko.add, "k18") + self.assertEqual(32, cko.capacity) # the failed expansion is rolled back + for key in added: + self.assertTrue(cko.check(key)) + def test_c_cuckoo_idx(self): """test that the indexing works correctly for counting cuckoo filter swap""" diff --git a/tests/cuckoo_test.py b/tests/cuckoo_test.py index 1ee34e0..68ff51f 100755 --- a/tests/cuckoo_test.py +++ b/tests/cuckoo_test.py @@ -3,6 +3,7 @@ import hashlib import os +import random import sys import unittest from pathlib import Path @@ -185,6 +186,34 @@ def test_cuckoo_full_msg(self): else: self.assertEqual(True, False) + def test_cuckoo_filter_full_preserves_elements(self): + """test that a failed insertion does not remove previously added elements""" + self.addCleanup(random.setstate, random.getstate()) + random.seed(0) + cko = CuckooFilter(capacity=4, bucket_size=1, max_swaps=1, auto_expand=False) + added = ["k24", "k26", "k2", "k16"] + for key in added: + cko.add(key) + self.assertRaises(CuckooFilterFullError, cko.add, "k32") + self.assertEqual(len(added), cko.elements_added) + for key in added: + self.assertTrue(cko.check(key)) + + def test_cuckoo_filter_failed_expand_preserves_elements(self): + """test that a failed expansion does not remove previously added elements""" + self.addCleanup(random.setstate, random.getstate()) + random.seed(1) + cko = CuckooFilter(capacity=4, bucket_size=1, max_swaps=1) + added = ["k8", "k36", "k54", "k51", "k7", "k31"] + for key in added: + cko.add(key) + self.assertEqual(8, cko.capacity) # an expansion already succeeded + self.assertRaises(CuckooFilterFullError, cko.add, "k48") + self.assertEqual(8, cko.capacity) # the failed expansion is rolled back + self.assertEqual(len(added), cko.elements_added) + for key in added: + self.assertTrue(cko.check(key)) + def test_cuckoo_idx(self): """test that the indexing works correctly for cuckoo filter swap""" cko = CuckooFilter(capacity=100, bucket_size=2, max_swaps=5)