1use core_query::Dir;
31use core_storage::{Direction, IdMap, Interner, Topology};
32use serde::{Deserialize, Serialize};
33use std::collections::BTreeMap;
34use std::time::{Duration, Instant};
35
36fn live_nodes(idmap: &IdMap, labels: &[u32]) -> (Vec<u32>, Vec<String>) {
45 let n = idmap.len() as u32;
46 let mut ids = Vec::new();
47 let mut keys = Vec::new();
48 for id in 0..n {
49 let Some(key) = idmap.key_of(id) else {
50 continue;
51 };
52 let Some(&sym) = labels.get(id as usize) else {
53 continue;
54 };
55 if sym == u32::MAX {
56 continue; }
58 ids.push(id);
59 keys.push(key.to_string());
60 }
61 (ids, keys)
62}
63
64fn resolve_etype(syms: &Interner, edge_type: Option<&str>) -> Option<Option<u32>> {
70 match edge_type {
71 None => Some(None), Some(name) => {
73 let sym = syms.get(name)?; Some(Some(sym))
75 }
76 }
77}
78
79fn etypes_filtered(topo: &Topology, filter: Option<u32>) -> Vec<u32> {
81 match filter {
82 Some(sym) => {
83 let all: Vec<u32> = topo.etypes().collect();
85 if all.contains(&sym) {
86 vec![sym]
87 } else {
88 vec![]
89 }
90 }
91 None => topo.etypes().collect(),
92 }
93}
94
95#[derive(Debug, Clone, Serialize, Deserialize)]
101#[serde(default)]
102pub struct PageRankConfig {
103 pub damping: f64,
106 pub max_iters: u32,
108 pub tol: f64,
110 pub edge_type: Option<String>,
112 pub direction: AlgoDir,
116 pub budget_ms: u64,
119}
120
121impl Default for PageRankConfig {
122 fn default() -> Self {
123 Self {
124 damping: 0.85,
125 max_iters: 50,
126 tol: 1e-6,
127 edge_type: None,
128 direction: AlgoDir::Out,
129 budget_ms: 5_000,
130 }
131 }
132}
133
134#[derive(Debug, Clone, Serialize, Deserialize)]
136pub struct PageRankReport {
137 pub scores: Vec<(String, f64)>,
140 pub converged: bool,
144}
145
146#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
148#[serde(rename_all = "lowercase")]
149pub enum AlgoDir {
150 Out,
152 In,
154 Both,
156}
157
158impl From<Dir> for AlgoDir {
159 fn from(d: Dir) -> Self {
160 match d {
161 Dir::Out => AlgoDir::Out,
162 Dir::In => AlgoDir::In,
163 Dir::Both => AlgoDir::Both,
164 }
165 }
166}
167
168pub(crate) fn pagerank(
172 topo: &Topology,
173 idmap: &IdMap,
174 syms: &Interner,
175 labels: &[u32],
176 config: &PageRankConfig,
177) -> PageRankReport {
178 let deadline = if config.budget_ms > 0 {
179 Some(Instant::now() + Duration::from_millis(config.budget_ms))
180 } else {
181 None
182 };
183
184 let (node_ids, node_keys) = live_nodes(idmap, labels);
185 let n = node_ids.len();
186
187 if n == 0 {
188 return PageRankReport {
189 scores: Vec::new(),
190 converged: true,
191 };
192 }
193
194 let max_id = topo.etypes().count(); let _ = max_id;
197 let mut id_to_idx: BTreeMap<u32, usize> = BTreeMap::new();
198 for (i, &id) in node_ids.iter().enumerate() {
199 id_to_idx.insert(id, i);
200 }
201
202 let etype_filter = match resolve_etype(syms, config.edge_type.as_deref()) {
204 None => {
205 let score = 1.0 / n as f64;
207 let mut scores: Vec<(String, f64)> =
208 node_keys.iter().map(|k| (k.clone(), score)).collect();
209 scores.sort_by(|(ka, sa), (kb, sb)| {
210 sb.partial_cmp(sa)
211 .unwrap_or(std::cmp::Ordering::Equal)
212 .then(ka.cmp(kb))
213 });
214 return PageRankReport {
215 scores,
216 converged: true,
217 };
218 }
219 Some(f) => f,
220 };
221
222 let etypes = etypes_filtered(topo, etype_filter);
223
224 let mut send_to: Vec<Vec<usize>> = vec![Vec::new(); n];
228
229 for &et in &etypes {
230 for (i, &id) in node_ids.iter().enumerate() {
231 match config.direction {
232 AlgoDir::Out => {
233 for &nbr in topo.neighbors(et, Direction::Out, id).as_ref() {
235 if let Some(&j) = id_to_idx.get(&nbr) {
236 if !send_to[i].contains(&j) {
237 send_to[i].push(j);
238 }
239 }
240 }
241 }
242 AlgoDir::In => {
243 for &nbr in topo.neighbors(et, Direction::In, id).as_ref() {
245 if let Some(&j) = id_to_idx.get(&nbr) {
246 if !send_to[i].contains(&j) {
247 send_to[i].push(j);
248 }
249 }
250 }
251 }
252 AlgoDir::Both => {
253 for dir in [Direction::Out, Direction::In] {
255 for &nbr in topo.neighbors(et, dir, id).as_ref() {
256 if let Some(&j) = id_to_idx.get(&nbr) {
257 if !send_to[i].contains(&j) {
258 send_to[i].push(j);
259 }
260 }
261 }
262 }
263 }
264 }
265 }
266 }
267
268 let mut receive_from: Vec<Vec<(usize, f64)>> = vec![Vec::new(); n];
271 let mut dangling: Vec<usize> = Vec::new();
272
273 for (i, send) in send_to.iter().enumerate() {
274 let out_deg = send.len();
275 if out_deg == 0 {
276 dangling.push(i);
277 } else {
278 let w = 1.0 / out_deg as f64;
279 for &j in send {
280 receive_from[j].push((i, w));
281 }
282 }
283 }
284
285 let nf = n as f64;
287 let d = config.damping;
288 let teleport = (1.0 - d) / nf;
289 let mut pr: Vec<f64> = vec![1.0 / nf; n];
290 let mut converged = false;
291
292 for _iter in 0..config.max_iters {
293 if let Some(dl) = deadline {
295 if Instant::now() >= dl {
296 break;
297 }
298 }
299
300 let dangling_sum: f64 = dangling.iter().map(|&i| pr[i]).sum::<f64>() * d / nf;
302
303 let mut new_pr = vec![teleport + dangling_sum; n];
304 for j in 0..n {
305 let received: f64 = receive_from[j].iter().map(|&(i, w)| pr[i] * w).sum();
306 new_pr[j] += d * received;
307 }
308
309 let delta: f64 = pr
311 .iter()
312 .zip(new_pr.iter())
313 .map(|(a, b)| (a - b).abs())
314 .sum();
315 pr = new_pr;
316
317 if delta < config.tol {
318 converged = true;
319 break;
320 }
321 }
322
323 let mut scores: Vec<(String, f64)> = node_keys.into_iter().zip(pr).collect();
325 scores.sort_by(|(ka, sa), (kb, sb)| {
326 sb.partial_cmp(sa)
327 .unwrap_or(std::cmp::Ordering::Equal)
328 .then(ka.cmp(kb))
329 });
330
331 PageRankReport { scores, converged }
332}
333
334#[derive(Debug, Clone, Serialize, Deserialize)]
340#[serde(default)]
341pub struct WccConfig {
342 pub edge_type: Option<String>,
344 pub budget_ms: u64,
346}
347
348impl Default for WccConfig {
349 fn default() -> Self {
350 Self {
351 edge_type: None,
352 budget_ms: 5_000,
353 }
354 }
355}
356
357#[derive(Debug, Clone, Serialize, Deserialize)]
359pub struct WccReport {
360 pub components: Vec<(String, String)>,
363 pub truncated: bool,
365}
366
367struct UnionFind {
369 parent: Vec<usize>,
370 rank: Vec<u8>,
371}
372
373impl UnionFind {
374 fn new(n: usize) -> Self {
375 Self {
376 parent: (0..n).collect(),
377 rank: vec![0; n],
378 }
379 }
380
381 fn find(&mut self, mut x: usize) -> usize {
382 while self.parent[x] != x {
383 self.parent[x] = self.parent[self.parent[x]]; x = self.parent[x];
385 }
386 x
387 }
388
389 fn union(&mut self, a: usize, b: usize) {
390 let ra = self.find(a);
391 let rb = self.find(b);
392 if ra == rb {
393 return;
394 }
395 match self.rank[ra].cmp(&self.rank[rb]) {
396 std::cmp::Ordering::Less => self.parent[ra] = rb,
397 std::cmp::Ordering::Greater => self.parent[rb] = ra,
398 std::cmp::Ordering::Equal => {
399 self.parent[rb] = ra;
400 self.rank[ra] += 1;
401 }
402 }
403 }
404}
405
406pub(crate) fn wcc(
410 topo: &Topology,
411 idmap: &IdMap,
412 syms: &Interner,
413 labels: &[u32],
414 config: &WccConfig,
415) -> WccReport {
416 let deadline = if config.budget_ms > 0 {
417 Some(Instant::now() + Duration::from_millis(config.budget_ms))
418 } else {
419 None
420 };
421
422 let (node_ids, node_keys) = live_nodes(idmap, labels);
423 let n = node_ids.len();
424
425 if n == 0 {
426 return WccReport {
427 components: Vec::new(),
428 truncated: false,
429 };
430 }
431
432 let mut id_to_idx: BTreeMap<u32, usize> = BTreeMap::new();
434 for (i, &id) in node_ids.iter().enumerate() {
435 id_to_idx.insert(id, i);
436 }
437
438 let etype_filter = match resolve_etype(syms, config.edge_type.as_deref()) {
440 None => {
441 let mut components: Vec<(String, String)> =
443 node_keys.iter().map(|k| (k.clone(), k.clone())).collect();
444 components.sort();
445 return WccReport {
446 components,
447 truncated: false,
448 };
449 }
450 Some(f) => f,
451 };
452
453 let etypes = etypes_filtered(topo, etype_filter);
454
455 let mut uf = UnionFind::new(n);
456 let mut truncated = false;
457
458 'outer: for &et in &etypes {
460 for (i, &id) in node_ids.iter().enumerate() {
461 if let Some(dl) = deadline {
462 if Instant::now() >= dl {
463 truncated = true;
464 break 'outer;
465 }
466 }
467 for &nbr in topo.neighbors(et, Direction::Out, id).as_ref() {
469 if let Some(&j) = id_to_idx.get(&nbr) {
470 uf.union(i, j);
471 }
472 }
473 for &nbr in topo.neighbors(et, Direction::In, id).as_ref() {
477 if let Some(&j) = id_to_idx.get(&nbr) {
478 uf.union(i, j);
479 }
480 }
481 }
482 }
483
484 let mut root_min_key: BTreeMap<usize, &str> = BTreeMap::new();
486 for (i, key_str) in node_keys.iter().enumerate() {
487 let root = uf.find(i);
488 let key = key_str.as_str();
489 let entry = root_min_key.entry(root).or_insert(key);
490 if key < *entry {
491 *entry = key;
492 }
493 }
494
495 let mut components: Vec<(String, String)> = node_keys
496 .iter()
497 .enumerate()
498 .map(|(i, key_str)| {
499 let root = uf.find(i);
500 let comp_id = root_min_key[&root].to_string();
501 (key_str.clone(), comp_id)
502 })
503 .collect();
504 components.sort_by(|(ka, ca), (kb, cb)| ca.cmp(cb).then(ka.cmp(kb)));
505
506 WccReport {
507 components,
508 truncated,
509 }
510}
511
512#[derive(Debug, Clone, Serialize, Deserialize)]
518#[serde(default)]
519pub struct DegreeConfig {
520 pub edge_type: Option<String>,
522 pub direction: AlgoDir,
524 pub budget_ms: u64,
526}
527
528impl Default for DegreeConfig {
529 fn default() -> Self {
530 Self {
531 edge_type: None,
532 direction: AlgoDir::Both,
533 budget_ms: 5_000,
534 }
535 }
536}
537
538#[derive(Debug, Clone, Serialize, Deserialize)]
540pub struct DegreeReport {
541 pub scores: Vec<(String, u64)>,
543 pub truncated: bool,
545}
546
547pub(crate) fn degree_centrality(
549 topo: &Topology,
550 idmap: &IdMap,
551 syms: &Interner,
552 labels: &[u32],
553 config: &DegreeConfig,
554) -> DegreeReport {
555 let deadline = if config.budget_ms > 0 {
556 Some(Instant::now() + Duration::from_millis(config.budget_ms))
557 } else {
558 None
559 };
560
561 let (node_ids, node_keys) = live_nodes(idmap, labels);
562 let n = node_ids.len();
563
564 if n == 0 {
565 return DegreeReport {
566 scores: Vec::new(),
567 truncated: false,
568 };
569 }
570
571 let etype_filter = match resolve_etype(syms, config.edge_type.as_deref()) {
573 None => {
574 let scores = node_keys.iter().map(|k| (k.clone(), 0u64)).collect();
576 return DegreeReport {
577 scores,
578 truncated: false,
579 };
580 }
581 Some(f) => f,
582 };
583
584 let etypes = etypes_filtered(topo, etype_filter);
585 let mut degrees: Vec<u64> = vec![0u64; n];
586 let mut truncated = false;
587
588 for (i, &id) in node_ids.iter().enumerate() {
589 if let Some(dl) = deadline {
590 if Instant::now() >= dl {
591 truncated = true;
592 break;
593 }
594 }
595 for &et in &etypes {
596 match config.direction {
597 AlgoDir::Out => {
598 degrees[i] += topo.neighbors(et, Direction::Out, id).len() as u64;
599 }
600 AlgoDir::In => {
601 degrees[i] += topo.neighbors(et, Direction::In, id).len() as u64;
602 }
603 AlgoDir::Both => {
604 degrees[i] += topo.neighbors(et, Direction::Out, id).len() as u64;
605 degrees[i] += topo.neighbors(et, Direction::In, id).len() as u64;
606 }
607 }
608 }
609 }
610
611 let mut scores: Vec<(String, u64)> = node_keys.into_iter().zip(degrees).collect();
612 scores.sort_by(|(ka, da), (kb, db)| db.cmp(da).then(ka.cmp(kb)));
613
614 DegreeReport { scores, truncated }
615}