use super::{group_id_for, local_vid, region_id_for_local, CsrShardSet};
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct RegionDirty {
pub inserted: bool,
pub deleted: bool,
pub column_updated: bool,
}
impl RegionDirty {
pub fn is_dirty(self) -> bool {
self.inserted || self.deleted
}
pub fn is_column_dirty(self) -> bool {
self.column_updated
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct GroupDirty {
pub inserted: bool,
pub deleted: bool,
pub column_updated: bool,
}
impl GroupDirty {
pub fn is_dirty(self) -> bool {
self.inserted || self.deleted
}
pub fn is_column_dirty(self) -> bool {
self.column_updated
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EdgeCheckpointKind {
AppendOnly,
Rebalance,
}
impl CsrShardSet {
pub fn group_dirty(&self, gid: usize) -> GroupDirty {
self.shards
.get(&gid)
.map(|shard| shard.dirty)
.unwrap_or_default()
}
pub fn dirty_group_ids(&self) -> Vec<usize> {
self.shards
.iter()
.filter_map(|(gid, shard)| shard.dirty.is_dirty().then_some(*gid))
.collect()
}
pub fn column_dirty_group_ids(&self) -> Vec<usize> {
self.shards
.iter()
.filter_map(|(gid, shard)| shard.dirty.is_column_dirty().then_some(*gid))
.collect()
}
pub fn checkpoint_kind(&self) -> EdgeCheckpointKind {
let rebalance = self.shards.values().any(|shard| shard.dirty.deleted);
if rebalance {
EdgeCheckpointKind::Rebalance
} else {
EdgeCheckpointKind::AppendOnly
}
}
pub fn mark_column_updated_for_group(&mut self, gid: usize) {
if !self.shards.contains_key(&gid) {
let _ = self.ensure_group_id(gid);
}
if let Some(shard) = self.shards.get_mut(&gid) {
shard.dirty.column_updated = true;
for region in shard.regions.iter_mut() {
region.column_updated = true;
}
}
}
pub fn mark_column_updated_for(&mut self, vid: u32) {
let gid = group_id_for(vid, self.group_bits);
let rid = region_id_for_local(local_vid(vid, self.group_bits));
if !self.shards.contains_key(&gid) {
let _ = self.ensure_group_id(gid);
}
if let Some(shard) = self.shards.get_mut(&gid) {
shard.dirty.column_updated = true;
if let Some(region) = shard.regions.get_mut(rid) {
region.column_updated = true;
}
}
}
pub fn clear_group_dirty(&mut self, gid: usize) {
if let Some(shard) = self.shards.get_mut(&gid) {
shard.dirty = GroupDirty::default();
for region in shard.regions.iter_mut() {
*region = RegionDirty::default();
}
}
}
pub fn clear_all_dirty(&mut self) {
for shard in self.shards.values_mut() {
shard.dirty = GroupDirty::default();
for region in shard.regions.iter_mut() {
*region = RegionDirty::default();
}
}
}
pub fn clear_all_column_dirty(&mut self) {
for shard in self.shards.values_mut() {
shard.dirty.column_updated = false;
for region in shard.regions.iter_mut() {
region.column_updated = false;
}
}
}
pub fn mark_all_dirty(&mut self) {
for shard in self.shards.values_mut() {
shard.dirty = GroupDirty {
inserted: true,
deleted: true,
column_updated: true,
};
for region in shard.regions.iter_mut() {
*region = RegionDirty {
inserted: true,
deleted: true,
column_updated: true,
};
}
}
}
pub fn needs_checkpoint(&self, gid: usize) -> bool {
self.group_dirty(gid).is_dirty()
}
pub fn group_needs_rebalance(&self, gid: usize) -> bool {
self.group_dirty(gid).deleted
}
pub fn region_dirty(&self, gid: usize, region: usize) -> RegionDirty {
self.shards
.get(&gid)
.and_then(|shard| shard.regions.get(region))
.copied()
.unwrap_or_default()
}
pub fn dirty_region_ids(&self, gid: usize) -> Vec<usize> {
self.shards
.get(&gid)
.map(|shard| {
shard
.regions
.iter()
.enumerate()
.filter_map(|(rid, region)| region.is_dirty().then_some(rid))
.collect()
})
.unwrap_or_default()
}
pub fn region_needs_checkpoint(&self, gid: usize, region: usize) -> bool {
self.region_dirty(gid, region).is_dirty()
}
pub fn region_needs_rebalance(&self, gid: usize, region: usize) -> bool {
self.region_dirty(gid, region).deleted
}
pub fn clear_region_dirty(&mut self, gid: usize, region: usize) {
if let Some(shard) = self.shards.get_mut(&gid) {
if let Some(slot) = shard.regions.get_mut(region) {
*slot = RegionDirty::default();
}
if shard.regions.iter().all(|r| !r.is_dirty()) {
shard.dirty.inserted = false;
shard.dirty.deleted = false;
} else {
shard.dirty.inserted = shard.regions.iter().any(|r| r.inserted);
shard.dirty.deleted = shard.regions.iter().any(|r| r.deleted);
}
}
}
pub(super) fn mark_region_insert(&mut self, gid: usize, local: u32) {
if let Some(shard) = self.shards.get_mut(&gid) {
shard.dirty.inserted = true;
let rid = region_id_for_local(local);
if let Some(region) = shard.regions.get_mut(rid) {
region.inserted = true;
}
}
}
pub(super) fn mark_region_delete(&mut self, gid: usize, local: u32) {
if let Some(shard) = self.shards.get_mut(&gid) {
shard.dirty.deleted = true;
let rid = region_id_for_local(local);
if let Some(region) = shard.regions.get_mut(rid) {
region.deleted = true;
}
}
}
}