<commit_before>extern crate clap;
extern crate rand;

pub const DEFAULT_ENGLISH_DICTIONARY_PATH: &'static str = "/usr/share/dict/words";

#[derive(Debug)]
pub struct EnglishPasswordOptions {
    pub num_words: u8,
    pub dictionary_path: String,
}

#[derive(Debug)]
pub struct KanaPasswordOptions {
    pub syllable_counts: Vec<u8>,
}

pub enum PasswordOptions {
    English(EnglishPasswordOptions),
    Japanese(KanaPasswordOptions),
}

mod args {
    use super::{
        EnglishPasswordOptions, KanaPasswordOptions, PasswordOptions,
        DEFAULT_ENGLISH_DICTIONARY_PATH,
    };

    // Aborts this process on error.
    pub fn parse_args() -> PasswordOptions {
        let matches = clap::App::new("xkpw")
            .version("0.1.1")
            .author("j39m")
            .about("Generates passwords")
            .subcommand(
                clap::App::new("en")
                    .about("generates dictionary passwords in English")
                    .arg(clap::Arg::new("num-words").required(true)),
            )
            .subcommand(
                clap::App::new("jp")
                    .about("generates random strings of Japanese syllables")
                    .arg(
                        clap::Arg::new("syllable-counts")
                            .required(true)
                            .multiple_values(true),
                    ),
            )
            .get_matches();

        match matches.subcommand() {
            Some(("en", en_matches)) => {
                return PasswordOptions::English(EnglishPasswordOptions {
                    num_words: en_matches
                        .value_of_t::<u8>("num-words")
                        .unwrap_or_else(|e| e.exit()),
                    dictionary_path: DEFAULT_ENGLISH_DICTIONARY_PATH.to_owned(),
                })
            }
            Some(("jp", jp_matches)) => {
                return PasswordOptions::Japanese(KanaPasswordOptions {
                    syllable_counts: jp_matches
                        .values_of_t("syllable-counts")
                        .unwrap_or_else(|e| e.exit()),
                })
            }
            _ => (),
        }
        PasswordOptions::Japanese(KanaPasswordOptions {
            syllable_counts: vec![4, 3, 3, 3],
        })
    }
} // mod args

mod kana {
    const BASE: &'static [&str] = &[
        "a", "i", "u", "e", "o", "ka", "ki", "ku", "ke", "ko", "sa", "shi", "su", "se", "so", "ta",
        "chi", "tsu", "te", "to", "na", "ni", "nu", "ne", "no", "ha", "hi", "hu", "he", "ho", "ma",
        "mi", "mu", "me", "mo", "ya", "yu", "yo", "ra", "ri", "ru", "re", "ro", "wa", "wo",
    ];

    const DIGRAPHS: &'static [&str] = &[
        "kya", "kyu", "kyo", "sha", "shu", "sho", "cha", "chu", "cho", "nya", "nyu", "nyo", "hya",
        "hyu", "hyo", "mya", "myu", "myo", "rya", "ryu", "ryo",
    ];

    const DIACRITICS: &'static [&str] = &[
        "ga", "gi", "gu", "ge", "go", "za", "ji", "zu", "ze", "zo", "da", "de", "do", "ba", "bi",
        "bu", "be", "bo", "pa", "pi", "pu", "pe", "po",
    ];

    const DIGRAPHS_WITH_DIACRITICS: &'static [&str] = &[
        "gya", "gyu", "gyo", "ja", "ju", "jo", "bya", "byu", "byo", "pya", "pyu", "pyo",
    ];

    fn all() -> Vec<&'static str> {
        let mut all_kana: Vec<&'static str> = Vec::new();
        all_kana.extend(BASE);
        all_kana.extend(DIGRAPHS);
        all_kana.extend(DIACRITICS);
        all_kana.extend(DIGRAPHS_WITH_DIACRITICS);
        all_kana
    }

    fn simple() -> Vec<&'static str> {
        let mut simple_kana: Vec<&'static str> = Vec::new();
        simple_kana.extend(BASE);
        simple_kana.extend(DIACRITICS);
        simple_kana
    }

    pub fn get(use_simple: bool) -> Vec<&'static str> {
        match use_simple {
            true => simple(),
            false => all(),
        }
    }
} // mod kana

mod helpers {
    use super::{EnglishPasswordOptions, KanaPasswordOptions, PasswordOptions};
    use rand::seq::SliceRandom;
    use std::convert::TryInto;
    use std::io::Read;

    // Reads the dictionary at |path| and returns a linewise vector of
    // its contents.
    fn ingest_english_dictionary(path: &std::path::Path) -> std::io::Result<Vec<String>> {
        let mut file = std::fs::File::open(path)?;
        let mut contents = String::new();
        file.read_to_string(&mut contents)?;
        Ok(contents
            .lines()
            .map(|word| word.trim().to_lowercase())
            .collect())
    }

    // Assumes ownership of a Vector of owned |words| and prints them
    // separated by spaces.
    fn print_words(words: Vec<String>) {
        let mut iter = words.into_iter();
        if let Some(word) = iter.next() {
            print!("{}", word);
            for word in iter {
                print!(" {}", word);
            }
            print!("\n");
        }
    }

    // Returns 1 if we fail to read the word dictionary.
    fn print_english_password(options: EnglishPasswordOptions) -> i32 {
        let all_words =
            match ingest_english_dictionary(std::path::Path::new(&options.dictionary_path)) {
                Ok(result) => result,
                Err(error) => {
                    eprintln!("Error reading dictionary: {}", error);
                    return 1;
                }
            };
        let mut rng = &mut rand::thread_rng();
        let selection: Vec<String> = all_words
            .choose_multiple(&mut rng, options.num_words.try_into().unwrap())
            .map(|borrowed| borrowed.to_owned())
            .collect();

        print_words(selection);
        0
    }

    // Builds a single random pseudo-Japanese word from |kana_set|.
    fn build_kana_word(kana_set: &Vec<&'static str>, syllable_count: &u8) -> String {
        let mut collected_kana: Vec<&'static str> = Vec::new();
        let mut rng = &mut rand::thread_rng();
        for _ in 0..*syllable_count {
            collected_kana.push(kana_set.choose(&mut rng).unwrap());
        }
        collected_kana.join("")
    }

    // Builds random pseudo-Japanese words from |kana_set|, observing
    // the word lengths specified by |options|.
    fn build_kana_words(kana_set: Vec<&'static str>, options: KanaPasswordOptions) -> Vec<String> {
        options
            .syllable_counts
            .iter()
            .map(|count| build_kana_word(&kana_set, count))
            .collect()
    }

    // Builds a vector of owned kana words, maps these to a vector of
    // borrowed kana words, and then prints the same.
    fn print_kana_password(options: KanaPasswordOptions) {
        print_words(build_kana_words(crate::kana::get(true), options));
    }

    pub fn main() -> i32 {
        let options = crate::args::parse_args();
        match options {
            PasswordOptions::English(options) => return print_english_password(options),
            PasswordOptions::Japanese(options) => print_kana_password(options),
        }
        0
    }
} // mod helpers

fn main() {
    std::process::exit(crate::helpers::main());
}
<commit_msg>xkpw: Reformat hiragana arrays<commit_after>extern crate clap;
extern crate rand;

pub const DEFAULT_ENGLISH_DICTIONARY_PATH: &'static str = "/usr/share/dict/words";

#[derive(Debug)]
pub struct EnglishPasswordOptions {
    pub num_words: u8,
    pub dictionary_path: String,
}

#[derive(Debug)]
pub struct KanaPasswordOptions {
    pub syllable_counts: Vec<u8>,
}

pub enum PasswordOptions {
    English(EnglishPasswordOptions),
    Japanese(KanaPasswordOptions),
}

mod args {
    use super::{
        EnglishPasswordOptions, KanaPasswordOptions, PasswordOptions,
        DEFAULT_ENGLISH_DICTIONARY_PATH,
    };

    // Aborts this process on error.
    pub fn parse_args() -> PasswordOptions {
        let matches = clap::App::new("xkpw")
            .version("0.1.1")
            .author("j39m")
            .about("Generates passwords")
            .subcommand(
                clap::App::new("en")
                    .about("generates dictionary passwords in English")
                    .arg(clap::Arg::new("num-words").required(true)),
            )
            .subcommand(
                clap::App::new("jp")
                    .about("generates random strings of Japanese syllables")
                    .arg(
                        clap::Arg::new("syllable-counts")
                            .required(true)
                            .multiple_values(true),
                    ),
            )
            .get_matches();

        match matches.subcommand() {
            Some(("en", en_matches)) => {
                return PasswordOptions::English(EnglishPasswordOptions {
                    num_words: en_matches
                        .value_of_t::<u8>("num-words")
                        .unwrap_or_else(|e| e.exit()),
                    dictionary_path: DEFAULT_ENGLISH_DICTIONARY_PATH.to_owned(),
                })
            }
            Some(("jp", jp_matches)) => {
                return PasswordOptions::Japanese(KanaPasswordOptions {
                    syllable_counts: jp_matches
                        .values_of_t("syllable-counts")
                        .unwrap_or_else(|e| e.exit()),
                })
            }
            _ => (),
        }
        PasswordOptions::Japanese(KanaPasswordOptions {
            syllable_counts: vec![4, 3, 3, 3],
        })
    }
} // mod args

mod kana {
    #[rustfmt::skip]
    const BASE: &'static [&str] = &[
        "a",    "i",    "u",    "e",    "o",
        "ka",   "ki",   "ku",   "ke",   "ko",
        "sa",   "shi",  "su",   "se",   "so",
        "ta",   "chi",  "tsu",  "te",   "to",
        "na",   "ni",   "nu",   "ne",   "no",
        "ha",   "hi",   "hu",   "he",   "ho",
        "ma",   "mi",   "mu",   "me",   "mo",
        "ya",           "yu",           "yo",
        "ra",   "ri",   "ru",   "re",   "ro",
        "wa",                           "wo",
    ];

    #[rustfmt::skip]
    const DIGRAPHS: &'static [&str] = &[
        "kya", "kyu", "kyo",
        "sha", "shu", "sho",
        "cha", "chu", "cho",
        "nya", "nyu", "nyo",
        "hya", "hyu", "hyo",
        "mya", "myu", "myo",
        "rya", "ryu", "ryo",
    ];

    #[rustfmt::skip]
    const DIACRITICS: &'static [&str] = &[
        "ga", "gi", "gu", "ge", "go",
        "za", "ji", "zu", "ze", "zo",
        "da",             "de", "do",
        "ba", "bi", "bu", "be", "bo",
        "pa", "pi", "pu", "pe", "po",
    ];

    #[rustfmt::skip]
    const DIGRAPHS_WITH_DIACRITICS: &'static [&str] = &[
        "gya",  "gyu",  "gyo",
        "ja",   "ju",   "jo",
        "bya",  "byu",  "byo",
        "pya",  "pyu",  "pyo",
    ];

    fn all() -> Vec<&'static str> {
        let mut all_kana: Vec<&'static str> = Vec::new();
        all_kana.extend(BASE);
        all_kana.extend(DIGRAPHS);
        all_kana.extend(DIACRITICS);
        all_kana.extend(DIGRAPHS_WITH_DIACRITICS);
        all_kana
    }

    fn simple() -> Vec<&'static str> {
        let mut simple_kana: Vec<&'static str> = Vec::new();
        simple_kana.extend(BASE);
        simple_kana.extend(DIACRITICS);
        simple_kana
    }

    pub fn get(use_simple: bool) -> Vec<&'static str> {
        match use_simple {
            true => simple(),
            false => all(),
        }
    }
} // mod kana

mod helpers {
    use super::{EnglishPasswordOptions, KanaPasswordOptions, PasswordOptions};
    use rand::seq::SliceRandom;
    use std::convert::TryInto;
    use std::io::Read;

    // Reads the dictionary at |path| and returns a linewise vector of
    // its contents.
    fn ingest_english_dictionary(path: &std::path::Path) -> std::io::Result<Vec<String>> {
        let mut file = std::fs::File::open(path)?;
        let mut contents = String::new();
        file.read_to_string(&mut contents)?;
        Ok(contents
            .lines()
            .map(|word| word.trim().to_lowercase())
            .collect())
    }

    // Assumes ownership of a Vector of owned |words| and prints them
    // separated by spaces.
    fn print_words(words: Vec<String>) {
        let mut iter = words.into_iter();
        if let Some(word) = iter.next() {
            print!("{}", word);
            for word in iter {
                print!(" {}", word);
            }
            print!("\n");
        }
    }

