use crate::storage::membership::{KeyedMemberships, MembershipBacking, MembershipKeys};
use crate::storage::warp_table::stable_slots::StableSlotsBacking;
use crate::storage::{
IndexKey,
membership::MembershipChunkSelection,
query_workspace::{MembershipOperand, MembershipOperation, chunk::read::ChunkRead},
warp_table::{
BulkSelection, BulkSetOutcome, BulkUpdateError,
inconsistency::invariant_violation,
stable_slots::{StableSlot, StableSlotResolveOutcome, StableSlots},
},
};
#[allow(
clippy::arithmetic_side_effects,
clippy::indexing_slicing,
reason = "selection belongs to the exclusively borrowed table; preflight bounds every pending chunk before mutation; field storage and stable slots remain fixed"
)]
pub(super) fn set_column_selection<
Key: IndexKey,
Keys: MembershipKeys<Key>,
Backing: MembershipBacking,
Slots: StableSlotsBacking,
>(
index: &mut KeyedMemberships<Key, Keys, Backing>,
column: &mut [Key],
stable_slots: &StableSlots<Slots>,
target: Key,
selection: &mut BulkSelection<'_>,
) -> Result<BulkSetOutcome, BulkUpdateError> {
let chunks = selection.selected.capacity().div_ceil(65536);
let selected_count = selection.selected.len();
let mut selected = MembershipOperand::from(selection.selected).chunks();
let mut maximum = 0;
for high in 0..chunks {
maximum = maximum.max(
selected
.at(high)
.unwrap_or_else(|reason| invariant_violation(reason))
.len as usize,
);
}
const DIRECT_SET_MAXIMUM_MEMBERS: u32 = 64;
if maximum <= DIRECT_SET_MAXIMUM_MEMBERS as usize {
let mut changed = 0;
for slot in selection.selected.slots() {
let dense = resolve(stable_slots, slot);
if column[dense] != target {
let _outcome = index
.prepare_move(column[dense], target, slot)
.unwrap_or_else(|reason| invariant_violation(reason))
.commit();
column[dense] = target;
changed += 1;
}
}
return Ok(if changed == 0 {
BulkSetOutcome::Unchanged {
selected: selected_count,
}
} else {
BulkSetOutcome::Updated {
selected: selected_count,
changed,
}
});
}
if selection.slots.len() < maximum {
let mut selected = MembershipOperand::from(selection.selected).chunks();
let target_members = index
.operand(target)
.unwrap_or_else(|reason| invariant_violation(reason));
let mut target_chunks = target_members.chunks();
let mut required = 0;
for high in 0..chunks {
selection
.pending
.copy_read(
selected
.at(high)
.unwrap_or_else(|reason| invariant_violation(reason)),
)
.unwrap_or_else(|reason| invariant_violation(reason));
selection
.pending
.apply_read(
MembershipOperation::Difference,
target_chunks
.at(high)
.unwrap_or_else(|reason| invariant_violation(reason)),
selection.group,
)
.unwrap_or_else(|reason| invariant_violation(reason));
required = required.max(
selection
.pending
.read()
.unwrap_or_else(|reason| invariant_violation(reason))
.len as usize,
);
}
if selection.slots.len() < required {
return Err(BulkUpdateError::InsufficientStorage {
required,
available: selection.slots.len(),
});
}
if required == 0 {
return Ok(BulkSetOutcome::Unchanged {
selected: selected_count,
});
}
}
let mut changed = 0_u64;
let mut selected = MembershipOperand::from(selection.selected).chunks();
for high in 0..chunks {
let target_chunk = index
.find(target)
.unwrap_or_else(|reason| invariant_violation(reason))
.map_or(ChunkRead::EMPTY, |members| members.chunk(high as u16).0);
let selected_chunk = selected
.at(high)
.unwrap_or_else(|reason| invariant_violation(reason));
if selected_chunk.is_subset(target_chunk) {
continue;
}
if selected_chunk.len > DIRECT_SET_MAXIMUM_MEMBERS
&& let Some(first) = selected_chunk.slots().next()
{
let source = column[resolve(stable_slots, ((high as u32) << 16) | u32::from(first))];
if source != target {
let source_chunk = index
.find(source)
.unwrap_or_else(|reason| invariant_violation(reason))
.map_or(ChunkRead::EMPTY, |members| members.chunk(high as u16).0);
if selected_chunk.is_subset(source_chunk) {
for (output, low) in selection.slots.iter_mut().zip(selected_chunk.slots()) {
*output = low;
}
let lows = &selection.slots[..selected_chunk.len as usize];
let admitted = MembershipChunkSelection::new(lows)
.unwrap_or_else(|reason| invariant_violation(reason));
let _outcome = index
.prepare_move_selection(source, target, high as u16, admitted)
.unwrap_or_else(|reason| invariant_violation(reason))
.commit();
for low in lows {
let slot = ((high as u32) << 16) | u32::from(*low);
column[resolve(stable_slots, slot)] = target;
}
changed += u64::from(selected_chunk.len);
continue;
}
}
}
selection
.pending
.copy_read(selected_chunk)
.unwrap_or_else(|reason| invariant_violation(reason));
selection
.pending
.apply_read(
MembershipOperation::Difference,
target_chunk,
selection.group,
)
.unwrap_or_else(|reason| invariant_violation(reason));
changed += u64::from(
selection
.pending
.read()
.unwrap_or_else(|reason| invariant_violation(reason))
.len,
);
while let Some(first) = selection
.pending
.read()
.unwrap_or_else(|reason| invariant_violation(reason))
.slots()
.next()
{
let pending = selection
.pending
.read()
.unwrap_or_else(|reason| invariant_violation(reason));
if pending.len > DIRECT_SET_MAXIMUM_MEMBERS {
let slot = ((high as u32) << 16) | u32::from(first);
let source = column[resolve(stable_slots, slot)];
let source_chunk = index
.find(source)
.unwrap_or_else(|reason| invariant_violation(reason))
.map_or(ChunkRead::EMPTY, |members| members.chunk(high as u16).0);
if source_chunk.len > DIRECT_SET_MAXIMUM_MEMBERS {
let complete = pending.is_subset(source_chunk);
let members = if complete {
pending
} else {
pending
.combine(
MembershipOperation::Intersection,
source_chunk,
selection.group,
)
.unwrap_or_else(|reason| invariant_violation(reason));
selection
.group
.read()
.unwrap_or_else(|reason| invariant_violation(reason))
};
let count = members.len;
const MAXIMUM_GROUP_SCAN_RATIO: u32 = 32;
if count > DIRECT_SET_MAXIMUM_MEMBERS
&& count * MAXIMUM_GROUP_SCAN_RATIO >= pending.len
{
for (output, low) in selection.slots.iter_mut().zip(members.slots()) {
*output = low;
}
if !complete {
selection
.pending
.apply_read(
MembershipOperation::Difference,
source_chunk,
selection.group,
)
.unwrap_or_else(|reason| invariant_violation(reason));
}
let lows = &selection.slots[..count as usize];
let admitted = MembershipChunkSelection::new(lows)
.unwrap_or_else(|reason| invariant_violation(reason));
let _outcome = index
.prepare_move_selection(source, target, high as u16, admitted)
.unwrap_or_else(|reason| invariant_violation(reason))
.commit();
for low in lows {
let slot = ((high as u32) << 16) | u32::from(*low);
column[resolve(stable_slots, slot)] = target;
}
if complete {
break;
}
continue;
}
}
}
for low in selection
.pending
.read()
.unwrap_or_else(|reason| invariant_violation(reason))
.slots()
{
let slot = ((high as u32) << 16) | u32::from(low);
let dense = resolve(stable_slots, slot);
let _outcome = index
.prepare_move(column[dense], target, slot)
.unwrap_or_else(|reason| invariant_violation(reason))
.commit();
column[dense] = target;
}
break;
}
}
Ok(if changed == 0 {
BulkSetOutcome::Unchanged {
selected: selected_count,
}
} else {
BulkSetOutcome::Updated {
selected: selected_count,
changed,
}
})
}
fn resolve<Slots: StableSlotsBacking>(stable_slots: &StableSlots<Slots>, slot: u32) -> usize {
let StableSlotResolveOutcome::Resolved { dense_row, .. } = stable_slots
.resolve(StableSlot::new(slot))
.unwrap_or_else(|reason| invariant_violation(reason));
dense_row
}