<commit_before>//! Types to build operators with general shapes.

use std::rc::Rc;
use std::cell::RefCell;
use std::default::Default;

use ::Data;

use progress::{ChangeBatch, Timestamp};
use progress::frontier::{Antichain, MutableAntichain};

use dataflow::{Stream, Scope};
use dataflow::channels::pushers::Tee;
use dataflow::channels::pushers::Counter as PushCounter;
use dataflow::channels::pushers::buffer::Buffer as PushBuffer;
use dataflow::channels::pact::ParallelizationContract;
use dataflow::channels::pullers::Counter as PullCounter;
use dataflow::operators::capability::Capability;
use dataflow::operators::capability::mint as mint_capability;

use dataflow::operators::generic::handles::{InputHandle, new_input_handle, OutputWrapper};

use logging::Logger;

use super::builder_raw::OperatorBuilder as OperatorBuilderRaw;

/// Builds operators with generic shape.
pub struct OperatorBuilder<G: Scope> {
    builder: OperatorBuilderRaw<G>,
    frontier: Vec<MutableAntichain<G::Timestamp>>,
    consumed: Vec<Rc<RefCell<ChangeBatch<G::Timestamp>>>>,
    internal: Rc<RefCell<Vec<Rc<RefCell<ChangeBatch<G::Timestamp>>>>>>,
    produced: Vec<Rc<RefCell<ChangeBatch<G::Timestamp>>>>,
    logging: Logger,
}

impl<G: Scope> OperatorBuilder<G> {

    /// Allocates a new generic operator builder from its containing scope.
    pub fn new(name: String, scope: G) -> Self {
        let logging = scope.logging();
        OperatorBuilder {
            builder: OperatorBuilderRaw::new(name, scope),
            frontier: Vec::new(),
            consumed: Vec::new(),
            internal: Rc::new(RefCell::new(Vec::new())),
            produced: Vec::new(),
            logging,
        }
    }

    /// Indicates whether the operator requires frontier information.
    pub fn set_notify(&mut self, notify: bool) {
        self.builder.set_notify(notify);
    }

    /// Adds a new input to a generic operator builder, returning the `Pull` implementor to use.
    pub fn new_input<D: Data, P>(&mut self, stream: &Stream<G, D>, pact: P) -> InputHandle<G::Timestamp, D, P::Puller>
    where
        P: ParallelizationContract<G::Timestamp, D> {

        let connection = vec![Antichain::from_elem(Default::default()); self.builder.shape().outputs()];
        self.new_input_connection(stream, pact, connection)
    }

    /// Adds a new input with connection information to a generic operator builder, returning the `Pull` implementor to use.
    pub fn new_input_connection<D: Data, P>(&mut self, stream: &Stream<G, D>, pact: P, connection: Vec<Antichain<<G::Timestamp as Timestamp>::Summary>>) -> InputHandle<G::Timestamp, D, P::Puller>
        where
            P: ParallelizationContract<G::Timestamp, D> {

        let puller = self.builder.new_input_connection(stream, pact, connection);

        let input = PullCounter::new(puller);
        self.frontier.push(MutableAntichain::new());
        self.consumed.push(input.consumed().clone());

        new_input_handle(input, self.internal.clone(), self.logging.clone())
    }

    /// Adds a new output to a generic operator builder, returning the `Pull` implementor to use.
    pub fn new_output<D: Data>(&mut self) -> (OutputWrapper<G::Timestamp, D, Tee<G::Timestamp, D>>, Stream<G, D>) {
        let connection = vec![Antichain::from_elem(Default::default()); self.builder.shape().inputs()];
        self.new_output_connection(connection)
    }

    /// Adds a new output with connection information to a generic operator builder, returning the `Pull` implementor to use.
    pub fn new_output_connection<D: Data>(&mut self, connection: Vec<Antichain<<G::Timestamp as Timestamp>::Summary>>) -> (OutputWrapper<G::Timestamp, D, Tee<G::Timestamp, D>>, Stream<G, D>) {

        let (tee, stream) = self.builder.new_output_connection(connection);

        let internal = Rc::new(RefCell::new(ChangeBatch::new()));
        self.internal.borrow_mut().push(internal.clone());

        let mut buffer = PushBuffer::new(PushCounter::new(tee));
        self.produced.push(buffer.inner().produced().clone());

        (OutputWrapper::new(buffer, internal), stream)
    }

    /// Creates an operator implementation from supplied logic constructor.
    pub fn build<B, L>(self, constructor: B)
    where
        B: FnOnce(Vec<Capability<G::Timestamp>>) -> L,
        L: FnMut(&[MutableAntichain<G::Timestamp>])+'static
    {
        // create capabilities, discard references to their creation.
        let mut capabilities = Vec::new();
        for output_index in 0  .. self.internal.borrow().len() {
            let borrow = &self.internal.borrow()[output_index];
            capabilities.push(mint_capability(Default::default(), borrow.clone()));
            // Discard evidence of creation, as we are assumed to start with one.
            borrow.borrow_mut().clear();
        }

        let mut logic = constructor(capabilities);

        let self_frontier1 = Rc::new(RefCell::new(self.frontier));
        let self_frontier2 = self_frontier1.clone();
        let self_consumed = self.consumed;
        let self_internal = self.internal;
        let self_produced = self.produced;

        let pep = move |changes: &mut [ChangeBatch<G::Timestamp>]| {
            let mut borrow = self_frontier1.borrow_mut();
            for index in 0 .. changes.len() {
                borrow[index].update_iter(changes[index].drain());
            }
        };

        let pip = move |consumed: &mut [ChangeBatch<G::Timestamp>],
                        internal: &mut [ChangeBatch<G::Timestamp>],
                        produced: &mut [ChangeBatch<G::Timestamp>]| {

            // invoke supplied logic
            let borrow = self_frontier2.borrow();
            logic(&*borrow);

            // move batches of consumed changes.
            for index in 0 .. consumed.len() {
                self_consumed[index].borrow_mut().drain_into(&mut consumed[index]);
            }

            // move batches of internal changes.
            let self_internal_borrow = self_internal.borrow_mut();
            for index in 0 .. self_internal_borrow.len() {
                let mut borrow = self_internal_borrow[index].borrow_mut();
                internal[index].extend(borrow.drain());
            }

            // move batches of produced changes.
            for index in 0 .. produced.len() {
                self_produced[index].borrow_mut().drain_into(&mut produced[index]);
            }

            false
        };

        self.builder.build(pep, pip);
    }

    /// Get the identifier assigned to the operator being constructed
    pub fn index(&self) -> usize {
        self.builder.index()
    }
}
<commit_msg>add passing and failing tests<commit_after>//! Types to build operators with general shapes.

use std::rc::Rc;
use std::cell::RefCell;
use std::default::Default;

use ::Data;

use progress::{ChangeBatch, Timestamp};
use progress::frontier::{Antichain, MutableAntichain};

use dataflow::{Stream, Scope};
use dataflow::channels::pushers::Tee;
use dataflow::channels::pushers::Counter as PushCounter;
use dataflow::channels::pushers::buffer::Buffer as PushBuffer;
use dataflow::channels::pact::ParallelizationContract;
use dataflow::channels::pullers::Counter as PullCounter;
use dataflow::operators::capability::Capability;
use dataflow::operators::capability::mint as mint_capability;

use dataflow::operators::generic::handles::{InputHandle, new_input_handle, OutputWrapper};

use logging::Logger;

use super::builder_raw::OperatorBuilder as OperatorBuilderRaw;

/// Builds operators with generic shape.
pub struct OperatorBuilder<G: Scope> {
    builder: OperatorBuilderRaw<G>,
    frontier: Vec<MutableAntichain<G::Timestamp>>,
    consumed: Vec<Rc<RefCell<ChangeBatch<G::Timestamp>>>>,
    internal: Rc<RefCell<Vec<Rc<RefCell<ChangeBatch<G::Timestamp>>>>>>,
    produced: Vec<Rc<RefCell<ChangeBatch<G::Timestamp>>>>,
    logging: Logger,
}

impl<G: Scope> OperatorBuilder<G> {

    /// Allocates a new generic operator builder from its containing scope.
    pub fn new(name: String, scope: G) -> Self {
        let logging = scope.logging();
        OperatorBuilder {
            builder: OperatorBuilderRaw::new(name, scope),
            frontier: Vec::new(),
            consumed: Vec::new(),
            internal: Rc::new(RefCell::new(Vec::new())),
            produced: Vec::new(),
            logging,
        }
    }

    /// Indicates whether the operator requires frontier information.
    pub fn set_notify(&mut self, notify: bool) {
        self.builder.set_notify(notify);
    }

    /// Adds a new input to a generic operator builder, returning the `Pull` implementor to use.
    pub fn new_input<D: Data, P>(&mut self, stream: &Stream<G, D>, pact: P) -> InputHandle<G::Timestamp, D, P::Puller>
    where
        P: ParallelizationContract<G::Timestamp, D> {

        let connection = vec![Antichain::from_elem(Default::default()); self.builder.shape().outputs()];
        self.new_input_connection(stream, pact, connection)
    }

    /// Adds a new input with connection information to a generic operator builder, returning the `Pull` implementor to use.
    pub fn new_input_connection<D: Data, P>(&mut self, stream: &Stream<G, D>, pact: P, connection: Vec<Antichain<<G::Timestamp as Timestamp>::Summary>>) -> InputHandle<G::Timestamp, D, P::Puller>
        where
            P: ParallelizationContract<G::Timestamp, D> {

        let puller = self.builder.new_input_connection(stream, pact, connection);

        let input = PullCounter::new(puller);
        self.frontier.push(MutableAntichain::new());
        self.consumed.push(input.consumed().clone());

        new_input_handle(input, self.internal.clone(), self.logging.clone())
    }

    /// Adds a new output to a generic operator builder, returning the `Pull` implementor to use.
    pub fn new_output<D: Data>(&mut self) -> (OutputWrapper<G::Timestamp, D, Tee<G::Timestamp, D>>, Stream<G, D>) {
        let connection = vec![Antichain::from_elem(Default::default()); self.builder.shape().inputs()];
        self.new_output_connection(connection)
    }

    /// Adds a new output with connection information to a generic operator builder, returning the `Pull` implementor to use.
    pub fn new_output_connection<D: Data>(&mut self, connection: Vec<Antichain<<G::Timestamp as Timestamp>::Summary>>) -> (OutputWrapper<G::Timestamp, D, Tee<G::Timestamp, D>>, Stream<G, D>) {

        let (tee, stream) = self.builder.new_output_connection(connection);

        let internal = Rc::new(RefCell::new(ChangeBatch::new()));
        self.internal.borrow_mut().push(internal.clone());

        let mut buffer = PushBuffer::new(PushCounter::new(tee));
        self.produced.push(buffer.inner().produced().clone());

        (OutputWrapper::new(buffer, internal), stream)
    }

