<commit_before>/*!

The afsort crate implements a sorting implementation based on
[American Flag sort](https://en.wikipedia.org/wiki/American_flag_sort). The implementation is
currently limited to sort byte slices, e.g. Strings. The main motivation is to sort strings of
text, so most of the benchmarks are based on English text strings. This implementation seems
to be about 40% faster than `sort_unstable` from the Rust standard library.

For small input, this method falls back to the standard library.

# Installation

Add the depndency to your `Cargo.toml`:

```ignore
[dependencies]
afsort = "0.1"
```
In your crate root:
```ignore
extern crate afsort;
```

# Usage

You can now use afsort to e.g. sort arrays of strings or string slices.

```rust
use afsort;
let mut strings = vec!("red", "green", "blue");
afsort::sort_unstable(strings);
assert_eq!(strings, vec!["blue", "green", "red"]);
```

You can also sort by an extractor function, e.g.:

```rust
use afsort;
let mut tuples = vec![("b", 2), ("a", 1)]
afsort::sort_unstable_by(&mut tuples, |t| t.0.as_bytes());
assert_eq!(strings, vec![("a", 1), ("b", 2)]);
```

# Motivation

Essentially, I noticed that sorting of strings took a long time when using the
[fst](https://github.com/BurntSushi/fst) crate, since it requires the input to be ordered.
Since sorting strings is a general problem, this is now a crate.

# Performance

As mentioned, this implementation seems to be about 40% faster than the sort in the standard
library, when sorting strings of random English words. That radix-style sorting algorithms can
be faster than quicksort is not anything new. The implementation is fairly naive, so I would
not be surprised if it could be improved further.

You can run the benchmark tests using `cargo bench`, like this:

```ignore
% cargo bench
   Compiling afsort v0.1.0 (file:///home/anton/dev/off/afsort)
    Finished release [optimized] target(s) in 5.66 secs
     Running target/release/deps/afsort-ca28db3ba0643253

running 12 tests
test tests::sorts_strings_same_as_unstable ... ignored
test tests::sorts_tuples_same_as_unstable ... ignored
test tests::sort_100000_en_af        ... bench:  20,402,968 ns/iter (+/- 1,512,907)
test tests::sort_100000_en_std       ... bench:  30,132,067 ns/iter (+/- 1,698,223)
test tests::sort_10000_en_af         ... bench:   1,377,303 ns/iter (+/- 114,481)
test tests::sort_10000_en_lower_af   ... bench:   1,371,022 ns/iter (+/- 95,391)
test tests::sort_10000_en_lower_std  ... bench:   2,227,486 ns/iter (+/- 127,281)
test tests::sort_10000_en_sorted_af  ... bench:     878,665 ns/iter (+/- 545,256)
test tests::sort_10000_en_sorted_std ... bench:     618,329 ns/iter (+/- 536,338)
test tests::sort_10000_en_std        ... bench:   2,221,089 ns/iter (+/- 157,461)
test tests::sort_1000_en_af          ... bench:     101,625 ns/iter (+/- 6,946)
test tests::sort_1000_en_std         ... bench:     171,655 ns/iter (+/- 9,844)

test result: ok. 0 passed; 0 failed; 2 ignored; 10 measured; 0 filtered out

```
# Limitations

Currently, this crate only supports sorting elements of `AsRef<[u8]>`. I think that this is a
good first step, since it supports sorting of Strings. There is however no reason why American
Flag sorting should be limited to this data type. Any kind of element that can deliver a radix
to a certain digit/depth can be sorted using this technique.

The American Flag algorithm is unstable, in the same way that sort_unstable in the standard
library. That is, equal elements might be re-ordered.

This crate can _only_ sort strings based on their `utf-8` byte values. For many problems, this
is fine. However, if you want to sort strings for display to a user, Locale might matter. This
crate does not try to address this issue.

# Testing

Testing is done using the [quickcheck](https://github.com/BurntSushi/quickcheck) crate. We run
about 50k different variations of input strings. We make sure that the afsort implementation
sorts in the same way as the standard library's sort_unstable methods.

*/

#![feature(test)]
#[cfg(test)]
extern crate test;
#[cfg(test)]
extern crate rand;
#[cfg(test)]
extern crate quickcheck;
#[cfg(test)]
extern crate regex;

/// Base module for the crate. See the crate documentation for more information.
pub mod afsort {

    /// Enhances slices of e.g. Strings to have a `af_sort_unstable` method, as a more idiomatic
    /// way to call sort.
    ///
    /// #Example
    ///
    /// ```rust
    /// let mut strings = vec!["c", "a", "b"];
    /// strings.af_sort_unstable();
    /// assert_eq!(strings, vec!["a", "b", "c"];
    /// ```
    pub trait AFSortable {
        fn af_sort_unstable(&mut self);
    }

    impl<T> AFSortable for [T]
    where
        T: AsRef<[u8]>,
    {
        fn af_sort_unstable(&mut self) {
            sort_unstable(self);
        }
    }

    /// Main sort method.
    ///
    /// #Example
    ///
    /// ```rust
    /// let mut strings = vec!["c", "a", "b"];
    /// strings.af_sort_unstable();
    /// assert_eq!(strings, vec!["a", "b", "c"];
    /// ```
    pub fn sort_unstable<'a, T>(vec: &'a mut [T])
    where
        T: AsRef<[u8]>,
    {
        sort_unstable_by(vec, (|s| s.as_ref()));
    }

    /// Sort method which accepts function to convert elements to &[u8].
    ///
    /// #Example
    ///
    /// ```rust
    /// let mut tuples = vec![("b", 2), ("a", 1)]
    ///afsort::sort_unstable_by(&mut tuples, |t| t.0.as_bytes());
    ///assert_eq!(strings, vec![("a", 1), ("b", 2)]);
    /// ```
    pub fn sort_unstable_by<T, F>(vec: &mut [T], to_slice: F)
    where
        F: Fn(&T) -> &[u8],
    {
        sort_req(vec, &to_slice, 0);
    }

    fn sort_req<T, F>(vec: &mut [T], to_slice: &F, depth: usize)
    where
        F: Fn(&T) -> &[u8],
    {
        if vec.len() <= 32 {
            vec.sort_unstable_by(|e1, e2| to_slice(e1).cmp(to_slice(e2)));
            return;
        }
        let mut min = u16::max_value();
        let mut max = 0u16;
        for elem in vec.iter() {
            let val = to_slice(elem);
            if val.len() > depth {
                let radix_val = val[depth] as u16;
                if radix_val < min {
                    min = radix_val;
                }
                if radix_val > max {
                    max = radix_val;
                }
            }
        }
        if min == u16::max_value() {
            return;
        }

        let num_items = (max - min + 2) as usize;
        let mut counts: Vec<usize> = vec![0usize; num_items as usize];

        for elem in vec.iter() {
            let radix_val = radix_for_str(to_slice(elem), depth, min);
            counts[radix_val as usize] += 1;
        }

        let mut sum = 0usize;
        let mut offsets: Vec<usize> = vec![0usize; num_items as usize];
        for i in 0..num_items {
            offsets[i as usize] = sum;
            sum += counts[i as usize];
        }
        let mut next_free = offsets.clone();

        let mut block = 0usize;
        let mut i = 0usize;
        while block < num_items - 1 {
            if i >= offsets[block as usize + 1] as usize {
                block += 1;
            } else {
                let radix_val = radix_for_str(to_slice(&vec[i]), depth, min);
                if radix_val == block as u16 {
                    i += 1;
                } else {
                    vec.swap(i as usize, next_free[radix_val as usize] as usize);
                    next_free[radix_val as usize] += 1;
                }
            }
        }
        for i in 0..num_items - 1 {
            sort_req(
                &mut vec[offsets[i as usize] as usize..offsets[i as usize + 1] as usize],
                to_slice,
                depth + 1,
            );
        }
        sort_req(&mut vec[offsets[num_items - 1]..], to_slice, depth + 1);
    }

    fn radix_for_str(s: &[u8], d: usize, base: u16) -> u16 {
        if s.len() > d {
            s[d] as u16 + 1 - base
        } else {
            0
        }
    }
}

#[cfg(test)]
mod tests {

    use afsort;
    use afsort::AFSortable;
    use quickcheck::QuickCheck;
    use test::Bencher;
    use std::fs::File;
    use std::io::{BufRead, BufReader};
    use rand::{self, Rng};
    use regex::Regex;
    use std::path::PathBuf;

    #[test]
    fn sorts_strings_same_as_unstable() {
        fn compare_sort(mut strings: Vec<String>) -> bool {
            let mut copy = strings.clone();
            copy.sort_unstable();
            strings.af_sort_unstable();
            strings == copy
        }
        QuickCheck::new().tests(50000).quickcheck(
            compare_sort as
                fn(Vec<String>) -> bool,
        );
    }

    #[test]
    fn sorts_tuples_same_as_unstable() {
        fn compare_sort(mut tuples: Vec<(String, u8)>) -> bool {
            let mut copy = tuples.clone();
            copy.sort_unstable_by(|t1, t2| t1.0.cmp(&t2.0));
            afsort::sort_unstable_by(&mut tuples, |t| t.0.as_bytes());
            tuples.into_iter().map(|t| t.0).collect::<Vec<String>>() ==
                copy.into_iter().map(|t| t.0).collect::<Vec<String>>()
        }
        QuickCheck::new().tests(50000).quickcheck(
            compare_sort as
                fn(Vec<(String, u8)>)
                   -> bool,
        );
    }