    // Returns 1 if we fail to read the word dictionary.
    fn print_english_password(options: EnglishPasswordOptions) -> i32 {
        let all_words =
            match ingest_english_dictionary(std::path::Path::new(&options.dictionary_path)) {
                Ok(result) => result,
                Err(error) => {
                    eprintln!("Error reading dictionary: {}", error);
                    return 1;
                }
            };
        let mut rng = &mut rand::thread_rng();
        let selection: Vec<String> = all_words
            .choose_multiple(&mut rng, options.num_words.try_into().unwrap())
            .map(|borrowed| borrowed.to_owned())
            .collect();

        print_words(selection);
        0
    }

    // Builds a single random pseudo-Japanese word from |kana_set|.
    fn build_kana_word(kana_set: &Vec<&'static str>, syllable_count: &u8) -> String {
        let mut collected_kana: Vec<&'static str> = Vec::new();
        let mut rng = &mut rand::thread_rng();
        for _ in 0..*syllable_count {
            collected_kana.push(kana_set.choose(&mut rng).unwrap());
        }
        collected_kana.join("")
    }

    // Builds random pseudo-Japanese words from |kana_set|, observing
    // the word lengths specified by |options|.
    fn build_kana_words(kana_set: Vec<&'static str>, options: KanaPasswordOptions) -> Vec<String> {
        options
            .syllable_counts
            .iter()
            .map(|count| build_kana_word(&kana_set, count))
            .collect()
    }

    // Builds a vector of owned kana words, maps these to a vector of
    // borrowed kana words, and then prints the same.
    fn print_kana_password(options: KanaPasswordOptions) {
        print_words(build_kana_words(crate::kana::get(true), options));
    }

    pub fn main() -> i32 {
        let options = crate::args::parse_args();
        match options {
            PasswordOptions::English(options) => return print_english_password(options),
            PasswordOptions::Japanese(options) => print_kana_password(options),
        }
        0
    }
} // mod helpers

fn main() {
    std::process::exit(crate::helpers::main());
}
<|endoftext|>


<commit_before>pub fn quick_sort<T: PartialOrd>(arr: &mut [T]) {
    if arr.is_empty() {
        return;
    }
    quick_sort_portion(arr, 0, arr.len() - 1);
}

fn quick_sort_portion<T: PartialOrd>(arr: &mut [T], left: usize, right: usize) {
    if left >= right {
        return;
    }

    let pivot_idx = partition(&mut arr[left..=right]);
    quick_sort_portion(
        arr,
        left,
        if left + pivot_idx > 0 {
            left + pivot_idx - 1
        } else {
            0
        },
    );
    quick_sort_portion(arr, left + pivot_idx + 1, right);
}

fn partition<T: PartialOrd>(arr: &mut [T]) -> usize {
    let (init, last) = arr.split_at_mut(arr.len() - 1);
    let pivot = &last[0];
    let mut temp_pivot = 0;
    for idx in 0..init.len() {
        if init[idx] <= *pivot {
            init.swap(idx, temp_pivot);
            temp_pivot += 1;
        }
    }
    arr.swap(temp_pivot, arr.len() - 1);
    temp_pivot
}

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

    #[test]
    fn test_empty() {
        let mut arr: Vec<i32> = vec![];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![]);
    }

    #[test]
    fn test_one() {
        let mut arr = vec![3];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![3]);
    }

    #[test]
    fn test_already_sorted() {
        let mut arr = vec![1, 2, 3, 4, 5, 6];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![1, 2, 3, 4, 5, 6]);
    }

    #[test]
    fn test_reverse_sorted() {
        let mut arr = vec![6, 5, 4, 3, 2, 1];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![1, 2, 3, 4, 5, 6]);
    }

    #[test]
    fn test_pivot_with_duplicate() {
        let mut arr = vec![3, 1, 1, 3, 7, 6, 3];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![1, 1, 3, 3, 3, 6, 7]);
    }

    #[test]
    fn test_pivot_no_move_should_not_stack_overflow() {
        let mut arr = vec![3, 5];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![3, 5]);
    }

    #[test]
    fn test_1() {
        let mut arr = vec![3, 5, 7, 8, 4, 2];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![2, 3, 4, 5, 7, 8]);
    }
}
<commit_msg>Fix "type annotations needed" caused by introducing serde_json. Ref: https://github.com/serde-rs/json/issues/357<commit_after>pub fn quick_sort<T: PartialOrd>(arr: &mut [T]) {
    if arr.is_empty() {
        return;
    }
    quick_sort_portion(arr, 0, arr.len() - 1);
}

fn quick_sort_portion<T: PartialOrd>(arr: &mut [T], left: usize, right: usize) {
    if left >= right {
        return;
    }

    let pivot_idx = partition(&mut arr[left..=right]);
    quick_sort_portion(
        arr,
        left,
        if left + pivot_idx > 0 {
            left + pivot_idx - 1
        } else {
            0
        },
    );
    quick_sort_portion(arr, left + pivot_idx + 1, right);
}

fn partition<T: PartialOrd>(arr: &mut [T]) -> usize {
    let (init, last) = arr.split_at_mut(arr.len() - 1);
    let pivot = &last[0];
    let mut temp_pivot = 0;
    for idx in 0..init.len() {
        if init[idx] <= *pivot {
            init.swap(idx, temp_pivot);
            temp_pivot += 1;
        }
    }
    arr.swap(temp_pivot, arr.len() - 1);
    temp_pivot
}

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

    #[test]
    fn test_empty() {
        let mut arr: Vec<i32> = vec![];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![] as Vec<i32>);
    }

    #[test]
    fn test_one() {
        let mut arr = vec![3];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![3]);
    }

    #[test]
    fn test_already_sorted() {
        let mut arr = vec![1, 2, 3, 4, 5, 6];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![1, 2, 3, 4, 5, 6]);
    }

    #[test]
    fn test_reverse_sorted() {
        let mut arr = vec![6, 5, 4, 3, 2, 1];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![1, 2, 3, 4, 5, 6]);
    }

    #[test]
    fn test_pivot_with_duplicate() {
        let mut arr = vec![3, 1, 1, 3, 7, 6, 3];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![1, 1, 3, 3, 3, 6, 7]);
    }

    #[test]
    fn test_pivot_no_move_should_not_stack_overflow() {
        let mut arr = vec![3, 5];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![3, 5]);
    }

    #[test]
    fn test_1() {
        let mut arr = vec![3, 5, 7, 8, 4, 2];
        quick_sort(&mut arr);
        assert_eq!(arr, vec![2, 3, 4, 5, 7, 8]);
    }
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate extra;

use std::env;
use std::io::{stdout, stderr, Write};
use extra::io::WriteExt;
use extra::option::OptionalExt;

const MAN_PAGE: &'static str = /* @MANSTART{yes} */ r#"
NAME
    yes - print the letter y or a given word, forever.

SYNOPSIS
    seq [ [ -h | --help ] | [ word ] ]

DESCRIPTION
    The yes utility prints the word passed as an operand, forever. If no operand is passed, then
    it defaults to the letter 'y'.

OPTIONS
    -h
    --help
        Print this manual page.
"#; /* @MANEND */

fn main() {
    let stdout = stdout();
    let mut stdout = stdout.lock();
    let mut stderr = stderr();

    if env::args().count() == 2 {
        if let Some(arg) = env::args().nth(1) {
            if arg == "--help" || arg == "-h" {
                stdout.write_all(MAN_PAGE.as_bytes()).try(&mut stderr);
                stdout.flush().try(&mut stderr);
                return;
            }
        }
    }

    if env::args().count() >= 2 {
        let answer = env::args().skip(1).collect::<Vec<_>>().join(" ");
        let print = answer.as_bytes();
        loop {
            stdout.writeln(print).try(&mut stderr);
        }
    } else {
        loop {
            stdout.writeln(b"y").try(&mut stderr);
        }
    };
}
<commit_msg>Correct command name in yes help<commit_after>#![deny(warnings)]

extern crate extra;

use std::env;
use std::io::{stdout, stderr, Write};
use extra::io::WriteExt;
use extra::option::OptionalExt;

const MAN_PAGE: &'static str = /* @MANSTART{yes} */ r#"
NAME
    yes - print the letter y or a given word, forever.

SYNOPSIS
    yes [ [ -h | --help ] | [ word ] ]

DESCRIPTION
    The yes utility prints the word passed as an operand, forever. If no operand is passed, then
    it defaults to the letter 'y'.

OPTIONS
    -h
    --help
        Print this manual page.
"#; /* @MANEND */

fn main() {
    let stdout = stdout();
    let mut stdout = stdout.lock();
    let mut stderr = stderr();

    if env::args().count() == 2 {
        if let Some(arg) = env::args().nth(1) {
            if arg == "--help" || arg == "-h" {
                stdout.write_all(MAN_PAGE.as_bytes()).try(&mut stderr);
                stdout.flush().try(&mut stderr);
                return;
            }
        }
    }

    if env::args().count() >= 2 {
        let answer = env::args().skip(1).collect::<Vec<_>>().join(" ");
        let print = answer.as_bytes();
        loop {
            stdout.writeln(print).try(&mut stderr);
        }
    } else {
        loop {
            stdout.writeln(b"y").try(&mut stderr);
        }
    };
}
<|endoftext|>


<commit_before>use std::{
    io,
    net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
    time::Duration,
};

use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpStream, UdpSocket},
};

use log::debug;

use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{Name, RData, RecordType},
};

use crate::{
    context::SharedContext,
    relay::{
        socks5::{
            Address,
            Command,
            HandshakeRequest,
            HandshakeResponse,
            Reply,
            TcpRequestHeader,
            TcpResponseHeader,
            SOCKS5_AUTH_METHOD_NONE,
        },
        utils::try_timeout,
    },
};

async fn udp_lookup(qname: &Name, qtype: RecordType, server: SocketAddr) -> io::Result<Message> {
    let bind_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0));
    let mut socket = UdpSocket::bind(bind_addr).await?;

    let mut message = Message::new();
    let mut query = Query::new();

    query.set_query_type(qtype);
    query.set_name(qname.clone());

    message.set_id(6666);
    message.set_recursion_desired(true);
    message.add_query(query);

    let req_buffer = message.to_vec()?;
    socket.send_to(&req_buffer, server).await?;

    let mut res_buffer = vec![0; 512];
    socket.recv_from(&mut res_buffer).await?;

    Ok(Message::from_vec(&mut res_buffer)?)
}

async fn socks5_lookup(qname: &Name, qtype: RecordType, socks5: SocketAddr, ns: SocketAddr) -> io::Result<Message> {
    let mut stream = TcpStream::connect(socks5).await?;

    // 1. Handshake
    let hs = HandshakeRequest::new(vec![SOCKS5_AUTH_METHOD_NONE]);
    hs.write_to(&mut stream).await?;
    stream.flush().await?;

    let hsp = HandshakeResponse::read_from(&mut stream).await?;
    assert_eq!(hsp.chosen_method, SOCKS5_AUTH_METHOD_NONE);

    // 2. Send request header
    let addr = Address::SocketAddress(ns);
    let h = TcpRequestHeader::new(Command::TcpConnect, addr);
    h.write_to(&mut stream).await?;
    stream.flush().await?;

    let hp = TcpResponseHeader::read_from(&mut stream).await?;
    match hp.reply {
        Reply::Succeeded => (),
        r => {
            let err = io::Error::new(io::ErrorKind::Other, format!("{}", r));
            return Err(err);
        }
    }

    let mut message = Message::new();
    let mut query = Query::new();

    query.set_query_type(qtype);
    query.set_name(qname.clone());

    message.set_id(6666);
    message.set_recursion_desired(true);
    message.add_query(query);

    let req_buffer = message.to_vec()?;
    let size = req_buffer.len();
    let mut size_buffer: [u8; 2] = [((size >> 8) & 0xFF) as u8, ((size >> 0) & 0xFF) as u8];
    let mut send_buffer: [u8; 512 + 2] = [0; 512 + 2];
    send_buffer[..2].copy_from_slice(&size_buffer[..2]);
    send_buffer[2..size + 2].copy_from_slice(&req_buffer[0..size]);
    stream.write_all(&send_buffer[0..size + 2]).await?;

    stream.read_exact(&mut size_buffer[0..2]).await?;

    let mut res_buffer = vec![0; 512];
    let size = ((size_buffer[0] as usize) << 8) + (size_buffer[1] as usize);
    stream.read_exact(&mut res_buffer[0..size]).await?;

    Ok(Message::from_vec(&mut res_buffer)?)
}

