<commit_before>use std::cell::RefCell;
use std::mem;
use std::rc::Rc;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::thread;

use {spawn_unsafe, FnBox};
use sync::AtomicOption;

pub struct Scope<'a> {
    dtors: RefCell<Option<DtorChain<'a>>>
}

struct DtorChain<'a> {
    dtor: Box<FnBox + 'a>,
    next: Option<Box<DtorChain<'a>>>
}

enum JoinState {
    Running(thread::JoinHandle<()>),
    Joined,
}

impl JoinState {
    fn join(&mut self) {
        let mut state = JoinState::Joined;
        mem::swap(self, &mut state);
        if let JoinState::Running(handle) = state {
            let res = handle.join();

            if !thread::panicking() { res.unwrap(); }
        }
    }
}

pub struct ScopedJoinHandle<T> {
    inner: Rc<RefCell<JoinState>>,
    packet: Arc<AtomicOption<T>>,
    thread: thread::Thread,
}

/// Create a new `scope`, for scoped concurrency operations.
///
/// Using crossbeam for concurrency involves creating a [`Scope`][scope]. That's the job of
/// `scope`.
///
/// [scope]: struct.Scope.html
///
/// # Examples
///
/// Creating a scope:
///
/// ```
/// crossbeam::scope(|scope| {
///     // do stuff with scope here
/// });
/// ```
pub fn scope<'a, F, R>(f: F) -> R where F: FnOnce(&Scope<'a>) -> R {
    let mut scope = Scope { dtors: RefCell::new(None) };
    let ret = f(&scope);
    scope.drop_all();
    ret
}

impl<'a> Scope<'a> {
    // This method is carefully written in a transactional style, so
    // that it can be called directly and, if any dtor panics, can be
    // resumed in the unwinding this causes. By initially running the
    // method outside of any destructor, we avoid any leakage problems
    // due to @rust-lang/rust#14875.
    fn drop_all(&mut self) {
        loop {
            // use a separate scope to ensure that the RefCell borrow
            // is relinquishe before running `dtor`
            let dtor = {
                let mut dtors = self.dtors.borrow_mut();
                if let Some(mut node) = dtors.take() {
                    *dtors = node.next.take().map(|b| *b);
                    node.dtor
                } else {
                    return
                }
            };
            dtor.call_box()
        }
    }

    pub fn defer<F>(&self, f: F) where F: FnOnce() + 'a {
        let mut dtors = self.dtors.borrow_mut();
        *dtors = Some(DtorChain {
            dtor: Box::new(f),
            next: dtors.take().map(Box::new)
        });
    }

    /// Create a scoped thread.
    ///
    /// `spawn` is similar to the [`spawn`][spawn] function in Rust's standard library. The
    /// difference is that this thread is scoped, meaning that it's guaranteed to terminate
    /// before the current stack frame goes away, allowing you to reference the parent stack frame
    /// directly. This is ensured by having the parent thread join on the child thread before the
    /// scope exits.
    ///
    /// [spawn]: http://doc.rust-lang.org/std/thread/fn.spawn.html
    ///
    /// # Examples
    ///
    /// A basic scoped thread:
    ///
    /// ```
    /// crossbeam::scope(|scope| {
    ///     scope.spawn(|| {
    ///         println!("Hello from a scoped thread!");
    ///     });
    /// });
    /// ```
    ///
    /// When writing concurrent Rust programs, you'll sometimes see a pattern like this, using
    /// [`std::thread::spawn`][spawn]:
    ///
    /// ```ignore
    /// let array = [1, 2, 3];
    /// let mut guards = vec![];
    ///
    /// for i in &array {
    ///     let guard = std::thread::spawn(move || {
    ///         println!("element: {}", i);
    ///     });
    ///
    ///     guards.push(guard);
    /// }
    ///
    /// for guard in guards {
    ///     guard.join().unwrap();
    /// }
    /// ```
    ///
    /// The basic pattern is:
    ///
    /// 1. Iterate over some collection.
    /// 2. Spin up a thread to operate on each part of the collection.
    /// 3. Join all the threads.
    ///
    /// However, this code actually gives an error:
    ///
    /// ```text
    /// error: `array` does not live long enough
    /// for i in &array {
    ///           ^~~~~
    /// in expansion of for loop expansion
    /// note: expansion site
    /// note: reference must be valid for the static lifetime...
    /// note: ...but borrowed value is only valid for the block suffix following statement 0 at ...
    ///     let array = [1, 2, 3];
    ///     let mut guards = vec![];
    ///
    ///     for i in &array {
    ///         let guard = std::thread::spawn(move || {
    ///             println!("element: {}", i);
    /// ...
    /// error: aborting due to previous error
    /// ```
    ///
    /// Because [`std::thread::spawn`][spawn] doesn't know about this scope, it requires a
    /// `'static` lifetime. One way of giving it a proper lifetime is to [`Box`][box] it up:
    ///
    /// [box]: http://doc.rust-lang.org/std/boxed/struct.Box.html
    ///
    /// ```
    /// let array = Box::new([1, 2, 3]);
    /// let mut guards = vec![];
    ///
    /// for i in (0..array.len()) {
    ///     let i = array[i];
    ///     let guard = std::thread::spawn(move || {
    ///         println!("element: {}", i);
    ///     });
    ///
    ///     guards.push(guard);
    /// }
    ///
    /// for guard in guards {
    ///     guard.join().unwrap();
    /// }
    /// ```
    ///
    /// But this introduces unnecessary allocation: we know that we're joining the threads before
    /// our function returns, so just taking a reference _should_ be safe. Rust can't know that,
    /// though.
    ///
    /// Enter scoped threads. Here's our original example, using `spawn` from crossbeam rather
    /// than from `std::thread`:
    ///
    /// ```
    /// let array = [1, 2, 3];
    ///
    /// crossbeam::scope(|scope| {
    ///     for i in &array {
    ///         scope.spawn(move || {
    ///             println!("element: {}", i);
    ///         });
    ///     }
    /// });
    /// ```
    ///
    /// Much more straightforward.
    pub fn spawn<F, T>(&self, f: F) -> ScopedJoinHandle<T> where
        F: FnOnce() -> T + Send + 'a, T: Send + 'a
    {
        let their_packet = Arc::new(AtomicOption::new());
        let my_packet = their_packet.clone();

        let join_handle = unsafe {
            spawn_unsafe(move || {
                their_packet.swap(f(), Ordering::Relaxed);
            })
        };

        let thread = join_handle.thread().clone();
        let deferred_handle = Rc::new(RefCell::new(JoinState::Running(join_handle)));
        let my_handle = deferred_handle.clone();

        self.defer(move || {
            let mut state = deferred_handle.borrow_mut();
            state.join();
        });

        ScopedJoinHandle {
            inner: my_handle,
            packet: my_packet,
            thread: thread,
        }
    }
}

impl<T> ScopedJoinHandle<T> {
    pub fn join(self) -> T {
        self.inner.borrow_mut().join();
        self.packet.take(Ordering::Relaxed).unwrap()
    }

    pub fn thread(&self) -> &thread::Thread {
        &self.thread
    }
}

impl<'a> Drop for Scope<'a> {
    fn drop(&mut self) {
        self.drop_all()
    }
}
<commit_msg>Change scoped example to use Arc instead of Box<commit_after>use std::cell::RefCell;
use std::mem;
use std::rc::Rc;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::thread;

use {spawn_unsafe, FnBox};
use sync::AtomicOption;

pub struct Scope<'a> {
    dtors: RefCell<Option<DtorChain<'a>>>
}

struct DtorChain<'a> {
    dtor: Box<FnBox + 'a>,
    next: Option<Box<DtorChain<'a>>>
}

enum JoinState {
    Running(thread::JoinHandle<()>),
    Joined,
}

impl JoinState {
    fn join(&mut self) {
        let mut state = JoinState::Joined;
        mem::swap(self, &mut state);
        if let JoinState::Running(handle) = state {
            let res = handle.join();

            if !thread::panicking() { res.unwrap(); }
        }
    }
}

pub struct ScopedJoinHandle<T> {
    inner: Rc<RefCell<JoinState>>,
    packet: Arc<AtomicOption<T>>,
    thread: thread::Thread,
}

/// Create a new `scope`, for scoped concurrency operations.
///
/// Using crossbeam for concurrency involves creating a [`Scope`][scope]. That's the job of
/// `scope`.
///
/// [scope]: struct.Scope.html
///
/// # Examples
///
/// Creating a scope:
///
/// ```
/// crossbeam::scope(|scope| {
///     // do stuff with scope here
/// });
/// ```
pub fn scope<'a, F, R>(f: F) -> R where F: FnOnce(&Scope<'a>) -> R {
    let mut scope = Scope { dtors: RefCell::new(None) };
    let ret = f(&scope);
    scope.drop_all();
    ret
}

impl<'a> Scope<'a> {
    // This method is carefully written in a transactional style, so
    // that it can be called directly and, if any dtor panics, can be
    // resumed in the unwinding this causes. By initially running the
    // method outside of any destructor, we avoid any leakage problems
    // due to @rust-lang/rust#14875.
    fn drop_all(&mut self) {
        loop {
            // use a separate scope to ensure that the RefCell borrow
            // is relinquishe before running `dtor`
            let dtor = {
                let mut dtors = self.dtors.borrow_mut();
                if let Some(mut node) = dtors.take() {
                    *dtors = node.next.take().map(|b| *b);
                    node.dtor
                } else {
                    return
                }
            };
            dtor.call_box()
        }
    }