    #[bench]
    fn sort_1000_en_std(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r".*").unwrap(), 1_000);
        b.iter(|| strings.clone().sort_unstable())
    }

    #[bench]
    fn sort_1000_en_af(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r".*").unwrap(), 1_000);
        b.iter(|| strings.clone().af_sort_unstable())
    }

    #[bench]
    fn sort_10000_en_std(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r".*").unwrap(), 10_000);
        b.iter(|| strings.clone().sort_unstable())
    }

    #[bench]
    fn sort_10000_en_af(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r".*").unwrap(), 10_000);
        b.iter(|| strings.clone().af_sort_unstable())
    }

    #[bench]
    fn sort_100000_en_std(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r".*").unwrap(), 100_000);
        b.iter(|| strings.clone().sort_unstable())
    }

    #[bench]
    fn sort_100000_en_af(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r".*").unwrap(), 100_000);
        b.iter(|| strings.clone().af_sort_unstable())
    }

    #[bench]
    fn sort_10000_en_sorted_std(b: &mut Bencher) {
        let mut strings = strings_en(&Regex::new(r".*").unwrap(), 10_000);
        strings.sort_unstable();
        b.iter(|| strings.clone().sort_unstable())
    }

    #[bench]
    fn sort_10000_en_sorted_af(b: &mut Bencher) {
        let mut strings = strings_en(&Regex::new(r".*").unwrap(), 10_000);
        strings.sort_unstable();
        b.iter(|| strings.clone().af_sort_unstable())
    }

    #[bench]
    fn sort_10000_en_lower_std(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r"^[a-z]+$").unwrap(), 10000);
        b.iter(|| strings.clone().sort_unstable())
    }

    #[bench]
    fn sort_10000_en_lower_af(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r"^[a-z]+$").unwrap(), 10000);
        b.iter(|| strings.clone().af_sort_unstable())
    }

    fn strings_en(re: &Regex, n: usize) -> Vec<String> {
        let d = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
        let f = File::open(d.join("test_resources/american-english.txt")).unwrap();
        let b = BufReader::new(f);
        let mut strings = b.lines()
            .map(|l| l.unwrap())
            .filter(|l| re.is_match(l))
            .collect::<Vec<String>>();
        let mut rng = rand::thread_rng();
        rng.shuffle(&mut strings);
        strings.into_iter().take(n).collect::<Vec<String>>()
    }

}
<commit_msg>Fixed some crate issues.<commit_after>/*!

The afsort crate implements a sorting implementation based on
[American Flag sort](https://en.wikipedia.org/wiki/American_flag_sort). The implementation is
currently limited to sort byte slices, e.g. Strings. The main motivation is to sort strings of
text, so most of the benchmarks are based on English text strings. This implementation seems
to be about 40% faster than `sort_unstable` from the Rust standard library.

For small input, this method falls back to the standard library.

# Installation

Add the depndency to your `Cargo.toml`:

```ignore
[dependencies]
afsort = "0.1"
```
In your crate root:
```ignore
extern crate afsort;
```

# Usage

You can now use afsort to e.g. sort arrays of strings or string slices.

```rust
use afsort;
let mut strings = vec!("red", "green", "blue");
afsort::sort_unstable(&mut strings);
assert_eq!(strings, vec!["blue", "green", "red"]);
```

You can also sort by an extractor function, e.g.:

```rust
use afsort;
let mut tuples = vec![("b", 2), ("a", 1)];
afsort::sort_unstable_by(&mut tuples, |t| t.0.as_bytes());
assert_eq!(tuples, vec![("a", 1), ("b", 2)]);
```

# Motivation

Essentially, I noticed that sorting of strings took a long time when using the
[fst](https://github.com/BurntSushi/fst) crate, since it requires the input to be ordered.
Since sorting strings is a general problem, this is now a crate.

# Performance

As mentioned, this implementation seems to be about 40% faster than the sort in the standard
library, when sorting strings of random English words.  The implementation is fairly naive,
so I would not be surprised if it could be improved further.

You can run the benchmark tests using `cargo bench`, like this:

```ignore
% cargo bench
   Compiling afsort v0.1.0 (file:///home/anton/dev/off/afsort)
    Finished release [optimized] target(s) in 5.66 secs
     Running target/release/deps/afsort-ca28db3ba0643253

running 12 tests
test tests::sorts_strings_same_as_unstable ... ignored
test tests::sorts_tuples_same_as_unstable ... ignored
test tests::sort_100000_en_af        ... bench:  20,402,968 ns/iter (+/- 1,512,907)
test tests::sort_100000_en_std       ... bench:  30,132,067 ns/iter (+/- 1,698,223)
test tests::sort_10000_en_af         ... bench:   1,377,303 ns/iter (+/- 114,481)
test tests::sort_10000_en_lower_af   ... bench:   1,371,022 ns/iter (+/- 95,391)
test tests::sort_10000_en_lower_std  ... bench:   2,227,486 ns/iter (+/- 127,281)
test tests::sort_10000_en_sorted_af  ... bench:     878,665 ns/iter (+/- 545,256)
test tests::sort_10000_en_sorted_std ... bench:     618,329 ns/iter (+/- 536,338)
test tests::sort_10000_en_std        ... bench:   2,221,089 ns/iter (+/- 157,461)
test tests::sort_1000_en_af          ... bench:     101,625 ns/iter (+/- 6,946)
test tests::sort_1000_en_std         ... bench:     171,655 ns/iter (+/- 9,844)

test result: ok. 0 passed; 0 failed; 2 ignored; 10 measured; 0 filtered out

```
# Limitations

Currently, this crate only supports sorting elements of `AsRef<[u8]>`. I think that this is a
good first step, since it supports sorting of Strings. There is however no reason why American
Flag sorting should be limited to this data type. Any kind of element that can deliver a radix
to a certain digit/depth can be sorted using this technique.

The American Flag algorithm is unstable, in the same way that sort_unstable in the standard
library. That is, equal elements might be re-ordered.

This crate can _only_ sort strings based on their `utf-8` byte values. For many problems, this
is fine. However, if you want to sort strings for display to a user, Locale might matter. This
crate does not try to address this issue.

# Testing

Testing is done using the [quickcheck](https://github.com/BurntSushi/quickcheck) crate. We run
about 50k different variations of input strings. We make sure that the afsort implementation
sorts in the same way as the standard library's sort_unstable methods.

*/

#![feature(test)]
#[cfg(test)]
extern crate test;
#[cfg(test)]
extern crate rand;
#[cfg(test)]
extern crate quickcheck;
#[cfg(test)]
extern crate regex;

/// Enhances slices of e.g. Strings to have a `af_sort_unstable` method, as a more idiomatic
/// way to call sort.
///
/// #Example
///
/// ```rust
/// use afsort::AFSortable;
///
/// let mut strings = vec!["c", "a", "b"];
/// strings.af_sort_unstable();
/// assert_eq!(strings, vec!["a", "b", "c"]);
/// ```

pub trait AFSortable {
    fn af_sort_unstable(&mut self);
}

impl<T> AFSortable for [T]
where
    T: AsRef<[u8]>,
{
    fn af_sort_unstable(&mut self) {
        sort_unstable(self);
    }
}

/// Main sort method.
///
/// #Example
///
/// ```rust
/// let mut strings = vec!["c", "a", "b"];
/// afsort::sort_unstable(&mut strings);
/// assert_eq!(strings, vec!["a", "b", "c"]);
/// ```
pub fn sort_unstable<'a, T>(vec: &'a mut [T])
where
    T: AsRef<[u8]>,
{
    sort_unstable_by(vec, (|s| s.as_ref()));
}

/// Sort method which accepts function to convert elements to &[u8].
///
/// #Example
///
/// ```rust
/// let mut tuples = vec![("b", 2), ("a", 1)];
///afsort::sort_unstable_by(&mut tuples, |ref t| t.0.as_bytes());
///assert_eq!(tuples, vec![("a", 1), ("b", 2)]);
/// ```
pub fn sort_unstable_by<T, F>(vec: &mut [T], to_slice: F)
where
    F: Fn(&T) -> &[u8],
{
    sort_req(vec, &to_slice, 0);
}

fn sort_req<T, F>(vec: &mut [T], to_slice: &F, depth: usize)
where
    F: Fn(&T) -> &[u8],
{
    if vec.len() <= 32 {
        vec.sort_unstable_by(|e1, e2| to_slice(e1).cmp(to_slice(e2)));
        return;
    }
    let mut min = u16::max_value();
    let mut max = 0u16;
    for elem in vec.iter() {
        let val = to_slice(elem);
        if val.len() > depth {
            let radix_val = val[depth] as u16;
            if radix_val < min {
                min = radix_val;
            }
            if radix_val > max {
                max = radix_val;
            }
        }
    }
    if min == u16::max_value() {
        return;
    }

    let num_items = (max - min + 2) as usize;
    let mut counts: Vec<usize> = vec![0usize; num_items as usize];

    for elem in vec.iter() {
        let radix_val = radix_for_str(to_slice(elem), depth, min);
        counts[radix_val as usize] += 1;
    }

    let mut sum = 0usize;
    let mut offsets: Vec<usize> = vec![0usize; num_items as usize];
    for i in 0..num_items {
        offsets[i as usize] = sum;
        sum += counts[i as usize];
    }
    let mut next_free = offsets.clone();

    let mut block = 0usize;
    let mut i = 0usize;
    while block < num_items - 1 {
        if i >= offsets[block as usize + 1] as usize {
            block += 1;
        } else {
            let radix_val = radix_for_str(to_slice(&vec[i]), depth, min);
            if radix_val == block as u16 {
                i += 1;
            } else {
                vec.swap(i as usize, next_free[radix_val as usize] as usize);
                next_free[radix_val as usize] += 1;
            }
        }
    }
    for i in 0..num_items - 1 {
        sort_req(
            &mut vec[offsets[i as usize] as usize..offsets[i as usize + 1] as usize],
            to_slice,
            depth + 1,
        );
    }
    sort_req(&mut vec[offsets[num_items - 1]..], to_slice, depth + 1);
}

fn radix_for_str(s: &[u8], d: usize, base: u16) -> u16 {
    if s.len() > d {
        s[d] as u16 + 1 - base
    } else {
        0
    }
}
//}

#[cfg(test)]
mod tests {
    use super::AFSortable;
    use quickcheck::QuickCheck;
    use test::Bencher;
    use std::fs::File;
    use std::io::{BufRead, BufReader};
    use rand::{self, Rng};
    use regex::Regex;
    use std::path::PathBuf;

    #[test]
    fn sorts_strings_same_as_unstable() {
        fn compare_sort(mut strings: Vec<String>) -> bool {
            let mut copy = strings.clone();
            copy.sort_unstable();
            strings.af_sort_unstable();
            strings == copy
        }
        QuickCheck::new().tests(50000).quickcheck(
            compare_sort as
                fn(Vec<String>) -> bool,
        );
    }

    #[test]
    fn sorts_tuples_same_as_unstable() {
        fn compare_sort(mut tuples: Vec<(String, u8)>) -> bool {
            let mut copy = tuples.clone();
            copy.sort_unstable_by(|t1, t2| t1.0.cmp(&t2.0));
            super::sort_unstable_by(&mut tuples, |t| t.0.as_bytes());
            tuples.into_iter().map(|t| t.0).collect::<Vec<String>>() ==
                copy.into_iter().map(|t| t.0).collect::<Vec<String>>()
        }
        QuickCheck::new().tests(50000).quickcheck(
            compare_sort as
                fn(Vec<(String, u8)>)
                   -> bool,
        );
    }

    #[bench]
    fn sort_1000_en_std(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r".*").unwrap(), 1_000);
        b.iter(|| strings.clone().sort_unstable())
    }

    #[bench]
    fn sort_1000_en_af(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r".*").unwrap(), 1_000);
        b.iter(|| strings.clone().af_sort_unstable())
    }

    #[bench]
    fn sort_10000_en_std(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r".*").unwrap(), 10_000);
        b.iter(|| strings.clone().sort_unstable())
    }

    #[bench]
    fn sort_10000_en_af(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r".*").unwrap(), 10_000);
        b.iter(|| strings.clone().af_sort_unstable())
    }

    #[bench]
    fn sort_100000_en_std(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r".*").unwrap(), 100_000);
        b.iter(|| strings.clone().sort_unstable())
    }

    #[bench]
    fn sort_100000_en_af(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r".*").unwrap(), 100_000);
        b.iter(|| strings.clone().af_sort_unstable())
    }

    #[bench]
    fn sort_10000_en_sorted_std(b: &mut Bencher) {
        let mut strings = strings_en(&Regex::new(r".*").unwrap(), 10_000);
        strings.sort_unstable();
        b.iter(|| strings.clone().sort_unstable())
    }

    #[bench]
    fn sort_10000_en_sorted_af(b: &mut Bencher) {
        let mut strings = strings_en(&Regex::new(r".*").unwrap(), 10_000);
        strings.sort_unstable();
        b.iter(|| strings.clone().af_sort_unstable())
    }

    #[bench]
    fn sort_10000_en_lower_std(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r"^[a-z]+$").unwrap(), 10000);
        b.iter(|| strings.clone().sort_unstable())
    }

    #[bench]
    fn sort_10000_en_lower_af(b: &mut Bencher) {
        let strings = strings_en(&Regex::new(r"^[a-z]+$").unwrap(), 10000);
        b.iter(|| strings.clone().af_sort_unstable())
    }

    fn strings_en(re: &Regex, n: usize) -> Vec<String> {
        let d = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
        let f = File::open(d.join("test_resources/american-english.txt")).unwrap();
        let b = BufReader::new(f);
        let mut strings = b.lines()
            .map(|l| l.unwrap())
            .filter(|l| re.is_match(l))
            .collect::<Vec<String>>();
        let mut rng = rand::thread_rng();
        rng.shuffle(&mut strings);
        strings.into_iter().take(n).collect::<Vec<String>>()
    }

}
<|endoftext|>


<commit_before>#![feature(globs)]