    /// Creates an operator implementation from supplied logic constructor.
    pub fn build<B, L>(self, constructor: B)
    where
        B: FnOnce(Vec<Capability<G::Timestamp>>) -> L,
        L: FnMut(&[MutableAntichain<G::Timestamp>])+'static
    {
        // create capabilities, discard references to their creation.
        let mut capabilities = Vec::new();
        for output_index in 0  .. self.internal.borrow().len() {
            let borrow = &self.internal.borrow()[output_index];
            capabilities.push(mint_capability(Default::default(), borrow.clone()));
            // Discard evidence of creation, as we are assumed to start with one.
            borrow.borrow_mut().clear();
        }

        let mut logic = constructor(capabilities);

        let self_frontier1 = Rc::new(RefCell::new(self.frontier));
        let self_frontier2 = self_frontier1.clone();
        let self_consumed = self.consumed;
        let self_internal = self.internal;
        let self_produced = self.produced;

        let pep = move |changes: &mut [ChangeBatch<G::Timestamp>]| {
            let mut borrow = self_frontier1.borrow_mut();
            for index in 0 .. changes.len() {
                borrow[index].update_iter(changes[index].drain());
            }
        };

        let pip = move |consumed: &mut [ChangeBatch<G::Timestamp>],
                        internal: &mut [ChangeBatch<G::Timestamp>],
                        produced: &mut [ChangeBatch<G::Timestamp>]| {

            // invoke supplied logic
            let borrow = self_frontier2.borrow();
            logic(&*borrow);

            // move batches of consumed changes.
            for index in 0 .. consumed.len() {
                self_consumed[index].borrow_mut().drain_into(&mut consumed[index]);
            }

            // move batches of internal changes.
            let self_internal_borrow = self_internal.borrow_mut();
            for index in 0 .. self_internal_borrow.len() {
                let mut borrow = self_internal_borrow[index].borrow_mut();
                internal[index].extend(borrow.drain());
            }

            // move batches of produced changes.
            for index in 0 .. produced.len() {
                self_produced[index].borrow_mut().drain_into(&mut produced[index]);
            }

            false
        };

        self.builder.build(pep, pip);
    }

    /// Get the identifier assigned to the operator being constructed
    pub fn index(&self) -> usize {
        self.builder.index()
    }
}

#[cfg(test)]
mod tests {

    #[test]
    #[should_panic]
    fn incorrect_capabilities() {

        // This tests that if we attempt to use a capability associated with the
        // wrong output, there is a run-time assertion.

        use ::dataflow::operators::generic::builder_rc::OperatorBuilder;

        ::example(|scope| {

            let mut builder = OperatorBuilder::new("Failure".to_owned(), scope.clone());

            // let mut input = builder.new_input(stream, Pipeline);
            let (mut output1, _stream1) = builder.new_output::<()>();
            let (mut output2, _stream2) = builder.new_output::<()>();

            builder.build(move |capabilities| {
                move |_frontiers| {

                    let mut output_handle1 = output1.activate();
                    let mut output_handle2 = output2.activate();

                    // NOTE: Using incorrect capabilities here.
                    output_handle2.session(&capabilities[0]);
                    output_handle1.session(&capabilities[1]);
                }
            });
        })
    }