    pub fn defer<F>(&self, f: F) where F: FnOnce() + 'a {
        let mut dtors = self.dtors.borrow_mut();
        *dtors = Some(DtorChain {
            dtor: Box::new(f),
            next: dtors.take().map(Box::new)
        });
    }

    /// Create a scoped thread.
    ///
    /// `spawn` is similar to the [`spawn`][spawn] function in Rust's standard library. The
    /// difference is that this thread is scoped, meaning that it's guaranteed to terminate
    /// before the current stack frame goes away, allowing you to reference the parent stack frame
    /// directly. This is ensured by having the parent thread join on the child thread before the
    /// scope exits.
    ///
    /// [spawn]: http://doc.rust-lang.org/std/thread/fn.spawn.html
    ///
    /// # Examples
    ///
    /// A basic scoped thread:
    ///
    /// ```
    /// crossbeam::scope(|scope| {
    ///     scope.spawn(|| {
    ///         println!("Hello from a scoped thread!");
    ///     });
    /// });
    /// ```
    ///
    /// When writing concurrent Rust programs, you'll sometimes see a pattern like this, using
    /// [`std::thread::spawn`][spawn]:
    ///
    /// ```ignore
    /// let array = [1, 2, 3];
    /// let mut guards = vec![];
    ///
    /// for i in &array {
    ///     let guard = std::thread::spawn(move || {
    ///         println!("element: {}", i);
    ///     });
    ///
    ///     guards.push(guard);
    /// }
    ///
    /// for guard in guards {
    ///     guard.join().unwrap();
    /// }
    /// ```
    ///
    /// The basic pattern is:
    ///
    /// 1. Iterate over some collection.
    /// 2. Spin up a thread to operate on each part of the collection.
    /// 3. Join all the threads.
    ///
    /// However, this code actually gives an error:
    ///
    /// ```text
    /// error: `array` does not live long enough
    /// for i in &array {
    ///           ^~~~~
    /// in expansion of for loop expansion
    /// note: expansion site
    /// note: reference must be valid for the static lifetime...
    /// note: ...but borrowed value is only valid for the block suffix following statement 0 at ...
    ///     let array = [1, 2, 3];
    ///     let mut guards = vec![];
    ///
    ///     for i in &array {
    ///         let guard = std::thread::spawn(move || {
    ///             println!("element: {}", i);
    /// ...
    /// error: aborting due to previous error
    /// ```
    ///
    /// Because [`std::thread::spawn`][spawn] doesn't know about this scope, it requires a
    /// `'static` lifetime. One way of giving it a proper lifetime is to use an [`Arc`][arc]:
    ///
    /// [arc]: http://doc.rust-lang.org/stable/std/sync/struct.Arc.html
    ///
    /// ```
    /// use std::sync::Arc;
    ///
    /// let array = Arc::new([1, 2, 3]);
    /// let mut guards = vec![];
    ///
    /// for i in (0..array.len()) {
    ///     let a = array.clone();
    ///
    ///     let guard = std::thread::spawn(move || {
    ///         println!("element: {}", a[i]);
    ///     });
    ///
    ///     guards.push(guard);
    /// }
    ///
    /// for guard in guards {
    ///     guard.join().unwrap();
    /// }
    /// ```
    ///
    /// But this introduces unnecessary allocation, as `Arc<T>` puts its data on the heap, and we
    /// also end up dealing with reference counts. We know that we're joining the threads before
    /// our function returns, so just taking a reference _should_ be safe. Rust can't know that,
    /// though.
    ///
    /// Enter scoped threads. Here's our original example, using `spawn` from crossbeam rather
    /// than from `std::thread`:
    ///
    /// ```
    /// let array = [1, 2, 3];
    ///
    /// crossbeam::scope(|scope| {
    ///     for i in &array {
    ///         scope.spawn(move || {
    ///             println!("element: {}", i);
    ///         });
    ///     }
    /// });
    /// ```
    ///
    /// Much more straightforward.
    pub fn spawn<F, T>(&self, f: F) -> ScopedJoinHandle<T> where
        F: FnOnce() -> T + Send + 'a, T: Send + 'a
    {
        let their_packet = Arc::new(AtomicOption::new());
        let my_packet = their_packet.clone();

        let join_handle = unsafe {
            spawn_unsafe(move || {
                their_packet.swap(f(), Ordering::Relaxed);
            })
        };

        let thread = join_handle.thread().clone();
        let deferred_handle = Rc::new(RefCell::new(JoinState::Running(join_handle)));
        let my_handle = deferred_handle.clone();

        self.defer(move || {
            let mut state = deferred_handle.borrow_mut();
            state.join();
        });

        ScopedJoinHandle {
            inner: my_handle,
            packet: my_packet,
            thread: thread,
        }
    }
}

impl<T> ScopedJoinHandle<T> {
    pub fn join(self) -> T {
        self.inner.borrow_mut().join();
        self.packet.take(Ordering::Relaxed).unwrap()
    }

    pub fn thread(&self) -> &thread::Thread {
        &self.thread
    }
}

impl<'a> Drop for Scope<'a> {
    fn drop(&mut self) {
        self.drop_all()
    }
}
<|endoftext|>


<commit_before>use nom::multispace;
use nom::{IResult, Err, ErrorKind, Needed};
use std::str;

use column::Column;
use common::FieldExpression;
use common::{field_expr, field_list, unsigned_number, statement_terminator, table_list};
use condition::{condition_expr, ConditionExpression};
use table::Table;

#[derive(Clone, Debug, PartialEq)]
pub struct GroupByClause {
    columns: Vec<Column>,
    having: Option<ConditionExpression>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct LimitClause {
    limit: u64,
    offset: u64,
}

#[derive(Clone, Debug, PartialEq)]
pub enum OrderType {
    OrderAscending,
    OrderDescending,
}

#[derive(Clone, Debug, PartialEq)]
pub struct OrderClause {
    cols: Vec<Column>, // TODO(malte): can this be an arbitrary expr?
    order: OrderType,
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct SelectStatement {
    pub tables: Vec<Table>,
    pub distinct: bool,
    pub fields: FieldExpression,
    pub where_clause: Option<ConditionExpression>,
    pub group_by: Option<GroupByClause>,
    pub order: Option<OrderClause>,
    pub limit: Option<LimitClause>,
}

/// Parse GROUP BY clause
named!(group_by_clause<&[u8], GroupByClause>,
    complete!(chain!(
        multispace? ~
        caseless_tag!("group by") ~
        multispace ~
        group_columns: field_list ~
        having_clause: opt!(
            complete!(chain!(
                multispace? ~
                caseless_tag!("having") ~
                multispace? ~
                ce: condition_expr,
                || { ce }
            ))
        ),
    || {
        GroupByClause {
            columns: group_columns,
            having: having_clause,
        }
    }))
);

/// Parse LIMIT clause
named!(limit_clause<&[u8], LimitClause>,
    complete!(chain!(
        multispace? ~
        caseless_tag!("limit") ~
        multispace ~
        limit_val: unsigned_number ~
        offset_val: opt!(
            complete!(chain!(
                multispace? ~
                caseless_tag!("offset") ~
                multispace ~
                val: unsigned_number,
                || { val }
            ))
        ),
    || {
        LimitClause {
            limit: limit_val,
            offset: match offset_val {
                None => 0,
                Some(v) => v,
            },
        }
    }))
);

named!(order_clause<&[u8], OrderClause>,
    complete!(chain!(
        multispace? ~
        caseless_tag!("order by") ~
        multispace ~
        order_expr: field_list ~
        ordering: opt!(
            complete!(chain!(
                multispace? ~
                ordering: alt_complete!(
                      map!(caseless_tag!("desc"), |_| OrderType::OrderDescending)
                    | map!(caseless_tag!("asc"), |_| OrderType::OrderAscending)
                ),
                || { ordering }
            ))
        ),
    || {
        OrderClause {
            cols: order_expr,
            order: match ordering {
                None => OrderType::OrderAscending,
                Some(ref o) => o.clone(),
            },
        }
    }))
);

/// Parse WHERE clause of a selection
named!(where_clause<&[u8], ConditionExpression>,
    complete!(chain!(
        multispace? ~
        caseless_tag!("where") ~
        multispace ~
        cond: condition_expr,
        || { cond }
    ))
);

/// Parse rule for a SQL selection query.
/// TODO(malte): support nested queries as selection targets
named!(pub selection<&[u8], SelectStatement>,
    chain!(
        caseless_tag!("select") ~
        multispace ~
        distinct: opt!(caseless_tag!("distinct")) ~
        multispace? ~
        fields: field_expr ~
        delimited!(opt!(multispace), caseless_tag!("from"), opt!(multispace)) ~
        tables: table_list ~
        cond: opt!(where_clause) ~
        group_by: opt!(group_by_clause) ~
        order: opt!(order_clause) ~
        limit: opt!(limit_clause) ~
        statement_terminator,
        || {
            SelectStatement {
                tables: tables,
                distinct: distinct.is_some(),
                fields: fields,
                where_clause: cond,
                group_by: group_by,
                order: order,
                limit: limit,
            }
        }
    )
);

#[cfg(test)]
mod tests {
    use super::*;
    use column::{AggregationExpression, Column};
    use common::{FieldExpression, Operator};
    use condition::{ConditionBase, ConditionExpression, ConditionTree};
    use table::Table;

    fn columns(cols: &[&str]) -> Vec<Column> {
        cols.iter()
            .map(|c| Column::from(*c))
            .collect()
    }

    #[test]
    fn simple_select() {
        let qstring = "SELECT id, name FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("users")],
                       fields: FieldExpression::Seq(columns(&["id", "name"])),
                       ..Default::default()
                   });
    }

