use cypherlite_core::{NodeId, SubgraphId};
use std::collections::BTreeMap;
pub struct MembershipIndex {
forward: BTreeMap<u64, Vec<u64>>,
reverse: BTreeMap<u64, Vec<u64>>,
}
impl MembershipIndex {
pub fn new() -> Self {
Self {
forward: BTreeMap::new(),
reverse: BTreeMap::new(),
}
}
pub fn add(&mut self, subgraph_id: SubgraphId, node_id: NodeId) {
let fwd = self.forward.entry(subgraph_id.0).or_default();
if !fwd.contains(&node_id.0) {
fwd.push(node_id.0);
}
let rev = self.reverse.entry(node_id.0).or_default();
if !rev.contains(&subgraph_id.0) {
rev.push(subgraph_id.0);
}
}
pub fn remove(&mut self, subgraph_id: SubgraphId, node_id: NodeId) -> bool {
let removed = if let Some(fwd) = self.forward.get_mut(&subgraph_id.0) {
let before = fwd.len();
fwd.retain(|&id| id != node_id.0);
fwd.len() < before
} else {
false
};
if removed {
if let Some(rev) = self.reverse.get_mut(&node_id.0) {
rev.retain(|&id| id != subgraph_id.0);
if rev.is_empty() {
self.reverse.remove(&node_id.0);
}
}
if let Some(fwd) = self.forward.get(&subgraph_id.0) {
if fwd.is_empty() {
self.forward.remove(&subgraph_id.0);
}
}
}
removed
}
pub fn remove_all(&mut self, subgraph_id: SubgraphId) -> Vec<NodeId> {
let node_ids = self.forward.remove(&subgraph_id.0).unwrap_or_default();
for &nid in &node_ids {
if let Some(rev) = self.reverse.get_mut(&nid) {
rev.retain(|&sid| sid != subgraph_id.0);
if rev.is_empty() {
self.reverse.remove(&nid);
}
}
}
node_ids.into_iter().map(NodeId).collect()
}
pub fn members(&self, subgraph_id: SubgraphId) -> Vec<NodeId> {
self.forward
.get(&subgraph_id.0)
.map(|ids| ids.iter().map(|&id| NodeId(id)).collect())
.unwrap_or_default()
}
pub fn memberships(&self, node_id: NodeId) -> Vec<SubgraphId> {
self.reverse
.get(&node_id.0)
.map(|ids| ids.iter().map(|&id| SubgraphId(id)).collect())
.unwrap_or_default()
}
}
impl Default for MembershipIndex {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use cypherlite_core::{NodeId, SubgraphId};
use super::*;
#[test]
fn test_membership_index_new_is_empty() {
let idx = MembershipIndex::new();
assert!(idx.members(SubgraphId(1)).is_empty());
assert!(idx.memberships(NodeId(1)).is_empty());
}
#[test]
fn test_add_member() {
let mut idx = MembershipIndex::new();
idx.add(SubgraphId(1), NodeId(10));
let members = idx.members(SubgraphId(1));
assert_eq!(members, vec![NodeId(10)]);
}
#[test]
fn test_add_multiple_members() {
let mut idx = MembershipIndex::new();
idx.add(SubgraphId(1), NodeId(10));
idx.add(SubgraphId(1), NodeId(20));
idx.add(SubgraphId(1), NodeId(30));
let members = idx.members(SubgraphId(1));
assert_eq!(members.len(), 3);
assert!(members.contains(&NodeId(10)));
assert!(members.contains(&NodeId(20)));
assert!(members.contains(&NodeId(30)));
}
#[test]
fn test_add_duplicate_member_is_idempotent() {
let mut idx = MembershipIndex::new();
idx.add(SubgraphId(1), NodeId(10));
idx.add(SubgraphId(1), NodeId(10)); let members = idx.members(SubgraphId(1));
assert_eq!(members.len(), 1);
}
#[test]
fn test_reverse_index() {
let mut idx = MembershipIndex::new();
idx.add(SubgraphId(1), NodeId(10));
idx.add(SubgraphId(2), NodeId(10));
let memberships = idx.memberships(NodeId(10));
assert_eq!(memberships.len(), 2);
assert!(memberships.contains(&SubgraphId(1)));
assert!(memberships.contains(&SubgraphId(2)));
}
#[test]
fn test_remove_member() {
let mut idx = MembershipIndex::new();
idx.add(SubgraphId(1), NodeId(10));
idx.add(SubgraphId(1), NodeId(20));
let removed = idx.remove(SubgraphId(1), NodeId(10));
assert!(removed);
let members = idx.members(SubgraphId(1));
assert_eq!(members, vec![NodeId(20)]);
assert!(idx.memberships(NodeId(10)).is_empty());
}
#[test]
fn test_remove_nonexistent_member() {
let mut idx = MembershipIndex::new();
let removed = idx.remove(SubgraphId(1), NodeId(10));
assert!(!removed);
}
#[test]
fn test_remove_all_members() {
let mut idx = MembershipIndex::new();
idx.add(SubgraphId(1), NodeId(10));
idx.add(SubgraphId(1), NodeId(20));
let removed = idx.remove_all(SubgraphId(1));
assert_eq!(removed.len(), 2);
assert!(idx.members(SubgraphId(1)).is_empty());
assert!(idx.memberships(NodeId(10)).is_empty());
assert!(idx.memberships(NodeId(20)).is_empty());
}
#[test]
fn test_remove_all_from_empty() {
let mut idx = MembershipIndex::new();
let removed = idx.remove_all(SubgraphId(999));
assert!(removed.is_empty());
}
#[test]
fn test_membership_index_default() {
let idx = MembershipIndex::default();
assert!(idx.members(SubgraphId(1)).is_empty());
}
}