Skip to content

Instantly share code, notes, and snippets.

@totechite

totechite/b_plus_tree.rs

Created January 29, 2020 15:23
Show Gist options
  • Save totechite/f4ea94b6095b31d4f62e71baa48135a1 to your computer and use it in GitHub Desktop.
Save totechite/f4ea94b6095b31d4f62e71baa48135a1 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
b_plus_tree.rs
use std::fmt::Debug;
use std::rc::Rc;
use std::cell::RefCell;
use std::any::{Any, TypeId};
const M: usize = 5;
#[derive(Debug)]
pub struct BPlusTree<K, V>
where
K: PartialEq + PartialOrd + Ord + Copy + Clone + Debug + 'static,
V: PartialEq + Eq + Clone + Debug + 'static
{
pub root_node: Box<dyn Node<K, V>>,
}
pub trait Node<K, V>: Debug
where
K: PartialEq + PartialOrd + Ord + Copy + Clone + Debug + 'static,
V: PartialEq + Eq + Clone + Debug + 'static
{
//core
fn delete(&mut self, key: K) -> Status;
fn print(&self, buf: Vec<(K, V)>) -> Vec<(K, V)>;
fn find(&self, key: K) -> Option<(K, V)>;
fn insert(&mut self, k: K, v: V) -> Option<(K, Rc<RefCell<dyn Node<K, V>>>)>;
fn update(&mut self, key: K, data: V) -> bool;
//utils
fn is_internal_node(&self) -> bool;
fn is_leaf_node(&self) -> bool;
fn marge(&mut self, from: Rc<RefCell<dyn Node<K, V>>>) -> bool;
fn right_share(&mut self, from: Rc<RefCell<dyn Node<K, V>>>) -> bool;
fn left_share(&mut self, from: Rc<RefCell<dyn Node<K, V>>>) -> bool;
fn split(&mut self) -> (K, Rc<RefCell<dyn Node<K, V>>>);
fn cmp_key(&self) -> K;
fn take_items(&mut self) -> Rc<RefCell<dyn Node<K, V>>>;
fn reorg_key(&mut self, delete_key_locate: usize) -> Option<K>;
}
#[derive(Debug)]
pub enum Status {
OK,
OkReorg,
OkReorgLeft,
OkReorgRight,
NotFound,
}
#[derive(Debug)]
pub struct InternalNode<K, V>
where
K: PartialEq + PartialOrd + Ord + Copy + Clone + Debug + 'static,
V: PartialEq + Eq + Clone + Debug + 'static
{
keys: Vec<K>,
node: Vec<Rc<RefCell<dyn Node<K, V>>>>,
}
#[derive(Default, Clone, Debug)]
pub struct LeafNode<K, V>
where
K: PartialEq + PartialOrd + Ord + Copy + Clone + Debug + 'static,
V: PartialEq + Eq + Clone + Debug + 'static
{
data: Vec<(K, V)>,
next_leaf: Option<Rc<RefCell<LeafNode<K, V>>>>,
}
impl<K, V> BPlusTree<K, V>
where
K: PartialEq + PartialOrd + Ord + Copy + Clone + Debug + 'static,
V: PartialEq + Eq + Clone + Debug + 'static
{
pub fn new() -> Self {
BPlusTree {
root_node: Box::new(LeafNode::new()),
}
}
pub fn print(&mut self) -> Vec<(K, V)> {
self.root_node.print(Vec::new())
}
pub fn insert(&mut self, k: K, v: V) {
if let Some((new_key, new_node)) = self.root_node.insert(k, v) {
self.split(new_key, new_node);
}
}
pub fn find(&self, key: K) -> Option<(K, V)> {
self.root_node.find(key)
}
pub fn update(&mut self, key: K, data: V) -> bool {
self.root_node.update(key, data)
}
fn split(&mut self, new_key: K, new_node: Rc<RefCell<dyn Node<K, V>>>) {
let mut new_root = InternalNode::new();
new_root.keys.push(new_key);
new_root.node.push(new_node);
new_root.node.push(self.root_node.take_items());
new_root.keys.sort();
new_root.node.sort_by_key(|a| a.borrow().cmp_key());
self.root_node = Box::new(new_root);
}
pub fn delete(&mut self, key: K) -> bool {
match self.root_node.delete(key) {
Status::NotFound => false,
_ => true
}
}
}
impl<K, V> Node<K, V> for InternalNode<K, V>
where
K: PartialEq + PartialOrd + Ord + Copy + Clone + Debug + 'static,
V: PartialEq + Eq + Clone + Debug + 'static
{
fn is_internal_node(&self) -> bool {
TypeId::of::<InternalNode<K, V>>() == self.type_id()
}
fn is_leaf_node(&self) -> bool {
TypeId::of::<LeafNode<K, V>>() == self.type_id()
}
fn marge(&mut self, from: Rc<RefCell<dyn Node<K, V>>>) -> bool {
unsafe {
let from = std::mem::transmute::<Box<Rc<RefCell<dyn Node<K, V>>>>, Box<Rc<RefCell<InternalNode<K, V>>>>>(Box::new(from));
while let Some(elem) = from.borrow_mut().node.pop() {
self.node.push(elem);
}
while let Some(elem) = from.borrow_mut().keys.pop() {
if !self.keys.contains(&elem) {
self.keys.push(elem);
};
}
self.node.sort_by_key(|a| a.borrow().cmp_key());
self.keys.sort();
}
true
}
fn right_share(&mut self, from: Rc<RefCell<dyn Node<K, V>>>) -> bool {
unsafe {
let from = std::mem::transmute::<Box<Rc<RefCell<dyn Node<K, V>>>>, Box<Rc<RefCell<InternalNode<K, V>>>>>(Box::new(from));
let len = from.borrow().node.len().clone();
if M <= (self.node.len() + len) {
for _ in 0..(self.node.len() + len) / 2 - self.node.len() {
self.node.push(from.borrow_mut().node.pop().unwrap())
};
self.node.sort_by_key(|a| a.borrow().cmp_key());
true
} else { false }
}
}
fn left_share(&mut self, from: Rc<RefCell<dyn Node<K, V>>>) -> bool {
unsafe {
let from = std::mem::transmute::<Box<Rc<RefCell<dyn Node<K, V>>>>, Box<Rc<RefCell<InternalNode<K, V>>>>>(Box::new(from));
let len = from.borrow().node.len().clone();
if M <= (self.node.len() + len) {
for _ in 0..(self.node.len() + len) / 2 - self.node.len() {
self.node.push(from.borrow_mut().node.remove(0))
};
self.node.sort_by_key(|a| a.borrow().cmp_key());
true
} else { false }
}
}
fn update(&mut self, key: K, data: V) -> bool {
let node_address = {
let mut result = self.keys.len();
for (i, comp_key) in self.keys.iter().enumerate() {
if &key < comp_key {
result = i;
break;
};
}
result
};
self.node[node_address].borrow_mut().update(key, data)
}
fn delete(&mut self, key: K) -> Status {
self.node.sort_by_key(|a| a.borrow().cmp_key());
let node_address = {
let mut result = self.keys.len();
for (i, comp_key) in self.keys.iter().enumerate() {
if &key < comp_key {
result = i;
break;
};
}
result
};
let delete_status = self.node[node_address].borrow_mut().delete(key);
let mut reorg_key_frag = None;
for (i, k) in self.keys.iter().enumerate() {
if &key == k {
reorg_key_frag = Some(i);
}
}
if let Some(i) = reorg_key_frag {
self.reorg_key(i);
}
let result = match delete_status {
Status::OkReorgLeft =>
{
let tmp = self.node.remove(node_address - 1);
if M - 1 <= self.keys.len() {
self.node[node_address].borrow_mut().marge(tmp);
} else {
self.marge(tmp);
// if self.keys.len() >= M {
// self.node.sort_by_key(|a| a.borrow().cmp_key());
// let left = self.split();
// self.node.push(left.1);
// }
}
Status::OK
}
Status::OkReorgRight =>
{
let tmp = self.node.remove(node_address + 1);
if M - 1 <= self.keys.len() {
self.node[node_address].borrow_mut().marge(tmp);
} else {
self.marge(tmp);
}
Status::OK
}
Status::OkReorg =>
{
if self.node[node_address].borrow().is_internal_node() {
if (self.node.len() - 1) == node_address {
// 対象が最右の場合
unsafe {
let right = self.node.remove(node_address);
let left = std::mem::transmute::<Box<Rc<RefCell<dyn Node<K, V>>>>, Box<Rc<RefCell<InternalNode<K, V>>>>>(Box::new(self.node.remove(node_address - 1)));
left.borrow_mut().marge(right);
if M - 1 <= self.keys.len() {
self.node.push(*left.clone());
} else {
if left.borrow().keys.len() >= M {
self.node.sort_by_key(|a| a.borrow().cmp_key());
let sp = left.borrow_mut().split();
self.node.push(sp.1);
self.keys.remove(node_address - 1);
self.keys.push(sp.0);
self.node.push(*left.clone());
} else {
self.marge(*left.clone());
}
}
self.key_conflict_resolver(*left);
self.node.sort_by_key(|a| a.borrow().cmp_key());
}
return Status::OkReorgLeft;
} else {
let right = self.node.remove(node_address + 1);
let mut tmp = InternalNode::new();
tmp.marge(right);
tmp.keys.push(tmp.node[0].borrow().cmp_key());
println!("{:?}", &tmp);
unsafe {
let left = std::mem::transmute::<Box<Rc<RefCell<dyn Node<K, V>>>>, Box<Rc<RefCell<InternalNode<K, V>>>>>(Box::new(self.node.remove(node_address)));
while let Some(elem) = left.borrow_mut().node.pop() {
tmp.node.push(elem);
}
tmp.node.sort_by_key(|a| a.borrow().cmp_key());
tmp.keys.sort();
println!("{:?}", &tmp);
}
let tmp = Rc::new(RefCell::new(tmp));
if M - 1 <= self.keys.len() {
self.node.push(tmp.clone());
} else {
self.marge(tmp.clone());
}
self.key_conflict_resolver(tmp);
self.node.sort_by_key(|a| a.borrow().cmp_key());
return Status::OkReorgRight;
}
} else {
let shared_node = if (self.node.len() - 1) == node_address { node_address - 1 } else { node_address + 1 };
let share_done = if (self.node.len() - 1) == node_address { self.node[node_address].borrow_mut().right_share(self.node[shared_node].clone()) } else { self.node[node_address].borrow_mut().left_share(self.node[shared_node].clone()) };
if !share_done {
if (self.node.len() - 1) == node_address {
// 対象が最右の場合、左隣りにmargeする
self.node[node_address - 1].borrow_mut().marge(self.node[node_address].clone());
self.node.remove(node_address);
} else {
// 右隣にmargeする
self.node[node_address].borrow_mut().marge(self.node[node_address + 1].clone());
self.node.remove(node_address + 1);
self.keys.remove(node_address);
}
} else {
self.node.sort_by_key(|a| a.borrow().cmp_key());
self.keys.remove(if (self.node.len() - 1) == node_address { node_address - 1 } else { node_address });
self.keys.push(self.node[if (self.node.len() - 1) == node_address { node_address } else { node_address + 1 }].borrow().cmp_key());
self.keys.sort();
}
};
self.node.sort_by_key(|a| a.borrow().cmp_key());
if (M / 2) - 1 >= self.node.len() || self.node.len() == 1 {
return Status::OkReorg;
}
return Status::OK;
}
Status::OK => delete_status,
Status::NotFound => delete_status
};
return result;
}
fn print(&self, buf: Vec<(K, V)>) -> Vec<(K, V)> {
//最小値をもつ節を辿る
self.node[0].borrow().print(buf)
}
fn find(&self, key: K) -> Option<(K, V)> {
let node_address = {
let mut result = self.keys.len();
for (i, comp_key) in self.keys.iter().enumerate() {
if &key < comp_key {
result = i;
break;
};
}
result
};
self.node[node_address].borrow().find(key)
}
fn take_items(&mut self) -> Rc<RefCell<dyn Node<K, V>>> {
let mut taken_keys = Vec::new();
let mut taken_node = Vec::new();
while let Some(key) = self.keys.pop() {
taken_keys.push(key);
}
while let Some(node) = self.node.pop() {
taken_node.push(node);
}
taken_keys.sort();
taken_node.sort_by_key(|a| a.clone().borrow().cmp_key());
Rc::new(RefCell::new(InternalNode { keys: taken_keys, node: taken_node }))
}
fn cmp_key(&self) -> K {
self.keys[0]
}
fn insert(&mut self, k: K, v: V) -> Option<(K, Rc<RefCell<dyn Node<K, V>>>)> {
self.node.sort_by_key(|a| a.borrow().cmp_key());
let mut node_address = None;
// insertする節を決定する
for (i, comp_key) in self.keys.iter().enumerate() {
if comp_key <= &k {
node_address = Some(self.keys.len());
continue;
} else {
node_address = Some(i);
break;
}
}
let split_info = self.node[node_address.unwrap_or(0)].borrow_mut().insert(k, v.clone());
if let Some((new_key, new_node)) = split_info {
self.node.push(new_node);
self.keys.push(new_key);
self.keys.sort();
}
if self.keys.len() >= M {
self.node.sort_by_key(|a| a.borrow().cmp_key());
return Some(self.split());
}
None
}
fn split(&mut self) -> (K, Rc<RefCell<dyn Node<K, V>>>) {
let mut new_key = Vec::new();
let mut new_node = Vec::new();
let return_key = self.keys.remove(M / 2);
self.keys.sort();
self.node.sort_by_key(|a| a.clone().borrow().cmp_key());
for _ in (M / 2)..self.keys.len() {
new_key.push(self.keys.pop().unwrap());
}
for _ in (M / 2) + 1..self.node.len() {
new_node.push(self.node.pop().unwrap())
}
new_key.sort();
new_node.sort_by_key(|a| a.clone().borrow().cmp_key());
let new_tree = Rc::new(RefCell::new(InternalNode { keys: new_key, node: new_node }));
(return_key, new_tree.clone())
}
fn reorg_key(&mut self, delete_key_locate: usize) -> Option<K> {
let result = self.keys.remove(delete_key_locate);
if let Some(key) = self.node[self.node.len() - 1].borrow_mut().reorg_key(0) {
if !self.keys.contains(&key) {
self.keys.push(key);
};
} else {
let key = self.node[self.node.len() - 2].borrow_mut().reorg_key(0).unwrap();
if !self.keys.contains(&key) {
self.keys.push(key);
};
}
self.keys.sort();
return Some(result);
}
}
impl<K, V> InternalNode<K, V>
where
K: PartialEq + PartialOrd + Ord + Copy + Clone + Debug + 'static,
V: PartialEq + Eq + Clone + Debug + 'static
{
fn new() -> Self {
InternalNode {
keys: Vec::new(),
node: Vec::new(),
}
}
fn key_conflict_resolver(&mut self, child_node: Rc<RefCell<dyn Node<K, V>>>) {
if child_node.borrow().is_internal_node() {
unsafe {
let internal = std::mem::transmute::<Box<Rc<RefCell<dyn Node<K, V>>>>, Box<Rc<RefCell<InternalNode<K, V>>>>>(Box::new(child_node));
for (i, key) in self.keys.iter().enumerate() {
if internal.borrow().keys.contains(key) {
self.keys.remove(i);
break;
}
}
}
}
}
}
impl<K, V> Node<K, V> for LeafNode<K, V>
where
K: PartialEq + PartialOrd + Ord + Copy + Clone + Debug + 'static,
V: PartialEq + Eq + Clone + Debug + 'static
{
fn is_internal_node(&self) -> bool {
TypeId::of::<InternalNode<K, V>>() == self.type_id()
}
fn is_leaf_node(&self) -> bool {
TypeId::of::<LeafNode<K, V>>() == self.type_id()
}
fn marge(&mut self, from: Rc<RefCell<dyn Node<K, V>>>) -> bool {
unsafe {
let from = std::mem::transmute::<Box<Rc<RefCell<dyn Node<K, V>>>>, Box<Rc<RefCell<LeafNode<K, V>>>>>(Box::new(from));
while let Some(elem) = from.borrow_mut().data.pop() {
self.data.push(elem);
}
self.data.sort_by_key(|a| a.0);
self.next_leaf = from.borrow_mut().next_leaf.take();
}
true
}
fn right_share(&mut self, from: Rc<RefCell<dyn Node<K, V>>>) -> bool {
unsafe {
let from = std::mem::transmute::<Box<Rc<RefCell<dyn Node<K, V>>>>, Box<Rc<RefCell<LeafNode<K, V>>>>>(Box::new(from));
let len = from.borrow().data.len().clone();
if 1 != (self.data.len() + len) && (M / 2) + 1 <= self.data.len() + len {
for _ in 0..(self.data.len() + len) / 2 - self.data.len() {
self.data.push(from.borrow_mut().data.pop().unwrap());
}
self.data.sort_by_key(|a| a.0);
true
} else { false }
}
}
fn left_share(&mut self, from: Rc<RefCell<dyn Node<K, V>>>) -> bool {
unsafe {
let from = std::mem::transmute::<Box<Rc<RefCell<dyn Node<K, V>>>>, Box<Rc<RefCell<LeafNode<K, V>>>>>(Box::new(from));
let len = from.borrow().data.len().clone();
if 1 != (self.data.len() + len) && (M / 2) + 1 <= self.data.len() + len {
for _ in 0..(self.data.len() + len) / 2 - self.data.len() {
self.data.push(from.borrow_mut().data.remove(0));
}
self.data.sort_by_key(|a| a.0);
true
} else { false }
}
}
fn update(&mut self, key: K, data: V) -> bool {
for i in 0..self.data.len() {
if key == self.data[i].0 {
self.data[i].1 = data;
return true;
};
};
false
}
fn delete(&mut self, key: K) -> Status {
for i in 0..self.data.len() {
if self.data[i].0 == key {
self.data.remove(i);
let threshold = (M / 2) - 1;
let status = if threshold >= self.data.len() || self.data.len() == 0 {
Status::OkReorg
} else {
Status::OK
};
return status;
}
}
Status::NotFound
}
fn print(&self, mut buf: Vec<(K, V)>) -> Vec<(K, V)> {
let s = self.data.as_slice().clone();
buf.append(&mut s.to_vec());
if let Some(nl) = &self.next_leaf {
nl.borrow().print(buf)
} else { buf }
}
fn find(&self, key: K) -> Option<(K, V)> {
for i in 0..self.data.len() {
if key == self.data[i].0 {
return Some(self.data[i].clone());
}
}
None
}
fn take_items(&mut self) -> Rc<RefCell<dyn Node<K, V>>> {
let mut taken_data = Vec::new();
while let Some(data) = self.data.pop() {
taken_data.push(data);
}
taken_data.sort_by_key(|a| a.0);
Rc::new(RefCell::new(LeafNode { data: taken_data, next_leaf: self.next_leaf.take() }))
}
fn cmp_key(&self) -> K {
self.data[0].0
}
fn insert(&mut self, k: K, v: V) -> Option<(K, Rc<RefCell<dyn Node<K, V>>>)> {
self.data.push((k.clone(), v));
self.data.sort_by(|prev, next| prev.0.cmp(&next.0));
if self.data.len() >= M {
return Some(self.split());
}
None
}
fn split(&mut self) -> (K, Rc<RefCell<dyn Node<K, V>>>) {
let mut new_data: Vec<(K, V)> = Vec::new();
let return_key = self.data[M / 2].0.clone();
for _ in (M / 2)..M {
let data = self.data.remove(self.data.len() - 1);
new_data.push(data);
}
self.data.sort_by_key(|k_v| k_v.0);
new_data.sort_by_key(|k_v| k_v.0);
let new_leaf = Rc::new(RefCell::new(LeafNode { data: new_data, next_leaf: self.next_leaf.take() }));
self.next_leaf = Some(new_leaf.clone());
(return_key, new_leaf)
}
fn reorg_key(&mut self, _: usize) -> Option<K> {
if let Some(k_v) = self.data.get(0) {
Some(k_v.0)
} else { None }
}
}
impl<K, V> LeafNode<K, V>
where
K: PartialEq + PartialOrd + Ord + Copy + Clone + Debug + 'static,
V: PartialEq + Eq + Clone + Debug + 'static
{
fn new() -> Self {
LeafNode {
data: Vec::new(),
next_leaf: None,
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment