use super::{
reconcile::{Reconciliation, StructuralRecords},
OwnedPopulationMutation,
};
use crate::mvcc::{
commit::RecordWriteKind, resolution::ResolutionMode, CommittedRecordSnapshot,
MergedRecordSnapshot, PreparedRecordCommit, PreparedRecordWrite, PrivateRecordChanges,
VersionError, VersionResult, VersionedPersistence,
};
use crate::read_control::StorageReadControl;
use std::sync::Arc;
use uqa_core::memory::BudgetedVec;
pub(in crate::mvcc) fn stage(
origins: &[PreparedRecordWrite],
lifecycle: &[OwnedPopulationMutation],
before: &MergedRecordSnapshot,
after: &MergedRecordSnapshot,
persistence: &dyn VersionedPersistence,
control: &StorageReadControl,
) -> VersionResult<PreparedRecordCommit> {
Reconciliation {
before,
after,
layout: persistence.diskann_population_record_layout().ok_or(
VersionError::InvalidEncoding("provider has no DiskANN population layout"),
)?,
history: persistence.database_id(),
control,
structural: None,
}
.run(origins, lifecycle)
}
pub(in crate::mvcc) fn resolve(
prepared: &PreparedRecordCommit,
resolved: Option<PreparedRecordCommit>,
current: &Arc<dyn CommittedRecordSnapshot>,
persistence: &dyn VersionedPersistence,
mode: ResolutionMode,
control: &StorageReadControl,
) -> VersionResult<Option<PreparedRecordCommit>> {
if prepared.populations.is_none()
&& !prepared.records().iter().any(|write| {
matches!(
write.kind(),
RecordWriteKind::DiskANNOrigin | RecordWriteKind::DiskANNPopulationPreview
)
})
{
return Ok(resolved);
}
reconcile(
resolved.as_ref().unwrap_or(prepared),
prepared
.populations
.as_ref()
.map_or(&[], |effects| &effects.operations),
current,
persistence,
mode,
control,
)
.map(Some)
}
fn reconcile(
original: &PreparedRecordCommit,
lifecycle: &[OwnedPopulationMutation],
current: &Arc<dyn CommittedRecordSnapshot>,
persistence: &dyn VersionedPersistence,
mode: ResolutionMode,
control: &StorageReadControl,
) -> VersionResult<PreparedRecordCommit> {
let layout =
persistence
.diskann_population_record_layout()
.ok_or(VersionError::InvalidEncoding(
"provider has no DiskANN population layout",
))?;
let changes = PrivateRecordChanges::new(control.memory());
let mut origins = BudgetedVec::new(control.memory());
let mut structural = StructuralRecords::new(control.memory());
for write in original.records() {
control.check()?;
if write.kind() == RecordWriteKind::Canonical {
structural.insert(write.key(), write)?;
}
}
for (position, write) in original.records().iter().enumerate() {
control.check()?;
let kind = match write.kind() {
RecordWriteKind::DiskANNOrigin => {
validate(current.as_ref(), position, write, control)?;
origins.push(write.clone())?;
mode.kind(RecordWriteKind::DiskANNOrigin)
}
RecordWriteKind::DiskANNPopulationPreview => {
let header = layout.preview_header(write.key(), control)?;
if structural
.get(&*header)
.is_none_or(|owner| owner.value().is_none())
{
continue;
}
validate(current.as_ref(), position, write, control)?;
RecordWriteKind::Canonical
}
RecordWriteKind::Canonical => {
validate(current.as_ref(), position, write, control)?;
RecordWriteKind::Canonical
}
kind => kind,
};
changes.apply_owned(&[write.clone().with_kind(kind)], control)?;
}
let empty = PrivateRecordChanges::new(control.memory());
let before = MergedRecordSnapshot::new(current.clone(), empty.snapshot()?);
let after = MergedRecordSnapshot::new(current.clone(), changes.snapshot()?);
let generated = Reconciliation {
before: &before,
after: &after,
layout,
history: persistence.database_id(),
control,
structural: Some(&structural),
}
.run(&origins, lifecycle)?;
for write in generated.records() {
control.check()?;
changes.apply_owned(
&[write
.clone()
.with_kind(mode.kind(RecordWriteKind::DiskANNPopulationPreview))],
control,
)?;
}
Ok(changes
.prepare(control)?
.retain_graph_effects(original, control)?
.resolved(original, current.sequence()))
}
fn validate(
current: &dyn CommittedRecordSnapshot,
mutation: usize,
write: &PreparedRecordWrite,
control: &StorageReadControl,
) -> VersionResult<()> {
let actual = current
.metadata(write.key(), control)?
.and_then(|row| row.revision);
if write.expected() != actual {
return Err(VersionError::WriteConflict {
mutation,
expected: write.expected(),
actual,
});
}
Ok(())
}