Skip to content

Instantly share code, notes, and snippets.

@totechite

totechite/b_plus_tree.rs

Created January 29, 2020 15:23
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save totechite/f4ea94b6095b31d4f62e71baa48135a1 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save totechite/f4ea94b6095b31d4f62e71baa48135a1 to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. totechite created this gist Jan 29, 2020.
    604 changes: 604 additions & 0 deletions b_plus_tree.rs
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,604 @@
    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,
    }
    }
    }