Rust習作 B木

main.rs

use std::fmt;

const N_VALS : usize = 4;

type KeyT = i64;

enum BTreeNode<T>
{
    Nil,
    Node
    {
        kvs : Vec<(KeyT, T)>,
        children : Vec< Box< BTreeNode<T> > >
    }
}

use BTreeNode::{Nil,Node};

fn bnode_dump< T : fmt::Display >( depth : usize, nd : &BTreeNode<T> )
{
    let (kvs, children) = match nd {
        Nil => return,
        Node {kvs, children} => (kvs, children)
    };

    for i in 0..children.len() {
        bnode_dump(depth + 1, &children[i]);
        if i < kvs.len() {
            let whs = String::from("  ").repeat(depth);
            let (k, v) = &kvs[i];
            println!("{}{}:{}", whs, k, v);
        }
    }
}

fn bnode_singleton<T>(k : KeyT, x : T) -> BTreeNode<T>
{
    let mut kvs = Vec::with_capacity(N_VALS + 1);
    kvs.push( (k, x) );
    let mut children = Vec::with_capacity(N_VALS + 2);
    children.push( Box::new(Nil) );
    children.push( Box::new(Nil) );
    Node { kvs : kvs, children : children }
}

fn bnode_insert<T>( mut nd : &mut BTreeNode<T>, k : KeyT, x : T )
                    -> Option< ( (KeyT, T), Box< BTreeNode<T> > ) >
{
    let (mut kvs, mut children) = match nd {
        Nil => { *nd = bnode_singleton(k, x); return None; },
        Node{kvs, children} => (kvs, children)
    };

    let mut pos = kvs.len();
    for i in 0..kvs.len() {
        if kvs[i].0 > k {
            pos = i;
            break;
        }
    }
    let mut ch = &mut *children[pos];
    let (pos, kv2, nd2) = match &mut ch {
        Nil => {
            (pos, (k, x), Box::new(Nil))
        }
        Node {..} => {
            match bnode_insert( &mut ch, k, x ) {
                None => return None,
                Some((kv2, nd2)) => (pos, kv2, nd2)
            }
        }
    };

    kvs.insert(pos, kv2);
    children.insert(pos + 1, nd2);

    if kvs.len() == N_VALS + 1 {
        let mut kvsDrain = kvs.drain( std::ops::RangeFrom{ start : N_VALS / 2 } );
        let kv_mid = match kvsDrain.next() {
            Some(kv_mid) => kv_mid,
            None => { return None; } // imposible in flavor of kvs.len() == N_VALS + 1
        };
        let chDrain = children.drain( std::ops::RangeFrom{ start : N_VALS / 2 + 1 } );
        let tail = Node { kvs : kvsDrain.collect(), children : chDrain.collect() };
        Some( (kv_mid, Box::new(tail)) )
    }
    else {
        None
    }
}

fn bnode_find<'a, T>( nd : &'a BTreeNode<T> , k0 : KeyT ) -> Option<&'a T>
{
    let (kvs, children) = match nd {
        Nil => { return None; },
        Node {kvs, children} => (kvs, children)
    };

    let mut pos = kvs.len();

    for i in 0..kvs.len() {
        let (k, x) = &kvs[i];
        if k == &k0 {
            return Some(&x);
        }
        else if k > &k0 {
            pos = i;
            break;
        }
    }

    bnode_find(&children[pos], k0)
}

struct BTree<T>
{
    root : Box< BTreeNode<T> >
}

fn btree_new<T>() -> BTree<T>
{
    BTree::<T>{ root : Box::new(Nil) }
}

impl<T> BTree<T>
{
    fn insert( &mut self, k : KeyT, x : T )
    {
        let (kv_mid, tail) = match bnode_insert( &mut self.root, k, x ) {
            None => { return {}; },
            Some( (kv_mid, tail) ) => (kv_mid, tail)
        };
        let newroot = Node {
            kvs : Vec::with_capacity(N_VALS + 1),
            children : Vec::with_capacity(N_VALS + 2)
        };
        let oldroot = std::mem::replace(&mut *self.root, newroot);
        match &mut *self.root {
            Nil => {}, // imposible
            Node{kvs, children} => {
                kvs.push(kv_mid);
                children.push(Box::new(oldroot));
                children.push(tail);
            }
        }
    }

    fn find<'a>( &'a self, k : KeyT ) -> Option<&'a T>
    {
        bnode_find(&self.root, k)
    }

    fn dump( &self ) where T : fmt::Display
    {
        bnode_dump(0, &self.root);
    }
}

fn main() {
    let mut b = btree_new::<String>();
    for i in 1..20 {
        let k = i * 2 - 1;
        let s = format!("n{}", k);
        b.insert(k, s);
    }
    for i in 1..20 {
        let k = i * 2;
        let s = format!("n{}", k);
        b.insert(k, s);
    }
    b.dump();
    let mx = b.find(10);
    println!("find : {:?}", mx);
}