<commit_before>extern crate crypto;
extern crate flate2;
extern crate num_cpus;

mod padding;

use crypto::digest::Digest;
use crypto::sha1;

use flate2::Compression;
use flate2::write::ZlibEncoder;

use std::env;
use std::fs;
use std::io::Write;
use std::ops;
use std::process;
use std::str;
use std::sync::mpsc;
use std::thread;
use std::u64;
use std::u8;

const SHA1_BYTE_LENGTH: usize = 20;

struct HashPrefix {
    data: Vec<u8>,
    half_byte: Option<u8>,
}

impl Clone for HashPrefix {
    fn clone(&self) -> Self {
        HashPrefix {
            data: self.data.to_owned(),
            half_byte: self.half_byte.to_owned()
        }
    }
}

struct SearchParams {
    current_message: String,
    desired_prefix: HashPrefix,
    counter_range: ops::Range<u64>,
    extension_word_length: usize
}

struct HashMatch {
    data: Vec<u8>,
    hash: [u8; SHA1_BYTE_LENGTH],
}

struct ProcessedCommitMessage {
    full_message: Vec<u8>,
    whitespace_index: usize,
}

fn main() {
    let args: Vec<String> = env::args().collect();

    match args.len() {
        1 => run_lucky_commit(&parse_prefix("0000000").unwrap()),
        2 => match parse_prefix(&args[1]) {
            Some(prefix) => run_lucky_commit(&prefix),
            None => print_usage_and_exit()
        },
        _ => print_usage_and_exit()
    }
}

fn print_usage_and_exit() {
    fail_with_message("Usage: lucky-commit [commit-hash-prefix]");
}

fn fail_with_message(message: &str) {
    eprintln!("{}", message);
    process::exit(1);
}

fn parse_prefix(prefix: &str) -> Option<HashPrefix> {
    let mut data = Vec::new();
    for index in 0..(prefix.len() / 2) {
        match u8::from_str_radix(&prefix[2 * index..2 * index + 2], 16) {
            Ok(value) => data.push(value),
            Err(_) => return None
        }
    }

    let parsed_prefix = HashPrefix {
        data,
        half_byte: if prefix.len() % 2 == 1 {
            match u8::from_str_radix(&prefix[prefix.len() - 1..], 16) {
                Ok(value) => Some(value << 4),
                Err(_) => return None
            }
        } else {
            None
        }
    };

    Some(parsed_prefix)
}

fn run_lucky_commit(desired_prefix: &HashPrefix) {
    let current_message_bytes = run_command("git", &["cat-file", "commit", "HEAD"]);
    let current_message = &String::from_utf8(current_message_bytes)
        .expect("Git commit contains invalid utf8");

    match find_match(current_message, desired_prefix) {
        Some(hash_match) => {
            create_git_object_file(&hash_match);
            git_reset_to_hash(&hash_match.hash);
        },
        None => fail_with_message("Failed to find a match")
    }
}

fn run_command(command: &str, args: &[&str]) -> Vec<u8> {
    let output = process::Command::new(command)
        .args(args)
        .stderr(process::Stdio::inherit())
        .output()
        .expect("Failed to run command");

    if !output.status.success() {
        process::exit(1);
    }

    output.stdout
}

fn find_match(current_message: &str, desired_prefix: &HashPrefix) -> Option<HashMatch> {
    let num_threads = num_cpus::get_physical();
    let (shared_sender, receiver) = mpsc::channel();
    let u32_ranges = split_range(0, 1u64 << 32, num_threads);
    let u64_ranges = split_range(0, u64::MAX, num_threads);

    for thread_index in 0..num_threads {
        spawn_hash_searcher(
            mpsc::Sender::clone(&shared_sender),
            SearchParams {
                current_message: current_message.to_owned(),
                desired_prefix: desired_prefix.clone(),
                counter_range: u32_ranges[thread_index].clone(),
                extension_word_length: 4
            }
        );
    }

    for thread_index in 0..num_threads * 2 {
        let result = receiver.recv().unwrap();
        if result.is_some() {
            return result;
        }

        if thread_index < num_threads {
            spawn_hash_searcher(
                mpsc::Sender::clone(&shared_sender),
                SearchParams {
                    current_message: current_message.to_owned(),
                    desired_prefix: desired_prefix.clone(),
                    counter_range: u64_ranges[thread_index].clone(),
                    extension_word_length: 8
                }
            );
        }
    }

    None
}

fn spawn_hash_searcher(result_sender: mpsc::Sender<Option<HashMatch>>, params: SearchParams) {
    thread::spawn(move || {
        match result_sender.send(iterate_for_match(&params)) {
            /*
             * If an error occurs when sending, then the receiver has already received
             * a match from another thread, so ignore the error.
             */
            Ok(_) => (),
            Err(_) => ()
        }
    });
}

fn split_range(min: u64, max: u64, num_segments: usize) -> Vec<ops::Range<u64>> {
    let segment_size = (max - min) / (num_segments as u64);

    let mut segments = Vec::new();
    let mut last_range_end = 0;
    for _ in 0..num_segments {
        segments.push(last_range_end..last_range_end + segment_size);
        last_range_end += segment_size;
    }

    segments
}

fn iterate_for_match(params: &SearchParams) -> Option<HashMatch> {
    let desired_prefix = &params.desired_prefix;
    let extension_length = params.extension_word_length * 8;
    let processed_message = process_commit_message(
        &params.current_message,
        extension_length
    );

    let mut hash_data = processed_message.full_message;
    let whitespace_index = processed_message.whitespace_index;

    let mut sha1_hash = sha1::Sha1::new();
    let mut hash_result: [u8; SHA1_BYTE_LENGTH] = [0; SHA1_BYTE_LENGTH];

    for counter in params.counter_range.clone() {
        let mut whitespace_word_index = 0;
        while whitespace_word_index < extension_length {
            let start_index = whitespace_index + whitespace_word_index;
            let padding_index = (counter >> whitespace_word_index) as u8 as usize;
            &mut hash_data[start_index..start_index + 8].copy_from_slice(&padding::PADDING_LIST[padding_index]);
            whitespace_word_index += 8;
        }

        sha1_hash.input(&hash_data);
        sha1_hash.result(&mut hash_result);

        if matches_desired_prefix(&hash_result, desired_prefix) {
            return Some(HashMatch {
                data: hash_data,
                hash: hash_result,
            })
        }

        sha1_hash.reset();
    }

    None
}

fn process_commit_message(original_message: &str, extension_length: usize) -> ProcessedCommitMessage {
    let mut message_object: Vec<u8> = format!("commit {}\x00", original_message.len() + extension_length)
        .into_bytes();

    let commit_split_index = get_commit_message_split_index(original_message);
    let whitespace_index = commit_split_index + message_object.len();

    for character in original_message[..commit_split_index].as_bytes() {
        message_object.push(*character);
    }

    for _ in 0..extension_length {
        message_object.push(padding::SPACE);
    }

    for character in original_message[commit_split_index..].as_bytes() {
        message_object.push(*character);
    }

    ProcessedCommitMessage {
        full_message: message_object,
        whitespace_index,
    }
}

fn get_commit_message_split_index(message: &str) -> usize {
    /*
     * If the commit has a GPG signature (detected by the presence of "gpgsig " at the start
     * of the fifth line), then add the padding whitespace immediately after the text "gpgsig ".
     * Otherwise, add the padding whitespace right before the end of the commit message.
     *
     * If a signature is present, modifying the commit message would make the signature invalid.
     */
    let mut current_line_index = 0;
    const SIGNATURE_MARKER: &str = "gpgsig ";
    for (index, character) in message.chars().enumerate() {
        if current_line_index == 4 {
            if message[index..].starts_with(SIGNATURE_MARKER) {
                return index + SIGNATURE_MARKER.len();
            } else {
                return message.len() - 1;
            }
        }
        if character == '\n' {
            current_line_index += 1;
        }
    }

    message.len() - 1
}

fn matches_desired_prefix(hash: &[u8; SHA1_BYTE_LENGTH], prefix: &HashPrefix) -> bool {
    prefix.data == &hash[..prefix.data.len()] && match prefix.half_byte {
        Some(half_byte) => (hash[prefix.data.len()] & 0xf0) == half_byte,
        None => true,
    }
}

fn create_git_object_file(search_result: &HashMatch) {
    let compressed_object = zlib_compress(&search_result.data);
    let dir_path = format!(".git/objects/{:02x}", search_result.hash[0]);
    let file_path = format!(
        "{}/{}",
        dir_path,
        to_hex_string(&search_result.hash[1..])
    );

    fs::DirBuilder::new()
        .recursive(true)
        .create(dir_path)
        .expect("Failed to create git object directory");

    fs::File::create(file_path)
        .expect("Failed to open git object file")
        .write_all(&compressed_object)
        .expect("Failed to write git object file");
}

fn zlib_compress(data: &[u8]) -> Vec<u8> {
    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
    encoder.write(data).expect("zlib compression failed");
    encoder.finish().expect("zlib compression failed")
}

fn git_reset_to_hash(hash: &[u8; SHA1_BYTE_LENGTH]) {
    run_command("git", &["reset", &to_hex_string(hash)]);
}