async fn acl_lookup(
    context: &SharedContext,
    local: SocketAddr,
    remote: SocketAddr,
    socks5: SocketAddr,
    qname: &Name,
    qtype: RecordType,
) -> io::Result<Message> {
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {} and {}",
        qtype, qname, local, remote
    );

    let timeout = Some(Duration::new(5, 0));

    let local_response = try_timeout(udp_lookup(qname, qtype.clone(), local), timeout)
        .await
        .unwrap_or(Message::new());
    let remote_response = try_timeout(socks5_lookup(qname, qtype.clone(), socks5, remote), timeout)
        .await
        .unwrap_or(Message::new());

    let addr = Address::DomainNameAddress(qname.to_string(), 0);
    let qname_bypassed = context.check_target_bypassed(&addr).await;

    let mut ip_bypassed = false;
    for rec in local_response.answers() {
        let bypassed = match rec.rdata() {
            RData::A(ref ip) => {
                let addr = Address::SocketAddress(SocketAddr::new(IpAddr::from(*ip), 0));
                context.check_target_bypassed(&addr).await
            }
            RData::AAAA(ref ip) => {
                let addr = Address::SocketAddress(SocketAddr::new(IpAddr::from(*ip), 0));
                context.check_target_bypassed(&addr).await
            }
            _ => false,
        };
        if bypassed {
            ip_bypassed = true;
        }
    }

    if local_response.answer_count() == 0 {
        return Ok(remote_response.clone());
    }

    if remote_response.answer_count() == 0 {
        return Ok(local_response.clone());
    }

    if qname_bypassed {
        debug!("Pick local response");
        Ok(local_response.clone())
    } else if ip_bypassed {
        debug!("Pick local response");
        Ok(local_response.clone())
    } else {
        debug!("Pick remote response");
        Ok(remote_response.clone())
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = context.config().local_dns_addr.expect("local dns");
    debug!("Local DNS server: {}", local_addr);

    let remote_addr = context.config().remote_dns_addr.expect("remote dns");
    debug!("Remote DNS server: {}", remote_addr);

    let socks5_config = context.config().local.as_ref().expect("socks5");
    let socks5_addr = socks5_config.bind_addr(&*context).await?;
    debug!("SOCKS5 server: {}", socks5_addr);

    let listen_addr = context.config().dns_relay_addr.expect("dns relay");
    debug!("Listen on {}", listen_addr);

    let mut socket = UdpSocket::bind(listen_addr).await?;

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match socket.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                debug!("Failed to read from UDP socket: {:?}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&mut req_buffer) {
            Ok(x) => x,
            Err(e) => {
                debug!("Failed to parse UDP query message: {:?}", e);
                continue;
            }
        };

        let mut message = Message::new();
        message.set_id(request.id());
        message.set_recursion_desired(true);
        message.set_recursion_available(true);
        message.set_message_type(MessageType::Response);

        if request.queries().is_empty() {
            message.set_response_code(ResponseCode::FormErr);
        } else {
            let question = &request.queries()[0];
            debug!("Received query: {:?}", question);

            if let Ok(result) = acl_lookup(
                &context,
                local_addr,
                remote_addr,
                socks5_addr,
                question.name(),
                question.query_type(),
            )
            .await
            {
                message.add_query(question.clone());
                message.set_response_code(result.response_code());

                for rec in result.answers() {
                    debug!("Answer: {:?}", rec);
                    match rec.rdata() {
                        RData::A(ref ip) => {
                            context.add_to_reverse_lookup_cache(IpAddr::from(*ip), question.name().to_ascii())
                        }
                        RData::AAAA(ref ip) => {
                            context.add_to_reverse_lookup_cache(IpAddr::from(*ip), question.name().to_ascii())
                        }
                        _ => (),
                    };
                    message.add_answer(rec.clone());
                }
                for rec in result.additionals() {
                    debug!("Additionals: {:?}", rec);
                    message.add_additional(rec.clone());
                }
            } else {
                message.set_response_code(ResponseCode::ServFail);
            }
        }

        let res_buffer = message.to_vec()?;
        match socket.send_to(&res_buffer, src).await {
            Ok(_) => {}
            Err(e) => {
                debug!("Failed to send response: {:?}", e);
                continue;
            }
        };
    }
}
<commit_msg>Refine message id<commit_after>use std::{
    io,
    net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
    time::Duration,
};

use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpStream, UdpSocket},
};

use byteorder::{BigEndian, ByteOrder};
use log::debug;
use rand::Rng;

use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{Name, RData, RecordType},
};

use crate::{
    context::SharedContext,
    relay::{
        socks5::{
            Address,
            Command,
            HandshakeRequest,
            HandshakeResponse,
            Reply,
            TcpRequestHeader,
            TcpResponseHeader,
            SOCKS5_AUTH_METHOD_NONE,
        },
        utils::try_timeout,
    },
};

async fn udp_lookup(qname: &Name, qtype: RecordType, server: SocketAddr) -> io::Result<Message> {
    let bind_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0));
    let mut socket = UdpSocket::bind(bind_addr).await?;

    let mut message = Message::new();
    let mut query = Query::new();

    query.set_query_type(qtype);
    query.set_name(qname.clone());

    let id = rand::thread_rng().gen();
    message.set_id(id);
    message.set_recursion_desired(true);
    message.add_query(query);

    let req_buffer = message.to_vec()?;
    socket.send_to(&req_buffer, server).await?;

    let mut res_buffer = vec![0; 512];
    socket.recv_from(&mut res_buffer).await?;

    Ok(Message::from_vec(&mut res_buffer)?)
}

async fn socks5_lookup(qname: &Name, qtype: RecordType, socks5: SocketAddr, ns: SocketAddr) -> io::Result<Message> {
    let mut stream = TcpStream::connect(socks5).await?;

    // 1. Handshake
    let hs = HandshakeRequest::new(vec![SOCKS5_AUTH_METHOD_NONE]);
    hs.write_to(&mut stream).await?;
    stream.flush().await?;

    let hsp = HandshakeResponse::read_from(&mut stream).await?;
    assert_eq!(hsp.chosen_method, SOCKS5_AUTH_METHOD_NONE);

    // 2. Send request header
    let addr = Address::SocketAddress(ns);
    let h = TcpRequestHeader::new(Command::TcpConnect, addr);
    h.write_to(&mut stream).await?;
    stream.flush().await?;

    let hp = TcpResponseHeader::read_from(&mut stream).await?;
    match hp.reply {
        Reply::Succeeded => (),
        r => {
            let err = io::Error::new(io::ErrorKind::Other, format!("{}", r));
            return Err(err);
        }
    }

    let mut message = Message::new();
    let mut query = Query::new();

    query.set_query_type(qtype);
    query.set_name(qname.clone());

    let id = rand::thread_rng().gen();
    message.set_id(id);
    message.set_recursion_desired(true);
    message.add_query(query);

    let req_buffer = message.to_vec()?;
    let size = req_buffer.len();
    let mut send_buffer = vec![];

    BigEndian::write_u16(&mut send_buffer[0..2], size as u16);
    send_buffer[2..size + 2].copy_from_slice(&req_buffer[0..size]);
    stream.write_all(&send_buffer[0..size + 2]).await?;

    let mut res_buffer = vec![];
    stream.read_exact(&mut res_buffer[0..2]).await?;

    let size = BigEndian::read_u16(&res_buffer[0..2]) as usize;
    stream.read_exact(&mut res_buffer[0..size]).await?;

    Ok(Message::from_vec(&mut res_buffer)?)
}

async fn acl_lookup(
    context: &SharedContext,
    local: SocketAddr,
    remote: SocketAddr,
    socks5: SocketAddr,
    qname: &Name,
    qtype: RecordType,
) -> io::Result<Message> {
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {} and {}",
        qtype, qname, local, remote
    );

    let timeout = Some(Duration::new(5, 0));

    let local_response = try_timeout(udp_lookup(qname, qtype.clone(), local), timeout)
        .await
        .unwrap_or(Message::new());
    let remote_response = try_timeout(socks5_lookup(qname, qtype.clone(), socks5, remote), timeout)
        .await
        .unwrap_or(Message::new());

    let addr = Address::DomainNameAddress(qname.to_string(), 0);
    let qname_bypassed = context.check_target_bypassed(&addr).await;

    let mut ip_bypassed = false;
    for rec in local_response.answers() {
        let bypassed = match rec.rdata() {
            RData::A(ref ip) => {
                let addr = Address::SocketAddress(SocketAddr::new(IpAddr::from(*ip), 0));
                context.check_target_bypassed(&addr).await
            }
            RData::AAAA(ref ip) => {
                let addr = Address::SocketAddress(SocketAddr::new(IpAddr::from(*ip), 0));
                context.check_target_bypassed(&addr).await
            }
            _ => false,
        };
        if bypassed {
            ip_bypassed = true;
        }
    }

    if local_response.answer_count() == 0 {
        return Ok(remote_response.clone());
    }

    if remote_response.answer_count() == 0 {
        return Ok(local_response.clone());
    }

    if qname_bypassed {
        debug!("Pick local response");
        Ok(local_response.clone())
    } else if ip_bypassed {
        debug!("Pick local response");
        Ok(local_response.clone())
    } else {
        debug!("Pick remote response");
        Ok(remote_response.clone())
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = context.config().local_dns_addr.expect("local dns");
    debug!("Local DNS server: {}", local_addr);

    let remote_addr = context.config().remote_dns_addr.expect("remote dns");
    debug!("Remote DNS server: {}", remote_addr);

    let socks5_config = context.config().local.as_ref().expect("socks5");
    let socks5_addr = socks5_config.bind_addr(&*context).await?;
    debug!("SOCKS5 server: {}", socks5_addr);

    let listen_addr = context.config().dns_relay_addr.expect("dns relay");
    debug!("Listen on {}", listen_addr);

    let mut socket = UdpSocket::bind(listen_addr).await?;

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match socket.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                debug!("Failed to read from UDP socket: {:?}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&mut req_buffer) {
            Ok(x) => x,
            Err(e) => {
                debug!("Failed to parse UDP query message: {:?}", e);
                continue;
            }
        };

        let mut message = Message::new();
        message.set_id(request.id());
        message.set_recursion_desired(true);
        message.set_recursion_available(true);
        message.set_message_type(MessageType::Response);

        if request.queries().is_empty() {
            message.set_response_code(ResponseCode::FormErr);
        } else {
            let question = &request.queries()[0];
            debug!("Received query: {:?}", question);

            if let Ok(result) = acl_lookup(
                &context,
                local_addr,
                remote_addr,
                socks5_addr,
                question.name(),
                question.query_type(),
            )
            .await
            {
                message.add_query(question.clone());
                message.set_response_code(result.response_code());

                for rec in result.answers() {
                    debug!("Answer: {:?}", rec);
                    match rec.rdata() {
                        RData::A(ref ip) => {
                            context.add_to_reverse_lookup_cache(IpAddr::from(*ip), question.name().to_ascii())
                        }
                        RData::AAAA(ref ip) => {
                            context.add_to_reverse_lookup_cache(IpAddr::from(*ip), question.name().to_ascii())
                        }
                        _ => (),
                    };
                    message.add_answer(rec.clone());
                }
                for rec in result.additionals() {
                    debug!("Additionals: {:?}", rec);
                    message.add_additional(rec.clone());
                }
            } else {
                message.set_response_code(ResponseCode::ServFail);
            }
        }

        let res_buffer = message.to_vec()?;
        match socket.send_to(&res_buffer, src).await {
            Ok(_) => {}
            Err(e) => {
                debug!("Failed to send response: {:?}", e);
                continue;
            }
        };
    }
}
<|endoftext|>


<commit_before>extern crate sdl2;
extern crate minifb;
use minifb::*;
pub mod cpu;
pub mod opcodes;
//mod clock;
mod memory;
mod cia;
mod vic;

pub const SCREEN_WIDTH:  usize = 384; // extend 20 pixels left and right for the borders
pub const SCREEN_HEIGHT: usize = 272; // extend 36 pixels top and down for the borders

pub struct C64
{
    pub window: minifb::Window,
    memory: memory::MemShared,
    //clock: clock::Clock,
    cpu: cpu::CPUShared,
    cia1: cia::CIAShared,
    cia2: cia::CIAShared,
    vic: vic::VICShared,

    cycle_count: u32,
}

