1use std::collections::{BTreeMap, BTreeSet};
13
14use serde::{Deserialize, Serialize};
15
16use crate::lines::Hash;
17use crate::{from_hex, to_hex};
18
19pub const LEAF: usize = 8;
21const MAX_OPEN: usize = 1 << 16;
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25struct Range {
26 lo: Hash,
27 hi: Option<Hash>,
29 depth: usize,
31}
32
33impl Range {
34 fn children(&self) -> Vec<Range> {
35 let (byte, shift) = (self.depth / 2, if self.depth.is_multiple_of(2) { 4 } else { 0 });
36 let with = |k: u8| {
37 let mut h = self.lo;
38 h[byte] |= k << shift;
39 h
40 };
41 (0..16u8).map(|k| Range { lo: with(k), hi: if k < 15 { Some(with(k + 1)) } else { self.hi }, depth: self.depth + 1 }).collect()
42 }
43
44 fn contains(&self, h: &Hash) -> bool {
45 *h >= self.lo && self.hi.is_none_or(|hi| *h < hi)
46 }
47}
48
49#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct Round {
53 pub fps: Vec<(String, u32)>,
54 pub leaves: Vec<Vec<String>>,
55}
56
57#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
59pub struct Cost {
60 pub rounds: usize,
61 pub fps: usize,
62 pub hashes: usize,
63}
64
65pub struct Recon<'a, V> {
66 set: &'a BTreeMap<Hash, V>,
67 open: Vec<Range>,
68 leaves: Vec<Range>,
69 pub want: BTreeSet<Hash>,
71 pub give: BTreeSet<Hash>,
72 pub cost: Cost,
73}
74
75impl<'a, V> Recon<'a, V> {
76 pub fn new(set: &'a BTreeMap<Hash, V>) -> Self {
77 Recon {
78 set,
79 open: vec![Range { lo: [0; 32], hi: None, depth: 0 }],
80 leaves: Vec::new(),
81 want: BTreeSet::new(),
82 give: BTreeSet::new(),
83 cost: Cost::default(),
84 }
85 }
86
87 fn keys(&self, r: &Range) -> impl Iterator<Item = &Hash> {
88 let range = match r.hi {
89 Some(hi) => self.set.range(r.lo..hi),
90 None => self.set.range(r.lo..),
91 };
92 range.map(|(k, _)| k)
93 }
94
95 fn fingerprint(&self, r: &Range) -> (String, u32) {
96 let mut h = blake3::Hasher::new();
97 let mut n = 0u32;
98 for k in self.keys(r) {
99 h.update(k);
100 n += 1;
101 }
102 (if n == 0 { "0".into() } else { to_hex(h.finalize().as_bytes()) }, n)
103 }
104
105 pub fn message(&mut self) -> Round {
107 let fps: Vec<(String, u32)> = self.open.iter().map(|r| self.fingerprint(r)).collect();
108 let leaves: Vec<Vec<String>> = self.leaves.iter().map(|r| self.keys(r).map(|k| to_hex(k)).collect()).collect();
109 self.cost.rounds += 1;
110 self.cost.fps += fps.len();
111 self.cost.hashes += leaves.iter().map(Vec::len).sum::<usize>();
112 Round { fps, leaves }
113 }
114
115 pub fn absorb(&mut self, mine: &Round, theirs: Round) -> Result<bool, String> {
117 if theirs.fps.len() != self.open.len() || theirs.leaves.len() != self.leaves.len() {
118 return Err("reconciliation: peer is out of step".into());
119 }
120 for (r, hashes) in self.leaves.iter().zip(theirs.leaves) {
121 let theirs: BTreeSet<Hash> = hashes.iter().map(|h| from_hex::<32>(h)).collect::<Result<_, _>>()?;
122 if !theirs.iter().all(|h| r.contains(h)) {
123 return Err("reconciliation: peer sent a hash outside its range".into());
124 }
125 for k in self.keys(r).copied().collect::<Vec<_>>() {
126 if !theirs.contains(&k) {
127 self.give.insert(k);
128 }
129 }
130 self.want.extend(theirs.into_iter().filter(|h| !self.set.contains_key(h)));
131 }
132 let (mut open, mut leaves) = (Vec::new(), Vec::new());
133 for ((r, my), their) in self.open.iter().zip(&mine.fps).zip(&theirs.fps) {
134 if my == their {
135 continue;
136 }
137 if (my.1 as usize) <= LEAF || (their.1 as usize) <= LEAF || r.depth >= 63 {
138 leaves.push(*r);
139 } else {
140 open.extend(r.children());
141 }
142 }
143 if open.len() > MAX_OPEN {
144 return Err("reconciliation: too many differing ranges".into());
145 }
146 (self.open, self.leaves) = (open, leaves);
147 Ok(self.open.is_empty() && self.leaves.is_empty())
148 }
149}
150
151#[cfg(test)]
152mod tests {
153 use super::*;
154
155 fn set(range: std::ops::Range<u32>) -> BTreeMap<Hash, ()> {
156 range.map(|i| (*blake3::hash(&i.to_le_bytes()).as_bytes(), ())).collect()
157 }
158
159 fn run(a: &BTreeMap<Hash, ()>, b: &BTreeMap<Hash, ()>) -> (BTreeSet<Hash>, BTreeSet<Hash>, Cost) {
161 let (mut ra, mut rb) = (Recon::new(a), Recon::new(b));
162 loop {
163 let (ma, mb) = (ra.message(), rb.message());
164 let da = ra.absorb(&ma, mb.clone()).unwrap();
165 let db = rb.absorb(&mb, ma).unwrap();
166 assert_eq!(da, db, "both sides stay in step");
167 if da {
168 break;
169 }
170 }
171 assert_eq!((&ra.want, &ra.give), (&rb.give, &rb.want));
172 let cost = Cost { rounds: ra.cost.rounds, fps: ra.cost.fps + rb.cost.fps, hashes: ra.cost.hashes + rb.cost.hashes };
173 (ra.want, ra.give, cost)
174 }
175
176 #[test]
177 fn equal_sets_cost_one_fingerprint() {
178 let a = set(0..5000);
179 let (want, give, cost) = run(&a, &a.clone());
180 assert!(want.is_empty() && give.is_empty());
181 assert_eq!(cost, Cost { rounds: 1, fps: 2, hashes: 0 });
182 }
183
184 #[test]
185 fn finds_exactly_the_difference_at_a_cost_bounded_by_it() {
186 let common = set(0..20_000);
187 let (mut a, mut b) = (common.clone(), common.clone());
188 a.extend(set(100_000..100_003));
189 b.extend(set(200_000..200_004));
190 let (want, give, cost) = run(&a, &b);
191 assert_eq!(want, set(200_000..200_004).into_keys().collect());
192 assert_eq!(give, set(100_000..100_003).into_keys().collect());
193 assert!(cost.hashes <= 2 * LEAF * 7 && cost.fps <= 2 * (1 + 16 * 7 * cost.rounds), "{cost:?}");
195 assert!(cost.rounds <= 6, "{cost:?}");
196 }
197
198 #[test]
199 fn one_side_empty_and_disjoint_small_sets() {
200 let a = set(0..5);
201 let (want, give, _) = run(&a, &BTreeMap::new());
202 assert_eq!((want.len(), give.len()), (0, 5));
203 let (want, give, _) = run(&set(0..300), &set(300..600));
204 assert_eq!((want.len(), give.len()), (300, 300));
205 }
206}