fn to_hex_string(hash: &[u8]) -> String {
    hash.iter().map(|byte| format!("{:02x}", *byte)).collect::<String>()
}

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

    #[test]
    fn matches_desired_prefix_empty() {
        assert!(
            matches_desired_prefix(
                &[0; SHA1_BYTE_LENGTH],
                &HashPrefix {
                    data: Vec::new(),
                    half_byte: None
                }
            )
        )
    }

    #[test]
    fn matches_desired_prefix_single_half() {
        assert!(
            matches_desired_prefix(
                &[0x1e; SHA1_BYTE_LENGTH],
                &HashPrefix {
                    data: Vec::new(),
                    half_byte: Some(0x10)
                }
            )
        )
    }

    #[test]
    fn matches_desired_prefix_single_half_mismatch() {
        assert!(
            !matches_desired_prefix(
                &[0x21; SHA1_BYTE_LENGTH],
                &HashPrefix {
                    data: Vec::new(),
                    half_byte: Some(0x10)
                }
            )
        )
    }

    #[test]
    fn matches_desired_prefix_data_without_half() {
        assert!(
            matches_desired_prefix(
                &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
                &HashPrefix {
                    data: vec![1, 2, 3],
                    half_byte: None
                }
            )
        )
    }

    #[test]
    fn matches_desired_prefix_matching_data_and_half() {
        assert!(
            matches_desired_prefix(
                &[1, 2, 3, 0x4f, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
                &HashPrefix {
                    data: vec![1, 2, 3],
                    half_byte: Some(0x40)
                }
            )
        )
    }

    #[test]
    fn matches_desired_prefix_matching_data_mismatching_half() {
        assert!(
            !matches_desired_prefix(
                &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
                &HashPrefix {
                    data: vec![1, 2, 3],
                    half_byte: Some(0x50)
                }
            )
        )
    }

    #[test]
    fn matches_desired_prefix_mismatching_data_matching_half() {
        assert!(
            !matches_desired_prefix(
                &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
                &HashPrefix {
                    data: vec![1, 5, 3],
                    half_byte: Some(0x40)
                }
            )
        )
    }

    #[test]
    fn to_hex_string_basic() {
        assert_eq!("00", to_hex_string(&[0]));
    }

    #[test]
    fn to_hex_string_multichar() {
        assert_eq!("00ff14", to_hex_string(&[0, 255, 20]));
    }
}
<commit_msg>Add tests for parse_prefix<commit_after>extern crate crypto;
extern crate flate2;
extern crate num_cpus;

mod padding;

use crypto::digest::Digest;
use crypto::sha1;

use flate2::Compression;
use flate2::write::ZlibEncoder;

use std::env;
use std::fs;
use std::io::Write;
use std::ops;
use std::process;
use std::str;
use std::sync::mpsc;
use std::thread;
use std::u64;
use std::u8;

const SHA1_BYTE_LENGTH: usize = 20;

#[derive(Debug)]
#[derive(PartialEq)]
struct HashPrefix {
    data: Vec<u8>,
    half_byte: Option<u8>,
}

impl Clone for HashPrefix {
    fn clone(&self) -> Self {
        HashPrefix {
            data: self.data.to_owned(),
            half_byte: self.half_byte.to_owned()
        }
    }
}

struct SearchParams {
    current_message: String,
    desired_prefix: HashPrefix,
    counter_range: ops::Range<u64>,
    extension_word_length: usize
}

struct HashMatch {
    data: Vec<u8>,
    hash: [u8; SHA1_BYTE_LENGTH],
}

struct ProcessedCommitMessage {
    full_message: Vec<u8>,
    whitespace_index: usize,
}

fn main() {
    let args: Vec<String> = env::args().collect();

    match args.len() {
        1 => run_lucky_commit(&parse_prefix("0000000").unwrap()),
        2 => match parse_prefix(&args[1]) {
            Some(prefix) => run_lucky_commit(&prefix),
            None => print_usage_and_exit()
        },
        _ => print_usage_and_exit()
    }
}

fn print_usage_and_exit() {
    fail_with_message("Usage: lucky-commit [commit-hash-prefix]");
}

fn fail_with_message(message: &str) {
    eprintln!("{}", message);
    process::exit(1);
}

fn parse_prefix(prefix: &str) -> Option<HashPrefix> {
    let mut data = Vec::new();
    for index in 0..(prefix.len() / 2) {
        match u8::from_str_radix(&prefix[2 * index..2 * index + 2], 16) {
            Ok(value) => data.push(value),
            Err(_) => return None
        }
    }

    let parsed_prefix = HashPrefix {
        data,
        half_byte: if prefix.len() % 2 == 1 {
            match u8::from_str_radix(&prefix[prefix.len() - 1..], 16) {
                Ok(value) => Some(value << 4),
                Err(_) => return None
            }
        } else {
            None
        }
    };

    Some(parsed_prefix)
}

fn run_lucky_commit(desired_prefix: &HashPrefix) {
    let current_message_bytes = run_command("git", &["cat-file", "commit", "HEAD"]);
    let current_message = &String::from_utf8(current_message_bytes)
        .expect("Git commit contains invalid utf8");

    match find_match(current_message, desired_prefix) {
        Some(hash_match) => {
            create_git_object_file(&hash_match);
            git_reset_to_hash(&hash_match.hash);
        },
        None => fail_with_message("Failed to find a match")
    }
}

fn run_command(command: &str, args: &[&str]) -> Vec<u8> {
    let output = process::Command::new(command)
        .args(args)
        .stderr(process::Stdio::inherit())
        .output()
        .expect("Failed to run command");

    if !output.status.success() {
        process::exit(1);
    }

    output.stdout
}

fn find_match(current_message: &str, desired_prefix: &HashPrefix) -> Option<HashMatch> {
    let num_threads = num_cpus::get_physical();
    let (shared_sender, receiver) = mpsc::channel();
    let u32_ranges = split_range(0, 1u64 << 32, num_threads);
    let u64_ranges = split_range(0, u64::MAX, num_threads);

    for thread_index in 0..num_threads {
        spawn_hash_searcher(
            mpsc::Sender::clone(&shared_sender),
            SearchParams {
                current_message: current_message.to_owned(),
                desired_prefix: desired_prefix.clone(),
                counter_range: u32_ranges[thread_index].clone(),
                extension_word_length: 4
            }
        );
    }

    for thread_index in 0..num_threads * 2 {
        let result = receiver.recv().unwrap();
        if result.is_some() {
            return result;
        }

        if thread_index < num_threads {
            spawn_hash_searcher(
                mpsc::Sender::clone(&shared_sender),
                SearchParams {
                    current_message: current_message.to_owned(),
                    desired_prefix: desired_prefix.clone(),
                    counter_range: u64_ranges[thread_index].clone(),
                    extension_word_length: 8
                }
            );
        }
    }

    None
}

fn spawn_hash_searcher(result_sender: mpsc::Sender<Option<HashMatch>>, params: SearchParams) {
    thread::spawn(move || {
        match result_sender.send(iterate_for_match(&params)) {
            /*
             * If an error occurs when sending, then the receiver has already received
             * a match from another thread, so ignore the error.
             */
            Ok(_) => (),
            Err(_) => ()
        }
    });
}

fn split_range(min: u64, max: u64, num_segments: usize) -> Vec<ops::Range<u64>> {
    let segment_size = (max - min) / (num_segments as u64);

    let mut segments = Vec::new();
    let mut last_range_end = 0;
    for _ in 0..num_segments {
        segments.push(last_range_end..last_range_end + segment_size);
        last_range_end += segment_size;
    }

    segments
}

fn iterate_for_match(params: &SearchParams) -> Option<HashMatch> {
    let desired_prefix = &params.desired_prefix;
    let extension_length = params.extension_word_length * 8;
    let processed_message = process_commit_message(
        &params.current_message,
        extension_length
    );

    let mut hash_data = processed_message.full_message;
    let whitespace_index = processed_message.whitespace_index;

    let mut sha1_hash = sha1::Sha1::new();
    let mut hash_result: [u8; SHA1_BYTE_LENGTH] = [0; SHA1_BYTE_LENGTH];

    for counter in params.counter_range.clone() {
        let mut whitespace_word_index = 0;
        while whitespace_word_index < extension_length {
            let start_index = whitespace_index + whitespace_word_index;
            let padding_index = (counter >> whitespace_word_index) as u8 as usize;
            &mut hash_data[start_index..start_index + 8].copy_from_slice(&padding::PADDING_LIST[padding_index]);
            whitespace_word_index += 8;
        }

        sha1_hash.input(&hash_data);
        sha1_hash.result(&mut hash_result);

        if matches_desired_prefix(&hash_result, desired_prefix) {
            return Some(HashMatch {
                data: hash_data,
                hash: hash_result,
            })
        }

        sha1_hash.reset();
    }

    None
}

fn process_commit_message(original_message: &str, extension_length: usize) -> ProcessedCommitMessage {
    let mut message_object: Vec<u8> = format!("commit {}\x00", original_message.len() + extension_length)
        .into_bytes();

    let commit_split_index = get_commit_message_split_index(original_message);
    let whitespace_index = commit_split_index + message_object.len();

    for character in original_message[..commit_split_index].as_bytes() {
        message_object.push(*character);
    }

    for _ in 0..extension_length {
        message_object.push(padding::SPACE);
    }

    for character in original_message[commit_split_index..].as_bytes() {
        message_object.push(*character);
    }

    ProcessedCommitMessage {
        full_message: message_object,
        whitespace_index,
    }
}

fn get_commit_message_split_index(message: &str) -> usize {
    /*
     * If the commit has a GPG signature (detected by the presence of "gpgsig " at the start
     * of the fifth line), then add the padding whitespace immediately after the text "gpgsig ".
     * Otherwise, add the padding whitespace right before the end of the commit message.
     *
     * If a signature is present, modifying the commit message would make the signature invalid.
     */
    let mut current_line_index = 0;
    const SIGNATURE_MARKER: &str = "gpgsig ";
    for (index, character) in message.chars().enumerate() {
        if current_line_index == 4 {
            if message[index..].starts_with(SIGNATURE_MARKER) {
                return index + SIGNATURE_MARKER.len();
            } else {
                return message.len() - 1;
            }
        }
        if character == '\n' {
            current_line_index += 1;
        }
    }

    message.len() - 1
}

fn matches_desired_prefix(hash: &[u8; SHA1_BYTE_LENGTH], prefix: &HashPrefix) -> bool {
    prefix.data == &hash[..prefix.data.len()] && match prefix.half_byte {
        Some(half_byte) => (hash[prefix.data.len()] & 0xf0) == half_byte,
        None => true,
    }
}

fn create_git_object_file(search_result: &HashMatch) {
    let compressed_object = zlib_compress(&search_result.data);
    let dir_path = format!(".git/objects/{:02x}", search_result.hash[0]);
    let file_path = format!(
        "{}/{}",
        dir_path,
        to_hex_string(&search_result.hash[1..])
    );

    fs::DirBuilder::new()
        .recursive(true)
        .create(dir_path)
        .expect("Failed to create git object directory");

    fs::File::create(file_path)
        .expect("Failed to open git object file")
        .write_all(&compressed_object)
        .expect("Failed to write git object file");
}

fn zlib_compress(data: &[u8]) -> Vec<u8> {
    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
    encoder.write(data).expect("zlib compression failed");
    encoder.finish().expect("zlib compression failed")
}

fn git_reset_to_hash(hash: &[u8; SHA1_BYTE_LENGTH]) {
    run_command("git", &["reset", &to_hex_string(hash)]);
}

