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}