extern crate semver;
extern crate conduit;

use std::io::net::ip::{IpAddr, Ipv4Addr};
use std::io::MemReader;
use std::collections::HashMap;
use std::fmt::Show;

use conduit::{Method, Scheme, Host, Extensions, Headers, TypeMap};

pub struct MockRequest {
    path: String,
    method: Method,
    query_string: Option<String>,
    string_body: Option<String>,
    build_headers: HashMap<String, String>,
    headers: MockHeaders,
    extensions: TypeMap,
    reader: Option<MemReader>
}

impl MockRequest {
    pub fn new(method: Method, path: &'static str) -> MockRequest {
        let headers = HashMap::new();

        MockRequest {
            path: path.to_string(),
            extensions: TypeMap::new(),
            query_string: None,
            string_body: None,
            build_headers: headers,
            headers: MockHeaders { headers: HashMap::new() },
            method: method,
            reader: None
        }
    }

    pub fn with_query<'a, S: Show>(&'a mut self, string: S) -> &'a mut MockRequest {
        self.query_string = Some(string.to_string());
        self
    }

    pub fn with_body<'a, S: Show>(&'a mut self, string: S) -> &'a mut MockRequest {
        self.string_body = Some(string.to_string());
        self
    }

    pub fn header<'a, S1: Show, S2: Show>(&'a mut self, name: S1, value: S2) -> &'a mut MockRequest {
        self.build_headers.insert(name.to_string(), value.to_string());
        let headers = MockHeaders { headers: self.build_headers.clone() };
        self.headers = headers;

        self
    }
}

pub struct MockHeaders {
    headers: HashMap<String, String>
}

impl Headers for MockHeaders {
    fn find<'a>(&'a self, key: &str) -> Option<Vec<&'a str>> {
        self.headers.find_equiv(&key).map(|v| vec!(v.as_slice()))
    }

    fn has(&self, key: &str) -> bool {
        self.headers.contains_key_equiv(&key)
    }

    fn iter<'a>(&'a self) -> conduit::HeaderEntries<'a> {
        box self.headers.iter().map(|(k,v)| (k.as_slice(), vec!(v.as_slice()))) as conduit::HeaderEntries<'a>
    }
}

impl<'a> conduit::Request for MockRequest {
    fn http_version(&self) -> semver::Version {
        semver::parse("1.1.0").unwrap()
    }

    fn conduit_version(&self) -> semver::Version {
        semver::parse("0.1.0").unwrap()
    }

    fn method(&self) -> Method { self.method }
    fn scheme(&self) -> Scheme { conduit::Http }
    fn host<'a>(&'a self) -> Host<'a> { conduit::HostName("example.com") }
    fn virtual_root<'a>(&'a self) -> Option<&'a str> { None }

    fn path<'a>(&'a self) -> &'a str {
        self.path.as_slice()
    }

    fn query_string<'a>(&'a self) -> Option<&'a str> {
        self.query_string.as_ref().map(|s| s.as_slice())
    }

    fn remote_ip(&self) -> IpAddr {
        Ipv4Addr(127, 0, 0, 1)
    }

    fn content_length(&self) -> Option<uint> {
        self.string_body.as_ref().map(|b| b.len())
    }

    fn headers<'a>(&'a self) -> &'a Headers {
        &self.headers as &Headers
    }

    fn body<'a>(&'a mut self) -> &'a mut Reader {
        let body = self.string_body.clone().unwrap_or("".to_string());
        self.reader = Some(MemReader::new(Vec::from_slice(body.as_bytes())));

        self.reader.get_mut_ref() as &mut Reader
    }

    fn extensions<'a>(&'a self) -> &'a Extensions {
        &self.extensions
    }
    fn mut_extensions<'a>(&'a mut self) -> &'a mut Extensions {
        &mut self.extensions
    }
}

#[cfg(test)]
mod tests {
    use super::MockRequest;
    use conduit;
    use semver;

    use std::io::net::ip::Ipv4Addr;

    use conduit::Request;

    #[test]
    fn simple_request_test() {
        let mut req = MockRequest::new(conduit::Get, "/");

        assert_eq!(req.http_version(), semver::parse("1.1.0").unwrap());
        assert_eq!(req.conduit_version(), semver::parse("0.1.0").unwrap());
        assert_eq!(req.method(), conduit::Get);
        assert_eq!(req.scheme(), conduit::Http);
        assert_eq!(req.host(), conduit::HostName("example.com"));
        assert_eq!(req.virtual_root(), None);
        assert_eq!(req.path(), "/");
        assert_eq!(req.query_string(), None);
        assert_eq!(req.remote_ip(), Ipv4Addr(127, 0, 0, 1));
        assert_eq!(req.content_length(), None);
        assert_eq!(req.headers().iter().count(), 0);
        assert_eq!(req.body().read_to_string().ok().expect("No body"),
                   "".to_string());
    }

    #[test]
    fn request_body_test() {
        let mut req = MockRequest::new(conduit::Post, "/articles");
        req.with_body("Hello world");

        assert_eq!(req.method(), conduit::Post);
        assert_eq!(req.path(), "/articles");
        assert_eq!(req.body().read_to_string().ok().expect("No body"),
                   "Hello world".to_string());
        assert_eq!(req.content_length(), Some(11));
    }

    #[test]
    fn request_query_test() {
        let mut req = MockRequest::new(conduit::Post, "/articles");
        req.with_query("foo=bar");

        assert_eq!(req.query_string().expect("No query string"), "foo=bar");
    }

    #[test]
    fn request_headers() {
        let mut req = MockRequest::new(conduit::Post, "/articles");
        req.header("User-Agent", "lulz");
        req.header("DNT", "1");

        assert_eq!(req.headers().iter().count(), 2);
        assert_eq!(req.headers().find("User-Agent").unwrap(), vec!("lulz"));
        assert_eq!(req.headers().find("DNT").unwrap(), vec!("1"));
    }
}
<commit_msg>Update to master<commit_after>#![feature(globs)]

extern crate semver;
extern crate conduit;

use std::io::net::ip::{IpAddr, Ipv4Addr};
use std::io::MemReader;
use std::collections::HashMap;
use std::fmt::Show;

use conduit::{Method, Scheme, Host, Extensions, Headers, TypeMap};

pub struct MockRequest {
    path: String,
    method: Method,
    query_string: Option<String>,
    string_body: Option<String>,
    build_headers: HashMap<String, String>,
    headers: MockHeaders,
    extensions: TypeMap,
    reader: Option<MemReader>
}

impl MockRequest {
    pub fn new(method: Method, path: &str) -> MockRequest {
        let headers = HashMap::new();

        MockRequest {
            path: path.to_string(),
            extensions: TypeMap::new(),
            query_string: None,
            string_body: None,
            build_headers: headers,
            headers: MockHeaders { headers: HashMap::new() },
            method: method,
            reader: None
        }
    }

    pub fn with_query<'a, S: Show>(&'a mut self, string: S) -> &'a mut MockRequest {
        self.query_string = Some(string.to_string());
        self
    }

    pub fn with_body<'a, S: Show>(&'a mut self, string: S) -> &'a mut MockRequest {
        self.string_body = Some(string.to_string());
        self
    }

    pub fn header<'a, S1: Show, S2: Show>(&'a mut self, name: S1, value: S2) -> &'a mut MockRequest {
        self.build_headers.insert(name.to_string(), value.to_string());
        let headers = MockHeaders { headers: self.build_headers.clone() };
        self.headers = headers;

        self
    }
}

pub struct MockHeaders {
    headers: HashMap<String, String>
}

impl Headers for MockHeaders {
    fn find<'a>(&'a self, key: &str) -> Option<Vec<&'a str>> {
        self.headers.find_equiv(&key).map(|v| vec!(v.as_slice()))
    }

    fn has(&self, key: &str) -> bool {
        self.headers.contains_key_equiv(&key)
    }

    fn iter<'a>(&'a self) -> conduit::HeaderEntries<'a> {
        box self.headers.iter().map(|(k,v)| (k.as_slice(), vec!(v.as_slice()))) as conduit::HeaderEntries<'a>
    }
}

impl<'a> conduit::Request for MockRequest {
    fn http_version(&self) -> semver::Version {
        semver::parse("1.1.0").unwrap()
    }

    fn conduit_version(&self) -> semver::Version {
        semver::parse("0.1.0").unwrap()
    }

    fn method(&self) -> Method { self.method }
    fn scheme(&self) -> Scheme { conduit::Http }
    fn host<'a>(&'a self) -> Host<'a> { conduit::HostName("example.com") }
    fn virtual_root<'a>(&'a self) -> Option<&'a str> { None }

    fn path<'a>(&'a self) -> &'a str {
        self.path.as_slice()
    }

    fn query_string<'a>(&'a self) -> Option<&'a str> {
        self.query_string.as_ref().map(|s| s.as_slice())
    }

    fn remote_ip(&self) -> IpAddr {
        Ipv4Addr(127, 0, 0, 1)
    }

    fn content_length(&self) -> Option<uint> {
        self.string_body.as_ref().map(|b| b.len())
    }

    fn headers<'a>(&'a self) -> &'a Headers {
        &self.headers as &Headers
    }

    fn body<'a>(&'a mut self) -> &'a mut Reader {
        let body = self.string_body.clone().unwrap_or("".to_string());
        self.reader = Some(MemReader::new(Vec::from_slice(body.as_bytes())));

        self.reader.as_mut().unwrap() as &mut Reader
    }

    fn extensions<'a>(&'a self) -> &'a Extensions {
        &self.extensions
    }
    fn mut_extensions<'a>(&'a mut self) -> &'a mut Extensions {
        &mut self.extensions
    }
}

#[cfg(test)]
mod tests {
    use super::MockRequest;
    use conduit;
    use semver;

    use std::io::net::ip::Ipv4Addr;

    use conduit::Request;

    #[test]
    fn simple_request_test() {
        let mut req = MockRequest::new(conduit::Get, "/");

        assert_eq!(req.http_version(), semver::parse("1.1.0").unwrap());
        assert_eq!(req.conduit_version(), semver::parse("0.1.0").unwrap());
        assert_eq!(req.method(), conduit::Get);
        assert_eq!(req.scheme(), conduit::Http);
        assert_eq!(req.host(), conduit::HostName("example.com"));
        assert_eq!(req.virtual_root(), None);
        assert_eq!(req.path(), "/");
        assert_eq!(req.query_string(), None);
        assert_eq!(req.remote_ip(), Ipv4Addr(127, 0, 0, 1));
        assert_eq!(req.content_length(), None);
        assert_eq!(req.headers().iter().count(), 0);
        assert_eq!(req.body().read_to_string().ok().expect("No body"),
                   "".to_string());
    }

    #[test]
    fn request_body_test() {
        let mut req = MockRequest::new(conduit::Post, "/articles");
        req.with_body("Hello world");

        assert_eq!(req.method(), conduit::Post);
        assert_eq!(req.path(), "/articles");
        assert_eq!(req.body().read_to_string().ok().expect("No body"),
                   "Hello world".to_string());
        assert_eq!(req.content_length(), Some(11));
    }

    #[test]
    fn request_query_test() {
        let mut req = MockRequest::new(conduit::Post, "/articles");
        req.with_query("foo=bar");

        assert_eq!(req.query_string().expect("No query string"), "foo=bar");
    }

    #[test]
    fn request_headers() {
        let mut req = MockRequest::new(conduit::Post, "/articles");
        req.header("User-Agent", "lulz");
        req.header("DNT", "1");

        assert_eq!(req.headers().iter().count(), 2);
        assert_eq!(req.headers().find("User-Agent").unwrap(), vec!("lulz"));
        assert_eq!(req.headers().find("DNT").unwrap(), vec!("1"));
    }
}
<|endoftext|>