    #[test]
    fn select_all() {
        let qstring = "SELECT * FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("users")],
                       fields: FieldExpression::All,
                       ..Default::default()
                   });
    }

    #[test]
    fn spaces_optional() {
        let qstring = "SELECT id,name FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("users")],
                       fields: FieldExpression::Seq(columns(&["id", "name"])),
                       ..Default::default()
                   });
    }

    #[test]
    fn case_sensitivity() {
        let qstring_lc = "select id, name from users;";
        let qstring_uc = "SELECT id, name FROM users;";

        assert_eq!(selection(qstring_lc.as_bytes()).unwrap(),
                   selection(qstring_uc.as_bytes()).unwrap());
    }

    #[test]
    fn termination() {
        let qstring_sem = "select id, name from users;";
        let qstring_linebreak = "select id, name from users\n";

        let r1 = selection(qstring_sem.as_bytes()).unwrap();
        let r2 = selection(qstring_linebreak.as_bytes()).unwrap();
        assert_eq!(r1, r2);
    }

    #[test]
    fn where_clause() {
        let qstring = "select * from ContactInfo where email=?;";

        let res = selection(qstring.as_bytes());

        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(Column::from("email"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("ContactInfo")],
                       fields: FieldExpression::All,
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn limit_clause() {
        let qstring1 = "select * from users limit 10\n";
        let qstring2 = "select * from users limit 10 offset 10\n";

        let expected_lim1 = LimitClause {
            limit: 10,
            offset: 0,
        };
        let expected_lim2 = LimitClause {
            limit: 10,
            offset: 10,
        };

        let res1 = selection(qstring1.as_bytes());
        let res2 = selection(qstring2.as_bytes());
        assert_eq!(res1.unwrap().1.limit, Some(expected_lim1));
        assert_eq!(res2.unwrap().1.limit, Some(expected_lim2));
    }

    #[test]
    fn order_clause() {
        let qstring1 = "select * from users order by name desc\n";
        let qstring2 = "select * from users order by name, age desc\n";
        let qstring3 = "select * from users order by name\n";

        let expected_ord1 = OrderClause {
            cols: columns(&["name"]),
            order: OrderType::OrderDescending,
        };
        let expected_ord2 = OrderClause {
            cols: columns(&["name", "age"]),
            order: OrderType::OrderDescending,
        };
        let expected_ord3 = OrderClause {
            cols: columns(&["name"]),
            order: OrderType::OrderAscending,
        };

        let res1 = selection(qstring1.as_bytes());
        let res2 = selection(qstring2.as_bytes());
        let res3 = selection(qstring3.as_bytes());
        assert_eq!(res1.unwrap().1.order, Some(expected_ord1));
        assert_eq!(res2.unwrap().1.order, Some(expected_ord2));
        assert_eq!(res3.unwrap().1.order, Some(expected_ord3));
    }

    #[test]
    fn table_alias() {
        let qstring1 = "select * from PaperTag as t;";
        // let qstring2 = "select * from PaperTag t;";

        let res1 = selection(qstring1.as_bytes());
        assert_eq!(res1.clone().unwrap().1,
                   SelectStatement {
                       tables: vec![Table {
                           name: String::from("PaperTag"),
                           alias: Some(String::from("t")),
                       }],
                       fields: FieldExpression::All,
                       ..Default::default()
                   });
        // let res2 = selection(qstring2.as_bytes());
        // assert_eq!(res1.unwrap().1, res2.unwrap().1);
    }

    #[test]
    fn distinct() {
        let qstring = "select distinct tag from PaperTag where paperId=?;";

        let res = selection(qstring.as_bytes());
        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(Column::from("paperId"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("PaperTag")],
                       distinct: true,
                       fields: FieldExpression::Seq(columns(&["tag"])),
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn simple_condition_expr() {
        let qstring = "select infoJson from PaperStorage where paperId=? and paperStorageId=?;";

        let res = selection(qstring.as_bytes());

        let left_comp = Some(Box::new(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(Column::from("paperId"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        })));
        let right_comp = Some(Box::new(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(Column::from("paperStorageId"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        })));
        let expected_where_cond = Some(ConditionExpression::LogicalOp(ConditionTree {
            left: left_comp,
            right: right_comp,
            operator: Operator::And,
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("PaperStorage")],
                       fields: FieldExpression::Seq(columns(&["infoJson"])),
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn where_and_limit_clauses() {
        let qstring = "select * from users where id = ? limit 10\n";
        let res = selection(qstring.as_bytes());

        let expected_lim = Some(LimitClause {
            limit: 10,
            offset: 0,
        });
        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(Column::from("id"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        }));

        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("users")],
                       fields: FieldExpression::All,
                       where_clause: expected_where_cond,
                       limit: expected_lim,
                       ..Default::default()
                   });
    }

    #[test]
    fn aggregation_column() {
        let qstring = "SELECT max(addr_id) FROM address;";

        let res = selection(qstring.as_bytes());
        let agg_expr =
            AggregationExpression::Max(FieldExpression::Seq(vec![Column::from("addr_id")]));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("address")],
                       fields: FieldExpression::Seq(vec![Column {
                                                             name: String::from("max"),
                                                             table: None,
                                                             aggregation: Some(agg_expr),
                                                         }]),
                       ..Default::default()
                   });
    }
}
<commit_msg>Aggregation and group-by test cases.<commit_after>use nom::multispace;
use nom::{IResult, Err, ErrorKind, Needed};
use std::str;

use column::Column;
use common::FieldExpression;
use common::{field_expr, field_list, unsigned_number, statement_terminator, table_list};
use condition::{condition_expr, ConditionExpression};
use table::Table;

#[derive(Clone, Debug, PartialEq)]
pub struct GroupByClause {
    columns: Vec<Column>,
    having: Option<ConditionExpression>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct LimitClause {
    limit: u64,
    offset: u64,
}

#[derive(Clone, Debug, PartialEq)]
pub enum OrderType {
    OrderAscending,
    OrderDescending,
}

#[derive(Clone, Debug, PartialEq)]
pub struct OrderClause {
    cols: Vec<Column>, // TODO(malte): can this be an arbitrary expr?
    order: OrderType,
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct SelectStatement {
    pub tables: Vec<Table>,
    pub distinct: bool,
    pub fields: FieldExpression,
    pub where_clause: Option<ConditionExpression>,
    pub group_by: Option<GroupByClause>,
    pub order: Option<OrderClause>,
    pub limit: Option<LimitClause>,
}

/// Parse GROUP BY clause
named!(group_by_clause<&[u8], GroupByClause>,
    complete!(chain!(
        multispace? ~
        caseless_tag!("group by") ~
        multispace ~
        group_columns: field_list ~
        having_clause: opt!(
            complete!(chain!(
                multispace? ~
                caseless_tag!("having") ~
                multispace? ~
                ce: condition_expr,
                || { ce }
            ))
        ),
    || {
        GroupByClause {
            columns: group_columns,
            having: having_clause,
        }
    }))
);

/// Parse LIMIT clause
named!(limit_clause<&[u8], LimitClause>,
    complete!(chain!(
        multispace? ~
        caseless_tag!("limit") ~
        multispace ~
        limit_val: unsigned_number ~
        offset_val: opt!(
            complete!(chain!(
                multispace? ~
                caseless_tag!("offset") ~
                multispace ~
                val: unsigned_number,
                || { val }
            ))
        ),
    || {
        LimitClause {
            limit: limit_val,
            offset: match offset_val {
                None => 0,
                Some(v) => v,
            },
        }
    }))
);

named!(order_clause<&[u8], OrderClause>,
    complete!(chain!(
        multispace? ~
        caseless_tag!("order by") ~
        multispace ~
        order_expr: field_list ~
        ordering: opt!(
            complete!(chain!(
                multispace? ~
                ordering: alt_complete!(
                      map!(caseless_tag!("desc"), |_| OrderType::OrderDescending)
                    | map!(caseless_tag!("asc"), |_| OrderType::OrderAscending)
                ),
                || { ordering }
            ))
        ),
    || {
        OrderClause {
            cols: order_expr,
            order: match ordering {
                None => OrderType::OrderAscending,
                Some(ref o) => o.clone(),
            },
        }
    }))
);

/// Parse WHERE clause of a selection
named!(where_clause<&[u8], ConditionExpression>,
    complete!(chain!(
        multispace? ~
        caseless_tag!("where") ~
        multispace ~
        cond: condition_expr,
        || { cond }
    ))
);

/// Parse rule for a SQL selection query.
/// TODO(malte): support nested queries as selection targets
named!(pub selection<&[u8], SelectStatement>,
    chain!(
        caseless_tag!("select") ~
        multispace ~
        distinct: opt!(caseless_tag!("distinct")) ~
        multispace? ~
        fields: field_expr ~
        delimited!(opt!(multispace), caseless_tag!("from"), opt!(multispace)) ~
        tables: table_list ~
        cond: opt!(where_clause) ~
        group_by: opt!(group_by_clause) ~
        order: opt!(order_clause) ~
        limit: opt!(limit_clause) ~
        statement_terminator,
        || {
            SelectStatement {
                tables: tables,
                distinct: distinct.is_some(),
                fields: fields,
                where_clause: cond,
                group_by: group_by,
                order: order,
                limit: limit,
            }
        }
    )
);

#[cfg(test)]
mod tests {
    use super::*;
    use column::{AggregationExpression, Column};
    use common::{FieldExpression, Operator};
    use condition::{ConditionBase, ConditionExpression, ConditionTree};
    use table::Table;

    fn columns(cols: &[&str]) -> Vec<Column> {
        cols.iter()
            .map(|c| Column::from(*c))
            .collect()
    }

    #[test]
    fn simple_select() {
        let qstring = "SELECT id, name FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("users")],
                       fields: FieldExpression::Seq(columns(&["id", "name"])),
                       ..Default::default()
                   });
    }

    #[test]
    fn select_all() {
        let qstring = "SELECT * FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("users")],
                       fields: FieldExpression::All,
                       ..Default::default()
                   });
    }

    #[test]
    fn spaces_optional() {
        let qstring = "SELECT id,name FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("users")],
                       fields: FieldExpression::Seq(columns(&["id", "name"])),
                       ..Default::default()
                   });
    }

    #[test]
    fn case_sensitivity() {
        let qstring_lc = "select id, name from users;";
        let qstring_uc = "SELECT id, name FROM users;";

        assert_eq!(selection(qstring_lc.as_bytes()).unwrap(),
                   selection(qstring_uc.as_bytes()).unwrap());
    }

    #[test]
    fn termination() {
        let qstring_sem = "select id, name from users;";
        let qstring_linebreak = "select id, name from users\n";

        let r1 = selection(qstring_sem.as_bytes()).unwrap();
        let r2 = selection(qstring_linebreak.as_bytes()).unwrap();
        assert_eq!(r1, r2);
    }

    #[test]
    fn where_clause() {
        let qstring = "select * from ContactInfo where email=?;";

        let res = selection(qstring.as_bytes());

        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(Column::from("email"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("ContactInfo")],
                       fields: FieldExpression::All,
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn limit_clause() {
        let qstring1 = "select * from users limit 10\n";
        let qstring2 = "select * from users limit 10 offset 10\n";

        let expected_lim1 = LimitClause {
            limit: 10,
            offset: 0,
        };
        let expected_lim2 = LimitClause {
            limit: 10,
            offset: 10,
        };

        let res1 = selection(qstring1.as_bytes());
        let res2 = selection(qstring2.as_bytes());
        assert_eq!(res1.unwrap().1.limit, Some(expected_lim1));
        assert_eq!(res2.unwrap().1.limit, Some(expected_lim2));
    }

    #[test]
    fn order_clause() {
        let qstring1 = "select * from users order by name desc\n";
        let qstring2 = "select * from users order by name, age desc\n";
        let qstring3 = "select * from users order by name\n";

        let expected_ord1 = OrderClause {
            cols: columns(&["name"]),
            order: OrderType::OrderDescending,
        };
        let expected_ord2 = OrderClause {
            cols: columns(&["name", "age"]),
            order: OrderType::OrderDescending,
        };
        let expected_ord3 = OrderClause {
            cols: columns(&["name"]),
            order: OrderType::OrderAscending,
        };

        let res1 = selection(qstring1.as_bytes());
        let res2 = selection(qstring2.as_bytes());
        let res3 = selection(qstring3.as_bytes());
        assert_eq!(res1.unwrap().1.order, Some(expected_ord1));
        assert_eq!(res2.unwrap().1.order, Some(expected_ord2));
        assert_eq!(res3.unwrap().1.order, Some(expected_ord3));
    }

    #[test]
    fn table_alias() {
        let qstring1 = "select * from PaperTag as t;";
        // let qstring2 = "select * from PaperTag t;";

        let res1 = selection(qstring1.as_bytes());
        assert_eq!(res1.clone().unwrap().1,
                   SelectStatement {
                       tables: vec![Table {
                           name: String::from("PaperTag"),
                           alias: Some(String::from("t")),
                       }],
                       fields: FieldExpression::All,
                       ..Default::default()
                   });
        // let res2 = selection(qstring2.as_bytes());
        // assert_eq!(res1.unwrap().1, res2.unwrap().1);
    }

    #[test]
    fn distinct() {
        let qstring = "select distinct tag from PaperTag where paperId=?;";

        let res = selection(qstring.as_bytes());
        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(Column::from("paperId"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("PaperTag")],
                       distinct: true,
                       fields: FieldExpression::Seq(columns(&["tag"])),
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn simple_condition_expr() {
        let qstring = "select infoJson from PaperStorage where paperId=? and paperStorageId=?;";

        let res = selection(qstring.as_bytes());

        let left_comp = Some(Box::new(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(Column::from("paperId"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        })));
        let right_comp = Some(Box::new(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(Column::from("paperStorageId"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        })));
        let expected_where_cond = Some(ConditionExpression::LogicalOp(ConditionTree {
            left: left_comp,
            right: right_comp,
            operator: Operator::And,
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("PaperStorage")],
                       fields: FieldExpression::Seq(columns(&["infoJson"])),
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn where_and_limit_clauses() {
        let qstring = "select * from users where id = ? limit 10\n";
        let res = selection(qstring.as_bytes());

        let expected_lim = Some(LimitClause {
            limit: 10,
            offset: 0,
        });
        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(Column::from("id"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        }));

        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("users")],
                       fields: FieldExpression::All,
                       where_clause: expected_where_cond,
                       limit: expected_lim,
                       ..Default::default()
                   });
    }

    #[test]
    fn aggregation_column() {
        let qstring = "SELECT max(addr_id) FROM address;";

        let res = selection(qstring.as_bytes());
        let agg_expr =
            AggregationExpression::Max(FieldExpression::Seq(vec![Column::from("addr_id")]));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("address")],
                       fields: FieldExpression::Seq(vec![Column {
                                                             name: String::from("anon_agg"),
                                                             table: None,
                                                             aggregation: Some(agg_expr),
                                                         }]),
                       ..Default::default()
                   });
    }

    #[test]
    fn aggregation_column_with_alias() {
        let qstring = "SELECT max(addr_id) AS max_addr FROM address;";

        let res = selection(qstring.as_bytes());
        let agg_expr =
            AggregationExpression::Max(FieldExpression::Seq(vec![Column::from("addr_id")]));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("address")],
                       fields: FieldExpression::Seq(vec![Column {
                                                             name: String::from("max_addr"),
                                                             table: None,
                                                             aggregation: Some(agg_expr),
                                                         }]),
                       ..Default::default()
                   });
    }

    #[test]
    fn count_all() {
        let qstring = "SELECT COUNT(*) FROM votes GROUP BY aid;";

        let res = selection(qstring.as_bytes());
        let agg_expr = AggregationExpression::Count(FieldExpression::All);
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![Table::from("votes")],
                       fields: FieldExpression::Seq(vec![Column {
                                                             name: String::from("anon_agg"),
                                                             table: None,
                                                             aggregation: Some(agg_expr),
                                                         }]),
                       group_by: Some(GroupByClause {
                           columns: vec![Column::from("aid")],
                           having: None,
                       }),
                       ..Default::default()
                   });
    }
}
<|endoftext|>