fn to_hex_string(hash: &[u8]) -> String {
    hash.iter().map(|byte| format!("{:02x}", *byte)).collect::<String>()
}

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

    #[test]
    fn parse_prefix_empty() {
        assert_eq!(
            Some(HashPrefix { data: Vec::new(), half_byte: None }),
            parse_prefix("")
        )
    }

    #[test]
    fn parse_prefix_single_char() {
        assert_eq!(
            Some(HashPrefix { data: Vec::new(), half_byte: Some(0xa0) }),
            parse_prefix("a")
        )
    }

    #[test]
    fn parse_prefix_even_chars() {
        assert_eq!(
            Some(HashPrefix { data: vec![0xab, 0xcd, 0xef], half_byte: None }),
            parse_prefix("abcdef")
        )
    }

    #[test]
    fn parse_prefix_odd_chars() {
        assert_eq!(
            Some(HashPrefix { data: vec![0xab, 0xcd, 0xef], half_byte: Some(0x50) }),
            parse_prefix("abcdef5")
        )
    }

    #[test]
    fn parse_prefix_capital_letters() {
        assert_eq!(
            Some(HashPrefix { data: vec![0xab, 0xcd, 0xef], half_byte: Some(0xb0) }),
            parse_prefix("ABCDEFB")
        )
    }

    #[test]
    fn parse_prefix_invalid_even_chars() {
        assert_eq!(None, parse_prefix("abcdgeb"))
    }

    #[test]
    fn parse_prefix_invalid_odd_char() {
        assert_eq!(None, parse_prefix("abcdefg"))
    }

    #[test]
    fn matches_desired_prefix_empty() {
        assert!(
            matches_desired_prefix(
                &[0; SHA1_BYTE_LENGTH],
                &HashPrefix {
                    data: Vec::new(),
                    half_byte: None
                }
            )
        )
    }

    #[test]
    fn matches_desired_prefix_single_half() {
        assert!(
            matches_desired_prefix(
                &[0x1e; SHA1_BYTE_LENGTH],
                &HashPrefix {
                    data: Vec::new(),
                    half_byte: Some(0x10)
                }
            )
        )
    }

    #[test]
    fn matches_desired_prefix_single_half_mismatch() {
        assert!(
            !matches_desired_prefix(
                &[0x21; SHA1_BYTE_LENGTH],
                &HashPrefix {
                    data: Vec::new(),
                    half_byte: Some(0x10)
                }
            )
        )
    }

    #[test]
    fn matches_desired_prefix_data_without_half() {
        assert!(
            matches_desired_prefix(
                &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
                &HashPrefix {
                    data: vec![1, 2, 3],
                    half_byte: None
                }
            )
        )
    }

    #[test]
    fn matches_desired_prefix_matching_data_and_half() {
        assert!(
            matches_desired_prefix(
                &[1, 2, 3, 0x4f, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
                &HashPrefix {
                    data: vec![1, 2, 3],
                    half_byte: Some(0x40)
                }
            )
        )
    }

    #[test]
    fn matches_desired_prefix_matching_data_mismatching_half() {
        assert!(
            !matches_desired_prefix(
                &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
                &HashPrefix {
                    data: vec![1, 2, 3],
                    half_byte: Some(0x50)
                }
            )
        )
    }

    #[test]
    fn matches_desired_prefix_mismatching_data_matching_half() {
        assert!(
            !matches_desired_prefix(
                &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
                &HashPrefix {
                    data: vec![1, 5, 3],
                    half_byte: Some(0x40)
                }
            )
        )
    }

    #[test]
    fn to_hex_string_basic() {
        assert_eq!("00", to_hex_string(&[0]));
    }

    #[test]
    fn to_hex_string_multichar() {
        assert_eq!("00ff14", to_hex_string(&[0, 255, 20]));
    }
}
<|endoftext|>


<commit_before>extern crate fuse;
extern crate libc;
extern crate time;
extern crate hyper;
extern crate rustc_serialize;
extern crate clap;

use std::path::Path;
use std::env;
use std::io::prelude::Read;
use libc::{ENOENT, ENOSYS};
use time::Timespec;
use fuse::{FileAttr, FileType, Filesystem, Request, ReplyAttr, ReplyData, ReplyEntry, ReplyDirectory};
use hyper::client::{Client, Response};
use rustc_serialize::json::Json;
use clap::{App, Arg};

struct TwitchFileSystem;

impl Filesystem for TwitchFileSystem {
    fn getattr(&mut self, _req: &Request, ino: u64, reply: ReplyAttr) {
        let ts = Timespec::new(0, 0);
        let attr = FileAttr {
            ino: 1,
            size: 0,
            blocks: 0,
            atime: ts,
            mtime: ts,
            ctime: ts,
            crtime: ts,
            kind: FileType::Directory,
            perm: 0o755,
            nlink: 0,
            uid: 0,
            gid: 0,
            rdev: 0,
            flags: 0
        };
        let ttl = Timespec::new(1, 0);

        if ino == 1 {
            reply.attr(&ttl, &attr);
        } else {
            reply.error(ENOSYS);
        }
    }

    fn readdir(&mut self, _req: &Request, ino: u64, fh: u64, offset: u64, mut reply: ReplyDirectory) {
        if offset == 0 {
            let mut body = String::new();
            Client::new()
                .get("https://api.twitch.tv/kraken/games/top")
                .send()
                .expect("Couldn't load twitch")
                .read_to_string(&mut body);

            match Json::from_str(&body) {
                Ok(data) => {
                    let games = data.find("top").unwrap().as_array().unwrap();
                    for (i, game) in games.iter().enumerate() {
                        let name = game
                            .find_path(&["game", "name"])
                            .unwrap()
                            .as_string()
                            .unwrap();
                        reply.add(1, 1 + 1 as u64,
                                  FileType::Directory,
                                  &Path::new(name));
                    }
                },
                Err(_) => println!("Twitch returned invalid json")
            }

            reply.add(1, 0, FileType::Directory, &Path::new("."));
            reply.add(1, 1, FileType::Directory, &Path::new(".."));
        }

        reply.ok();
    }
}

fn main() {
    let matches = App::new("twitch-fs")
        .version(option_env!("CARGO_PKG_VERSION").unwrap_or("unknown version"))
        .arg(Arg::with_name("mountpoint")
             .index(1)
             .required(true))
        .get_matches();

    // unwrap() is safe here because the argument is set as required
    let mountpoint = matches.value_of_os("mountpoint").unwrap();

    fuse::mount(TwitchFileSystem, &mountpoint, &[])
}
<commit_msg>Create a constructor for TwitchFileSystem<commit_after>extern crate fuse;
extern crate libc;
extern crate time;
extern crate hyper;
extern crate rustc_serialize;
extern crate clap;

use std::collections::BTreeMap;
use std::path::Path;
use std::env;
use std::io::prelude::Read;
use libc::{ENOENT, ENOSYS};
use time::Timespec;
use fuse::{FileAttr, FileType, Filesystem, Request, ReplyAttr, ReplyData, ReplyEntry, ReplyDirectory, ReplyOpen};
use hyper::client::{Client, Response};
use rustc_serialize::json::Json;
use clap::{App, Arg};

struct TwitchFileSystem {
    attrs: BTreeMap<u64, FileAttr>,
    inodes: BTreeMap<String, u64>
}

impl TwitchFileSystem {
    fn new() -> TwitchFileSystem {
        let mut attrs = BTreeMap::new();
        let mut inodes = BTreeMap::new();
        let ts = time::now().to_timespec();
        let attr = FileAttr {
            ino: 1,
            size: 0,
            blocks: 0,
            atime: ts,
            mtime: ts,
            ctime: ts,
            crtime: ts,
            kind: FileType::Directory,
            perm: 0o755,
            nlink: 0,
            uid: 0,
            gid: 0,
            rdev: 0,
            flags: 0
        };
        attrs.insert(1, attr);
        inodes.insert("/".to_owned(), 1);
        let mut body = String::new();
        Client::new()
            .get("https://api.twitch.tv/kraken/games/top")
            .send()
            .expect("Couldn't load twitch")
            .read_to_string(&mut body);

        match Json::from_str(&body) {
            Ok(data) => {
                let games = data.find("top").unwrap().as_array().unwrap();
                for (i, game) in games.iter().enumerate() {
                    let attr = FileAttr {
                        ino: i as u64 + 2,
                        size: 0,
                        blocks: 0,
                        atime: ts,
                        mtime: ts,
                        ctime: ts,
                        crtime: ts,
                        kind: FileType::RegularFile,
                        perm: 0o644,
                        nlink: 0,
                        uid: 0,
                        gid: 0,
                        rdev: 0,
                        flags: 0
                    };
                    let name = game
                        .find_path(&["game", "name"])
                        .unwrap()
                        .as_string()
                        .unwrap();

                    attrs.insert(attr.ino, attr);
                    inodes.insert(name.to_owned(), attr.ino);
                }
            },
            Err(_) => println!("Twitch returned invalid json")
        }
        TwitchFileSystem {attrs: attrs, inodes: inodes}
    }
}

impl Filesystem for TwitchFileSystem {
    fn getattr(&mut self, _req: &Request, ino: u64, reply: ReplyAttr) {
        match self.attrs.get(&ino) {
            Some(attr) => {
                let ttl = Timespec::new(1, 0);
                reply.attr(&ttl, attr);
            },
            None => reply.error(ENOENT)
        }
    }

    fn readdir(&mut self, _req: &Request, ino: u64, fh: u64, offset: u64, mut reply: ReplyDirectory) {
        if offset == 0 {

            reply.add(1, 0, FileType::Directory, &Path::new("."));
            reply.add(1, 1, FileType::Directory, &Path::new(".."));
            reply.add(1, 29, FileType::RegularFile, &Path::new("test"));
        }

        reply.ok();
    }
    fn read(&mut self, _req: &Request, ino: u64, fh: u64, offset: u64, size: u32, reply: ReplyData) {
        reply.data("test".as_bytes());
    }
}

fn main() {
    let matches = App::new("twitch-fs")
        .version(option_env!("CARGO_PKG_VERSION").unwrap_or("unknown version"))
        .arg(Arg::with_name("mountpoint")
             .index(1)
             .required(true))
        .get_matches();

    // unwrap() is safe here because the argument is set as required
    let mountpoint = matches.value_of_os("mountpoint").unwrap();
    let fs = TwitchFileSystem::new();

    fuse::mount(fs , &mountpoint, &[])
}
<|endoftext|>


