1use std::collections::VecDeque;
2use std::time::{SystemTime, UNIX_EPOCH};
3
4use indexmap::{IndexMap, IndexSet};
5
6const BINARY_FORMAT_VERSION: u64 = 2;
7const BINARY_CHUNK_SIZE: usize = 16 * 1024;
8const MAX_FUTURE_EVENT_SECONDS: u64 = 10 * 60;
9const UNKNOWN_FOLLOW_DISTANCE: u32 = 1000;
10
11#[derive(Debug, thiserror::Error)]
12pub enum SocialGraphError {
13 #[error("cannot store empty or whitespace-only strings")]
14 EmptyString,
15 #[error("invalid id {0}")]
16 InvalidId(u32),
17 #[error("invalid binary version {0}")]
18 InvalidVersion(u64),
19 #[error("unexpected end of binary data")]
20 UnexpectedEof,
21 #[error("invalid hex string for id {0}: {1}")]
22 InvalidHex(u32, String),
23}
24
25pub type Result<T> = std::result::Result<T, SocialGraphError>;
26
27pub trait SocialGraphBackend {
28 type Error: std::error::Error + Send + Sync + 'static;
29
30 fn get_root(&self) -> std::result::Result<String, Self::Error>;
31
32 fn set_root(&mut self, root: &str) -> std::result::Result<(), Self::Error>;
33
34 fn handle_event(
35 &mut self,
36 event: &NostrEvent,
37 allow_unknown_authors: bool,
38 overmute_threshold: f64,
39 ) -> std::result::Result<(), Self::Error>;
40
41 fn get_follow_distance(&self, user: &str) -> std::result::Result<u32, Self::Error>;
42
43 fn is_following(
44 &self,
45 follower: &str,
46 followed_user: &str,
47 ) -> std::result::Result<bool, Self::Error>;
48
49 fn get_followed_by_user(&self, user: &str) -> std::result::Result<Vec<String>, Self::Error>;
50
51 fn get_followers_by_user(&self, user: &str) -> std::result::Result<Vec<String>, Self::Error>;
52
53 fn get_muted_by_user(&self, user: &str) -> std::result::Result<Vec<String>, Self::Error>;
54
55 fn get_user_muted_by(&self, user: &str) -> std::result::Result<Vec<String>, Self::Error>;
56
57 fn get_follow_list_created_at(
58 &self,
59 user: &str,
60 ) -> std::result::Result<Option<u64>, Self::Error>;
61
62 fn get_mute_list_created_at(&self, user: &str)
63 -> std::result::Result<Option<u64>, Self::Error>;
64
65 fn is_overmuted(&self, user: &str, threshold: f64) -> std::result::Result<bool, Self::Error>;
66
67 fn flush(&mut self) -> std::result::Result<(), Self::Error> {
68 Ok(())
69 }
70
71 fn has_unflushed_changes(&self) -> bool {
72 false
73 }
74}
75
76#[derive(Debug, Clone, PartialEq, Eq)]
77pub struct SocialGraphState {
78 pub root: String,
79 pub unique_ids: Vec<(String, u32)>,
80 pub follow_distance_by_user: Vec<(u32, u32)>,
81 pub users_by_follow_distance: Vec<(u32, Vec<u32>)>,
82 pub followed_by_user: Vec<(u32, Vec<u32>)>,
83 pub followers_by_user: Vec<(u32, Vec<u32>)>,
84 pub follow_list_created_at: Vec<(u32, u64)>,
85 pub muted_by_user: Vec<(u32, Vec<u32>)>,
86 pub user_muted_by: Vec<(u32, Vec<u32>)>,
87 pub mute_list_created_at: Vec<(u32, u64)>,
88}
89
90#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
91pub struct BinaryBudget {
92 pub max_nodes: Option<usize>,
93 pub max_edges: Option<usize>,
94 pub max_distance: Option<u32>,
95 pub max_edges_per_node: Option<usize>,
96}
97
98impl BinaryBudget {
99 fn has_active_limits(self) -> bool {
100 self.max_nodes.is_some_and(|value| value > 0)
101 || self.max_edges.is_some_and(|value| value > 0)
102 || self.max_distance.is_some()
103 || self.max_edges_per_node.is_some_and(|value| value > 0)
104 }
105}
106
107#[derive(Debug, Clone, PartialEq, Eq)]
108pub struct GraphStats {
109 pub users: usize,
110 pub follows: usize,
111 pub mutes: usize,
112 pub size_by_distance: IndexMap<u32, usize>,
113}
114
115#[derive(Debug, Clone, PartialEq, Eq)]
116pub struct NostrEvent {
117 pub created_at: u64,
118 pub content: String,
119 pub tags: Vec<Vec<String>>,
120 pub kind: u32,
121 pub pubkey: String,
122 pub id: String,
123 pub sig: String,
124}
125
126#[derive(Debug, Clone, Default)]
127struct UniqueIds {
128 str_to_unique_id: IndexMap<String, u32>,
129 unique_id_to_str: IndexMap<u32, String>,
130 current_unique_id: u32,
131}
132
133impl UniqueIds {
134 fn id(&mut self, value: &str) -> Result<u32> {
135 if value.trim().is_empty() {
136 return Err(SocialGraphError::EmptyString);
137 }
138 if let Some(id) = self.str_to_unique_id.get(value) {
139 return Ok(*id);
140 }
141 let id = self.current_unique_id;
142 self.current_unique_id = self.current_unique_id.saturating_add(1);
143 self.str_to_unique_id.insert(value.to_owned(), id);
144 self.unique_id_to_str.insert(id, value.to_owned());
145 Ok(id)
146 }
147
148 fn existing_id(&self, value: &str) -> Option<u32> {
149 self.str_to_unique_id.get(value).copied()
150 }
151
152 fn str(&self, id: u32) -> Result<&str> {
153 self.unique_id_to_str
154 .get(&id)
155 .map(String::as_str)
156 .ok_or(SocialGraphError::InvalidId(id))
157 }
158
159 fn clear(&mut self) {
160 self.str_to_unique_id.clear();
161 self.unique_id_to_str.clear();
162 self.current_unique_id = 0;
163 }
164
165 fn insert_with_id(&mut self, value: String, id: u32) {
166 self.current_unique_id = self.current_unique_id.max(id.saturating_add(1));
167 self.str_to_unique_id.insert(value.clone(), id);
168 self.unique_id_to_str.insert(id, value);
169 }
170
171 fn remove(&mut self, id: u32) {
172 if let Some(value) = self.unique_id_to_str.shift_remove(&id) {
173 self.str_to_unique_id.shift_remove(&value);
174 }
175 }
176}
177
178#[derive(Debug, Clone, Default)]
179pub struct SocialGraph {
180 root: u32,
181 follow_distance_by_user: IndexMap<u32, u32>,
182 users_by_follow_distance: IndexMap<u32, IndexSet<u32>>,
183 followed_by_user: IndexMap<u32, IndexSet<u32>>,
184 followers_by_user: IndexMap<u32, IndexSet<u32>>,
185 follow_list_created_at: IndexMap<u32, u64>,
186 muted_by_user: IndexMap<u32, IndexSet<u32>>,
187 user_muted_by: IndexMap<u32, IndexSet<u32>>,
188 mute_list_created_at: IndexMap<u32, u64>,
189 ids: UniqueIds,
190}
191
192#[derive(Debug)]
193struct BinaryPlan {
194 used_ids: IndexSet<u32>,
195 follow_edge_count: IndexMap<u32, usize>,
196 mute_edge_count: IndexMap<u32, usize>,
197 follow_owners: Vec<u32>,
198 mute_owners: Vec<u32>,
199}
200
201#[derive(Debug, Clone, Copy)]
202struct PotentialEdge {
203 owner: u32,
204 target: u32,
205 is_follow: bool,
206}
207
208impl SocialGraph {
209 pub fn new(root: &str) -> Self {
210 let mut ids = UniqueIds::default();
211 let root_id = ids.id(root).expect("root must not be empty");
212 let mut follow_distance_by_user = IndexMap::new();
213 follow_distance_by_user.insert(root_id, 0);
214 let mut users_by_follow_distance = IndexMap::new();
215 let mut bucket = IndexSet::new();
216 bucket.insert(root_id);
217 users_by_follow_distance.insert(0, bucket);
218 Self {
219 root: root_id,
220 follow_distance_by_user,
221 users_by_follow_distance,
222 followed_by_user: IndexMap::new(),
223 followers_by_user: IndexMap::new(),
224 follow_list_created_at: IndexMap::new(),
225 muted_by_user: IndexMap::new(),
226 user_muted_by: IndexMap::new(),
227 mute_list_created_at: IndexMap::new(),
228 ids,
229 }
230 }
231
232 pub fn get_root(&self) -> &str {
233 self.ids
234 .str(self.root)
235 .expect("root id should always exist")
236 }
237
238 pub fn set_root(&mut self, root: &str) -> Result<()> {
239 let root_id = self.ids.id(root)?;
240 if root_id == self.root {
241 return Ok(());
242 }
243 self.root = root_id;
244 self.recalculate_follow_distances();
245 Ok(())
246 }
247
248 pub fn handle_event(
249 &mut self,
250 event: &NostrEvent,
251 allow_unknown_authors: bool,
252 overmute_threshold: f64,
253 ) {
254 if !matches!(event.kind, 3 | 10_000) {
255 return;
256 }
257
258 let now = SystemTime::now()
259 .duration_since(UNIX_EPOCH)
260 .unwrap_or_default()
261 .as_secs();
262 if event.created_at > now.saturating_add(MAX_FUTURE_EVENT_SECONDS) {
263 return;
264 }
265
266 let author = if allow_unknown_authors {
267 let Ok(author) = self.ids.id(&event.pubkey) else {
268 return;
269 };
270 author
271 } else {
272 let Some(author) = self.ids.existing_id(&event.pubkey) else {
273 return;
274 };
275 if !self.follow_distance_by_user.contains_key(&author) {
276 return;
277 }
278 author
279 };
280
281 if self.is_overmuted(&event.pubkey, overmute_threshold) {
282 return;
283 }
284
285 match event.kind {
286 3 => self.handle_follow_list(author, event.created_at, &event.tags),
287 10_000 => self.handle_mute_list(author, event.created_at, &event.tags),
288 _ => {}
289 }
290 }
291
292 pub fn recalculate_follow_distances(&mut self) {
293 self.follow_distance_by_user.clear();
294 self.users_by_follow_distance.clear();
295 self.follow_distance_by_user.insert(self.root, 0);
296 let mut root_bucket = IndexSet::new();
297 root_bucket.insert(self.root);
298 self.users_by_follow_distance.insert(0, root_bucket);
299
300 let mut queue = VecDeque::from([self.root]);
301 while let Some(user) = queue.pop_front() {
302 let Some(distance) = self.follow_distance_by_user.get(&user).copied() else {
303 continue;
304 };
305 let Some(followed) = self.followed_by_user.get(&user).cloned() else {
306 continue;
307 };
308 let next_distance = distance.saturating_add(1);
309 for target in followed {
310 if self.follow_distance_by_user.contains_key(&target) {
311 continue;
312 }
313 self.follow_distance_by_user.insert(target, next_distance);
314 self.users_by_follow_distance
315 .entry(next_distance)
316 .or_default()
317 .insert(target);
318 queue.push_back(target);
319 }
320 }
321 }
322
323 pub fn get_follow_distance(&self, user: &str) -> u32 {
324 let Some(user_id) = self.ids.existing_id(user) else {
325 return UNKNOWN_FOLLOW_DISTANCE;
326 };
327 self.follow_distance_by_user
328 .get(&user_id)
329 .copied()
330 .unwrap_or(UNKNOWN_FOLLOW_DISTANCE)
331 }
332
333 pub fn is_following(&self, follower: &str, followed_user: &str) -> bool {
334 let Some(follower_id) = self.ids.existing_id(follower) else {
335 return false;
336 };
337 let Some(followed_id) = self.ids.existing_id(followed_user) else {
338 return false;
339 };
340 self.followed_by_user
341 .get(&follower_id)
342 .is_some_and(|set| set.contains(&followed_id))
343 }
344
345 pub fn get_followed_by_user(&self, user: &str) -> Vec<String> {
346 let Some(user_id) = self.ids.existing_id(user) else {
347 return Vec::new();
348 };
349 self.followed_by_user
350 .get(&user_id)
351 .into_iter()
352 .flat_map(|set| set.iter())
353 .filter_map(|id| self.ids.str(*id).ok().map(ToOwned::to_owned))
354 .collect()
355 }
356
357 pub fn get_followers_by_user(&self, user: &str) -> Vec<String> {
358 let Some(user_id) = self.ids.existing_id(user) else {
359 return Vec::new();
360 };
361 self.followers_by_user
362 .get(&user_id)
363 .into_iter()
364 .flat_map(|set| set.iter())
365 .filter_map(|id| self.ids.str(*id).ok().map(ToOwned::to_owned))
366 .collect()
367 }
368
369 pub fn get_muted_by_user(&self, user: &str) -> Vec<String> {
370 let Some(user_id) = self.ids.existing_id(user) else {
371 return Vec::new();
372 };
373 self.muted_by_user
374 .get(&user_id)
375 .into_iter()
376 .flat_map(|set| set.iter())
377 .filter_map(|id| self.ids.str(*id).ok().map(ToOwned::to_owned))
378 .collect()
379 }
380
381 pub fn get_user_muted_by(&self, user: &str) -> Vec<String> {
382 let Some(user_id) = self.ids.existing_id(user) else {
383 return Vec::new();
384 };
385 self.user_muted_by
386 .get(&user_id)
387 .into_iter()
388 .flat_map(|set| set.iter())
389 .filter_map(|id| self.ids.str(*id).ok().map(ToOwned::to_owned))
390 .collect()
391 }
392
393 pub fn get_follow_list_created_at(&self, user: &str) -> Option<u64> {
394 let user_id = self.ids.existing_id(user)?;
395 self.follow_list_created_at.get(&user_id).copied()
396 }
397
398 pub fn get_mute_list_created_at(&self, user: &str) -> Option<u64> {
399 let user_id = self.ids.existing_id(user)?;
400 self.mute_list_created_at.get(&user_id).copied()
401 }
402
403 pub fn size(&self) -> GraphStats {
404 let follows = self
405 .followed_by_user
406 .values()
407 .map(IndexSet::len)
408 .sum::<usize>();
409 let mutes = self
410 .muted_by_user
411 .values()
412 .map(IndexSet::len)
413 .sum::<usize>();
414 let size_by_distance = self
415 .users_by_follow_distance
416 .iter()
417 .map(|(distance, users)| (*distance, users.len()))
418 .collect();
419
420 GraphStats {
421 users: if self.follow_distance_by_user.is_empty() {
422 self.ids.unique_id_to_str.len()
423 } else {
424 self.follow_distance_by_user.len()
425 },
426 follows,
427 mutes,
428 size_by_distance,
429 }
430 }
431
432 pub fn get_users_by_follow_distance(&self, distance: u32) -> Vec<String> {
433 self.users_by_follow_distance
434 .get(&distance)
435 .into_iter()
436 .flat_map(|users| users.iter())
437 .filter_map(|id| self.ids.str(*id).ok().map(ToOwned::to_owned))
438 .collect()
439 }
440
441 pub fn users_in_distance_order(&self, up_to_distance: Option<u32>) -> Vec<String> {
442 let mut distances: Vec<u32> = self.users_by_follow_distance.keys().copied().collect();
443 distances.sort_unstable();
444 let mut users = Vec::new();
445 for distance in distances {
446 if up_to_distance.is_some_and(|max_distance| distance > max_distance) {
447 break;
448 }
449 users.extend(self.get_users_by_follow_distance(distance));
450 }
451 users
452 }
453
454 pub fn remove_muted_not_followed_users(&mut self) -> usize {
455 let mut has_followers = IndexSet::new();
456 for followed_users in self.followed_by_user.values() {
457 for user in followed_users {
458 has_followers.insert(*user);
459 }
460 }
461
462 let users_to_remove: Vec<u32> = self
463 .user_muted_by
464 .iter()
465 .filter_map(|(user, muters)| {
466 if *user != self.root && !muters.is_empty() && !has_followers.contains(user) {
467 Some(*user)
468 } else {
469 None
470 }
471 })
472 .collect();
473
474 if users_to_remove.is_empty() {
475 return 0;
476 }
477
478 for user in users_to_remove.iter().copied() {
479 if let Some(distance) = self.follow_distance_by_user.shift_remove(&user)
480 && let Some(bucket) = self.users_by_follow_distance.get_mut(&distance)
481 {
482 bucket.shift_remove(&user);
483 }
484 self.followed_by_user.shift_remove(&user);
485 self.followers_by_user.shift_remove(&user);
486 self.follow_list_created_at.shift_remove(&user);
487 self.muted_by_user.shift_remove(&user);
488 self.user_muted_by.shift_remove(&user);
489 self.mute_list_created_at.shift_remove(&user);
490 self.ids.remove(user);
491 }
492
493 for followed_users in self.followed_by_user.values_mut() {
494 for user in &users_to_remove {
495 followed_users.shift_remove(user);
496 }
497 }
498 for followers in self.followers_by_user.values_mut() {
499 for user in &users_to_remove {
500 followers.shift_remove(user);
501 }
502 }
503 for muted_users in self.muted_by_user.values_mut() {
504 for user in &users_to_remove {
505 muted_users.shift_remove(user);
506 }
507 }
508 for muters in self.user_muted_by.values_mut() {
509 for user in &users_to_remove {
510 muters.shift_remove(user);
511 }
512 }
513
514 users_to_remove.len()
515 }
516
517 pub fn export_state(&self) -> SocialGraphState {
518 SocialGraphState {
519 root: self.get_root().to_string(),
520 unique_ids: self
521 .ids
522 .unique_id_to_str
523 .iter()
524 .map(|(id, value)| (value.clone(), *id))
525 .collect(),
526 follow_distance_by_user: self
527 .follow_distance_by_user
528 .iter()
529 .map(|(id, distance)| (*id, *distance))
530 .collect(),
531 users_by_follow_distance: self
532 .users_by_follow_distance
533 .iter()
534 .map(|(distance, users)| (*distance, users.iter().copied().collect()))
535 .collect(),
536 followed_by_user: self
537 .followed_by_user
538 .iter()
539 .map(|(user, followed)| (*user, followed.iter().copied().collect()))
540 .collect(),
541 followers_by_user: self
542 .followers_by_user
543 .iter()
544 .map(|(user, followers)| (*user, followers.iter().copied().collect()))
545 .collect(),
546 follow_list_created_at: self
547 .follow_list_created_at
548 .iter()
549 .map(|(user, created_at)| (*user, *created_at))
550 .collect(),
551 muted_by_user: self
552 .muted_by_user
553 .iter()
554 .map(|(user, muted)| (*user, muted.iter().copied().collect()))
555 .collect(),
556 user_muted_by: self
557 .user_muted_by
558 .iter()
559 .map(|(user, muters)| (*user, muters.iter().copied().collect()))
560 .collect(),
561 mute_list_created_at: self
562 .mute_list_created_at
563 .iter()
564 .map(|(user, created_at)| (*user, *created_at))
565 .collect(),
566 }
567 }
568
569 pub fn from_state(state: SocialGraphState) -> Result<Self> {
570 let mut graph = Self::new(&state.root);
571 graph.ids.clear();
572 graph.follow_distance_by_user.clear();
573 graph.users_by_follow_distance.clear();
574 graph.followed_by_user.clear();
575 graph.followers_by_user.clear();
576 graph.follow_list_created_at.clear();
577 graph.muted_by_user.clear();
578 graph.user_muted_by.clear();
579 graph.mute_list_created_at.clear();
580
581 for (value, id) in state.unique_ids {
582 graph.ids.insert_with_id(value, id);
583 }
584
585 graph.root = graph.ids.id(&state.root)?;
586
587 for (user, distance) in state.follow_distance_by_user {
588 graph.follow_distance_by_user.insert(user, distance);
589 }
590 for (distance, users) in state.users_by_follow_distance {
591 graph
592 .users_by_follow_distance
593 .insert(distance, users.into_iter().collect());
594 }
595 for (user, followed) in state.followed_by_user {
596 graph
597 .followed_by_user
598 .insert(user, followed.into_iter().collect());
599 }
600 for (user, followers) in state.followers_by_user {
601 graph
602 .followers_by_user
603 .insert(user, followers.into_iter().collect());
604 }
605 for (user, created_at) in state.follow_list_created_at {
606 graph.follow_list_created_at.insert(user, created_at);
607 }
608 for (user, muted) in state.muted_by_user {
609 graph
610 .muted_by_user
611 .insert(user, muted.into_iter().collect());
612 }
613 for (user, muters) in state.user_muted_by {
614 graph
615 .user_muted_by
616 .insert(user, muters.into_iter().collect());
617 }
618 for (user, created_at) in state.mute_list_created_at {
619 graph.mute_list_created_at.insert(user, created_at);
620 }
621
622 Ok(graph)
623 }
624
625 pub fn to_binary(&self) -> Result<Vec<u8>> {
626 self.to_binary_with_budget(BinaryBudget::default())
627 }
628
629 pub fn to_binary_with_budget(&self, budget: BinaryBudget) -> Result<Vec<u8>> {
630 let plan = self.plan_binary(budget)?;
631 self.serialize_binary(&plan)
632 }
633
634 pub fn to_binary_chunks(&self) -> Result<Vec<Vec<u8>>> {
635 self.to_binary_chunks_with_budget(BinaryBudget::default())
636 }
637
638 pub fn to_binary_chunks_with_budget(&self, budget: BinaryBudget) -> Result<Vec<Vec<u8>>> {
639 let binary = self.to_binary_with_budget(budget)?;
640 Ok(binary
641 .chunks(BINARY_CHUNK_SIZE)
642 .map(|chunk| chunk.to_vec())
643 .collect())
644 }
645
646 fn serialize_binary(&self, plan: &BinaryPlan) -> Result<Vec<u8>> {
647 let BinaryPlan {
648 used_ids,
649 follow_edge_count,
650 mute_edge_count,
651 follow_owners,
652 mute_owners,
653 } = plan;
654
655 let mut out = Vec::new();
656 write_varint(&mut out, BINARY_FORMAT_VERSION);
657 write_varint(&mut out, used_ids.len() as u64);
658 for id in used_ids.iter().copied() {
659 let key = self.ids.str(id)?;
660 let bytes = decode_hex_32(key, id)?;
661 out.extend_from_slice(&bytes);
662 write_varint(&mut out, id as u64);
663 }
664
665 write_varint(&mut out, follow_owners.len() as u64);
666 for owner in follow_owners.iter().copied() {
667 let limit = follow_edge_count
668 .get(&owner)
669 .copied()
670 .expect("follow owner must have edge count");
671 write_varint(&mut out, owner as u64);
672 write_varint(
673 &mut out,
674 self.follow_list_created_at
675 .get(&owner)
676 .copied()
677 .unwrap_or(0),
678 );
679 write_varint(&mut out, limit as u64);
680 if let Some(targets) = self.followed_by_user.get(&owner) {
681 for target in targets.iter().take(limit) {
682 write_varint(&mut out, *target as u64);
683 }
684 }
685 }
686
687 write_varint(&mut out, mute_owners.len() as u64);
688 for owner in mute_owners.iter().copied() {
689 let limit = mute_edge_count
690 .get(&owner)
691 .copied()
692 .expect("mute owner must have edge count");
693 write_varint(&mut out, owner as u64);
694 write_varint(
695 &mut out,
696 self.mute_list_created_at.get(&owner).copied().unwrap_or(0),
697 );
698 write_varint(&mut out, limit as u64);
699 if let Some(targets) = self.muted_by_user.get(&owner) {
700 for target in targets.iter().take(limit) {
701 write_varint(&mut out, *target as u64);
702 }
703 }
704 }
705
706 Ok(out)
707 }
708
709 fn plan_binary(&self, budget: BinaryBudget) -> Result<BinaryPlan> {
710 if !budget.has_active_limits() {
711 return Ok(self.plan_full_binary());
712 }
713
714 let max_nodes = budget.max_nodes.filter(|value| *value > 0);
715 let max_edges = budget.max_edges.filter(|value| *value > 0);
716 let max_edges_per_node = budget.max_edges_per_node.filter(|value| *value > 0);
717
718 let mut distances: Vec<u32> = self.users_by_follow_distance.keys().copied().collect();
719 distances.sort_unstable();
720
721 let mut potential_edges = Vec::new();
722 for distance in distances {
723 if budget
724 .max_distance
725 .is_some_and(|max_distance| distance > max_distance)
726 {
727 continue;
728 }
729 let Some(users) = self.users_by_follow_distance.get(&distance) else {
730 continue;
731 };
732
733 for owner in users.iter().copied() {
734 let mut owner_edge_count = 0usize;
735
736 if let Some(followed_users) = self.followed_by_user.get(&owner) {
737 for target in followed_users.iter().copied() {
738 if max_edges_per_node.is_none_or(|limit| owner_edge_count < limit) {
739 potential_edges.push(PotentialEdge {
740 owner,
741 target,
742 is_follow: true,
743 });
744 owner_edge_count += 1;
745 }
746 }
747 }
748
749 if let Some(muted_users) = self.muted_by_user.get(&owner) {
750 for target in muted_users.iter().copied() {
751 if max_edges_per_node.is_none_or(|limit| owner_edge_count < limit) {
752 potential_edges.push(PotentialEdge {
753 owner,
754 target,
755 is_follow: false,
756 });
757 owner_edge_count += 1;
758 }
759 }
760 }
761 }
762 }
763
764 let mut used_ids = IndexSet::new();
765 let mut follow_edge_count = IndexMap::new();
766 let mut mute_edge_count = IndexMap::new();
767 let mut edge_count = 0usize;
768
769 for edge in potential_edges {
770 if max_edges.is_some_and(|limit| edge_count >= limit) {
771 break;
772 }
773
774 if self.ids.str(edge.owner).is_err() || self.ids.str(edge.target).is_err() {
775 continue;
776 }
777
778 if let Some(limit) = max_nodes {
779 let new_nodes_count = usize::from(!used_ids.contains(&edge.owner))
780 + usize::from(!used_ids.contains(&edge.target));
781 if used_ids.len() + new_nodes_count > limit {
782 break;
783 }
784 }
785
786 used_ids.insert(edge.owner);
787 used_ids.insert(edge.target);
788 edge_count += 1;
789
790 let edge_counts = if edge.is_follow {
791 &mut follow_edge_count
792 } else {
793 &mut mute_edge_count
794 };
795 *edge_counts.entry(edge.owner).or_insert(0) += 1;
796 }
797
798 let follow_owners = follow_edge_count.keys().copied().collect();
799 let mute_owners = mute_edge_count.keys().copied().collect();
800 Ok(BinaryPlan {
801 used_ids,
802 follow_edge_count,
803 mute_edge_count,
804 follow_owners,
805 mute_owners,
806 })
807 }
808
809 fn plan_full_binary(&self) -> BinaryPlan {
810 let mut used_ids = IndexSet::new();
811 let mut follow_edge_count = IndexMap::new();
812 let mut mute_edge_count = IndexMap::new();
813
814 for (user, followed_users) in &self.followed_by_user {
815 used_ids.insert(*user);
816 follow_edge_count.insert(*user, followed_users.len());
817 for followed in followed_users {
818 used_ids.insert(*followed);
819 }
820 }
821
822 for (user, muted_users) in &self.muted_by_user {
823 used_ids.insert(*user);
824 mute_edge_count.insert(*user, muted_users.len());
825 for muted in muted_users {
826 used_ids.insert(*muted);
827 }
828 }
829
830 let follow_owners = follow_edge_count.keys().copied().collect();
831 let mute_owners = mute_edge_count.keys().copied().collect();
832
833 BinaryPlan {
834 used_ids,
835 follow_edge_count,
836 mute_edge_count,
837 follow_owners,
838 mute_owners,
839 }
840 }
841
842 pub fn from_binary(root: &str, data: &[u8]) -> Result<Self> {
843 let mut offset = 0usize;
844 let version = read_varint(data, &mut offset)?;
845 if version != BINARY_FORMAT_VERSION {
846 return Err(SocialGraphError::InvalidVersion(version));
847 }
848
849 let ids_count = read_varint(data, &mut offset)? as usize;
850 let mut unique_ids = Vec::with_capacity(ids_count);
851 for _ in 0..ids_count {
852 let hex_bytes = read_bytes(data, &mut offset, 32)?;
853 let id = read_varint(data, &mut offset)? as u32;
854 unique_ids.push((hex::encode(hex_bytes), id));
855 }
856
857 let follow_lists_count = read_varint(data, &mut offset)? as usize;
858 let mut follow_lists = Vec::with_capacity(follow_lists_count);
859 for _ in 0..follow_lists_count {
860 let user = read_varint(data, &mut offset)? as u32;
861 let timestamp = read_varint(data, &mut offset)?;
862 let followed_count = read_varint(data, &mut offset)? as usize;
863 let mut followed = Vec::with_capacity(followed_count);
864 for _ in 0..followed_count {
865 followed.push(read_varint(data, &mut offset)? as u32);
866 }
867 follow_lists.push((user, followed, timestamp));
868 }
869
870 let mute_lists_count = read_varint(data, &mut offset)? as usize;
871 let mut mute_lists = Vec::with_capacity(mute_lists_count);
872 for _ in 0..mute_lists_count {
873 let user = read_varint(data, &mut offset)? as u32;
874 let timestamp = read_varint(data, &mut offset)?;
875 let muted_count = read_varint(data, &mut offset)? as usize;
876 let mut muted = Vec::with_capacity(muted_count);
877 for _ in 0..muted_count {
878 muted.push(read_varint(data, &mut offset)? as u32);
879 }
880 mute_lists.push((user, muted, timestamp));
881 }
882
883 let mut graph = Self::new(root);
884 graph.ids.clear();
885 graph.follow_distance_by_user.clear();
886 graph.users_by_follow_distance.clear();
887 graph.followed_by_user.clear();
888 graph.followers_by_user.clear();
889 graph.follow_list_created_at.clear();
890 graph.muted_by_user.clear();
891 graph.user_muted_by.clear();
892 graph.mute_list_created_at.clear();
893
894 for (value, id) in unique_ids {
895 graph.ids.insert_with_id(value, id);
896 }
897
898 graph.root = graph.ids.id(root)?;
899 graph.follow_distance_by_user.insert(graph.root, 0);
900 let mut root_bucket = IndexSet::new();
901 root_bucket.insert(graph.root);
902 graph.users_by_follow_distance.insert(0, root_bucket);
903
904 for (follower, followed_users, created_at) in follow_lists {
905 for followed_user in followed_users {
906 graph.private_add_follower(followed_user, follower);
907 }
908 graph.follow_list_created_at.insert(follower, created_at);
909 }
910
911 for (muter, muted_users, created_at) in mute_lists {
912 let entry = graph.muted_by_user.entry(muter).or_default();
913 for muted_user in muted_users {
914 entry.insert(muted_user);
915 graph
916 .user_muted_by
917 .entry(muted_user)
918 .or_default()
919 .insert(muter);
920 }
921 graph.mute_list_created_at.insert(muter, created_at);
922 }
923
924 Ok(graph)
925 }
926
927 pub fn is_overmuted(&self, user: &str, threshold: f64) -> bool {
928 if user == self.get_root() {
929 return false;
930 }
931
932 let Some(user_id) = self.ids.existing_id(user) else {
933 return false;
934 };
935 let Some(muters) = self.user_muted_by.get(&user_id) else {
936 return false;
937 };
938 if muters.is_empty() {
939 return false;
940 }
941 if muters.contains(&self.root) {
942 return true;
943 }
944
945 let mut stats = IndexMap::<u32, (u32, u32)>::new();
946 if let Some(followers) = self.followers_by_user.get(&user_id) {
947 for follower in followers {
948 if let Some(distance) = self.follow_distance_by_user.get(follower).copied() {
949 let entry = stats.entry(distance).or_insert((0, 0));
950 entry.0 += 1;
951 }
952 }
953 }
954 for muter in muters {
955 if let Some(distance) = self.follow_distance_by_user.get(muter).copied() {
956 let entry = stats.entry(distance).or_insert((0, 0));
957 entry.1 += 1;
958 }
959 }
960
961 let mut distances: Vec<u32> = stats.keys().copied().collect();
962 distances.sort_unstable();
963 for distance in distances {
964 let (followers, muters) = stats[&distance];
965 if followers + muters > 0 {
966 return (muters as f64) * threshold > followers as f64;
967 }
968 }
969 false
970 }
971
972 fn handle_follow_list(&mut self, author: u32, created_at: u64, tags: &[Vec<String>]) {
973 if self
974 .follow_list_created_at
975 .get(&author)
976 .is_some_and(|existing| created_at <= *existing)
977 {
978 return;
979 }
980 self.follow_list_created_at.insert(author, created_at);
981
982 let mut followed_in_event = IndexSet::new();
983 for tag in tags {
984 if tag.first().is_some_and(|value| value == "p")
985 && tag.get(1).is_some_and(|pk| is_valid_pubkey(pk))
986 {
987 let Ok(followed_user) = self.ids.id(&tag[1]) else {
988 continue;
989 };
990 if followed_user != author {
991 followed_in_event.insert(followed_user);
992 }
993 }
994 }
995
996 let currently_followed = self
997 .followed_by_user
998 .get(&author)
999 .cloned()
1000 .unwrap_or_default();
1001 for user in currently_followed {
1002 if !followed_in_event.contains(&user) {
1003 self.private_remove_follower(user, author);
1004 }
1005 }
1006 for user in followed_in_event {
1007 self.private_add_follower(user, author);
1008 }
1009 }
1010
1011 fn handle_mute_list(&mut self, author: u32, created_at: u64, tags: &[Vec<String>]) {
1012 if self
1013 .mute_list_created_at
1014 .get(&author)
1015 .is_some_and(|existing| created_at <= *existing)
1016 {
1017 return;
1018 }
1019 self.mute_list_created_at.insert(author, created_at);
1020
1021 let mut muted_in_event = IndexSet::new();
1022 for tag in tags {
1023 if tag.first().is_some_and(|value| value == "p")
1024 && tag.get(1).is_some_and(|pk| is_valid_pubkey(pk))
1025 {
1026 let Ok(muted_user) = self.ids.id(&tag[1]) else {
1027 continue;
1028 };
1029 if muted_user != author {
1030 muted_in_event.insert(muted_user);
1031 }
1032 }
1033 }
1034
1035 let currently_muted = self.muted_by_user.get(&author).cloned().unwrap_or_default();
1036 for user in currently_muted {
1037 if !muted_in_event.contains(&user) {
1038 if let Some(set) = self.muted_by_user.get_mut(&author) {
1039 set.shift_remove(&user);
1040 }
1041 if let Some(set) = self.user_muted_by.get_mut(&user) {
1042 set.shift_remove(&author);
1043 }
1044 }
1045 }
1046
1047 for user in muted_in_event {
1048 self.muted_by_user.entry(author).or_default().insert(user);
1049 self.user_muted_by.entry(user).or_default().insert(author);
1050 }
1051 }
1052
1053 fn private_add_follower(&mut self, followed_user: u32, follower: u32) {
1054 self.followed_by_user
1055 .entry(follower)
1056 .or_default()
1057 .insert(followed_user);
1058 self.followers_by_user
1059 .entry(followed_user)
1060 .or_default()
1061 .insert(follower);
1062
1063 if followed_user == self.root {
1064 return;
1065 }
1066
1067 if follower == self.root {
1068 self.follow_distance_by_user.insert(followed_user, 1);
1069 self.add_user_by_follow_distance(1, followed_user);
1070 return;
1071 }
1072
1073 let existing = self.follow_distance_by_user.get(&followed_user).copied();
1074 let follower_distance = self.follow_distance_by_user.get(&follower).copied();
1075 let new_distance = follower_distance.map(|distance| distance.saturating_add(1));
1076 if let Some(distance) = new_distance
1077 .filter(|distance| existing.is_none() || *distance < existing.unwrap_or(u32::MAX))
1078 {
1079 self.follow_distance_by_user.insert(followed_user, distance);
1080 self.add_user_by_follow_distance(distance, followed_user);
1081 }
1082 }
1083
1084 fn private_remove_follower(&mut self, unfollowed_user: u32, follower: u32) {
1085 if let Some(set) = self.followed_by_user.get_mut(&follower) {
1086 set.shift_remove(&unfollowed_user);
1087 }
1088 if let Some(set) = self.followers_by_user.get_mut(&unfollowed_user) {
1089 set.shift_remove(&follower);
1090 }
1091
1092 if unfollowed_user == self.root {
1093 return;
1094 }
1095
1096 let mut smallest = None;
1097 if let Some(followers) = self.followers_by_user.get(&unfollowed_user) {
1098 for follower in followers {
1099 if let Some(distance) = self.follow_distance_by_user.get(follower).copied() {
1100 let candidate = distance.saturating_add(1);
1101 smallest =
1102 Some(smallest.map_or(candidate, |current: u32| current.min(candidate)));
1103 }
1104 }
1105 }
1106
1107 match smallest {
1108 Some(distance) => {
1109 self.follow_distance_by_user
1110 .insert(unfollowed_user, distance);
1111 self.add_user_by_follow_distance(distance, unfollowed_user);
1112 }
1113 None => {
1114 self.follow_distance_by_user.shift_remove(&unfollowed_user);
1115 for bucket in self.users_by_follow_distance.values_mut() {
1116 bucket.shift_remove(&unfollowed_user);
1117 }
1118 }
1119 }
1120 }
1121
1122 fn add_user_by_follow_distance(&mut self, distance: u32, user: u32) {
1123 self.users_by_follow_distance
1124 .entry(distance)
1125 .or_default()
1126 .insert(user);
1127 let keys: Vec<u32> = self.users_by_follow_distance.keys().copied().collect();
1128 for key in keys.into_iter().filter(|key| *key > distance) {
1129 if let Some(bucket) = self.users_by_follow_distance.get_mut(&key) {
1130 bucket.shift_remove(&user);
1131 }
1132 }
1133 }
1134}
1135
1136impl SocialGraphBackend for SocialGraph {
1137 type Error = SocialGraphError;
1138
1139 fn get_root(&self) -> std::result::Result<String, Self::Error> {
1140 Ok(SocialGraph::get_root(self).to_string())
1141 }
1142
1143 fn set_root(&mut self, root: &str) -> std::result::Result<(), Self::Error> {
1144 SocialGraph::set_root(self, root)
1145 }
1146
1147 fn handle_event(
1148 &mut self,
1149 event: &NostrEvent,
1150 allow_unknown_authors: bool,
1151 overmute_threshold: f64,
1152 ) -> std::result::Result<(), Self::Error> {
1153 SocialGraph::handle_event(self, event, allow_unknown_authors, overmute_threshold);
1154 Ok(())
1155 }
1156
1157 fn get_follow_distance(&self, user: &str) -> std::result::Result<u32, Self::Error> {
1158 Ok(SocialGraph::get_follow_distance(self, user))
1159 }
1160
1161 fn is_following(
1162 &self,
1163 follower: &str,
1164 followed_user: &str,
1165 ) -> std::result::Result<bool, Self::Error> {
1166 Ok(SocialGraph::is_following(self, follower, followed_user))
1167 }
1168
1169 fn get_followed_by_user(&self, user: &str) -> std::result::Result<Vec<String>, Self::Error> {
1170 Ok(SocialGraph::get_followed_by_user(self, user))
1171 }
1172
1173 fn get_followers_by_user(&self, user: &str) -> std::result::Result<Vec<String>, Self::Error> {
1174 Ok(SocialGraph::get_followers_by_user(self, user))
1175 }
1176
1177 fn get_muted_by_user(&self, user: &str) -> std::result::Result<Vec<String>, Self::Error> {
1178 Ok(SocialGraph::get_muted_by_user(self, user))
1179 }
1180
1181 fn get_user_muted_by(&self, user: &str) -> std::result::Result<Vec<String>, Self::Error> {
1182 Ok(SocialGraph::get_user_muted_by(self, user))
1183 }
1184
1185 fn get_follow_list_created_at(
1186 &self,
1187 user: &str,
1188 ) -> std::result::Result<Option<u64>, Self::Error> {
1189 Ok(SocialGraph::get_follow_list_created_at(self, user))
1190 }
1191
1192 fn get_mute_list_created_at(
1193 &self,
1194 user: &str,
1195 ) -> std::result::Result<Option<u64>, Self::Error> {
1196 Ok(SocialGraph::get_mute_list_created_at(self, user))
1197 }
1198
1199 fn is_overmuted(&self, user: &str, threshold: f64) -> std::result::Result<bool, Self::Error> {
1200 Ok(SocialGraph::is_overmuted(self, user, threshold))
1201 }
1202}
1203
1204fn is_valid_pubkey(key: &str) -> bool {
1205 key.len() == 64 && key.bytes().all(|byte| byte.is_ascii_hexdigit())
1206}
1207
1208fn decode_hex_32(hex_value: &str, id: u32) -> Result<[u8; 32]> {
1209 if hex_value.len() != 64 || !hex_value.bytes().all(|byte| byte.is_ascii_hexdigit()) {
1210 return Err(SocialGraphError::InvalidHex(id, hex_value.to_owned()));
1211 }
1212 let bytes = hex::decode(hex_value)
1213 .map_err(|_| SocialGraphError::InvalidHex(id, hex_value.to_owned()))?;
1214 let mut output = [0u8; 32];
1215 output.copy_from_slice(&bytes);
1216 Ok(output)
1217}
1218
1219fn write_varint(out: &mut Vec<u8>, mut value: u64) {
1220 while value >= 0x80 {
1221 out.push(((value as u8) & 0x7f) | 0x80);
1222 value >>= 7;
1223 }
1224 out.push((value & 0x7f) as u8);
1225}
1226
1227fn read_varint(data: &[u8], offset: &mut usize) -> Result<u64> {
1228 let mut value = 0u64;
1229 let mut shift = 0u32;
1230 loop {
1231 let byte = *data.get(*offset).ok_or(SocialGraphError::UnexpectedEof)?;
1232 *offset += 1;
1233 value |= u64::from(byte & 0x7f) << shift;
1234 if byte & 0x80 == 0 {
1235 return Ok(value);
1236 }
1237 shift += 7;
1238 }
1239}
1240
1241fn read_bytes<'a>(data: &'a [u8], offset: &mut usize, len: usize) -> Result<&'a [u8]> {
1242 let end = offset.saturating_add(len);
1243 if end > data.len() {
1244 return Err(SocialGraphError::UnexpectedEof);
1245 }
1246 let slice = &data[*offset..end];
1247 *offset = end;
1248 Ok(slice)
1249}
1250
1251#[cfg(test)]
1252mod tests {
1253 use super::*;
1254
1255 const ADAM: &str = "020f2d21ae09bf35fcdfb65decf1478b846f5f728ab30c5eaabcd6d081a81c3e";
1256 const FIATJAF: &str = "3bf0c63fcb93463407af97a5e5ee64fa883d107ef9e558472c4eb9aaaefa459d";
1257
1258 fn event(pubkey: &str, kind: u32, created_at: u64, tagged: Vec<&str>) -> NostrEvent {
1259 NostrEvent {
1260 created_at,
1261 content: String::new(),
1262 tags: tagged
1263 .into_iter()
1264 .map(|pk| vec!["p".to_string(), pk.to_string()])
1265 .collect(),
1266 kind,
1267 pubkey: pubkey.to_string(),
1268 id: format!("{pubkey}:{kind}:{created_at}"),
1269 sig: "00".repeat(64),
1270 }
1271 }
1272
1273 #[test]
1274 fn read_queries_do_not_intern_unknown_users() {
1275 let graph = SocialGraph::new(ADAM);
1276 let before = graph.export_state();
1277
1278 assert_eq!(
1279 graph.get_follow_distance("ff".repeat(32).as_str()),
1280 UNKNOWN_FOLLOW_DISTANCE
1281 );
1282 assert!(!graph.is_following(ADAM, &"ee".repeat(32)));
1283 assert!(graph.get_followed_by_user(&"dd".repeat(32)).is_empty());
1284 assert!(graph.get_followers_by_user(&"cc".repeat(32)).is_empty());
1285 assert!(graph.get_muted_by_user(&"bb".repeat(32)).is_empty());
1286 assert!(graph.get_user_muted_by(&"aa".repeat(32)).is_empty());
1287 assert_eq!(graph.get_follow_list_created_at(&"11".repeat(32)), None);
1288 assert_eq!(graph.get_mute_list_created_at(&"22".repeat(32)), None);
1289 assert!(!graph.is_overmuted(&"33".repeat(32), 1.0));
1290
1291 let after = graph.export_state();
1292 assert_eq!(after.unique_ids, before.unique_ids);
1293 }
1294
1295 #[test]
1296 fn unique_ids_round_trip_and_reuse_existing_ids() {
1297 let mut ids = UniqueIds::default();
1298
1299 let adam = ids.id(ADAM).unwrap();
1300 let fiatjaf = ids.id(FIATJAF).unwrap();
1301
1302 assert_eq!(adam, 0);
1303 assert_eq!(fiatjaf, 1);
1304 assert_eq!(ids.id(ADAM).unwrap(), adam);
1305 assert_eq!(ids.str(adam).unwrap(), ADAM);
1306 assert_eq!(ids.str(fiatjaf).unwrap(), FIATJAF);
1307 }
1308
1309 #[test]
1310 fn unique_ids_reject_empty_strings_and_invalid_ids() {
1311 let mut ids = UniqueIds::default();
1312
1313 assert_eq!(
1314 ids.id(" ").unwrap_err().to_string(),
1315 "cannot store empty or whitespace-only strings"
1316 );
1317 assert_eq!(ids.str(99).unwrap_err().to_string(), "invalid id 99");
1318 }
1319
1320 #[test]
1321 fn from_binary_rejects_invalid_version() {
1322 let error = SocialGraph::from_binary(ADAM, &[99]).unwrap_err();
1323 assert!(matches!(error, SocialGraphError::InvalidVersion(99)));
1324 }
1325
1326 #[test]
1327 fn from_binary_rejects_truncated_binary() {
1328 let error = SocialGraph::from_binary(ADAM, &[BINARY_FORMAT_VERSION as u8]).unwrap_err();
1329 assert!(matches!(error, SocialGraphError::UnexpectedEof));
1330 }
1331
1332 #[test]
1333 fn to_binary_rejects_non_hex_pubkeys_when_they_are_serialized() {
1334 let mut graph = SocialGraph::new(ADAM);
1335 graph.handle_event(
1336 &event("not-a-hex-pubkey", 3, 1_000, vec![FIATJAF]),
1337 true,
1338 1.0,
1339 );
1340
1341 let error = graph.to_binary().unwrap_err();
1342 assert!(matches!(
1343 error,
1344 SocialGraphError::InvalidHex(1, ref value) if value == "not-a-hex-pubkey"
1345 ));
1346 }
1347}