impl C64
{
    pub fn new() -> C64
    {
        let memory = memory::Memory::new_shared();
        let vic    = vic::VIC::new_shared();
        let cia1   = cia::CIA::new_shared(true);
        let cia2   = cia::CIA::new_shared(false);
        let cpu    = cpu::CPU::new_shared();

        let c64 = C64
        {
            window: Window::new("Rust64", SCREEN_WIDTH, SCREEN_HEIGHT, Scale::X4, Vsync::No).unwrap(),
            memory: memory.clone(), // shared system memory (RAM, ROM, IO registers)
            //clock: clock::Clock::new(),
            cpu: cpu.clone(),
            cia1: cia1.clone(),
            cia2: cia2.clone(),
            vic: vic.clone(),
            cycle_count: 0,
        };

        // cyclic dependencies are not possible in Rust (yet?), so we have
        // to resort to setting references manually
        c64.cia1.borrow_mut().set_references(memory.clone(), cpu.clone(), vic.clone());
        c64.cia2.borrow_mut().set_references(memory.clone(), cpu.clone(), vic.clone());
        c64.vic.borrow_mut().set_references(memory.clone(), cpu.clone());
        c64.cpu.borrow_mut().set_references(memory.clone(), vic.clone(), cia1.clone(), cia2.clone());

        drop(memory);
        drop(cia1);
        drop(cia2);
        drop(vic);
        drop(cpu);

        c64
    }

    pub fn reset(&mut self)
    {
        self.memory.borrow_mut().reset();
        self.cpu.borrow_mut().reset();
        self.cia1.borrow_mut().reset();
        self.cia2.borrow_mut().reset();
    }

    pub fn update(&mut self)
    {
        let mut should_trigger_vblank = false;
        //if self.clock.tick() { println!("Clock tick"); }
        self.vic.borrow_mut().update(self.cycle_count, &mut should_trigger_vblank);
        // update sid here when it's done

        self.cia1.borrow_mut().process_irq();
        self.cia2.borrow_mut().process_irq();
        self.cia1.borrow_mut().update();
        self.cia2.borrow_mut().update();

        self.cpu.borrow_mut().update();

        if should_trigger_vblank
        {
            //let w = self.vic.borrow_mut().window_buffer[0];
            //println!("w: {}", w);
            //self.window.update(&self.vic.borrow_mut().window_buffer);
            self.cia1.borrow_mut().count_tod();
            self.cia2.borrow_mut().count_tod();
        }

        self.cycle_count += 1;
    }

    // debug
    pub fn render(&mut self)
    {
        //true
        /*for key in self.window.get_keys().iter()
        {
            match *key { _ => () }
        }*/
        self.window.update(&self.vic.borrow_mut().window_buffer);
    }
}
<commit_msg>default to X1 window scale for now<commit_after>extern crate sdl2;
extern crate minifb;
use minifb::*;
pub mod cpu;
pub mod opcodes;
//mod clock;
mod memory;
mod cia;
mod vic;

pub const SCREEN_WIDTH:  usize = 384; // extend 20 pixels left and right for the borders
pub const SCREEN_HEIGHT: usize = 272; // extend 36 pixels top and down for the borders

pub struct C64
{
    pub window: minifb::Window,
    memory: memory::MemShared,
    //clock: clock::Clock,
    cpu: cpu::CPUShared,
    cia1: cia::CIAShared,
    cia2: cia::CIAShared,
    vic: vic::VICShared,

    cycle_count: u32,
}

impl C64
{
    pub fn new() -> C64
    {
        let memory = memory::Memory::new_shared();
        let vic    = vic::VIC::new_shared();
        let cia1   = cia::CIA::new_shared(true);
        let cia2   = cia::CIA::new_shared(false);
        let cpu    = cpu::CPU::new_shared();

        let c64 = C64
        {
            window: Window::new("Rust64", SCREEN_WIDTH, SCREEN_HEIGHT, Scale::X1, Vsync::No).unwrap(),
            memory: memory.clone(), // shared system memory (RAM, ROM, IO registers)
            //clock: clock::Clock::new(),
            cpu: cpu.clone(),
            cia1: cia1.clone(),
            cia2: cia2.clone(),
            vic: vic.clone(),
            cycle_count: 0,
        };

        // cyclic dependencies are not possible in Rust (yet?), so we have
        // to resort to setting references manually
        c64.cia1.borrow_mut().set_references(memory.clone(), cpu.clone(), vic.clone());
        c64.cia2.borrow_mut().set_references(memory.clone(), cpu.clone(), vic.clone());
        c64.vic.borrow_mut().set_references(memory.clone(), cpu.clone());
        c64.cpu.borrow_mut().set_references(memory.clone(), vic.clone(), cia1.clone(), cia2.clone());

        drop(memory);
        drop(cia1);
        drop(cia2);
        drop(vic);
        drop(cpu);

        c64
    }

    pub fn reset(&mut self)
    {
        self.memory.borrow_mut().reset();
        self.cpu.borrow_mut().reset();
        self.cia1.borrow_mut().reset();
        self.cia2.borrow_mut().reset();
    }

    pub fn update(&mut self)
    {
        let mut should_trigger_vblank = false;
        //if self.clock.tick() { println!("Clock tick"); }
        self.vic.borrow_mut().update(self.cycle_count, &mut should_trigger_vblank);
        // update sid here when it's done

        self.cia1.borrow_mut().process_irq();
        self.cia2.borrow_mut().process_irq();
        self.cia1.borrow_mut().update();
        self.cia2.borrow_mut().update();

        self.cpu.borrow_mut().update();

        if should_trigger_vblank
        {
            //let w = self.vic.borrow_mut().window_buffer[0];
            //println!("w: {}", w);
            //self.window.update(&self.vic.borrow_mut().window_buffer);
            self.cia1.borrow_mut().count_tod();
            self.cia2.borrow_mut().count_tod();
        }

        self.cycle_count += 1;
    }

    // debug
    pub fn render(&mut self)
    {
        //true
        /*for key in self.window.get_keys().iter()
        {
            match *key { _ => () }
        }*/
        self.window.update(&self.vic.borrow_mut().window_buffer);
    }
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::io::BufReader;
use std::io::Error as IoError;
use std::ascii::AsciiExt;

use color::{ColorType};
use image::{DecodingResult, ImageDecoder, ImageResult, ImageError};
extern crate byteorder;
use self::byteorder::{BigEndian, ByteOrder};

enum DecodeStrategy {
    Bytes,
    Ascii,
}

/// PPM decoder
pub struct PPMDecoder<R> {
    reader: BufReader<R>,
    width: u32,
    height: u32,
    maxwhite: u32,
    samples: u32,
    decoder: DecodeStrategy,
}

impl<R: Read> PPMDecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(read: R) -> ImageResult<PPMDecoder<R>> {
        let mut buf = BufReader::new(read);
        let mut magic: [u8; 2] = [0, 0];
        try!(buf.read_exact(&mut magic[..])); // Skip magic constant
        if magic[0] != b'P' || (magic[1] != b'3' && magic[1] != b'6') {
            return Err(ImageError::FormatError("Expected magic constant for ppm, P3 or P6".to_string()));
        }

        // Remove this once the reader can read plain ppm
        if magic[1] == b'3' {
            return Err(ImageError::FormatError("Plain format is not yet supported".to_string()))
        }

        let width = try!(PPMDecoder::read_next_u32(&mut buf));
        let height = try!(PPMDecoder::read_next_u32(&mut buf));
        let maxwhite = try!(PPMDecoder::read_next_u32(&mut buf));

        if !(maxwhite <= u16::max_value() as u32) {
            return Err(ImageError::FormatError("Image maxval is not less or equal to 65535".to_string()))
        }

        Ok(PPMDecoder {
            reader: buf,
            width: width,
            height: height,
            maxwhite: maxwhite,
            samples: 3,
            decoder: DecodeStrategy::Bytes,
        })
    }

    /// Reads a string as well as a single whitespace after it, ignoring comments
    fn read_next_string(reader: &mut BufReader<R>) -> ImageResult<String> {
        let mut bytes = Vec::new();

        // pair input bytes with a bool mask to remove comments
        let mark_comments = reader
            .bytes()
            .scan(true, |partof, read| {
                let byte = match read {
                    Err(err) => return Some((*partof, Err(err))),
                    Ok(byte) => byte,
                };
                let cur_enabled = *partof && byte != b'#';
                let next_enabled = cur_enabled || (byte == b'\r' || byte == b'\n');
                *partof = next_enabled;
                return Some((cur_enabled, Ok(byte)));
            });

        for (_, byte) in mark_comments.filter(|ref e| e.0) {
            match byte {
                Ok(b'\t') | Ok(b'\n') | Ok(b'\x0b') | Ok(b'\x0c') | Ok(b'\r') | Ok(b' ') => {
                    if !bytes.is_empty() {
                        break // We're done as we already have some content
                    }
                },
                Ok(byte) => {
                    bytes.push(byte);
                },
                Err(_) => break,
            }
        }

        if bytes.is_empty() {
            return Err(ImageError::FormatError("Unexpected eof".to_string()))
        }

        if !bytes.as_slice().is_ascii() {
            return Err(ImageError::FormatError("Non ascii character in preamble".to_string()))
        }

        String::from_utf8(bytes).map_err(|_| ImageError::FormatError("Couldn't read preamble".to_string()))
    }

    fn read_next_u32(reader: &mut BufReader<R>) -> ImageResult<u32> {
        let s = try!(PPMDecoder::read_next_string(reader));
        s.parse::<u32>().map_err(|_| ImageError::FormatError("Invalid number in preamble".to_string()))
    }
}

impl<R: Read> ImageDecoder for PPMDecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        Ok((self.width, self.height))
    }

    fn colortype(&mut self) -> ImageResult<ColorType> {
        match (self.bytewidth(), self.components()) {
            (1, 3) => Ok(ColorType::RGB(8)),
            (2, 3) => Ok(ColorType::RGB(16)),
            _ => Err(ImageError::FormatError("Don't know how to decode PPM with more than 16 bits".to_string())),
        }
    }

    fn row_len(&mut self) -> ImageResult<usize> {
        Ok((self.width*self.components()*self.bytewidth()) as usize)
    }

    fn read_scanline(&mut self, _buf: &mut [u8]) -> ImageResult<u32> {
        unimplemented!();
    }

    fn read_image(&mut self) -> ImageResult<DecodingResult> {
        let pixelcount = self.width
            .checked_mul(self.height)
            .and_then(|v| v.checked_mul(self.components()));

        let pixelcount = match pixelcount {
            Some(v) => v,
            _ => return Err(ImageError::DimensionError),
        };

        self.read(pixelcount)
    }
}

impl<R: Read> PPMDecoder<R> {
    fn bytewidth(&self) -> u32 {
        if self.maxwhite < 256 { 1 } else { 2 }
    }

    fn components(&self) -> u32 {
        self.samples
    }

    fn read(&mut self, count: u32) -> ImageResult<DecodingResult> {
        if self.bytewidth() == 1 {
            let mut data = vec![0 as u8; count as usize];
            self.read_u8(&mut data)?;
            Ok(DecodingResult::U8(data))
        } else if self.bytewidth() == 2 {
            let mut data = vec![0 as u16; count as usize];
            self.read_u16(&mut data)?;
            Ok(DecodingResult::U16(data))
        } else {
            Err(ImageError::FormatError("Invalid sample bitwidth".to_string()))
        }
    }

    fn read_u8(&mut self, mut buffer: &mut [u8]) -> ImageResult<()> {
        match self.decoder {
            DecodeStrategy::Bytes => match self.reader.read_exact(&mut buffer) {
                    Ok(_) => Ok(()),
                    Err(e) => Err(ImageError::IoError(e)),
                },
            DecodeStrategy::Ascii => {
                    unimplemented!()
                }
        }
    }

    fn read_u16(&mut self, mut buffer: &mut [u16]) -> ImageResult<()> {
        match self.decoder {
            DecodeStrategy::Bytes => {
                    let mut bytebuffer = vec![0 as u8; buffer.len() * 2];
                    match self.reader.read_exact(&mut bytebuffer) {
                        Err(e) => return Err(ImageError::IoError(e)),
                        Ok(_) => {},
                    }
                    BigEndian::read_u16_into(&mut bytebuffer, &mut buffer);
                    Ok(())
                },
            DecodeStrategy::Ascii => {
                    unimplemented!()
                }
        }
    }
}

