use crate::{
error::{IndexChangeError, PIRes},
id::RecRef,
index::{
config::{IndexOrd, IndexTypeInternal, ValueMode},
tree::PosRef,
},
};
use std::{
cmp::Ordering,
fmt::Display,
iter::{Peekable, Rev},
ops::Bound,
vec::IntoIter,
};
pub(crate) fn block_size_for_split(len: usize, limit: usize) -> usize {
len / (len / limit + 1) + 1
}
pub(crate) fn block_counts_for_split(len: usize, limit: usize) -> usize {
let size = block_size_for_split(len, limit);
let count = len / size;
if len % size != 0 { count + 1 } else { count }
}
pub type TreeNodeRef = RecRef;
#[derive(Clone)]
pub enum TreeNode<K, V> {
Node(Node<K>),
Leaf(Leaf<K, V>),
}
impl<K, V> TreeNode<K, V> {
pub fn get_prev(&self) -> &Option<K> {
match self {
TreeNode::Node(n) => n.get_prev(),
TreeNode::Leaf(l) => l.get_prev(),
}
}
pub fn get_next(&self) -> &Option<K> {
match self {
TreeNode::Node(n) => n.get_next(),
TreeNode::Leaf(l) => l.get_next(),
}
}
pub fn len(&self) -> usize {
match self {
TreeNode::Node(n) => n.len(),
TreeNode::Leaf(l) => l.len(),
}
}
pub fn as_mut_node(&mut self) -> &mut Node<K> {
match self {
TreeNode::Node(n) => n,
_ => panic!("not a node"),
}
}
pub fn as_mut_leaf(&mut self) -> &mut Leaf<K, V> {
match self {
TreeNode::Node(_) => panic!("not a leaf"),
TreeNode::Leaf(l) => l,
}
}
}
impl<K: IndexOrd + Clone, V> TreeNode<K, V> {
pub fn merge_right(&mut self, k: &K, node: &mut TreeNode<K, V>) {
match self {
TreeNode::Node(n) => match node {
TreeNode::Node(n1) => {
n.merge_right(k, n1);
}
TreeNode::Leaf(_) => {
panic!("impossible merge a leaf to node");
}
},
TreeNode::Leaf(l) => match node {
TreeNode::Node(_) => {
panic!("impossible merge a node to leaf");
}
TreeNode::Leaf(l1) => {
l.merge_right(k, l1);
}
},
}
}
pub fn split(&mut self, top_limit: usize) -> Vec<(K, TreeNode<K, V>)> {
match self {
TreeNode::Node(n) => n
.split(top_limit)
.into_iter()
.map(|x| (x.0, TreeNode::Node(x.1)))
.collect(),
TreeNode::Leaf(l) => l
.split(top_limit)
.into_iter()
.map(|x| (x.0, TreeNode::Leaf(x.1)))
.collect(),
}
}
pub fn check_range(&self, k: &K) -> bool {
match self {
TreeNode::Node(n) => n.check_range(k),
TreeNode::Leaf(l) => l.check_range(k),
}
}
}
impl<K: Display, V: Display> Display for TreeNode<K, V> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match &self {
TreeNode::Node(n) => write!(f, "Node: {}", n)?,
TreeNode::Leaf(l) => write!(f, "Leaf: {}", l)?,
}
Ok(())
}
}
pub(crate) fn compare<T: IndexOrd>(first: &T, second: &T) -> Ordering {
first.cmp(second)
}
pub struct PosIter<'a, K> {
nodes: &'a Node<K>,
pos: usize,
}
impl<'a, K> PosIter<'a, K> {
fn new(nodes: &'a Node<K>, pos: usize) -> Self {
Self { nodes, pos }
}
}
impl<K> std::iter::Iterator for PosIter<'_, K>
where
K: Clone,
{
type Item = PosRef<K>;
fn next(&mut self) -> Option<Self::Item> {
if self.pos < self.nodes.len() {
let pointer = self.pos;
self.pos += 1;
let k = if pointer == 0 {
self.nodes.prev.clone().expect("should be called on middle nodes")
} else {
self.nodes.keys[pointer - 1].clone()
};
Some(PosRef::new(&k, pointer, self.nodes.pointers[pointer]))
} else {
None
}
}
}
#[derive(Clone)]
pub struct Node<K> {
pub keys: Vec<K>,
pub pointers: Vec<TreeNodeRef>,
pub prev: Option<K>,
pub next: Option<K>,
}
impl<K> Node<K> {
pub fn len(&self) -> usize {
self.pointers.len()
}
pub(crate) fn get_next(&self) -> &Option<K> {
&self.next
}
pub(crate) fn get_prev(&self) -> &Option<K> {
&self.prev
}
}
impl<K: IndexOrd + Clone> Node<K> {
pub fn new_from_split(left: TreeNodeRef, values: &[(K, TreeNodeRef)]) -> Node<K> {
let keys = values.iter().map(|z| z.0.clone()).collect();
let mut pointers: Vec<TreeNodeRef> = values.iter().map(|z| z.1).collect();
pointers.insert(0, left);
Node {
keys,
pointers,
prev: None,
next: None,
}
}
pub fn remove_key(&mut self, k: &K) -> Option<K> {
match self.keys.binary_search_by(|x| compare(x, k)) {
Ok(index) => {
self.keys.remove(index);
self.pointers.remove(index + 1);
None
}
Err(index) => {
if index == 0 {
let prev = if !self.keys.is_empty() {
Some(self.keys.remove(0))
} else {
self.next.clone()
};
self.prev = prev.clone();
self.pointers.remove(0);
prev
} else {
unreachable!("when removing a key it should be missing only for index 0");
}
}
}
}
pub fn add(&mut self, pos: usize, k: &K, node_ref: TreeNodeRef) {
self.keys.insert(pos, k.clone());
self.pointers.insert(pos + 1, node_ref);
}
pub fn find(&self, k: &K) -> PosRef<K> {
match self.keys.binary_search_by(|x| compare(x, k)) {
Ok(index) => PosRef::new(&self.keys[index], index + 1, self.pointers[index + 1]),
Err(index) => {
if index == 0 {
PosRef::new(self.prev.as_ref().unwrap_or(k), index, self.pointers[index])
} else {
PosRef::new(&self.keys[index - 1], index, self.pointers[index])
}
}
}
}
pub fn iter_from<'a>(&'a self, pos: Option<&PosRef<K>>) -> Option<PosIter<'a, K>> {
if pos.is_none() && self.prev.is_none() {
return None;
}
let base = pos.map(|x| x.pos + 1).unwrap_or(0);
if base > self.len() {
return None;
}
Some(PosIter::new(self, base))
}
pub fn swap_next(&mut self, new: &K) -> bool {
if let Some(next) = &self.next {
if compare(next, new) != Ordering::Equal {
self.next = Some(new.clone());
true
} else {
false
}
} else {
false
}
}
pub fn find_pre_write(&self, k: &K) -> Option<PosRef<K>> {
let pos = match self.keys.binary_search_by(|x| compare(x, k)) {
Ok(index) => {
if index == 0 {
PosRef::new(self.prev.as_ref().unwrap_or(k), index, self.pointers[index])
} else {
PosRef::new(&self.keys[index - 1], index, self.pointers[index])
}
}
Err(index) => {
if index == 0 {
PosRef::new(self.prev.as_ref().unwrap_or(k), index, self.pointers[index])
} else {
PosRef::new(&self.keys[index - 1], index, self.pointers[index])
}
}
};
if pos.pos == 0 {
if let Some(pk) = &self.prev {
if compare(k, pk) == Ordering::Less {
return None;
}
}
} else if pos.pos == self.pointers.len() {
if let Some(nk) = &self.next {
if compare(k, nk) != Ordering::Less {
return None;
}
}
}
Some(pos)
}
pub fn find_write(&self, k: &K) -> Option<PosRef<K>> {
let pos = self.find(k);
if pos.pos == 0 {
if let Some(pk) = &self.prev {
if compare(k, pk) == Ordering::Less {
return None;
}
}
} else if pos.pos == self.pointers.len() {
if let Some(nk) = &self.next {
if compare(k, nk) != Ordering::Less {
return None;
}
}
}
Some(pos)
}
pub fn get(&self, pos: usize) -> TreeNodeRef {
self.pointers[pos]
}
pub fn insert_after(&mut self, pos: usize, values: &[(K, TreeNodeRef)]) {
for (key, node) in values.iter().rev() {
self.add(pos, key, *node);
}
}
pub fn insert_after_key(&mut self, values: &[(K, TreeNodeRef)]) {
if let Some((k, _)) = values.first() {
let pos = self.find(k);
self.insert_after(pos.pos, values);
}
}
pub fn split(&mut self, max: usize) -> Vec<(K, Node<K>)> {
let mut split_result: Vec<(K, Node<K>)> = Vec::new();
let size = self.keys.len();
let split_offset = block_size_for_split(size, max);
let mut others = self.keys.split_off(split_offset - 1);
let mut other_pointers = self.pointers.split_off(split_offset);
let pre_next = self.next.clone();
while others.len() > max {
let new = others.split_off(split_offset);
let new_pointers = other_pointers.split_off(split_offset);
let key = others.remove(0);
if let Some((_, x)) = split_result.last_mut() {
x.next = Some(key.clone());
} else {
self.next = Some(key.clone());
}
let leaf = Node {
keys: others,
pointers: other_pointers,
prev: Some(key.clone()),
next: None,
};
split_result.push((key, leaf));
others = new;
other_pointers = new_pointers;
}
let key = others.remove(0);
if let Some((_, x)) = split_result.last_mut() {
x.next = Some(key.clone());
} else {
self.next = Some(key.clone());
}
let leaf = Node {
keys: others,
pointers: other_pointers,
prev: Some(key.clone()),
next: pre_next,
};
split_result.push((key, leaf));
split_result
}
#[cfg(test)]
pub fn merge_left(&mut self, owner: K, nodes: &mut Node<K>) {
let mut keys = std::mem::take(&mut nodes.keys);
let mut pointers = std::mem::take(&mut nodes.pointers);
keys.push(owner);
keys.append(&mut self.keys);
pointers.append(&mut self.pointers);
self.keys = keys;
self.pointers = pointers;
}
pub fn merge_right(&mut self, k: &K, nodes: &mut Node<K>) {
self.keys.push(k.clone());
self.keys.append(&mut nodes.keys);
self.pointers.append(&mut nodes.pointers);
self.next = nodes.next.clone();
}
pub fn replace_key(&mut self, node: &TreeNodeRef, new_key: &K) -> (bool, bool) {
for (p, r) in self.pointers.iter().enumerate() {
if r == node {
return if p == 0 {
if let Some(p) = &self.prev {
if compare(p, new_key) != Ordering::Equal {
self.prev = Some(new_key.clone());
(true, true)
} else {
(false, false)
}
} else {
(false, false)
}
} else if compare(&self.keys[p - 1], new_key) != Ordering::Equal {
self.keys[p - 1] = new_key.clone();
(true, false)
} else {
(false, false)
};
}
}
(false, false)
}
pub fn find_key(&self, node: &TreeNodeRef) -> Option<&K> {
for (p, r) in self.pointers.iter().enumerate() {
if r == node {
if p == 0 {
return self.prev.as_ref();
} else {
return Some(&self.keys[p - 1]);
}
}
}
None
}
fn check_range(&self, k: &K) -> bool {
if let Some(x) = &self.prev {
if compare(x, k) == Ordering::Greater {
return false;
}
}
if let Some(x) = &self.next {
if compare(x, k) == Ordering::Less {
return false;
}
}
true
}
}
impl<K: Display> Display for Node<K> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "[")?;
for x in 0..self.pointers.len() {
if x == 0 {
write!(
f,
"{:?}={}",
self.prev.as_ref().map(|p| format!("{}", p)),
self.pointers[x]
)?;
} else {
write!(f, ",{}={}", self.keys[x - 1], self.pointers[x])?;
}
}
write!(f, "]")?;
Ok(())
}
}
#[derive(Clone, PartialEq, Debug, Eq)]
pub enum Value<V> {
Cluster(Vec<V>),
Single(V),
}
impl<V> IntoIterator for Value<V> {
type Item = V;
type IntoIter = IntoIter<V>;
fn into_iter(self) -> IntoIter<V> {
match self {
Value::Single(v) => vec![v].into_iter(),
Value::Cluster(v) => v.into_iter(),
}
}
}
pub struct PageIter<K, V> {
pub iter: Peekable<IntoIter<LeafEntry<K, V>>>,
}
pub struct PageIterBack<K, V> {
pub iter: Peekable<Rev<IntoIter<LeafEntry<K, V>>>>,
}
#[derive(Clone)]
pub struct Leaf<K, V> {
pub entries: Vec<LeafEntry<K, V>>,
pub prev: Option<K>,
pub next: Option<K>,
}
#[derive(Clone)]
pub struct LeafEntry<K, V> {
pub key: K,
pub value: Value<V>,
}
impl<K, V> Leaf<K, V> {
pub fn new() -> Leaf<K, V> {
Leaf {
entries: Vec::new(),
prev: None,
next: None,
}
}
pub fn len(&self) -> usize {
self.entries.len()
}
#[cfg(test)]
pub fn merge_left(&mut self, leaf: &mut Leaf<K, V>) {
let mut entries = std::mem::take(&mut leaf.entries);
entries.append(&mut self.entries);
self.entries = entries;
}
pub(crate) fn get_next(&self) -> &Option<K> {
&self.next
}
pub(crate) fn get_prev(&self) -> &Option<K> {
&self.prev
}
}
impl<K: IndexOrd + Clone, V> Leaf<K, V> {
fn check_range(&self, k: &K) -> bool {
if let Some(x) = &self.prev {
if compare(x, k) == Ordering::Greater {
return false;
}
}
if let Some(x) = &self.next {
if compare(x, k) == Ordering::Less {
return false;
}
}
true
}
pub fn merge_right(&mut self, _k: &K, leaf: &mut Leaf<K, V>) {
self.entries.append(&mut leaf.entries);
self.next = leaf.next.clone();
}
pub fn split(&mut self, max: usize) -> Vec<(K, Leaf<K, V>)> {
let mut split_result: Vec<(K, Leaf<K, V>)> = Vec::new();
let size = self.entries.len();
let split_offset = block_size_for_split(size, max);
let mut others = self.entries.split_off(split_offset);
let pre_next = self.next.clone();
while others.len() > max {
let new = others.split_off(split_offset);
let key = others[0].key.clone();
if let Some((_, x)) = split_result.last_mut() {
x.next = Some(key.clone());
} else {
self.next = Some(key.clone());
}
let leaf = Leaf {
entries: others,
prev: Some(key.clone()),
next: None,
};
split_result.push((key, leaf));
others = new;
}
let key = others[0].key.clone();
if let Some((_, x)) = split_result.last_mut() {
x.next = Some(key.clone());
} else {
self.next = Some(key.clone());
}
let leaf = Leaf {
entries: others,
prev: Some(key.clone()),
next: pre_next,
};
split_result.push((key, leaf));
split_result
}
}
impl<K: IndexOrd + Clone, V: Clone> Leaf<K, V> {
pub fn find(&self, k: &K) -> Result<(K, Value<V>), usize> {
self.entries
.binary_search_by(|n| compare(&n.key, k))
.map(|index| (self.entries[index].key.clone(), self.entries[index].value.clone()))
}
pub fn add_at(&mut self, pos: usize, k: &K, v: &V, _value_mode: ValueMode) {
self.entries.insert(
pos,
LeafEntry {
key: k.clone(),
value: Value::Single(v.clone()),
},
);
}
pub fn iter_from(&self, bound: Bound<&K>) -> IntoIter<LeafEntry<K, V>> {
let index = match bound {
Bound::Included(k) => match self.entries.binary_search_by(|n| compare(&n.key, k)) {
Ok(index) => index,
Err(index) => index,
},
Bound::Excluded(k) => match self.entries.binary_search_by(|n| compare(&n.key, k)) {
Ok(index) => index + 1,
Err(index) => index,
},
Bound::Unbounded => 0,
};
self.entries[index..].to_vec().into_iter()
}
pub fn back_iter_from(&self, bound: Bound<&K>) -> Rev<IntoIter<LeafEntry<K, V>>> {
let index = match bound {
Bound::Included(k) => match self.entries.binary_search_by(|n| compare(&n.key, k)) {
Ok(index) => index + 1,
Err(index) => index,
},
Bound::Excluded(k) => match self.entries.binary_search_by(|n| compare(&n.key, k)) {
Ok(index) => index,
Err(index) => index,
},
Bound::Unbounded => self.len(),
};
self.entries[..index].to_vec().into_iter().rev()
}
}
impl<K: IndexTypeInternal, V: IndexOrd + Clone> Leaf<K, V> {
pub fn add(&mut self, k: &K, v: &V, value_mode: ValueMode, index_name: &str) -> PIRes<()> {
if self.len() > 0 {
self.insert_or_update(k, v, value_mode, index_name)
} else {
self.add_at(0, k, v, value_mode);
Ok(())
}
}
pub fn insert_or_update(&mut self, k: &K, v: &V, value_mode: ValueMode, index_name: &str) -> PIRes<()> {
match self.entries.binary_search_by(|n| compare(&n.key, k)) {
Ok(index) => {
let entry = &mut self.entries[index];
match value_mode {
ValueMode::Replace => {
entry.value = Value::Single(v.clone());
}
ValueMode::Exclusive => match entry.value {
Value::Single(ref ev) => {
if compare(ev, v) != Ordering::Equal {
return Err(IndexChangeError::IndexDuplicateKey(
index_name.to_string(),
format!("{}", k),
));
}
}
_ => unreachable!("Exclusive leaves never have cluster values"),
},
ValueMode::Cluster => {
let mut new_value = None;
match entry.value {
Value::Single(ref ev) => match compare(ev, v) {
Ordering::Equal => {}
Ordering::Less => {
new_value = Some(Value::Cluster(vec![ev.clone(), v.clone()]));
}
Ordering::Greater => {
new_value = Some(Value::Cluster(vec![v.clone(), ev.clone()]));
}
},
Value::Cluster(ref mut cl) => {
if let Err(index) = cl.binary_search_by(|x| compare(x, v)) {
cl.insert(index, v.clone());
}
}
}
if let Some(v) = new_value {
entry.value = v;
}
}
}
}
Err(index) => self.add_at(index, k, v, value_mode),
}
Ok(())
}
pub fn remove(&mut self, k: &K, v: &Option<V>) -> bool {
match self.entries.binary_search_by(|n| compare(&n.key, k)) {
Ok(index) => {
if let Some(rv) = v {
let mut removed = false;
let remove_entry = {
let mut new_value = None;
let entry = &mut self.entries[index];
let remove_entry = match &mut entry.value {
Value::Single(val) => {
if compare(val, rv) == Ordering::Equal {
removed = true;
true
} else {
false
}
}
Value::Cluster(cl) => {
if let Ok(index) = cl.binary_search_by(|x| compare(x, rv)) {
removed = true;
cl.remove(index);
}
if cl.len() == 1 {
new_value = Some(Value::Single(cl.pop().unwrap()));
false
} else {
cl.is_empty()
}
}
};
if let Some(new) = new_value {
entry.value = new;
}
remove_entry
};
if remove_entry {
self.entries.remove(index);
}
removed
} else {
self.entries.remove(index);
true
}
}
Err(_) => false,
}
}
}
impl<K: Display, V: Display> Display for Leaf<K, V> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "[")?;
for x in 0..self.entries.len() {
write!(f, "{}=[", self.entries[x].key)?;
match &self.entries[x].value {
Value::Single(v) => write!(f, "{}", v)?,
Value::Cluster(c) => {
for (c, ev) in c.iter().enumerate() {
write!(f, "{}", ev)?;
if c != 0 {
write!(f, ",")?;
}
}
}
}
write!(f, "]")?;
if x != 0 {
write!(f, ",")?;
}
}
write!(f, "]")?;
Ok(())
}
}