<commit_before>//! The `source` module provides a representation of `ast::structure::Program`s
//!
//! It defines a trait `Source` that sources a `ast::structure::Program` to a `String`.

use super::structure::{Program,Condition,Command,Expression,Sensor};

/// Representation of an `ast::structure::Program`
pub trait Source {
	fn source(&self) -> String;
}

impl Source for Program {
	fn source(&self) -> String {
		match *self {
			Program::If(ref condition, ref left, ref right) => format!("iff!({},{},{})", condition.source(), left.source(), right.source()),
			Program::Command(ref command)                   => command.source(),
		}
	}
}

impl Source for Condition {
	fn source(&self) -> String {
		match *self {
			Condition::True                              => format!("T!()"),
			Condition::False                             => format!("F!()"),
			Condition::Not(ref condition)                => format!("not!({})", (*condition).source()),
			Condition::Or(ref left, ref right)           => format!("or!({},{})", (*left).source(), (*right).source()),
			Condition::And(ref left, ref right)          => format!("and!({},{})", (*left).source(), (*right).source()),

			Condition::Less(ref left, ref right)         => format!("less!({},{})", (*left).source(), (*right).source()),
			Condition::LessEqual(ref left, ref right)    => format!("less_equal!({},{})", (*left).source(), (*right).source()),
			Condition::Equal(ref left, ref right)        => format!("equal!({},{})", (*left).source(), (*right).source()),
			Condition::GreaterEqual(ref left, ref right) => format!("greater_equal!({},{})", (*left).source(), (*right).source()),
			Condition::Greater(ref left, ref right)      => format!("greater!({},{})", (*left).source(), (*right).source()),
		}
	}
}

impl Source for Expression {
	fn source(&self) -> String {
		match *self {
			Expression::Constant(value)               => format!("constant!({:.*})", 4, value),
			Expression::Sensor(ref sensor)            => sensor.source(),
			Expression::Plus(ref left, ref right)     => format!("plus!({},{})", (*left).source(), (*right).source()),
			Expression::Minus(ref left, ref right)    => format!("minus!({},{})", (*left).source(), (*right).source()),
			Expression::Multiply(ref left, ref right) => format!("multiply!({},{})", (*left).source(), (*right).source()),
			Expression::Divide(ref left, ref right)   => format!("divide!({},{})", (*left).source(), (*right).source()),
		}
	}
}

impl Source for Sensor {
	fn source(&self) -> String {
		match *self {
			Sensor::X  => format!("x!()"),
			Sensor::Y  => format!("u!()"),
			Sensor::Vx => format!("vx!()"),
			Sensor::Vy => format!("vy!()"),
			Sensor::O  => format!("o!()"),
			Sensor::W  => format!("w!()"),
		}
	}
}

impl Source for Command {
	fn source(&self) -> String {
		match *self {
			Command::Skip  => format!("skip!()"),
			Command::Left  => format!("left!()"),
			Command::Right => format!("right!()"),
			Command::Up    => format!("up!()"),
		}
	}
}

#[test]
fn should_correctly_source_programs() {
	assert_eq!("iff!(T!(),skip!(),left!())",
		Program::If(Condition::True, Box::new(Program::Command(Command::Skip)), Box::new(Program::Command(Command::Left))).source());
	assert_eq!("skip!()", Program::Command(Command::Skip).source());
}

#[test]
fn should_correctly_source_conditions() {
	assert_eq!("T!()", Condition::True.source());
	assert_eq!("F!()", Condition::False.source());
	assert_eq!("not!(T!())", Condition::Not(Box::new(Condition::True)).source());
	assert_eq!("or!(T!(),F!())", Condition::Or(Box::new(Condition::True), Box::new(Condition::False)).source());
	assert_eq!("and!(F!(),T!())", Condition::And(Box::new(Condition::False), Box::new(Condition::True)).source());
	assert_eq!("less!(constant!(1.0000),constant!(2.0000))",
		Condition::Less(Box::new(Expression::Constant(1.0)), Box::new(Expression::Constant(2.0))).source());
	assert_eq!("less_equal!(constant!(1.0000),constant!(2.0000))",
		Condition::LessEqual(Box::new(Expression::Constant(1.0)), Box::new(Expression::Constant(2.0))).source());
	assert_eq!("equal!(constant!(1.0000),constant!(2.0000))",
		Condition::Equal(Box::new(Expression::Constant(1.0)), Box::new(Expression::Constant(2.0))).source());
	assert_eq!("greater_equal!(constant!(1.0000),constant!(2.0000))",
		Condition::GreaterEqual(Box::new(Expression::Constant(1.0)), Box::new(Expression::Constant(2.0))).source());
	assert_eq!("greater!(constant!(1.0000),constant!(2.0000))",
		Condition::Greater(Box::new(Expression::Constant(1.0)), Box::new(Expression::Constant(2.0))).source());
}

#[test]
fn should_correctly_source_expressions() {
	assert_eq!("constant!(1.0000)", Expression::Constant(1.0).source());
	assert_eq!("vx!()", Expression::Sensor(Sensor::Vx).source());
	assert_eq!("plus!(constant!(1.0000),vy!())",
		Expression::Plus(Box::new(Expression::Constant(1.0)),Box::new(Expression::Sensor(Sensor::Vy))).source());
	assert_eq!("minus!(constant!(1.0000),constant!(2.0000))",
		Expression::Minus(Box::new(Expression::Constant(1.0)),Box::new(Expression::Constant(2.0))).source());
	assert_eq!("multiply!(constant!(1.0000),constant!(2.0000))",
		Expression::Multiply(Box::new(Expression::Constant(1.0)),Box::new(Expression::Constant(2.0))).source());
	assert_eq!("divide!(constant!(1.0000),constant!(2.0000))",
		Expression::Divide(Box::new(Expression::Constant(1.0)),Box::new(Expression::Constant(2.0))).source());
}

#[test]
fn should_correctly_source_sensors() {
	assert_eq!("x!()",  Sensor::X.source());
	assert_eq!("u!()",  Sensor::Y.source());
	assert_eq!("vx!()", Sensor::Vx.source());
	assert_eq!("vy!()", Sensor::Vy.source());
	assert_eq!("o!()",  Sensor::O.source());
	assert_eq!("w!()",  Sensor::W.source());
}

#[test]
fn should_correctly_source_commands() {
	assert_eq!("skip!()",  Command::Skip.source());
	assert_eq!("left!()",  Command::Left.source());
	assert_eq!("right!()", Command::Right.source());
	assert_eq!("up!()",    Command::Up.source());
}
<commit_msg>Move source tests into tests module<commit_after>//! The `source` module provides a representation of `ast::structure::Program`s
//!
//! It defines a trait `Source` that sources a `ast::structure::Program` to a `String`.

use super::structure::{Program,Condition,Command,Expression,Sensor};

/// Representation of an `ast::structure::Program`
pub trait Source {
	fn source(&self) -> String;
}

impl Source for Program {
	fn source(&self) -> String {
		match *self {
			Program::If(ref condition, ref left, ref right) => format!("iff!({},{},{})", condition.source(), left.source(), right.source()),
			Program::Command(ref command)                   => command.source(),
		}
	}
}

impl Source for Condition {
	fn source(&self) -> String {
		match *self {
			Condition::True                              => format!("T!()"),
			Condition::False                             => format!("F!()"),
			Condition::Not(ref condition)                => format!("not!({})", (*condition).source()),
			Condition::Or(ref left, ref right)           => format!("or!({},{})", (*left).source(), (*right).source()),
			Condition::And(ref left, ref right)          => format!("and!({},{})", (*left).source(), (*right).source()),

			Condition::Less(ref left, ref right)         => format!("less!({},{})", (*left).source(), (*right).source()),
			Condition::LessEqual(ref left, ref right)    => format!("less_equal!({},{})", (*left).source(), (*right).source()),
			Condition::Equal(ref left, ref right)        => format!("equal!({},{})", (*left).source(), (*right).source()),
			Condition::GreaterEqual(ref left, ref right) => format!("greater_equal!({},{})", (*left).source(), (*right).source()),
			Condition::Greater(ref left, ref right)      => format!("greater!({},{})", (*left).source(), (*right).source()),
		}
	}
}