/// Tests parsing binary buffers were written based on and validated against `pamfile` from
/// netpbm (http://netpbm.sourceforge.net/).
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn minimal_form() {
        // Violates current specification (October 2016 ) but accepted by both netpbm and ImageMagick
        decode_minimal_image(&b"P61 1 255 123"[..]);
        decode_minimal_image(&b"P6 1 1 255 123"[..]);
        decode_minimal_image(&b"P6 1 1 255 123\xFF"[..]); // Too long should not be an issue
    }

    #[test]
    fn comment_in_token() {
        decode_minimal_image(&b"P6 1 1 2#comment\n55 123"[..]); // Terminating LF
        decode_minimal_image(&b"P6 1 1 2#comment\r55 123"[..]); // Terminating CR
        decode_minimal_image(&b"P6 1 1#comment\n 255 123"[..]); // Comment after token
        decode_minimal_image(&b"P6 1 1 #comment\n255 123"[..]); // Comment before token
        decode_minimal_image(&b"P6#comment\n 1 1 255 123"[..]); // Begin of header
        decode_minimal_image(&b"P6 1 1 255#comment\n 123"[..]); // End of header
    }

    #[test]
    fn whitespace() {
        decode_minimal_image(&b"P6\x091\x091\x09255\x09123"[..]); // TAB
        decode_minimal_image(&b"P6\x0a1\x0a1\x0a255\x0a123"[..]); // LF
        decode_minimal_image(&b"P6\x0b1\x0b1\x0b255\x0b123"[..]); // VT
        decode_minimal_image(&b"P6\x0c1\x0c1\x0c255\x0c123"[..]); // FF
        decode_minimal_image(&b"P6\x0d1\x0d1\x0d255\x0d123"[..]); // CR
        // Spaces tested before
        decode_minimal_image(&b"P61\x09\x0a\x0b\x0c\x0d1 255 123"[..]); // All whitespace, combined
    }

    /// Tests for decoding error, assuming `encoded` is ppm encoding for the very simplistic image
    /// containing a single pixel with one byte values (1, 2, 3).
    fn decode_minimal_image(encoded: &[u8]) {
        let content = vec![49 as u8, 50, 51];
        let mut decoder = PPMDecoder::new(encoded).unwrap();

        assert_eq!(decoder.dimensions().unwrap(), (1, 1));
        assert_eq!(decoder.colortype().unwrap(), ColorType::RGB(8));
        assert_eq!(decoder.row_len().unwrap(), 3);
        assert_eq!(decoder.bytewidth(), 1);

        match decoder.read_image().unwrap() {
            DecodingResult::U8(image) => assert_eq!(image, content),
            _ => assert!(false),
        }
    }

    #[test]
    fn wrong_tag() {
        assert!(PPMDecoder::new(&b"P5 1 1 255 1"[..]).is_err());
    }

    #[test]
    fn invalid_characters() {
        assert!(PPMDecoder::new(&b"P6 1chars1 255 1"[..]).is_err()); // No text outside of comments
        assert!(PPMDecoder::new(&b"P6 1\xFF1 255 1"[..]).is_err()); // No invalid ascii chars
        assert!(PPMDecoder::new(&b"P6 0x01 1 255 1"[..]).is_err()); // Numbers only as decimal
    }

    /// These violate the narrow specification of ppm but are commonly supported in other programs.
    /// Fail fast and concise is important here as these might be received as input files.
    #[test]
    fn unsupported_extensions() {
        assert!(PPMDecoder::new(&b"P6 1 1 65536 1"[..]).is_err()); // No bitwidth above 16
    }
}
<commit_msg>Add support for ascii decoding<commit_after>use std::io::Read;
use std::io::BufReader;
use std::ascii::AsciiExt;

use color::{ColorType};
use image::{DecodingResult, ImageDecoder, ImageResult, ImageError};
extern crate byteorder;
use self::byteorder::{BigEndian, ByteOrder};

enum DecodeStrategy {
    Bytes,
    Ascii,
}

/// PPM decoder
pub struct PPMDecoder<R> {
    reader: BufReader<R>,
    width: u32,
    height: u32,
    maxwhite: u32,
    samples: u32,
    decoder: DecodeStrategy,
}

impl<R: Read> PPMDecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(read: R) -> ImageResult<PPMDecoder<R>> {
        let mut buf = BufReader::new(read);
        let mut magic: [u8; 2] = [0, 0];
        try!(buf.read_exact(&mut magic[..])); // Skip magic constant
        if magic[0] != b'P' || (magic[1] != b'3' && magic[1] != b'6') {
            return Err(ImageError::FormatError("Expected magic constant for ppm, P3 or P6".to_string()));
        }

        let decoder = match magic[1] {
            b'3' => DecodeStrategy::Ascii,
            b'6' => DecodeStrategy::Bytes,
            _ => unreachable!(),
        };

        let width = try!(PPMDecoder::read_next_u32(&mut buf));
        let height = try!(PPMDecoder::read_next_u32(&mut buf));
        let maxwhite = try!(PPMDecoder::read_next_u32(&mut buf));

        if !(maxwhite <= u16::max_value() as u32) {
            return Err(ImageError::FormatError("Image maxval is not less or equal to 65535".to_string()))
        }

        Ok(PPMDecoder {
            reader: buf,
            width: width,
            height: height,
            maxwhite: maxwhite,
            samples: 3,
            decoder: decoder,
        })
    }

    /// Reads a string as well as a single whitespace after it, ignoring comments
    fn read_next_string(reader: &mut BufReader<R>) -> ImageResult<String> {
        let mut bytes = Vec::new();

        // pair input bytes with a bool mask to remove comments
        let mark_comments = reader
            .bytes()
            .scan(true, |partof, read| {
                let byte = match read {
                    Err(err) => return Some((*partof, Err(err))),
                    Ok(byte) => byte,
                };
                let cur_enabled = *partof && byte != b'#';
                let next_enabled = cur_enabled || (byte == b'\r' || byte == b'\n');
                *partof = next_enabled;
                return Some((cur_enabled, Ok(byte)));
            });

        for (_, byte) in mark_comments.filter(|ref e| e.0) {
            match byte {
                Ok(b'\t') | Ok(b'\n') | Ok(b'\x0b') | Ok(b'\x0c') | Ok(b'\r') | Ok(b' ') => {
                    if !bytes.is_empty() {
                        break // We're done as we already have some content
                    }
                },
                Ok(byte) => {
                    bytes.push(byte);
                },
                Err(_) => break,
            }
        }

        if bytes.is_empty() {
            return Err(ImageError::FormatError("Unexpected eof".to_string()))
        }

        if !bytes.as_slice().is_ascii() {
            return Err(ImageError::FormatError("Non ascii character in preamble".to_string()))
        }

        String::from_utf8(bytes).map_err(|_| ImageError::FormatError("Couldn't read preamble".to_string()))
    }

    fn read_next_u32(reader: &mut BufReader<R>) -> ImageResult<u32> {
        let s = try!(PPMDecoder::read_next_string(reader));
        s.parse::<u32>().map_err(|_| ImageError::FormatError("Invalid number in preamble".to_string()))
    }
}

impl<R: Read> ImageDecoder for PPMDecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        Ok((self.width, self.height))
    }

    fn colortype(&mut self) -> ImageResult<ColorType> {
        match (self.bytewidth(), self.components()) {
            (1, 1) => Ok(ColorType::Gray(8)),
            (2, 1) => Ok(ColorType::Gray(16)),
            (1, 3) => Ok(ColorType::RGB(8)),
            (2, 3) => Ok(ColorType::RGB(16)),
            _ => Err(ImageError::FormatError("Don't know how to decode PPM with more than 16 bits".to_string())),
        }
    }

    fn row_len(&mut self) -> ImageResult<usize> {
        Ok((self.width*self.components()*self.bytewidth()) as usize)
    }

    fn read_scanline(&mut self, _buf: &mut [u8]) -> ImageResult<u32> {
        unimplemented!();
    }

    fn read_image(&mut self) -> ImageResult<DecodingResult> {
        let pixelcount = self.width
            .checked_mul(self.height)
            .and_then(|v| v.checked_mul(self.components()));

        let pixelcount = match pixelcount {
            Some(v) => v,
            _ => return Err(ImageError::DimensionError),
        };

        self.read(pixelcount)
    }
}

impl<R: Read> PPMDecoder<R> {
    fn bytewidth(&self) -> u32 {
        if self.maxwhite < 256 { 1 } else { 2 }
    }

    fn components(&self) -> u32 {
        self.samples
    }

    fn read(&mut self, count: u32) -> ImageResult<DecodingResult> {
        if self.bytewidth() == 1 {
            let mut data = vec![0 as u8; count as usize];
            self.read_u8(&mut data)?;
            Ok(DecodingResult::U8(data))
        } else if self.bytewidth() == 2 {
            let mut data = vec![0 as u16; count as usize];
            self.read_u16(&mut data)?;
            Ok(DecodingResult::U16(data))
        } else {
            Err(ImageError::FormatError("Invalid sample bitwidth".to_string()))
        }
    }

    fn read_u8(&mut self, mut buffer: &mut [u8]) -> ImageResult<()> {
        match self.decoder {
            DecodeStrategy::Bytes => match self.reader.read_exact(&mut buffer) {
                    Ok(_) => Ok(()),
                    Err(e) => Err(ImageError::IoError(e)),
                },
            DecodeStrategy::Ascii => {
                    self.read_ascii::<u8>(&mut buffer)
                }
        }
    }

    fn read_u16(&mut self, mut buffer: &mut [u16]) -> ImageResult<()> {
        match self.decoder {
            DecodeStrategy::Bytes => {
                    let mut bytebuffer = vec![0 as u8; buffer.len() * 2];
                    match self.reader.read_exact(&mut bytebuffer) {
                        Err(e) => return Err(ImageError::IoError(e)),
                        Ok(_) => {},
                    }
                    BigEndian::read_u16_into(&mut bytebuffer, &mut buffer);
                    Ok(())
                },
            DecodeStrategy::Ascii => {
                    self.read_ascii::<u16>(&mut buffer)
                }
        }
    }

    fn read_ascii<Basic: FromSample>(&mut self, buffer: &mut [Basic::T]) -> ImageResult<()> {
        for pixel in buffer {
            let value = self.read_sample()?;
            *pixel = Basic::try_convert(value)?;
        }
        Ok(())
    }

    fn read_sample(&mut self) -> ImageResult<u32> {
        let token = (&mut self.reader).bytes()
            .take_while(|v| match v {
                    &Err(_) => false,
                    &Ok(b'\t') | &Ok(b'\n') | &Ok(b'\x0b') | &Ok(b'\x0c') | &Ok(b'\r') | &Ok(b' ') => false,
                    _ => true,
                })
            .collect::<Result<Vec<u8>, _>>()?;
        if !token.is_ascii() {
            return Err(ImageError::FormatError("Non ascii character where sample value was expected".to_string()))
        }
        let string = String::from_utf8(token).map_err(|_| ImageError::FormatError("Error parsing sample".to_string()))?;
        string.parse::<u32>().map_err(|_| ImageError::FormatError("Error parsing sample value".to_string()))
    }
}

trait FromSample {
    type T;
    fn try_convert(u32) -> ImageResult<Self::T>;
}

impl FromSample for u8 {
    type T = u8;
    fn try_convert(val: u32) -> ImageResult<Self::T> {
        if val > u8::max_value() as u32 {
            Err(ImageError::FormatError("Sample value outside of bounds".to_string()))
        } else {
            Ok(val as u8)
        }
    }
}

impl FromSample for u16 {
    type T = u16;
    fn try_convert(val: u32) -> ImageResult<Self::T> {
        if val > u16::max_value() as u32 {
            Err(ImageError::FormatError("Sample value outside of bounds".to_string()))
        } else {
            Ok(val as u16)
        }
    }
}