    #[test]
    fn correct_capabilities() {

        // This tests that if we attempt to use capabilities with the correct outputs
        // there is no runtime assertion

        use ::dataflow::operators::generic::builder_rc::OperatorBuilder;

        ::example(|scope| {

            let mut builder = OperatorBuilder::new("Failure".to_owned(), scope.clone());

            // let mut input = builder.new_input(stream, Pipeline);
            let (mut output1, _stream1) = builder.new_output::<()>();
            let (mut output2, _stream2) = builder.new_output::<()>();

            builder.build(move |mut capabilities| {
                move |_frontiers| {

                    let mut output_handle1 = output1.activate();
                    let mut output_handle2 = output2.activate();

                    // Avoid second call.
                    if !capabilities.is_empty() {

                        // NOTE: Using correct capabilities here.
                        output_handle1.session(&capabilities[0]);
                        output_handle2.session(&capabilities[1]);

                        capabilities.clear();
                    }
                }
            });

            "Hello".to_owned()
        });
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::hash::Hash;

///////////////////////////////////////////////////////////////////////////////

trait Deserialize<S, E> {
    fn deserialize(state: &mut S) -> Result<Self, E>;
}

///////////////////////////////////////////////////////////////////////////////

trait VisitorState<E> {
    fn visit_int(&mut self) -> Result<int, E>;

    fn visit_string(&mut self) -> Result<String, E>;

    fn visit_seq<
        C: Deserialize<Self, E>,
        V: Visitor<C, Self, E>
    >(&mut self) -> Result<C, E>;

    fn visit_map<
        C: Deserialize<Self, E>,
        V: Visitor<C, Self, E>
    >(&mut self) -> Result<C, E>;
}

trait Visitor<C, S, E> {
    fn new(len: uint) -> Self;

    fn visit(&mut self, state: &mut S) -> Result<(), E>;

    fn unwrap(self) -> Result<C, E>;
}

///////////////////////////////////////////////////////////////////////////////

impl<S: VisitorState<E>, E> Deserialize<S, E> for int {
    fn deserialize(state: &mut S) -> Result<int, E> {
        state.visit_int()
    }
}

impl<S: VisitorState<E>, E> Deserialize<S, E> for String {
    fn deserialize(state: &mut S) -> Result<String, E> {
        state.visit_string()
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    T: Deserialize<S, E>,
    S: VisitorState<E>,
    E
> Deserialize<S, E> for Vec<T> {
    fn deserialize(state: &mut S) -> Result<Vec<T>, E> {
        struct Visitor<T> {
            value: Vec<T>,
        }

        impl<
            T: Deserialize<S, E>,
            S: VisitorState<E>,
            E
        > ::Visitor<Vec<T>, S, E> for Visitor<T> {
            fn new(len: uint) -> Visitor<T> {
                Visitor {
                    value: Vec::with_capacity(len),
                }
            }

            fn visit(&mut self, state: &mut S) -> Result<(), E> {
                let value = try!(Deserialize::deserialize(state));
                self.value.push(value);
                Ok(())
            }

            fn unwrap(self) -> Result<Vec<T>, E> {
                Ok(self.value)
            }
        }

        state.visit_seq::<Vec<T>, Visitor<T>>()
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    T0: Deserialize<S, E>,
    T1: Deserialize<S, E>,
    S: VisitorState<E>,
    E
> Deserialize<S, E> for (T0, T1) {
    fn deserialize(state: &mut S) -> Result<(T0, T1), E> {
        struct Visitor<T0, T1> {
            state: uint,
            t0: Option<T0>,
            t1: Option<T1>,
        }

        impl<
            T0: Deserialize<S, E>,
            T1: Deserialize<S, E>,
            S,
            E
        > ::Visitor<(T0, T1), S, E> for Visitor<T0, T1> {
            fn new(_: uint) -> Visitor<T0, T1> {
                Visitor {
                    state: 0,
                    t0: None,
                    t1: None,
                }
            }

            fn visit(&mut self, state: &mut S) -> Result<(), E> {
                match self.state {
                    0 => {
                        self.state += 1;
                        self.t0 = Some(try!(Deserialize::deserialize(state)));
                    }
                    1 => {
                        self.state += 1;
                        self.t1 = Some(try!(Deserialize::deserialize(state)));
                    }
                    _ => fail!()
                }

                Ok(())
            }

            fn unwrap(self) -> Result<(T0, T1), E> {
                let t0 = match self.t0 {
                    Some(t0) => t0,
                    None => { fail!(); }
                };

                let t1 = match self.t1 {
                    Some(t1) => t1,
                    None => { fail!(); }
                };

                Ok((t0, t1))
            }
        }

        state.visit_seq::<(T0, T1), Visitor<T0, T1>>()
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    K: Deserialize<S, E> + Eq + Hash,
    V: Deserialize<S, E>,
    S: VisitorState<E>,
    E
> Deserialize<S, E> for HashMap<K, V> {
    fn deserialize(state: &mut S) -> Result<HashMap<K, V>, E> {
        struct Visitor<K, V> {
            value: HashMap<K, V>,
        }

        impl<
            K: Deserialize<S, E> + Eq + Hash,
            V: Deserialize<S, E>,
            S: VisitorState<E>,
            E
        > ::Visitor<HashMap<K, V>, S, E> for Visitor<K, V> {
            fn new(len: uint) -> Visitor<K, V> {
                Visitor {
                    value: HashMap::with_capacity(len),
                }
            }

            fn visit(&mut self, state: &mut S) -> Result<(), E> {
                let key = try!(Deserialize::deserialize(state));
                let value = try!(Deserialize::deserialize(state));
                self.value.insert(key, value);
                Ok(())
            }

            fn unwrap(self) -> Result<HashMap<K, V>, E> {
                Ok(self.value)
            }
        }

        state.visit_map::<HashMap<K, V>, Visitor<K, V>>()
    }
}

///////////////////////////////////////////////////////////////////////////////

trait Deserializer<S, E> {
    fn deserialize<T: Deserialize<S, E>>(&mut self) -> Result<T, E>;
}

///////////////////////////////////////////////////////////////////////////////

enum Token {
    Int(int),
    String(String),
    SeqStart(uint),
    MapStart(uint),
    End,
}

struct MyDeserializerState<Iter> {
    tokens: Iter,
    peeked: Option<Token>,
}

impl<Iter: Iterator<Token>> MyDeserializerState<Iter> {
    fn new(tokens: Iter) -> MyDeserializerState<Iter> {
        MyDeserializerState {
            tokens: tokens,
            peeked: None,
        }
    }

    fn next(&mut self) -> Option<Token> {
        match self.peeked.take() {
            Some(token) => { return Some(token); }
            None => { }
        }

        self.tokens.next()
    }
}

impl<Iter: Iterator<Token>> VisitorState<()> for MyDeserializerState<Iter> {
    fn visit_int(&mut self) -> Result<int, ()> {
        match self.next() {
            Some(Int(v)) => Ok(v),
            _ => Err(())
        }
    }

    fn visit_string(&mut self) -> Result<String, ()> {
        match self.next() {
            Some(String(v)) => Ok(v),
            _ => Err(())
        }
    }

    fn visit_seq<
        T: Deserialize<MyDeserializerState<Iter>, ()>,
        V: Visitor<T, MyDeserializerState<Iter>, ()>
    >(&mut self) -> Result<T, ()> {
        let len = match self.next() {
            Some(SeqStart(len)) => len,
            _ => { return Err(()); }
        };

        let mut visitor: V = Visitor::new(len);

        loop {
            match self.next() {
                Some(End) => { break; }
                Some(token) => {
                    self.peeked = Some(token);
                    try!(visitor.visit(self));
                }
                _ => { return Err(()); }
            }
        }

        visitor.unwrap()
    }

    fn visit_map<
        T: Deserialize<MyDeserializerState<Iter>, ()>,
        V: Visitor<T, MyDeserializerState<Iter>, ()>
    >(&mut self) -> Result<T, ()> {
        let len = match self.next() {
            Some(MapStart(len)) => len,
            _ => { return Err(()); }
        };

        let mut visitor: V = Visitor::new(len);

        loop {
            match self.next() {
                Some(End) => { break; }
                Some(token) => {
                    self.peeked = Some(token);
                    try!(visitor.visit(self));
                }
                _ => { return Err(()); }
            }
        }

        visitor.unwrap()
    }
}

///////////////////////////////////////////////////////////////////////////////

fn main() {
    let tokens = vec!(
        SeqStart(2),
        Int(1),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<Vec<int>, ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);

    let tokens = vec!(
        SeqStart(2),
        Int(1),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<(int, int), ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);

    let tokens = vec!(
        MapStart(2),
        String("a".to_string()),
        Int(1),
        String("b".to_string()),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<HashMap<String, int>, ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);
}
<commit_msg>Add visit_*_elt to de2<commit_after>use std::collections::HashMap;
use std::hash::Hash;

///////////////////////////////////////////////////////////////////////////////

trait Deserialize<S, E> {
    fn deserialize(state: &mut S) -> Result<Self, E>;
}

///////////////////////////////////////////////////////////////////////////////

trait VisitorState<E> {
    fn visit_int(&mut self) -> Result<int, E>;

    fn visit_string(&mut self) -> Result<String, E>;

    fn visit_seq<
        T: Deserialize<Self, E>,
        V: Visitor<T, Self, E>
    >(&mut self) -> Result<T, E>;

    fn visit_seq_elt<
        T: Deserialize<Self, E>,
    >(&mut self) -> Result<T, E>;

    fn visit_map<
        T: Deserialize<Self, E>,
        V: Visitor<T, Self, E>
    >(&mut self) -> Result<T, E>;

    fn visit_map_elt<
        K: Deserialize<Self, E>,
        V: Deserialize<Self, E>
    >(&mut self) -> Result<(K, V), E>;
}

trait Visitor<C, S, E> {
    fn new(len: uint) -> Self;

    fn visit(&mut self, state: &mut S) -> Result<(), E>;

    fn unwrap(self) -> Result<C, E>;
}

///////////////////////////////////////////////////////////////////////////////

impl<S: VisitorState<E>, E> Deserialize<S, E> for int {
    fn deserialize(state: &mut S) -> Result<int, E> {
        state.visit_int()
    }
}

impl<S: VisitorState<E>, E> Deserialize<S, E> for String {
    fn deserialize(state: &mut S) -> Result<String, E> {
        state.visit_string()
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    T: Deserialize<S, E>,
    S: VisitorState<E>,
    E
> Deserialize<S, E> for Vec<T> {
    fn deserialize(state: &mut S) -> Result<Vec<T>, E> {
        struct Visitor<T> {
            value: Vec<T>,
        }

        impl<
            T: Deserialize<S, E>,
            S: VisitorState<E>,
            E
        > ::Visitor<Vec<T>, S, E> for Visitor<T> {
            fn new(len: uint) -> Visitor<T> {
                Visitor {
                    value: Vec::with_capacity(len),
                }
            }

            fn visit(&mut self, state: &mut S) -> Result<(), E> {
                let value = try!(state.visit_seq_elt());
                self.value.push(value);
                Ok(())
            }

            fn unwrap(self) -> Result<Vec<T>, E> {
                Ok(self.value)
            }
        }

        state.visit_seq::<Vec<T>, Visitor<T>>()
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    T0: Deserialize<S, E>,
    T1: Deserialize<S, E>,
    S: VisitorState<E>,
    E
> Deserialize<S, E> for (T0, T1) {
    fn deserialize(state: &mut S) -> Result<(T0, T1), E> {
        struct Visitor<T0, T1> {
            state: uint,
            t0: Option<T0>,
            t1: Option<T1>,
        }

        impl<
            T0: Deserialize<S, E>,
            T1: Deserialize<S, E>,
            S: VisitorState<E>,
            E
        > ::Visitor<(T0, T1), S, E> for Visitor<T0, T1> {
            fn new(_: uint) -> Visitor<T0, T1> {
                Visitor {
                    state: 0,
                    t0: None,
                    t1: None,
                }
            }

            fn visit(&mut self, state: &mut S) -> Result<(), E> {
                match self.state {
                    0 => {
                        self.state += 1;
                        self.t0 = Some(try!(state.visit_seq_elt()));
                    }
                    1 => {
                        self.state += 1;
                        self.t1 = Some(try!(state.visit_seq_elt()));
                    }
                    _ => fail!()
                }

                Ok(())
            }

            fn unwrap(self) -> Result<(T0, T1), E> {
                let t0 = match self.t0 {
                    Some(t0) => t0,
                    None => { fail!(); }
                };

                let t1 = match self.t1 {
                    Some(t1) => t1,
                    None => { fail!(); }
                };

                Ok((t0, t1))
            }
        }

        state.visit_seq::<(T0, T1), Visitor<T0, T1>>()
    }
}

///////////////////////////////////////////////////////////////////////////////

impl<
    K: Deserialize<S, E> + Eq + Hash,
    V: Deserialize<S, E>,
    S: VisitorState<E>,
    E
> Deserialize<S, E> for HashMap<K, V> {
    fn deserialize(state: &mut S) -> Result<HashMap<K, V>, E> {
        struct Visitor<K, V> {
            value: HashMap<K, V>,
        }

        impl<
            K: Deserialize<S, E> + Eq + Hash,
            V: Deserialize<S, E>,
            S: VisitorState<E>,
            E
        > ::Visitor<HashMap<K, V>, S, E> for Visitor<K, V> {
            fn new(len: uint) -> Visitor<K, V> {
                Visitor {
                    value: HashMap::with_capacity(len),
                }
            }

            fn visit(&mut self, state: &mut S) -> Result<(), E> {
                let (key, value) = try!(state.visit_map_elt());
                self.value.insert(key, value);
                Ok(())
            }

            fn unwrap(self) -> Result<HashMap<K, V>, E> {
                Ok(self.value)
            }
        }

        state.visit_map::<HashMap<K, V>, Visitor<K, V>>()
    }
}

///////////////////////////////////////////////////////////////////////////////

trait Deserializer<S, E> {
    fn deserialize<T: Deserialize<S, E>>(&mut self) -> Result<T, E>;
}

///////////////////////////////////////////////////////////////////////////////

enum Token {
    Int(int),
    String(String),
    SeqStart(uint),
    MapStart(uint),
    End,
}

struct MyDeserializerState<Iter> {
    tokens: Iter,
    peeked: Option<Token>,
}

impl<Iter: Iterator<Token>> MyDeserializerState<Iter> {
    fn new(tokens: Iter) -> MyDeserializerState<Iter> {
        MyDeserializerState {
            tokens: tokens,
            peeked: None,
        }
    }

    fn next(&mut self) -> Option<Token> {
        match self.peeked.take() {
            Some(token) => { return Some(token); }
            None => { }
        }

        self.tokens.next()
    }
}

impl<Iter: Iterator<Token>> VisitorState<()> for MyDeserializerState<Iter> {
    fn visit_int(&mut self) -> Result<int, ()> {
        match self.next() {
            Some(Int(v)) => Ok(v),
            _ => Err(())
        }
    }

    fn visit_string(&mut self) -> Result<String, ()> {
        match self.next() {
            Some(String(v)) => Ok(v),
            _ => Err(())
        }
    }

    fn visit_seq<
        T: Deserialize<MyDeserializerState<Iter>, ()>,
        V: Visitor<T, MyDeserializerState<Iter>, ()>
    >(&mut self) -> Result<T, ()> {
        let len = match self.next() {
            Some(SeqStart(len)) => len,
            _ => { return Err(()); }
        };

        let mut visitor: V = Visitor::new(len);

        loop {
            match self.next() {
                Some(End) => { break; }
                Some(token) => {
                    self.peeked = Some(token);
                    try!(visitor.visit(self));
                }
                _ => { return Err(()); }
            }
        }

        visitor.unwrap()
    }

    fn visit_seq_elt<
        T: Deserialize<MyDeserializerState<Iter>, ()>,
    >(&mut self) -> Result<T, ()> {
        Deserialize::deserialize(self)
    }

    fn visit_map<
        T: Deserialize<MyDeserializerState<Iter>, ()>,
        V: Visitor<T, MyDeserializerState<Iter>, ()>
    >(&mut self) -> Result<T, ()> {
        let len = match self.next() {
            Some(MapStart(len)) => len,
            _ => { return Err(()); }
        };

        let mut visitor: V = Visitor::new(len);

        loop {
            match self.next() {
                Some(End) => { break; }
                Some(token) => {
                    self.peeked = Some(token);
                    try!(visitor.visit(self));
                }
                _ => { return Err(()); }
            }
        }

        visitor.unwrap()
    }

    fn visit_map_elt<
        K: Deserialize<MyDeserializerState<Iter>, ()>,
        V: Deserialize<MyDeserializerState<Iter>, ()>
    >(&mut self) -> Result<(K, V), ()> {
        let k = try!(Deserialize::deserialize(self));
        let v = try!(Deserialize::deserialize(self));
        Ok((k, v))
    }
}

///////////////////////////////////////////////////////////////////////////////

fn main() {
    let tokens = vec!(
        SeqStart(2),
        Int(1),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<Vec<int>, ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);

    let tokens = vec!(
        SeqStart(2),
        Int(1),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<(int, int), ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);

    let tokens = vec!(
        MapStart(2),
        String("a".to_string()),
        Int(1),
        String("b".to_string()),
        Int(2),
        End
    );
    let mut state = MyDeserializerState::new(tokens.move_iter());

    let v: Result<HashMap<String, int>, ()> = Deserialize::deserialize(&mut state);
    println!("{}", v);
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate time;
extern crate timely;
extern crate differential_dataflow;

use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::progress::timestamp::RootTimestamp;

use rand::{Rng, SeedableRng, StdRng};

use differential_dataflow::operators::*;
use differential_dataflow::collection::LeastUpperBound;

type Node = u32;
type Edge = (Node, Node);

fn main() {

    let nodes: u32 = std::env::args().nth(1).unwrap().parse().unwrap();
    let edges: u32 = std::env::args().nth(2).unwrap().parse().unwrap();
    let batch: u32 = std::env::args().nth(3).unwrap().parse().unwrap();

    // define a new computational scope, in which to run BFS
    timely::execute_from_args(std::env::args().skip(4), move |computation| {
        let start = time::precise_time_s();

        // define BFS dataflow; return handles to roots and edges inputs
        let (mut graph, probe) = computation.scoped(|scope| {

            let roots = vec![(0,1)].into_iter().to_stream(scope);
            let (edge_input, graph) = scope.new_input();
            let probe = bfs(&graph, &roots).probe().0;

            (edge_input, probe)
        });

        let seed: &[_] = &[1, 2, 3, 4];
        let mut rng1: StdRng = SeedableRng::from_seed(seed);    // rng for edge additions
        let mut rng2: StdRng = SeedableRng::from_seed(seed);    // rng for edge deletions

        println!("performing BFS on {} nodes, {} edges:", nodes, edges);

        if computation.index() == 0 {
            // trickle edges in to dataflow
            for _ in 0..(edges/1000) {
                for _ in 0..1000 {
                    graph.send(((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)), 1));
                }
                computation.step();
            }
        }

        println!("loaded; elapsed: {}s", time::precise_time_s() - start);

        graph.advance_to(1);
        while probe.le(&RootTimestamp::new(0)) { computation.step(); }

        println!("stable; elapsed: {}s", time::precise_time_s() - start);

        if batch > 0 {
            let mut changes = Vec::new();
            for wave in 0.. {
                for _ in 0..batch {
                    changes.push(((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)), 1));
                    changes.push(((rng2.gen_range(0, nodes), rng2.gen_range(0, nodes)),-1));
                }

                let start = time::precise_time_s();
                let round = *graph.epoch();
                if computation.index() == 0 {
                    while let Some(change) = changes.pop() {
                        graph.send(change);
                    }
                }
                graph.advance_to(round + 1);

                while probe.le(&RootTimestamp::new(round)) { computation.step(); }

                if computation.index() == 0 {
                    println!("wave {}: avg {}", wave, (time::precise_time_s() - start) / (batch as f64));
                }
            }
        }
    });
}

// returns pairs (n, s) indicating node n can be reached from a root in s steps.
fn bfs<G: Scope>(edges: &Stream<G, (Edge, i32)>, roots: &Stream<G, (Node, i32)>) -> Stream<G, ((Node, u32), i32)>
where G::Timestamp: LeastUpperBound {

    // initialize roots as reaching themselves at distance 0
    let nodes = roots.map(|(x,w)| ((x, 0), w));
    // let edges = edges.map_in_place(|x| x.0 = ((x.0).1, (x.0).0))
    //                  .concat(&edges);

    // repeatedly update minimal distances each node can be reached from each root
    nodes.iterate(|inner| {

        let edges = inner.scope().enter(&edges);
        let nodes = inner.scope().enter(&nodes);

        inner.join_map(&edges, |_k,l,d| (*d, l+1))
             .concat(&nodes)
             .group(|_, s, t| t.push((*s.peek().unwrap().0, 1)))
     })
}
<commit_msg>made operators unsigned<commit_after>extern crate rand;
extern crate time;
extern crate timely;
extern crate differential_dataflow;

use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::progress::timestamp::RootTimestamp;

use rand::{Rng, SeedableRng, StdRng};

use differential_dataflow::operators::*;
use differential_dataflow::operators::join::JoinUnsigned;
use differential_dataflow::operators::group::GroupUnsigned;
use differential_dataflow::collection::LeastUpperBound;

type Node = u32;
type Edge = (Node, Node);

fn main() {

    let nodes: u32 = std::env::args().nth(1).unwrap().parse().unwrap();
    let edges: u32 = std::env::args().nth(2).unwrap().parse().unwrap();
    let batch: u32 = std::env::args().nth(3).unwrap().parse().unwrap();

    // define a new computational scope, in which to run BFS
    timely::execute_from_args(std::env::args().skip(4), move |computation| {
        let start = time::precise_time_s();

        // define BFS dataflow; return handles to roots and edges inputs
        let (mut graph, probe) = computation.scoped(|scope| {

            let roots = vec![(0,1)].into_iter().to_stream(scope);
            let (edge_input, graph) = scope.new_input();
            let probe = bfs(&graph, &roots).probe().0;

            (edge_input, probe)
        });

        let seed: &[_] = &[1, 2, 3, 4];
        let mut rng1: StdRng = SeedableRng::from_seed(seed);    // rng for edge additions
        let mut rng2: StdRng = SeedableRng::from_seed(seed);    // rng for edge deletions

        println!("performing BFS on {} nodes, {} edges:", nodes, edges);

        if computation.index() == 0 {
            // trickle edges in to dataflow
            for _ in 0..(edges/1000) {
                for _ in 0..1000 {
                    graph.send(((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)), 1));
                }
                computation.step();
            }
        }

        println!("loaded; elapsed: {}s", time::precise_time_s() - start);

        graph.advance_to(1);
        while probe.le(&RootTimestamp::new(0)) { computation.step(); }

        println!("stable; elapsed: {}s", time::precise_time_s() - start);

        if batch > 0 {
            let mut changes = Vec::new();
            for wave in 0.. {
                for _ in 0..batch {
                    changes.push(((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)), 1));
                    changes.push(((rng2.gen_range(0, nodes), rng2.gen_range(0, nodes)),-1));
                }

                let start = time::precise_time_s();
                let round = *graph.epoch();
                if computation.index() == 0 {
                    while let Some(change) = changes.pop() {
                        graph.send(change);
                    }
                }
                graph.advance_to(round + 1);

                while probe.le(&RootTimestamp::new(round)) { computation.step(); }

                if computation.index() == 0 {
                    println!("wave {}: avg {}", wave, (time::precise_time_s() - start) / (batch as f64));
                }
            }
        }
    });
}

// returns pairs (n, s) indicating node n can be reached from a root in s steps.
fn bfs<G: Scope>(edges: &Stream<G, (Edge, i32)>, roots: &Stream<G, (Node, i32)>) -> Stream<G, ((Node, u32), i32)>
where G::Timestamp: LeastUpperBound {

    // initialize roots as reaching themselves at distance 0
    let nodes = roots.map(|(x,w)| ((x, 0), w));
    // let edges = edges.map_in_place(|x| x.0 = ((x.0).1, (x.0).0))
    //                  .concat(&edges);

    // repeatedly update minimal distances each node can be reached from each root
    nodes.iterate(|inner| {

        let edges = inner.scope().enter(&edges);
        let nodes = inner.scope().enter(&nodes);

        inner.join_map_u(&edges, |_k,l,d| (*d, l+1))
             .concat(&nodes)
             .group_u(|_, s, t| t.push((*s.peek().unwrap().0, 1)))
     })
}
<|endoftext|>


<commit_before>use super::btree::{BPlusTree, RTCursor as BPlusTreeCursor};
use super::*;
use client::AsyncClient;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
use itertools::Itertools;
use parking_lot::RwLock;
use ram::segs::MAX_SEGMENT_SIZE;
use std::collections::BTreeSet;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use std::{mem, ptr};

pub const LEVEL_ELEMENTS_MULTIPLIER: usize = 10;
pub const LEVEL_PAGE_DIFF_MULTIPLIER: usize = 10;

const LEVEL_M_MAX_ELEMENTS_COUNT: usize = LEVEL_M * LEVEL_M * LEVEL_M;
const LEVEL_M: usize = 24;
const LEVEL_1: usize = LEVEL_M * LEVEL_PAGE_DIFF_MULTIPLIER;
const LEVEL_2: usize = LEVEL_1 * LEVEL_PAGE_DIFF_MULTIPLIER;
const LEVEL_3: usize = LEVEL_2 * LEVEL_PAGE_DIFF_MULTIPLIER;
const LEVEL_4: usize = LEVEL_3 * LEVEL_PAGE_DIFF_MULTIPLIER;

type LevelTrees = Vec<Box<LevelTree>>;
pub type Ptr = NodeCellRef;
pub type Key = EntryKey;

with_levels! {
    LM, LEVEL_M;
    L1, LEVEL_1;
    L2, LEVEL_2;
    L3, LEVEL_3;
    // L4, LEVEL_4; // See https://github.com/rust-lang/rust/issues/58164
}

pub struct LSMTree {
    trees: LevelTrees,
    // use Vec here for convenience
    max_sizes: Vec<usize>,
}

unsafe impl Send for LSMTree {}
unsafe impl Sync for LSMTree {}

impl LSMTree {
    pub fn new(neb_client: &Arc<AsyncClient>) -> Arc<Self> {
        debug!("Initializing LSM-tree...");
        let (trees, max_sizes) = init_lsm_level_trees(neb_client);
        debug!("Initialized LSM-tree");
        Arc::new(LSMTree { trees, max_sizes })
    }

    pub fn insert(&self, mut key: EntryKey, id: &Id) {
        key_with_id(&mut key, id);
        self.trees[0].insert_into(&key)
    }

    pub fn remove(&self, mut key: EntryKey, id: &Id) -> bool {
        key_with_id(&mut key, id);
        self.trees
            .iter()
            .map(|tree| tree.mark_key_deleted(&key))
            .collect_vec() // collect here to prevent short circuit
            .into_iter()
            .any(|d| d)
    }

    pub fn seek(&self, mut key: EntryKey, ordering: Ordering) -> LSMTreeCursor {
        match ordering {
            Ordering::Forward => key_with_id(&mut key, &Id::unit_id()),
            Ordering::Backward => key_with_id(&mut key, &Id::new(::std::u64::MAX, ::std::u64::MAX)),
        };
        let mut cursors: Vec<Box<Cursor>> = vec![];
        for tree in &self.trees {
            cursors.push(tree.seek_for(&key, ordering));
        }
        return LSMTreeCursor::new(cursors);
    }

    pub fn check_and_merge(&self) {
        for i in 0..self.trees.len() - 1 {
            debug!("Checking tree merge {}", i);
            let lower = &*self.trees[i];
            let upper = &*self.trees[i + 1];
            if lower.count() > self.max_sizes[i] {
                lower.merge_to(upper);
            }
        }
    }

    pub fn start_sentinel(this: &Arc<Self>) {
        let this = this.clone();
        thread::Builder::new().name("LSM-Tree Sentinel".to_string()).spawn(move || {
            loop {
                this.check_and_merge();
                thread::sleep(Duration::from_millis(500));
            }
        });
    }

    pub fn level_sizes(&self) -> Vec<usize> {
        self.trees.iter().map(|t| t.count()).collect()
    }

    pub fn len(&self) -> usize {
        self.trees.iter().map(|tree| tree.count()).sum::<usize>()
    }
}

pub struct LSMTreeCursor {
    level_cursors: Vec<Box<Cursor>>,
}

impl LSMTreeCursor {
    fn new(cursors: Vec<Box<Cursor>>) -> Self {
        LSMTreeCursor {
            level_cursors: cursors,
        }
    }
}

impl Cursor for LSMTreeCursor {
    fn next(&mut self) -> bool {
        if let Some(prev_key) = self.current().map(|k| k.to_owned()) {
            self.level_cursors
                .iter_mut()
                .filter(|c| c.current().is_some())
                .min_by(|a, b| a.current().unwrap().cmp(b.current().unwrap()))
                .map(|c| c.next());
            let dedupe_current = if let Some(current) = self.current() {
                if current <= &prev_key {
                    None
                } else {
                    Some(true)
                }
            } else {
                Some(false)
            };
            if let Some(has_next) = dedupe_current {
                has_next
            } else {
                self.next()
            }
        } else {
            false
        }
    }

    fn current(&self) -> Option<&EntryKey> {
        self.level_cursors.iter().filter_map(|c| c.current()).min()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use byteorder::BigEndian;
    use byteorder::WriteBytesExt;
    use client;
    use futures::prelude::*;
    use rand::distributions::Uniform;
    use rand::thread_rng;
    use rand::Rng;
    use rayon::prelude::*;
    use server::NebServer;
    use server::ServerOptions;
    use std::env;
    use std::io::Cursor as StdCursor;

    fn u64_to_slice(n: u64) -> [u8; 8] {
        let mut key_slice = [0u8; 8];
        {
            let mut cursor = StdCursor::new(&mut key_slice[..]);
            cursor.write_u64::<BigEndian>(n);
        };
        key_slice
    }

    fn dump_trees(lsm_tree: &LSMTree, name: &str) {
        for i in 0..lsm_tree.trees.len() {
            lsm_tree.trees[i].dump(&format!("{}_lsm_{}_dump.json", name, i));
        }
    }

    #[test]
    pub fn insertions() {
        env_logger::init();
        let server_group = "sstable_index_init";
        let server_addr = String::from("127.0.0.1:5700");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 1,
                memory_size: 3 * 1024 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
            },
            &server_addr,
            &server_group,
        );
        let client = Arc::new(
            client::AsyncClient::new(&server.rpc, &vec![server_addr], server_group).unwrap(),
        );
        client.new_schema_with_id(btree::page_schema()).wait();
        let tree = LSMTree::new(&client);
        let num = env::var("LSM_TREE_TEST_ITEMS")
            // this value cannot do anything useful to the test
            // must arrange a long-term test to cover every levels
            .unwrap_or("1000".to_string())
            .parse::<u64>()
            .unwrap();

        let tree_clone = tree.clone();
        debug!("Testing LSM-tree");
        thread::spawn(move || loop {
            thread::sleep(Duration::from_secs(10));
            let tree_len = tree_clone.len();
            debug!(
                "LSM-Tree now have {}/{} elements, total {:.2}%",
                tree_len,
                num,
                tree_len as f32 / num as f32 * 100.0
            );
        });

        (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let key = SmallVec::from_slice(&key_slice);
            tree.insert(key, &id);
        });

        dump_trees(&*tree, "stage_1_before_insertion");
        tree.check_and_merge();
        dump_trees(&*tree, "stage_1_after_insertion");

        debug!("Start point search validations");
        (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let mut key = SmallVec::from_slice(&key_slice);
            debug!("checking: {}", i);
            let mut cursor = tree.seek(key.clone(), Ordering::Forward);
            key_with_id(&mut key, &id);
            assert_eq!(cursor.current(), Some(&key), "{}", i);
        });

        LSMTree::start_sentinel(&tree);

        (num..num * 2).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let key = SmallVec::from_slice(&key_slice);
            tree.insert(key, &id);
        });

        dump_trees(&*tree, "stage_2_after_insertion");
        thread::sleep(Duration::new(30, 0));
        dump_trees(&*tree, "stage_2_waited_insertion");

        (0..num * 2).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let mut key = SmallVec::from_slice(&key_slice);
            debug!("checking: {}", i);
            let mut cursor = tree.seek(key.clone(), Ordering::Forward);
            key_with_id(&mut key, &id);
            assert_eq!(cursor.current(), Some(&key), "{}", i);
        });
    }
}
<commit_msg>reset test data for lsm-tree<commit_after>use super::btree::{BPlusTree, RTCursor as BPlusTreeCursor};
use super::*;
use client::AsyncClient;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
use itertools::Itertools;
use parking_lot::RwLock;
use ram::segs::MAX_SEGMENT_SIZE;
use std::collections::BTreeSet;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use std::{mem, ptr};

pub const LEVEL_ELEMENTS_MULTIPLIER: usize = 10;
pub const LEVEL_PAGE_DIFF_MULTIPLIER: usize = 10;

const LEVEL_M_MAX_ELEMENTS_COUNT: usize = LEVEL_M * LEVEL_M * LEVEL_M;
const LEVEL_M: usize = 24;
const LEVEL_1: usize = LEVEL_M * LEVEL_PAGE_DIFF_MULTIPLIER;
const LEVEL_2: usize = LEVEL_1 * LEVEL_PAGE_DIFF_MULTIPLIER;
const LEVEL_3: usize = LEVEL_2 * LEVEL_PAGE_DIFF_MULTIPLIER;
const LEVEL_4: usize = LEVEL_3 * LEVEL_PAGE_DIFF_MULTIPLIER;

type LevelTrees = Vec<Box<LevelTree>>;
pub type Ptr = NodeCellRef;
pub type Key = EntryKey;

with_levels! {
    LM, LEVEL_M;
    L1, LEVEL_1;
    L2, LEVEL_2;
    L3, LEVEL_3;
    // L4, LEVEL_4; // See https://github.com/rust-lang/rust/issues/58164
}

pub struct LSMTree {
    trees: LevelTrees,
    // use Vec here for convenience
    max_sizes: Vec<usize>,
}

unsafe impl Send for LSMTree {}
unsafe impl Sync for LSMTree {}

impl LSMTree {
    pub fn new(neb_client: &Arc<AsyncClient>) -> Arc<Self> {
        debug!("Initializing LSM-tree...");
        let (trees, max_sizes) = init_lsm_level_trees(neb_client);
        debug!("Initialized LSM-tree");
        Arc::new(LSMTree { trees, max_sizes })
    }

    pub fn insert(&self, mut key: EntryKey, id: &Id) {
        key_with_id(&mut key, id);
        self.trees[0].insert_into(&key)
    }

    pub fn remove(&self, mut key: EntryKey, id: &Id) -> bool {
        key_with_id(&mut key, id);
        self.trees
            .iter()
            .map(|tree| tree.mark_key_deleted(&key))
            .collect_vec() // collect here to prevent short circuit
            .into_iter()
            .any(|d| d)
    }

    pub fn seek(&self, mut key: EntryKey, ordering: Ordering) -> LSMTreeCursor {
        match ordering {
            Ordering::Forward => key_with_id(&mut key, &Id::unit_id()),
            Ordering::Backward => key_with_id(&mut key, &Id::new(::std::u64::MAX, ::std::u64::MAX)),
        };
        let mut cursors: Vec<Box<Cursor>> = vec![];
        for tree in &self.trees {
            cursors.push(tree.seek_for(&key, ordering));
        }
        return LSMTreeCursor::new(cursors);
    }

    pub fn check_and_merge(&self) {
        for i in 0..self.trees.len() - 1 {
            debug!("Checking tree merge {}", i);
            let lower = &*self.trees[i];
            let upper = &*self.trees[i + 1];
            if lower.count() > self.max_sizes[i] {
                lower.merge_to(upper);
            }
        }
    }

    pub fn start_sentinel(this: &Arc<Self>) {
        let this = this.clone();
        thread::Builder::new().name("LSM-Tree Sentinel".to_string()).spawn(move || {
            loop {
                this.check_and_merge();
                thread::sleep(Duration::from_millis(500));
            }
        });
    }

    pub fn level_sizes(&self) -> Vec<usize> {
        self.trees.iter().map(|t| t.count()).collect()
    }

    pub fn len(&self) -> usize {
        self.trees.iter().map(|tree| tree.count()).sum::<usize>()
    }
}

pub struct LSMTreeCursor {
    level_cursors: Vec<Box<Cursor>>,
}

impl LSMTreeCursor {
    fn new(cursors: Vec<Box<Cursor>>) -> Self {
        LSMTreeCursor {
            level_cursors: cursors,
        }
    }
}

impl Cursor for LSMTreeCursor {
    fn next(&mut self) -> bool {
        if let Some(prev_key) = self.current().map(|k| k.to_owned()) {
            self.level_cursors
                .iter_mut()
                .filter(|c| c.current().is_some())
                .min_by(|a, b| a.current().unwrap().cmp(b.current().unwrap()))
                .map(|c| c.next());
            let dedupe_current = if let Some(current) = self.current() {
                if current <= &prev_key {
                    None
                } else {
                    Some(true)
                }
            } else {
                Some(false)
            };
            if let Some(has_next) = dedupe_current {
                has_next
            } else {
                self.next()
            }
        } else {
            false
        }
    }

    fn current(&self) -> Option<&EntryKey> {
        self.level_cursors.iter().filter_map(|c| c.current()).min()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use byteorder::BigEndian;
    use byteorder::WriteBytesExt;
    use client;
    use futures::prelude::*;
    use rand::distributions::Uniform;
    use rand::thread_rng;
    use rand::Rng;
    use rayon::prelude::*;
    use server::NebServer;
    use server::ServerOptions;
    use std::env;
    use std::io::Cursor as StdCursor;

    fn u64_to_slice(n: u64) -> [u8; 8] {
        let mut key_slice = [0u8; 8];
        {
            let mut cursor = StdCursor::new(&mut key_slice[..]);
            cursor.write_u64::<BigEndian>(n);
        };
        key_slice
    }

    fn dump_trees(lsm_tree: &LSMTree, name: &str) {
        for i in 0..lsm_tree.trees.len() {
            lsm_tree.trees[i].dump(&format!("{}_lsm_{}_dump.json", name, i));
        }
    }

    #[test]
    pub fn insertions() {
        env_logger::init();
        let server_group = "sstable_index_init";
        let server_addr = String::from("127.0.0.1:5700");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 1,
                memory_size: 3 * 1024 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
            },
            &server_addr,
            &server_group,
        );
        let client = Arc::new(
            client::AsyncClient::new(&server.rpc, &vec![server_addr], server_group).unwrap(),
        );
        client.new_schema_with_id(btree::page_schema()).wait();
        let tree = LSMTree::new(&client);
        let num = env::var("LSM_TREE_TEST_ITEMS")
            // this value cannot do anything useful to the test
            // must arrange a long-term test to cover every levels
            .unwrap_or("331776".to_string())
            .parse::<u64>()
            .unwrap();

        let tree_clone = tree.clone();
        debug!("Testing LSM-tree");
        thread::spawn(move || loop {
            thread::sleep(Duration::from_secs(10));
            let tree_len = tree_clone.len();
            debug!(
                "LSM-Tree now have {}/{} elements, total {:.2}%",
                tree_len,
                num,
                tree_len as f32 / num as f32 * 100.0
            );
        });

        (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let key = SmallVec::from_slice(&key_slice);
            tree.insert(key, &id);
        });

        dump_trees(&*tree, "stage_1_before_insertion");
        tree.check_and_merge();
        dump_trees(&*tree, "stage_1_after_insertion");

        debug!("Start point search validations");
        (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let mut key = SmallVec::from_slice(&key_slice);
            debug!("checking: {}", i);
            let mut cursor = tree.seek(key.clone(), Ordering::Forward);
            key_with_id(&mut key, &id);
            assert_eq!(cursor.current(), Some(&key), "{}", i);
        });

