Skip to main content

bdk_labels/
changeset.rs

1use bip329::{Label, LabelRef};
2use serde::{Deserialize, Serialize};
3use std::collections::{BTreeMap, BTreeSet};
4
5/// Defines the strategy for resolving conflicts when merging two label sets.
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum MergeStrategy {
8    /// The incoming label will overwrite the existing label if they share the same reference.
9    Overwrite,
10    /// The existing label is preserved; the incoming label is ignored if they share the same reference.
11    KeepExisting,
12}
13
14/// An in-memory, deterministic collection of BIP-329 wallet labels.
15///
16/// Backed by a `BTreeMap`, this structure ensures $O(\log n)$ deduplication and
17/// guarantees that labels are deterministically sorted by their reference key.
18/// This prevents noisy diffs when exporting to version-controlled JSONL files.
19#[derive(Debug, Default, Clone, Serialize, Deserialize)]
20pub struct LabelChangeset {
21    labels: BTreeMap<LabelRef, Label>,
22    // Refs that have not been persisted to the wallet's database
23    #[serde(skip)]
24    staged: BTreeSet<LabelRef>,
25}
26
27impl LabelChangeset {
28    /// Creates a new, empty `LabelChangeset`.
29    pub fn new() -> Self {
30        Self {
31            labels: BTreeMap::new(),
32            staged: BTreeSet::new(),
33        }
34    }
35
36    /// Inserts a new label into the changeset making its ref staged.
37    /// If a label with the same reference already exists, it is overwritten.
38    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    /// Removes the label associated with the given target reference, returning it if it existed.
45    pub fn remove(&mut self, target: &LabelRef) -> Option<Label> {
46        self.staged.insert(target.clone());
47        self.labels.remove(target)
48    }
49
50    /// Retrieves a reference to the label associated with the given target.
51    pub fn get(&self, target: &LabelRef) -> Option<&Label> {
52        self.labels.get(target)
53    }
54
55    /// Returns `true` if the changeset contains no labels.
56    pub fn is_empty(&self) -> bool {
57        self.labels.is_empty()
58    }
59
60    /// Returns the total number of labels in the changeset.
61    pub fn len(&self) -> usize {
62        self.labels.len()
63    }
64
65    /// Returns an iterator over the labels, yielded in deterministic order based on their reference.
66    pub fn iter(&self) -> impl Iterator<Item = &Label> {
67        self.labels.values()
68    }
69
70    /// Merges an incoming `LabelChangeset` into the current one, resolving conflicts based on the provided `MergeStrategy`.
71    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    /// Returns the unpersisted diff of the changeset.
86    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    /// Clears the staged label refs
98    pub fn clear_staged(&mut self) {
99        self.staged.clear();
100    }
101
102    /// Checks and returns true if there are changes staged
103    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}