<commit_before>#![feature(deque_extras)]
#![feature(box_syntax)]
#![feature(plugin)]
#![plugin(peg_syntax_ext)]

extern crate glob;

use std::collections::HashMap;
use std::fs::File;
use std::io::{stdout, Read, Write};
use std::env;
use std::process;
use std::thread;

use self::directory_stack::DirectoryStack;
use self::input_editor::readln;
use self::peg::{parse, Job};
use self::variables::Variables;
use self::history::History;
use self::flow_control::{FlowControl, is_flow_control_command, Statement};
use self::status::{SUCCESS, NO_SUCH_COMMAND, TERMINATED};

pub mod directory_stack;
pub mod to_num;
pub mod input_editor;
pub mod peg;
pub mod variables;
pub mod history;
pub mod flow_control;
pub mod status;

/// This struct will contain all of the data structures related to this
/// instance of the shell.
pub struct Shell {
    variables: Variables,
    flow_control: FlowControl,
    directory_stack: DirectoryStack,
    history: History,
}

impl Shell {
    /// Panics if DirectoryStack construction fails
    pub fn new() -> Self {
        Shell {
            variables: Variables::new(),
            flow_control: FlowControl::new(),
            directory_stack: DirectoryStack::new().expect(""),
            history: History::new(),
        }
    }

    pub fn print_prompt(&self) {
        self.print_prompt_prefix();
        match self.flow_control.current_statement {
            Statement::For(_, _) => self.print_for_prompt(),
            Statement::Default => self.print_default_prompt(),
        }
        if let Err(message) = stdout().flush() {
            println!("{}: failed to flush prompt to stdout", message);
        }

    }

    // TODO eventually this thing should be gone
    fn print_prompt_prefix(&self) {
        let prompt_prefix = self.flow_control.modes.iter().rev().fold(String::new(), |acc, mode| {
            acc +
            if mode.value {
                "+ "
            } else {
                "- "
            }
        });
        print!("{}", prompt_prefix);
    }

    fn print_for_prompt(&self) {
        print!("for> ");
    }

    fn print_default_prompt(&self) {
        let cwd = env::current_dir().ok().map_or("?".to_string(),
                                                 |ref p| p.to_str().unwrap_or("?").to_string());
        print!("ion:{}# ", cwd);
    }

    fn on_command(&mut self, command_string: &str, commands: &HashMap<&str, Command>) {
        self.history.add(command_string.to_string());

        let mut jobs = parse(command_string);

        // Execute commands
        for job in jobs.drain(..) {
            if self.flow_control.collecting_block {
                // TODO move this logic into "end" command
                if job.command == "end" {
                    self.flow_control.collecting_block = false;
                    let block_jobs: Vec<Job> = self.flow_control
                                                   .current_block
                                                   .jobs
                                                   .drain(..)
                                                   .collect();
                    let mut variable = String::new();
                    let mut values: Vec<String> = vec![];
                    if let Statement::For(ref var, ref vals) = self.flow_control.current_statement {
                        variable = var.clone();
                        values = vals.clone();
                    }
                    for value in values {
                        self.variables.set_var(&variable, &value);
                        for job in block_jobs.iter() {
                            self.run_job(job, commands);
                        }
                    }
                    self.flow_control.current_statement = Statement::Default;
                } else {
                    self.flow_control.current_block.jobs.push(job);
                }
            } else {
                if self.flow_control.skipping() && !is_flow_control_command(&job.command) {
                    continue;
                }
                self.run_job(&job, commands);
            }
        }
    }

    fn run_job(&mut self, job: &Job, commands: &HashMap<&str, Command>) {
        let mut job = self.variables.expand_job(job);
        job.expand_globs();
        let exit_status = if let Some(command) = commands.get(job.command.as_str()) {
            Some((*command.main)(job.args.as_slice(), self))
        } else {
            self.run_external_commmand(job)
        };
        if let Some(code) = exit_status {
            self.variables.set_var("?", &code.to_string());
        }
    }

    /// Returns an exit code if a command was run
    fn run_external_commmand(&mut self, job: Job) -> Option<i32> {
        if job.background {
            thread::spawn(move || {
                let mut command = Shell::build_command(&job);
                command.stdin(process::Stdio::null());
                if let Ok(mut child) = command.spawn() {
                    Shell::wait_and_get_status(&mut child, &job.command);
                }
            });
            None
        } else {
            let mut command = Shell::build_command(&job);
            match command.spawn() {
                Ok(mut child) => Some(Shell::wait_and_get_status(&mut child, &job.command)),
                Err(err) => {
                    println!("{}: Failed to execute: {}", job.command, err);
                    Some(NO_SUCH_COMMAND)
                }
            }
        }
    }

    fn build_command(job: &Job) -> process::Command {
        let mut command = process::Command::new(&job.command);
        for i in 1..job.args.len() {
            if let Some(arg) = job.args.get(i) {
                command.arg(arg);
            }
        }
        command
    }

    // TODO don't pass in command and do printing outside this function
    fn wait_and_get_status(child: &mut process::Child, command: &str) -> i32 {
        match child.wait() {
            Ok(status) => {
                if let Some(code) = status.code() {
                    code
                } else {
                    println!("{}: child ended by signal", command);
                    TERMINATED
                }
            }
            Err(err) => {
                println!("{}: Failed to wait: {}", command, err);
                100 // TODO what should we return here?
            }
        }
    }
}

/// Structure which represents a Terminal's command.
/// This command structure contains a name, and the code which run the
/// functionnality associated to this one, with zero, one or several argument(s).
/// # Example
/// ```
/// let my_command = Command {
///     name: "my_command",
///     help: "Describe what my_command does followed by a newline showing usage",
///     main: box|args: &[String], &mut Shell| -> i32 {
///         println!("Say 'hello' to my command! :-D");
///     }
/// }
/// ```
pub struct Command {
    pub name: &'static str,
    pub help: &'static str,
    pub main: Box<Fn(&[String], &mut Shell) -> i32>,
}

impl Command {
    /// Return the map from command names to commands
    pub fn map() -> HashMap<&'static str, Self> {
        let mut commands: HashMap<&str, Self> = HashMap::new();

        commands.insert("cd",
                        Command {
                            name: "cd",
                            help: "Change the current directory\n    cd <path>",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.directory_stack.cd(args)
                            },
                        });

        commands.insert("dirs",
                        Command {
                            name: "dirs",
                            help: "Display the current directory stack",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.directory_stack.dirs(args)
                            },
                        });

        commands.insert("exit",
                        Command {
                            name: "exit",
                            help: "To exit the curent session",
                            main: box |args: &[String], _: &mut Shell| -> i32 {
                                if let Some(status) = args.get(1) {
                                    if let Ok(status) = status.parse::<i32>() {
                                        process::exit(status);
                                    }
                                }
                                // TODO should use exit status of previously run command, not 0
                                process::exit(0);
                            },
                        });

        commands.insert("let",
                        Command {
                            name: "let",
                            help: "View, set or unset variables",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.variables.let_(args)
                            },
                        });

        commands.insert("read",
                        Command {
                            name: "read",
                            help: "Read some variables\n    read <variable>",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.variables.read(args)
                            },
                        });

        commands.insert("pushd",
                        Command {
                            name: "pushd",
                            help: "Push a directory to the stack",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.directory_stack.pushd(args)
                            },
                        });

        commands.insert("popd",
                        Command {
                            name: "popd",
                            help: "Pop a directory from the stack",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.directory_stack.popd(args)
                            },
                        });

        commands.insert("history",
                        Command {
                            name: "history",
                            help: "Display a log of all commands previously executed",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.history.history(args)
                            },
                        });

        commands.insert("if",
                        Command {
                            name: "if",
                            help: "Conditionally execute code",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.flow_control.if_(args)
                            },
                        });

        commands.insert("else",
                        Command {
                            name: "else",
                            help: "Execute code if a previous condition was false",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.flow_control.else_(args)
                            },
                        });

        commands.insert("end",
                        Command {
                            name: "end",
                            help: "End a code block",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.flow_control.end(args)
                            },
                        });

        commands.insert("for",
                        Command {
                            name: "for",
                            help: "Iterate through a list",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.flow_control.for_(args)
                            },
                        });

        let command_helper: HashMap<&'static str, &'static str> = commands.iter()
                                                                          .map(|(k, v)| {
                                                                              (*k, v.help)
                                                                          })
                                                                          .collect();

        commands.insert("help",
                        Command {
                            name: "help",
                            help: "Display helpful information about a given command, or list \
                                   commands if none specified\n    help <command>",
                            main: box move |args: &[String], _: &mut Shell| -> i32 {
                                if let Some(command) = args.get(1) {
                                    if command_helper.contains_key(command.as_str()) {
                                        match command_helper.get(command.as_str()) {
                                            Some(help) => println!("{}", help),
                                            None => {
                                                println!("Command helper not found [run 'help']...")
                                            }
                                        }
                                    } else {
                                        println!("Command helper not found [run 'help']...");
                                    }
                                } else {
                                    for (command, _help) in command_helper.iter() {
                                        println!("{}", command);
                                    }
                                }
                                SUCCESS
                            },
                        });

        commands
    }
}

fn main() {
    let commands = Command::map();
    let mut shell = Shell::new();

    for arg in env::args().skip(1) {
        let mut command_list = String::new();
        if let Ok(mut file) = File::open(&arg) {
            if let Err(message) = file.read_to_string(&mut command_list) {
                println!("{}: Failed to read {}", message, arg);
            }
        }
        shell.on_command(&command_list, &commands);
        return;
    }

    loop {
        shell.print_prompt();

        if let Some(command) = readln() {
            let command = command.trim();
            if !command.is_empty() {
                shell.on_command(command, &commands);
            }
        } else {
            break;
        }
    }
}
<commit_msg>Less Obnoxious Error Instead of the current verbose error message when an invalid command is entered, what if we were to print a traditional error message like so:<commit_after>#![feature(deque_extras)]
#![feature(box_syntax)]
#![feature(plugin)]
#![plugin(peg_syntax_ext)]

extern crate glob;

use std::collections::HashMap;
use std::fs::File;
use std::io::{stdout, Read, Write};
use std::env;
use std::process;
use std::thread;