        LSMTree::start_sentinel(&tree);

        (num..num * 2).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let key = SmallVec::from_slice(&key_slice);
            tree.insert(key, &id);
        });

        dump_trees(&*tree, "stage_2_after_insertion");
        thread::sleep(Duration::new(30, 0));
        dump_trees(&*tree, "stage_2_waited_insertion");

        (0..num * 2).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let mut key = SmallVec::from_slice(&key_slice);
            debug!("checking: {}", i);
            let mut cursor = tree.seek(key.clone(), Ordering::Forward);
            key_with_id(&mut key, &id);
            assert_eq!(cursor.current(), Some(&key), "{}", i);
        });
    }
}
<|endoftext|>


<commit_before>//! Tests for various types of file (video, image, compressed, etc).
//!
//! Currently this is dependent on the file’s name and extension, because
//! those are the only metadata that we have access to without reading the
//! file’s contents.

use ansi_term::Style;

use crate::fs::File;
use crate::output::file_name::FileColours;
use crate::output::icons::FileIcon;

#[derive(Debug, Default, PartialEq)]
pub struct FileExtensions;

impl FileExtensions {

    /// An “immediate” file is something that can be run or activated somehow
    /// in order to kick off the build of a project. It’s usually only present
    /// in directories full of source code.
    fn is_immediate(&self, file: &File) -> bool {
        file.name.to_lowercase().starts_with("readme") ||
        file.name.ends_with(".ninja") ||
        file.name_is_one_of( &[
            "Makefile", "Cargo.toml", "SConstruct", "CMakeLists.txt",
            "build.gradle", "pom.xml", "Rakefile", "package.json", "Gruntfile.js",
            "Gruntfile.coffee", "BUILD", "BUILD.bazel", "WORKSPACE", "build.xml",
            "webpack.config.js", "meson.build",
        ])
    }

