equix/
collision.rs

1//! Utilities for representing and finding partial sum collisions in the solver
2
3use crate::bucket_array::hash::{
4    Count, Insert, Key, KeyLookup, KeyStorage, Shape, ValueBucketArray, ValueLookup,
5};
6use crate::bucket_array::mem::BucketArrayMemory;
7use std::fmt::Debug;
8use std::ops::{BitAnd, BitOr, Shl, Shr};
9
10/// Look for partial sum collisions between items in one bucket array.
11///
12/// The items in each bucket are not sorted. This uses an additional small
13/// hash table, with the supplied backing memory, to collect matches.
14///
15/// The temporary memory can have an arbitrary shape. Capacity of the
16/// buffer will affect how may potential collisions we have to discard,
17/// and bucket count will affect how much of the key we are operating on.
18/// Its value type must match the [`Count`] type of the input table, since
19/// it will store item-in-bucket indices.
20///
21/// For each collision, calls the supplied predicate with the remaining portion
22/// of the hash sum and a [`CollisionLocation`] describing the two items.
23#[inline(always)]
24pub(crate) fn search<const TEMP_N: usize, const TEMP_CAP: usize, A, F, C, K, KS>(
25    array: &A,
26    scratchpad: &mut BucketArrayMemory<TEMP_N, TEMP_CAP, C>,
27    num_bits: usize,
28    mut predicate: F,
29) where
30    A: Shape<K> + KeyLookup<KS, K>,
31    F: FnMut(K, CollisionLocation),
32    C: Count,
33    K: Key,
34    KS: KeyStorage<K>,
35{
36    for first_bucket in 0..=(A::NUM_BUCKETS / 2) {
37        let second_bucket = first_bucket.wrapping_neg() % A::NUM_BUCKETS;
38        let mut paired_item_hash =
39            ValueBucketArray::<'_, { TEMP_N }, { TEMP_CAP }, u8, K, C>::new(scratchpad);
40
41        // Collect into paired_item_hash a mapping from the remainder portion
42        // of the key, to the item index in the bucket where we found that key.
43        for first_item in array.item_range(first_bucket) {
44            let first_hash_remainder = array.item_stored_key(first_bucket, first_item);
45            let _ = paired_item_hash.insert(
46                first_hash_remainder.into_key(),
47                C::from_item_index(first_item),
48            );
49        }
50
51        // For each item in the second bucket, scan its (small) complementary
52        // bucket in paired_item_hash to find partial sum collisions.
53        for second_item in array.item_range(second_bucket) {
54            let second_hash = array.item_full_key(second_bucket, second_item);
55            let hash_complement = second_hash.wrapping_neg();
56
57            let (_, hash_complement_remainder) = array.split_wide_key(hash_complement);
58            let (bucket_in_paired_hash, _) =
59                paired_item_hash.split_wide_key(hash_complement_remainder);
60
61            for first_item in paired_item_hash
62                .item_range(bucket_in_paired_hash)
63                .map(|item| paired_item_hash.item_value(bucket_in_paired_hash, item))
64            {
65                let first_item: usize = first_item.into();
66                let first_hash = array.item_full_key(first_bucket, first_item);
67                let sum = first_hash.wrapping_add(&second_hash);
68
69                // Compare two items that are in complementary buckets, to see
70                // if they actually have a matching sum in all of num_bits.
71                if sum.low_bits_are_zero(num_bits) {
72                    predicate(
73                        sum >> num_bits,
74                        CollisionLocation {
75                            first_bucket,
76                            first_item,
77                            second_item,
78                        },
79                    );
80                }
81            }
82        }
83    }
84}
85
86/// Locating information for one partial sum collision between items in
87/// complementary buckets.
88#[derive(Debug, Clone)]
89pub(crate) struct CollisionLocation {
90    /// Bucket index for the first colliding item, and the additive inverse of
91    /// the bucket index for the second colliding item.
92    pub(crate) first_bucket: usize,
93    /// Index of the first colliding item within its bucket
94    pub(crate) first_item: usize,
95    /// Index of the second colliding item within its bucket
96    pub(crate) second_item: usize,
97}
98
99impl CollisionLocation {
100    /// Return values associated with both colliding items, as a 2-element array.
101    #[inline(always)]
102    pub(crate) fn pair<A: ValueLookup<T> + Shape<K>, K: Key, T: Copy>(&self, array: &A) -> [T; 2] {
103        [
104            array.item_value(self.first_bucket, self.first_item),
105            array.item_value(
106                self.first_bucket.wrapping_neg() % A::NUM_BUCKETS,
107                self.second_item,
108            ),
109        ]
110    }
111}
112
113/// Packed representation of a [`CollisionLocation`]
114#[derive(Debug, Copy, Clone)]
115pub(crate) struct PackedCollision<
116    T: Copy
117        + TryFrom<usize>
118        + TryInto<usize>
119        + Shl<usize, Output = T>
120        + Shr<usize, Output = T>
121        + BitAnd<T, Output = T>
122        + BitOr<T, Output = T>,
123    const BUCKET_BITS: usize,
124    const ITEM_BITS: usize,
125>(T);
126
127impl<
128        T: Copy
129            + TryFrom<usize>
130            + TryInto<usize>
131            + Shl<usize, Output = T>
132            + Shr<usize, Output = T>
133            + BitAnd<T, Output = T>
134            + BitOr<T, Output = T>,
135        const BUCKET_BITS: usize,
136        const ITEM_BITS: usize,
137    > PackedCollision<T, BUCKET_BITS, ITEM_BITS>
138{
139    /// Construct a new [`PackedCollision`] from its inner type.
140    #[inline(always)]
141    pub(crate) fn new(inner: T) -> Self {
142        Self(inner)
143    }
144
145    /// Unwrap this [`PackedCollision`] into its inner type.
146    #[inline(always)]
147    pub(crate) fn into_inner(self) -> T {
148        self.0
149    }
150
151    /// Cast to the inner type from [`usize`], with panic on overflow.
152    #[inline(always)]
153    fn from_usize(i: usize) -> T {
154        i.try_into()
155            .map_err(|_| ())
156            .expect("masked collision field always fits into bitfield type")
157    }
158
159    /// Cast the inner type to [`usize`], with panic on overflow.
160    #[inline(always)]
161    fn to_usize(i: T) -> usize {
162        i.try_into()
163            .map_err(|_| ())
164            .expect("masked collision field always fits in usize")
165    }
166
167    /// Construct a new packed location from a [`CollisionLocation`].
168    ///
169    /// Packs all members into a bitfield. Panics if any of the indices
170    /// are larger than the selected field widths can represent.
171    #[inline(always)]
172    pub(crate) fn pack(loc: &CollisionLocation) -> Self {
173        assert!(loc.first_bucket < (1 << BUCKET_BITS));
174        assert!(loc.first_item < (1 << ITEM_BITS));
175        assert!(loc.second_item < (1 << ITEM_BITS));
176
177        let first_bucket: T = Self::from_usize(loc.first_bucket) << (ITEM_BITS * 2);
178        let first_item: T = Self::from_usize(loc.first_item) << ITEM_BITS;
179        let second_item: T = Self::from_usize(loc.second_item);
180        Self(first_bucket | first_item | second_item)
181    }
182
183    /// Unpack a bitfield into its [`CollisionLocation`].
184    #[inline(always)]
185    pub(crate) fn unpack(&self) -> CollisionLocation {
186        let bucket_mask = Self::from_usize((1_usize << BUCKET_BITS) - 1);
187        let item_mask = Self::from_usize((1_usize << ITEM_BITS) - 1);
188
189        CollisionLocation {
190            first_bucket: Self::to_usize((self.0 >> (ITEM_BITS * 2)) & bucket_mask),
191            first_item: Self::to_usize((self.0 >> ITEM_BITS) & item_mask),
192            second_item: Self::to_usize(self.0 & item_mask),
193        }
194    }
195}