summaryrefslogtreecommitdiff
path: root/crates/equix/src/collision.rs
blob: 7fc587a4af4300e5b9bfa7055c161a6ec2cdb7e9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
//! Utilities for representing and finding partial sum collisions in the solver

use crate::bucket_array::hash::{
    Count, Insert, Key, KeyLookup, KeyStorage, Shape, ValueBucketArray, ValueLookup,
};
use crate::bucket_array::mem::BucketArrayMemory;
use std::fmt::Debug;
use std::ops::{BitAnd, BitOr, Shl, Shr};

/// Look for partial sum collisions between items in one bucket array.
///
/// The items in each bucket are not sorted. This uses an additional small
/// hash table, with the supplied backing memory, to collect matches.
///
/// The temporary memory can have an arbitrary shape. Capacity of the
/// buffer will affect how may potential collisions we have to discard,
/// and bucket count will affect how much of the key we are operating on.
/// Its value type must match the [`Count`] type of the input table, since
/// it will store item-in-bucket indices.
///
/// For each collision, calls the supplied predicate with the remaining portion
/// of the hash sum and a [`CollisionLocation`] describing the two items.
#[inline(always)]
pub(crate) fn search<const TEMP_N: usize, const TEMP_CAP: usize, A, F, C, K, KS>(
    array: &A,
    scratchpad: &mut BucketArrayMemory<TEMP_N, TEMP_CAP, C>,
    num_bits: usize,
    mut predicate: F,
) where
    A: Shape<K> + KeyLookup<KS, K>,
    F: FnMut(K, CollisionLocation),
    C: Count,
    K: Key,
    KS: KeyStorage<K>,
{
    for first_bucket in 0..=(A::NUM_BUCKETS / 2) {
        let second_bucket = first_bucket.wrapping_neg() % A::NUM_BUCKETS;
        let mut paired_item_hash =
            ValueBucketArray::<'_, { TEMP_N }, { TEMP_CAP }, u8, K, C>::new(scratchpad);

        // Collect into paired_item_hash a mapping from the remainder portion
        // of the key, to the item index in the bucket where we found that key.
        for first_item in array.item_range(first_bucket) {
            let first_hash_remainder = array.item_stored_key(first_bucket, first_item);
            let _ = paired_item_hash.insert(
                first_hash_remainder.into_key(),
                C::from_item_index(first_item),
            );
        }

        // For each item in the second bucket, scan its (small) complementary
        // bucket in paired_item_hash to find partial sum collisions.
        for second_item in array.item_range(second_bucket) {
            let second_hash = array.item_full_key(second_bucket, second_item);
            let hash_complement = second_hash.wrapping_neg();

            let (_, hash_complement_remainder) = array.split_wide_key(hash_complement);
            let (bucket_in_paired_hash, _) =
                paired_item_hash.split_wide_key(hash_complement_remainder);

            for first_item in paired_item_hash
                .item_range(bucket_in_paired_hash)
                .map(|item| paired_item_hash.item_value(bucket_in_paired_hash, item))
            {
                let first_item: usize = first_item.into();
                let first_hash = array.item_full_key(first_bucket, first_item);
                let sum = first_hash.wrapping_add(&second_hash);

                // Compare two items that are in complementary buckets, to see
                // if they actually have a matching sum in all of num_bits.
                if sum.low_bits_are_zero(num_bits) {
                    predicate(
                        sum >> num_bits,
                        CollisionLocation {
                            first_bucket,
                            first_item,
                            second_item,
                        },
                    );
                }
            }
        }
    }
}

/// Locating information for one partial sum collision between items in
/// complementary buckets.
#[derive(Debug, Clone)]
pub(crate) struct CollisionLocation {
    /// Bucket index for the first colliding item, and the additive inverse of
    /// the bucket index for the second colliding item.
    pub(crate) first_bucket: usize,
    /// Index of the first colliding item within its bucket
    pub(crate) first_item: usize,
    /// Index of the second colliding item within its bucket
    pub(crate) second_item: usize,
}

impl CollisionLocation {
    /// Return values associated with both colliding items, as a 2-element array.
    #[inline(always)]
    pub(crate) fn pair<A: ValueLookup<T> + Shape<K>, K: Key, T: Copy>(&self, array: &A) -> [T; 2] {
        [
            array.item_value(self.first_bucket, self.first_item),
            array.item_value(
                self.first_bucket.wrapping_neg() % A::NUM_BUCKETS,
                self.second_item,
            ),
        ]
    }
}

/// Packed representation of a [`CollisionLocation`]
#[derive(Debug, Copy, Clone)]
pub(crate) struct PackedCollision<
    T: Copy
        + TryFrom<usize>
        + TryInto<usize>
        + Shl<usize, Output = T>
        + Shr<usize, Output = T>
        + BitAnd<T, Output = T>
        + BitOr<T, Output = T>,
    const BUCKET_BITS: usize,
    const ITEM_BITS: usize,
>(T);

impl<
    T: Copy
        + TryFrom<usize>
        + TryInto<usize>
        + Shl<usize, Output = T>
        + Shr<usize, Output = T>
        + BitAnd<T, Output = T>
        + BitOr<T, Output = T>,
    const BUCKET_BITS: usize,
    const ITEM_BITS: usize,
> PackedCollision<T, BUCKET_BITS, ITEM_BITS>
{
    /// Construct a new [`PackedCollision`] from its inner type.
    #[inline(always)]
    pub(crate) fn new(inner: T) -> Self {
        Self(inner)
    }

    /// Unwrap this [`PackedCollision`] into its inner type.
    #[inline(always)]
    pub(crate) fn into_inner(self) -> T {
        self.0
    }

    /// Cast to the inner type from [`usize`], with panic on overflow.
    #[inline(always)]
    fn from_usize(i: usize) -> T {
        i.try_into()
            .map_err(|_| ())
            .expect("masked collision field always fits into bitfield type")
    }

    /// Cast the inner type to [`usize`], with panic on overflow.
    #[inline(always)]
    fn to_usize(i: T) -> usize {
        i.try_into()
            .map_err(|_| ())
            .expect("masked collision field always fits in usize")
    }

    /// Construct a new packed location from a [`CollisionLocation`].
    ///
    /// Packs all members into a bitfield. Panics if any of the indices
    /// are larger than the selected field widths can represent.
    #[inline(always)]
    pub(crate) fn pack(loc: &CollisionLocation) -> Self {
        assert!(loc.first_bucket < (1 << BUCKET_BITS));
        assert!(loc.first_item < (1 << ITEM_BITS));
        assert!(loc.second_item < (1 << ITEM_BITS));

        let first_bucket: T = Self::from_usize(loc.first_bucket) << (ITEM_BITS * 2);
        let first_item: T = Self::from_usize(loc.first_item) << ITEM_BITS;
        let second_item: T = Self::from_usize(loc.second_item);
        Self(first_bucket | first_item | second_item)
    }

    /// Unpack a bitfield into its [`CollisionLocation`].
    #[inline(always)]
    pub(crate) fn unpack(&self) -> CollisionLocation {
        let bucket_mask = Self::from_usize((1_usize << BUCKET_BITS) - 1);
        let item_mask = Self::from_usize((1_usize << ITEM_BITS) - 1);

        CollisionLocation {
            first_bucket: Self::to_usize((self.0 >> (ITEM_BITS * 2)) & bucket_mask),
            first_item: Self::to_usize((self.0 >> ITEM_BITS) & item_mask),
            second_item: Self::to_usize(self.0 & item_mask),
        }
    }
}