<commit_before>// tag_safe
//
// A linting plugin to flag calls to methods not marked "tag_safe"
// from methods marked "tag_safe".
//
// Author: John Hodge (thePowersGang/Mutabah)
//
// TODO: Support '#[tag_unsafe(type)]' which is used when a method has no marker
// - Allows default safe fallback, with upwards propagation.
//
#![crate_name="tag_safe"]
#![crate_type="dylib"]
#![feature(plugin_registrar, rustc_private)]

#[macro_use]
extern crate log;

extern crate syntax;
#[macro_use]
extern crate rustc;

#[macro_use]
extern crate rustc_plugin;

use syntax::ast;
use rustc::hir::def_id::{self,DefId};
use rustc::hir::def;
use syntax::codemap::Span;
use rustc::lint::{self, LintContext, LintPass, LateLintPass, LintArray};
use rustc_plugin::Registry;
use rustc::ty::{self, TyCtxt};
use rustc::hir;

declare_lint!(NOT_TAGGED_SAFE, Warn, "Warn about use of non-tagged methods within tagged function");

#[derive(Copy,Clone,Debug)]
enum SafetyType
{
    Safe,
    Unsafe,
    Unknown,
}

#[derive(Default)]
struct Pass
{
    /// Cache of flag types
    flag_types: Vec<String>,
    /// Node => (Type => IsSafe)
    flag_cache: ::rustc::util::nodemap::NodeMap< ::rustc::util::nodemap::FnvHashMap<usize, SafetyType> >,

    lvl: usize,
}

struct Visitor<'a, 'gcx: 'a + 'tcx, 'tcx: 'a, F: FnMut(&Span) + 'a>
{
    pass: &'a mut Pass,
    tcx: &'a TyCtxt<'a, 'gcx, 'tcx>,
    name: &'a str,
    unknown_assume: bool,
    cb: F,
}

// Hack to provide indenting in debug calls
struct Indent(usize);
impl ::std::fmt::Display for Indent {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        for s in ::std::iter::repeat(" ").take(self.0) {
            try!(write!(f, "{}", s));
        }
        Ok( () )
    }
}

impl LintPass for Pass {
    fn get_lints(&self) -> LintArray {
        lint_array!(NOT_TAGGED_SAFE)
    }
}

impl LateLintPass for Pass {
    fn check_fn(&mut self, cx: &lint::LateContext, _kind: ::rustc::hir::intravisit::FnKind, _decl: &hir::FnDecl, body: &hir::Block, _: Span, id: ast::NodeId) {
        let attrs = cx.tcx.map.attrs(id);
        for ty in attrs.iter()
            .filter(|a| a.check_name("tag_safe"))
            .filter_map(|a| a.meta_item_list())
            .flat_map(|x| x.iter())
        {
            if let Some(ty_name) = ty.name()
            {
                // Search body for calls to non safe methods
                let mut v = Visitor{
                        pass: self, tcx: &cx.tcx, name: &ty_name,
                        // - Assumes an untagged method is unsafe
                        unknown_assume: false,
                        cb: |span| {
                                cx.span_lint(NOT_TAGGED_SAFE, *span,
                                    &format!("Calling {}-unsafe method from a #[tag_safe({})] method", ty_name, ty_name)[..]
                                    );
                            },
                        };
                debug!("Method {:?} is marked safe '{}'", id, ty_name);
                hir::intravisit::walk_block(&mut v, body);
            }
        }
    }
}

impl Pass
{
    // Searches for the relevant marker
    fn check_for_marker(tcx: &TyCtxt, id: ast::NodeId, marker: &str, name: &str) -> bool
    {
        debug!("Checking for marker {}({}) on {:?}", marker, name, id);
        tcx.map.attrs(id).iter()
            .filter_map( |a| if a.check_name(marker) { a.meta_item_list() } else { None })
            .flat_map(|x| x.iter())
            .any(|a| a.name().as_ref().map(|x| &x[..]) == Some(name))
    }

    /// Recursively check that the provided function is either safe or unsafe.
    // Used to avoid excessive annotating
    fn recurse_fcn_body(&mut self, tcx: &TyCtxt, node_id: ast::NodeId, name_id: usize, name: &str, unknown_assume: bool) -> bool
    {
        // Cache this method as unknown (to prevent infinite recursion)
        self.flag_cache.entry(node_id)
            .or_insert(Default::default())
            .insert(name_id, SafetyType::Unknown)
            ;

        // and apply a visitor to all
        match tcx.map.get(node_id)
        {
        rustc::hir::map::NodeItem(i) =>
            match i.node {
            hir::ItemFn(_, _, _, _, _, ref body) => {
                // Enumerate this function's code, recursively checking for a call to an unsafe method
                let mut is_safe = true;
                {
                    let mut v = Visitor {
                        pass: self, tcx: tcx, name: name,
                        unknown_assume: true,
                        cb: |_| { is_safe = false; }
                        };
                    hir::intravisit::walk_block(&mut v, body);
                }
                is_safe
                },
            _ => unknown_assume,
            },
        rustc::hir::map::NodeImplItem(i) =>
            match i.node {
            hir::ImplItemKind::Method(_, ref body) => {
                let mut is_safe = true;
                {
                    let mut v = Visitor {
                        pass: self, tcx: tcx, name: name,
                        unknown_assume: true,
                        cb: |_| { is_safe = false; }
                        };
                    hir::intravisit::walk_block(&mut v, body);
                }
                is_safe
                },
            _ => unknown_assume,
            },
        rustc::hir::map::NodeForeignItem(i) =>
            if Self::check_for_marker(tcx, i.id, "tag_safe", name) {
                true
            }
            else if Self::check_for_marker(tcx, i.id, "tag_unsafe", name) {
                false
            }
            else {
                unknown_assume
            },
        v @ _ => {
            error!("Node ID {} points to non-item {:?}", node_id, v);
            unknown_assume
            }
        }
    }

    /// Check that a method within this crate is safe with the provided tag
    fn crate_method_is_safe(&mut self, tcx: &TyCtxt, node_id: ast::NodeId, name: &str, unknown_assume: bool) -> bool
    {
        // Obtain tag name ID (avoids storing a string in the map)
        let name_id =
            match self.flag_types.iter().position(|a| *a == name)
            {
            Some(v) => v,
            None => {
                self.flag_types.push( String::from(name) );
                self.flag_types.len() - 1
                },
            };

        // Check cache first
        if let Some(&st) = self.flag_cache.get(&node_id).and_then(|a| a.get(&name_id))
        {
            match st
            {
            SafetyType::Safe => true,
            SafetyType::Unsafe => false,
            SafetyType::Unknown => unknown_assume,
            }
        }
        else
        {
            // Search for a safety marker, possibly recursing
            let is_safe =
                if Self::check_for_marker(tcx, node_id, "tag_safe", name) {
                    true
                }
                else if Self::check_for_marker(tcx, node_id, "tag_unsafe", name) {
                    false
                }
                else {
                    self.recurse_fcn_body(tcx, node_id, name_id, name, unknown_assume)
                };
            // Save resultant value
            self.flag_cache.entry(node_id)
                .or_insert(Default::default())
                .insert(name_id, if is_safe { SafetyType::Safe } else { SafetyType::Unsafe })
                ;
            is_safe
        }
    }

    /// Locate a #[tag_safe(<name>)] attribute on the passed item
    pub fn method_is_safe(&mut self, tcx: &TyCtxt, id: DefId, name: &str, unknown_assume: bool) -> bool
    {
        debug!("{}Checking method {:?} (A {})", Indent(self.lvl), id, unknown_assume);
        self.lvl += 1;
        let rv = if id.krate == def_id::LOCAL_CRATE {
                self.crate_method_is_safe(tcx, tcx.map.as_local_node_id(id).unwrap(), name, unknown_assume)
            }
            else {
                for a in tcx.get_attrs(id).iter()
                {
                    if a.check_name("tag_safe") {
                        if a.meta_item_list().iter().flat_map(|a| a.iter()).any(|a| a.name().as_ref().map(|x| &x[..]) == Some(name)) {
                            return true;
                        }
                    }
                    if a.check_name("tag_unsafe") {
                        if a.meta_item_list().iter().flat_map(|a| a.iter()).any(|a| a.name().as_ref().map(|x| &x[..]) == Some(name)) {
                            return false;
                        }
                    }
                }
                warn!("TODO: Crate ID non-zero {:?} (assuming safe)", id);
                // TODO: Check the crate import for an annotation
                true
            };
        self.lvl -= 1;
        debug!("{}Checking method {:?} = {}", Indent(self.lvl), id, rv);
        rv
    }
}

impl<'a, 'gcx: 'tcx + 'a, 'tcx: 'a, F: FnMut(&Span)> hir::intravisit::Visitor<'a> for Visitor<'a,'gcx, 'tcx, F>
{
    // Locate function/method calls in a code block
    // - uses visit_expr_post because it doesn't _need_ to do anything
    fn visit_expr_post(&mut self, ex: &'a hir::Expr) {
        match ex.node
        {
        // Call expressions - check that it's a path call
        hir::ExprCall(ref fcn, _) =>
            match fcn.node
            {
            hir::ExprPath(ref _qs, ref _p) => {
                    if let def::Def::Fn(did) = self.tcx.expect_def(fcn.id) {
                        // Check for a safety tag
                        if !self.pass.method_is_safe(self.tcx, did, self.name, self.unknown_assume)
                        {
                            (self.cb)(&ex.span);
                        }
                    }
                },
            _ => {},
            },

        // Method call expressions - get the relevant method
        hir::ExprMethodCall(ref _id, ref _tys, ref _exprs) =>
            {
                let tables = self.tcx.tables.borrow();
                let mm = &tables.method_map;

                let callee = mm.get( &ty::MethodCall::expr(ex.id) ).unwrap();
                let id = callee.def_id;

                //if let ty::MethodStatic(id) = callee.origin {
                        // Check for a safety tag
                        if !self.pass.method_is_safe(self.tcx, id, self.name, self.unknown_assume) {
                            (self.cb)(&ex.span);
                        }
                //}
            },

        // Ignore any other type of node
        _ => {},
        }
    }
}

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    use syntax::feature_gate::AttributeType;
    reg.register_late_lint_pass( Box::new(Pass::default()) );

    reg.register_attribute(String::from("tag_safe"),   AttributeType::Whitelisted);
    reg.register_attribute(String::from("tag_unsafe"), AttributeType::Whitelisted);
}

// vim: ts=4 expandtab sw=4
<commit_msg>Nighly fix<commit_after>// tag_safe
//
// A linting plugin to flag calls to methods not marked "tag_safe"
// from methods marked "tag_safe".
//
// Author: John Hodge (thePowersGang/Mutabah)
//
// TODO: Support '#[tag_unsafe(type)]' which is used when a method has no marker
// - Allows default safe fallback, with upwards propagation.
//
#![crate_name="tag_safe"]
#![crate_type="dylib"]
#![feature(plugin_registrar, rustc_private)]

#[macro_use]
extern crate log;

extern crate syntax;
#[macro_use]
extern crate rustc;

#[macro_use]
extern crate rustc_plugin;

use syntax::ast;
use rustc::hir::def_id::{self,DefId};
use rustc::hir::def;
use syntax::codemap::Span;
use rustc::lint::{self, LintContext, LintPass, LateLintPass, LintArray};
use rustc_plugin::Registry;
use rustc::ty::{self, TyCtxt};
use rustc::hir;

declare_lint!(NOT_TAGGED_SAFE, Warn, "Warn about use of non-tagged methods within tagged function");

#[derive(Copy,Clone,Debug)]
enum SafetyType
{
    Safe,
    Unsafe,
    Unknown,
}