    fn is_image(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "png", "jpeg", "jpg", "gif", "bmp", "tiff", "tif",
            "ppm", "pgm", "pbm", "pnm", "webp", "raw", "arw",
            "svg", "stl", "eps", "dvi", "ps", "cbr", "jpf",
            "cbz", "xpm", "ico", "cr2", "orf", "nef",
        ])
    }

    fn is_video(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "avi", "flv", "m2v", "m4v", "mkv", "mov", "mp4", "mpeg",
            "mpg", "ogm", "ogv", "vob", "wmv", "webm", "m2ts",
        ])
    }

    fn is_music(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "aac", "m4a", "mp3", "ogg", "wma", "mka", "opus",
        ])
    }

    // Lossless music, rather than any other kind of data...
    fn is_lossless(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "alac", "ape", "flac", "wav",
        ])
    }

    fn is_crypto(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "asc", "enc", "gpg", "pgp", "sig", "signature", "pfx", "p12",
        ])
    }

    fn is_document(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "djvu", "doc", "docx", "dvi", "eml", "eps", "fotd",
            "odp", "odt", "pdf", "ppt", "pptx", "rtf",
            "xls", "xlsx",
        ])
    }

    fn is_compressed(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "zip", "tar", "Z", "z", "gz", "bz2", "a", "ar", "7z",
            "iso", "dmg", "tc", "rar", "par", "tgz", "xz", "txz",
            "lz", "tlz", "lzma", "deb", "rpm", "zst",
        ])
    }

    fn is_temp(&self, file: &File) -> bool {
        file.name.ends_with('~')
            || (file.name.starts_with('#') && file.name.ends_with('#'))
            || file.extension_is_one_of( &[ "tmp", "swp", "swo", "swn", "bak", "bk" ])
    }

    fn is_compiled(&self, file: &File) -> bool {
        if file.extension_is_one_of( &[ "class", "elc", "hi", "o", "pyc", "zwc" ]) {
            true
        }
        else if let Some(dir) = file.parent_dir {
            file.get_source_files().iter().any(|path| dir.contains(path))
        }
        else {
            false
        }
    }
}

