equix/
solver.rs

1//! Find Equi-X solutions.
2//!
3//! This is a particular instance of Equihash, a tree-structured search for
4//! partial hash collisions using temporary memory. Equi-X modifies it to use
5//! sums instead of XOR, and chooses specific parameters.
6
7use crate::bucket_array::{
8    hash::Insert, hash::KeyValueBucketArray, mem::BucketArrayMemory, mem::Uninit,
9};
10use crate::collision::{self, PackedCollision};
11use crate::solution::{self, HashValue, Solution, SolutionArray, SolutionItem, EQUIHASH_N};
12use arrayvec::ArrayVec;
13use hashx::HashX;
14
15// The hash table bucket counts here are mostly constrained by the shape of
16// the Equihash tree, but the bucket capacities are somewhat arbitrary. Larger
17// buckets use more memory, smaller buckets may discard solutions if there are
18// sufficient collisions during the earlier layers. This uses the same bucket
19// counts as the original Equi-X implementation, in order to exactly match
20// its solution discarding behavior and thus its output.
21
22/// First layer hash table, needs room for the full set of [`HashValue`]s and
23/// [`SolutionItem`]s. Key remainder is `(EQUIHASH_N - 8 == 52)` bits here.
24type Layer0<'k, 'v> =
25    KeyValueBucketArray<'k, 'v, 256, 336, u16, HashValue, HashValue, SolutionItem>;
26/// Key backing storage for [`Layer0`]
27type Layer0KeyMem = BucketArrayMemory<256, 336, HashValue>;
28/// Value backing storage for [`Layer0`]
29type Layer0ValueMem = BucketArrayMemory<256, 336, SolutionItem>;
30/// Packed collision type for [`Layer0`]
31type Layer0Collision = PackedCollision<u32, 8, 9>;
32
33/// Next layer maps residual hash (3/4 full width == 45 bits) to packed
34/// [`Layer0Collision`]. Key remainder is 37 bits after this.
35type Layer1<'k, 'v> = KeyValueBucketArray<'k, 'v, 256, 336, u16, u64, u64, u32>;
36/// Key backing storage for [`Layer1`]
37type Layer1KeyMem = BucketArrayMemory<256, 336, u64>;
38/// Value backing storage for [`Layer1`]
39type Layer1ValueMem = BucketArrayMemory<256, 336, u32>;
40/// Packed collision type for [`Layer1`]
41type Layer1Collision = PackedCollision<u32, 8, 9>;
42
43/// Final layer maps the (N/2 == 30 bits) residual hash sum to packed
44/// [`Layer1Collision`]. Key remainder is 22 bits.
45type Layer2<'k, 'v> = KeyValueBucketArray<'k, 'v, 256, 336, u16, u32, u32, u32>;
46/// Key backing storage for [`Layer2`]
47type Layer2KeyMem = BucketArrayMemory<256, 336, u32>;
48/// Value backing storage for [`Layer2`]
49type Layer2ValueMem = BucketArrayMemory<256, 336, u32>;
50
51/// Temporary value hash for resolving collisions between buckets.
52/// This removes 7 bits of key between each of the other layers.
53/// Value type may hold a [`SolutionItem`] or a temporary item-in-bucket index.
54type TempMem = BucketArrayMemory<128, 12, u16>;
55
56/// Internal solver memory, inside the heap allocation
57///
58/// Contains exclusively [`std::mem::MaybeUninit`] members. Other than the
59/// temporary memory used between each pair of tree layers, the memory regions
60/// are organized by which layer of the solution tree they form starting
61/// from leaf `(hash, index)` items and moving toward the root of the tree.
62#[derive(Copy, Clone)]
63struct SolverMemoryInner {
64    /// Temporary memory for each [`collision::search()`]
65    temp: TempMem,
66    /// Memory overlay union, access controlled with mutable references
67    overlay: Overlay,
68    /// Temporary value memory for [`Layer0`]
69    layer0_values: Layer0ValueMem,
70    /// Temporary value memory for [`Layer1`]
71    layer1_values: Layer1ValueMem,
72    /// Temporary key storage memory for [`Layer1`]
73    layer1_keys: Layer1KeyMem,
74}
75
76/// SAFETY: We are proimising that [`SolverMemoryInner`] is
77///         made only from [`Uninit`] types like [`BucketArrayMemory`].
78unsafe impl Uninit for SolverMemoryInner {}
79
80/// Union of overlay layouts used in solver memory
81///
82/// As a memory optimization, some of the allocations in [`SolverMemoryInner`]
83/// are overlapped. We only use one of these at a time, checked statically
84/// using a mutable borrow.
85#[derive(Copy, Clone)]
86union Overlay {
87    /// Initial memory overlay layout
88    first: OverlayFirst,
89    /// Layout which replaces the first once we can drop `layer0_keys`
90    second: OverlaySecond,
91}
92
93impl Overlay {
94    /// Access the first overlay layout, via a mutable reference
95    fn first(&mut self) -> &mut OverlayFirst {
96        // SAFETY: This union and its members implement [`Uninit`], promising
97        //         that we can soundly create an instance out of uninitialized
98        //         or reused memory. Initialized state is controlled via a
99        //         mutable reference, and by reborrowing a union field we ensure
100        //         exclusive access using the borrow checker.
101        unsafe { &mut self.first }
102    }
103
104    /// Access the second overlay layout, via a mutable reference
105    fn second(&mut self) -> &mut OverlaySecond {
106        // SAFETY: As above, we implement [`Uninit`] and use &mut to control access
107        unsafe { &mut self.second }
108    }
109}
110
111/// First memory overlay, contains the key portion of [`Layer0`]
112#[derive(Copy, Clone)]
113struct OverlayFirst {
114    /// Key remainder table for [`Layer0`], which we can drop
115    /// after running [`collision::search()`] on that layer.
116    layer0_keys: Layer0KeyMem,
117}
118
119/// Second overlay, with both parts of Layer2
120#[derive(Copy, Clone)]
121struct OverlaySecond {
122    /// Temporary key storage memory for [`Layer2`]
123    layer2_keys: Layer2KeyMem,
124    /// Temporary value memory for [`Layer2`]
125    layer2_values: Layer2ValueMem,
126}
127
128/// Search for solutions, iterating the entire [`SolutionItem`] space and using
129/// temporary memory to locate partial sum collisions at each tree layer.
130pub(crate) fn find_solutions(func: &HashX, mem: &mut SolverMemory, results: &mut SolutionArray) {
131    // Use the first memory overlay layout.
132    let overlay = mem.heap.overlay.first();
133
134    // Enumerate all hash values into the first layer
135    let mut layer0 = Layer0::new(&mut overlay.layer0_keys, &mut mem.heap.layer0_values);
136    for item in SolutionItem::MIN..=SolutionItem::MAX {
137        let hash = solution::item_hash(func, item);
138        let _ = layer0.insert(hash, item);
139    }
140
141    // Now form the first layer of the Equihash tree,
142    // with collisions in the low N/4 (15) bits
143    let layer1_n = EQUIHASH_N / 4;
144    let mut layer1 = Layer1::new(&mut mem.heap.layer1_keys, &mut mem.heap.layer1_values);
145    collision::search(&layer0, &mut mem.heap.temp, layer1_n, |sum, loc| {
146        let _ = layer1.insert(sum, Layer0Collision::pack(&loc).into_inner());
147    });
148
149    // Once we finish searching a layer for collisions,
150    // we can drop the key data and make the rest immutable.
151    // This drops the last mutable reference into the first overlay.
152    let layer0 = layer0.drop_key_storage();
153
154    // Now switch to the second memory overlay layout.
155    let overlay = mem.heap.overlay.second();
156
157    // Next Equihash layer, collisions in the low N/2 (30) bits
158    let layer2_n = EQUIHASH_N / 2;
159    let mut layer2 = Layer2::new(&mut overlay.layer2_keys, &mut overlay.layer2_values);
160    collision::search(
161        &layer1,
162        &mut mem.heap.temp,
163        layer2_n - layer1_n,
164        |sum, loc| {
165            let _ = layer2.insert(sum as u32, Layer1Collision::pack(&loc).into_inner());
166        },
167    );
168
169    // Final layer, match the entire N bits and assemble complete solutions
170    let layer3_n = EQUIHASH_N;
171    collision::search(
172        &layer2,
173        &mut mem.heap.temp,
174        layer3_n - layer2_n,
175        |_sum, loc| {
176            let mut items = ArrayVec::<SolutionItem, { Solution::NUM_ITEMS }>::new();
177
178            // Walk the binary tree of collision locations, in order.
179            // The leaf layer will have our SolutionItems.
180            loc.pair(&layer2).into_iter().for_each(|loc| {
181                Layer1Collision::new(loc)
182                    .unpack()
183                    .pair(&layer1)
184                    .into_iter()
185                    .for_each(|loc| {
186                        Layer0Collision::new(loc)
187                            .unpack()
188                            .pair(&layer0)
189                            .into_iter()
190                            .for_each(|item| items.push(item));
191                    });
192            });
193
194            // Apply the ordering constraints and check for duplicates
195            let solution = Solution::sort_from_array(
196                items
197                    .into_inner()
198                    .expect("always collected a full SolutionItem tree"),
199            );
200            if results.last() != Some(&solution) {
201                let _ = results.try_push(solution);
202            }
203        },
204    );
205}
206
207/// Temporary memory used by the Equi-X solver
208///
209/// This space is needed temporarily during a solver run. It will be
210/// allocated on the heap by [`SolverMemory::new()`], and the solver
211/// provides a [`crate::EquiX::solve_with_memory()`] interface for reusing
212/// this memory between runs.
213pub struct SolverMemory {
214    /// Inner heap allocation which holds the actual solver memory
215    heap: Box<SolverMemoryInner>,
216}
217
218impl SolverMemory {
219    /// Size of the solver memory region, in bytes
220    pub const SIZE: usize = std::mem::size_of::<SolverMemoryInner>();
221
222    /// New uninitialized memory, usable as solver temporary space.
223    pub fn new() -> Self {
224        Self {
225            heap: SolverMemoryInner::alloc(),
226        }
227    }
228}
229
230impl Default for SolverMemory {
231    fn default() -> Self {
232        Self::new()
233    }
234}
235
236#[cfg(test)]
237mod test {
238    #[test]
239    fn solver_memory_size() {
240        // Regression test for memory usage. Our actual memory size is very
241        // similar to the original C implementation, the only difference that
242        // our bucket counters are stored outside this structure.
243        let size = super::SolverMemory::SIZE;
244        assert_eq!(size, 1_895_424);
245    }
246}