impl Source for Expression {
	fn source(&self) -> String {
		match *self {
			Expression::Constant(value)               => format!("constant!({:.*})", 4, value),
			Expression::Sensor(ref sensor)            => sensor.source(),
			Expression::Plus(ref left, ref right)     => format!("plus!({},{})", (*left).source(), (*right).source()),
			Expression::Minus(ref left, ref right)    => format!("minus!({},{})", (*left).source(), (*right).source()),
			Expression::Multiply(ref left, ref right) => format!("multiply!({},{})", (*left).source(), (*right).source()),
			Expression::Divide(ref left, ref right)   => format!("divide!({},{})", (*left).source(), (*right).source()),
		}
	}
}

impl Source for Sensor {
	fn source(&self) -> String {
		match *self {
			Sensor::X  => format!("x!()"),
			Sensor::Y  => format!("u!()"),
			Sensor::Vx => format!("vx!()"),
			Sensor::Vy => format!("vy!()"),
			Sensor::O  => format!("o!()"),
			Sensor::W  => format!("w!()"),
		}
	}
}

impl Source for Command {
	fn source(&self) -> String {
		match *self {
			Command::Skip  => format!("skip!()"),
			Command::Left  => format!("left!()"),
			Command::Right => format!("right!()"),
			Command::Up    => format!("up!()"),
		}
	}
}

#[cfg(test)]
mod tests {
    use super::*;
    use structure::{Program,Condition,Command,Expression,Sensor};

    #[test]
    fn should_correctly_source_programs() {
	      assert_eq!("iff!(T!(),skip!(),left!())",
		               Program::If(Condition::True, Box::new(Program::Command(Command::Skip)), Box::new(Program::Command(Command::Left))).source());
	      assert_eq!("skip!()", Program::Command(Command::Skip).source());
    }

    #[test]
    fn should_correctly_source_conditions() {
	      assert_eq!("T!()", Condition::True.source());
	      assert_eq!("F!()", Condition::False.source());
	      assert_eq!("not!(T!())", Condition::Not(Box::new(Condition::True)).source());
	      assert_eq!("or!(T!(),F!())", Condition::Or(Box::new(Condition::True), Box::new(Condition::False)).source());
	      assert_eq!("and!(F!(),T!())", Condition::And(Box::new(Condition::False), Box::new(Condition::True)).source());
	      assert_eq!("less!(constant!(1.0000),constant!(2.0000))",
		               Condition::Less(Box::new(Expression::Constant(1.0)), Box::new(Expression::Constant(2.0))).source());
	      assert_eq!("less_equal!(constant!(1.0000),constant!(2.0000))",
		               Condition::LessEqual(Box::new(Expression::Constant(1.0)), Box::new(Expression::Constant(2.0))).source());
	      assert_eq!("equal!(constant!(1.0000),constant!(2.0000))",
		               Condition::Equal(Box::new(Expression::Constant(1.0)), Box::new(Expression::Constant(2.0))).source());
	      assert_eq!("greater_equal!(constant!(1.0000),constant!(2.0000))",
		               Condition::GreaterEqual(Box::new(Expression::Constant(1.0)), Box::new(Expression::Constant(2.0))).source());
	      assert_eq!("greater!(constant!(1.0000),constant!(2.0000))",
		               Condition::Greater(Box::new(Expression::Constant(1.0)), Box::new(Expression::Constant(2.0))).source());
    }

    #[test]
    fn should_correctly_source_expressions() {
	      assert_eq!("constant!(1.0000)", Expression::Constant(1.0).source());
	      assert_eq!("vx!()", Expression::Sensor(Sensor::Vx).source());
	      assert_eq!("plus!(constant!(1.0000),vy!())",
		               Expression::Plus(Box::new(Expression::Constant(1.0)),Box::new(Expression::Sensor(Sensor::Vy))).source());
	      assert_eq!("minus!(constant!(1.0000),constant!(2.0000))",
		               Expression::Minus(Box::new(Expression::Constant(1.0)),Box::new(Expression::Constant(2.0))).source());
	      assert_eq!("multiply!(constant!(1.0000),constant!(2.0000))",
		               Expression::Multiply(Box::new(Expression::Constant(1.0)),Box::new(Expression::Constant(2.0))).source());
	      assert_eq!("divide!(constant!(1.0000),constant!(2.0000))",
		               Expression::Divide(Box::new(Expression::Constant(1.0)),Box::new(Expression::Constant(2.0))).source());
    }

    #[test]
    fn should_correctly_source_sensors() {
	      assert_eq!("x!()",  Sensor::X.source());
	      assert_eq!("u!()",  Sensor::Y.source());
	      assert_eq!("vx!()", Sensor::Vx.source());
	      assert_eq!("vy!()", Sensor::Vy.source());
	      assert_eq!("o!()",  Sensor::O.source());
	      assert_eq!("w!()",  Sensor::W.source());
    }