impl FileColours for FileExtensions {
    fn colour_file(&self, file: &File) -> Option<Style> {
        use ansi_term::Colour::*;

        Some(match file {
            f if self.is_temp(f)        => Fixed(244).normal(),
            f if self.is_immediate(f)   => Yellow.bold().underline(),
            f if self.is_image(f)       => Fixed(133).normal(),
            f if self.is_video(f)       => Fixed(135).normal(),
            f if self.is_music(f)       => Fixed(92).normal(),
            f if self.is_lossless(f)    => Fixed(93).normal(),
            f if self.is_crypto(f)      => Fixed(109).normal(),
            f if self.is_document(f)    => Fixed(105).normal(),
            f if self.is_compressed(f)  => Red.normal(),
            f if self.is_compiled(f)    => Fixed(137).normal(),
            _                           => return None,
        })
    }
}

impl FileIcon for FileExtensions {
    fn icon_file(&self, file: &File) -> Option<char> {
        use crate::output::icons::Icons;

        Some(match file {
            f if self.is_music(f) || self.is_lossless(f) => Icons::Audio.value(),
            f if self.is_image(f) => Icons::Image.value(),
            f if self.is_video(f) => Icons::Video.value(),
            _ => return None,
        })
    }
}
<commit_msg>Add the documents extension for Apple file formats<commit_after>//! Tests for various types of file (video, image, compressed, etc).
//!
//! Currently this is dependent on the file’s name and extension, because
//! those are the only metadata that we have access to without reading the
//! file’s contents.

use ansi_term::Style;

use crate::fs::File;
use crate::output::file_name::FileColours;
use crate::output::icons::FileIcon;

#[derive(Debug, Default, PartialEq)]
pub struct FileExtensions;

impl FileExtensions {

    /// An “immediate” file is something that can be run or activated somehow
    /// in order to kick off the build of a project. It’s usually only present
    /// in directories full of source code.
    fn is_immediate(&self, file: &File) -> bool {
        file.name.to_lowercase().starts_with("readme") ||
        file.name.ends_with(".ninja") ||
        file.name_is_one_of( &[
            "Makefile", "Cargo.toml", "SConstruct", "CMakeLists.txt",
            "build.gradle", "pom.xml", "Rakefile", "package.json", "Gruntfile.js",
            "Gruntfile.coffee", "BUILD", "BUILD.bazel", "WORKSPACE", "build.xml",
            "webpack.config.js", "meson.build",
        ])
    }

    fn is_image(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "png", "jpeg", "jpg", "gif", "bmp", "tiff", "tif",
            "ppm", "pgm", "pbm", "pnm", "webp", "raw", "arw",
            "svg", "stl", "eps", "dvi", "ps", "cbr", "jpf",
            "cbz", "xpm", "ico", "cr2", "orf", "nef",
        ])
    }

    fn is_video(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "avi", "flv", "m2v", "m4v", "mkv", "mov", "mp4", "mpeg",
            "mpg", "ogm", "ogv", "vob", "wmv", "webm", "m2ts",
        ])
    }

    fn is_music(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "aac", "m4a", "mp3", "ogg", "wma", "mka", "opus",
        ])
    }

    // Lossless music, rather than any other kind of data...
    fn is_lossless(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "alac", "ape", "flac", "wav",
        ])
    }

    fn is_crypto(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "asc", "enc", "gpg", "pgp", "sig", "signature", "pfx", "p12",
        ])
    }

    fn is_document(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "djvu", "doc", "docx", "dvi", "eml", "eps", "fotd", "key",
            "keynote", "numbers", "odp", "odt", "pages", "pdf", "ppt",
            "pptx", "rtf", "xls", "xlsx",
        ])
    }

    fn is_compressed(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "zip", "tar", "Z", "z", "gz", "bz2", "a", "ar", "7z",
            "iso", "dmg", "tc", "rar", "par", "tgz", "xz", "txz",
            "lz", "tlz", "lzma", "deb", "rpm", "zst",
        ])
    }

    fn is_temp(&self, file: &File) -> bool {
        file.name.ends_with('~')
            || (file.name.starts_with('#') && file.name.ends_with('#'))
            || file.extension_is_one_of( &[ "tmp", "swp", "swo", "swn", "bak", "bk" ])
    }

    fn is_compiled(&self, file: &File) -> bool {
        if file.extension_is_one_of( &[ "class", "elc", "hi", "o", "pyc", "zwc" ]) {
            true
        }
        else if let Some(dir) = file.parent_dir {
            file.get_source_files().iter().any(|path| dir.contains(path))
        }
        else {
            false
        }
    }
}

