1use std::{collections::BTreeMap, num::NonZeroU64};
4
5use anyhow::Result;
6
7use crate::AuthorId;
8
9type Timestamp = u64;
10
11#[derive(Debug, Clone, Eq, PartialEq, Default)]
13pub struct AuthorHeads {
14 heads: BTreeMap<AuthorId, Timestamp>,
15}
16
17impl AuthorHeads {
18 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 pub fn len(&self) -> usize {
28 self.heads.len()
29 }
30
31 pub fn is_empty(&self) -> bool {
33 self.heads.is_empty()
34 }
35
36 pub fn get(&self, author: &AuthorId) -> Option<Timestamp> {
38 self.heads.get(author).copied()
39 }
40
41 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 pub fn merge(&mut self, other: &Self) {
58 for (a, t) in other.iter() {
59 self.insert(*a, *t);
60 }
61 }
62
63 pub fn iter(&self) -> std::collections::btree_map::Iter<AuthorId, Timestamp> {
65 self.heads.iter()
66 }
67
68 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 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}