raphy/
fast_csr.rs

1/*
2Copyright 2020 Brandon Lucia <blucia@gmail.com>
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6http://www.apache.org/licenses/LICENSE-2.0
7Unless required by applicable law or agreed to in writing, software
8distributed under the License is distributed on an "AS IS" BASIS,
9WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10See the License for the specific language governing permissions and
11limitations under the License.
12*/
13
14use 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} /*impl FastCSR*/