/// Tests parsing binary buffers were written based on and validated against `pamfile` from
/// netpbm (http://netpbm.sourceforge.net/).
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn minimal_form() {
        // Violates current specification (October 2016 ) but accepted by both netpbm and ImageMagick
        decode_minimal_image(&b"P61 1 255 123"[..]);
        decode_minimal_image(&b"P6 1 1 255 123"[..]);
        decode_minimal_image(&b"P6 1 1 255 123\xFF"[..]); // Too long should not be an issue
    }

    #[test]
    fn comment_in_token() {
        decode_minimal_image(&b"P6 1 1 2#comment\n55 123"[..]); // Terminating LF
        decode_minimal_image(&b"P6 1 1 2#comment\r55 123"[..]); // Terminating CR
        decode_minimal_image(&b"P6 1 1#comment\n 255 123"[..]); // Comment after token
        decode_minimal_image(&b"P6 1 1 #comment\n255 123"[..]); // Comment before token
        decode_minimal_image(&b"P6#comment\n 1 1 255 123"[..]); // Begin of header
        decode_minimal_image(&b"P6 1 1 255#comment\n 123"[..]); // End of header
    }

    #[test]
    fn whitespace() {
        decode_minimal_image(&b"P6\x091\x091\x09255\x09123"[..]); // TAB
        decode_minimal_image(&b"P6\x0a1\x0a1\x0a255\x0a123"[..]); // LF
        decode_minimal_image(&b"P6\x0b1\x0b1\x0b255\x0b123"[..]); // VT
        decode_minimal_image(&b"P6\x0c1\x0c1\x0c255\x0c123"[..]); // FF
        decode_minimal_image(&b"P6\x0d1\x0d1\x0d255\x0d123"[..]); // CR
        // Spaces tested before
        decode_minimal_image(&b"P61\x09\x0a\x0b\x0c\x0d1 255 123"[..]); // All whitespace, combined
    }

    #[test]
    fn ascii_encoded() {
        decode_minimal_image(&b"P31 1 255 49 50 51"[..]);
        assert!(PPMDecoder::new(&b"P31 1 65535 65535 65535 65535"[..]).unwrap()
            .read_image().is_ok());
    }

    /// Tests for decoding error, assuming `encoded` is ppm encoding for the very simplistic image
    /// containing a single pixel with one byte values (1, 2, 3).
    fn decode_minimal_image(encoded: &[u8]) {
        let content = vec![49 as u8, 50, 51];
        let mut decoder = PPMDecoder::new(encoded).unwrap();

        assert_eq!(decoder.dimensions().unwrap(), (1, 1));
        assert_eq!(decoder.colortype().unwrap(), ColorType::RGB(8));
        assert_eq!(decoder.row_len().unwrap(), 3);
        assert_eq!(decoder.bytewidth(), 1);

        match decoder.read_image().unwrap() {
            DecodingResult::U8(image) => assert_eq!(image, content),
            _ => assert!(false),
        }
    }

    #[test]
    fn wrong_tag() {
        assert!(PPMDecoder::new(&b"P5 1 1 255 1"[..]).is_err());
    }

    #[test]
    fn invalid_characters() {
        assert!(PPMDecoder::new(&b"P6 1chars1 255 1"[..]).is_err()); // No text outside of comments
        assert!(PPMDecoder::new(&b"P6 1\xFF1 255 1"[..]).is_err()); // No invalid ascii chars
        assert!(PPMDecoder::new(&b"P6 0x01 1 255 1"[..]).is_err()); // Numbers only as decimal
    }

    /// These violate the narrow specification of ppm but are commonly supported in other programs.
    /// Fail fast and concise is important here as these might be received as input files.
    #[test]
    fn unsupported_extensions() {
        assert!(PPMDecoder::new(&b"P6 1 1 65536 1"[..]).is_err()); // No bitwidth above 16
    }
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::io::BufReader;
use std::io::Error as IoError;
use std::ascii::AsciiExt;

use color::{ColorType};
use image::{DecodingResult, ImageDecoder, ImageResult, ImageError};
extern crate byteorder;
use self::byteorder::{BigEndian, ByteOrder};

/// PPM decoder
pub struct PPMDecoder<R> {
    reader: BufReader<R>,
    width: u32,
    height: u32,
    maxwhite: u32,
}

impl<R: Read> PPMDecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(read: R) -> ImageResult<PPMDecoder<R>> {
        let mut buf = BufReader::new(read);
        let mut magic: [u8; 2] = [0, 0];
        try!(buf.read_exact(&mut magic[..])); // Skip magic constant
        if magic[0] != b'P' || (magic[1] != b'3' && magic[1] != b'6') {
            return Err(ImageError::FormatError("Expected magic constant for ppm, P3 or P6".to_string()));
        }

        let width = try!(PPMDecoder::read_next_u32(&mut buf));
        let height = try!(PPMDecoder::read_next_u32(&mut buf));
        let maxwhite = try!(PPMDecoder::read_next_u32(&mut buf));

        Ok(PPMDecoder {
            reader: buf,
            width: width,
            height: height,
            maxwhite: maxwhite,
        })
    }

    /// Reads a string as well as a single whitespace after it, ignoring comments
    fn read_next_string(reader: &mut BufReader<R>) -> ImageResult<String> {
        let mut bytes = Vec::new();

        // pair input bytes with a bool mask to remove comments
        let mark_comments = reader
            .bytes()
            .scan(true, |partof, read| {
                let byte = match read {
                    Err(err) => return Some((*partof, Err(err))),
                    Ok(byte) => byte,
                };
                let cur_enabled = *partof && byte != b'#';
                let next_enabled = cur_enabled || (byte == b'\r' || byte == b'\n');
                *partof = next_enabled;
                return Some((cur_enabled, Ok(byte)));
            });

        for (_, byte) in mark_comments.filter(|ref e| e.0) {
            match byte {
                Ok(b'\t') | Ok(b'\n') | Ok(b'\x0b') | Ok(b'\x0c') | Ok(b'\r') | Ok(b' ') => {
                    if !bytes.is_empty() {
                        break // We're done as we already have some content
                    }
                },
                Ok(byte) => {
                    bytes.push(byte);
                },
                Err(_) => break,
            }
        }

        if bytes.is_empty() {
            return Err(ImageError::FormatError("Unexpected eof".to_string()))
        }

        if !bytes.as_slice().is_ascii() {
            return Err(ImageError::FormatError("Non ascii character in preamble".to_string()))
        }

        String::from_utf8(bytes).map_err(|_| ImageError::FormatError("Couldn't read preamble".to_string()))
    }

    fn read_next_u32(reader: &mut BufReader<R>) -> ImageResult<u32> {
        let s = try!(PPMDecoder::read_next_string(reader));
        s.parse::<u32>().map_err(|_| ImageError::FormatError("Invalid number in preamble".to_string()))
    }
}

impl<R: Read> ImageDecoder for PPMDecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        Ok((self.width, self.height))
    }

    fn colortype(&mut self) -> ImageResult<ColorType> {
        match self.bytewidth() {
            1 => Ok(ColorType::RGB(8)),
            2 => Ok(ColorType::RGB(16)),
            _ => Err(ImageError::FormatError("Don't know how to decode PPM with more than 16 bits".to_string())),
        }
    }

    fn row_len(&mut self) -> ImageResult<usize> {
        Ok((self.width*3*self.bytewidth()) as usize)
    }

    fn read_scanline(&mut self, _buf: &mut [u8]) -> ImageResult<u32> {
        unimplemented!();
    }

    fn read_image(&mut self) -> ImageResult<DecodingResult> {
        let opt_size = self.width.checked_mul(self.height)
            .map_or(None, |v| v.checked_mul(3))
            .map_or(None, |v| v.checked_mul(self.bytewidth()));

        let size = match opt_size {
            Some(v) => v,
            None => return Err(ImageError::DimensionError),
        };

        let mut data = vec![0 as u8; size as usize];

        match self.reader.read_exact(&mut data) {
            Ok(_) => {},
            Err(e) => return Err(ImageError::IoError(e)),
        };

        if self.bytewidth() == 1 {
            Ok(DecodingResult::U8(data))
        } else {
            let mut out = vec![0 as u16; (self.width*self.height*3) as usize];
            for (o, i) in out.chunks_mut(1).zip(data.chunks(2)) {
                o[0] = BigEndian::read_u16(i);
            }
            Ok(DecodingResult::U16(out))
        }
    }
}

impl<R: Read> PPMDecoder<R> {
    fn bytewidth(&self) -> u32 {
        if self.maxwhite < 256 { 1 } else { 2 }
    }
}

/// Tests parsing binary buffers were written based on and validated against `pamfile` from
/// netpbm (http://netpbm.sourceforge.net/).
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn minimal_form() {
        // Violates current specification (October 2016 ) but accepted by both netpbm and ImageMagick
        decode_minimal_image(&b"P61 1 255 123"[..]);
        decode_minimal_image(&b"P6 1 1 255 123"[..]);
        decode_minimal_image(&b"P6 1 1 255 123\xFF"[..]); // Too long should not be an issue
    }

    #[test]
    fn comment_in_token() {
        decode_minimal_image(&b"P6 1 1 2#comment\n55 123"[..]); // Terminating LF
        decode_minimal_image(&b"P6 1 1 2#comment\r55 123"[..]); // Terminating CR
        decode_minimal_image(&b"P6 1 1#comment\n 255 123"[..]); // Comment after token
        decode_minimal_image(&b"P6 1 1 #comment\n255 123"[..]); // Comment before token
        decode_minimal_image(&b"P6#comment\n 1 1 255 123"[..]); // Begin of header
        decode_minimal_image(&b"P6 1 1 255#comment\n 123"[..]); // End of header
    }

    #[test]
    fn whitespace() {
        decode_minimal_image(&b"P6\x091\x091\x09255\x09123"[..]); // TAB
        decode_minimal_image(&b"P6\x0a1\x0a1\x0a255\x0a123"[..]); // LF
        decode_minimal_image(&b"P6\x0b1\x0b1\x0b255\x0b123"[..]); // VT
        decode_minimal_image(&b"P6\x0c1\x0c1\x0c255\x0c123"[..]); // FF
        decode_minimal_image(&b"P6\x0d1\x0d1\x0d255\x0d123"[..]); // CR
        // Spaces tested before
        decode_minimal_image(&b"P61\x09\x0a\x0b\x0c\x0d1 255 123"[..]); // All whitespace, combined
    }

    /// Tests for decoding error, assuming `encoded` is ppm encoding for the very simplistic image
    /// containing a single pixel with one byte values (1, 2, 3).
    fn decode_minimal_image(encoded: &[u8]) {
        let content = vec![49 as u8, 50, 51];
        let mut decoder = PPMDecoder::new(encoded).unwrap();

        assert_eq!(decoder.dimensions().unwrap(), (1, 1));
        assert_eq!(decoder.colortype().unwrap(), ColorType::RGB(8));
        assert_eq!(decoder.row_len().unwrap(), 3);
        assert_eq!(decoder.bytewidth(), 1);

        match decoder.read_image().unwrap() {
            DecodingResult::U8(image) => assert_eq!(image, content),
            _ => assert!(false),
        }
    }

    #[test]
    fn wrong_tag() {
        assert!(PPMDecoder::new(&b"P5 1 1 255 1"[..]).is_err());
    }

    #[test]
    fn invalid_characters() {
        assert!(PPMDecoder::new(&b"P6 1chars1 255 1"[..]).is_err()); // No text outside of comments
        assert!(PPMDecoder::new(&b"P6 1\xFF1 255 1"[..]).is_err()); // No invalid ascii chars
        assert!(PPMDecoder::new(&b"P6 0x01 1 255 1"[..]).is_err()); // Numbers only as decimal
    }

    /// These violate the narrow specification of ppm but are commonly supported in other programs.
    /// Fail fast and concise is important here as these might be received as input files.
    #[test]
    fn unsupported_extensions() {
        assert!(PPMDecoder::new(&b"P6 1 1 65536 1"[..]).is_err()); // No bitwidth above 16
    }
}
<commit_msg>Check for violation of maxval<commit_after>use std::io::Read;
use std::io::BufReader;
use std::io::Error as IoError;
use std::ascii::AsciiExt;

use color::{ColorType};
use image::{DecodingResult, ImageDecoder, ImageResult, ImageError};
extern crate byteorder;
use self::byteorder::{BigEndian, ByteOrder};

/// PPM decoder
pub struct PPMDecoder<R> {
    reader: BufReader<R>,
    width: u32,
    height: u32,
    maxwhite: u32,
}

impl<R: Read> PPMDecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(read: R) -> ImageResult<PPMDecoder<R>> {
        let mut buf = BufReader::new(read);
        let mut magic: [u8; 2] = [0, 0];
        try!(buf.read_exact(&mut magic[..])); // Skip magic constant
        if magic[0] != b'P' || (magic[1] != b'3' && magic[1] != b'6') {
            return Err(ImageError::FormatError("Expected magic constant for ppm, P3 or P6".to_string()));
        }

        let width = try!(PPMDecoder::read_next_u32(&mut buf));
        let height = try!(PPMDecoder::read_next_u32(&mut buf));
        let maxwhite = try!(PPMDecoder::read_next_u32(&mut buf));

        if !(maxwhite <= u16::max_value() as u32) {
            return Err(ImageError::FormatError("Image maxval is not less or equal to 65535".to_string()))
        }

