1use byte_slice_cast::*;
15use memmap2::Mmap;
16use rayon::prelude::*;
17use std::fs::OpenOptions;
18use std::path::PathBuf;
19
20pub struct FastCSR {
21 v: usize,
22 e: usize,
23 obase: usize,
24 nbase: usize,
25 raw: Box<Mmap>,
26}
27
28impl FastCSR {
29 pub fn getv(&self) -> usize {
30 self.v
31 }
32
33 pub fn gete(&self) -> usize {
34 self.e
35 }
36
37 pub fn new(s: String) -> FastCSR {
38 let path = PathBuf::from(s);
39 let file = OpenOptions::new().read(true).open(&path).unwrap();
40
41 let mmap = Box::new(unsafe { Mmap::map(&file).unwrap() });
42
43 assert!(mmap.len() >= 8);
44 let csr = mmap[..]
45 .as_slice_of::<usize>()
46 .unwrap();
47
48 let v = csr[0];
49 let e = csr[1];
50
51 println!("{} edges total", e);
52 FastCSR {
53 v: v,
54 e: e,
55 obase: 16,
56 nbase: 16 + v * 8,
57 raw: mmap,
58 }
59 }
60
61 pub fn offset(&self, i: usize) -> usize {
62 let offsets = &self.raw[self.obase..self.nbase]
63 .as_slice_of::<usize>()
64 .unwrap();
65
66 offsets[i]
67 }
68
69 pub fn neighbors(&self, i: usize) -> &[usize] {
70 let (n0, nn) = self.vtx_offset_range(i);
71 let edges = &self.raw[self.nbase..].as_slice_of::<usize>().unwrap();
72 &edges[n0..nn]
73 }
74
75 fn vtx_offset_range(&self, v: usize) -> (usize, usize) {
76 (
77 self.offset(v),
78 match v {
79 v if v == self.v - 1 => self.e,
80 _ => self.offset(v + 1),
81 },
82 )
83 }
84 pub fn neighbor_scan_prop(&self, f: impl Fn(usize, &[usize]) -> f64 + std::marker::Sync, prop: &mut [f64]) {
85 prop.par_iter_mut().enumerate().for_each(|(v,p)| {
86 let (n0, nn) = self.vtx_offset_range(v);
87 let edges = &self.raw[self.nbase..].as_slice_of::<usize>().unwrap();
88 let res = f(v, &edges[n0..nn]);
89 *p = res;
90 });
91 }
92
93 pub fn neighbor_scan(&self, f: impl Fn(usize, &[usize]) -> () + std::marker::Sync) {
94 (0..self.v).into_par_iter().for_each(|v| {
95 let (n0, nn) = self.vtx_offset_range(v);
96 let edges = &self.raw[self.nbase..].as_slice_of::<usize>().unwrap();
97 f(v, &edges[n0..nn]);
98 });
99 }
100
101 pub fn read_only_scan(&self, f: impl Fn(usize, usize) -> () + std::marker::Sync) {
102 (0..self.v).into_par_iter().for_each(|v| {
103 let (n0, nn) = self.vtx_offset_range(v);
104 let edges = self.raw[self.nbase..].as_slice_of::<usize>().unwrap();
105 edges[n0..nn].into_iter().for_each(|n| {
106 f(v, *n);
107 });
108 });
109 }
110}