1use rusqlite::{CachedStatement, Connection, Result, Row};
6use std::collections::HashMap;
7
8use crate::{
9 error::table::TableError,
10 tables::{
11 diagnostic::{HandleDiagnostic, column_exists, count_query},
12 table::{HANDLE, ME, Table},
13 },
14 util::union_find::UnionFind,
15};
16
17#[derive(Debug)]
20pub struct Handle {
21 pub rowid: i32,
23 pub id: String,
25 pub person_centric_id: Option<String>,
28}
29
30impl Table for Handle {
32 fn from_row(row: &Row) -> Result<Handle> {
33 Ok(Handle {
34 rowid: row.get("rowid")?,
35 id: row.get("id")?,
36 person_centric_id: row.get("person_centric_id").unwrap_or(None),
37 })
38 }
39}
40
41impl Handle {
43 pub fn get(db: &'_ Connection) -> Result<CachedStatement<'_>, TableError> {
45 Ok(db.prepare_cached(&format!("SELECT * from {HANDLE}"))?)
46 }
47
48 pub fn load(db: &Connection) -> Result<Vec<Handle>, TableError> {
67 let mut statement = Handle::get(db)?;
68 Handle::rows(&mut statement, [])?.collect()
69 }
70
71 #[must_use]
88 pub fn ids_by_rowid(handles: Vec<Handle>) -> HashMap<i32, String> {
89 let mut map = HashMap::with_capacity(handles.len() + 1);
90 map.insert(0, ME.to_string());
91
92 for handle in handles {
93 map.insert(handle.rowid, handle.id);
94 }
95
96 map
97 }
98}
99
100#[derive(Debug)]
134pub struct HandleGroups {
135 sets: UnionFind,
136}
137
138impl HandleGroups {
139 pub fn from_db(db: &Connection) -> Result<Self, TableError> {
144 Ok(Self::from_rows(&Handle::load(db)?))
145 }
146
147 #[must_use]
152 pub fn from_rows(handles: &[Handle]) -> Self {
153 let mut sets = UnionFind::new();
154 sets.make_set(0);
155
156 let mut first_by_id: HashMap<&str, i32> = HashMap::new();
157 let mut first_by_person: HashMap<&str, i32> = HashMap::new();
158
159 for handle in handles {
160 sets.make_set(handle.rowid);
161
162 if !handle.id.is_empty() {
163 let first = *first_by_id
164 .entry(handle.id.as_str())
165 .or_insert(handle.rowid);
166 sets.union(handle.rowid, first);
167 }
168
169 if let Some(person) = handle
170 .person_centric_id
171 .as_deref()
172 .filter(|p| !p.is_empty())
173 {
174 let first = *first_by_person.entry(person).or_insert(handle.rowid);
175 sets.union(handle.rowid, first);
176 }
177 }
178
179 Self { sets }
180 }
181
182 pub fn union(&mut self, a: i32, b: i32) {
187 if self.sets.contains(a) && self.sets.contains(b) {
188 self.sets.union(a, b);
189 }
190 }
191
192 pub fn finalize(mut self) -> HashMap<i32, i32> {
198 let ids: Vec<i32> = self.sets.ids().collect();
199 self.sets.label(ids)
200 }
201}
202
203impl Handle {
205 pub fn run_diagnostic(db: &Connection) -> Result<HandleDiagnostic, TableError> {
223 let query = concat!(
224 "SELECT COUNT(DISTINCT person_centric_id) ",
225 "FROM handle ",
226 "WHERE person_centric_id NOT NULL AND person_centric_id != ''"
227 );
228
229 let person_centric_ids = if column_exists(db, HANDLE, "person_centric_id")? {
230 Some(count_query(db, query)?)
231 } else {
232 None
233 };
234
235 let group_by_handle = HandleGroups::from_db(db)?.finalize();
238 let total_handles = group_by_handle.len() - 1;
239
240 let mut rows_per_group: HashMap<i32, usize> = HashMap::new();
241 for (rowid, group) in group_by_handle {
242 if rowid != 0 {
243 *rows_per_group.entry(group).or_insert(0) += 1;
244 }
245 }
246
247 Ok(HandleDiagnostic {
248 total_handles,
249 person_centric_ids,
250 multi_handle_groups: rows_per_group.values().filter(|rows| **rows > 1).count(),
251 merged_handles: total_handles - rows_per_group.len(),
252 })
253 }
254}
255
256#[cfg(test)]
258mod tests {
259 use crate::tables::handle::{Handle, HandleGroups};
260 use rusqlite::Connection;
261 use std::collections::{HashMap, HashSet};
262
263 fn db_with(rows: &[(i32, &str, &str, Option<&str>)]) -> Connection {
265 let db = Connection::open_in_memory().unwrap();
266 db.execute(
267 "CREATE TABLE handle (
268 ROWID INTEGER PRIMARY KEY,
269 id TEXT NOT NULL,
270 service TEXT NOT NULL,
271 person_centric_id TEXT
272 )",
273 [],
274 )
275 .unwrap();
276 for (rowid, id, service, person) in rows {
277 db.execute(
278 "INSERT INTO handle (ROWID, id, service, person_centric_id) VALUES (?1, ?2, ?3, ?4)",
279 (rowid, id, service, person),
280 )
281 .unwrap();
282 }
283 db
284 }
285
286 fn groups_of(map: &HashMap<i32, i32>, rowids: &[i32]) -> HashSet<i32> {
287 rowids.iter().map(|r| map[r]).collect()
288 }
289
290 #[test]
291 fn load_returns_every_row() {
292 let db = db_with(&[
293 (1, "+15551234567", "iMessage", Some("p1")),
294 (2, "a@example.com", "SMS", None),
295 (3, "b@example.com", "iMessage", Some("")),
296 ]);
297
298 let handles = Handle::load(&db).unwrap();
299
300 let rows: Vec<(i32, &str, Option<&str>)> = handles
301 .iter()
302 .map(|h| (h.rowid, h.id.as_str(), h.person_centric_id.as_deref()))
303 .collect();
304 assert_eq!(
305 rows,
306 vec![
307 (1, "+15551234567", Some("p1")),
308 (2, "a@example.com", None),
309 (3, "b@example.com", Some("")),
310 ]
311 );
312 }
313
314 #[test]
315 fn ids_by_rowid_includes_me_and_every_row() {
316 let db = db_with(&[
317 (1, "+15551234567", "iMessage", None),
318 (2, "a@example.com", "SMS", None),
319 ]);
320
321 let ids = Handle::ids_by_rowid(Handle::load(&db).unwrap());
322
323 assert_eq!(ids.len(), 3);
324 assert_eq!(ids[&0], "Me");
325 assert_eq!(ids[&1], "+15551234567");
326 assert_eq!(ids[&2], "a@example.com");
327 }
328
329 #[test]
330 fn identical_ids_across_services_share_a_group() {
331 let db = db_with(&[
332 (1, "+15551234567", "iMessage", None),
333 (2, "+15551234567", "SMS", None),
334 (3, "+15551234567", "RCS", None),
335 (4, "a@example.com", "iMessage", None),
336 ]);
337 let map = HandleGroups::from_db(&db).unwrap().finalize();
338
339 assert_eq!(groups_of(&map, &[1, 2, 3]).len(), 1);
340 assert_ne!(map[&1], map[&4]);
341 assert_eq!(map[&0], 0, "the local user is group 0");
342 assert_eq!(map.len(), 5);
343 }
344
345 #[test]
346 fn ids_are_not_normalized() {
347 let db = db_with(&[
348 (1, "+15551234567", "iMessage", None),
349 (2, "15551234567", "SMS", None),
350 (3, "A@example.com", "iMessage", None),
351 (4, "a@example.com", "iMessage", None),
352 ]);
353 let map = HandleGroups::from_db(&db).unwrap().finalize();
354
355 assert_ne!(map[&1], map[&2]);
356 assert_ne!(map[&3], map[&4]);
357 }
358
359 #[test]
360 fn person_centric_id_joins_different_ids() {
361 let db = db_with(&[
362 (1, "+15551234567", "iMessage", Some("p1")),
363 (2, "a@example.com", "iMessage", Some("p1")),
364 (3, "b@example.com", "iMessage", Some("p2")),
365 (4, "+15559876543", "iMessage", None),
366 ]);
367 let map = HandleGroups::from_db(&db).unwrap().finalize();
368
369 assert_eq!(map[&1], map[&2]);
370 assert_ne!(map[&1], map[&3]);
371 assert_ne!(map[&1], map[&4]);
372 }
373
374 #[test]
375 fn edges_are_transitive() {
376 let db = db_with(&[
379 (1, "+15551234567", "iMessage", Some("p1")),
380 (2, "a@example.com", "iMessage", Some("p1")),
381 (3, "a@example.com", "SMS", None),
382 ]);
383 let map = HandleGroups::from_db(&db).unwrap().finalize();
384
385 assert_eq!(groups_of(&map, &[1, 2, 3]).len(), 1);
386 }
387
388 #[test]
389 fn empty_person_centric_id_is_not_an_edge() {
390 let db = db_with(&[
391 (1, "+15551234567", "iMessage", Some("")),
392 (2, "a@example.com", "iMessage", Some("")),
393 ]);
394 let map = HandleGroups::from_db(&db).unwrap().finalize();
395
396 assert_ne!(map[&1], map[&2]);
397 }
398
399 #[test]
400 fn external_union_merges_groups() {
401 let db = db_with(&[
402 (1, "+15551234567", "iMessage", None),
403 (2, "a@example.com", "iMessage", None),
404 (3, "b@example.com", "iMessage", None),
405 ]);
406 let mut groups = HandleGroups::from_db(&db).unwrap();
407 groups.union(1, 2);
408 let map = groups.finalize();
409
410 assert_eq!(map[&1], map[&2]);
411 assert_ne!(map[&1], map[&3]);
412 }
413
414 #[test]
415 fn external_union_ignores_unknown_rowids() {
416 let db = db_with(&[(1, "+15551234567", "iMessage", None)]);
417 let mut groups = HandleGroups::from_db(&db).unwrap();
418 groups.union(1, 99);
419 groups.union(98, 99);
420 let map = groups.finalize();
421
422 assert_eq!(map.len(), 2);
423 assert!(!map.contains_key(&99));
424 }
425
426 #[test]
427 fn group_ids_are_sequential_in_rowid_order() {
428 let db = db_with(&[
429 (5, "c", "iMessage", None),
430 (3, "b", "iMessage", None),
431 (9, "b", "SMS", None),
432 (1, "a", "iMessage", None),
433 ]);
434 let map = HandleGroups::from_db(&db).unwrap().finalize();
435
436 assert_eq!(map[&0], 0);
437 assert_eq!(map[&1], 1);
438 assert_eq!(map[&3], 2);
439 assert_eq!(map[&5], 3);
440 assert_eq!(map[&9], 2);
441 }
442
443 #[test]
444 fn diagnostic_omits_person_centric_count_when_column_is_missing() {
445 let db = Connection::open_in_memory().unwrap();
446 db.execute(
447 "CREATE TABLE handle (ROWID INTEGER PRIMARY KEY, id TEXT NOT NULL)",
448 [],
449 )
450 .unwrap();
451 db.execute(
452 "INSERT INTO handle (ROWID, id) VALUES (1, 'first'), (2, 'second')",
453 [],
454 )
455 .unwrap();
456
457 let diagnostic = Handle::run_diagnostic(&db).unwrap();
458
459 assert_eq!(diagnostic.total_handles, 2);
460 assert_eq!(diagnostic.person_centric_ids, None);
461 assert_eq!(diagnostic.multi_handle_groups, 0);
462 assert_eq!(diagnostic.merged_handles, 0);
463 }
464
465 #[test]
466 fn diagnostic_counts_groups_and_merged_rows() {
467 let db = db_with(&[
468 (1, "+15551234567", "iMessage", Some("p1")),
469 (2, "+15551234567", "SMS", Some("p1")),
470 (3, "a@example.com", "iMessage", Some("p1")),
471 (4, "b@example.com", "iMessage", None),
472 (5, "b@example.com", "SMS", None),
473 (6, "+15559876543", "iMessage", Some("p2")),
474 ]);
475
476 let diagnostic = Handle::run_diagnostic(&db).unwrap();
477
478 assert_eq!(diagnostic.total_handles, 6);
479 assert_eq!(diagnostic.person_centric_ids, Some(2));
480 assert_eq!(diagnostic.multi_handle_groups, 2);
482 assert_eq!(diagnostic.merged_handles, 3);
483 }
484
485 #[test]
486 fn diagnostic_ignores_empty_person_centric_ids() {
487 let db = db_with(&[
488 (1, "a@example.com", "iMessage", Some("")),
489 (2, "b@example.com", "iMessage", Some("")),
490 ]);
491
492 let diagnostic = Handle::run_diagnostic(&db).unwrap();
493
494 assert_eq!(diagnostic.person_centric_ids, Some(0));
495 assert_eq!(diagnostic.multi_handle_groups, 0);
496 assert_eq!(diagnostic.merged_handles, 0);
497 }
498}