impl FileColours for FileExtensions {
    fn colour_file(&self, file: &File) -> Option<Style> {
        use ansi_term::Colour::*;

        Some(match file {
            f if self.is_temp(f)        => Fixed(244).normal(),
            f if self.is_immediate(f)   => Yellow.bold().underline(),
            f if self.is_image(f)       => Fixed(133).normal(),
            f if self.is_video(f)       => Fixed(135).normal(),
            f if self.is_music(f)       => Fixed(92).normal(),
            f if self.is_lossless(f)    => Fixed(93).normal(),
            f if self.is_crypto(f)      => Fixed(109).normal(),
            f if self.is_document(f)    => Fixed(105).normal(),
            f if self.is_compressed(f)  => Red.normal(),
            f if self.is_compiled(f)    => Fixed(137).normal(),
            _                           => return None,
        })
    }
}

impl FileIcon for FileExtensions {
    fn icon_file(&self, file: &File) -> Option<char> {
        use crate::output::icons::Icons;

        Some(match file {
            f if self.is_music(f) || self.is_lossless(f) => Icons::Audio.value(),
            f if self.is_image(f) => Icons::Image.value(),
            f if self.is_video(f) => Icons::Video.value(),
            _ => return None,
        })
    }
}
<|endoftext|>


<commit_before>use ::kdtree::*;

pub struct Bounds {
    pub bounds: Vec<(f64, f64)>,

    widest_dim : usize,
    midvalue_of_widest_dim : f64,
}

impl Bounds {
    pub fn new_from_points<T: KdtreePointTrait>(points: &Vec<T>) -> Bounds {
        let mut bounds = Bounds {
            bounds: vec![],
            widest_dim : 0,
            midvalue_of_widest_dim : 0.,
        };

        bounds.bounds.resize(points[0].dims().len(), (0., 0.));

        for i in 0..points[0].dims().len() {
            bounds.bounds[i].0 = points[0].dims()[i];
            bounds.bounds[i].1 = points[0].dims()[i];
        }

        for v in points.iter() {
            for dim in 0..v.dims().len() {
                bounds.bounds[dim].0 = bounds.bounds[dim].0.min(v.dims()[dim]);
                bounds.bounds[dim].1 = bounds.bounds[dim].1.max(v.dims()[dim]);
            }
        }

        bounds.calculate_variables();

        bounds
    }

    pub fn get_widest_dim(&self) -> usize {
        self.widest_dim
    }

    pub fn get_midvalue_of_widest_dim(&self) -> f64 {
        self.midvalue_of_widest_dim
    }

    pub fn clone_moving_max(&self, value: f64, dimension: usize) -> Bounds {
        let mut cloned = Bounds {
            bounds: self.bounds.clone(),
            ..*self
        };
        cloned.bounds[dimension].1 = value;

        cloned.calculate_variables();

        cloned
    }

    pub fn clone_moving_min(&self, value: f64, dimension: usize) -> Bounds {
        let mut cloned = Bounds {
            bounds: self.bounds.clone(),
            ..*self
        };
        cloned.bounds[dimension].0 = value;

        cloned.calculate_variables();

        cloned
    }

    fn calculate_widest_dim(&mut self)  {
        let mut widest_dimension = 0usize;
        let mut max_found_spread = self.bounds[0].1 - self.bounds[0].0;

        for i in 0..self.bounds.len() {
            let dimension_spread = self.bounds[i].1 - self.bounds[i].0;

            if dimension_spread > max_found_spread {
                max_found_spread = dimension_spread;
                widest_dimension = i;
            }
        }

        self.widest_dim = widest_dimension;
    }