    #[test]
    fn should_correctly_source_commands() {
	      assert_eq!("skip!()",  Command::Skip.source());
	      assert_eq!("left!()",  Command::Left.source());
	      assert_eq!("right!()", Command::Right.source());
	      assert_eq!("up!()",    Command::Up.source());
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::mem;
use std::sync::{Arc, mpsc};
use std::sync::atomic::{AtomicBool, Ordering};

use mio::{Token, EventSet, TryRead, TryWrite};
use ::{Action, Queued, Protocol, Transport};

pub type Queued_ = Queued;

pub struct Stream<P: Protocol, T: Transport> {
    transport: T,
    protocol: P,
    rx: mpsc::Receiver<Queued>,
    is_queued: Arc<AtomicBool>,
    last_action: Option<Action>,
    reading: Reading,
    writing: Writing,
}

impl<P: Protocol, T: Transport> Stream<P, T> {

    pub fn new(transport: T, protocol: P, rx: mpsc::Receiver<Queued_>, is_queued: Arc<AtomicBool>) -> Stream<P, T> {
        Stream {
            transport: transport,
            protocol: protocol,
            rx: rx,
            is_queued: is_queued,
            last_action: None,
            reading: Reading::Open(vec![0; 4096]),
            writing: Writing::Waiting(io::Cursor::new(vec![])),
        }
    }

    fn queue_writing(&mut self, data: Option<Vec<u8>>) {
        match (self.writing.can_write(), data) {
            (true, Some(bytes)) => {
                trace!("queue writing {} bytes", bytes.len());
                let mut buf = self.writing.close();
                buf.get_mut().extend(&bytes);
                self.writing.open(buf);
            },
            (true, None) => {
                trace!("queueing closing");
                self.writing.closing();
            },
            (false, data) => {
                trace!("cannot queue writing: {:?}", data);
            },
        }
    }

    fn pause(&mut self) {
        trace!("pause");
        self.reading.pause();
    }

    fn resume(&mut self) {
        trace!("resume");
        self.reading.resume();
    }

    pub fn queue(&mut self) {
        while let Ok(q) = self.rx.try_recv() {
            match q {
                Queued::Write(data) => self.queue_writing(data),
                Queued::Pause => self.pause(),
                Queued::Resume => self.resume(),
            }
        }
        self.is_queued.store(false, Ordering::Release);
    }

    pub fn close(&mut self) {
        self.queue_writing(None);
        self.reading = Reading::Closed;
    }

    fn abort(&mut self) {
        self.reading = Reading::Closed;
        self.writing = Writing::Closed;
        self.protocol.on_end(None);
    }

    pub fn ready(&mut self, token: Token, events: EventSet) {
        self.last_action = None;
        self.queue();
        if events.is_error() {
            debug!("error event on {:?}", token);
        }

        if events.is_readable() {
            loop {
                //TODO: check if protocol paused
                if self.reading.is_paused() {
                    trace!("reading is paused");
                    break;
                }
                let mut buf = self.reading.close();
                match self.transport.try_read(&mut buf) {
                    Ok(Some(0)) => {
                        trace!("read eof {:?}", token);
                        self.protocol.on_eof();
                        self.abort();
                        break;
                    }
                    Ok(Some(n)) => {
                        trace!("read {} bytes {:?}", n, token);
                        self.protocol.on_data(&buf[..n]);
                        self.reading.open(buf);
                    }
                    Ok(None) => {
                        trace!("read would block {:?}", token);
                        self.reading.open(buf);
                        break;
                    }
                    Err(e) => {
                        self.on_error(e.into());
                        return;
                    }
                }
                self.queue();
            }
        }

        if events.is_writable() {
            let is_closing = !self.writing.can_write();
            let mut buf = self.writing.close();
            while buf.position() < buf.get_ref().len() as u64 {
                match self.transport.try_write_buf(&mut buf) {
                    Ok(Some(0)) => {
                        error!("===> write 0 means what again? {:?}", token);
                        break;
                    },
                    Ok(Some(n)) => {
                        trace!("wrote {} bytes {:?}", n, token);
                    },
                    Ok(None) => {
                        trace!("write would block {:?}", token);
                        break;
                    }
                    Err(e) => {
                        self.on_error(e.into());
                        return;
                    }
                }
            }

            if buf.position() == buf.get_ref().len() as u64 {
                buf.set_position(0);
                buf.get_mut().truncate(0);
                if is_closing {
                    //leave closed
                } else {
                    self.writing.open(buf);
                    self.writing.wait();
                }
            } else {
                self.writing.open(buf);
                if is_closing {
                    self.writing.closing();
                }
            }
        }
    }

    pub fn transport(&self) -> &T {
        &self.transport
    }

    pub fn action(&mut self) -> Action {
        let action = match (&self.reading, &self.writing) {
            (&Reading::Open(_), &Writing::Open(_)) |(&Reading::Open(_), &Writing::Closing(_)) => Action::Register(EventSet::readable() | EventSet::writable()),
            (&Reading::Closed, &Writing::Closed) => Action::Remove,

            (&Reading::Open(_), _) => Action::Register(EventSet::readable()),
            (_, &Writing::Open(_)) | (_, &Writing::Closing(_)) => Action::Register(EventSet::writable()),

            _ => Action::Wait
        };

        if let Some(ref last) = self.last_action {
            if *last == action {
                trace!("< Action is the same {:?}, Waiting", action);
                return Action::Wait
            }
        }
        self.last_action = Some(action.clone());
        action
    }

    fn on_error(&mut self, e: ::Error) {
        self.writing = Writing::Closed;
        self.reading = Reading::Closed;
        self.protocol.on_end(Some(e));
    }
}

enum Reading {
    Open(Vec<u8>),
    Paused(Vec<u8>),
    Closed
}

impl Reading {
    fn open(&mut self, buf: Vec<u8>) {
        mem::replace(self, Reading::Open(buf));
    }

    fn resume(&mut self) {
        let buf = self.close();
        mem::replace(self, Reading::Open(buf));
    }

    fn pause(&mut self) {
        let buf = self.close();
        mem::replace(self, Reading::Paused(buf));
    }

    fn is_paused(&self) -> bool {
        match *self {
            Reading::Paused(..) => true,
            _ => false
        }
    }

    fn close(&mut self) -> Vec<u8> {
        match mem::replace(self, Reading::Closed) {
            Reading::Open(buf) => buf,
            Reading::Paused(buf) => buf,
            _ => panic!("already closed")
        }
    }
}

enum Writing {
    Open(io::Cursor<Vec<u8>>),
    Waiting(io::Cursor<Vec<u8>>),
    Closing(io::Cursor<Vec<u8>>),
    Closed
}

impl Writing {
    fn open(&mut self, buf: io::Cursor<Vec<u8>>) {
        mem::replace(self, Writing::Open(buf));
    }

    fn wait(&mut self) {
        let buf = self.close();
        mem::replace(self, Writing::Waiting(buf));
    }

    fn closing(&mut self) {
        let buf = self.close();
        mem::replace(self, Writing::Closing(buf));
    }

    fn can_write(&self) -> bool {
        match *self {
            Writing::Open(_) | Writing::Waiting(_) => true,
            _ => false
        }
    }

    fn close(&mut self) -> io::Cursor<Vec<u8>> {
        match mem::replace(self, Writing::Closed) {
            Writing::Open(buf) => buf,
            Writing::Waiting(buf) => buf,
            Writing::Closing(buf) => buf,
            _ => panic!("already closed")
        }
    }
}

<commit_msg>todo: take_socket_error<commit_after>use std::io;
use std::mem;
use std::sync::{Arc, mpsc};
use std::sync::atomic::{AtomicBool, Ordering};

use mio::{Token, EventSet, TryRead, TryWrite};
use ::{Action, Queued, Protocol, Transport};

pub type Queued_ = Queued;

pub struct Stream<P: Protocol, T: Transport> {
    transport: T,
    protocol: P,
    rx: mpsc::Receiver<Queued>,
    is_queued: Arc<AtomicBool>,
    last_action: Option<Action>,
    reading: Reading,
    writing: Writing,
}

impl<P: Protocol, T: Transport> Stream<P, T> {

    pub fn new(transport: T, protocol: P, rx: mpsc::Receiver<Queued_>, is_queued: Arc<AtomicBool>) -> Stream<P, T> {
        Stream {
            transport: transport,
            protocol: protocol,
            rx: rx,
            is_queued: is_queued,
            last_action: None,
            reading: Reading::Open(vec![0; 4096]),
            writing: Writing::Waiting(io::Cursor::new(vec![])),
        }
    }

    fn queue_writing(&mut self, data: Option<Vec<u8>>) {
        match (self.writing.can_write(), data) {
            (true, Some(bytes)) => {
                trace!("queue writing {} bytes", bytes.len());
                let mut buf = self.writing.close();
                buf.get_mut().extend(&bytes);
                self.writing.open(buf);
            },
            (true, None) => {
                trace!("queueing closing");
                self.writing.closing();
            },
            (false, data) => {
                trace!("cannot queue writing: {:?}", data);
            },
        }
    }

    fn pause(&mut self) {
        trace!("pause");
        self.reading.pause();
    }

    fn resume(&mut self) {
        trace!("resume");
        self.reading.resume();
    }

    pub fn queue(&mut self) {
        while let Ok(q) = self.rx.try_recv() {
            match q {
                Queued::Write(data) => self.queue_writing(data),
                Queued::Pause => self.pause(),
                Queued::Resume => self.resume(),
            }
        }
        self.is_queued.store(false, Ordering::Release);
    }

    pub fn close(&mut self) {
        self.queue_writing(None);
        self.reading = Reading::Closed;
    }

    fn abort(&mut self) {
        self.reading = Reading::Closed;
        self.writing = Writing::Closed;
        self.protocol.on_end(None);
    }

    pub fn ready(&mut self, token: Token, events: EventSet) {
        self.last_action = None;
        self.queue();
        if events.is_error() {
            debug!("error event on {:?}", token);
            //TODO: try to use self.transport.take_socket_error()
            //requires adding a method to the Transport trait
        }

        if events.is_readable() {
            loop {
                //TODO: check if protocol paused
                if self.reading.is_paused() {
                    trace!("reading is paused");
                    break;
                }
                let mut buf = self.reading.close();
                match self.transport.try_read(&mut buf) {
                    Ok(Some(0)) => {
                        trace!("read eof {:?}", token);
                        self.protocol.on_eof();
                        self.abort();
                        break;
                    }
                    Ok(Some(n)) => {
                        trace!("read {} bytes {:?}", n, token);
                        self.protocol.on_data(&buf[..n]);
                        self.reading.open(buf);
                    }
                    Ok(None) => {
                        trace!("read would block {:?}", token);
                        self.reading.open(buf);
                        break;
                    }
                    Err(e) => {
                        self.on_error(e.into());
                        return;
                    }
                }
                self.queue();
            }
        }

        if events.is_writable() {
            let is_closing = !self.writing.can_write();
            let mut buf = self.writing.close();
            while buf.position() < buf.get_ref().len() as u64 {
                match self.transport.try_write_buf(&mut buf) {
                    Ok(Some(0)) => {
                        error!("===> write 0 means what again? {:?}", token);
                        break;
                    },
                    Ok(Some(n)) => {
                        trace!("wrote {} bytes {:?}", n, token);
                    },
                    Ok(None) => {
                        trace!("write would block {:?}", token);
                        break;
                    }
                    Err(e) => {
                        self.on_error(e.into());
                        return;
                    }
                }
            }

            if buf.position() == buf.get_ref().len() as u64 {
                buf.set_position(0);
                buf.get_mut().truncate(0);
                if is_closing {
                    //leave closed
                } else {
                    self.writing.open(buf);
                    self.writing.wait();
                }
            } else {
                self.writing.open(buf);
                if is_closing {
                    self.writing.closing();
                }
            }
        }
    }

    pub fn transport(&self) -> &T {
        &self.transport
    }

    pub fn action(&mut self) -> Action {
        let action = match (&self.reading, &self.writing) {
            (&Reading::Open(_), &Writing::Open(_)) |(&Reading::Open(_), &Writing::Closing(_)) => Action::Register(EventSet::readable() | EventSet::writable()),
            (&Reading::Closed, &Writing::Closed) => Action::Remove,

            (&Reading::Open(_), _) => Action::Register(EventSet::readable()),
            (_, &Writing::Open(_)) | (_, &Writing::Closing(_)) => Action::Register(EventSet::writable()),

            _ => Action::Wait
        };

        if let Some(ref last) = self.last_action {
            if *last == action {
                trace!("< Action is the same {:?}, Waiting", action);
                return Action::Wait
            }
        }
        self.last_action = Some(action.clone());
        action
    }

    fn on_error(&mut self, e: ::Error) {
        self.writing = Writing::Closed;
        self.reading = Reading::Closed;
        self.protocol.on_end(Some(e));
    }
}

enum Reading {
    Open(Vec<u8>),
    Paused(Vec<u8>),
    Closed
}

impl Reading {
    fn open(&mut self, buf: Vec<u8>) {
        mem::replace(self, Reading::Open(buf));
    }

    fn resume(&mut self) {
        let buf = self.close();
        mem::replace(self, Reading::Open(buf));
    }

    fn pause(&mut self) {
        let buf = self.close();
        mem::replace(self, Reading::Paused(buf));
    }

    fn is_paused(&self) -> bool {
        match *self {
            Reading::Paused(..) => true,
            _ => false
        }
    }

    fn close(&mut self) -> Vec<u8> {
        match mem::replace(self, Reading::Closed) {
            Reading::Open(buf) => buf,
            Reading::Paused(buf) => buf,
            _ => panic!("already closed")
        }
    }
}

enum Writing {
    Open(io::Cursor<Vec<u8>>),
    Waiting(io::Cursor<Vec<u8>>),
    Closing(io::Cursor<Vec<u8>>),
    Closed
}

impl Writing {
    fn open(&mut self, buf: io::Cursor<Vec<u8>>) {
        mem::replace(self, Writing::Open(buf));
    }

    fn wait(&mut self) {
        let buf = self.close();
        mem::replace(self, Writing::Waiting(buf));
    }

    fn closing(&mut self) {
        let buf = self.close();
        mem::replace(self, Writing::Closing(buf));
    }

    fn can_write(&self) -> bool {
        match *self {
            Writing::Open(_) | Writing::Waiting(_) => true,
            _ => false
        }
    }

    fn close(&mut self) -> io::Cursor<Vec<u8>> {
        match mem::replace(self, Writing::Closed) {
            Writing::Open(buf) => buf,
            Writing::Waiting(buf) => buf,
            Writing::Closing(buf) => buf,
            _ => panic!("already closed")
        }
    }
}

<|endoftext|>


<commit_before>//! A byte-offset based string table.
//! Commonly used in ELF binaries, Unix archives, and even PE binaries.

use core::fmt;
use core::ops::Index;
use core::str;
use scroll::{ctx, Pread};
if_alloc! {
    use crate::error;
    use alloc::vec::Vec;
}

/// A common string table format which is indexed by byte offsets (and not
/// member index). Constructed using [`parse`](#method.parse)
/// with your choice of delimiter. Please be careful.
pub struct Strtab<'a> {
    delim: ctx::StrCtx,
    bytes: &'a [u8],
    #[cfg(feature = "alloc")]
    strings: Vec<(usize, &'a str)>,
}

#[inline(always)]
fn get_str(offset: usize, bytes: &[u8], delim: ctx::StrCtx) -> scroll::Result<&str> {
    bytes.pread_with::<&str>(offset, delim)
}

impl<'a> Strtab<'a> {
    /// Creates a `Strtab` with `bytes` as the backing string table, using `delim` as the delimiter between entries.
    ///
    /// NB: this does *not* preparse the string table, which can have non-optimal access patterns.
    /// See https://github.com/m4b/goblin/pull/275#issue-660364025
    pub fn new(bytes: &'a [u8], delim: u8) -> Self {
        Self::from_slice_unparsed(bytes, 0, bytes.len(), delim)
    }

    /// Creates a `Strtab` directly without bounds check and without parsing it.
    ///
    /// This is potentially unsafe and should only be used if `feature = "alloc"` is disabled.
    pub fn from_slice_unparsed(bytes: &'a [u8], offset: usize, len: usize, delim: u8) -> Self {
        Self {
            delim: ctx::StrCtx::Delimiter(delim),
            bytes: &bytes[offset..offset + len],
            #[cfg(feature = "alloc")]
            strings: Vec::new(),
        }
    }
    /// Gets a str reference from the backing bytes starting at byte `offset`.
    ///
    /// If the index is out of bounds, `None` is returned. Panics if bytes are invalid UTF-8.
    /// Use this method if the `Strtab` was created using `from_slice_unparsed()`.
    pub fn get_unsafe(&self, offset: usize) -> Option<&'a str> {
        if offset >= self.bytes.len() {
            None
        } else {
            Some(get_str(offset, self.bytes, self.delim).unwrap())
        }
    }
    #[cfg(feature = "alloc")]
    /// Parses a `Strtab` from `bytes` at `offset` with `len` size as the backing string table, using `delim` as the delimiter.
    ///
    /// Errors if bytes are invalid UTF-8.
    /// Requires `feature = "alloc"`
    pub fn parse(bytes: &'a [u8], offset: usize, len: usize, delim: u8) -> error::Result<Self> {
        let (end, overflow) = offset.overflowing_add(len);
        if overflow || end > bytes.len() {
            return Err(error::Error::Malformed(format!(
                "Strtable size ({}) + offset ({}) is out of bounds for {} #bytes. Overflowed: {}",
                len,
                offset,
                bytes.len(),
                overflow
            )));
        }
        let mut result = Self::from_slice_unparsed(bytes, offset, len, delim);
        let mut i = 0;
        while i < result.bytes.len() {
            let string = get_str(i, result.bytes, result.delim)?;
            result.strings.push((i, string));
            i += string.len() + 1;
        }
        Ok(result)
    }
    #[cfg(feature = "alloc")]
    /// Parses a `Strtab` with `bytes` as the backing string table, using `delim` as the delimiter between entries.
    ///
    /// Requires `feature = "alloc"`
    pub fn new_preparsed(bytes: &'a [u8], delim: u8) -> error::Result<Self> {
        Self::parse(bytes, 0, bytes.len(), delim)
    }
    #[cfg(feature = "alloc")]
    /// Converts the string table to a vector of parsed strings.
    ///
    /// Requires `feature = "alloc"`
    pub fn to_vec(&self) -> error::Result<Vec<&'a str>> {
        // Fallback in case `Strtab` was created using `from_slice_unparsed()`.
        if self.strings.is_empty() {
            let mut result = Vec::new();
            let mut i = 0;
            while i < self.bytes.len() {
                let string = get_str(i, self.bytes, self.delim)?;
                result.push(string);
                i += string.len() + 1;
            }
            return Ok(result);
        }
        Ok(self.strings.iter().map(|&(_key, value)| value).collect())
    }
    #[cfg(feature = "alloc")]
    /// Safely gets a str reference from the parsed table starting at byte `offset`.
    ///
    /// If the index is out of bounds, `None` is returned.
    /// Requires `feature = "alloc"`
    pub fn get_at(&self, offset: usize) -> Option<&'a str> {
        match self
            .strings
            .binary_search_by_key(&offset, |&(key, _value)| key)
        {
            Ok(index) => Some(self.strings[index].1),
            Err(index) => {
                if index == 0 {
                    return None;
                }
                let (string_begin_offset, entire_string) = self.strings[index - 1];
                entire_string.get(offset - string_begin_offset..)
            }
        }
    }
    #[deprecated(since = "0.4.2", note = "Use from_slice_unparsed() instead")]
    /// Construct a strtab from a `ptr`, and a `size`, using `delim` as the delimiter
    ///
    /// # Safety
    /// This function creates a `Strtab` directly from a raw pointer and size
    pub unsafe fn from_raw(ptr: *const u8, len: usize, delim: u8) -> Strtab<'a> {
        Self::from_slice_unparsed(core::slice::from_raw_parts(ptr, len), 0, len, delim)
    }
    #[deprecated(since = "0.4.2", note = "Bad performance, use get_at() instead")]
    #[cfg(feature = "alloc")]
    /// Parses a str reference from the parsed table starting at byte `offset`.
    ///
    /// If the index is out of bounds, `None` is returned.
    /// Requires `feature = "alloc"`
    pub fn get(&self, offset: usize) -> Option<error::Result<&'a str>> {
        if offset >= self.bytes.len() {
            None
        } else {
            Some(get_str(offset, self.bytes, self.delim).map_err(core::convert::Into::into))
        }
    }
}

