iroh_docs/
heads.rs

1//! Author heads
2
3use std::{collections::BTreeMap, num::NonZeroU64};
4
5use anyhow::Result;
6
7use crate::AuthorId;
8
9type Timestamp = u64;
10
11/// Timestamps of the latest entry for each author.
12#[derive(Debug, Clone, Eq, PartialEq, Default)]
13pub struct AuthorHeads {
14    heads: BTreeMap<AuthorId, Timestamp>,
15}
16
17impl AuthorHeads {
18    /// Insert a new timestamp.
19    pub fn insert(&mut self, author: AuthorId, timestamp: Timestamp) {
20        self.heads
21            .entry(author)
22            .and_modify(|t| *t = (*t).max(timestamp))
23            .or_insert(timestamp);
24    }
25
26    /// Number of author-timestamp pairs.
27    pub fn len(&self) -> usize {
28        self.heads.len()
29    }
30
31    /// Whether this [`AuthorHeads`] is empty.
32    pub fn is_empty(&self) -> bool {
33        self.heads.is_empty()
34    }
35
36    /// Get the timestamp for an author.
37    pub fn get(&self, author: &AuthorId) -> Option<Timestamp> {
38        self.heads.get(author).copied()
39    }
40
41    /// Can this state offer newer stuff to `other`?
42    pub fn has_news_for(&self, other: &Self) -> Option<NonZeroU64> {
43        let mut updates = 0;
44        for (author, ts_ours) in self.iter() {
45            if other
46                .get(author)
47                .map(|ts_theirs| *ts_ours > ts_theirs)
48                .unwrap_or(true)
49            {
50                updates += 1;
51            }
52        }
53        NonZeroU64::new(updates)
54    }
55
56    /// Merge another author head state into this one.
57    pub fn merge(&mut self, other: &Self) {
58        for (a, t) in other.iter() {
59            self.insert(*a, *t);
60        }
61    }
62
63    /// Create an iterator over the entries in this state.
64    pub fn iter(&self) -> std::collections::btree_map::Iter<AuthorId, Timestamp> {
65        self.heads.iter()
66    }
67
68    /// Encode into a byte array with a limited size.
69    ///
70    /// Will skip oldest entries if the size limit is reached.
71    /// Returns a byte array with a maximum length of `size_limit`.
72    pub fn encode(&self, size_limit: Option<usize>) -> Result<Vec<u8>> {
73        let mut by_timestamp = BTreeMap::new();
74        for (author, ts) in self.iter() {
75            by_timestamp.insert(*ts, *author);
76        }
77        let mut items = Vec::new();
78        for (ts, author) in by_timestamp.into_iter().rev() {
79            items.push((ts, author));
80            if let Some(size_limit) = size_limit {
81                if postcard::experimental::serialized_size(&items)? > size_limit {
82                    items.pop();
83                    break;
84                }
85            }
86        }
87        let encoded = postcard::to_stdvec(&items)?;
88        debug_assert!(size_limit.map(|s| encoded.len() <= s).unwrap_or(true));
89        Ok(encoded)
90    }
91
92    /// Decode from byte slice created with [`Self::encode`].
93    pub fn decode(bytes: &[u8]) -> Result<Self> {
94        let items: Vec<(Timestamp, AuthorId)> = postcard::from_bytes(bytes)?;
95        let mut heads = AuthorHeads::default();
96        for (ts, author) in items {
97            heads.insert(author, ts);
98        }
99        Ok(heads)
100    }
101}
102
103impl FromIterator<(AuthorId, Timestamp)> for AuthorHeads {
104    fn from_iter<T: IntoIterator<Item = (AuthorId, Timestamp)>>(iter: T) -> Self {
105        Self {
106            heads: iter.into_iter().collect(),
107        }
108    }
109}
110
111impl FromIterator<(Timestamp, AuthorId)> for AuthorHeads {
112    fn from_iter<T: IntoIterator<Item = (Timestamp, AuthorId)>>(iter: T) -> Self {
113        Self {
114            heads: iter.into_iter().map(|(ts, author)| (author, ts)).collect(),
115        }
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122    use crate::Record;
123    #[test]
124    fn author_heads_encode_decode() -> Result<()> {
125        let mut heads = AuthorHeads::default();
126        let start = Record::empty_current().timestamp();
127        for i in 0..10u64 {
128            heads.insert(AuthorId::from(&[i as u8; 32]), start + i);
129        }
130        let encoded = heads.encode(Some(256))?;
131        let decoded = AuthorHeads::decode(&encoded)?;
132        assert_eq!(decoded.len(), 6);
133        let expected: AuthorHeads = (0u64..6)
134            .map(|n| (AuthorId::from(&[9 - n as u8; 32]), start + (9 - n)))
135            .collect();
136        assert_eq!(expected, decoded);
137        Ok(())
138    }
139
140    #[test]
141    fn author_heads_compare() -> Result<()> {
142        let a = [
143            (AuthorId::from(&[0u8; 32]), 5),
144            (AuthorId::from(&[1u8; 32]), 7),
145        ];
146        let b = [
147            (AuthorId::from(&[0u8; 32]), 4),
148            (AuthorId::from(&[1u8; 32]), 6),
149            (AuthorId::from(&[2u8; 32]), 7),
150        ];
151        let a: AuthorHeads = a.into_iter().collect();
152        let b: AuthorHeads = b.into_iter().collect();
153        assert_eq!(a.has_news_for(&b), NonZeroU64::new(2));
154        assert_eq!(b.has_news_for(&a), NonZeroU64::new(1));
155        Ok(())
156    }
157}