use self::directory_stack::DirectoryStack;
use self::input_editor::readln;
use self::peg::{parse, Job};
use self::variables::Variables;
use self::history::History;
use self::flow_control::{FlowControl, is_flow_control_command, Statement};
use self::status::{SUCCESS, NO_SUCH_COMMAND, TERMINATED};

pub mod directory_stack;
pub mod to_num;
pub mod input_editor;
pub mod peg;
pub mod variables;
pub mod history;
pub mod flow_control;
pub mod status;

/// This struct will contain all of the data structures related to this
/// instance of the shell.
pub struct Shell {
    variables: Variables,
    flow_control: FlowControl,
    directory_stack: DirectoryStack,
    history: History,
}

impl Shell {
    /// Panics if DirectoryStack construction fails
    pub fn new() -> Self {
        Shell {
            variables: Variables::new(),
            flow_control: FlowControl::new(),
            directory_stack: DirectoryStack::new().expect(""),
            history: History::new(),
        }
    }

    pub fn print_prompt(&self) {
        self.print_prompt_prefix();
        match self.flow_control.current_statement {
            Statement::For(_, _) => self.print_for_prompt(),
            Statement::Default => self.print_default_prompt(),
        }
        if let Err(message) = stdout().flush() {
            println!("{}: failed to flush prompt to stdout", message);
        }

    }

    // TODO eventually this thing should be gone
    fn print_prompt_prefix(&self) {
        let prompt_prefix = self.flow_control.modes.iter().rev().fold(String::new(), |acc, mode| {
            acc +
            if mode.value {
                "+ "
            } else {
                "- "
            }
        });
        print!("{}", prompt_prefix);
    }

    fn print_for_prompt(&self) {
        print!("for> ");
    }

    fn print_default_prompt(&self) {
        let cwd = env::current_dir().ok().map_or("?".to_string(),
                                                 |ref p| p.to_str().unwrap_or("?").to_string());
        print!("ion:{}# ", cwd);
    }

    fn on_command(&mut self, command_string: &str, commands: &HashMap<&str, Command>) {
        self.history.add(command_string.to_string());

        let mut jobs = parse(command_string);

        // Execute commands
        for job in jobs.drain(..) {
            if self.flow_control.collecting_block {
                // TODO move this logic into "end" command
                if job.command == "end" {
                    self.flow_control.collecting_block = false;
                    let block_jobs: Vec<Job> = self.flow_control
                                                   .current_block
                                                   .jobs
                                                   .drain(..)
                                                   .collect();
                    let mut variable = String::new();
                    let mut values: Vec<String> = vec![];
                    if let Statement::For(ref var, ref vals) = self.flow_control.current_statement {
                        variable = var.clone();
                        values = vals.clone();
                    }
                    for value in values {
                        self.variables.set_var(&variable, &value);
                        for job in block_jobs.iter() {
                            self.run_job(job, commands);
                        }
                    }
                    self.flow_control.current_statement = Statement::Default;
                } else {
                    self.flow_control.current_block.jobs.push(job);
                }
            } else {
                if self.flow_control.skipping() && !is_flow_control_command(&job.command) {
                    continue;
                }
                self.run_job(&job, commands);
            }
        }
    }

    fn run_job(&mut self, job: &Job, commands: &HashMap<&str, Command>) {
        let mut job = self.variables.expand_job(job);
        job.expand_globs();
        let exit_status = if let Some(command) = commands.get(job.command.as_str()) {
            Some((*command.main)(job.args.as_slice(), self))
        } else {
            self.run_external_commmand(job)
        };
        if let Some(code) = exit_status {
            self.variables.set_var("?", &code.to_string());
        }
    }

    /// Returns an exit code if a command was run
    fn run_external_commmand(&mut self, job: Job) -> Option<i32> {
        if job.background {
            thread::spawn(move || {
                let mut command = Shell::build_command(&job);
                command.stdin(process::Stdio::null());
                if let Ok(mut child) = command.spawn() {
                    Shell::wait_and_get_status(&mut child, &job.command);
                }
            });
            None
        } else {
            let mut command = Shell::build_command(&job);
            match command.spawn() {
                Ok(mut child) => Some(Shell::wait_and_get_status(&mut child, &job.command)),
                Err(_) => {
                    println!("ion: command not found: {}", job.command);
                    Some(NO_SUCH_COMMAND)
                }
            }
        }
    }

    fn build_command(job: &Job) -> process::Command {
        let mut command = process::Command::new(&job.command);
        for i in 1..job.args.len() {
            if let Some(arg) = job.args.get(i) {
                command.arg(arg);
            }
        }
        command
    }

    // TODO don't pass in command and do printing outside this function
    fn wait_and_get_status(child: &mut process::Child, command: &str) -> i32 {
        match child.wait() {
            Ok(status) => {
                if let Some(code) = status.code() {
                    code
                } else {
                    println!("{}: child ended by signal", command);
                    TERMINATED
                }
            }
            Err(err) => {
                println!("{}: Failed to wait: {}", command, err);
                100 // TODO what should we return here?
            }
        }
    }
}

/// Structure which represents a Terminal's command.
/// This command structure contains a name, and the code which run the
/// functionnality associated to this one, with zero, one or several argument(s).
/// # Example
/// ```
/// let my_command = Command {
///     name: "my_command",
///     help: "Describe what my_command does followed by a newline showing usage",
///     main: box|args: &[String], &mut Shell| -> i32 {
///         println!("Say 'hello' to my command! :-D");
///     }
/// }
/// ```
pub struct Command {
    pub name: &'static str,
    pub help: &'static str,
    pub main: Box<Fn(&[String], &mut Shell) -> i32>,
}

impl Command {
    /// Return the map from command names to commands
    pub fn map() -> HashMap<&'static str, Self> {
        let mut commands: HashMap<&str, Self> = HashMap::new();

        commands.insert("cd",
                        Command {
                            name: "cd",
                            help: "Change the current directory\n    cd <path>",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.directory_stack.cd(args)
                            },
                        });

        commands.insert("dirs",
                        Command {
                            name: "dirs",
                            help: "Display the current directory stack",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.directory_stack.dirs(args)
                            },
                        });

        commands.insert("exit",
                        Command {
                            name: "exit",
                            help: "To exit the curent session",
                            main: box |args: &[String], _: &mut Shell| -> i32 {
                                if let Some(status) = args.get(1) {
                                    if let Ok(status) = status.parse::<i32>() {
                                        process::exit(status);
                                    }
                                }
                                // TODO should use exit status of previously run command, not 0
                                process::exit(0);
                            },
                        });

        commands.insert("let",
                        Command {
                            name: "let",
                            help: "View, set or unset variables",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.variables.let_(args)
                            },
                        });

        commands.insert("read",
                        Command {
                            name: "read",
                            help: "Read some variables\n    read <variable>",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.variables.read(args)
                            },
                        });

        commands.insert("pushd",
                        Command {
                            name: "pushd",
                            help: "Push a directory to the stack",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.directory_stack.pushd(args)
                            },
                        });

        commands.insert("popd",
                        Command {
                            name: "popd",
                            help: "Pop a directory from the stack",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.directory_stack.popd(args)
                            },
                        });

        commands.insert("history",
                        Command {
                            name: "history",
                            help: "Display a log of all commands previously executed",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.history.history(args)
                            },
                        });

        commands.insert("if",
                        Command {
                            name: "if",
                            help: "Conditionally execute code",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.flow_control.if_(args)
                            },
                        });

        commands.insert("else",
                        Command {
                            name: "else",
                            help: "Execute code if a previous condition was false",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.flow_control.else_(args)
                            },
                        });

        commands.insert("end",
                        Command {
                            name: "end",
                            help: "End a code block",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.flow_control.end(args)
                            },
                        });

        commands.insert("for",
                        Command {
                            name: "for",
                            help: "Iterate through a list",
                            main: box |args: &[String], shell: &mut Shell| -> i32 {
                                shell.flow_control.for_(args)
                            },
                        });

        let command_helper: HashMap<&'static str, &'static str> = commands.iter()
                                                                          .map(|(k, v)| {
                                                                              (*k, v.help)
                                                                          })
                                                                          .collect();

        commands.insert("help",
                        Command {
                            name: "help",
                            help: "Display helpful information about a given command, or list \
                                   commands if none specified\n    help <command>",
                            main: box move |args: &[String], _: &mut Shell| -> i32 {
                                if let Some(command) = args.get(1) {
                                    if command_helper.contains_key(command.as_str()) {
                                        match command_helper.get(command.as_str()) {
                                            Some(help) => println!("{}", help),
                                            None => {
                                                println!("Command helper not found [run 'help']...")
                                            }
                                        }
                                    } else {
                                        println!("Command helper not found [run 'help']...");
                                    }
                                } else {
                                    for (command, _help) in command_helper.iter() {
                                        println!("{}", command);
                                    }
                                }
                                SUCCESS
                            },
                        });

        commands
    }
}

fn main() {
    let commands = Command::map();
    let mut shell = Shell::new();

    for arg in env::args().skip(1) {
        let mut command_list = String::new();
        if let Ok(mut file) = File::open(&arg) {
            if let Err(message) = file.read_to_string(&mut command_list) {
                println!("{}: Failed to read {}", message, arg);
            }
        }
        shell.on_command(&command_list, &commands);
        return;
    }

    loop {
        shell.print_prompt();

        if let Some(command) = readln() {
            let command = command.trim();
            if !command.is_empty() {
                shell.on_command(command, &commands);
            }
        } else {
            break;
        }
    }
}
<|endoftext|>


<commit_before>extern crate rand;

#[macro_use]
extern crate lazy_static;

use rand::{XorShiftRng, Rng, thread_rng, SeedableRng};
use std::fmt;
use std::fmt::{Formatter, Debug};

use std::sync::Mutex;

trait Specimen: Sized {
    fn score(&self) -> f32;
    fn reevaluate(&mut self);
    fn mutate(generation: &mut[Self], rng: &mut XorShiftRng);
    fn procrastinate(parents: &[Self], rng: &mut XorShiftRng) -> Self;
    fn filter_strongest(species: &mut Vec<Self>);

    fn initial() -> Vec<Self>;
}

fn genetic<S: Specimen>(threshold: f32, max_iters: usize, nparents: usize, nchildren: usize) -> S {
    let mut species: Vec<S> = S::initial();
    let mut best = (0f32, 0);
    for _ in 0..max_iters {
        species = next_gen(species, nparents, nchildren);

        for (i, specimen) in species.iter().enumerate() {
            let score = specimen.score();
            if best.0 <= score {
                best = (score, i)
            }
        }

        if best.0 >= threshold {
            println!("found!");
            break;
        }
    }

    species.swap_remove(best.1)
}

