1use bip329::{Label, LabelRef};
2use serde::{Deserialize, Serialize};
3use std::collections::{BTreeMap, BTreeSet};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum MergeStrategy {
8 Overwrite,
10 KeepExisting,
12}
13
14#[derive(Debug, Default, Clone, Serialize, Deserialize)]
20pub struct LabelChangeset {
21 labels: BTreeMap<LabelRef, Label>,
22 #[serde(skip)]
24 staged: BTreeSet<LabelRef>,
25}
26
27impl LabelChangeset {
28 pub fn new() -> Self {
30 Self {
31 labels: BTreeMap::new(),
32 staged: BTreeSet::new(),
33 }
34 }
35
36 pub fn insert(&mut self, label: Label) {
39 let target = label.ref_();
40 self.staged.insert(target.clone());
41 self.labels.insert(target, label);
42 }
43
44 pub fn remove(&mut self, target: &LabelRef) -> Option<Label> {
46 self.staged.insert(target.clone());
47 self.labels.remove(target)
48 }
49
50 pub fn get(&self, target: &LabelRef) -> Option<&Label> {
52 self.labels.get(target)
53 }
54
55 pub fn is_empty(&self) -> bool {
57 self.labels.is_empty()
58 }
59
60 pub fn len(&self) -> usize {
62 self.labels.len()
63 }
64
65 pub fn iter(&self) -> impl Iterator<Item = &Label> {
67 self.labels.values()
68 }
69
70 pub fn merge(&mut self, incoming: LabelChangeset, strategy: MergeStrategy) {
72 for (_, incoming_label) in incoming.labels {
73 let target = incoming_label.ref_();
74 match strategy {
75 MergeStrategy::Overwrite => {
76 self.labels.insert(target, incoming_label);
77 }
78 MergeStrategy::KeepExisting => {
79 self.labels.entry(target).or_insert(incoming_label);
80 }
81 }
82 }
83 }
84
85 pub fn diff(&self) -> LabelChangeset {
87 let mut d = LabelChangeset::new();
88
89 for target in &self.staged {
90 if let Some(label) = self.labels.get(target) {
91 d.labels.insert(target.clone(), label.clone());
92 }
93 }
94 d
95 }
96
97 pub fn clear_staged(&mut self) {
99 self.staged.clear();
100 }
101
102 pub fn has_staged_changes(&self) -> bool {
104 !self.staged.is_empty()
105 }
106}
107
108#[cfg(test)]
109mod tests {
110 use super::*;
111 use bip329::{AddressRecord, Label, TransactionRecord};
112 use bitcoin::{Address, Txid};
113 use std::str::FromStr;
114
115 #[test]
116 fn test_basic_crud_operations() {
117 let mut changeset = LabelChangeset::new();
118 assert!(changeset.is_empty(), "Expected new changeset to be empty");
119 assert_eq!(changeset.len(), 0);
120
121 let dummy_txid =
122 Txid::from_str("0000000000000000000000000000000000000000000000000000000000000000")
123 .unwrap();
124
125 let dummy_label = Label::Transaction(TransactionRecord {
126 ref_: dummy_txid,
127 label: Some("Machinery".to_string()),
128 origin: None,
129 });
130
131 changeset.insert(dummy_label.clone());
132
133 assert_eq!(changeset.len(), 1);
134 assert_eq!(changeset.get(&dummy_label.ref_()), Some(&dummy_label));
135
136 let updated_label = Label::Transaction(TransactionRecord {
137 ref_: dummy_txid,
138 label: Some("Heavy Machinery".to_string()),
139 origin: None,
140 });
141 changeset.insert(updated_label.clone());
142
143 assert_eq!(changeset.len(), 1);
144 assert_eq!(changeset.get(&dummy_label.ref_()), Some(&updated_label));
145
146 let removed_label = changeset.remove(&dummy_label.ref_());
147
148 assert_eq!(removed_label, Some(updated_label));
149 assert!(changeset.is_empty());
150 assert_eq!(changeset.get(&dummy_label.ref_()), None);
151 }
152
153 #[test]
154 fn test_implicit_deduplication() {
155 let mut changeset = LabelChangeset::new();
156
157 let dummy_txid =
158 Txid::from_str("0000000000000000000000000000000000000000000000000000000000000000")
159 .unwrap();
160
161 let dummy_label = Label::Transaction(TransactionRecord {
162 ref_: dummy_txid,
163 label: Some("Machinery".to_string()),
164 origin: None,
165 });
166
167 changeset.insert(dummy_label.clone());
168 changeset.insert(dummy_label.clone());
169 changeset.insert(dummy_label.clone());
170
171 assert_eq!(changeset.len(), 1);
172 }
173
174 #[test]
175 fn test_merge_strategy_overwrite() {
176 let mut base_changeset = LabelChangeset::new();
177 let mut incoming_changeset = LabelChangeset::new();
178
179 let dummy_txid =
180 Txid::from_str("0000000000000000000000000000000000000000000000000000000000000000")
181 .unwrap();
182
183 let dummy_label = Label::Transaction(TransactionRecord {
184 ref_: dummy_txid,
185 label: Some("Machinery".to_string()),
186 origin: None,
187 });
188
189 let another_dummy_label = Label::Transaction(TransactionRecord {
190 ref_: dummy_txid,
191 label: Some("Heavy Machinery".to_string()),
192 origin: None,
193 });
194
195 base_changeset.insert(dummy_label.clone());
196
197 incoming_changeset.insert(another_dummy_label.clone());
198
199 base_changeset.merge(incoming_changeset, MergeStrategy::Overwrite);
200
201 assert_eq!(
202 base_changeset.get(&dummy_label.ref_()),
203 Some(&another_dummy_label)
204 );
205 }
206
207 #[test]
208 fn test_merge_strategy_keep_existing() {
209 let mut base_changeset = LabelChangeset::new();
210 let mut incoming_changeset = LabelChangeset::new();
211
212 let dummy_txid =
213 Txid::from_str("0000000000000000000000000000000000000000000000000000000000000000")
214 .unwrap();
215
216 let dummy_label = Label::Transaction(TransactionRecord {
217 ref_: dummy_txid,
218 label: Some("Machinery".to_string()),
219 origin: None,
220 });
221
222 let another_dummy_label = Label::Transaction(TransactionRecord {
223 ref_: dummy_txid,
224 label: Some("Heavy Machinery".to_string()),
225 origin: None,
226 });
227
228 base_changeset.insert(dummy_label.clone());
229
230 incoming_changeset.insert(another_dummy_label.clone());
231
232 base_changeset.merge(incoming_changeset, MergeStrategy::KeepExisting);
233
234 assert_eq!(base_changeset.get(&dummy_label.ref_()), Some(&dummy_label));
235 }
236
237 #[test]
238 fn test_non_overlapping_merges() {
239 let mut base_changeset = LabelChangeset::new();
240 let mut incoming_changeset = LabelChangeset::new();
241
242 let dummy_txid =
243 Txid::from_str("0000000000000000000000000000000000000000000000000000000000000000")
244 .unwrap();
245
246 let another_dummy_txid =
247 Txid::from_str("0000000000000000000000000000000000000000000000000000000000000110")
248 .unwrap();
249
250 let dummy_label = Label::Transaction(TransactionRecord {
251 ref_: dummy_txid,
252 label: Some("Machinery".to_string()),
253 origin: None,
254 });
255
256 let another_dummy_label = Label::Transaction(TransactionRecord {
257 ref_: another_dummy_txid,
258 label: Some("Heavy Machinery".to_string()),
259 origin: None,
260 });
261
262 base_changeset.insert(dummy_label.clone());
263
264 incoming_changeset.insert(another_dummy_label.clone());
265
266 base_changeset.merge(incoming_changeset, MergeStrategy::Overwrite);
267
268 assert_eq!(base_changeset.len(), 2);
269 assert_eq!(base_changeset.get(&dummy_label.ref_()), Some(&dummy_label));
270 assert_eq!(
271 base_changeset.get(&another_dummy_label.ref_()),
272 Some(&another_dummy_label)
273 );
274 }
275
276 #[test]
277 fn test_empty_merges() {
278 let mut populated_base_changeset = LabelChangeset::new();
279 let empty_incoming_changeset = LabelChangeset::new();
280
281 let dummy_txid =
282 Txid::from_str("0000000000000000000000000000000000000000000000000000000000000000")
283 .unwrap();
284
285 let dummy_label = Label::Transaction(TransactionRecord {
286 ref_: dummy_txid,
287 label: Some("Machinery".to_string()),
288 origin: None,
289 });
290
291 populated_base_changeset.insert(dummy_label.clone());
292
293 populated_base_changeset.merge(empty_incoming_changeset, MergeStrategy::Overwrite);
294
295 assert_eq!(populated_base_changeset.len(), 1);
296 assert_eq!(
297 populated_base_changeset.get(&dummy_label.ref_()),
298 Some(&dummy_label)
299 );
300
301 let mut empty_base_changeset = LabelChangeset::new();
302 let mut populated_incoming_changeset = LabelChangeset::new();
303
304 let dummy_address =
305 Address::from_str("bc1p0dq0tzg2r780hldthn5mrznmpxsxc0jux5f20fwj0z3wqxxk6fpqm7q0va")
306 .expect("Failed to parse address");
307
308 let another_dummy_label = Label::Address(AddressRecord {
309 ref_: dummy_address,
310 label: Some("Heavy Machinery".to_string()),
311 });
312
313 populated_incoming_changeset.insert(another_dummy_label.clone());
314
315 empty_base_changeset.merge(populated_incoming_changeset, MergeStrategy::Overwrite);
316
317 assert_eq!(empty_base_changeset.len(), 1);
318 assert_eq!(
319 empty_base_changeset.get(&another_dummy_label.ref_()),
320 Some(&another_dummy_label)
321 );
322 }
323
324 #[test]
325 fn test_diff_only_contains_staged_entries_and_clears_on_demand() {
326 let mut changeset = LabelChangeset::new();
327
328 let first_txid =
329 Txid::from_str("0000000000000000000000000000000000000000000000000000000000000000")
330 .unwrap();
331 let second_txid =
332 Txid::from_str("0000000000000000000000000000000000000000000000000000000000000110")
333 .unwrap();
334
335 let first_label = Label::Transaction(TransactionRecord {
336 ref_: first_txid,
337 label: Some("First Transaction".to_string()),
338 origin: None,
339 });
340 let second_label = Label::Transaction(TransactionRecord {
341 ref_: second_txid,
342 label: Some("Second Transaction".to_string()),
343 origin: None,
344 });
345
346 changeset.insert(first_label.clone());
347 assert!(changeset.has_staged_changes());
348
349 let first_diff = changeset.diff();
350
351 assert_eq!(first_diff.labels.len(), 1);
352
353 assert_eq!(first_diff.get(&first_label.ref_()), Some(&first_label));
354
355 changeset.clear_staged();
356 assert!(!changeset.has_staged_changes());
357
358 let empty_diff = changeset.diff();
359 assert!(empty_diff.is_empty());
360 assert_eq!(first_diff.labels.len(), 1);
361
362 changeset.insert(second_label.clone());
363 assert!(changeset.has_staged_changes());
364 let second_diff = changeset.diff();
365 assert_eq!(second_diff.labels.len(), 1);
366 assert_eq!(second_diff.get(&second_label.ref_()), Some(&second_label));
367 assert_eq!(second_diff.get(&first_label.ref_()), None);
368 }
369}