use crate::graph::VertexId;
use std::collections::{HashMap, HashSet, VecDeque};
#[derive(Debug, Clone)]
pub struct CacheOptAdjacency {
neighbors: Vec<(VertexId, f64)>,
offsets: Vec<usize>,
vertex_count: usize,
}
impl CacheOptAdjacency {
pub fn from_edges(edges: &[(VertexId, VertexId, f64)], max_vertex: VertexId) -> Self {
let vertex_count = (max_vertex + 1) as usize;
let mut adj: Vec<Vec<(VertexId, f64)>> = vec![Vec::new(); vertex_count];
for &(u, v, w) in edges {
adj[u as usize].push((v, w));
adj[v as usize].push((u, w));
}
let mut neighbors = Vec::with_capacity(edges.len() * 2);
let mut offsets = Vec::with_capacity(vertex_count + 1);
offsets.push(0);
for vertex_neighbors in &adj {
neighbors.extend_from_slice(vertex_neighbors);
offsets.push(neighbors.len());
}
Self {
neighbors,
offsets,
vertex_count,
}
}
#[inline]
pub fn neighbors(&self, v: VertexId) -> &[(VertexId, f64)] {
let v = v as usize;
if v >= self.vertex_count {
return &[];
}
&self.neighbors[self.offsets[v]..self.offsets[v + 1]]
}
#[inline]
pub fn prefetch_neighbors(&self, v: VertexId) {
let v = v as usize;
if v < self.vertex_count {
let _start = self.offsets[v];
}
}
pub fn vertex_count(&self) -> usize {
self.vertex_count
}
}
pub struct CacheOptBFS<'a> {
adj: &'a CacheOptAdjacency,
visited: Vec<bool>,
queue: VecDeque<VertexId>,
prefetch_distance: usize,
}
impl<'a> CacheOptBFS<'a> {
pub fn new(adj: &'a CacheOptAdjacency, start: VertexId) -> Self {
let mut visited = vec![false; adj.vertex_count()];
let mut queue = VecDeque::with_capacity(adj.vertex_count());
if (start as usize) < adj.vertex_count() {
visited[start as usize] = true;
queue.push_back(start);
}
Self {
adj,
visited,
queue,
prefetch_distance: 4,
}
}
pub fn run(mut self) -> HashSet<VertexId> {
let mut result = HashSet::new();
while let Some(v) = self.queue.pop_front() {
result.insert(v);
if let Some(&prefetch_v) = self.queue.get(self.prefetch_distance) {
self.adj.prefetch_neighbors(prefetch_v);
}
for &(neighbor, _) in self.adj.neighbors(v) {
let idx = neighbor as usize;
if idx < self.visited.len() && !self.visited[idx] {
self.visited[idx] = true;
self.queue.push_back(neighbor);
}
}
}
result
}
pub fn connected_to(mut self, target: VertexId) -> bool {
if (target as usize) >= self.adj.vertex_count() {
return false;
}
while let Some(v) = self.queue.pop_front() {
if v == target {
return true;
}
if let Some(&prefetch_v) = self.queue.get(self.prefetch_distance) {
self.adj.prefetch_neighbors(prefetch_v);
}
for &(neighbor, _) in self.adj.neighbors(v) {
let idx = neighbor as usize;
if idx < self.visited.len() && !self.visited[idx] {
self.visited[idx] = true;
self.queue.push_back(neighbor);
}
}
}
false
}
}
pub struct BatchProcessor {
batch_size: usize,
}
impl BatchProcessor {
pub fn new() -> Self {
Self { batch_size: 32 }
}
pub fn with_batch_size(batch_size: usize) -> Self {
Self { batch_size }
}
pub fn process_batched<F>(&self, vertices: &[VertexId], mut f: F)
where
F: FnMut(&[VertexId]),
{
for chunk in vertices.chunks(self.batch_size) {
f(chunk);
}
}
pub fn compute_degrees(
&self,
adj: &CacheOptAdjacency,
vertices: &[VertexId],
) -> HashMap<VertexId, usize> {
let mut degrees = HashMap::with_capacity(vertices.len());
for chunk in vertices.chunks(self.batch_size) {
for &v in chunk {
adj.prefetch_neighbors(v);
}
for &v in chunk {
degrees.insert(v, adj.neighbors(v).len());
}
}
degrees
}
}
impl Default for BatchProcessor {
fn default() -> Self {
Self::new()
}
}
#[repr(C, align(64))]
pub struct AlignedBuffer<T, const N: usize> {
data: [T; N],
}
impl<T: Default + Copy, const N: usize> AlignedBuffer<T, N> {
pub fn new() -> Self
where
T: Default + Copy,
{
Self {
data: [T::default(); N],
}
}
pub fn as_slice(&self) -> &[T] {
&self.data
}
pub fn as_mut_slice(&mut self) -> &mut [T] {
&mut self.data
}
}
impl<T: Default + Copy, const N: usize> Default for AlignedBuffer<T, N> {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache_opt_adjacency() {
let edges = vec![(0, 1, 1.0), (1, 2, 1.0), (2, 3, 1.0)];
let adj = CacheOptAdjacency::from_edges(&edges, 3);
assert_eq!(adj.vertex_count(), 4);
assert_eq!(adj.neighbors(0).len(), 1);
assert_eq!(adj.neighbors(1).len(), 2);
assert_eq!(adj.neighbors(2).len(), 2);
assert_eq!(adj.neighbors(3).len(), 1);
}
#[test]
fn test_cache_opt_bfs() {
let edges = vec![(0, 1, 1.0), (1, 2, 1.0), (2, 3, 1.0)];
let adj = CacheOptAdjacency::from_edges(&edges, 3);
let bfs = CacheOptBFS::new(&adj, 0);
let visited = bfs.run();
assert!(visited.contains(&0));
assert!(visited.contains(&1));
assert!(visited.contains(&2));
assert!(visited.contains(&3));
}
#[test]
fn test_bfs_connectivity() {
let edges = vec![(0, 1, 1.0), (2, 3, 1.0)];
let adj = CacheOptAdjacency::from_edges(&edges, 3);
assert!(CacheOptBFS::new(&adj, 0).connected_to(1));
assert!(!CacheOptBFS::new(&adj, 0).connected_to(2));
}
#[test]
fn test_batch_processor() {
let edges = vec![(0, 1, 1.0), (1, 2, 1.0), (2, 3, 1.0)];
let adj = CacheOptAdjacency::from_edges(&edges, 3);
let processor = BatchProcessor::new();
let vertices: Vec<VertexId> = (0..4).collect();
let degrees = processor.compute_degrees(&adj, &vertices);
assert_eq!(degrees.get(&0), Some(&1));
assert_eq!(degrees.get(&1), Some(&2));
assert_eq!(degrees.get(&2), Some(&2));
assert_eq!(degrees.get(&3), Some(&1));
}
#[test]
fn test_aligned_buffer() {
let buffer: AlignedBuffer<u64, 8> = AlignedBuffer::new();
let ptr = buffer.as_slice().as_ptr();
assert_eq!(ptr as usize % 64, 0);
assert_eq!(buffer.as_slice().len(), 8);
}
}