#[derive(Default)]
struct Pass
{
    /// Cache of flag types
    flag_types: Vec<String>,
    /// Node => (Type => IsSafe)
    flag_cache: ::rustc::util::nodemap::NodeMap< ::std::collections::HashMap<usize, SafetyType> >,

    lvl: usize,
}

struct Visitor<'a, 'gcx: 'a + 'tcx, 'tcx: 'a, F: FnMut(&Span) + 'a>
{
    pass: &'a mut Pass,
    tcx: &'a TyCtxt<'a, 'gcx, 'tcx>,
    name: &'a str,
    unknown_assume: bool,
    cb: F,
}

// Hack to provide indenting in debug calls
struct Indent(usize);
impl ::std::fmt::Display for Indent {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        for s in ::std::iter::repeat(" ").take(self.0) {
            try!(write!(f, "{}", s));
        }
        Ok( () )
    }
}

impl LintPass for Pass {
    fn get_lints(&self) -> LintArray {
        lint_array!(NOT_TAGGED_SAFE)
    }
}

impl LateLintPass for Pass {
    fn check_fn(&mut self, cx: &lint::LateContext, _kind: ::rustc::hir::intravisit::FnKind, _decl: &hir::FnDecl, body: &hir::Expr, _: Span, id: ast::NodeId) {
        let attrs = cx.tcx.map.attrs(id);
        for ty in attrs.iter()
            .filter(|a| a.check_name("tag_safe"))
            .filter_map(|a| a.meta_item_list())
            .flat_map(|x| x.iter())
        {
            if let Some(ty_name) = ty.name()
            {
                // Search body for calls to non safe methods
                let mut v = Visitor{
                        pass: self, tcx: &cx.tcx, name: &ty_name,
                        // - Assumes an untagged method is unsafe
                        unknown_assume: false,
                        cb: |span| {
                                cx.span_lint(NOT_TAGGED_SAFE, *span,
                                    &format!("Calling {}-unsafe method from a #[tag_safe({})] method", ty_name, ty_name)[..]
                                    );
                            },
                        };
                debug!("Method {:?} is marked safe '{}'", id, ty_name);
                hir::intravisit::walk_expr(&mut v, body);
            }
        }
    }
}

impl Pass
{
    // Searches for the relevant marker
    fn check_for_marker(tcx: &TyCtxt, id: ast::NodeId, marker: &str, name: &str) -> bool
    {
        debug!("Checking for marker {}({}) on {:?}", marker, name, id);
        tcx.map.attrs(id).iter()
            .filter_map( |a| if a.check_name(marker) { a.meta_item_list() } else { None })
            .flat_map(|x| x.iter())
            .any(|a| a.name().as_ref().map(|x| &x[..]) == Some(name))
    }

    /// Recursively check that the provided function is either safe or unsafe.
    // Used to avoid excessive annotating
    fn recurse_fcn_body(&mut self, tcx: &TyCtxt, node_id: ast::NodeId, name_id: usize, name: &str, unknown_assume: bool) -> bool
    {
        // Cache this method as unknown (to prevent infinite recursion)
        self.flag_cache.entry(node_id)
            .or_insert(Default::default())
            .insert(name_id, SafetyType::Unknown)
            ;

        // and apply a visitor to all
        match tcx.map.get(node_id)
        {
        rustc::hir::map::NodeItem(i) =>
            match i.node {
            hir::ItemFn(_, _, _, _, _, ref body) => {
                // Enumerate this function's code, recursively checking for a call to an unsafe method
                let mut is_safe = true;
                {
                    let mut v = Visitor {
                        pass: self, tcx: tcx, name: name,
                        unknown_assume: true,
                        cb: |_| { is_safe = false; }
                        };
                    hir::intravisit::walk_expr(&mut v, body);
                }
                is_safe
                },
            _ => unknown_assume,
            },
        rustc::hir::map::NodeImplItem(i) =>
            match i.node {
            hir::ImplItemKind::Method(_, ref body) => {
                let mut is_safe = true;
                {
                    let mut v = Visitor {
                        pass: self, tcx: tcx, name: name,
                        unknown_assume: true,
                        cb: |_| { is_safe = false; }
                        };
                    hir::intravisit::walk_expr(&mut v, body);
                }
                is_safe
                },
            _ => unknown_assume,
            },
        rustc::hir::map::NodeForeignItem(i) =>
            if Self::check_for_marker(tcx, i.id, "tag_safe", name) {
                true
            }
            else if Self::check_for_marker(tcx, i.id, "tag_unsafe", name) {
                false
            }
            else {
                unknown_assume
            },
        v @ _ => {
            error!("Node ID {} points to non-item {:?}", node_id, v);
            unknown_assume
            }
        }
    }

    /// Check that a method within this crate is safe with the provided tag
    fn crate_method_is_safe(&mut self, tcx: &TyCtxt, node_id: ast::NodeId, name: &str, unknown_assume: bool) -> bool
    {
        // Obtain tag name ID (avoids storing a string in the map)
        let name_id =
            match self.flag_types.iter().position(|a| *a == name)
            {
            Some(v) => v,
            None => {
                self.flag_types.push( String::from(name) );
                self.flag_types.len() - 1
                },
            };

        // Check cache first
        if let Some(&st) = self.flag_cache.get(&node_id).and_then(|a| a.get(&name_id))
        {
            match st
            {
            SafetyType::Safe => true,
            SafetyType::Unsafe => false,
            SafetyType::Unknown => unknown_assume,
            }
        }
        else
        {
            // Search for a safety marker, possibly recursing
            let is_safe =
                if Self::check_for_marker(tcx, node_id, "tag_safe", name) {
                    true
                }
                else if Self::check_for_marker(tcx, node_id, "tag_unsafe", name) {
                    false
                }
                else {
                    self.recurse_fcn_body(tcx, node_id, name_id, name, unknown_assume)
                };
            // Save resultant value
            self.flag_cache.entry(node_id)
                .or_insert(Default::default())
                .insert(name_id, if is_safe { SafetyType::Safe } else { SafetyType::Unsafe })
                ;
            is_safe
        }
    }

    /// Locate a #[tag_safe(<name>)] attribute on the passed item
    pub fn method_is_safe(&mut self, tcx: &TyCtxt, id: DefId, name: &str, unknown_assume: bool) -> bool
    {
        debug!("{}Checking method {:?} (A {})", Indent(self.lvl), id, unknown_assume);
        self.lvl += 1;
        let rv = if id.krate == def_id::LOCAL_CRATE {
                self.crate_method_is_safe(tcx, tcx.map.as_local_node_id(id).unwrap(), name, unknown_assume)
            }
            else {
                for a in tcx.get_attrs(id).iter()
                {
                    if a.check_name("tag_safe") {
                        if a.meta_item_list().iter().flat_map(|a| a.iter()).any(|a| a.name().as_ref().map(|x| &x[..]) == Some(name)) {
                            return true;
                        }
                    }
                    if a.check_name("tag_unsafe") {
                        if a.meta_item_list().iter().flat_map(|a| a.iter()).any(|a| a.name().as_ref().map(|x| &x[..]) == Some(name)) {
                            return false;
                        }
                    }
                }
                warn!("TODO: Crate ID non-zero {:?} (assuming safe)", id);
                // TODO: Check the crate import for an annotation
                true
            };
        self.lvl -= 1;
        debug!("{}Checking method {:?} = {}", Indent(self.lvl), id, rv);
        rv
    }
}

impl<'a, 'gcx: 'tcx + 'a, 'tcx: 'a, F: FnMut(&Span)> hir::intravisit::Visitor<'a> for Visitor<'a,'gcx, 'tcx, F>
{
    // Locate function/method calls in a code block
    // - uses visit_expr_post because it doesn't _need_ to do anything
    fn visit_expr_post(&mut self, ex: &'a hir::Expr) {
        match ex.node
        {
        // Call expressions - check that it's a path call
        hir::ExprCall(ref fcn, _) =>
            match fcn.node
            {
            hir::ExprPath(ref _qs, ref _p) => {
                    if let def::Def::Fn(did) = self.tcx.expect_def(fcn.id) {
                        // Check for a safety tag
                        if !self.pass.method_is_safe(self.tcx, did, self.name, self.unknown_assume)
                        {
                            (self.cb)(&ex.span);
                        }
                    }
                },
            _ => {},
            },

        // Method call expressions - get the relevant method
        hir::ExprMethodCall(ref _id, ref _tys, ref _exprs) =>
            {
                let tables = self.tcx.tables.borrow();
                let mm = &tables.method_map;

                let callee = mm.get( &ty::MethodCall::expr(ex.id) ).unwrap();
                let id = callee.def_id;

                //if let ty::MethodStatic(id) = callee.origin {
                        // Check for a safety tag
                        if !self.pass.method_is_safe(self.tcx, id, self.name, self.unknown_assume) {
                            (self.cb)(&ex.span);
                        }
                //}
            },

        // Ignore any other type of node
        _ => {},
        }
    }
}

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    use syntax::feature_gate::AttributeType;
    reg.register_late_lint_pass( Box::new(Pass::default()) );

    reg.register_attribute(String::from("tag_safe"),   AttributeType::Whitelisted);
    reg.register_attribute(String::from("tag_unsafe"), AttributeType::Whitelisted);
}

// vim: ts=4 expandtab sw=4
<|endoftext|>


<commit_before>//! The package provides a parser for the
//! [TGFF](http://ziyang.eecs.umich.edu/~dickrp/tgff/) (Task Graphs For Free)
//! format, which is a format for storing task graphs and accompanying data
//! used in scheduling and allocation research.

#![feature(macro_rules, if_let)]

use std::collections::HashMap;
use std::iter::Peekable;
use std::str::CharOffsets;

pub use content::Content;
pub use content::{Graph, Task, Arc, Deadline};
pub use content::{Table, Column};

mod content;

static READ_CAPACITY: uint = 20;

pub type Result<T> = std::result::Result<T, Error>;

pub struct Error {
    line: uint,
    message: &'static str,
}

impl std::fmt::Show for Error {
    fn fmt(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(formatter, "{} on line {}", self.message, self.line)
    }
}

pub struct Parser<'a> {
    line: uint,
    cursor: Peekable<(uint, char), CharOffsets<'a>>,
    content: Content,
}

macro_rules! raise(
    ($parser:expr, $message:expr) => (
        return Err(Error { line: $parser.line, message: $message });
    );
)

macro_rules! some(
    ($parser:expr, $result:expr, $message:expr) => (
        match $result {
            Some(result) => result,
            None => raise!($parser, $message),
        }
    );
)