        Ok(PPMDecoder {
            reader: buf,
            width: width,
            height: height,
            maxwhite: maxwhite,
        })
    }

    /// Reads a string as well as a single whitespace after it, ignoring comments
    fn read_next_string(reader: &mut BufReader<R>) -> ImageResult<String> {
        let mut bytes = Vec::new();

        // pair input bytes with a bool mask to remove comments
        let mark_comments = reader
            .bytes()
            .scan(true, |partof, read| {
                let byte = match read {
                    Err(err) => return Some((*partof, Err(err))),
                    Ok(byte) => byte,
                };
                let cur_enabled = *partof && byte != b'#';
                let next_enabled = cur_enabled || (byte == b'\r' || byte == b'\n');
                *partof = next_enabled;
                return Some((cur_enabled, Ok(byte)));
            });

        for (_, byte) in mark_comments.filter(|ref e| e.0) {
            match byte {
                Ok(b'\t') | Ok(b'\n') | Ok(b'\x0b') | Ok(b'\x0c') | Ok(b'\r') | Ok(b' ') => {
                    if !bytes.is_empty() {
                        break // We're done as we already have some content
                    }
                },
                Ok(byte) => {
                    bytes.push(byte);
                },
                Err(_) => break,
            }
        }

        if bytes.is_empty() {
            return Err(ImageError::FormatError("Unexpected eof".to_string()))
        }

        if !bytes.as_slice().is_ascii() {
            return Err(ImageError::FormatError("Non ascii character in preamble".to_string()))
        }

        String::from_utf8(bytes).map_err(|_| ImageError::FormatError("Couldn't read preamble".to_string()))
    }

    fn read_next_u32(reader: &mut BufReader<R>) -> ImageResult<u32> {
        let s = try!(PPMDecoder::read_next_string(reader));
        s.parse::<u32>().map_err(|_| ImageError::FormatError("Invalid number in preamble".to_string()))
    }
}

impl<R: Read> ImageDecoder for PPMDecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        Ok((self.width, self.height))
    }

    fn colortype(&mut self) -> ImageResult<ColorType> {
        match self.bytewidth() {
            1 => Ok(ColorType::RGB(8)),
            2 => Ok(ColorType::RGB(16)),
            _ => Err(ImageError::FormatError("Don't know how to decode PPM with more than 16 bits".to_string())),
        }
    }

    fn row_len(&mut self) -> ImageResult<usize> {
        Ok((self.width*3*self.bytewidth()) as usize)
    }

    fn read_scanline(&mut self, _buf: &mut [u8]) -> ImageResult<u32> {
        unimplemented!();
    }

    fn read_image(&mut self) -> ImageResult<DecodingResult> {
        let opt_size = self.width.checked_mul(self.height)
            .map_or(None, |v| v.checked_mul(3))
            .map_or(None, |v| v.checked_mul(self.bytewidth()));

        let size = match opt_size {
            Some(v) => v,
            None => return Err(ImageError::DimensionError),
        };

        let mut data = vec![0 as u8; size as usize];

        match self.reader.read_exact(&mut data) {
            Ok(_) => {},
            Err(e) => return Err(ImageError::IoError(e)),
        };

        if self.bytewidth() == 1 {
            Ok(DecodingResult::U8(data))
        } else {
            let mut out = vec![0 as u16; (self.width*self.height*3) as usize];
            for (o, i) in out.chunks_mut(1).zip(data.chunks(2)) {
                o[0] = BigEndian::read_u16(i);
            }
            Ok(DecodingResult::U16(out))
        }
    }
}

impl<R: Read> PPMDecoder<R> {
    fn bytewidth(&self) -> u32 {
        if self.maxwhite < 256 { 1 } else { 2 }
    }
}

/// Tests parsing binary buffers were written based on and validated against `pamfile` from
/// netpbm (http://netpbm.sourceforge.net/).
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn minimal_form() {
        // Violates current specification (October 2016 ) but accepted by both netpbm and ImageMagick
        decode_minimal_image(&b"P61 1 255 123"[..]);
        decode_minimal_image(&b"P6 1 1 255 123"[..]);
        decode_minimal_image(&b"P6 1 1 255 123\xFF"[..]); // Too long should not be an issue
    }

    #[test]
    fn comment_in_token() {
        decode_minimal_image(&b"P6 1 1 2#comment\n55 123"[..]); // Terminating LF
        decode_minimal_image(&b"P6 1 1 2#comment\r55 123"[..]); // Terminating CR
        decode_minimal_image(&b"P6 1 1#comment\n 255 123"[..]); // Comment after token
        decode_minimal_image(&b"P6 1 1 #comment\n255 123"[..]); // Comment before token
        decode_minimal_image(&b"P6#comment\n 1 1 255 123"[..]); // Begin of header
        decode_minimal_image(&b"P6 1 1 255#comment\n 123"[..]); // End of header
    }

    #[test]
    fn whitespace() {
        decode_minimal_image(&b"P6\x091\x091\x09255\x09123"[..]); // TAB
        decode_minimal_image(&b"P6\x0a1\x0a1\x0a255\x0a123"[..]); // LF
        decode_minimal_image(&b"P6\x0b1\x0b1\x0b255\x0b123"[..]); // VT
        decode_minimal_image(&b"P6\x0c1\x0c1\x0c255\x0c123"[..]); // FF
        decode_minimal_image(&b"P6\x0d1\x0d1\x0d255\x0d123"[..]); // CR
        // Spaces tested before
        decode_minimal_image(&b"P61\x09\x0a\x0b\x0c\x0d1 255 123"[..]); // All whitespace, combined
    }

    /// Tests for decoding error, assuming `encoded` is ppm encoding for the very simplistic image
    /// containing a single pixel with one byte values (1, 2, 3).
    fn decode_minimal_image(encoded: &[u8]) {
        let content = vec![49 as u8, 50, 51];
        let mut decoder = PPMDecoder::new(encoded).unwrap();

        assert_eq!(decoder.dimensions().unwrap(), (1, 1));
        assert_eq!(decoder.colortype().unwrap(), ColorType::RGB(8));
        assert_eq!(decoder.row_len().unwrap(), 3);
        assert_eq!(decoder.bytewidth(), 1);

        match decoder.read_image().unwrap() {
            DecodingResult::U8(image) => assert_eq!(image, content),
            _ => assert!(false),
        }
    }

    #[test]
    fn wrong_tag() {
        assert!(PPMDecoder::new(&b"P5 1 1 255 1"[..]).is_err());
    }

    #[test]
    fn invalid_characters() {
        assert!(PPMDecoder::new(&b"P6 1chars1 255 1"[..]).is_err()); // No text outside of comments
        assert!(PPMDecoder::new(&b"P6 1\xFF1 255 1"[..]).is_err()); // No invalid ascii chars
        assert!(PPMDecoder::new(&b"P6 0x01 1 255 1"[..]).is_err()); // Numbers only as decimal
    }

    /// These violate the narrow specification of ppm but are commonly supported in other programs.
    /// Fail fast and concise is important here as these might be received as input files.
    #[test]
    fn unsupported_extensions() {
        assert!(PPMDecoder::new(&b"P6 1 1 65536 1"[..]).is_err()); // No bitwidth above 16
    }
}
<|endoftext|>


<commit_before>//! Memo runner of printf
//! Takes a format string and arguments
//! 1. tokenizes format string into tokens, consuming
//! any subst. arguments along the way.
//! 2. feeds remaining arguments into function
//! that prints tokens.

use cli;
use itertools::put_back_n;
use std::iter::Peekable;
use std::slice::Iter;
use tokenize::sub::Sub;
use tokenize::token::{Token, Tokenizer};
use tokenize::unescaped_text::UnescapedText;

pub struct Memo {
    tokens: Vec<Box<dyn Token>>,
}

fn warn_excess_args(first_arg: &str) {
    cli::err_msg(&format!(
        "warning: ignoring excess arguments, starting with '{}'",
        first_arg
    ));
}

impl Memo {
    pub fn new(pf_string: &str, pf_args_it: &mut Peekable<Iter<String>>) -> Memo {
        let mut pm = Memo { tokens: Vec::new() };
        let mut tmp_token: Option<Box<dyn Token>>;
        let mut it = put_back_n(pf_string.chars());
        let mut has_sub = false;
        loop {
            tmp_token = UnescapedText::from_it(&mut it, pf_args_it);
            if let Some(x) = tmp_token {
                pm.tokens.push(x);
            }
            tmp_token = Sub::from_it(&mut it, pf_args_it);
            if let Some(x) = tmp_token {
                if !has_sub {
                    has_sub = true;
                }
                pm.tokens.push(x);
            }
            if let Some(x) = it.next() {
                it.put_back(x);
            } else {
                break;
            }
        }
        if !has_sub {
            let mut drain = false;
            if let Some(first_arg) = pf_args_it.peek() {
                warn_excess_args(first_arg);
                drain = true;
            }
            if drain {
                loop {
                    // drain remaining args;
                    if pf_args_it.next().is_none() {
                        break;
                    }
                }
            }
        }
        pm
    }
    pub fn apply(&self, pf_args_it: &mut Peekable<Iter<String>>) {
        for tkn in self.tokens.iter() {
            tkn.print(pf_args_it);
        }
    }
    pub fn run_all(pf_string: &str, pf_args: &[String]) {
        let mut arg_it = pf_args.iter().peekable();
        let pm = Memo::new(&pf_string.to_string(), &mut arg_it);
        loop {
            if arg_it.peek().is_none() {
                break;
            }
            pm.apply(&mut arg_it);
        }
    }
}
<commit_msg>printf: refactor ~ remove redundant string allocation<commit_after>//! Memo runner of printf
//! Takes a format string and arguments
//! 1. tokenizes format string into tokens, consuming
//! any subst. arguments along the way.
//! 2. feeds remaining arguments into function
//! that prints tokens.

use cli;
use itertools::put_back_n;
use std::iter::Peekable;
use std::slice::Iter;
use tokenize::sub::Sub;
use tokenize::token::{Token, Tokenizer};
use tokenize::unescaped_text::UnescapedText;

pub struct Memo {
    tokens: Vec<Box<dyn Token>>,
}

fn warn_excess_args(first_arg: &str) {
    cli::err_msg(&format!(
        "warning: ignoring excess arguments, starting with '{}'",
        first_arg
    ));
}

impl Memo {
    pub fn new(pf_string: &str, pf_args_it: &mut Peekable<Iter<String>>) -> Memo {
        let mut pm = Memo { tokens: Vec::new() };
        let mut tmp_token: Option<Box<dyn Token>>;
        let mut it = put_back_n(pf_string.chars());
        let mut has_sub = false;
        loop {
            tmp_token = UnescapedText::from_it(&mut it, pf_args_it);
            if let Some(x) = tmp_token {
                pm.tokens.push(x);
            }
            tmp_token = Sub::from_it(&mut it, pf_args_it);
            if let Some(x) = tmp_token {
                if !has_sub {
                    has_sub = true;
                }
                pm.tokens.push(x);
            }
            if let Some(x) = it.next() {
                it.put_back(x);
            } else {
                break;
            }
        }
        if !has_sub {
            let mut drain = false;
            if let Some(first_arg) = pf_args_it.peek() {
                warn_excess_args(first_arg);
                drain = true;
            }
            if drain {
                loop {
                    // drain remaining args;
                    if pf_args_it.next().is_none() {
                        break;
                    }
                }
            }
        }
        pm
    }
    pub fn apply(&self, pf_args_it: &mut Peekable<Iter<String>>) {
        for tkn in self.tokens.iter() {
            tkn.print(pf_args_it);
        }
    }
    pub fn run_all(pf_string: &str, pf_args: &[String]) {
        let mut arg_it = pf_args.iter().peekable();
        let pm = Memo::new(pf_string, &mut arg_it);
        loop {
            if arg_it.peek().is_none() {
                break;
            }
            pm.apply(&mut arg_it);
        }
    }
}
<|endoftext|>


<commit_before>use hardware::cartridge::Cartridge;
use hardware::handler_holder::GBHandlerHolder;
use hardware::apu::NoisePattern;
use hardware::cpu::Cpu;

const VOLUME_MAX: i16 = 32000;
pub const AUDIO_BUFFER_SIZE: usize = 1470;

pub struct Emulator {
    pub cpu: Cpu,
    noise_7_bit: [u8; 127],
    noise_15_bit: [u8; 32767],
    phase: Phase,
    frequency: f64,
}