#[inline(never)]
fn next_gen<S: Specimen>(mut species: Vec<S>, nparents: usize, nchildren: usize) -> Vec<S> {
    let mut rng = RNG.lock().unwrap();

    // 1. filter strongest
    S::filter_strongest(&mut species);

    // 2. bear children
    let mut children = vec![];
    let mut family = vec![];
    for _ in 0..nchildren {
        make_family(&mut species, nparents, &mut family, &mut rng);
        let child = S::procrastinate(&family, &mut rng);
        children.push(child);
        species.extend(family.drain(..));
    }

    // 3. mutate
    S::mutate(&mut children, &mut rng);
    for child in children.iter_mut() {
        child.reevaluate();
    }

    if KILL_PARENTS {
        children
    } else {
        species.extend(children);
        species
    }
}

fn make_family<S: Specimen>(species: &mut Vec<S>, nparents: usize, family: &mut Vec<S>, rng: &mut XorShiftRng) {
    assert!(family.is_empty());
    for _ in 0..nparents {
        let next = rng.gen_range(0, species.len());
        let s = species.swap_remove(next);
        family.push(s);
    }
}

//---------------------------------------------------

// BEST FOR SIZE=8
//const KILL_PARENTS: bool = true;
//const SIZE: usize = 8;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const POPULATION: usize = 3;
//const MAX_ITERS: usize = 1000;
//const NPARENTS: usize = 2;
//const NCHILDREN: usize = 100;

// BEST FOR SIZE=12
//const KILL_PARENTS: bool = true;
//const SIZE: usize = 12;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const POPULATION: usize = 3;
//const MAX_ITERS: usize = 1000;
//const NPARENTS: usize = 2;
//const NCHILDREN: usize = 100;

// BEST FOR SIZE=15
//const KILL_PARENTS: bool = true;
//const SIZE: usize = 15;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const POPULATION: usize = 4;
//const MAX_ITERS: usize = 1000;
//const NPARENTS: usize = 2;
//const NCHILDREN: usize = 280;

// BEST FOR SIZE=30
//const KILL_PARENTS: bool = true;
//const SIZE: usize = 30;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const POPULATION: usize = 4;
//const MAX_ITERS: usize = 5000;
//const NPARENTS: usize = 2;
//const NCHILDREN: usize = 480;

// BEST FOR SIZE=50
const KILL_PARENTS: bool = true;
const SIZE: usize = 50;
const MUTATION_PROBABILITY: f32 = 1.0f32;
const POPULATION: usize = 5;
const MAX_ITERS: usize = 5000;
const NPARENTS: usize = 2;
const NCHILDREN: usize = 500;

//const KILL_PARENTS: bool = true;
//const SIZE: usize = 100;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const POPULATION: usize = 5;
//const MAX_ITERS: usize = 5000;
//const NPARENTS: usize = 2;
//const NCHILDREN: usize = 1200;

lazy_static! {
    static ref RNG: Mutex<XorShiftRng> = Mutex::new(XorShiftRng::new_unseeded());
}

struct Board {
    score: f32,
    queens: Vec<(usize, usize)>,
    cells: Vec<bool>
}

impl Board {
    fn new(n: usize, queens: Vec<(usize, usize)>) -> Board {
        let mut cells = vec![false; n*n];
        for &(ref x, ref y) in queens.iter() {
            cells[n*y + x] = true;
        }

        Board { score: 0f32, queens: queens, cells: cells }
    }

    fn at(&mut self, x: usize, y: usize) -> &mut bool {
        &mut self.cells[SIZE * y + x]
    }

    #[inline(never)]
    fn do_mutate(&mut self, rng: &mut XorShiftRng) {
        let q = rng.gen_range(0, SIZE);
        loop {
            let x = rng.gen_range(0, SIZE);
            let y = rng.gen_range(0, SIZE);
            if !*self.at(x, y) {
                let (oldx, oldy) = self.queens[q].clone();
                *self.at(oldx, oldy) = false;
                *self.at(x, y) = true;
                self.queens[q] = (x, y);
                break;
            }
        }
    }

    fn eval_conflicts(&self, q1: usize, conflicts: &mut [usize]) {
        let &(ref x1, ref y1) = &self.queens[q1];

        for q2 in q1+1..SIZE {
            let &(ref x2, ref y2) = &self.queens[q2];
            if Self::conflict(*x1, *y1, *x2, *y2) {
                conflicts[q1] += 1;
                conflicts[q2] += 1;
            }
        }
    }

    fn conflict(x1: usize, y1: usize, x2: usize, y2: usize) -> bool {
        x1 == x2
            || y1 == y2
            || x1.wrapping_sub(x2) == y1.wrapping_sub(y2)
            || x1.wrapping_sub(x2) == y2.wrapping_sub(y1)
    }
}

impl Specimen for Board {
    #[inline(never)]
    fn reevaluate(&mut self) {
        let mut score: usize = 0;
        let mut conflicts = vec![0; SIZE];
        for i in 0 .. SIZE {
            self.eval_conflicts(i, &mut conflicts);

            if conflicts[i] == 0 {
                score += 1;
            }
        }

        self.score = score as f32 + 0.000001;
    }

    fn score(&self) -> f32 {
        self.score
    }

    #[inline(never)]
    fn mutate(new_gen: &mut [Self], rng: &mut XorShiftRng) {
        for board in new_gen.iter_mut() {
            if rng.next_f32() < MUTATION_PROBABILITY {
                board.do_mutate(rng);
            }
        }
    }

    #[inline(never)]
    fn procrastinate(parents: &[Self], rng: &mut XorShiftRng) -> Self {
        let nqueens = parents.len() * SIZE;

        let mut child = Board::new(SIZE, Vec::with_capacity(SIZE));

        let mut all_queens = (0..nqueens).collect::<Vec<_>>();

        loop {
            let next = rng.gen_range(0, all_queens.len());
            let mut queen = all_queens.swap_remove(next);
            let parent = queen / SIZE;
            queen %= SIZE;

            let (x, y) = parents[parent].queens[queen].clone();
            if !*child.at(x, y) {
                child.queens.push((x, y));
                *child.at(x, y) = true;
                if child.queens.len() == SIZE {
                    break;
                }
            }
        }

        child
    }

    #[inline(never)]
    fn filter_strongest(species: &mut Vec<Self>) {
        if species.len() <= POPULATION {
            return;
        }

        species.sort_by(|s, t| t.score.partial_cmp(&s.score).unwrap());
        species.truncate(POPULATION);
    }

    #[inline(never)]
    fn initial() -> Vec<Self> {
        let mut boards = Vec::with_capacity(POPULATION);

        let mut rng = RNG.lock().unwrap();
        for _ in 0..POPULATION {
            let mut board = Board::new(SIZE, vec![]);

            loop {
                let x = rng.gen_range(0, SIZE);
                let y = rng.gen_range(0, SIZE);
                if !*board.at(x, y) {
                    board.queens.push((x, y));
                    *board.at(x, y) = true;
                    if board.queens.len() == SIZE {
                        break;
                    }
                }
            }

            board.reevaluate();
            boards.push(board);
        }
        boards
    }
}

impl Debug for Board {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "score={}, queens=[{:?}]", self.score, self.queens)
    }
}

fn genetic_queens() {
    {
        let mut rng = RNG.lock().unwrap();
        let mut trng = thread_rng();
        rng.reseed([trng.next_u32(), trng.next_u32(), trng.next_u32(), trng.next_u32()]);
    }

    let best = genetic::<Board>(SIZE as f32, MAX_ITERS, NPARENTS, NCHILDREN);
    println!("best: {:?}", &best);
}

fn main() {
    for _ in 0..10 {
        genetic_queens()
    }
}

<commit_msg>achieved a huge speed boost by   a) requiring that every row contain exactly 1 queen (=> faster convergence)   b) exploiting the above to remove all quadratic factors in the algorithm<commit_after>extern crate rand;

#[macro_use]
extern crate lazy_static;

use rand::{XorShiftRng, Rng, thread_rng, SeedableRng};
use std::fmt;
use std::fmt::{Formatter, Debug};

use std::sync::Mutex;

trait Specimen: Sized {
    fn score(&self) -> f32;
    fn reevaluate(&mut self);
    fn mutate(generation: &mut[Self], rng: &mut XorShiftRng);
    fn procrastinate(parents: &[Self], rng: &mut XorShiftRng) -> Self;
    fn filter_strongest(species: &mut Vec<Self>);

    fn initial() -> Vec<Self>;
}

fn genetic<S: Specimen>(threshold: f32, max_iters: usize, nparents: usize, nchildren: usize) -> S {
    let mut species: Vec<S> = S::initial();
    let mut best = (0f32, 0);
    for _ in 0..max_iters {
        species = next_gen(species, nparents, nchildren);

        for (i, specimen) in species.iter().enumerate() {
            let score = specimen.score();
            if best.0 <= score {
                best = (score, i)
            }
        }

        if best.0 >= threshold {
            println!("found!");
            break;
        }
    }

    species.swap_remove(best.1)
}

#[inline(never)]
fn next_gen<S: Specimen>(mut species: Vec<S>, nparents: usize, nchildren: usize) -> Vec<S> {
    let mut rng = RNG.lock().unwrap();

    // 1. filter strongest
    S::filter_strongest(&mut species);

    // 2. bear children
    let mut children = vec![];
    let mut family = vec![];
    for _ in 0..nchildren {
        make_family(&mut species, nparents, &mut family, &mut rng);
        let child = S::procrastinate(&family, &mut rng);
        children.push(child);
        species.extend(family.drain(..));
    }

    // 3. mutate
    S::mutate(&mut children, &mut rng);
    for child in children.iter_mut() {
        child.reevaluate();
    }

    if KILL_PARENTS {
        children
    } else {
        species.extend(children);
        species
    }
}

fn make_family<S: Specimen>(species: &mut Vec<S>, nparents: usize, family: &mut Vec<S>, rng: &mut XorShiftRng) {
    assert!(family.is_empty());
    for _ in 0..nparents {
        let next = rng.gen_range(0, species.len());
        let s = species.swap_remove(next);
        family.push(s);
    }
}

//---------------------------------------------------