impl<'a> Parser<'a> {
    /// Create a new `Parser` for processing the content of a TGFF file
    /// generated by the `tgff` command-line utility and given in `input`.
    pub fn new(input: &'a str) -> Parser<'a> {
        Parser {
            line: 1,
            cursor: input.char_indices().peekable(),
            content: Content::new(),
        }
    }

    /// Perform parsing of the data passed to `new`.
    pub fn process<'a>(&'a mut self) -> Result<&'a Content> {
        loop {
            match self.peek() {
                Some('@') => try!(self.process_at()),
                Some(_) => raise!(self, "found an unknown statement"),
                None => return Ok(&self.content),
            }
        }
    }

    fn process_at(&mut self) -> Result<()> {
        self.next(); // @

        let name = some!(self, self.read_token(),
                         "found an @-statement without a name");
        let number = some!(self, self.read_natural(),
                           "found an @-statement not followed by a natural number");

        if let Some('{') = self.peek() {
            self.process_block(name, number)
        } else {
            self.content.set_attribute(name, number);
            Ok(())
        }
    }

    fn process_block(&mut self, name: String, id: uint) -> Result<()> {
        self.next(); // {
        self.skip_void();

        if let Some('#') = self.peek() {
            try!(self.process_table(name, id));
        } else {
            try!(self.process_graph(name, id));
        }

        if let Some('}') = self.peek() {
            self.next();
            return Ok(());
        }

        raise!(self, "cannot find the end of a {}-block");
    }

    fn process_graph(&mut self, name: String, id: uint) -> Result<()> {
        let mut graph = Graph::new(name, id);

        loop {
            match self.read_token() {
                Some(ref token) => match token.as_slice() {
                    "TASK" => {
                        let id = some!(self, self.read_id(),
                                       "found a task without an id");

                        some!(self, self.skip_sequence("TYPE"),
                              "found a task without a TYPE");

                        let kind = some!(self, self.read_natural(),
                                         "found a task type without a value");

                        graph.add_task(Task::new(id, kind));
                    },
                    "ARC" => {
                        let id = some!(self, self.read_id(),
                                       "found an arc without an id");

                        some!(self, self.skip_sequence("FROM"),
                              "found an arc without a source");

                        let from = some!(self, self.read_id(),
                                         "found an arc source without a value");

                        some!(self, self.skip_sequence("TO"),
                              "found an arc without a destination");

                        let to = some!(self, self.read_id(),
                                       "found an arc destination without a value");

                        some!(self, self.skip_sequence("TYPE"),
                              "found an arc without a type");

                        let kind = some!(self, self.read_natural(),
                                         "found an arc type without a value");

                        graph.add_arc(Arc::new(id, from, to, kind));
                    },
                    "HARD_DEADLINE" => {
                        let id = some!(self, self.read_id(),
                                       "found a deadline without an id");

                        some!(self, self.skip_sequence("ON"),
                              "found a deadline without a target");

                        let on = some!(self, self.read_id(),
                                       "found a deadline target without a value");

                        some!(self, self.skip_sequence("AT"),
                              "found a deadline without a time stamp");

                        let at = some!(self, self.read_natural(),
                                       "found a deadline time stamp without a value");

                        graph.add_deadline(Deadline::new(id, on, at));
                    },
                    _ => {
                        let value = some!(self, self.read_natural(),
                                          "found an attribute not followed by a natural number");

                        graph.set_attribute(token.clone(), value);
                    },
                },
                None => break,
            }
        }

        self.content.add_graph(graph);
        Ok(())
    }

    fn process_table(&mut self, name: String, id: uint) -> Result<()> {
        let mut table = Table::new(name, id);

        self.content.add_table(table);
        Ok(())
    }

    fn skip(&mut self, accept: |uint, char| -> bool) -> uint {
        let mut count = 0;

        loop {
            match self.peek() {
                Some(c) => {
                    if !accept(count, c) { break; }
                    self.next();
                    count += 1;
                },
                None => break,
            }
        }

        count
    }

    #[inline]
    fn skip_void(&mut self) {
        self.skip(|_, c| c == ' ' || c == '\t' || c == '\n');
    }

    fn skip_sequence(&mut self, chars: &str) -> Option<()> {
        let len = chars.len();
        let count = self.skip(|i, c| i < len && c == chars.char_at(i));
        self.skip_void();
        if len == count { Some(()) } else { None }
    }

    fn read(&mut self, accept: |uint, char| -> bool) -> Option<String> {
        let mut result = std::string::String::with_capacity(READ_CAPACITY);
        let mut count = 0;

        loop {
            match self.peek() {
                Some(c) => {
                    if !accept(count, c) { break; }
                    result.push(c);
                    self.next();
                    count += 1;
                },
                None => break,
            }
        }

        if count == 0 {
            None
        } else {
            Some(result)
        }
    }

    fn read_token(&mut self) -> Option<String> {
        let result = self.read(|i, c| {
            match c {
                'A'...'Z' | 'a'...'z' if i == 0 => true,
                'A'...'Z' | 'a'...'z' | '_' | '0'...'9' if i > 0 => true,
                _ => false,
            }
        });
        self.skip_void();
        result
    }

    fn read_natural(&mut self) -> Option<uint> {
        let result = match self.read(|_, c| c >= '0' && c <= '9') {
            Some(ref number) => std::num::from_str_radix(number.as_slice(), 10),
            None => None,
        };
        self.skip_void();
        result
    }

    fn read_id(&mut self) -> Option<uint> {
        match self.read_token() {
            Some(ref token) => match token.as_slice().split('_').nth(1) {
                Some(id) => std::num::from_str_radix(id, 10),
                None => None,
            },
            None => None,
        }
    }

    #[inline]
    fn peek(&mut self) -> Option<char> {
        match self.cursor.peek() {
            Some(&(_, c)) => Some(c),
            None => None,
        }
    }
}

impl<'a> std::iter::Iterator<char> for Parser<'a> {
    fn next(&mut self) -> Option<char> {
        match self.cursor.next() {
            Some((_, '\n')) => {
                self.line += 1;
                Some('\n')
            },
            Some((_, c)) => Some(c),
            None => None,
        }
    }
}

#[cfg(test)]
mod tests {
    macro_rules! assert_ok(
        ($result: expr) => (
            if let Err(err) = $result {
                assert!(false, "{}", err);
            }
        );
    )

    macro_rules! assert_error(
        ($result: expr) => (
            if let Ok(_) = $result {
                assert!(false, "expected an error");
            }
        );
    )

    macro_rules! parser(
        ($input:expr) => (super::Parser::new($input));
    )

    #[test]
    fn process_at() {
        assert_ok!(parser!("@abc 12").process_at());
        assert_error!(parser!("@ ").process_at());
        assert_error!(parser!("@abc").process_at());
    }

    #[test]
    fn process_block() {
        assert_ok!(parser!("{}").process_block(String::from_str("life"), 42));
    }

    #[test]
    fn process_graph() {
        let mut parser = parser!("TASK t0_0\tTYPE 2   ");
        parser.process_graph(String::new(), 0);
        {
            let ref task = parser.content.graphs[0].tasks[0];
            assert_eq!(task.id, 0);
            assert_eq!(task.kind, 2);
        }

        parser = parser!("ARC a0_42 \tFROM t0_0  TO  t0_1 TYPE 35   ");
        parser.process_graph(String::new(), 0);
        {
            let ref arc = parser.content.graphs[0].arcs[0];
            assert_eq!(arc.id, 42);
            assert_eq!(arc.from, 0);
            assert_eq!(arc.to, 1);
            assert_eq!(arc.kind, 35);
        }

        parser = parser!("HARD_DEADLINE d0_9 ON t0_12 AT 1000   ");
        parser.process_graph(String::new(), 0);
        {
            let ref deadline = parser.content.graphs[0].deadlines[0];
            assert_eq!(deadline.id, 9);
            assert_eq!(deadline.on, 12);
            assert_eq!(deadline.at, 1000);
        }
    }

    #[test]
    fn process_table() {
    }

    #[test]
    fn skip_void() {
        let mut parser = parser!("  \t  abc");
        parser.skip_void();
        assert_eq!(parser.next().unwrap(), 'a');
    }

    #[test]
    fn read_token() {
        macro_rules! test(
            ($input:expr, $output:expr) => (
                assert_eq!(parser!($input).read_token().unwrap(),
                           String::from_str($output));
            );
        )

        test!("AZ xyz", "AZ");
        test!("az xyz", "az");
        test!("AZ_az_09 xyz", "AZ_az_09");
    }

    #[test]
    fn read_natural() {
        assert_eq!(parser!("09").read_natural().unwrap(), 9);
    }

    #[test]
    fn read_id() {
        assert_eq!(parser!("t0_42").read_id().unwrap(), 42);
    }
}
<commit_msg>Simplified error handling<commit_after>//! The package provides a parser for the
//! [TGFF](http://ziyang.eecs.umich.edu/~dickrp/tgff/) (Task Graphs For Free)
//! format, which is a format for storing task graphs and accompanying data
//! used in scheduling and allocation research.

#![feature(macro_rules, if_let)]

use std::collections::HashMap;
use std::iter::Peekable;
use std::str::CharOffsets;

pub use content::Content;
pub use content::{Graph, Task, Arc, Deadline};
pub use content::{Table, Column};

mod content;

static READ_CAPACITY: uint = 20;

pub type Result<T> = std::result::Result<T, Error>;

pub struct Error {
    line: uint,
    message: String,
}

impl std::fmt::Show for Error {
    fn fmt(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(formatter, "{} on line {}", self.message, self.line)
    }
}

pub struct Parser<'a> {
    line: uint,
    cursor: Peekable<(uint, char), CharOffsets<'a>>,
    content: Content,
}

macro_rules! raise(
    ($parser:expr, $($arg:tt)*) => (
        return Err(Error { line: $parser.line, message: format!($($arg)*) });
    );
)

macro_rules! some(
    ($parser:expr, $result:expr, $($arg:tt)*) => (
        match $result {
            Some(result) => result,
            None => raise!($parser, $($arg)*),
        }
    );
)

impl<'a> Parser<'a> {
    /// Create a new `Parser` for processing the content of a TGFF file
    /// generated by the `tgff` command-line utility and given in `input`.
    pub fn new(input: &'a str) -> Parser<'a> {
        Parser {
            line: 1,
            cursor: input.char_indices().peekable(),
            content: Content::new(),
        }
    }

    /// Perform parsing of the data passed to `new`.
    pub fn process<'a>(&'a mut self) -> Result<&'a Content> {
        loop {
            match self.peek() {
                Some('@') => try!(self.process_at()),
                Some(_) => raise!(self, "found an unknown statement"),
                None => return Ok(&self.content),
            }
        }
    }

    fn process_at(&mut self) -> Result<()> {
        self.next(); // @

        let name = try!(self.get_token());
        let number = try!(self.get_natural());

        if let Some('{') = self.peek() {
            self.process_block(name, number)
        } else {
            self.content.set_attribute(name, number);
            Ok(())
        }
    }

    fn process_block(&mut self, name: String, id: uint) -> Result<()> {
        self.next(); // {
        self.skip_void();

        if let Some('#') = self.peek() {
            try!(self.process_table(name, id));
        } else {
            try!(self.process_graph(name, id));
        }

        if let Some('}') = self.peek() {
            self.next();
            return Ok(());
        }

        raise!(self, "cannot find the end of a block");
    }

    fn process_graph(&mut self, name: String, id: uint) -> Result<()> {
        let mut graph = Graph::new(name, id);

        loop {
            match self.read_token() {
                Some(ref token) => match token.as_slice() {
                    "TASK" => {
                        let id = try!(self.get_id());
                        try!(self.skip_str("TYPE"));
                        let kind = try!(self.get_natural());

                        graph.add_task(Task::new(id, kind));
                    },
                    "ARC" => {
                        let id = try!(self.get_id());
                        try!(self.skip_str("FROM"));
                        let from = try!(self.get_id());
                        try!(self.skip_str("TO"));
                        let to = try!(self.get_id());
                        try!(self.skip_str("TYPE"));
                        let kind = try!(self.get_natural());

                        graph.add_arc(Arc::new(id, from, to, kind));
                    },
                    "HARD_DEADLINE" => {
                        let id = try!(self.get_id());
                        try!(self.skip_str("ON"));
                        let on = try!(self.get_id());
                        try!(self.skip_str("AT"));
                        let at = try!(self.get_natural());

                        graph.add_deadline(Deadline::new(id, on, at));
                    },
                    _ => {
                        let value = try!(self.get_natural());

                        graph.set_attribute(token.clone(), value);
                    },
                },
                None => break,
            }
        }

        self.content.add_graph(graph);
        Ok(())
    }

    fn process_table(&mut self, name: String, id: uint) -> Result<()> {
        let mut table = Table::new(name, id);

        self.content.add_table(table);
        Ok(())
    }

    fn skip(&mut self, accept: |uint, char| -> bool) -> uint {
        let mut count = 0;

        loop {
            match self.peek() {
                Some(c) => {
                    if !accept(count, c) { break; }
                    self.next();
                    count += 1;
                },
                None => break,
            }
        }

        count
    }

    #[inline]
    fn skip_void(&mut self) {
        self.skip(|_, c| c == ' ' || c == '\t' || c == '\n');
    }

    fn skip_str(&mut self, chars: &str) -> Result<()> {
        let len = chars.len();
        if self.skip(|i, c| i < len && c == chars.char_at(i)) != len {
            raise!(self, "expected `{}`", chars);
        }
        self.skip_void();
        Ok(())
    }

    fn read(&mut self, accept: |uint, char| -> bool) -> Option<String> {
        let mut result = std::string::String::with_capacity(READ_CAPACITY);
        let mut count = 0;

        loop {
            match self.peek() {
                Some(c) => {
                    if !accept(count, c) { break; }
                    result.push(c);
                    self.next();
                    count += 1;
                },
                None => break,
            }
        }

        if count == 0 {
            None
        } else {
            Some(result)
        }
    }

    fn read_token(&mut self) -> Option<String> {
        let result = self.read(|i, c| {
            match c {
                'A'...'Z' | 'a'...'z' if i == 0 => true,
                'A'...'Z' | 'a'...'z' | '_' | '0'...'9' if i > 0 => true,
                _ => false,
            }
        });
        self.skip_void();
        result
    }

    fn read_natural(&mut self) -> Option<uint> {
        let result = match self.read(|_, c| c >= '0' && c <= '9') {
            Some(ref number) => std::num::from_str_radix(number.as_slice(), 10),
            None => None,
        };
        self.skip_void();
        result
    }

    fn read_id(&mut self) -> Option<uint> {
        match self.read_token() {
            Some(ref token) => match token.as_slice().split('_').nth(1) {
                Some(id) => std::num::from_str_radix(id, 10),
                None => None,
            },
            None => None,
        }
    }

    fn get_token(&mut self) -> Result<String> {
        match self.read_token() {
            Some(token) => Ok(token),
            None => raise!(self, "expected a token"),
        }
    }

    fn get_natural(&mut self) -> Result<uint> {
        match self.read_natural() {
            Some(number) => Ok(number),
            None => raise!(self, "expected a natural number"),
        }
    }

    fn get_id(&mut self) -> Result<uint> {
        match self.read_id() {
            Some(id) => Ok(id),
            None => raise!(self, "expected an id"),
        }
    }

    #[inline]
    fn peek(&mut self) -> Option<char> {
        match self.cursor.peek() {
            Some(&(_, c)) => Some(c),
            None => None,
        }
    }
}

