use rudb_common::{Error, Result};
pub type Rid = u64;
pub const NO_PARENT: Rid = u64::MAX;
pub const PART_ROWS: usize = 1024;
pub const STRIPE_PARTS: usize = 64;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Place {
pub stripe: u32,
pub part: u32,
pub offset: u32,
}
#[derive(Debug, Clone)]
struct StripeSum {
base: u64,
rows: u64,
parts: Vec<u32>,
uniform: bool,
}
#[derive(Debug, Clone)]
pub struct Places {
stripes: Vec<StripeSum>,
rows: u64,
}
impl Places {
pub fn build(per_stripe: &[Vec<u32>]) -> Result<Self> {
let mut stripes = Vec::with_capacity(per_stripe.len());
let mut base = 0_u64;
for (at, parts) in per_stripe.iter().enumerate() {
if parts.len() > STRIPE_PARTS {
return Err(malformed(format!(
"stripe {at} has {} parts and a stripe holds at most {STRIPE_PARTS}",
parts.len()
)));
}
let mut cumulative = Vec::with_capacity(parts.len() + 1);
cumulative.push(0);
let mut total = 0_u32;
for (which, &rows) in parts.iter().enumerate() {
if rows as usize > PART_ROWS {
return Err(malformed(format!(
"part {which} of stripe {at} holds {rows} rows and a part holds at most \
{PART_ROWS}"
)));
}
total = total.checked_add(rows).ok_or_else(|| {
malformed(format!("stripe {at} overflows a thirty two bit row count"))
})?;
cumulative.push(total);
}
let uniform = parts.iter().all(|&rows| rows as usize == PART_ROWS);
if let Some((_, earlier)) = parts.split_last()
&& let Some(which) = earlier.iter().position(|&rows| rows as usize != PART_ROWS)
{
return Err(malformed(format!(
"part {which} of stripe {at} holds {} rows and only the last part of a \
stripe may be short",
earlier[which]
)));
}
let rows = u64::from(total);
stripes.push(StripeSum { base, rows, parts: cumulative, uniform });
base = base
.checked_add(rows)
.ok_or_else(|| malformed("the table overflows a sixty four bit row count"))?;
}
Ok(Self { stripes, rows: base })
}
#[must_use]
pub fn rows(&self) -> u64 {
self.rows
}
#[must_use]
pub fn bytes(&self) -> usize {
let per_stripe = size_of::<StripeSum>();
self.stripes.iter().map(|stripe| per_stripe + stripe.parts.len() * size_of::<u32>()).sum()
}
#[must_use]
pub fn place(&self, rid: Rid) -> Option<Place> {
if rid >= self.rows {
return None;
}
let at = self.stripes.partition_point(|stripe| stripe.base <= rid) - 1;
let stripe = &self.stripes[at];
let within = rid - stripe.base;
#[expect(
clippy::cast_possible_truncation,
reason = "a stripe holds at most 65,536 rows, so `within` fits a u32"
)]
let within = within as u32;
let (part, offset) = if stripe.uniform {
(within >> PART_SHIFT, within & PART_MASK)
} else {
let part = stripe.parts.partition_point(|&before| before <= within) - 1;
#[expect(
clippy::cast_possible_truncation,
reason = "a stripe holds at most sixty four parts"
)]
let part = part as u32;
(part, within - stripe.parts[part as usize])
};
#[expect(
clippy::cast_possible_truncation,
reason = "a table holds at most 4,294,967,295 stripes and the build checked the count"
)]
let stripe = at as u32;
Some(Place { stripe, part, offset })
}
#[must_use]
pub fn rid(&self, place: Place) -> Option<Rid> {
let stripe = self.stripes.get(place.stripe as usize)?;
let before = *stripe.parts.get(place.part as usize)?;
let rows = stripe.parts.get(place.part as usize + 1)? - before;
if place.offset >= rows {
return None;
}
Some(stripe.base + u64::from(before + place.offset))
}
#[must_use]
pub fn stripe_rows(&self, stripe: u32) -> Option<u64> {
self.stripes.get(stripe as usize).map(|held| held.rows)
}
#[must_use]
pub fn stripes(&self) -> usize {
self.stripes.len()
}
}
const PART_SHIFT: u32 = PART_ROWS.trailing_zeros();
#[expect(
clippy::cast_possible_truncation,
reason = "PART_ROWS is 1024, so the mask fits a u32 with room to spare"
)]
const PART_MASK: u32 = (PART_ROWS - 1) as u32;
fn malformed(message: impl Into<String>) -> Error {
Error::invalid_input(format!("invalid rudb row id prefix sum: {}", message.into()))
}
#[cfg(test)]
mod tests {
use super::*;
fn ladder(stripes: usize, tail: u32) -> Vec<Vec<u32>> {
#[expect(clippy::cast_possible_truncation, reason = "PART_ROWS is 1024")]
let full = PART_ROWS as u32;
let mut out = vec![vec![full; STRIPE_PARTS]; stripes.saturating_sub(1)];
if stripes > 0 {
let whole = (tail / full) as usize;
let mut last = vec![full; whole];
if !tail.is_multiple_of(full) {
last.push(tail % full);
}
out.push(last);
}
out
}
#[test]
fn a_rid_resolves_to_the_part_and_the_offset_append_order_gave_it() {
let places = Places::build(&ladder(1, 3000)).expect("build");
assert_eq!(places.rows(), 3000);
assert_eq!(places.place(0), Some(Place { stripe: 0, part: 0, offset: 0 }));
assert_eq!(places.place(1023), Some(Place { stripe: 0, part: 0, offset: 1023 }));
assert_eq!(places.place(1024), Some(Place { stripe: 0, part: 1, offset: 0 }));
assert_eq!(places.place(2999), Some(Place { stripe: 0, part: 2, offset: 951 }));
assert_eq!(places.place(3000), None);
}
#[test]
fn every_rid_of_a_multi_stripe_table_round_trips_through_its_place() {
let places = Places::build(&ladder(3, 5000)).expect("build");
assert_eq!(places.rows(), 65_536 * 2 + 5000);
for rid in 0..places.rows() {
let place = places.place(rid).expect("every rid under the row count resolves");
assert_eq!(places.rid(place), Some(rid), "rid {rid} did not round trip");
}
assert_eq!(places.place(places.rows()), None);
}
#[test]
fn the_uniform_path_and_the_searching_path_agree_on_the_same_stripe() {
#[expect(clippy::cast_possible_truncation, reason = "PART_ROWS is 1024")]
let full = PART_ROWS as u32;
let fast = Places::build(&[vec![full; 8]]).expect("build");
let slow = Places::build(&[{
let mut parts = vec![full; 7];
parts.push(full - 1);
parts
}])
.expect("build");
for rid in 0..slow.rows() {
assert_eq!(fast.place(rid), slow.place(rid), "rid {rid}");
}
}
#[test]
fn a_hole_in_the_middle_of_a_stripe_is_refused_rather_than_resolved() {
#[expect(clippy::cast_possible_truncation, reason = "PART_ROWS is 1024")]
let full = PART_ROWS as u32;
let complaint = Places::build(&[vec![full, 7, full]]).expect_err("a short middle part");
let complaint = complaint.to_string();
assert!(complaint.contains("part 1"), "{complaint}");
assert!(complaint.contains("only the last part"), "{complaint}");
}
#[test]
fn a_part_wider_than_a_part_is_refused() {
#[expect(clippy::cast_possible_truncation, reason = "PART_ROWS is 1024")]
let over = PART_ROWS as u32 + 1;
let complaint = Places::build(&[vec![over]]).expect_err("an oversized part");
assert!(complaint.to_string().contains("at most"), "{complaint}");
}
#[test]
fn a_stripe_wider_than_a_stripe_is_refused() {
#[expect(clippy::cast_possible_truncation, reason = "PART_ROWS is 1024")]
let full = PART_ROWS as u32;
let complaint =
Places::build(&[vec![full; STRIPE_PARTS + 1]]).expect_err("an oversized stripe");
assert!(complaint.to_string().contains("at most 64"), "{complaint}");
}
#[test]
fn an_empty_table_resolves_nothing_and_says_so_rather_than_panicking() {
let places = Places::build(&[]).expect("build");
assert_eq!(places.rows(), 0);
assert_eq!(places.place(0), None);
assert_eq!(places.stripes(), 0);
}
#[test]
fn the_prefix_sum_of_a_hundred_million_rows_is_about_four_hundred_kilobytes() {
let stripes = 100_000_000 / (PART_ROWS * STRIPE_PARTS) + 1;
let places = Places::build(&ladder(stripes, 4096)).expect("build");
let bytes = places.bytes();
assert!(bytes < 600 * 1024, "the prefix sum took {bytes} bytes");
assert!(bytes > 200 * 1024, "the prefix sum took {bytes} bytes, which is suspiciously few");
}
#[test]
fn no_parent_is_not_a_rid_any_table_can_resolve() {
let places = Places::build(&ladder(2, 1)).expect("build");
assert_eq!(places.place(NO_PARENT), None);
}
}