// BEST FOR SIZE=8
//const KILL_PARENTS: bool = true;
//const SIZE: usize = 8;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const POPULATION: usize = 3;
//const MAX_ITERS: usize = 1000;
//const NPARENTS: usize = 2;
//const NCHILDREN: usize = 100;

// BEST FOR SIZE=12
//const KILL_PARENTS: bool = true;
//const SIZE: usize = 12;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const POPULATION: usize = 3;
//const MAX_ITERS: usize = 1000;
//const NPARENTS: usize = 2;
//const NCHILDREN: usize = 100;

// BEST FOR SIZE=15
//const KILL_PARENTS: bool = true;
//const SIZE: usize = 15;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const POPULATION: usize = 4;
//const MAX_ITERS: usize = 1000;
//const NPARENTS: usize = 2;
//const NCHILDREN: usize = 280;

// BEST FOR SIZE=30
//const KILL_PARENTS: bool = true;
//const SIZE: usize = 30;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const POPULATION: usize = 4;
//const MAX_ITERS: usize = 5000;
//const NPARENTS: usize = 2;
//const NCHILDREN: usize = 480;

// BEST FOR SIZE=50
const KILL_PARENTS: bool = true;
const SIZE: usize = 400;
const MUTATION_PROBABILITY: f32 = 1.0f32;
const POPULATION: usize = 5;
const MAX_ITERS: usize = 5000;
const NPARENTS: usize = 2;
const NCHILDREN: usize = 640;

//const KILL_PARENTS: bool = true;
//const SIZE: usize = 100;
//const MUTATION_PROBABILITY: f32 = 1.0f32;
//const POPULATION: usize = 5;
//const MAX_ITERS: usize = 5000;
//const NPARENTS: usize = 2;
//const NCHILDREN: usize = 1200;

lazy_static! {
    static ref RNG: Mutex<XorShiftRng> = Mutex::new(XorShiftRng::new_unseeded());
}

struct Board {
    score: f32,
    queens: Vec<usize>,
}

impl Board {
    fn new(queens: Vec<usize>) -> Board {
        Board { score: 0f32, queens: queens }
    }

    #[inline(never)]
    fn do_mutate(&mut self, rng: &mut XorShiftRng) {
        let x = rng.gen_range(0, SIZE);

        let mut y = rng.gen_range(0, SIZE-1);
        if y >= self.queens[x] {
            y += 1;
        }
        self.queens[x] = y;
    }
}

impl Specimen for Board {
    #[inline(never)]
    fn reevaluate(&mut self) {
        let mut occurences = vec![0; 2*SIZE-1];
        let mut nonconflicting = Vec::with_capacity(SIZE);

        // columns
        for x in 0..SIZE {
            occurences[self.queens[x]] += 1;
        }
        for x in 0..SIZE {
            if occurences[self.queens[x]] == 1 {
                nonconflicting.push(x);
            }
        }

        // forward diagonals
        for y in 0..SIZE {
            occurences[y] = 0;
        }
        for x in 0..SIZE {
            let d = SIZE-1 + x - self.queens[x];
            occurences[d] += 1;
        }
        let mut i = 0;
        while i < nonconflicting.len() {
            let x = nonconflicting[i];
            let d = SIZE-1 + x - self.queens[x];
            if occurences[d] == 1 {
                i += 1;
            } else {
                nonconflicting.swap_remove(i);
            }
        }

        // backward diagonals
        for y in 0..2*SIZE-1 {
            occurences[y] = 0;
        }
        for x in 0..SIZE {
            let d = x + self.queens[x];
            occurences[d] += 1;
        }
        i = 0;
        while i < nonconflicting.len() {
            let x = nonconflicting[i];
            let d = x + self.queens[x];
            if occurences[d] == 1 {
                i += 1;
            } else {
                nonconflicting.swap_remove(i);
            }
        }

        self.score = nonconflicting.len() as f32 + 0.000001;
    }

    fn score(&self) -> f32 {
        self.score
    }

    #[inline(never)]
    fn mutate(new_gen: &mut [Self], rng: &mut XorShiftRng) {
        for board in new_gen.iter_mut() {
            if rng.next_f32() < MUTATION_PROBABILITY {
                board.do_mutate(rng);
            }
        }
    }

    #[inline(never)]
    fn procrastinate(parents: &[Self], rng: &mut XorShiftRng) -> Self {
        let mut child = Board::new(Vec::with_capacity(SIZE));

        for x in 0..SIZE {
            let parent = rng.gen_range(0, parents.len());
            child.queens.push(parents[parent].queens[x]);
        }

        child
    }

    #[inline(never)]
    fn filter_strongest(species: &mut Vec<Self>) {
        if species.len() <= POPULATION {
            return;
        }

        species.sort_by(|s, t| t.score.partial_cmp(&s.score).unwrap());
        species.truncate(POPULATION);
    }

    #[inline(never)]
    fn initial() -> Vec<Self> {
        let mut boards = Vec::with_capacity(POPULATION);

        let mut rng = RNG.lock().unwrap();
        for _ in 0..POPULATION {
            let mut board = Board::new(Vec::with_capacity(SIZE));

            for _ in 0..SIZE {
                let y = rng.gen_range(0, SIZE);
                board.queens.push(y);
            }

            board.reevaluate();
            boards.push(board);
        }
        boards
    }
}

impl Debug for Board {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "score={}, queens=[{:?}]", self.score, self.queens)
    }
}

fn genetic_queens() {
    {
        let mut rng = RNG.lock().unwrap();
        let mut trng = thread_rng();
        rng.reseed([trng.next_u32(), trng.next_u32(), trng.next_u32(), trng.next_u32()]);
    }

    let best = genetic::<Board>(SIZE as f32, MAX_ITERS, NPARENTS, NCHILDREN);
    println!("best: {:?}", &best);
}

fn main() {
    for _ in 0..10 {
        genetic_queens()
    }
}

<|endoftext|>


<commit_before>mod bxdf;
mod fresnel;
mod lambertian;
mod oren_nayar;
mod microfacet;

pub use self::bxdf::*;
pub use self::fresnel::*;
pub use self::lambertian::*;
pub use self::oren_nayar::*;
pub use self::microfacet::*;

use std::cmp;

use na::{Cross, Dot, Norm};

use ::{Vector3f, Point2f, ONE_MINUS_EPSILON};
use spectrum::Spectrum;
use interaction::SurfaceInteraction;

bitflags! {
    pub flags BxDFType: u32 {
        const BSDF_REFLECTION   = 0b_00000001,
        const BSDF_TRANSMISSION = 0b_00000010,
        const BSDF_DIFFUSE      = 0b_00000100,
        const BSDF_GLOSSY       = 0b_00001000,
        const BSDF_SPECULAR     = 0b_00010000,
    }
}

/// Represents the Bidirectional Scattering Distribution Function.
/// It represents the properties of a material at a given point.
pub struct BSDF {
    /// Index of refraction of the surface
    eta: f32,
    /// Shading normal (i.e. potentially affected by bump-mapping)
    ns: Vector3f,
    /// Geometry normal
    ng: Vector3f,
    ss: Vector3f,
    ts: Vector3f,
    bxdfs: Vec<Box<BxDF + Sync + Send>>,
}

impl BSDF {
    pub fn new(isect: &SurfaceInteraction, eta: f32, bxdfs: Vec<Box<BxDF + Sync + Send>>) -> BSDF {
        let ss = isect.dpdu.normalize();
        BSDF {
            eta: eta,
            ns: isect.shading.n,
            ng: isect.n,
            ss: ss,
            ts: isect.shading.n.cross(&ss),
            bxdfs: bxdfs,
        }
    }

    /// Evaluate the BSDF for the given incoming light direction and outgoing light direction.
    pub fn f(&self, wo_w: &Vector3f, wi_w: &Vector3f, flags: BxDFType) -> Spectrum {
        let wi = self.world_to_local(wi_w);
        let wo = self.world_to_local(wo_w);
        if wo.z == 0.0 {
            return Spectrum::black();
        }
        let reflect = wi_w.dot(&self.ng) * wo_w.dot(&self.ng) > 0.0;
        self.bxdfs
            .iter()
            .filter(|b| {
                // Make sure we only evaluate reflection or transmission based on whether wi and wo
                // lie in the same hemisphere.
                b.matches(flags) &&
                ((reflect && (b.get_type().contains(BSDF_REFLECTION))) ||
                 (!reflect && (b.get_type().contains(BSDF_TRANSMISSION))))
            })
            .fold(Spectrum::black(), |c, b| c + b.f(&wo, &wi))
    }

    pub fn pdf(&self, wo_w: &Vector3f, wi_w: &Vector3f, flags: BxDFType) -> f32 {
        if self.bxdfs.is_empty() {
            return 0.0;
        }
        let wo = self.world_to_local(wo_w);
        if wo.z == 0.0 {
            return 0.0;
        }
        let wi = self.world_to_local(wi_w);

        let mut matched_comps = 0;
        let mut pdf = 0.0;
        for bxdf in &self.bxdfs {
            if bxdf.matches(flags) {
                matched_comps += 1;
                pdf += bxdf.pdf(&wo, &wi);
            }
        }
        if matched_comps == 0 {
            0.0
        } else {
            pdf / matched_comps as f32
        }
    }