impl Emulator {
    pub fn from_data(data: &[u8], frequency: f64) -> Result<Emulator, String> {
        let handler = GBHandlerHolder::new(Cartridge::from_data(data));

        Ok(Emulator {
            cpu: Cpu::new(Box::new(handler)),
            frequency: frequency,
            noise_7_bit: generate_noise_7_bit(),
            noise_15_bit: generate_noise_15_bit(),
            phase: Phase {
                channel_1: 0.0,
                channel_2: 0.0,
                channel_3: 0.0,
                channel_4: 0.0,
            },
        })
    }

    pub fn reset(&mut self) {
        self.cpu.reset();
    }

    pub fn generate_sound(&mut self) -> [i16; AUDIO_BUFFER_SIZE] {
        let audio = self.cpu.handler_holder.get_audio_buffer();

        let mut i = 0;

        let mut channel_1_phase = self.phase.channel_1;
        let channel_1_frequency = 131072 / (2048 - audio.sound_1().frequency());
        let channel_1_phase_inc = channel_1_frequency as f64 / self.frequency;
        let channel_1_volume = (VOLUME_MAX as f64 * audio.sound_1().volume() as f64
                                / 16.0 / 4.0) as i16;

        let mut channel_2_phase = self.phase.channel_2;
        let channel_2_frequency = 131072 / (2048 - audio.sound_2().frequency());
        let channel_2_phase_inc = channel_2_frequency as f64 / self.frequency;
        let channel_2_volume = (VOLUME_MAX as f64 * audio.sound_2().volume() as f64
                                / 16.0 / 4.0) as i16;

        let mut channel_3_phase = self.phase.channel_3;
        let mut channel_3_pattern = [0; 32];
        for i in 0..16 {
            channel_3_pattern[i * 2]     = (audio.sound_3().wave_pattern()[i] & 0xF0) >> 4;
            channel_3_pattern[i * 2 + 1] =  audio.sound_3().wave_pattern()[i] & 0x0F;
        }

        let channel_3_frequency = 65536 / (2048 - audio.sound_3().frequency());
        let channel_3_volume = (VOLUME_MAX as f64 * audio.sound_3().volume().to_volume() as f64
                                / 4.0) as i16;
        let channel_3_phase_inc = channel_3_frequency as f64 / self.frequency;

        let channel_4_pattern = match audio.sound_4().pattern() {
            NoisePattern::C15 => &self.noise_15_bit[..],
            NoisePattern::C7  => &self.noise_7_bit[..],
        };
        let mut channel_4_phase = self.phase.channel_4 % channel_4_pattern.len() as f64;
        let channel_4_phase_inc = audio.sound_4().frequency() as f64 / self.frequency;
        let channel_4_volume = (VOLUME_MAX as f64 * audio.sound_4().volume() as f64
                                / 16.0 / 4.0) as i16;

        let mut out = [0; AUDIO_BUFFER_SIZE];

        while i < out.len() {
            let channel_1 = if channel_1_phase < audio.sound_1().wave_duty() as f64 {
                channel_1_volume
            } else {
                -channel_1_volume
            };

            let channel_2 = if channel_2_phase < audio.sound_2().wave_duty() as f64 {
                channel_2_volume
            } else {
                -channel_2_volume
            };

            let channel_3_index = (channel_3_phase * 32.0) as usize;
            let channel_3 = ((channel_3_pattern[channel_3_index] as f32 - 7.0)
                            / 8.0 * channel_3_volume as f32) as i16;

            let channel_4_index = channel_4_phase as usize;
            let channel_4_bit = channel_4_pattern[channel_4_index];
            let channel_4 = if channel_4_bit == 1 {
                channel_4_volume
            } else {
                -channel_4_volume
            };

            out[i] =
                  ((if audio.sound_1().playing_left() { channel_1 } else { 0 }) +
                   (if audio.sound_2().playing_left() { channel_2 } else { 0 }) +
                   (if audio.sound_3().playing_left() { channel_3 } else { 0 }) +
                   (if audio.sound_4().playing_left() { channel_4 } else { 0 })) / 4;

            out[i + 1] =
               ((if audio.sound_1().playing_right() { channel_1 } else { 0 }) +
                (if audio.sound_2().playing_right() { channel_2 } else { 0 }) +
                (if audio.sound_3().playing_right() { channel_3 } else { 0 }) +
                (if audio.sound_4().playing_right() { channel_4 } else { 0 })) / 4;

            channel_1_phase = (channel_1_phase + channel_1_phase_inc) % 1.0;
            channel_2_phase = (channel_2_phase + channel_2_phase_inc) % 1.0;
            channel_3_phase = (channel_3_phase + channel_3_phase_inc) % 1.0;
            channel_4_phase = (channel_4_phase + channel_4_phase_inc)
                            % channel_4_pattern.len() as f64;

            i += 2;
        }

        self.phase = Phase {
            channel_1: channel_1_phase,
            channel_2: channel_2_phase,
            channel_3: channel_3_phase,
            channel_4: channel_4_phase,
        };

        out
    }
}

#[derive(Clone, Copy)]
pub struct Phase {
    channel_1: f64,
    channel_2: f64,
    channel_3: f64,
    channel_4: f64,
}

fn sound_rng(v: u16, generator: u16) -> u16 {
    let bit = v & 1;
    let next = v >> 1;
    if bit != 0 {
        next ^ generator
    } else {
        next
    }
}

fn generate_noise_array(out: &mut [u8], generator: u16, bits: u32) {
    assert!(bits < 16);

    let mut v = 2u16.pow(bits) - 1;

    for i in 0..out.len() {
        out[i] = (v & 1) as u8;
        v = sound_rng(v, generator);
    }

    // Let's verify that we have reached the end of the cycle
    assert_eq!(2u16.pow(bits) - 1, v);
}

fn generate_noise_7_bit() -> [u8; 127] {
    let mut pattern = [0; 127];
    generate_noise_array(&mut pattern, 0b1000001, 7);
    pattern
}

fn generate_noise_15_bit() -> [u8; 32767] {
    let mut pattern = [0; 32767];
    generate_noise_array(&mut pattern, 0b100000000000001, 15);
    pattern
}
<commit_msg>Refactor channel_3 sound handling.<commit_after>use hardware::cartridge::Cartridge;
use hardware::handler_holder::GBHandlerHolder;
use hardware::apu::NoisePattern;
use hardware::cpu::Cpu;

const VOLUME_MAX: i16 = 32000;
pub const AUDIO_BUFFER_SIZE: usize = 1470;

pub struct Emulator {
    pub cpu: Cpu,
    noise_7_bit: [u8; 127],
    noise_15_bit: [u8; 32767],
    phase: Phase,
    frequency: f64,
}

impl Emulator {
    pub fn from_data(data: &[u8], frequency: f64) -> Result<Emulator, String> {
        let handler = GBHandlerHolder::new(Cartridge::from_data(data));

        Ok(Emulator {
            cpu: Cpu::new(Box::new(handler)),
            frequency: frequency,
            noise_7_bit: generate_noise_7_bit(),
            noise_15_bit: generate_noise_15_bit(),
            phase: Phase {
                channel_1: 0.0,
                channel_2: 0.0,
                channel_3: 0.0,
                channel_4: 0.0,
            },
        })
    }

    pub fn reset(&mut self) {
        self.cpu.reset();
    }

    pub fn generate_sound(&mut self) -> [i16; AUDIO_BUFFER_SIZE] {
        let audio = self.cpu.handler_holder.get_audio_buffer();

        let mut i = 0;

        let mut channel_1_phase = self.phase.channel_1;
        let channel_1_frequency = 131072 / (2048 - audio.sound_1().frequency());
        let channel_1_phase_inc = channel_1_frequency as f64 / self.frequency;
        let channel_1_volume = (VOLUME_MAX as f64 * audio.sound_1().volume() as f64
                                / 16.0 / 4.0) as i16;

        let mut channel_2_phase = self.phase.channel_2;
        let channel_2_frequency = 131072 / (2048 - audio.sound_2().frequency());
        let channel_2_phase_inc = channel_2_frequency as f64 / self.frequency;
        let channel_2_volume = (VOLUME_MAX as f64 * audio.sound_2().volume() as f64
                                / 16.0 / 4.0) as i16;

        let mut channel_3_phase = self.phase.channel_3;
        let mut channel_3_pattern = [0.0; 32];
        for i in 0..16 {
            channel_3_pattern[i * 2]     = ((audio.sound_3().wave_pattern()[i] & 0xF0) >> 4) as f32;
            channel_3_pattern[i * 2 + 1] = (audio.sound_3().wave_pattern()[i] & 0x0F) as f32;
        }

        let channel_3_frequency = 65536 / (2048 - audio.sound_3().frequency());
        let channel_3_volume = VOLUME_MAX as f32
            * audio.sound_3().volume().to_volume() as f32 / 4.0;
        let channel_3_phase_inc = channel_3_frequency as f64 / self.frequency;

        let channel_4_pattern = match audio.sound_4().pattern() {
            NoisePattern::C15 => &self.noise_15_bit[..],
            NoisePattern::C7  => &self.noise_7_bit[..],
        };
        let mut channel_4_phase = self.phase.channel_4 % channel_4_pattern.len() as f64;
        let channel_4_phase_inc = audio.sound_4().frequency() as f64 / self.frequency;
        let channel_4_volume = (VOLUME_MAX as f64 * audio.sound_4().volume() as f64
                                / 16.0 / 4.0) as i16;

        let mut out = [0; AUDIO_BUFFER_SIZE];

        while i < out.len() {
            let channel_1 = if channel_1_phase < audio.sound_1().wave_duty() as f64 {
                channel_1_volume
            } else {
                -channel_1_volume
            };

            let channel_2 = if channel_2_phase < audio.sound_2().wave_duty() as f64 {
                channel_2_volume
            } else {
                -channel_2_volume
            };

            let channel_3_index = (channel_3_phase * 32.0) as usize;
            let channel_3 = ((channel_3_pattern[channel_3_index] - 7.0) / 8.0 * channel_3_volume) as i16;

            let channel_4_index = channel_4_phase as usize;
            let channel_4_bit = channel_4_pattern[channel_4_index];
            let channel_4 = if channel_4_bit == 1 {
                channel_4_volume
            } else {
                -channel_4_volume
            };

            out[i] =
                  ((if audio.sound_1().playing_left() { channel_1 } else { 0 }) +
                   (if audio.sound_2().playing_left() { channel_2 } else { 0 }) +
                   (if audio.sound_3().playing_left() { channel_3 } else { 0 }) +
                   (if audio.sound_4().playing_left() { channel_4 } else { 0 })) / 4;

            out[i + 1] =
               ((if audio.sound_1().playing_right() { channel_1 } else { 0 }) +
                (if audio.sound_2().playing_right() { channel_2 } else { 0 }) +
                (if audio.sound_3().playing_right() { channel_3 } else { 0 }) +
                (if audio.sound_4().playing_right() { channel_4 } else { 0 })) / 4;

            channel_1_phase = (channel_1_phase + channel_1_phase_inc) % 1.0;
            channel_2_phase = (channel_2_phase + channel_2_phase_inc) % 1.0;
            channel_3_phase = (channel_3_phase + channel_3_phase_inc) % 1.0;
            channel_4_phase = (channel_4_phase + channel_4_phase_inc)
                            % channel_4_pattern.len() as f64;

            i += 2;
        }

        self.phase = Phase {
            channel_1: channel_1_phase,
            channel_2: channel_2_phase,
            channel_3: channel_3_phase,
            channel_4: channel_4_phase,
        };

        out
    }
}

#[derive(Clone, Copy)]
pub struct Phase {
    channel_1: f64,
    channel_2: f64,
    channel_3: f64,
    channel_4: f64,
}

fn sound_rng(v: u16, generator: u16) -> u16 {
    let bit = v & 1;
    let next = v >> 1;
    if bit != 0 {
        next ^ generator
    } else {
        next
    }
}

fn generate_noise_array(out: &mut [u8], generator: u16, bits: u32) {
    assert!(bits < 16);

    let mut v = 2u16.pow(bits) - 1;

    for i in 0..out.len() {
        out[i] = (v & 1) as u8;
        v = sound_rng(v, generator);
    }

    // Let's verify that we have reached the end of the cycle
    assert_eq!(2u16.pow(bits) - 1, v);
}

fn generate_noise_7_bit() -> [u8; 127] {
    let mut pattern = [0; 127];
    generate_noise_array(&mut pattern, 0b1000001, 7);
    pattern
}

fn generate_noise_15_bit() -> [u8; 32767] {
    let mut pattern = [0; 32767];
    generate_noise_array(&mut pattern, 0b100000000000001, 15);
    pattern
}
<|endoftext|>