use std::cmp::Ordering;
use std::sync::atomic::{AtomicU32, Ordering as AtomicOrdering};
use std::sync::Arc;
use rand::Rng;
use crate::error::Result as CrateResult;
use crate::memtable::arena::Arena;
use crate::{InternalKey, InternalKeyRef, LSMIterator};
const MAX_HEIGHT: usize = 20;
const P_VALUE: f64 = 1.0 / std::f64::consts::E;
const LINKS_SIZE: usize = std::mem::size_of::<Links>();
const MAX_NODE_SIZE: usize = std::mem::size_of::<Node>() + (MAX_HEIGHT - 1) * LINKS_SIZE;
fn probabilities() -> &'static [u32; MAX_HEIGHT] {
static PROBABILITIES: std::sync::OnceLock<[u32; MAX_HEIGHT]> = std::sync::OnceLock::new();
PROBABILITIES.get_or_init(|| {
let mut p = [0u32; MAX_HEIGHT];
let mut prob = 1.0f64;
for p_item in p.iter_mut() {
*p_item = (u32::MAX as f64 * prob) as u32;
prob *= P_VALUE;
}
p
})
}
#[derive(Debug, Clone, PartialEq)]
pub enum Error {
RecordExists,
ArenaFull,
}
pub type Compare = fn(&[u8], &[u8]) -> Ordering;
#[repr(C)]
struct Links {
next_offset: AtomicU32,
prev_offset: AtomicU32,
}
impl Links {
#[inline]
fn init(&self, prev_offset: u32, next_offset: u32) {
self.next_offset.store(next_offset, AtomicOrdering::Release);
self.prev_offset.store(prev_offset, AtomicOrdering::Release);
}
}
#[repr(C)]
struct Node {
key_offset: u32,
key_size: u32,
key_trailer: u64,
key_timestamp: u64,
value_size: u32,
_padding: u32,
tower: [Links; 1],
}
impl Node {
#[inline]
fn get_key_bytes<'a>(&self, arena: &'a Arena) -> &'a [u8] {
arena.get_bytes(self.key_offset, self.key_size)
}
#[inline]
fn get_value<'a>(&self, arena: &'a Arena) -> &'a [u8] {
arena.get_bytes(self.key_offset + self.key_size, self.value_size)
}
#[inline]
fn next_offset(&self, h: usize) -> u32 {
unsafe {
let links = (self.tower.as_ptr()).add(h);
(*links).next_offset.load(AtomicOrdering::Acquire)
}
}
#[inline]
fn prev_offset(&self, h: usize) -> u32 {
unsafe {
let links = (self.tower.as_ptr()).add(h);
(*links).prev_offset.load(AtomicOrdering::Acquire)
}
}
#[inline]
fn cas_next_offset(&self, h: usize, old: u32, val: u32) -> bool {
unsafe {
let links = (self.tower.as_ptr()).add(h);
(*links)
.next_offset
.compare_exchange(old, val, AtomicOrdering::AcqRel, AtomicOrdering::Acquire)
.is_ok()
}
}
#[inline]
fn cas_prev_offset(&self, h: usize, old: u32, val: u32) -> bool {
unsafe {
let links = (self.tower.as_ptr()).add(h);
(*links)
.prev_offset
.compare_exchange(old, val, AtomicOrdering::AcqRel, AtomicOrdering::Acquire)
.is_ok()
}
}
#[inline]
fn tower_init(&self, h: usize, prev_offset: u32, next_offset: u32) {
unsafe {
let links = (self.tower.as_ptr()).add(h);
(*links).init(prev_offset, next_offset);
}
}
}
fn new_node(
arena: &Arena,
height: u32,
key: &[u8],
trailer: u64,
timestamp: u64,
value: &[u8],
) -> Option<*mut Node> {
if height < 1 || height > MAX_HEIGHT as u32 {
panic!("height cannot be less than one or greater than the max height");
}
let nd = new_raw_node(arena, height, key.len() as u32, value.len() as u32)?;
unsafe {
(*nd).key_trailer = trailer;
(*nd).key_timestamp = timestamp;
let key_bytes = arena.get_bytes_mut((*nd).key_offset, (*nd).key_size);
key_bytes.copy_from_slice(key);
let val_bytes = arena.get_bytes_mut((*nd).key_offset + (*nd).key_size, (*nd).value_size);
val_bytes.copy_from_slice(value);
}
Some(nd)
}
fn new_raw_node(arena: &Arena, height: u32, key_size: u32, value_size: u32) -> Option<*mut Node> {
let unused_size = (MAX_HEIGHT - height as usize) * LINKS_SIZE;
let node_size = MAX_NODE_SIZE - unused_size;
let node_offset = arena.alloc(
(node_size as u32) + key_size + value_size,
8, unused_size as u32,
)?;
let nd = arena.get_pointer(node_offset) as *mut Node;
unsafe {
(*nd).key_offset = node_offset + node_size as u32;
(*nd).key_size = key_size;
(*nd).value_size = value_size;
}
Some(nd)
}
#[derive(Clone, Copy, Default)]
struct Splice {
prev: *mut Node,
next: *mut Node,
}
impl Splice {
#[inline]
fn init(&mut self, prev: *mut Node, next: *mut Node) {
self.prev = prev;
self.next = next;
}
}
#[derive(Default)]
pub struct Inserter {
spl: [Splice; MAX_HEIGHT],
height: u32,
}
impl Inserter {
pub fn new() -> Self {
Self::default()
}
}
pub struct Skiplist {
arena: Arc<Arena>,
cmp: Compare,
head: *mut Node,
tail: *mut Node,
height: AtomicU32,
}
unsafe impl Send for Skiplist {}
unsafe impl Sync for Skiplist {}
impl Skiplist {
pub fn new(arena: Arc<Arena>, cmp: Compare) -> Self {
let head = new_raw_node(&arena, MAX_HEIGHT as u32, 0, 0)
.expect("arenaSize is not large enough to hold the head node");
unsafe {
(*head).key_offset = 0;
}
let tail = new_raw_node(&arena, MAX_HEIGHT as u32, 0, 0)
.expect("arenaSize is not large enough to hold the tail node");
unsafe {
(*tail).key_offset = 0;
}
let head_offset = arena.get_pointer_offset(head as *const u8);
let tail_offset = arena.get_pointer_offset(tail as *const u8);
for i in 0..MAX_HEIGHT {
unsafe {
(*head).tower_init(i, 0, tail_offset);
(*tail).tower_init(i, head_offset, 0);
}
}
Self {
arena,
cmp,
head,
tail,
height: AtomicU32::new(1),
}
}
#[inline]
pub fn height(&self) -> u32 {
self.height.load(AtomicOrdering::Acquire)
}
#[inline]
pub fn size(&self) -> u32 {
self.arena.size() as u32
}
pub fn add(&self, key: &[u8], trailer: u64, timestamp: u64, value: &[u8]) -> Result<(), Error> {
let mut ins = Inserter::new();
self.add_internal(key, trailer, timestamp, value, &mut ins)
}
fn add_internal(
&self,
key: &[u8],
trailer: u64,
timestamp: u64,
value: &[u8],
ins: &mut Inserter,
) -> Result<(), Error> {
if self.find_splice(key, trailer, timestamp, ins) {
return Err(Error::RecordExists);
}
let (nd, height) = self.new_node(key, trailer, timestamp, value)?;
let nd_offset = self.arena.get_pointer_offset(nd as *const u8);
let mut invalidate_splice = false;
for i in 0..height as usize {
let mut prev = ins.spl[i].prev;
let mut next = ins.spl[i].next;
if prev.is_null() {
if !next.is_null() {
panic!("next is expected to be nil, since prev is nil");
}
prev = self.head;
next = self.tail;
}
loop {
let prev_offset = self.arena.get_pointer_offset(prev as *const u8);
let next_offset = self.arena.get_pointer_offset(next as *const u8);
unsafe {
(*nd).tower_init(i, prev_offset, next_offset);
let next_prev_offset = (*next).prev_offset(i);
if next_prev_offset != prev_offset {
let prev_next_offset = (*prev).next_offset(i);
if prev_next_offset == next_offset {
(*next).cas_prev_offset(i, next_prev_offset, prev_offset);
}
}
if (*prev).cas_next_offset(i, next_offset, nd_offset) {
(*next).cas_prev_offset(i, prev_offset, nd_offset);
break;
}
}
let (new_prev, new_next, found) =
self.find_splice_for_level(key, trailer, timestamp, i, prev);
if found {
if i != 0 {
panic!("how can another thread have inserted a node at a non-base level?");
}
return Err(Error::RecordExists);
}
prev = new_prev;
next = new_next;
invalidate_splice = true;
}
}
if invalidate_splice {
ins.height = 0;
} else {
for i in 0..height as usize {
ins.spl[i].prev = nd;
}
}
Ok(())
}
fn new_node(
&self,
key: &[u8],
trailer: u64,
timestamp: u64,
value: &[u8],
) -> Result<(*mut Node, u32), Error> {
let height = self.random_height();
let nd = new_node(&self.arena, height, key, trailer, timestamp, value)
.ok_or(Error::ArenaFull)?;
let mut list_height = self.height();
while height > list_height {
match self.height.compare_exchange_weak(
list_height,
height,
AtomicOrdering::AcqRel,
AtomicOrdering::Acquire,
) {
Ok(_) => break,
Err(h) => list_height = h,
}
}
Ok((nd, height))
}
fn random_height(&self) -> u32 {
let rnd: u32 = rand::rng().random();
let mut h = 1u32;
let probs = probabilities();
while h < MAX_HEIGHT as u32 && rnd <= probs[h as usize] {
h += 1;
}
h
}
fn find_splice(&self, key: &[u8], trailer: u64, timestamp: u64, ins: &mut Inserter) -> bool {
let list_height = self.height();
let mut level: i32;
let mut prev = self.head;
if ins.height < list_height {
ins.height = list_height;
level = ins.height as i32;
} else {
level = 0;
for l in 0..list_height as usize {
let spl = &ins.spl[l];
if self.get_next(spl.prev, l) != spl.next {
continue;
}
if (spl.prev != self.head
&& !self.key_is_after_node(spl.prev, key, trailer, timestamp))
|| (spl.next != self.tail
&& self.key_is_after_node(spl.next, key, trailer, timestamp))
{
level = list_height as i32;
} else {
prev = spl.prev;
}
break;
}
}
let mut found = false;
let mut next: *mut Node = std::ptr::null_mut();
for l in (0..level as usize).rev() {
let prev_level_next = next;
loop {
next = self.get_next(prev, l);
if next == prev_level_next {
break;
}
if next == self.tail {
break;
}
let next_key = unsafe { (*next).get_key_bytes(&self.arena) };
let cmp = (self.cmp)(key, next_key);
if cmp == Ordering::Less {
break;
}
if cmp == Ordering::Equal {
let next_trailer = unsafe { (*next).key_trailer };
if trailer == next_trailer {
let next_timestamp = unsafe { (*next).key_timestamp };
if timestamp == next_timestamp {
found = true;
break;
}
if timestamp > next_timestamp {
break;
}
} else if trailer > next_trailer {
break;
}
}
prev = next;
}
ins.spl[l].init(prev, next);
}
found
}
fn find_splice_for_level(
&self,
key: &[u8],
trailer: u64,
timestamp: u64,
level: usize,
start: *mut Node,
) -> (*mut Node, *mut Node, bool) {
let mut prev = start;
loop {
let next = self.get_next(prev, level);
if next == self.tail {
return (prev, next, false);
}
let next_key = unsafe { (*next).get_key_bytes(&self.arena) };
let cmp = (self.cmp)(key, next_key);
if cmp == Ordering::Less {
return (prev, next, false);
}
if cmp == Ordering::Equal {
let next_trailer = unsafe { (*next).key_trailer };
if trailer == next_trailer {
let next_timestamp = unsafe { (*next).key_timestamp };
if timestamp == next_timestamp {
return (prev, next, true); }
if timestamp > next_timestamp {
return (prev, next, false);
}
} else if trailer > next_trailer {
return (prev, next, false);
}
}
prev = next;
}
}
#[inline]
fn key_is_after_node(&self, nd: *mut Node, key: &[u8], trailer: u64, timestamp: u64) -> bool {
if nd == self.head {
return true;
}
if nd == self.tail {
return false;
}
let nd_key = unsafe { (*nd).get_key_bytes(&self.arena) };
let cmp = (self.cmp)(nd_key, key);
match cmp {
Ordering::Less => true,
Ordering::Greater => false,
Ordering::Equal => {
let nd_trailer = unsafe { (*nd).key_trailer };
if trailer == nd_trailer {
let nd_timestamp = unsafe { (*nd).key_timestamp };
if timestamp == nd_timestamp {
false } else {
timestamp < nd_timestamp
}
} else {
trailer < nd_trailer
}
}
}
}
#[inline]
fn get_next(&self, nd: *mut Node, h: usize) -> *mut Node {
let offset = unsafe { (*nd).next_offset(h) };
self.arena.get_pointer(offset) as *mut Node
}
#[inline]
fn get_prev(&self, nd: *mut Node, h: usize) -> *mut Node {
let offset = unsafe { (*nd).prev_offset(h) };
self.arena.get_pointer(offset) as *mut Node
}
pub(crate) fn iter(&self) -> SkiplistIterator<'_> {
self.new_iter(None, None)
}
pub(crate) fn new_iter<'b>(
&'b self,
lower: Option<&[u8]>,
upper: Option<&[u8]>,
) -> SkiplistIterator<'b> {
SkiplistIterator {
list: self,
nd: self.head,
lower: lower.map(|s| s.to_vec()),
upper: upper.map(|s| s.to_vec()),
lower_node: std::ptr::null_mut(),
upper_node: std::ptr::null_mut(),
encoded_key_buf: Vec::new(),
}
}
}
pub(crate) struct SkiplistIterator<'a> {
list: &'a Skiplist,
nd: *mut Node,
lower: Option<Vec<u8>>, upper: Option<Vec<u8>>, lower_node: *mut Node, upper_node: *mut Node, encoded_key_buf: Vec<u8>, }
impl<'a> SkiplistIterator<'a> {
#[inline]
pub fn is_valid(&self) -> bool {
self.nd != self.list.head
&& self.nd != self.list.tail
&& !self.nd.is_null()
&& self.nd != self.lower_node
&& self.nd != self.upper_node
}
#[inline]
pub fn key_bytes(&self) -> &[u8] {
debug_assert!(self.is_valid());
unsafe { (*self.nd).get_key_bytes(&self.list.arena) }
}
#[inline]
pub fn trailer(&self) -> u64 {
debug_assert!(self.is_valid());
unsafe { (*self.nd).key_trailer }
}
#[inline]
pub fn value_bytes(&self) -> &[u8] {
debug_assert!(self.is_valid());
unsafe { (*self.nd).get_value(&self.list.arena) }
}
fn populate_encoded_key(&mut self) {
if !self.is_valid() {
return;
}
let (key_bytes, trailer, timestamp) = unsafe {
let node = &*self.nd;
(node.get_key_bytes(&self.list.arena), node.key_trailer, node.key_timestamp)
};
self.encoded_key_buf.clear();
self.encoded_key_buf.extend_from_slice(key_bytes);
self.encoded_key_buf.extend_from_slice(&trailer.to_be_bytes());
self.encoded_key_buf.extend_from_slice(×tamp.to_be_bytes());
}
pub fn first(&mut self) {
if let Some(ref lower) = self.lower.clone() {
self.seek_ge(lower);
} else {
self.nd = self.list.get_next(self.list.head, 0);
}
if self.nd == self.list.tail || self.nd == self.upper_node {
return;
}
if let Some(upper) = self.upper.as_deref() {
if self.is_valid() {
let key = self.key_bytes();
if (self.list.cmp)(upper, key) <= Ordering::Equal {
self.upper_node = self.nd;
self.nd = self.list.tail;
}
}
}
}
pub fn last(&mut self) {
self.nd = self.list.get_prev(self.list.tail, 0);
if self.nd == self.list.head || self.nd == self.lower_node {
return;
}
if let Some(upper) = self.upper.as_deref() {
while self.is_valid() {
let key = self.key_bytes();
if (self.list.cmp)(upper, key) == Ordering::Greater {
break;
}
self.nd = self.list.get_prev(self.nd, 0);
if self.nd == self.list.head || self.nd == self.lower_node {
return;
}
}
}
if let Some(lower) = self.lower.as_deref() {
if self.is_valid() {
let key = self.key_bytes();
if (self.list.cmp)(lower, key) == Ordering::Greater {
self.lower_node = self.nd;
self.nd = self.list.head;
}
}
}
}
pub fn advance(&mut self) {
debug_assert!(self.is_valid());
self.nd = self.list.get_next(self.nd, 0);
if self.nd == self.list.tail || self.nd == self.upper_node {
return;
}
if let Some(upper) = self.upper.as_deref() {
if self.is_valid() {
let key = self.key_bytes();
if (self.list.cmp)(upper, key) <= Ordering::Equal {
self.upper_node = self.nd;
self.nd = self.list.tail;
}
}
}
}
pub fn prev_internal(&mut self) {
debug_assert!(self.is_valid());
self.nd = self.list.get_prev(self.nd, 0);
if self.nd == self.list.head || self.nd == self.lower_node {
return;
}
if let Some(lower) = self.lower.as_deref() {
if self.is_valid() {
let key = self.key_bytes();
if (self.list.cmp)(lower, key) == Ordering::Greater {
self.lower_node = self.nd;
self.nd = self.list.head;
}
}
}
}
pub fn seek_ge(&mut self, key: &[u8]) {
let (_, next) = self.seek_for_base_splice(key);
self.nd = next;
if self.nd == self.list.tail || self.nd == self.upper_node {
return;
}
if let Some(upper) = self.upper.as_deref() {
if self.is_valid() {
let current_key = self.key_bytes();
if (self.list.cmp)(upper, current_key) <= Ordering::Equal {
self.upper_node = self.nd;
self.nd = self.list.tail;
}
}
}
}
fn seek_for_base_splice(&self, key: &[u8]) -> (*mut Node, *mut Node) {
let mut prev = self.list.head;
let mut next: *mut Node = std::ptr::null_mut();
for level in (0..self.list.height() as usize).rev() {
let prev_level_next = next;
loop {
next = self.list.get_next(prev, level);
if next == prev_level_next {
break;
}
if next == self.list.tail {
break;
}
let next_key = unsafe { (*next).get_key_bytes(&self.list.arena) };
let cmp = (self.list.cmp)(key, next_key);
if cmp <= Ordering::Equal {
break;
}
prev = next;
}
}
(prev, next)
}
}
impl LSMIterator for SkiplistIterator<'_> {
fn seek(&mut self, target: &[u8]) -> CrateResult<bool> {
let user_key = InternalKey::user_key_from_encoded(target);
self.seek_ge(user_key);
self.populate_encoded_key();
Ok(self.is_valid())
}
fn seek_first(&mut self) -> CrateResult<bool> {
self.first();
self.populate_encoded_key();
Ok(self.is_valid())
}
fn seek_last(&mut self) -> CrateResult<bool> {
self.last();
self.populate_encoded_key();
Ok(self.is_valid())
}
fn next(&mut self) -> CrateResult<bool> {
if !self.is_valid() {
return Ok(false);
}
self.advance();
self.populate_encoded_key();
Ok(self.is_valid())
}
fn prev(&mut self) -> CrateResult<bool> {
if !self.is_valid() {
return Ok(false);
}
self.prev_internal();
self.populate_encoded_key();
Ok(self.is_valid())
}
fn valid(&self) -> bool {
self.is_valid()
}
fn key(&self) -> InternalKeyRef<'_> {
debug_assert!(self.is_valid());
InternalKeyRef::from_encoded(&self.encoded_key_buf)
}
fn value_encoded(&self) -> CrateResult<&[u8]> {
debug_assert!(self.is_valid());
Ok(self.value_bytes())
}
}