    pub fn sample_f(&self,
                    wo_w: &Vector3f,
                    u: &Point2f,
                    flags: BxDFType)
                    -> (Spectrum, Vector3f, f32, BxDFType) {

        let matching_comps = self.bxdfs
            .iter()
            .filter(|b| b.matches(flags))
            .collect::<Vec<&Box<BxDF + Sync + Send>>>();
        if matching_comps.is_empty() {
            return (Spectrum::black(), Vector3f::new(0.0, 0.0, 0.0), 0.0, BxDFType::empty());
        }
        // Chose which BxDF to sample
        let comp = cmp::min((u[0] * matching_comps.len() as f32).floor() as usize,
                            matching_comps.len() - 1);
        let bxdf = matching_comps.get(comp).expect("Was expecting a BxDF with this index!");

        debug!("BDDF::sample_f chose comp = {} / matching {}",
               comp,
               matching_comps.len());

        // Remap BxDF sample u to [0,1)^2
        let u_remapped = Point2f::new((u[0] * matching_comps.len() as f32 - comp as f32)
                                          .min(ONE_MINUS_EPSILON),
                                      u[1]);
        // Sample chosen BxDF
        let wo = self.world_to_local(wo_w);
        if wo.z == 0.0 {
            return (Spectrum::black(), Vector3f::new(0.0, 0.0, 0.0), 0.0, bxdf.get_type());
        }
        let (mut f, wi, mut pdf, sampled_type) = bxdf.sample_f(&wo, &u_remapped);
        debug!("For wo = {:?}, sampled f = {}, pdf = {}, ratio = {}, wi = {:?}",
               wo,
               f,
               pdf,
               if pdf > 0.0 {
                   f / pdf
               } else {
                   Spectrum::black()
               },
               wi);
        if pdf == 0.0 {
            return (Spectrum::black(), Vector3f::new(0.0, 0.0, 0.0), 0.0, BxDFType::empty());
        }
        let wi_w = self.local_to_world(&wi);

        // Compute overall PDF with all matching BxDF
        if !bxdf.get_type().contains(BSDF_SPECULAR) && matching_comps.len() > 1 {
            for i in 0..matching_comps.len() {
                if i != comp {
                    pdf += matching_comps[i].pdf(&wo, &wi);
                }
            }
        }
        if matching_comps.len() > 1 {
            pdf /= matching_comps.len() as f32;
        }

        // Compute value of BSDF for sampled direction
        if !bxdf.get_type().contains(BSDF_SPECULAR) && matching_comps.len() > 1 {
            let reflect = wi_w.dot(&self.ng) * wo_w.dot(&self.ng) > 0.0;
            f = matching_comps.iter()
                .filter(|b| {
                    (reflect && b.get_type().contains(BSDF_REFLECTION)) ||
                    (!reflect && b.get_type().contains(BSDF_TRANSMISSION))
                })
                .fold(Spectrum::black(), |f, b| f + b.f(&wo, &wi));
        }

        debug!("Overall f = {}, pdf = {}, ratio = {}",
               f,
               pdf,
               if pdf > 0.0 {
                   f / pdf
               } else {
                   Spectrum::black()
               });
        (f, wi, pdf, sampled_type)
    }

    fn world_to_local(&self, v: &Vector3f) -> Vector3f {
        Vector3f::new(v.dot(&self.ss), v.dot(&self.ts), v.dot(&self.ns))
    }

    fn local_to_world(&self, v: &Vector3f) -> Vector3f {
        Vector3f::new(self.ss.x * v.x + self.ts.x * v.y + self.ns.x * v.z,
                      self.ss.y * v.x + self.ts.y * v.y + self.ns.y * v.z,
                      self.ss.z * v.z + self.ts.z * v.y + self.ns.z * v.z)
    }

    /// Return the number of BxDFs matching the given flags
    fn num_components(&self, flags: BxDFType) -> usize {
        self.bxdfs.iter().filter(|b| b.matches(flags)).count()
    }
}
<commit_msg>Return wi in world space, not local space<commit_after>mod bxdf;
mod fresnel;
mod lambertian;
mod oren_nayar;
mod microfacet;

pub use self::bxdf::*;
pub use self::fresnel::*;
pub use self::lambertian::*;
pub use self::oren_nayar::*;
pub use self::microfacet::*;

use std::cmp;

use na::{Cross, Dot, Norm};

use ::{Vector3f, Point2f, ONE_MINUS_EPSILON};
use spectrum::Spectrum;
use interaction::SurfaceInteraction;

bitflags! {
    pub flags BxDFType: u32 {
        const BSDF_REFLECTION   = 0b_00000001,
        const BSDF_TRANSMISSION = 0b_00000010,
        const BSDF_DIFFUSE      = 0b_00000100,
        const BSDF_GLOSSY       = 0b_00001000,
        const BSDF_SPECULAR     = 0b_00010000,
    }
}

/// Represents the Bidirectional Scattering Distribution Function.
/// It represents the properties of a material at a given point.
pub struct BSDF {
    /// Index of refraction of the surface
    eta: f32,
    /// Shading normal (i.e. potentially affected by bump-mapping)
    ns: Vector3f,
    /// Geometry normal
    ng: Vector3f,
    ss: Vector3f,
    ts: Vector3f,
    bxdfs: Vec<Box<BxDF + Sync + Send>>,
}

impl BSDF {
    pub fn new(isect: &SurfaceInteraction, eta: f32, bxdfs: Vec<Box<BxDF + Sync + Send>>) -> BSDF {
        let ss = isect.dpdu.normalize();
        BSDF {
            eta: eta,
            ns: isect.shading.n,
            ng: isect.n,
            ss: ss,
            ts: isect.shading.n.cross(&ss),
            bxdfs: bxdfs,
        }
    }

    /// Evaluate the BSDF for the given incoming light direction and outgoing light direction.
    pub fn f(&self, wo_w: &Vector3f, wi_w: &Vector3f, flags: BxDFType) -> Spectrum {
        let wi = self.world_to_local(wi_w);
        let wo = self.world_to_local(wo_w);
        if wo.z == 0.0 {
            return Spectrum::black();
        }
        let reflect = wi_w.dot(&self.ng) * wo_w.dot(&self.ng) > 0.0;
        self.bxdfs
            .iter()
            .filter(|b| {
                // Make sure we only evaluate reflection or transmission based on whether wi and wo
                // lie in the same hemisphere.
                b.matches(flags) &&
                ((reflect && (b.get_type().contains(BSDF_REFLECTION))) ||
                 (!reflect && (b.get_type().contains(BSDF_TRANSMISSION))))
            })
            .fold(Spectrum::black(), |c, b| c + b.f(&wo, &wi))
    }

    pub fn pdf(&self, wo_w: &Vector3f, wi_w: &Vector3f, flags: BxDFType) -> f32 {
        if self.bxdfs.is_empty() {
            return 0.0;
        }
        let wo = self.world_to_local(wo_w);
        if wo.z == 0.0 {
            return 0.0;
        }
        let wi = self.world_to_local(wi_w);

        let mut matched_comps = 0;
        let mut pdf = 0.0;
        for bxdf in &self.bxdfs {
            if bxdf.matches(flags) {
                matched_comps += 1;
                pdf += bxdf.pdf(&wo, &wi);
            }
        }
        if matched_comps == 0 {
            0.0
        } else {
            pdf / matched_comps as f32
        }
    }

    pub fn sample_f(&self,
                    wo_w: &Vector3f,
                    u: &Point2f,
                    flags: BxDFType)
                    -> (Spectrum, Vector3f, f32, BxDFType) {

        let matching_comps = self.bxdfs
            .iter()
            .filter(|b| b.matches(flags))
            .collect::<Vec<&Box<BxDF + Sync + Send>>>();
        if matching_comps.is_empty() {
            return (Spectrum::black(), Vector3f::new(0.0, 0.0, 0.0), 0.0, BxDFType::empty());
        }
        // Chose which BxDF to sample
        let comp = cmp::min((u[0] * matching_comps.len() as f32).floor() as usize,
                            matching_comps.len() - 1);
        let bxdf = matching_comps.get(comp).expect("Was expecting a BxDF with this index!");

        debug!("BDDF::sample_f chose comp = {} / matching {}",
               comp,
               matching_comps.len());

        // Remap BxDF sample u to [0,1)^2
        let u_remapped = Point2f::new((u[0] * matching_comps.len() as f32 - comp as f32)
                                          .min(ONE_MINUS_EPSILON),
                                      u[1]);
        // Sample chosen BxDF
        let wo = self.world_to_local(wo_w);
        if wo.z == 0.0 {
            return (Spectrum::black(), Vector3f::new(0.0, 0.0, 0.0), 0.0, bxdf.get_type());
        }
        let (mut f, wi, mut pdf, sampled_type) = bxdf.sample_f(&wo, &u_remapped);
        debug!("For wo = {:?}, sampled f = {}, pdf = {}, ratio = {}, wi = {:?}",
               wo,
               f,
               pdf,
               if pdf > 0.0 {
                   f / pdf
               } else {
                   Spectrum::black()
               },
               wi);
        if pdf == 0.0 {
            return (Spectrum::black(), Vector3f::new(0.0, 0.0, 0.0), 0.0, BxDFType::empty());
        }
        let wi_w = self.local_to_world(&wi);

        // Compute overall PDF with all matching BxDF
        if !bxdf.get_type().contains(BSDF_SPECULAR) && matching_comps.len() > 1 {
            for i in 0..matching_comps.len() {
                if i != comp {
                    pdf += matching_comps[i].pdf(&wo, &wi);
                }
            }
        }
        if matching_comps.len() > 1 {
            pdf /= matching_comps.len() as f32;
        }

        // Compute value of BSDF for sampled direction
        if !bxdf.get_type().contains(BSDF_SPECULAR) && matching_comps.len() > 1 {
            let reflect = wi_w.dot(&self.ng) * wo_w.dot(&self.ng) > 0.0;
            f = matching_comps.iter()
                .filter(|b| {
                    (reflect && b.get_type().contains(BSDF_REFLECTION)) ||
                    (!reflect && b.get_type().contains(BSDF_TRANSMISSION))
                })
                .fold(Spectrum::black(), |f, b| f + b.f(&wo, &wi));
        }

        debug!("Overall f = {}, pdf = {}, ratio = {}",
               f,
               pdf,
               if pdf > 0.0 {
                   f / pdf
               } else {
                   Spectrum::black()
               });
        (f, wi_w, pdf, sampled_type)
    }

    fn world_to_local(&self, v: &Vector3f) -> Vector3f {
        Vector3f::new(v.dot(&self.ss), v.dot(&self.ts), v.dot(&self.ns))
    }

    fn local_to_world(&self, v: &Vector3f) -> Vector3f {
        Vector3f::new(self.ss.x * v.x + self.ts.x * v.y + self.ns.x * v.z,
                      self.ss.y * v.x + self.ts.y * v.y + self.ns.y * v.z,
                      self.ss.z * v.z + self.ts.z * v.y + self.ns.z * v.z)
    }

    /// Return the number of BxDFs matching the given flags
    fn num_components(&self, flags: BxDFType) -> usize {
        self.bxdfs.iter().filter(|b| b.matches(flags)).count()
    }
}

#[test]
fn test_flags() {
    let flags = BSDF_SPECULAR | BSDF_REFLECTION;
    let bxdf_type = BSDF_SPECULAR | BSDF_REFLECTION | BSDF_TRANSMISSION;

    assert!((bxdf_type & flags) == flags);
}
<|endoftext|>