impl<'a> fmt::Debug for Strtab<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Strtab")
            .field("delim", &self.delim)
            .field("bytes", &str::from_utf8(self.bytes))
            .finish()
    }
}

impl<'a> Default for Strtab<'a> {
    fn default() -> Self {
        Self {
            delim: ctx::StrCtx::default(),
            bytes: &[],
            #[cfg(feature = "alloc")]
            strings: Vec::new(),
        }
    }
}

impl<'a> Index<usize> for Strtab<'a> {
    type Output = str;
    /// Gets str reference at starting at byte `offset`.
    /// **NB**: this will panic if the underlying bytes are not valid utf8, or the offset is invalid
    #[inline(always)]
    fn index(&self, offset: usize) -> &Self::Output {
        // This can't delegate to get() because get() requires #[cfg(features = "alloc")]
        // It's also slightly less useful than get() because the lifetime -- specified by the Index
        // trait -- matches &self, even though we could return &'a instead
        get_str(offset, self.bytes, self.delim).unwrap()
    }
}

#[test]
fn as_vec_no_final_null() {
    let strtab = Strtab::new_preparsed(b"\0printf\0memmove\0busta", 0x0).unwrap();
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 4);
    assert_eq!(vec, vec!["", "printf", "memmove", "busta"]);
}

#[test]
fn as_vec_no_first_null_no_final_null() {
    let strtab = Strtab::new_preparsed(b"printf\0memmove\0busta", 0x0).unwrap();
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 3);
    assert_eq!(vec, vec!["printf", "memmove", "busta"]);
}

#[test]
fn to_vec_final_null() {
    let strtab = Strtab::new_preparsed(b"\0printf\0memmove\0busta\0", 0x0).unwrap();
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 4);
    assert_eq!(vec, vec!["", "printf", "memmove", "busta"]);
}

#[test]
fn to_vec_newline_delim() {
    let strtab = Strtab::new_preparsed(b"\nprintf\nmemmove\nbusta\n", b'\n').unwrap();
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 4);
    assert_eq!(vec, vec!["", "printf", "memmove", "busta"]);
}

#[test]
fn parse_utf8() {
    assert!(match Strtab::new_preparsed(&[0x80, 0x80], b'\n') {
        Err(error::Error::Scroll(scroll::Error::BadInput {
            size: 2,
            msg: "invalid utf8",
        })) => true,
        _ => false,
    });
    assert!(
        match Strtab::new_preparsed(&[0xC6, 0x92, 0x6F, 0x6F], b'\n') {
            Ok(_) => true,
            _ => false,
        }
    );
}

#[test]
fn get_at_utf8() {
    let strtab = Strtab::new_preparsed("\nƒoo\nmemmove\n🅱usta\n".as_bytes(), b'\n').unwrap();
    assert_eq!(strtab.get_at(0), Some(""));
    assert_eq!(strtab.get_at(5), Some(""));
    assert_eq!(strtab.get_at(6), Some("memmove"));
    assert_eq!(strtab.get_at(14), Some("\u{1f171}\u{fe0f}usta"));
    assert_eq!(strtab.get_at(16), None);
    assert_eq!(strtab.get_at(18), Some("\u{fe0f}usta"));
    assert_eq!(strtab.get_at(21), Some("usta"));
    assert_eq!(strtab.get_at(25), Some(""));
    assert_eq!(strtab.get_at(26), None);
}
<commit_msg>strtab: document to_vec behavior; give example<commit_after>//! A byte-offset based string table.
//! Commonly used in ELF binaries, Unix archives, and even PE binaries.

use core::fmt;
use core::ops::Index;
use core::str;
use scroll::{ctx, Pread};
if_alloc! {
    use crate::error;
    use alloc::vec::Vec;
}

/// A common string table format which is indexed by byte offsets (and not
/// member index). Constructed using [`parse`](#method.parse)
/// with your choice of delimiter. Please be careful.
pub struct Strtab<'a> {
    delim: ctx::StrCtx,
    bytes: &'a [u8],
    #[cfg(feature = "alloc")]
    strings: Vec<(usize, &'a str)>,
}

#[inline(always)]
fn get_str(offset: usize, bytes: &[u8], delim: ctx::StrCtx) -> scroll::Result<&str> {
    bytes.pread_with::<&str>(offset, delim)
}

impl<'a> Strtab<'a> {
    /// Creates a `Strtab` with `bytes` as the backing string table, using `delim` as the delimiter between entries.
    ///
    /// NB: this does *not* preparse the string table, which can have non-optimal access patterns.
    /// See https://github.com/m4b/goblin/pull/275#issue-660364025
    pub fn new(bytes: &'a [u8], delim: u8) -> Self {
        Self::from_slice_unparsed(bytes, 0, bytes.len(), delim)
    }

