Skip to content

Commit 5a3c6ef

Browse files
yukibtcyukibtc
committed
mls-storage: replace Vec with BTreeSet for admins and relays
Use a `BTreeSet` instead of a `Vec` for storing admin public keys and relay URLs. This ensures unique elements, faster lookups and sorted results. Signed-off-by: Yuki Kishimoto <[email protected]>
1 parent 53d0723 commit 5a3c6ef

File tree

9 files changed

+56
-40
lines changed

9 files changed

+56
-40
lines changed

‎crates/nostr-mls-memory-storage/src/groups.rs‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
//! Memory-based storage implementation of the NostrMlsStorageProvider trait for Nostr MLS groups
22
3+
use std::collections::BTreeSet;
4+
35
use nostr::PublicKey;
46
use nostr_mls_storage::groups::error::{GroupError, InvalidGroupState};
57
use nostr_mls_storage::groups::types::*;
@@ -57,14 +59,14 @@ impl GroupStorage for NostrMlsMemoryStorage {
5759
}
5860
}
5961

60-
fn admins(&self, mls_group_id: &[u8]) -> Result<Vec<PublicKey>, GroupError> {
62+
fn admins(&self, mls_group_id: &[u8]) -> Result<BTreeSet<PublicKey>, GroupError> {
6163
match self.find_group_by_mls_group_id(mls_group_id)? {
6264
Some(group) => Ok(group.admin_pubkeys),
6365
None => Err(GroupError::InvalidState(InvalidGroupState::NoAdmins)),
6466
}
6567
}
6668

67-
fn group_relays(&self, mls_group_id: &[u8]) -> Result<Vec<GroupRelay>, GroupError> {
69+
fn group_relays(&self, mls_group_id: &[u8]) -> Result<BTreeSet<GroupRelay>, GroupError> {
6870
// Check if the group exists first
6971
self.find_group_by_mls_group_id(mls_group_id)?;
7072

@@ -85,17 +87,16 @@ impl GroupStorage for NostrMlsMemoryStorage {
8587
match cache.get_mut(&group_relay.mls_group_id) {
8688
// If the group exists, add the new relay to the vector
8789
Some(existing_relays) => {
88-
// TODO: the time complexity here is O(n). The number of relays is likely to be small, but it's probably better to use a BTreeSet anyway.
89-
9090
// Add the new relay if it doesn't already exist
91-
if !existing_relays.contains(&group_relay) {
92-
existing_relays.push(group_relay);
93-
}
91+
existing_relays.insert(group_relay);
9492
}
9593
// If the group doesn't exist, create a new vector with the new relay
9694
None => {
9795
// Update the cache with the new vector
98-
cache.put(group_relay.mls_group_id.clone(), vec![group_relay]);
96+
cache.put(
97+
group_relay.mls_group_id.clone(),
98+
BTreeSet::from([group_relay]),
99+
);
99100
}
100101
};
101102

‎crates/nostr-mls-memory-storage/src/lib.rs‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
#![warn(missing_docs)]
1111
#![warn(rustdoc::bare_urls)]
1212

13+
use std::collections::BTreeSet;
1314
use std::num::NonZeroUsize;
1415

1516
use lru::LruCache;
@@ -58,7 +59,7 @@ pub struct NostrMlsMemoryStorage {
5859
/// LRU Cache for Group objects, keyed by Nostr group ID (String)
5960
groups_by_nostr_id_cache: RwLock<LruCache<String, Group>>,
6061
/// LRU Cache for GroupRelay objects, keyed by MLS group ID (`Vec<u8>`)
61-
group_relays_cache: RwLock<LruCache<Vec<u8>, Vec<GroupRelay>>>,
62+
group_relays_cache: RwLock<LruCache<Vec<u8>, BTreeSet<GroupRelay>>>,
6263
/// LRU Cache for Welcome objects, keyed by Event ID
6364
welcomes_cache: RwLock<LruCache<EventId, Welcome>>,
6465
/// LRU Cache for ProcessedWelcome objects, keyed by Event ID
@@ -164,6 +165,8 @@ impl NostrMlsStorageProvider for NostrMlsMemoryStorage {
164165

165166
#[cfg(test)]
166167
mod tests {
168+
use std::collections::BTreeSet;
169+
167170
use nostr::{EventId, Kind, PublicKey, RelayUrl, Tags, Timestamp, UnsignedEvent};
168171
use nostr_mls_storage::groups::types::{Group, GroupState, GroupType};
169172
use nostr_mls_storage::groups::GroupStorage;
@@ -230,7 +233,7 @@ mod tests {
230233
nostr_group_id: "test_group_123".to_string(),
231234
name: "Test Group".to_string(),
232235
description: "A test group".to_string(),
233-
admin_pubkeys: vec![],
236+
admin_pubkeys: BTreeSet::new(),
234237
last_message_id: None,
235238
last_message_at: None,
236239
group_type: GroupType::Group,
@@ -273,7 +276,7 @@ mod tests {
273276
nostr_group_id: "test_group_456".to_string(),
274277
name: "Another Test Group".to_string(),
275278
description: "Another test group".to_string(),
276-
admin_pubkeys: vec![],
279+
admin_pubkeys: BTreeSet::new(),
277280
last_message_id: None,
278281
last_message_at: None,
279282
group_type: GroupType::Group,
@@ -358,8 +361,8 @@ mod tests {
358361
nostr_group_id: "test_welcome_group".to_string(),
359362
group_name: "Test Welcome Group".to_string(),
360363
group_description: "A test welcome group".to_string(),
361-
group_admin_pubkeys: vec![pubkey.to_hex()],
362-
group_relays: vec!["wss://relay.example.com".to_string()],
364+
group_admin_pubkeys: BTreeSet::from([pubkey]),
365+
group_relays: BTreeSet::from([RelayUrl::parse("wss://relay.example.com").unwrap()]),
363366
welcomer: pubkey,
364367
member_count: 2,
365368
state: WelcomeState::Pending,
@@ -422,7 +425,7 @@ mod tests {
422425
nostr_group_id: "message_test_group".to_string(),
423426
name: "Message Test Group".to_string(),
424427
description: "A group for testing messages".to_string(),
425-
admin_pubkeys: vec![],
428+
admin_pubkeys: BTreeSet::new(),
426429
last_message_id: None,
427430
last_message_at: None,
428431
group_type: GroupType::Group,
@@ -531,7 +534,7 @@ mod tests {
531534
nostr_group_id: "custom_cache_group".to_string(),
532535
name: "Custom Cache Group".to_string(),
533536
description: "A group for testing custom cache size".to_string(),
534-
admin_pubkeys: vec![],
537+
admin_pubkeys: BTreeSet::new(),
535538
last_message_id: None,
536539
last_message_at: None,
537540
group_type: GroupType::Group,
@@ -560,7 +563,7 @@ mod tests {
560563
nostr_group_id: "default_impl_group".to_string(),
561564
name: "Default Implementation Group".to_string(),
562565
description: "A group for testing default implementation".to_string(),
563-
admin_pubkeys: vec![],
566+
admin_pubkeys: BTreeSet::new(),
564567
last_message_id: None,
565568
last_message_at: None,
566569
group_type: GroupType::Group,

‎crates/nostr-mls-sqlite-storage/src/db.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
//! Database utilities for SQLite storage.
22
3+
use std::collections::BTreeSet;
34
use std::io::{Error as IoError, ErrorKind};
45
use std::str::FromStr;
56

@@ -47,7 +48,7 @@ pub fn row_to_group(row: &Row) -> SqliteResult<Group> {
4748

4849
// Parse admin pubkeys from JSON
4950
let admin_pubkeys_json: &str = row.get_ref("admin_pubkeys")?.as_str()?;
50-
let admin_pubkeys: Vec<PublicKey> =
51+
let admin_pubkeys: BTreeSet<PublicKey> =
5152
serde_json::from_str(admin_pubkeys_json).map_err(map_to_text_boxed_error)?;
5253

5354
let last_message_id: Option<&[u8]> = row.get_ref("last_message_id")?.as_blob_or_null()?;
@@ -195,10 +196,10 @@ pub fn row_to_welcome(row: &Row) -> SqliteResult<Welcome> {
195196
let event: UnsignedEvent =
196197
UnsignedEvent::from_json(event_json).map_err(map_to_text_boxed_error)?;
197198

198-
let group_admin_pubkeys: Vec<String> =
199+
let group_admin_pubkeys: BTreeSet<PublicKey> =
199200
serde_json::from_str(group_admin_pubkeys_json).map_err(map_to_text_boxed_error)?;
200201

201-
let group_relays: Vec<String> =
202+
let group_relays: BTreeSet<RelayUrl> =
202203
serde_json::from_str(group_relays_json).map_err(map_to_text_boxed_error)?;
203204

204205
let welcomer: PublicKey = PublicKey::from_slice(welcomer_blob)

‎crates/nostr-mls-sqlite-storage/src/groups.rs‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
//! Implementation of GroupStorage trait for SQLite storage.
22
3+
use std::collections::BTreeSet;
4+
35
use nostr::PublicKey;
46
use nostr_mls_storage::groups::error::GroupError;
57
use nostr_mls_storage::groups::types::{Group, GroupRelay};
@@ -132,7 +134,7 @@ impl GroupStorage for NostrMlsSqliteStorage {
132134
Ok(messages)
133135
}
134136

135-
fn admins(&self, mls_group_id: &[u8]) -> Result<Vec<PublicKey>, GroupError> {
137+
fn admins(&self, mls_group_id: &[u8]) -> Result<BTreeSet<PublicKey>, GroupError> {
136138
// Get the group which contains the admin_pubkeys
137139
match self.find_group_by_mls_group_id(mls_group_id)? {
138140
Some(group) => Ok(group.admin_pubkeys),
@@ -143,7 +145,7 @@ impl GroupStorage for NostrMlsSqliteStorage {
143145
}
144146
}
145147

146-
fn group_relays(&self, mls_group_id: &[u8]) -> Result<Vec<GroupRelay>, GroupError> {
148+
fn group_relays(&self, mls_group_id: &[u8]) -> Result<BTreeSet<GroupRelay>, GroupError> {
147149
// First verify the group exists
148150
if self.find_group_by_mls_group_id(mls_group_id)?.is_none() {
149151
return Err(GroupError::InvalidParameters(format!(
@@ -162,11 +164,11 @@ impl GroupStorage for NostrMlsSqliteStorage {
162164
.query_map(params![mls_group_id], db::row_to_group_relay)
163165
.map_err(into_group_err)?;
164166

165-
let mut relays: Vec<GroupRelay> = Vec::new();
167+
let mut relays: BTreeSet<GroupRelay> = BTreeSet::new();
166168

167169
for relay_result in relays_iter {
168170
let relay: GroupRelay = relay_result.map_err(into_group_err)?;
169-
relays.push(relay);
171+
relays.insert(relay);
170172
}
171173

172174
Ok(relays)
@@ -215,7 +217,7 @@ mod tests {
215217
nostr_group_id: "test_group_123".to_string(),
216218
name: "Test Group".to_string(),
217219
description: "A test group".to_string(),
218-
admin_pubkeys: vec![],
220+
admin_pubkeys: BTreeSet::new(),
219221
last_message_id: None,
220222
last_message_at: None,
221223
group_type: GroupType::Group,
@@ -257,7 +259,7 @@ mod tests {
257259
nostr_group_id: "test_group_123".to_string(),
258260
name: "Test Group".to_string(),
259261
description: "A test group".to_string(),
260-
admin_pubkeys: vec![],
262+
admin_pubkeys: BTreeSet::new(),
261263
last_message_id: None,
262264
last_message_at: None,
263265
group_type: GroupType::Group,
@@ -283,6 +285,9 @@ mod tests {
283285
// Get group relays
284286
let relays = storage.group_relays(&mls_group_id).unwrap();
285287
assert_eq!(relays.len(), 1);
286-
assert_eq!(relays[0].relay_url.to_string(), "wss://relay.example.com");
288+
assert_eq!(
289+
relays.first().unwrap().relay_url.to_string(),
290+
"wss://relay.example.com"
291+
);
287292
}
288293
}

‎crates/nostr-mls-sqlite-storage/src/messages.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,8 @@ impl MessageStorage for NostrMlsSqliteStorage {
109109

110110
#[cfg(test)]
111111
mod tests {
112+
use std::collections::BTreeSet;
113+
112114
use nostr::{EventId, Kind, PublicKey, Tags, Timestamp, UnsignedEvent};
113115
use nostr_mls_storage::groups::types::{Group, GroupState, GroupType};
114116
use nostr_mls_storage::groups::GroupStorage;
@@ -127,7 +129,7 @@ mod tests {
127129
nostr_group_id: "test_group_123".to_string(),
128130
name: "Test Group".to_string(),
129131
description: "A test group".to_string(),
130-
admin_pubkeys: vec![],
132+
admin_pubkeys: BTreeSet::new(),
131133
last_message_id: None,
132134
last_message_at: None,
133135
group_type: GroupType::Group,

‎crates/nostr-mls-sqlite-storage/src/welcomes.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,9 @@ impl WelcomeStorage for NostrMlsSqliteStorage {
141141

142142
#[cfg(test)]
143143
mod tests {
144-
use nostr::{EventId, Kind, PublicKey, Timestamp, UnsignedEvent};
144+
use std::collections::BTreeSet;
145+
146+
use nostr::{EventId, Kind, PublicKey, RelayUrl, Timestamp, UnsignedEvent};
145147
use nostr_mls_storage::groups::types::{Group, GroupState, GroupType};
146148
use nostr_mls_storage::groups::GroupStorage;
147149
use nostr_mls_storage::welcomes::types::{ProcessedWelcomeState, WelcomeState};
@@ -159,7 +161,7 @@ mod tests {
159161
nostr_group_id: "test_group_123".to_string(),
160162
name: "Test Group".to_string(),
161163
description: "A test group".to_string(),
162-
admin_pubkeys: vec![],
164+
admin_pubkeys: BTreeSet::new(),
163165
last_message_id: None,
164166
last_message_at: None,
165167
group_type: GroupType::Group,
@@ -195,10 +197,8 @@ mod tests {
195197
nostr_group_id: "test_group_123".to_string(),
196198
group_name: "Test Group".to_string(),
197199
group_description: "A test group".to_string(),
198-
group_admin_pubkeys: vec![
199-
"79be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798".to_string(),
200-
],
201-
group_relays: vec!["wss://relay.example.com".to_string()],
200+
group_admin_pubkeys: BTreeSet::from([pubkey]),
201+
group_relays: BTreeSet::from([RelayUrl::parse("wss://relay.example.com").unwrap()]),
202202
welcomer: pubkey,
203203
member_count: 3,
204204
state: WelcomeState::Pending,

‎crates/nostr-mls-storage/src/groups/mod.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
//!
88
//! Here we also define the storage traits that are used to store and retrieve groups
99
10+
use std::collections::BTreeSet;
11+
1012
use nostr::PublicKey;
1113

1214
pub mod error;
@@ -37,10 +39,10 @@ pub trait GroupStorage {
3739
fn messages(&self, mls_group_id: &[u8]) -> Result<Vec<Message>, GroupError>;
3840

3941
/// Get all admins for a group
40-
fn admins(&self, mls_group_id: &[u8]) -> Result<Vec<PublicKey>, GroupError>;
42+
fn admins(&self, mls_group_id: &[u8]) -> Result<BTreeSet<PublicKey>, GroupError>;
4143

4244
/// Get all relays for a group
43-
fn group_relays(&self, mls_group_id: &[u8]) -> Result<Vec<GroupRelay>, GroupError>;
45+
fn group_relays(&self, mls_group_id: &[u8]) -> Result<BTreeSet<GroupRelay>, GroupError>;
4446

4547
/// Save a group relay
4648
fn save_group_relay(&self, group_relay: GroupRelay) -> Result<(), GroupError>;

‎crates/nostr-mls-storage/src/groups/types.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
//! Types for the groups module
22
3+
use std::collections::BTreeSet;
34
use std::fmt;
45
use std::str::FromStr;
56

@@ -140,7 +141,7 @@ pub struct Group {
140141
/// UTF-8 encoded (same value as the NostrGroupDataExtension)
141142
pub description: String,
142143
/// Hex encoded (same value as the NostrGroupDataExtension)
143-
pub admin_pubkeys: Vec<PublicKey>,
144+
pub admin_pubkeys: BTreeSet<PublicKey>,
144145
/// Hex encoded Nostr event ID of the last message in the group
145146
pub last_message_id: Option<EventId>,
146147
/// Timestamp of the last message in the group
@@ -268,7 +269,7 @@ mod tests {
268269
nostr_group_id: "test_id".to_string(),
269270
name: "Test Group".to_string(),
270271
description: "Test Description".to_string(),
271-
admin_pubkeys: Vec::new(),
272+
admin_pubkeys: BTreeSet::new(),
272273
last_message_id: None,
273274
last_message_at: None,
274275
group_type: GroupType::Group,

‎crates/nostr-mls-storage/src/welcomes/types.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
//! Types for the welcomes module
22
3+
use std::collections::BTreeSet;
34
use std::fmt;
45
use std::str::FromStr;
56

6-
use nostr::{EventId, PublicKey, Timestamp, UnsignedEvent};
7+
use nostr::{EventId, PublicKey, RelayUrl, Timestamp, UnsignedEvent};
78
use serde::{Deserialize, Deserializer, Serialize, Serializer};
89

910
use super::error::WelcomeError;
@@ -39,9 +40,9 @@ pub struct Welcome {
3940
/// Group description (from NostrGroupDataExtension)
4041
pub group_description: String,
4142
/// Group admin pubkeys (from NostrGroupDataExtension)
42-
pub group_admin_pubkeys: Vec<String>,
43+
pub group_admin_pubkeys: BTreeSet<PublicKey>,
4344
/// Group relays (from NostrGroupDataExtension)
44-
pub group_relays: Vec<String>,
45+
pub group_relays: BTreeSet<RelayUrl>,
4546
/// Pubkey of the user that sent the welcome
4647
pub welcomer: PublicKey,
4748
/// Member count of the group

0 commit comments

Comments
 (0)