    fn calculate_variables(&mut self) {
        self.calculate_widest_dim();
        self.midvalue_of_widest_dim = (self.bounds[self.get_widest_dim()].0 + self.bounds[self.get_widest_dim()].1) / 2.0;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ::kdtree::test_common::tests_utils::*;

    #[test]
    fn bounds_test() {
        let p1 = Point2WithId::new(1, 1.0, 0.5);
        let p2 = Point2WithId::new(1, 3.0, 4.0);
        let v = vec![p1, p2];

        let bounds = Bounds::new_from_points(&v);

        assert_eq!((1., 3.0), bounds.bounds[0]);
        assert_eq!((0.5, 4.0), bounds.bounds[1]);

        assert_eq!(1, bounds.get_widest_dim());
    }
}<commit_msg>Issue1 (#3)<commit_after>use ::kdtree::*;

pub struct Bounds {
    pub bounds: [(f64,f64);3],

    widest_dim : usize,
    midvalue_of_widest_dim : f64,
}

impl Bounds {
    pub fn new_from_points<T: KdtreePointTrait>(points: &Vec<T>) -> Bounds {
        let mut bounds = Bounds {
            bounds: [(0.,0.),(0.,0.),(0.,0.)],
            widest_dim : 0,
            midvalue_of_widest_dim : 0.,
        };

        for i in 0..points[0].dims().len() {
            bounds.bounds[i].0 = points[0].dims()[i];
            bounds.bounds[i].1 = points[0].dims()[i];
        }

        for v in points.iter() {
            for dim in 0..v.dims().len() {
                bounds.bounds[dim].0 = bounds.bounds[dim].0.min(v.dims()[dim]);
                bounds.bounds[dim].1 = bounds.bounds[dim].1.max(v.dims()[dim]);
            }
        }

        bounds.calculate_variables();

        bounds
    }

    pub fn get_widest_dim(&self) -> usize {
        self.widest_dim
    }

    pub fn get_midvalue_of_widest_dim(&self) -> f64 {
        self.midvalue_of_widest_dim
    }

    pub fn clone_moving_max(&self, value: f64, dimension: usize) -> Bounds {
        let mut cloned = Bounds {
            bounds: self.bounds.clone(),
            ..*self
        };
        cloned.bounds[dimension].1 = value;

        cloned.calculate_variables();

        cloned
    }

    pub fn clone_moving_min(&self, value: f64, dimension: usize) -> Bounds {
        let mut cloned = Bounds {
            bounds: self.bounds.clone(),
            ..*self
        };
        cloned.bounds[dimension].0 = value;

        cloned.calculate_variables();

        cloned
    }

    fn calculate_widest_dim(&mut self)  {
        let mut widest_dimension = 0usize;
        let mut max_found_spread = self.bounds[0].1 - self.bounds[0].0;

        for i in 0..self.bounds.len() {
            let dimension_spread = self.bounds[i].1 - self.bounds[i].0;

            if dimension_spread > max_found_spread {
                max_found_spread = dimension_spread;
                widest_dimension = i;
            }
        }

        self.widest_dim = widest_dimension;
    }

    fn calculate_variables(&mut self) {
        self.calculate_widest_dim();
        self.midvalue_of_widest_dim = (self.bounds[self.get_widest_dim()].0 + self.bounds[self.get_widest_dim()].1) / 2.0;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ::kdtree::test_common::tests_utils::*;

    #[test]
    fn bounds_test() {
        let p1 = Point2WithId::new(1, 1.0, 0.5);
        let p2 = Point2WithId::new(1, 3.0, 4.0);
        let v = vec![p1, p2];

        let bounds = Bounds::new_from_points(&v);

        assert_eq!((1., 3.0), bounds.bounds[0]);
        assert_eq!((0.5, 4.0), bounds.bounds[1]);

        assert_eq!(1, bounds.get_widest_dim());
    }
}
<|endoftext|>


<commit_before>//! Art objects represent both tiles and static graphics.

use mul_reader::MulReader;
use std::io::{IoResult, MemReader, IoError, OtherIoError, SeekStyle};
use color::{Color, Color16, Color32};
use utils::DataBuffer;
use utils::MEMWRITER_ERROR;

pub trait Art {
    /**
     * Convert to a 32bit array
     *
     * Returns (width, height, colors)
     */
    fn to_32bit(&self) -> (u32, u32, Vec<Color32>);
    fn serialize(&self) -> Vec<u8>;
}

pub const TILE_SIZE: u32 = 2048;
pub const STATIC_OFFSET: u32 = 0x4000;

pub struct RunPair {
    pub offset: u16,
    pub run: Vec<Color16>
}

pub type StaticRow = Vec<RunPair>;

pub struct Tile {
    pub header: u32,
    pub image_data: [Color16, ..1022]
}

impl Art for Tile {
    fn to_32bit(&self) -> (u32, u32, Vec<Color32>) {
        let mut image: Vec<Color32> = vec![];

        let mut reader = DataBuffer::new(&self.image_data);

        for i in range(0, 44) {

            let slice_size = if i >= 22 {
                (44 - i) * 2
            } else {
                (i + 1) * 2
            };

            image.grow((22 - (slice_size / 2)), 0);
            for _pixel_idx in range(0, slice_size) {
                let (r, g, b, a) = reader.read().to_rgba();
                image.push(Color::from_rgba(r, g, b, a));
            }
            image.grow((22 - (slice_size / 2)), 0);
        };
        (44, 44, image)
    }

    fn serialize(&self) -> Vec<u8> {
        let mut writer = vec![];
        writer.write_le_u32(self.header).ok().expect(MEMWRITER_ERROR);
        for &pixel in self.image_data.iter() {
            writer.write_le_u16(pixel).ok().expect(MEMWRITER_ERROR);
        }
        writer
    }
}

impl Art for Static {
    fn to_32bit(&self) -> (u32, u32, Vec<Color32>) {
        let mut image: Vec<Color32> = vec![];

        for row in self.rows.iter() {
            let mut current_width = 0;
            for run_pair in row.iter() {
                image.grow(run_pair.offset as uint, 0);
                for pixel in run_pair.run.iter() {
                    let (r, g, b, a) = pixel.to_rgba();
                    image.push(Color::from_rgba(r, g, b, a));
                }
                current_width += run_pair.offset as uint + run_pair.run.len();
                assert!(current_width <= self.width as uint)
            }
            if current_width < self.width as uint {
                image.grow((self.width as uint) - current_width, 0)
            }
        };

        (self.width as u32, self.height as u32, image)
    }

    fn serialize(&self) -> Vec<u8> {
        let mut writer = vec![];
        writer.write_le_u16(self.size).ok().expect(MEMWRITER_ERROR);
        writer.write_le_u16(self.trigger).ok().expect(MEMWRITER_ERROR);
        writer.write_le_u16(self.width).ok().expect(MEMWRITER_ERROR);
        writer.write_le_u16(self.height).ok().expect(MEMWRITER_ERROR);
        unimplemented!();
    }
}

pub struct Static {
    pub size: u16,
    pub trigger: u16,
    pub width: u16,
    pub height: u16,
    pub rows: Vec<StaticRow>
}

pub enum TileOrStatic {
    Tile(Tile),
    Static(Static)
}

pub struct ArtReader {
    mul_reader: MulReader
}

impl ArtReader {
    pub fn new(index_path: &Path, mul_path: &Path) -> IoResult<ArtReader> {
        let mul_reader = try!(MulReader::new(index_path, mul_path));
        Ok(ArtReader {
            mul_reader: mul_reader
        })
    }

    pub fn read(&mut self, id: u32) -> IoResult<TileOrStatic> {
        let raw = try!(self.mul_reader.read(id));
        let mut reader = MemReader::new(raw.data);
        if id >= STATIC_OFFSET {
            //It's a static, so deal with accordingly
            let size = try!(reader.read_le_u16());
            let trigger = try!(reader.read_le_u16());
            let width = try!(reader.read_le_u16());
            let height = try!(reader.read_le_u16());
            if width == 0 || height >= 1024 || height == 0 || height >= 1024 {
                Err(IoError{
                    kind: OtherIoError,
                    desc: "Invalid image dimensions",
                    detail: Some(format!("Got invalid width and height of {}, {}", width, height))
                })
            } else {
                //Load our offset table
                let mut offset_table = vec![];
                for _index in range(0, height) {
                    offset_table.push(try!(reader.read_le_u16()));
                }

                let data_start_pos = try!(reader.tell());
                let mut rows = vec![];

                for &offset in offset_table.iter() {
                    try!(reader.seek((data_start_pos + offset as u64 * 2) as i64, SeekStyle::SeekSet));
                    let mut current_row_width = 0;
                    let mut row = vec![];

                    loop {
                        let x_offset = try!(reader.read_le_u16());
                        let run_length = try!(reader.read_le_u16());
                        if x_offset + run_length == 0 {
                            break
                        } else {
                            let mut run = vec![];
                            for _index in range(0, run_length) {
                                run.push(try!(reader.read_le_u16()));
                            }

                            row.push(RunPair {
                                offset: x_offset,
                                run: run
                            });
                            current_row_width += x_offset + run_length;
                            //assert!(current_row_width < width);
                        }
                    }
                    rows.push(row);
                }

                Ok(TileOrStatic::Static(Static {
                    size: size,
                    trigger: trigger,
                    width: width,
                    height: height,
                    rows: rows
                }))
            }
        } else {
            //It's a map tile
            if raw.length != TILE_SIZE {
                Err(IoError{
                    kind: OtherIoError,
                    desc: "Invalid tile size",
                    detail: Some(format!("Got tile size of {}, expected {}", raw.length, TILE_SIZE))
                })
            } else {
                let header = try!(reader.read_le_u32());
                let mut body = [0, ..1022];
                for idx in range(0, 1022) {
                    body[idx] = try!(reader.read_le_u16());
                }
                Ok(TileOrStatic::Tile(Tile {
                    header: header, image_data: body
                }))
            }
        }
    }

    pub fn read_tile(&mut self, id: u32) -> IoResult<Tile> {
        match self.read(id) {
            Ok(TileOrStatic::Tile(tile)) => Ok(tile),
            Ok(_) => Err(IoError {
                kind: OtherIoError,
                desc: "Index out of bounds",
                detail: None
            }),
            Err(e) => Err(e)
        }
    }

    pub fn read_static(&mut self, id: u32)-> IoResult<Static> {
        match self.read(id + STATIC_OFFSET) {
            Ok(TileOrStatic::Static(stat)) => Ok(stat),
            Ok(_) => Err(IoError {
                kind: OtherIoError,
                desc: "Index out of bounds",
                detail: None
            }),
            Err(e) => Err(e)
        }
    }
}
<commit_msg>Blocking out logic for serializing art<commit_after>//! Art objects represent both tiles and static graphics.

use mul_reader::MulReader;
use std::io::{IoResult, MemReader, IoError, OtherIoError, SeekStyle};
use color::{Color, Color16, Color32};
use utils::DataBuffer;
use utils::MEMWRITER_ERROR;

pub trait Art {
    /**
     * Convert to a 32bit array
     *
     * Returns (width, height, colors)
     */
    fn to_32bit(&self) -> (u32, u32, Vec<Color32>);
    fn serialize(&self) -> Vec<u8>;
}

pub const TILE_SIZE: u32 = 2048;
pub const STATIC_OFFSET: u32 = 0x4000;

pub struct RunPair {
    pub offset: u16,
    pub run: Vec<Color16>
}

pub type StaticRow = Vec<RunPair>;

pub struct Tile {
    pub header: u32,
    pub image_data: [Color16, ..1022]
}

impl Art for Tile {
    fn to_32bit(&self) -> (u32, u32, Vec<Color32>) {
        let mut image: Vec<Color32> = vec![];

        let mut reader = DataBuffer::new(&self.image_data);

        for i in range(0, 44) {

            let slice_size = if i >= 22 {
                (44 - i) * 2
            } else {
                (i + 1) * 2
            };

            image.grow((22 - (slice_size / 2)), 0);
            for _pixel_idx in range(0, slice_size) {
                let (r, g, b, a) = reader.read().to_rgba();
                image.push(Color::from_rgba(r, g, b, a));
            }
            image.grow((22 - (slice_size / 2)), 0);
        };
        (44, 44, image)
    }

    fn serialize(&self) -> Vec<u8> {
        let mut writer = vec![];
        writer.write_le_u32(self.header).ok().expect(MEMWRITER_ERROR);
        for &pixel in self.image_data.iter() {
            writer.write_le_u16(pixel).ok().expect(MEMWRITER_ERROR);
        }
        writer
    }
}

impl Art for Static {
    fn to_32bit(&self) -> (u32, u32, Vec<Color32>) {
        let mut image: Vec<Color32> = vec![];

        for row in self.rows.iter() {
            let mut current_width = 0;
            for run_pair in row.iter() {
                image.grow(run_pair.offset as uint, 0);
                for pixel in run_pair.run.iter() {
                    let (r, g, b, a) = pixel.to_rgba();
                    image.push(Color::from_rgba(r, g, b, a));
                }
                current_width += run_pair.offset as uint + run_pair.run.len();
                assert!(current_width <= self.width as uint)
            }
            if current_width < self.width as uint {
                image.grow((self.width as uint) - current_width, 0)
            }
        };

        (self.width as u32, self.height as u32, image)
    }

    fn serialize(&self) -> Vec<u8> {
        let mut writer = vec![];
        writer.write_le_u16(self.size).ok().expect(MEMWRITER_ERROR);
        writer.write_le_u16(self.trigger).ok().expect(MEMWRITER_ERROR);
        writer.write_le_u16(self.width).ok().expect(MEMWRITER_ERROR);
        writer.write_le_u16(self.height).ok().expect(MEMWRITER_ERROR);

        //TODO: Generate a lookup table

        //Write our rows
        for row in self.rows.iter() {
            for pair in row.iter() {
                writer.write(pair.serialize());
            }
            //We write a "newline" after each row
            writer.write_le_u16(0).ok().expect(MEMWRITER_ERROR);
            writer.write_le_u16(0).ok().expect(MEMWRITER_ERROR);
        }
        unimplemented!();
    }
}

pub struct Static {
    pub size: u16,
    pub trigger: u16,
    pub width: u16,
    pub height: u16,
    pub rows: Vec<StaticRow>
}

pub enum TileOrStatic {
    Tile(Tile),
    Static(Static)
}

pub struct ArtReader {
    mul_reader: MulReader
}

impl ArtReader {
    pub fn new(index_path: &Path, mul_path: &Path) -> IoResult<ArtReader> {
        let mul_reader = try!(MulReader::new(index_path, mul_path));
        Ok(ArtReader {
            mul_reader: mul_reader
        })
    }

    pub fn read(&mut self, id: u32) -> IoResult<TileOrStatic> {
        let raw = try!(self.mul_reader.read(id));
        let mut reader = MemReader::new(raw.data);
        if id >= STATIC_OFFSET {
            //It's a static, so deal with accordingly
            let size = try!(reader.read_le_u16());
            let trigger = try!(reader.read_le_u16());
            let width = try!(reader.read_le_u16());
            let height = try!(reader.read_le_u16());
            if width == 0 || height >= 1024 || height == 0 || height >= 1024 {
                Err(IoError{
                    kind: OtherIoError,
                    desc: "Invalid image dimensions",
                    detail: Some(format!("Got invalid width and height of {}, {}", width, height))
                })
            } else {
                //Load our offset table
                let mut offset_table = vec![];
                for _index in range(0, height) {
                    offset_table.push(try!(reader.read_le_u16()));
                }

                let data_start_pos = try!(reader.tell());
                let mut rows = vec![];

                for &offset in offset_table.iter() {
                    try!(reader.seek((data_start_pos + offset as u64 * 2) as i64, SeekStyle::SeekSet));
                    let mut current_row_width = 0;
                    let mut row = vec![];

                    loop {
                        let x_offset = try!(reader.read_le_u16());
                        let run_length = try!(reader.read_le_u16());
                        if x_offset + run_length == 0 {
                            break
                        } else {
                            let mut run = vec![];
                            for _index in range(0, run_length) {
                                run.push(try!(reader.read_le_u16()));
                            }

                            row.push(RunPair {
                                offset: x_offset,
                                run: run
                            });
                            current_row_width += x_offset + run_length;
                            //assert!(current_row_width < width);
                        }
                    }
                    rows.push(row);
                }

                Ok(TileOrStatic::Static(Static {
                    size: size,
                    trigger: trigger,
                    width: width,
                    height: height,
                    rows: rows
                }))
            }
        } else {
            //It's a map tile
            if raw.length != TILE_SIZE {
                Err(IoError{
                    kind: OtherIoError,
                    desc: "Invalid tile size",
                    detail: Some(format!("Got tile size of {}, expected {}", raw.length, TILE_SIZE))
                })
            } else {
                let header = try!(reader.read_le_u32());
                let mut body = [0, ..1022];
                for idx in range(0, 1022) {
                    body[idx] = try!(reader.read_le_u16());
                }
                Ok(TileOrStatic::Tile(Tile {
                    header: header, image_data: body
                }))
            }
        }
    }

    pub fn read_tile(&mut self, id: u32) -> IoResult<Tile> {
        match self.read(id) {
            Ok(TileOrStatic::Tile(tile)) => Ok(tile),
            Ok(_) => Err(IoError {
                kind: OtherIoError,
                desc: "Index out of bounds",
                detail: None
            }),
            Err(e) => Err(e)
        }
    }

    pub fn read_static(&mut self, id: u32)-> IoResult<Static> {
        match self.read(id + STATIC_OFFSET) {
            Ok(TileOrStatic::Static(stat)) => Ok(stat),
            Ok(_) => Err(IoError {
                kind: OtherIoError,
                desc: "Index out of bounds",
                detail: None
            }),
            Err(e) => Err(e)
        }
    }
}
<|endoftext|>