impl<'a> std::iter::Iterator<char> for Parser<'a> {
    fn next(&mut self) -> Option<char> {
        match self.cursor.next() {
            Some((_, '\n')) => {
                self.line += 1;
                Some('\n')
            },
            Some((_, c)) => Some(c),
            None => None,
        }
    }
}

#[cfg(test)]
mod tests {
    macro_rules! assert_ok(
        ($result: expr) => (
            if let Err(err) = $result {
                assert!(false, "{}", err);
            }
        );
    )

    macro_rules! assert_error(
        ($result: expr) => (
            if let Ok(_) = $result {
                assert!(false, "expected an error");
            }
        );
    )

    macro_rules! parser(
        ($input:expr) => (super::Parser::new($input));
    )

    #[test]
    fn process_at() {
        assert_ok!(parser!("@abc 12").process_at());
        assert_error!(parser!("@ ").process_at());
        assert_error!(parser!("@abc").process_at());
    }

    #[test]
    fn process_block() {
        assert_ok!(parser!("{}").process_block(String::from_str("life"), 42));
    }

    #[test]
    fn process_graph() {
        let mut parser = parser!("TASK t0_0\tTYPE 2   ");
        parser.process_graph(String::new(), 0);
        {
            let ref task = parser.content.graphs[0].tasks[0];
            assert_eq!(task.id, 0);
            assert_eq!(task.kind, 2);
        }

        parser = parser!("ARC a0_42 \tFROM t0_0  TO  t0_1 TYPE 35   ");
        parser.process_graph(String::new(), 0);
        {
            let ref arc = parser.content.graphs[0].arcs[0];
            assert_eq!(arc.id, 42);
            assert_eq!(arc.from, 0);
            assert_eq!(arc.to, 1);
            assert_eq!(arc.kind, 35);
        }

        parser = parser!("HARD_DEADLINE d0_9 ON t0_12 AT 1000   ");
        parser.process_graph(String::new(), 0);
        {
            let ref deadline = parser.content.graphs[0].deadlines[0];
            assert_eq!(deadline.id, 9);
            assert_eq!(deadline.on, 12);
            assert_eq!(deadline.at, 1000);
        }
    }

    #[test]
    fn process_table() {
    }

    #[test]
    fn skip_void() {
        let mut parser = parser!("  \t  abc");
        parser.skip_void();
        assert_eq!(parser.next().unwrap(), 'a');
    }

    #[test]
    fn get_token() {
        macro_rules! test(
            ($input:expr, $output:expr) => (
                assert_eq!(parser!($input).get_token().unwrap(),
                           String::from_str($output));
            );
        )

        test!("AZ xyz", "AZ");
        test!("az xyz", "az");
        test!("AZ_az_09 xyz", "AZ_az_09");
    }

    #[test]
    fn get_natural() {
        assert_eq!(parser!("09").get_natural().unwrap(), 9);
    }

    #[test]
    fn get_id() {
        assert_eq!(parser!("t0_42").get_id().unwrap(), 42);
    }
}
<|endoftext|>


<commit_before>use std::collections::{BTreeSet, VecDeque};

pub fn puzzle(input: &str) -> u32 {
    // let mut floor1 = BTreeSet::new();
    // floor1.insert(Component::Generator(Element::Thulium));
    // floor1.insert(Component::Microchip(Element::Thulium));
    // floor1.insert(Component::Generator(Element::Plutonium));
    // floor1.insert(Component::Generator(Element::Strontium));
    //
    // let mut floor2 = BTreeSet::new();
    // floor2.insert(Component::Microchip(Element::Plutonium));
    // floor2.insert(Component::Microchip(Element::Strontium));
    //
    // let mut floor3 = BTreeSet::new();
    // floor3.insert(Component::Generator(Element::Promethium));
    // floor3.insert(Component::Microchip(Element::Promethium));
    // floor3.insert(Component::Generator(Element::Ruthenium));
    // floor3.insert(Component::Microchip(Element::Ruthenium));

    let mut floor1 = BTreeSet::new();
    floor1.insert(Component::Microchip(Element::Hydrogen));
    floor1.insert(Component::Microchip(Element::Lithium));

    let mut floor2 = BTreeSet::new();
    floor2.insert(Component::Generator(Element::Hydrogen));

    let mut floor3 = BTreeSet::new();
    floor3.insert(Component::Generator(Element::Lithium));

    let initial_building_state = BuildingState {
        elevator_floor: 0,
        floors: [
            floor1,
            floor2,
            floor3,
            BTreeSet::new(),
        ],
    };
    let initial_world_state = WorldState {
        steps: 0,
        building: initial_building_state.clone()
    };
    let mut queue = VecDeque::new();
    queue.push_back(initial_world_state);
    let mut seen = BTreeSet::new();
    seen.insert(initial_building_state);
    let mut steps = 0;

    while let Some(world) = queue.pop_front() {
        if world.steps != steps {
            println!("{} step(s), queue: {}", world.steps, queue.len());
            steps = world.steps;
        }

        seen.insert(world.building.clone());

        if world.steps == 100 {
            panic!("too far!");
        }
        if world.in_end_state() {
            return world.steps;
        }

        for next_move in world.next_moves(&seen) {
            queue.push_back(next_move);
        }
    }

    panic!("Exhausted all possible moves without finding end condition!");
}

#[derive(Debug)]
pub struct WorldState {
    steps: u32,
    building: BuildingState,
}

impl WorldState {
    pub fn in_end_state(&self) -> bool {
        self.building.in_end_state()
    }

    pub fn next_moves(&self, seen: &BTreeSet<BuildingState>) -> Vec<WorldState> {
        self.building.next_moves(seen).into_iter().map(|b| {
            WorldState {
                steps: self.steps + 1,
                building: b,
            }
        }).collect()
    }
}

#[derive(PartialEq, Eq, Copy, Clone, Debug, PartialOrd, Ord, Hash)]
pub enum Component {
    Microchip(Element),
    Generator(Element),
}

#[derive(PartialEq, Eq, Copy, Clone, Debug, PartialOrd, Ord, Hash)]
pub enum Element {
    Hydrogen,
    Lithium,
    Thulium,
    Plutonium,
    Strontium,
    Promethium,
    Ruthenium,
}

#[derive(Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct BuildingState {
    elevator_floor: usize,
    floors: [BTreeSet<Component>; 4],
}

impl Clone for BuildingState {
    fn clone(&self) -> BuildingState {
        BuildingState {
            elevator_floor: self.elevator_floor,
            floors: [
                self.floors[0].clone(),
                self.floors[1].clone(),
                self.floors[2].clone(),
                self.floors[3].clone(),
            ]
        }
    }
}

impl BuildingState {
    pub fn in_end_state(&self) -> bool {
        self.floors[0].is_empty() &&
            self.floors[1].is_empty() &&
            self.floors[2].is_empty()
    }

    pub fn next_floors(&self) -> Vec<usize> {
        match self.elevator_floor {
            0 => vec![1],
            1 => vec![0, 2],
            2 => vec![1, 3],
            3 => vec![2],
            other => panic!("How did you get on floor {}, elevator?", other),
        }
    }

    pub fn elevator_combos(&self) -> BTreeSet<BTreeSet<Component>> {
        let mut result = BTreeSet::new();
        for &component0 in self.floors[self.elevator_floor].iter() {
            // Could take just this item on the elevator
            let mut set = BTreeSet::new();
            set.insert(component0);
            result.insert(set.clone());

            // Could take this item and one other, if they don't fry
            for &component1 in self.floors[self.elevator_floor].iter() {
                let mut scoped_set = set.clone();
                scoped_set.insert(component1);
                if !is_fried(&scoped_set) {
                    result.insert(scoped_set);
                }
            }
        }
        result
    }

    pub fn next_moves(&self, seen: &BTreeSet<BuildingState>)
                      -> Vec<BuildingState> {
        let mut valid_next_moves = vec![];
        for next_floor in self.next_floors() {
            for combo in self.elevator_combos() {
                let mut bs = (*self).clone();
                for &c in combo.iter() {
                    bs.floors[self.elevator_floor].remove(&c);
                    bs.floors[next_floor].insert(c);
                }
                bs.elevator_floor = next_floor;
                if !bs.has_fried_chips() && !seen.contains(&bs) {
                    valid_next_moves.push(bs);
                }
            }
        }
        valid_next_moves
    }

    pub fn has_fried_chips(&self) -> bool {
        for floor in self.floors.iter() {
            if is_fried(floor) {
                return true
            }
        }
        false
    }
}

pub fn is_fried(components: &BTreeSet<Component>) -> bool {
    let mut generator_elements = BTreeSet::new();
    let mut microchip_elements = BTreeSet::new();

    for &component in components.iter() {
        match component {
            Component::Generator(el) => {
                generator_elements.insert(el);
            },
            Component::Microchip(el) => {
                microchip_elements.insert(el);
            },
        }
    }

    if generator_elements.is_empty() {
        return false;
    }

    microchip_elements.difference(&generator_elements).next().is_some()
}

#[cfg(test)]
mod test {
    use super::*;
    use std::collections::BTreeSet;