    /// Creates a `Strtab` directly without bounds check and without parsing it.
    ///
    /// This is potentially unsafe and should only be used if `feature = "alloc"` is disabled.
    pub fn from_slice_unparsed(bytes: &'a [u8], offset: usize, len: usize, delim: u8) -> Self {
        Self {
            delim: ctx::StrCtx::Delimiter(delim),
            bytes: &bytes[offset..offset + len],
            #[cfg(feature = "alloc")]
            strings: Vec::new(),
        }
    }
    /// Gets a str reference from the backing bytes starting at byte `offset`.
    ///
    /// If the index is out of bounds, `None` is returned. Panics if bytes are invalid UTF-8.
    /// Use this method if the `Strtab` was created using `from_slice_unparsed()`.
    pub fn get_unsafe(&self, offset: usize) -> Option<&'a str> {
        if offset >= self.bytes.len() {
            None
        } else {
            Some(get_str(offset, self.bytes, self.delim).unwrap())
        }
    }
    #[cfg(feature = "alloc")]
    /// Parses a `Strtab` from `bytes` at `offset` with `len` size as the backing string table, using `delim` as the delimiter.
    ///
    /// Errors if bytes are invalid UTF-8.
    /// Requires `feature = "alloc"`
    pub fn parse(bytes: &'a [u8], offset: usize, len: usize, delim: u8) -> error::Result<Self> {
        let (end, overflow) = offset.overflowing_add(len);
        if overflow || end > bytes.len() {
            return Err(error::Error::Malformed(format!(
                "Strtable size ({}) + offset ({}) is out of bounds for {} #bytes. Overflowed: {}",
                len,
                offset,
                bytes.len(),
                overflow
            )));
        }
        let mut result = Self::from_slice_unparsed(bytes, offset, len, delim);
        let mut i = 0;
        while i < result.bytes.len() {
            let string = get_str(i, result.bytes, result.delim)?;
            result.strings.push((i, string));
            i += string.len() + 1;
        }
        Ok(result)
    }
    #[cfg(feature = "alloc")]
    /// Parses a `Strtab` with `bytes` as the backing string table, using `delim` as the delimiter between entries.
    ///
    /// Requires `feature = "alloc"`
    pub fn new_preparsed(bytes: &'a [u8], delim: u8) -> error::Result<Self> {
        Self::parse(bytes, 0, bytes.len(), delim)
    }
    #[cfg(feature = "alloc")]
    /// Converts the string table to a vector of parsed strings.
    ///
    /// Note: This method is used to check the parsed contents of `strtab`.
    /// If you want to get the correct contents of `strtab` as `Vec`, use the following example.
    ///
    /// # Examples
    /// ```rust
    /// use goblin::error::Error;
    ///
    /// pub fn show_shdr_strtab(bytes: &[u8]) -> Result<(), Error> {
    ///     let elf = goblin::elf::Elf::parse(&bytes)?;
    ///
    ///     for section in elf.section_headers {
    ///         println!("{}", elf.shdr_strtab.get_at(section.sh_name).unwrap_or(""));
    ///     }
    ///
    ///     Ok(())
    /// }
    /// ```
    ///
    /// Requires `feature = "alloc"`
    pub fn to_vec(&self) -> error::Result<Vec<&'a str>> {
        // Fallback in case `Strtab` was created using `from_slice_unparsed()`.
        if self.strings.is_empty() {
            let mut result = Vec::new();
            let mut i = 0;
            while i < self.bytes.len() {
                let string = get_str(i, self.bytes, self.delim)?;
                result.push(string);
                i += string.len() + 1;
            }
            return Ok(result);
        }
        Ok(self.strings.iter().map(|&(_key, value)| value).collect())
    }
    #[cfg(feature = "alloc")]
    /// Safely gets a str reference from the parsed table starting at byte `offset`.
    ///
    /// If the index is out of bounds, `None` is returned.
    /// Requires `feature = "alloc"`
    pub fn get_at(&self, offset: usize) -> Option<&'a str> {
        match self
            .strings
            .binary_search_by_key(&offset, |&(key, _value)| key)
        {
            Ok(index) => Some(self.strings[index].1),
            Err(index) => {
                if index == 0 {
                    return None;
                }
                let (string_begin_offset, entire_string) = self.strings[index - 1];
                entire_string.get(offset - string_begin_offset..)
            }
        }
    }
    #[deprecated(since = "0.4.2", note = "Use from_slice_unparsed() instead")]
    /// Construct a strtab from a `ptr`, and a `size`, using `delim` as the delimiter
    ///
    /// # Safety
    /// This function creates a `Strtab` directly from a raw pointer and size
    pub unsafe fn from_raw(ptr: *const u8, len: usize, delim: u8) -> Strtab<'a> {
        Self::from_slice_unparsed(core::slice::from_raw_parts(ptr, len), 0, len, delim)
    }
    #[deprecated(since = "0.4.2", note = "Bad performance, use get_at() instead")]
    #[cfg(feature = "alloc")]
    /// Parses a str reference from the parsed table starting at byte `offset`.
    ///
    /// If the index is out of bounds, `None` is returned.
    /// Requires `feature = "alloc"`
    pub fn get(&self, offset: usize) -> Option<error::Result<&'a str>> {
        if offset >= self.bytes.len() {
            None
        } else {
            Some(get_str(offset, self.bytes, self.delim).map_err(core::convert::Into::into))
        }
    }
}

impl<'a> fmt::Debug for Strtab<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Strtab")
            .field("delim", &self.delim)
            .field("bytes", &str::from_utf8(self.bytes))
            .finish()
    }
}

impl<'a> Default for Strtab<'a> {
    fn default() -> Self {
        Self {
            delim: ctx::StrCtx::default(),
            bytes: &[],
            #[cfg(feature = "alloc")]
            strings: Vec::new(),
        }
    }
}

impl<'a> Index<usize> for Strtab<'a> {
    type Output = str;
    /// Gets str reference at starting at byte `offset`.
    /// **NB**: this will panic if the underlying bytes are not valid utf8, or the offset is invalid
    #[inline(always)]
    fn index(&self, offset: usize) -> &Self::Output {
        // This can't delegate to get() because get() requires #[cfg(features = "alloc")]
        // It's also slightly less useful than get() because the lifetime -- specified by the Index
        // trait -- matches &self, even though we could return &'a instead
        get_str(offset, self.bytes, self.delim).unwrap()
    }
}

#[test]
fn as_vec_no_final_null() {
    let strtab = Strtab::new_preparsed(b"\0printf\0memmove\0busta", 0x0).unwrap();
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 4);
    assert_eq!(vec, vec!["", "printf", "memmove", "busta"]);
}

#[test]
fn as_vec_no_first_null_no_final_null() {
    let strtab = Strtab::new_preparsed(b"printf\0memmove\0busta", 0x0).unwrap();
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 3);
    assert_eq!(vec, vec!["printf", "memmove", "busta"]);
}

#[test]
fn to_vec_final_null() {
    let strtab = Strtab::new_preparsed(b"\0printf\0memmove\0busta\0", 0x0).unwrap();
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 4);
    assert_eq!(vec, vec!["", "printf", "memmove", "busta"]);
}

#[test]
fn to_vec_newline_delim() {
    let strtab = Strtab::new_preparsed(b"\nprintf\nmemmove\nbusta\n", b'\n').unwrap();
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 4);
    assert_eq!(vec, vec!["", "printf", "memmove", "busta"]);
}

#[test]
fn parse_utf8() {
    assert!(match Strtab::new_preparsed(&[0x80, 0x80], b'\n') {
        Err(error::Error::Scroll(scroll::Error::BadInput {
            size: 2,
            msg: "invalid utf8",
        })) => true,
        _ => false,
    });
    assert!(
        match Strtab::new_preparsed(&[0xC6, 0x92, 0x6F, 0x6F], b'\n') {
            Ok(_) => true,
            _ => false,
        }
    );
}

#[test]
fn get_at_utf8() {
    let strtab = Strtab::new_preparsed("\nƒoo\nmemmove\n🅱usta\n".as_bytes(), b'\n').unwrap();
    assert_eq!(strtab.get_at(0), Some(""));
    assert_eq!(strtab.get_at(5), Some(""));
    assert_eq!(strtab.get_at(6), Some("memmove"));
    assert_eq!(strtab.get_at(14), Some("\u{1f171}\u{fe0f}usta"));
    assert_eq!(strtab.get_at(16), None);
    assert_eq!(strtab.get_at(18), Some("\u{fe0f}usta"));
    assert_eq!(strtab.get_at(21), Some("usta"));
    assert_eq!(strtab.get_at(25), Some(""));
    assert_eq!(strtab.get_at(26), None);
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use model::{Color, GameResult, Model, Move};
use view::Event;

pub fn update(model: &mut Model, event: Option<Event>) {
    if let Some(event) = event {
        use view::Event::*;

        match event {
            Click(clicked) => match model.selected_piece {
                Some(selected) => {
                    if clicked.color() != model.board.turn() || selected == clicked {
                        model.clear_selection();
                    } else if model.board.is_piece_on_field(&clicked) {
                        model.available_moves = Some(model.board.get_available_moves(&clicked));
                        model.selected_piece = Some(clicked);
                    } else {
                        try_move(model, Move::Move(selected, clicked));
                        model.clear_selection();
                    }
                }
                None => {
                    if model.exchanging && try_move(model, Move::Exchange(clicked)) {
                        model.exchanging = false;
                    } else if !model.exchanging && clicked.color() == model.board.turn()
                        && model.board.is_piece_on_field(&clicked)
                    {
                        model.available_moves = Some(model.board.get_available_moves(&clicked));
                        model.selected_piece = Some(clicked);
                    }
                }
            },
            NewGame(map) => *model = Model::new(map),
            Exchange => if model.board.can_exchange() {
                model.exchanging = !model.exchanging;
                model.clear_selection();
            },
            Resign => {
                model.commit_state();
                model.game_result = GameResult::Win(model.board.turn().switch());
            }
            Undo => model.undo_move(),
            Redo => model.redo_move(),
        }
    }
}

fn try_move(model: &mut Model, mv: Move) -> bool {
    if model.board.can_apply_move(&mv) {
        model.commit_state();
        model.board.apply_move(&mv);
        model.last_move = Some(mv);
        check_win(model);
        true
    } else {
        false
    }
}

fn check_win(model: &mut Model) {
    if model.board.pieces(Color::White) == 0 {
        model.game_result = GameResult::Win(Color::White);
    } else if model.board.pieces(Color::Black) == 0 {
        model.game_result = GameResult::Win(Color::Black);
    }
}
<commit_msg>Rename NewGame(map) to NewGame(players)<commit_after>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use model::{Color, GameResult, Model, Move};
use view::Event;

pub fn update(model: &mut Model, event: Option<Event>) {
    if let Some(event) = event {
        use view::Event::*;

        match event {
            Click(clicked) => match model.selected_piece {
                Some(selected) => {
                    if clicked.color() != model.board.turn() || selected == clicked {
                        model.clear_selection();
                    } else if model.board.is_piece_on_field(&clicked) {
                        model.available_moves = Some(model.board.get_available_moves(&clicked));
                        model.selected_piece = Some(clicked);
                    } else {
                        try_move(model, Move::Move(selected, clicked));
                        model.clear_selection();
                    }
                }
                None => {
                    if model.exchanging && try_move(model, Move::Exchange(clicked)) {
                        model.exchanging = false;
                    } else if !model.exchanging && clicked.color() == model.board.turn()
                        && model.board.is_piece_on_field(&clicked)
                    {
                        model.available_moves = Some(model.board.get_available_moves(&clicked));
                        model.selected_piece = Some(clicked);
                    }
                }
            },
            NewGame(players) => *model = Model::new(players),
            Exchange => if model.board.can_exchange() {
                model.exchanging = !model.exchanging;
                model.clear_selection();
            },
            Resign => {
                model.commit_state();
                model.game_result = GameResult::Win(model.board.turn().switch());
            }
            Undo => model.undo_move(),
            Redo => model.redo_move(),
        }
    }
}

fn try_move(model: &mut Model, mv: Move) -> bool {
    if model.board.can_apply_move(&mv) {
        model.commit_state();
        model.board.apply_move(&mv);
        model.last_move = Some(mv);
        check_win(model);
        true
    } else {
        false
    }
}

fn check_win(model: &mut Model) {
    if model.board.pieces(Color::White) == 0 {
        model.game_result = GameResult::Win(Color::White);
    } else if model.board.pieces(Color::Black) == 0 {
        model.game_result = GameResult::Win(Color::Black);
    }
}
<|endoftext|>