use std::mem::size_of;
use std::sync::Arc;
use nodedb_mem::{EngineId, MemoryGovernor};
use super::types::{CsrIndex, Direction};
use crate::GraphError;
use crate::csr::LocalNodeId;
pub(crate) struct DenseAdjacency {
pub(crate) offsets: Vec<u32>,
pub(crate) targets: Vec<u32>,
pub(crate) labels: Vec<u32>,
}
impl CsrIndex {
#[inline]
pub fn partition_tag(&self) -> u32 {
self.partition_tag
}
#[inline]
pub fn local(&self, id: u32) -> LocalNodeId {
LocalNodeId::new(id, self.partition_tag)
}
pub fn neighbors(
&self,
node: &str,
label_filter: Option<&str>,
direction: Direction,
) -> Vec<(String, String)> {
let Some(&node_id) = self.node_to_id.get(node) else {
return Vec::new();
};
self.record_access(node_id);
let label_id = label_filter.and_then(|l| self.label_to_id.get(l).copied());
let mut result = Vec::new();
if matches!(direction, Direction::Out | Direction::Both) {
for (lid, dst) in self.dense_iter_out(node_id) {
if label_id.is_none_or(|f| f == lid) {
result.push((
self.id_to_label[lid as usize].clone(),
self.id_to_node[dst as usize].clone(),
));
}
}
}
if matches!(direction, Direction::In | Direction::Both) {
for (lid, src) in self.dense_iter_in(node_id) {
if label_id.is_none_or(|f| f == lid) {
result.push((
self.id_to_label[lid as usize].clone(),
self.id_to_node[src as usize].clone(),
));
}
}
}
result
}
pub fn neighbors_multi(
&self,
node: &str,
label_filters: &[&str],
direction: Direction,
) -> Vec<(String, String)> {
let Some(&node_id) = self.node_to_id.get(node) else {
return Vec::new();
};
self.record_access(node_id);
let label_ids: Vec<u32> = label_filters
.iter()
.filter_map(|l| self.label_to_id.get(*l).copied())
.collect();
let match_label = |lid: u32| label_ids.is_empty() || label_ids.contains(&lid);
let mut result = Vec::new();
if matches!(direction, Direction::Out | Direction::Both) {
for (lid, dst) in self.dense_iter_out(node_id) {
if match_label(lid) {
result.push((
self.id_to_label[lid as usize].clone(),
self.id_to_node[dst as usize].clone(),
));
}
}
}
if matches!(direction, Direction::In | Direction::Both) {
for (lid, src) in self.dense_iter_in(node_id) {
if match_label(lid) {
result.push((
self.id_to_label[lid as usize].clone(),
self.id_to_node[src as usize].clone(),
));
}
}
}
result
}
pub fn add_node(&mut self, name: &str) -> Result<LocalNodeId, crate::GraphError> {
let raw = self.ensure_node(name)?;
Ok(LocalNodeId::new(raw, self.partition_tag))
}
pub fn node_count(&self) -> usize {
self.id_to_node.len()
}
pub fn contains_node(&self, node: &str) -> bool {
self.node_to_id.contains_key(node)
}
pub fn node_name(&self, id: LocalNodeId) -> &str {
&self.id_to_node[id.raw(self.partition_tag) as usize]
}
pub fn node_id(&self, name: &str) -> Option<LocalNodeId> {
self.node_to_id
.get(name)
.copied()
.map(|raw| LocalNodeId::new(raw, self.partition_tag))
}
pub fn label_name(&self, label_id: u32) -> &str {
&self.id_to_label[label_id as usize]
}
pub fn label_id(&self, name: &str) -> Option<u32> {
self.label_to_id.get(name).copied()
}
pub fn out_degree(&self, id: LocalNodeId) -> usize {
self.dense_iter_out(id.raw(self.partition_tag)).count()
}
pub fn in_degree(&self, id: LocalNodeId) -> usize {
self.dense_iter_in(id.raw(self.partition_tag)).count()
}
pub fn edge_count(&self) -> usize {
let n = self.id_to_node.len();
(0..n as u32).map(|i| self.out_degree(self.local(i))).sum()
}
pub(crate) fn build_dense(
edges: &[Vec<(u32, u32)>],
governor: Option<&Arc<MemoryGovernor>>,
) -> Result<DenseAdjacency, GraphError> {
let n = edges.len();
let total: usize = edges.iter().map(|e| e.len()).sum();
let reserve_bytes = (n + 1 + 2 * total) * size_of::<u32>();
let _budget_guard = governor
.map(|g| g.reserve(EngineId::Graph, reserve_bytes))
.transpose()?;
let mut offsets = Vec::with_capacity(n + 1);
let mut targets = Vec::with_capacity(total);
let mut labels = Vec::with_capacity(total);
let mut offset = 0u32;
for node_edges in edges {
offsets.push(offset);
for &(lid, target) in node_edges {
targets.push(target);
labels.push(lid);
}
offset += node_edges.len() as u32;
}
offsets.push(offset);
Ok(DenseAdjacency {
offsets,
targets,
labels,
})
}
pub(crate) fn dense_has_edge(&self, src: u32, label: u32, dst: u32) -> bool {
for (lid, target) in self.dense_out_edges(src) {
if lid == label && target == dst {
return true;
}
}
false
}
pub(crate) fn dense_out_edges(&self, node: u32) -> impl Iterator<Item = (u32, u32)> + '_ {
let idx = node as usize;
if idx + 1 >= self.out_offsets.len() {
return Vec::new().into_iter();
}
let start = self.out_offsets[idx] as usize;
let end = self.out_offsets[idx + 1] as usize;
(start..end)
.map(move |i| (self.out_labels[i], self.out_targets[i]))
.collect::<Vec<_>>()
.into_iter()
}
pub(crate) fn dense_in_edges(&self, node: u32) -> impl Iterator<Item = (u32, u32)> + '_ {
let idx = node as usize;
if idx + 1 >= self.in_offsets.len() {
return Vec::new().into_iter();
}
let start = self.in_offsets[idx] as usize;
let end = self.in_offsets[idx + 1] as usize;
(start..end)
.map(move |i| (self.in_labels[i], self.in_targets[i]))
.collect::<Vec<_>>()
.into_iter()
}
pub(crate) fn dense_iter_out(&self, node: u32) -> impl Iterator<Item = (u32, u32)> + '_ {
let dense = self
.dense_out_edges(node)
.filter(move |&(lid, dst)| !self.deleted_edges.contains(&(node, lid, dst)));
dense.chain(self.buffer_out_iter(node))
}
pub(crate) fn dense_iter_in(&self, node: u32) -> impl Iterator<Item = (u32, u32)> + '_ {
let dense = self
.dense_in_edges(node)
.filter(move |&(lid, src)| !self.deleted_edges.contains(&(src, lid, node)));
dense.chain(self.buffer_in_iter(node))
}
pub(crate) fn buffer_out_iter(&self, node: u32) -> std::vec::IntoIter<(u32, u32)> {
let idx = node as usize;
if idx < self.buffer_out.len() {
self.buffer_out[idx].clone().into_iter()
} else {
Vec::new().into_iter()
}
}
pub(crate) fn buffer_in_iter(&self, node: u32) -> std::vec::IntoIter<(u32, u32)> {
let idx = node as usize;
if idx < self.buffer_in.len() {
self.buffer_in[idx].clone().into_iter()
} else {
Vec::new().into_iter()
}
}
pub fn iter_out_edges(
&self,
node: LocalNodeId,
) -> impl Iterator<Item = (u32, LocalNodeId)> + '_ {
let raw = node.raw(self.partition_tag);
let tag = self.partition_tag;
self.dense_iter_out(raw)
.map(move |(lid, dst)| (lid, LocalNodeId::new(dst, tag)))
}
pub fn iter_in_edges(
&self,
node: LocalNodeId,
) -> impl Iterator<Item = (u32, LocalNodeId)> + '_ {
let raw = node.raw(self.partition_tag);
let tag = self.partition_tag;
self.dense_iter_in(raw)
.map(move |(lid, src)| (lid, LocalNodeId::new(src, tag)))
}
pub fn iter_out_edges_raw(&self, node: u32) -> impl Iterator<Item = (u32, u32)> + '_ {
self.dense_iter_out(node)
}
pub fn iter_in_edges_raw(&self, node: u32) -> impl Iterator<Item = (u32, u32)> + '_ {
self.dense_iter_in(node)
}
pub fn out_degree_raw(&self, node: u32) -> usize {
self.dense_iter_out(node).count()
}
pub fn in_degree_raw(&self, node: u32) -> usize {
self.dense_iter_in(node).count()
}
pub fn node_name_raw(&self, id: u32) -> &str {
&self.id_to_node[id as usize]
}
pub fn node_id_raw(&self, name: &str) -> Option<u32> {
self.node_to_id.get(name).copied()
}
}