    #[test]
    fn not_in_end_state() {
        let mut floor1 = BTreeSet::new();
        floor1.insert(Component::Microchip(Element::Hydrogen));
        floor1.insert(Component::Microchip(Element::Lithium));

        let mut floor2 = BTreeSet::new();
        floor2.insert(Component::Generator(Element::Hydrogen));

        let mut floor3 = BTreeSet::new();
        floor3.insert(Component::Generator(Element::Lithium));

        let world_state = WorldState {
            steps: 0,
            building: BuildingState {
                elevator_floor: 0,
                floors: [
                    floor1,
                    floor2,
                    floor3,
                    BTreeSet::new(),
                ],
            }
        };
        assert!( ! world_state.in_end_state() );
    }

    #[test]
    fn in_end_state() {
        let mut floor4 = BTreeSet::new();
        floor4.insert(Component::Generator(Element::Hydrogen));
        floor4.insert(Component::Microchip(Element::Hydrogen));
        floor4.insert(Component::Microchip(Element::Lithium));
        floor4.insert(Component::Generator(Element::Lithium));
        let world_state = WorldState {
            steps: 0,
            building: BuildingState {
                elevator_floor: 0,
                floors: [
                    BTreeSet::new(),
                    BTreeSet::new(),
                    BTreeSet::new(),
                    floor4,
                ],
            }
        };

        assert!(world_state.in_end_state());
    }

    #[test]
    fn two_microchips_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Microchip(Element::Lithium));
        hs.insert(Component::Microchip(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn microchip_alone_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Microchip(Element::Lithium));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn two_generators_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Generator(Element::Lithium));
        hs.insert(Component::Generator(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn microchip_and_matching_generator_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Generator(Element::Hydrogen));
        hs.insert(Component::Microchip(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn microchip_attached_to_matching_generator_not_fried_by_other_generator() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Generator(Element::Lithium));
        hs.insert(Component::Generator(Element::Hydrogen));
        hs.insert(Component::Microchip(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn microchip_without_its_generator_with_another_generator_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Generator(Element::Lithium));
        hs.insert(Component::Microchip(Element::Hydrogen));
        assert!(is_fried(&hs));
    }

    #[test]
    fn generator_with_a_microchip_without_its_generator_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Microchip(Element::Lithium));
        hs.insert(Component::Generator(Element::Hydrogen));
        assert!(is_fried(&hs));
    }

    #[test]
    fn generator_with_a_generator_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Generator(Element::Lithium));
        hs.insert(Component::Generator(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn generator_with_its_microchip_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Microchip(Element::Hydrogen));
        hs.insert(Component::Generator(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn generator_with_a_microchip_with_its_generator_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Generator(Element::Lithium));
        hs.insert(Component::Microchip(Element::Lithium));
        hs.insert(Component::Generator(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }
}
<commit_msg>Use a nice iterator method<commit_after>use std::collections::{BTreeSet, VecDeque};

pub fn puzzle(input: &str) -> u32 {
    // let mut floor1 = BTreeSet::new();
    // floor1.insert(Component::Generator(Element::Thulium));
    // floor1.insert(Component::Microchip(Element::Thulium));
    // floor1.insert(Component::Generator(Element::Plutonium));
    // floor1.insert(Component::Generator(Element::Strontium));
    //
    // let mut floor2 = BTreeSet::new();
    // floor2.insert(Component::Microchip(Element::Plutonium));
    // floor2.insert(Component::Microchip(Element::Strontium));
    //
    // let mut floor3 = BTreeSet::new();
    // floor3.insert(Component::Generator(Element::Promethium));
    // floor3.insert(Component::Microchip(Element::Promethium));
    // floor3.insert(Component::Generator(Element::Ruthenium));
    // floor3.insert(Component::Microchip(Element::Ruthenium));

    let mut floor1 = BTreeSet::new();
    floor1.insert(Component::Microchip(Element::Hydrogen));
    floor1.insert(Component::Microchip(Element::Lithium));

    let mut floor2 = BTreeSet::new();
    floor2.insert(Component::Generator(Element::Hydrogen));

    let mut floor3 = BTreeSet::new();
    floor3.insert(Component::Generator(Element::Lithium));

    let initial_building_state = BuildingState {
        elevator_floor: 0,
        floors: [
            floor1,
            floor2,
            floor3,
            BTreeSet::new(),
        ],
    };
    let initial_world_state = WorldState {
        steps: 0,
        building: initial_building_state.clone()
    };
    let mut queue = VecDeque::new();
    queue.push_back(initial_world_state);
    let mut seen = BTreeSet::new();
    seen.insert(initial_building_state);
    let mut steps = 0;

    while let Some(world) = queue.pop_front() {
        if world.steps != steps {
            println!("{} step(s), queue: {}", world.steps, queue.len());
            steps = world.steps;
        }

        seen.insert(world.building.clone());

        if world.steps == 100 {
            panic!("too far!");
        }
        if world.in_end_state() {
            return world.steps;
        }

        for next_move in world.next_moves(&seen) {
            queue.push_back(next_move);
        }
    }

    panic!("Exhausted all possible moves without finding end condition!");
}

#[derive(Debug)]
pub struct WorldState {
    steps: u32,
    building: BuildingState,
}

impl WorldState {
    pub fn in_end_state(&self) -> bool {
        self.building.in_end_state()
    }

    pub fn next_moves(&self, seen: &BTreeSet<BuildingState>) -> Vec<WorldState> {
        self.building.next_moves(seen).into_iter().map(|b| {
            WorldState {
                steps: self.steps + 1,
                building: b,
            }
        }).collect()
    }
}

#[derive(PartialEq, Eq, Copy, Clone, Debug, PartialOrd, Ord, Hash)]
pub enum Component {
    Microchip(Element),
    Generator(Element),
}

#[derive(PartialEq, Eq, Copy, Clone, Debug, PartialOrd, Ord, Hash)]
pub enum Element {
    Hydrogen,
    Lithium,
    Thulium,
    Plutonium,
    Strontium,
    Promethium,
    Ruthenium,
}

#[derive(Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct BuildingState {
    elevator_floor: usize,
    floors: [BTreeSet<Component>; 4],
}

impl Clone for BuildingState {
    fn clone(&self) -> BuildingState {
        BuildingState {
            elevator_floor: self.elevator_floor,
            floors: [
                self.floors[0].clone(),
                self.floors[1].clone(),
                self.floors[2].clone(),
                self.floors[3].clone(),
            ]
        }
    }
}

impl BuildingState {
    pub fn in_end_state(&self) -> bool {
        self.floors[0].is_empty() &&
            self.floors[1].is_empty() &&
            self.floors[2].is_empty()
    }

    pub fn next_floors(&self) -> Vec<usize> {
        match self.elevator_floor {
            0 => vec![1],
            1 => vec![0, 2],
            2 => vec![1, 3],
            3 => vec![2],
            other => panic!("How did you get on floor {}, elevator?", other),
        }
    }

    pub fn elevator_combos(&self) -> BTreeSet<BTreeSet<Component>> {
        let mut result = BTreeSet::new();
        for &component0 in self.floors[self.elevator_floor].iter() {
            // Could take just this item on the elevator
            let mut set = BTreeSet::new();
            set.insert(component0);
            result.insert(set.clone());

            // Could take this item and one other, if they don't fry
            for &component1 in self.floors[self.elevator_floor].iter() {
                let mut scoped_set = set.clone();
                scoped_set.insert(component1);
                if !is_fried(&scoped_set) {
                    result.insert(scoped_set);
                }
            }
        }
        result
    }

    pub fn next_moves(&self, seen: &BTreeSet<BuildingState>)
                      -> Vec<BuildingState> {
        let mut valid_next_moves = vec![];
        for next_floor in self.next_floors() {
            for combo in self.elevator_combos() {
                let mut bs = (*self).clone();
                for &c in combo.iter() {
                    bs.floors[self.elevator_floor].remove(&c);
                    bs.floors[next_floor].insert(c);
                }
                bs.elevator_floor = next_floor;
                if !bs.has_fried_chips() && !seen.contains(&bs) {
                    valid_next_moves.push(bs);
                }
            }
        }
        valid_next_moves
    }

    pub fn has_fried_chips(&self) -> bool {
        self.floors.iter().any(is_fried)
    }
}

pub fn is_fried(components: &BTreeSet<Component>) -> bool {
    let mut generator_elements = BTreeSet::new();
    let mut microchip_elements = BTreeSet::new();

    for &component in components.iter() {
        match component {
            Component::Generator(el) => {
                generator_elements.insert(el);
            },
            Component::Microchip(el) => {
                microchip_elements.insert(el);
            },
        }
    }

    if generator_elements.is_empty() {
        return false;
    }

    microchip_elements.difference(&generator_elements).next().is_some()
}

#[cfg(test)]
mod test {
    use super::*;
    use std::collections::BTreeSet;

    #[test]
    fn not_in_end_state() {
        let mut floor1 = BTreeSet::new();
        floor1.insert(Component::Microchip(Element::Hydrogen));
        floor1.insert(Component::Microchip(Element::Lithium));

        let mut floor2 = BTreeSet::new();
        floor2.insert(Component::Generator(Element::Hydrogen));

        let mut floor3 = BTreeSet::new();
        floor3.insert(Component::Generator(Element::Lithium));

        let world_state = WorldState {
            steps: 0,
            building: BuildingState {
                elevator_floor: 0,
                floors: [
                    floor1,
                    floor2,
                    floor3,
                    BTreeSet::new(),
                ],
            }
        };
        assert!( ! world_state.in_end_state() );
    }

    #[test]
    fn in_end_state() {
        let mut floor4 = BTreeSet::new();
        floor4.insert(Component::Generator(Element::Hydrogen));
        floor4.insert(Component::Microchip(Element::Hydrogen));
        floor4.insert(Component::Microchip(Element::Lithium));
        floor4.insert(Component::Generator(Element::Lithium));
        let world_state = WorldState {
            steps: 0,
            building: BuildingState {
                elevator_floor: 0,
                floors: [
                    BTreeSet::new(),
                    BTreeSet::new(),
                    BTreeSet::new(),
                    floor4,
                ],
            }
        };

        assert!(world_state.in_end_state());
    }

    #[test]
    fn two_microchips_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Microchip(Element::Lithium));
        hs.insert(Component::Microchip(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn microchip_alone_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Microchip(Element::Lithium));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn two_generators_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Generator(Element::Lithium));
        hs.insert(Component::Generator(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn microchip_and_matching_generator_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Generator(Element::Hydrogen));
        hs.insert(Component::Microchip(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn microchip_attached_to_matching_generator_not_fried_by_other_generator() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Generator(Element::Lithium));
        hs.insert(Component::Generator(Element::Hydrogen));
        hs.insert(Component::Microchip(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn microchip_without_its_generator_with_another_generator_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Generator(Element::Lithium));
        hs.insert(Component::Microchip(Element::Hydrogen));
        assert!(is_fried(&hs));
    }

    #[test]
    fn generator_with_a_microchip_without_its_generator_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Microchip(Element::Lithium));
        hs.insert(Component::Generator(Element::Hydrogen));
        assert!(is_fried(&hs));
    }

    #[test]
    fn generator_with_a_generator_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Generator(Element::Lithium));
        hs.insert(Component::Generator(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn generator_with_its_microchip_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Microchip(Element::Hydrogen));
        hs.insert(Component::Generator(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }

    #[test]
    fn generator_with_a_microchip_with_its_generator_not_fried() {
        let mut hs = BTreeSet::new();
        hs.insert(Component::Generator(Element::Lithium));
        hs.insert(Component::Microchip(Element::Lithium));
        hs.insert(Component::Generator(Element::Hydrogen));
        assert!( ! is_fried(&hs) );
    }
}
<|endoftext|>