<commit_before>// Copyright (C) 2020, Cloudflare, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use super::common::alpns;

pub trait Args {
    fn with_docopt(docopt: &docopt::Docopt) -> Self;
}

/// Contains commons arguments for creating a quiche QUIC connection.
pub struct CommonArgs {
    pub alpns: Vec<u8>,
    pub max_data: u64,
    pub max_stream_data: u64,
    pub max_streams_bidi: u64,
    pub max_streams_uni: u64,
    pub idle_timeout: u64,
    pub early_data: bool,
    pub dump_packet_path: Option<String>,
    pub no_grease: bool,
    pub cc_algorithm: String,
    pub disable_hystart: bool,
    pub dgrams_enabled: bool,
    pub dgram_count: u64,
    pub dgram_data: String,
}

/// Creates a new `CommonArgs` structure using the provided [`Docopt`].
///
/// The `Docopt` usage String needs to include the following:
///
/// --http-version VERSION      HTTP version to use.
/// --max-data BYTES            Connection-wide flow control limit.
/// --max-stream-data BYTES     Per-stream flow control limit.
/// --max-streams-bidi STREAMS  Number of allowed concurrent streams.
/// --max-streams-uni STREAMS   Number of allowed concurrent streams.
/// --dump-packets PATH         Dump the incoming packets in PATH.
/// --no-grease                 Don't send GREASE.
/// --cc-algorithm NAME         Set a congestion control algorithm.
/// --disable-hystart           Disable HyStart++.
/// --dgram-proto PROTO         DATAGRAM application protocol.
/// --dgram-count COUNT         Number of DATAGRAMs to send.
///  --dgram-data DATA          DATAGRAM data to send.
///
/// [`Docopt`]: https://docs.rs/docopt/1.1.0/docopt/
impl Args for CommonArgs {
    fn with_docopt(docopt: &docopt::Docopt) -> Self {
        let args = docopt.parse().unwrap_or_else(|e| e.exit());

        let http_version = args.get_str("--http-version");
        let dgram_proto = args.get_str("--dgram-proto");
        let (alpns, dgrams_enabled) = match (http_version, dgram_proto) {
            ("HTTP/0.9", "none") =>
                (alpns::length_prefixed(&alpns::HTTP_09), false),

            ("HTTP/0.9", _) =>
                panic!("Unsupported HTTP version and DATAGRAM protocol."),

            ("HTTP/3", "none") => (alpns::length_prefixed(&alpns::HTTP_3), false),

            ("HTTP/3", "oneway") =>
                (alpns::length_prefixed(&alpns::HTTP_3), true),

            ("all", "none") => (
                [
                    alpns::length_prefixed(&alpns::HTTP_3),
                    alpns::length_prefixed(&alpns::HTTP_09),
                ]
                .concat(),
                false,
            ),

            // SiDuck is it's own application protocol.
            (_, "siduck") => (alpns::length_prefixed(&alpns::SIDUCK), true),

            (..) => panic!("Unsupported HTTP version and DATAGRAM protocol."),
        };

        let dgram_count = args.get_str("--dgram-count");
        let dgram_count = dgram_count.parse::<u64>().unwrap();

        let dgram_data = args.get_str("--dgram-data").to_string();

        let max_data = args.get_str("--max-data");
        let max_data = max_data.parse::<u64>().unwrap();

        let max_stream_data = args.get_str("--max-stream-data");
        let max_stream_data = max_stream_data.parse::<u64>().unwrap();

        let max_streams_bidi = args.get_str("--max-streams-bidi");
        let max_streams_bidi = max_streams_bidi.parse::<u64>().unwrap();

        let max_streams_uni = args.get_str("--max-streams-uni");
        let max_streams_uni = max_streams_uni.parse::<u64>().unwrap();

        let idle_timeout = args.get_str("--idle-timeout");
        let idle_timeout = idle_timeout.parse::<u64>().unwrap();

        let early_data = args.get_bool("--early-data");

        let dump_packet_path = if args.get_str("--dump-packets") != "" {
            Some(args.get_str("--dump-packets").to_string())
        } else {
            None
        };

        let no_grease = args.get_bool("--no-grease");

        let cc_algorithm = args.get_str("--cc-algorithm");

        let disable_hystart = args.get_bool("--disable-hystart");

        CommonArgs {
            alpns,
            max_data,
            max_stream_data,
            max_streams_bidi,
            max_streams_uni,
            idle_timeout,
            early_data,
            dump_packet_path,
            no_grease,
            cc_algorithm: cc_algorithm.to_string(),
            disable_hystart,
            dgrams_enabled,
            dgram_count,
            dgram_data,
        }
    }
}

impl Default for CommonArgs {
    fn default() -> Self {
        CommonArgs {
            alpns: alpns::length_prefixed(&alpns::HTTP_3),
            max_data: 10000000,
            max_stream_data: 1000000,
            max_streams_bidi: 100,
            max_streams_uni: 100,
            idle_timeout: 30000,
            early_data: false,
            dump_packet_path: None,
            no_grease: false,
            cc_algorithm: "cubic".to_string(),
            disable_hystart: false,
            dgrams_enabled: false,
            dgram_count: 0,
            dgram_data: "quack".to_string(),
        }
    }
}

pub const CLIENT_USAGE: &str = "Usage:
  quiche-client [options] URL...
  quiche-client -h | --help

Options:
  --method METHOD          Use the given HTTP request method [default: GET].
  --body FILE              Send the given file as request body.
  --max-data BYTES         Connection-wide flow control limit [default: 10000000].
  --max-stream-data BYTES  Per-stream flow control limit [default: 1000000].
  --max-streams-bidi STREAMS  Number of allowed concurrent streams [default: 100].
  --max-streams-uni STREAMS   Number of allowed concurrent streams [default: 100].
  --idle-timeout TIMEOUT   Idle timeout in milliseconds [default: 30000].
  --wire-version VERSION   The version number to send to the server [default: babababa].
  --http-version VERSION   HTTP version to use [default: all].
  --early-data             Enable sending early data.
  --dgram-proto PROTO      DATAGRAM application protocol to use [default: none].
  --dgram-count COUNT      Number of DATAGRAMs to send [default: 0].
  --dgram-data DATA        Data to send for certain types of DATAGRAM application protocol [default: quack].
  --dump-packets PATH      Dump the incoming packets as files in the given directory.
  --dump-responses PATH    Dump response payload as files in the given directory.
  --dump-json              Dump response headers and payload to stdout in JSON format.
  --max-json-payload BYTES  Per-response payload limit when dumping JSON [default: 10000].
  --connect-to ADDRESS     Override ther server's address.
  --no-verify              Don't verify server's certificate.
  --no-grease              Don't send GREASE.
  --cc-algorithm NAME      Specify which congestion control algorithm to use [default: cubic].
  --disable-hystart        Disable HyStart++.
  -H --header HEADER ...   Add a request header.
  -n --requests REQUESTS   Send the given number of identical requests [default: 1].
  --session-file PATH      File used to cache a TLS session for resumption.
  -h --help                Show this screen.
";

/// Application-specific arguments that compliment the `CommonArgs`.
pub struct ClientArgs {
    pub version: u32,
    pub dump_response_path: Option<String>,
    pub dump_json: Option<usize>,
    pub urls: Vec<url::Url>,
    pub reqs_cardinal: u64,
    pub req_headers: Vec<String>,
    pub no_verify: bool,
    pub body: Option<Vec<u8>>,
    pub method: String,
    pub connect_to: Option<String>,
    pub session_file: Option<String>,
}

impl Args for ClientArgs {
    fn with_docopt(docopt: &docopt::Docopt) -> Self {
        let args = docopt.parse().unwrap_or_else(|e| e.exit());

        let version = args.get_str("--wire-version");
        let version = u32::from_str_radix(version, 16).unwrap();

        let dump_response_path = if args.get_str("--dump-responses") != "" {
            Some(args.get_str("--dump-responses").to_string())
        } else {
            None
        };

        let dump_json = args.get_bool("--dump-json");
        let dump_json = if dump_json {
            let max_payload = args.get_str("--max-json-payload");
            let max_payload = max_payload.parse::<usize>().unwrap();
            Some(max_payload)
        } else {
            None
        };

        // URLs (can be multiple).
        let urls: Vec<url::Url> = args
            .get_vec("URL")
            .into_iter()
            .map(|x| url::Url::parse(x).unwrap())
            .collect();

        // Request headers (can be multiple).
        let req_headers = args
            .get_vec("--header")
            .into_iter()
            .map(|x| x.to_string())
            .collect();

        let reqs_cardinal = args.get_str("--requests");
        let reqs_cardinal = reqs_cardinal.parse::<u64>().unwrap();

        let no_verify = args.get_bool("--no-verify");

        let body = if args.get_bool("--body") {
            std::fs::read(args.get_str("--body")).ok()
        } else {
            None
        };

        let method = args.get_str("--method").to_string();

        let connect_to = if args.get_bool("--connect-to") {
            Some(args.get_str("--connect-to").to_string())
        } else {
            None
        };

        let session_file = if args.get_bool("--session-file") {
            Some(args.get_str("--session-file").to_string())
        } else {
            None
        };

        ClientArgs {
            version,
            dump_response_path,
            dump_json,
            urls,
            req_headers,
            reqs_cardinal,
            no_verify,
            body,
            method,
            connect_to,
            session_file,
        }
    }
}

impl Default for ClientArgs {
    fn default() -> Self {
        ClientArgs {
            version: 0xbabababa,
            dump_response_path: None,
            dump_json: None,
            urls: vec![],
            req_headers: vec![],
            reqs_cardinal: 1,
            no_verify: false,
            body: None,
            method: "GET".to_string(),
            connect_to: None,
            session_file: None,
        }
    }
}

pub const SERVER_USAGE: &str = "Usage:
  quiche-server [options]
  quiche-server -h | --help

Options:
  --listen <addr>             Listen on the given IP:port [default: 127.0.0.1:4433]
  --cert <file>               TLS certificate path [default: src/bin/cert.crt]
  --key <file>                TLS certificate key path [default: src/bin/cert.key]
  --root <dir>                Root directory [default: src/bin/root/]
  --index <name>              The file that will be used as index [default: index.html].
  --name <str>                Name of the server [default: quic.tech]
  --max-data BYTES            Connection-wide flow control limit [default: 10000000].
  --max-stream-data BYTES     Per-stream flow control limit [default: 1000000].
  --max-streams-bidi STREAMS  Number of allowed concurrent streams [default: 100].
  --max-streams-uni STREAMS   Number of allowed concurrent streams [default: 100].
  --idle-timeout TIMEOUT   Idle timeout in milliseconds [default: 30000].
  --dump-packets PATH         Dump the incoming packets as files in the given directory.
  --early-data                Enable receiving early data.
  --no-retry                  Disable stateless retry.
  --no-grease                 Don't send GREASE.
  --http-version VERSION      HTTP version to use [default: all].
  --dgram-proto PROTO         DATAGRAM application protocol to use [default: none].
  --dgram-count COUNT         Number of DATAGRAMs to send [default: 0].
  --dgram-data DATA           Data to send for certain types of DATAGRAM application protocol [default: brrr].
  --cc-algorithm NAME         Specify which congestion control algorithm to use [default: cubic].
  --disable-hystart           Disable HyStart++.
  -h --help                   Show this screen.
";

// Application-specific arguments that compliment the `CommonArgs`.
pub struct ServerArgs {
    pub listen: String,
    pub no_retry: bool,
    pub root: String,
    pub index: String,
    pub cert: String,
    pub key: String,
}

impl Args for ServerArgs {
    fn with_docopt(docopt: &docopt::Docopt) -> Self {
        let args = docopt.parse().unwrap_or_else(|e| e.exit());

        let listen = args.get_str("--listen").to_string();
        let no_retry = args.get_bool("--no-retry");
        let root = args.get_str("--root").to_string();
        let index = args.get_str("--index").to_string();
        let cert = args.get_str("--cert").to_string();
        let key = args.get_str("--key").to_string();

        ServerArgs {
            listen,
            no_retry,
            root,
            index,
            cert,
            key,
        }
    }
}
<commit_msg>apps: fix inconsistent struct constructor<commit_after>// Copyright (C) 2020, Cloudflare, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use super::common::alpns;

pub trait Args {
    fn with_docopt(docopt: &docopt::Docopt) -> Self;
}

/// Contains commons arguments for creating a quiche QUIC connection.
pub struct CommonArgs {
    pub alpns: Vec<u8>,
    pub max_data: u64,
    pub max_stream_data: u64,
    pub max_streams_bidi: u64,
    pub max_streams_uni: u64,
    pub idle_timeout: u64,
    pub early_data: bool,
    pub dump_packet_path: Option<String>,
    pub no_grease: bool,
    pub cc_algorithm: String,
    pub disable_hystart: bool,
    pub dgrams_enabled: bool,
    pub dgram_count: u64,
    pub dgram_data: String,
}

/// Creates a new `CommonArgs` structure using the provided [`Docopt`].
///
/// The `Docopt` usage String needs to include the following:
///
/// --http-version VERSION      HTTP version to use.
/// --max-data BYTES            Connection-wide flow control limit.
/// --max-stream-data BYTES     Per-stream flow control limit.
/// --max-streams-bidi STREAMS  Number of allowed concurrent streams.
/// --max-streams-uni STREAMS   Number of allowed concurrent streams.
/// --dump-packets PATH         Dump the incoming packets in PATH.
/// --no-grease                 Don't send GREASE.
/// --cc-algorithm NAME         Set a congestion control algorithm.
/// --disable-hystart           Disable HyStart++.
/// --dgram-proto PROTO         DATAGRAM application protocol.
/// --dgram-count COUNT         Number of DATAGRAMs to send.
///  --dgram-data DATA          DATAGRAM data to send.
///
/// [`Docopt`]: https://docs.rs/docopt/1.1.0/docopt/
impl Args for CommonArgs {
    fn with_docopt(docopt: &docopt::Docopt) -> Self {
        let args = docopt.parse().unwrap_or_else(|e| e.exit());

        let http_version = args.get_str("--http-version");
        let dgram_proto = args.get_str("--dgram-proto");
        let (alpns, dgrams_enabled) = match (http_version, dgram_proto) {
            ("HTTP/0.9", "none") =>
                (alpns::length_prefixed(&alpns::HTTP_09), false),

            ("HTTP/0.9", _) =>
                panic!("Unsupported HTTP version and DATAGRAM protocol."),

            ("HTTP/3", "none") => (alpns::length_prefixed(&alpns::HTTP_3), false),

            ("HTTP/3", "oneway") =>
                (alpns::length_prefixed(&alpns::HTTP_3), true),

            ("all", "none") => (
                [
                    alpns::length_prefixed(&alpns::HTTP_3),
                    alpns::length_prefixed(&alpns::HTTP_09),
                ]
                .concat(),
                false,
            ),

            // SiDuck is it's own application protocol.
            (_, "siduck") => (alpns::length_prefixed(&alpns::SIDUCK), true),

            (..) => panic!("Unsupported HTTP version and DATAGRAM protocol."),
        };

        let dgram_count = args.get_str("--dgram-count");
        let dgram_count = dgram_count.parse::<u64>().unwrap();

        let dgram_data = args.get_str("--dgram-data").to_string();

        let max_data = args.get_str("--max-data");
        let max_data = max_data.parse::<u64>().unwrap();

        let max_stream_data = args.get_str("--max-stream-data");
        let max_stream_data = max_stream_data.parse::<u64>().unwrap();

        let max_streams_bidi = args.get_str("--max-streams-bidi");
        let max_streams_bidi = max_streams_bidi.parse::<u64>().unwrap();

        let max_streams_uni = args.get_str("--max-streams-uni");
        let max_streams_uni = max_streams_uni.parse::<u64>().unwrap();

        let idle_timeout = args.get_str("--idle-timeout");
        let idle_timeout = idle_timeout.parse::<u64>().unwrap();

        let early_data = args.get_bool("--early-data");

        let dump_packet_path = if args.get_str("--dump-packets") != "" {
            Some(args.get_str("--dump-packets").to_string())
        } else {
            None
        };

        let no_grease = args.get_bool("--no-grease");

        let cc_algorithm = args.get_str("--cc-algorithm");

        let disable_hystart = args.get_bool("--disable-hystart");

        CommonArgs {
            alpns,
            max_data,
            max_stream_data,
            max_streams_bidi,
            max_streams_uni,
            idle_timeout,
            early_data,
            dump_packet_path,
            no_grease,
            cc_algorithm: cc_algorithm.to_string(),
            disable_hystart,
            dgrams_enabled,
            dgram_count,
            dgram_data,
        }
    }
}

impl Default for CommonArgs {
    fn default() -> Self {
        CommonArgs {
            alpns: alpns::length_prefixed(&alpns::HTTP_3),
            max_data: 10000000,
            max_stream_data: 1000000,
            max_streams_bidi: 100,
            max_streams_uni: 100,
            idle_timeout: 30000,
            early_data: false,
            dump_packet_path: None,
            no_grease: false,
            cc_algorithm: "cubic".to_string(),
            disable_hystart: false,
            dgrams_enabled: false,
            dgram_count: 0,
            dgram_data: "quack".to_string(),
        }
    }
}

pub const CLIENT_USAGE: &str = "Usage:
  quiche-client [options] URL...
  quiche-client -h | --help

Options:
  --method METHOD          Use the given HTTP request method [default: GET].
  --body FILE              Send the given file as request body.
  --max-data BYTES         Connection-wide flow control limit [default: 10000000].
  --max-stream-data BYTES  Per-stream flow control limit [default: 1000000].
  --max-streams-bidi STREAMS  Number of allowed concurrent streams [default: 100].
  --max-streams-uni STREAMS   Number of allowed concurrent streams [default: 100].
  --idle-timeout TIMEOUT   Idle timeout in milliseconds [default: 30000].
  --wire-version VERSION   The version number to send to the server [default: babababa].
  --http-version VERSION   HTTP version to use [default: all].
  --early-data             Enable sending early data.
  --dgram-proto PROTO      DATAGRAM application protocol to use [default: none].
  --dgram-count COUNT      Number of DATAGRAMs to send [default: 0].
  --dgram-data DATA        Data to send for certain types of DATAGRAM application protocol [default: quack].
  --dump-packets PATH      Dump the incoming packets as files in the given directory.
  --dump-responses PATH    Dump response payload as files in the given directory.
  --dump-json              Dump response headers and payload to stdout in JSON format.
  --max-json-payload BYTES  Per-response payload limit when dumping JSON [default: 10000].
  --connect-to ADDRESS     Override ther server's address.
  --no-verify              Don't verify server's certificate.
  --no-grease              Don't send GREASE.
  --cc-algorithm NAME      Specify which congestion control algorithm to use [default: cubic].
  --disable-hystart        Disable HyStart++.
  -H --header HEADER ...   Add a request header.
  -n --requests REQUESTS   Send the given number of identical requests [default: 1].
  --session-file PATH      File used to cache a TLS session for resumption.
  -h --help                Show this screen.
";

/// Application-specific arguments that compliment the `CommonArgs`.
pub struct ClientArgs {
    pub version: u32,
    pub dump_response_path: Option<String>,
    pub dump_json: Option<usize>,
    pub urls: Vec<url::Url>,
    pub reqs_cardinal: u64,
    pub req_headers: Vec<String>,
    pub no_verify: bool,
    pub body: Option<Vec<u8>>,
    pub method: String,
    pub connect_to: Option<String>,
    pub session_file: Option<String>,
}

impl Args for ClientArgs {
    fn with_docopt(docopt: &docopt::Docopt) -> Self {
        let args = docopt.parse().unwrap_or_else(|e| e.exit());

        let version = args.get_str("--wire-version");
        let version = u32::from_str_radix(version, 16).unwrap();

        let dump_response_path = if args.get_str("--dump-responses") != "" {
            Some(args.get_str("--dump-responses").to_string())
        } else {
            None
        };

        let dump_json = args.get_bool("--dump-json");
        let dump_json = if dump_json {
            let max_payload = args.get_str("--max-json-payload");
            let max_payload = max_payload.parse::<usize>().unwrap();
            Some(max_payload)
        } else {
            None
        };

        // URLs (can be multiple).
        let urls: Vec<url::Url> = args
            .get_vec("URL")
            .into_iter()
            .map(|x| url::Url::parse(x).unwrap())
            .collect();

        // Request headers (can be multiple).
        let req_headers = args
            .get_vec("--header")
            .into_iter()
            .map(|x| x.to_string())
            .collect();

        let reqs_cardinal = args.get_str("--requests");
        let reqs_cardinal = reqs_cardinal.parse::<u64>().unwrap();

        let no_verify = args.get_bool("--no-verify");

        let body = if args.get_bool("--body") {
            std::fs::read(args.get_str("--body")).ok()
        } else {
            None
        };

        let method = args.get_str("--method").to_string();

        let connect_to = if args.get_bool("--connect-to") {
            Some(args.get_str("--connect-to").to_string())
        } else {
            None
        };

        let session_file = if args.get_bool("--session-file") {
            Some(args.get_str("--session-file").to_string())
        } else {
            None
        };

        ClientArgs {
            version,
            dump_response_path,
            dump_json,
            urls,
            reqs_cardinal,
            req_headers,
            no_verify,
            body,
            method,
            connect_to,
            session_file,
        }
    }
}

impl Default for ClientArgs {
    fn default() -> Self {
        ClientArgs {
            version: 0xbabababa,
            dump_response_path: None,
            dump_json: None,
            urls: vec![],
            req_headers: vec![],
            reqs_cardinal: 1,
            no_verify: false,
            body: None,
            method: "GET".to_string(),
            connect_to: None,
            session_file: None,
        }
    }
}

pub const SERVER_USAGE: &str = "Usage:
  quiche-server [options]
  quiche-server -h | --help

Options:
  --listen <addr>             Listen on the given IP:port [default: 127.0.0.1:4433]
  --cert <file>               TLS certificate path [default: src/bin/cert.crt]
  --key <file>                TLS certificate key path [default: src/bin/cert.key]
  --root <dir>                Root directory [default: src/bin/root/]
  --index <name>              The file that will be used as index [default: index.html].
  --name <str>                Name of the server [default: quic.tech]
  --max-data BYTES            Connection-wide flow control limit [default: 10000000].
  --max-stream-data BYTES     Per-stream flow control limit [default: 1000000].
  --max-streams-bidi STREAMS  Number of allowed concurrent streams [default: 100].
  --max-streams-uni STREAMS   Number of allowed concurrent streams [default: 100].
  --idle-timeout TIMEOUT   Idle timeout in milliseconds [default: 30000].
  --dump-packets PATH         Dump the incoming packets as files in the given directory.
  --early-data                Enable receiving early data.
  --no-retry                  Disable stateless retry.
  --no-grease                 Don't send GREASE.
  --http-version VERSION      HTTP version to use [default: all].
  --dgram-proto PROTO         DATAGRAM application protocol to use [default: none].
  --dgram-count COUNT         Number of DATAGRAMs to send [default: 0].
  --dgram-data DATA           Data to send for certain types of DATAGRAM application protocol [default: brrr].
  --cc-algorithm NAME         Specify which congestion control algorithm to use [default: cubic].
  --disable-hystart           Disable HyStart++.
  -h --help                   Show this screen.
";

// Application-specific arguments that compliment the `CommonArgs`.
pub struct ServerArgs {
    pub listen: String,
    pub no_retry: bool,
    pub root: String,
    pub index: String,
    pub cert: String,
    pub key: String,
}

impl Args for ServerArgs {
    fn with_docopt(docopt: &docopt::Docopt) -> Self {
        let args = docopt.parse().unwrap_or_else(|e| e.exit());

        let listen = args.get_str("--listen").to_string();
        let no_retry = args.get_bool("--no-retry");
        let root = args.get_str("--root").to_string();
        let index = args.get_str("--index").to_string();
        let cert = args.get_str("--cert").to_string();
        let key = args.get_str("--key").to_string();

        ServerArgs {
            listen,
            no_retry,
            root,
            index,
            cert,
            key,
        }
    }
}
<|endoftext|>


<commit_before>use atom::*;
use env::Env;
use errors::*;
use symbol::*;
use funcs::*;
use eval::*;
use self::Instruction::*;

#[derive (Debug, Clone)]
struct Frame {
    program: CompiledFunction,
    pc: usize
}

impl Frame {
    fn new(f: CompiledFunction) -> Frame {
        Frame {
            program: f,
            pc: 0
        }
    }

    fn current_instruction(&self) -> Option<&Instruction> {
        self.program.body.get(self.pc)
    }
}

#[derive (Debug)]
pub struct VirtualMachine {
    stack: Vec<Atom>,
    frames: Vec<Frame>,
    fp: usize,
    sp: usize,
}

#[derive (Debug, Clone, PartialEq)]
#[allow(dead_code)]
pub enum Instruction {
    CONST(Atom), // pushes the atom on the stack
    LOAD(InternedStr), // loads then pushes the value
    DEFINE(InternedStr), // sets the symbol to value on the top of the stack
    POP, // pops one off the stack
    // FUNCTION,
    JUMP,
    RETURN, // pops the frame
    CALL(Function, usize), // calls a noncompiled or native function
    DCALL(usize) // calls the function at the top of the stack
}

fn compile_node(node: Atom, out: &mut Vec<Instruction>, env: &Env) -> Result<(), Error> {
    match node {
        Atom::Symbol(sym) => {
            out.push(Instruction::LOAD(sym))
        },
        Atom::List(ref list) => {
            for n in list.iter() {
                try!(compile(n.clone(), out, env));
            }
        },
        // handle constants
        _ => {
            out.push(Instruction::CONST(node.clone()));
        }
    }
    Ok(())
}

pub fn compile(node: Atom, out: &mut Vec<Instruction>, env: &Env) -> Result<(), Error> {
    match node {
        Atom::List(ref list) => {
            if list.is_empty() {
                out.push(Instruction::CONST(node.clone()));
                return Ok(())
            }
            ()
        },
        _ => return compile_node(node, out, env)
    }

    let list = try!(node.as_list());

    if let Atom::List(_) = list[0] {
        for i in list.iter().skip(1) {
            try!(compile(i.clone(), out, env));
        }
        try!(compile(list[0].clone(), out, env));
        out.push(DCALL(list.len() - 1));
        return Ok(())
    }

    match list[0] {
        Atom::Form(f) => {
            match f {
                Form::Let => {
                    let mut new_env = Env::new(Some(env.clone()));
                    let bind_list = try!(list[1].as_list());
                    let mut bindings = Vec::new();

                    for binding in bind_list.iter() {
                        let bind_exp = try!(binding.as_list());
                        try!(expect_arg_length(bind_exp, 2));
                        bindings.push(bind_exp[0].clone());
                        try!(compile(bind_exp[1].clone(), out, env));
                    }
                    let mut body = Vec::new();
                    try!(compile(list[2].clone(), &mut body, &mut new_env));
                    body.push(RETURN);

                    let arity = bindings.len();

                    let func = CompiledFunction {
                        body: body,
                        params: to_list(bindings),
                        env: new_env
                    };

                    out.push(CONST(Atom::Function(Function::Compiled(func))));
                    out.push(DCALL(arity));
                    return Ok(())
                },
                Form::Set => {
                },
                Form::Def => {
                    let sym = try!(list[1].as_symbol());
                    try!(compile(list[2].clone(), out, env));
                    out.push(DEFINE(*sym));
                    return Ok(())
                },
                Form::Do => {
                    if list.len() > 2 {
                        for i in 1..list.len()-1 {
                            try!(compile(list[i].clone(), out, env));
                            out.push(POP);
                        }
                    }
                    return compile(list[list.len()-1].clone(), out, env)
                },
                Form::Fn => {
                    let mut body = Vec::new();
                    let mut env = Env::new(Some(env.clone()));
                    try!(compile(list[2].clone(), &mut body, &mut env));
                    body.push(RETURN);
                    let func = CompiledFunction {
                        body: body,
                        params: try!(list[1].as_list()).clone(),
                        env: env
                    };
                    out.push(CONST(Atom::Function(Function::Compiled(func))));
                    return Ok(())
                }
                _ => {
                    return Err(Error::NotImplemented)
                }
            }
        },
        Atom::Symbol(_) => {
            for n in list.iter().skip(1) {
                try!(compile(n.clone(), out, env));
            }
            try!(compile_node(list[0].clone(), out, env));
            out.push( DCALL(list.len() -1 ));
            return Ok(())
        }
        Atom::Function(ref func) => {
            for n in list.iter().skip(1) {
                try!(compile(n.clone(), out, env));
            }
            match *func {
                Function::Native(_) => {
                    out.push(Instruction::CALL(func.clone(), list.len() -1));
                    return Ok(())
                },
                Function::Compiled(_) => {
                    out.push(CALL(func.clone(), list.len() -1));
                    return Ok(())
                },
                _ => {
                    return Err(Error::NotImplemented)
                }
            }
        },
        _ => {
            return Err(Error::NotAFunction)
        }
    }

    Err(Error::InvalidArguments)
}

fn eval_native_borrow(v: &mut VirtualMachine, n: Native, len: usize) -> AtomResult {
    let &mut VirtualMachine { ref mut stack, ref mut frames, fp, .. } = v;
    eval_native(n, &stack[len..], &mut frames[fp].program.env)
}

impl VirtualMachine {
    pub fn new() -> VirtualMachine {
        VirtualMachine {
            stack: vec![],
            frames: vec![],
            sp: 0,
            fp: 0
        }
    }

    fn next_instruction(&self) -> Option<Instruction> {
        self.frames[self.fp].current_instruction().map(|n| n.clone())
    }

    fn current_frame(&mut self) -> &mut Frame {
        &mut self.frames[self.fp]
    }

    fn current_env(&mut self) -> &mut Env {
        &mut self.current_frame().program.env
    }

    pub fn run(&mut self) -> AtomResult {
        while let Some(instruction) = self.next_instruction() {
            match instruction {
                POP => {
                    self.pop();
                }
                RETURN => {
                    if self.fp > 0 {
                        self.fp -= 1;
                        self.frames.pop();
                    } else {
                        return Ok(self.pop())
                    }
                },
                LOAD(sym) => {
                    let value = self.current_frame().program.env.get(sym.as_ref()).unwrap_or(Atom::Nil);
                    self.push(value);
                },
                CONST(ref atom) => {
                    self.push(atom.clone())
                },
                DEFINE(sym) => {
                    let v = self.pop();
                    let a = try!(self.current_env().define(sym, v));
                    self.push(a);
                },
                DCALL(arity) => {
                    let func = self.pop();

                    match func {
                        Atom::Function(Function::Compiled(mut f)) => {
                            f.env.bind_mut(&f.params, &self.stack[self.sp-arity..]);
                            for _ in 0..arity {
                                self.pop();
                            }
                            self.frames.push(Frame::new(f.clone()));
                            self.fp += 1;
                            continue; // don't advance PC of the new frame
                        },
                        _ => {
                            println!("not a function: {:?}", func);
                            return Err(Error::NotAFunction)
                        }
                    }
                },
                CALL(func, arity) => {
                    match func {
                        Function::Native(n) => {
                            let len = self.stack.len();
                            let r = try!(eval_native_borrow(self, n, len - arity));
                            for _ in 1..arity {
                                self.pop();
                            }
                            self.push(r);
                        },
                        _ => {
                            return Err(Error::NotImplemented)
                        }
                    }
                },
                _ => ()
            }
            self.current_frame().pc += 1;
        }
        Ok(self.pop())
    }

    fn push(&mut self, a: Atom)  {
        self.stack.push(a);
        self.sp += 1
    }

    fn pop(&mut self) -> Atom {
        self.sp -= 1;
        self.stack.pop().unwrap()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use parser::*;
    use atom::*;
    use env::*;
    use super::Frame;

    fn empty_frame() -> Frame {
        Frame::new(CompiledFunction{
            body: Vec::new(),
            params: empty_list(),
            env: Env::new(None)
        })
    }

    fn run_expr(s: &str) -> Atom {
        let p = tokenize(s).unwrap();
        let mut out = Vec::new();
        let env = Env::new(None);
        compile(p, &mut out, &env).unwrap();
        println!("output: {:?}", out);
        let mut vm = VirtualMachine::new();
        let mut f = empty_frame();
        f.program.env = env;
        f.program.body = out;
        vm.frames.push(f);
        vm.run().unwrap()
    }

    #[test]
    fn test() {
        assert_eq!(Atom::from(6), run_expr("(+ (+ 2 2) 2)"));
        assert_eq!(Atom::symbol("a"), run_expr("(def a 1)"));
        assert_eq!(Atom::from(3), run_expr("(do (def foo (fn (x) (+ x 2))) (foo 1))"));
    }

    #[test]
    fn test_invoke_fn_inline() {
        assert_eq!(Atom::from(1), run_expr("((fn (x) 1) 0)"));
    }

    #[test]
    fn test_closure() {
        assert_eq!(Atom::from(2), run_expr("(do (def foo (fn (x) (fn () (* 2 x)))) (def foo2 (foo 1)) (foo2))"))
    }

    #[test]
    fn let_binding() {
        assert_eq!(Atom::from(-1), run_expr("(let* ((x 2) (y 3)) (- x y))"));
    }

    #[test]
    fn let_doesnt_leak() {
        assert_eq!(Atom::from(0), run_expr("(do (let* ((x 1)) (let* ((x 0)) x)))"))
    }
}
<commit_msg>rough take on if<commit_after>use atom::*;
use env::Env;
use errors::*;
use symbol::*;
use funcs::*;
use eval::*;
use self::Instruction::*;

#[derive (Debug, Clone)]
struct Frame {
    program: CompiledFunction,
    pc: usize
}

impl Frame {
    fn new(f: CompiledFunction) -> Frame {
        Frame {
            program: f,
            pc: 0
        }
    }

    fn current_instruction(&self) -> Option<&Instruction> {
        self.program.body.get(self.pc)
    }
}

#[derive (Debug)]
pub struct VirtualMachine {
    stack: Vec<Atom>,
    frames: Vec<Frame>,
    fp: usize,
    sp: usize,
}

#[derive (Debug, Clone, PartialEq)]
#[allow(non_camel_case_types)]
pub enum Instruction {
    CONST(Atom), // pushes the atom on the stack
    LOAD(InternedStr), // loads then pushes the value
    DEFINE(InternedStr), // sets the symbol to value on the top of the stack
    POP, // pops one off the stack
    // FUNCTION,
    JUMP_IF(usize),
    JUMP(usize),
    RETURN, // pops the frame
    CALL(Function, usize), // calls a noncompiled or native function
    DCALL(usize) // calls the function at the top of the stack
}

fn compile_node(node: Atom, out: &mut Vec<Instruction>, env: &Env) -> Result<(), Error> {
    match node {
        Atom::Symbol(sym) => {
            out.push(Instruction::LOAD(sym))
        },
        Atom::List(ref list) => {
            for n in list.iter() {
                try!(compile(n.clone(), out, env));
            }
        },
        // handle constants
        _ => {
            out.push(Instruction::CONST(node.clone()));
        }
    }
    Ok(())
}

pub fn compile(node: Atom, out: &mut Vec<Instruction>, env: &Env) -> Result<(), Error> {
    match node {
        Atom::List(ref list) => {
            if list.is_empty() {
                out.push(Instruction::CONST(node.clone()));
                return Ok(())
            }
            ()
        },
        _ => return compile_node(node, out, env)
    }

    let list = try!(node.as_list());

    if let Atom::List(_) = list[0] {
        for i in list.iter().skip(1) {
            try!(compile(i.clone(), out, env));
        }
        try!(compile(list[0].clone(), out, env));
        out.push(DCALL(list.len() - 1));
        return Ok(())
    }

    match list[0] {
        Atom::Form(f) => {
            match f {
                Form::If => {
                    if list.len() < 2 {
                        return Err(Error::NotEnoughArguments)
                    }
                    try!(compile(list[1].clone(), out, env));
                    // predicate on stack
                    out.push(JUMP_IF(0));
                    let else_jump_pos = out.len() -1;

                    try!(compile(list[2].clone(), out, env));
                    out.push(JUMP(0));
                    let out_jump = out.len() - 1;
                    out[else_jump_pos] = JUMP_IF(out.len());

                    try!(compile(list[3].clone(), out , env));

                    out[out_jump] = JUMP(out.len()-1);
                    return Ok(())
                },
                Form::Let => {
                    let mut new_env = Env::new(Some(env.clone()));
                    let bind_list = try!(list[1].as_list());
                    let mut bindings = Vec::new();

                    for binding in bind_list.iter() {
                        let bind_exp = try!(binding.as_list());
                        try!(expect_arg_length(bind_exp, 2));
                        bindings.push(bind_exp[0].clone());
                        try!(compile(bind_exp[1].clone(), out, env));
                    }
                    let mut body = Vec::new();
                    try!(compile(list[2].clone(), &mut body, &mut new_env));
                    body.push(RETURN);

                    let arity = bindings.len();

                    let func = CompiledFunction {
                        body: body,
                        params: to_list(bindings),
                        env: new_env
                    };

                    out.push(CONST(Atom::Function(Function::Compiled(func))));
                    out.push(DCALL(arity));
                    return Ok(())
                },
                Form::Set => {
                },
                Form::Def => {
                    let sym = try!(list[1].as_symbol());
                    try!(compile(list[2].clone(), out, env));
                    out.push(DEFINE(*sym));
                    return Ok(())
                },
                Form::Do => {
                    if list.len() > 2 {
                        for i in 1..list.len()-1 {
                            try!(compile(list[i].clone(), out, env));
                            out.push(POP);
                        }
                    }
                    return compile(list[list.len()-1].clone(), out, env)
                },
                Form::Fn => {
                    let mut body = Vec::new();
                    let mut env = Env::new(Some(env.clone()));
                    try!(compile(list[2].clone(), &mut body, &mut env));
                    body.push(RETURN);
                    let func = CompiledFunction {
                        body: body,
                        params: try!(list[1].as_list()).clone(),
                        env: env
                    };
                    out.push(CONST(Atom::Function(Function::Compiled(func))));
                    return Ok(())
                }
                _ => {
                    return Err(Error::NotImplemented)
                }
            }
        },
        Atom::Symbol(_) => {
            for n in list.iter().skip(1) {
                try!(compile(n.clone(), out, env));
            }
            try!(compile_node(list[0].clone(), out, env));
            out.push( DCALL(list.len() -1 ));
            return Ok(())
        }
        Atom::Function(ref func) => {
            for n in list.iter().skip(1) {
                try!(compile(n.clone(), out, env));
            }
            match *func {
                Function::Native(_) => {
                    out.push(Instruction::CALL(func.clone(), list.len() -1));
                    return Ok(())
                },
                Function::Compiled(_) => {
                    out.push(CALL(func.clone(), list.len() -1));
                    return Ok(())
                },
                _ => {
                    return Err(Error::NotImplemented)
                }
            }
        },
        _ => {
            return Err(Error::NotAFunction)
        }
    }

    Err(Error::InvalidArguments)
}

fn eval_native_borrow(v: &mut VirtualMachine, n: Native, len: usize) -> AtomResult {
    let &mut VirtualMachine { ref mut stack, ref mut frames, fp, .. } = v;
    eval_native(n, &stack[len..], &mut frames[fp].program.env)
}

impl VirtualMachine {
    pub fn new() -> VirtualMachine {
        VirtualMachine {
            stack: vec![],
            frames: vec![],
            sp: 0,
            fp: 0
        }
    }

    fn next_instruction(&self) -> Option<Instruction> {
        self.frames[self.fp].current_instruction().map(|n| n.clone())
    }

    fn current_frame(&mut self) -> &mut Frame {
        &mut self.frames[self.fp]
    }

    fn current_env(&mut self) -> &mut Env {
        &mut self.current_frame().program.env
    }

    pub fn run(&mut self) -> AtomResult {
        while let Some(instruction) = self.next_instruction() {
            match instruction {
                JUMP_IF(addr) => {
                    let pred = match self.pop() {
                        Atom::Boolean(b) => b,
                        Atom::Nil => false,
                        Atom::List(ref l) if l.is_empty() => false,
                        _ => true
                    };
                    if pred {
                        self.current_frame().pc = addr;
                        continue;
                    }
                },
                JUMP(addr) => {
                    self.current_frame().pc = addr;
                    continue;
                }
                POP => {
                    self.pop();
                }
                RETURN => {
                    if self.fp > 0 {
                        self.fp -= 1;
                        self.frames.pop();
                    } else {
                        return Ok(self.pop())
                    }
                },
                LOAD(sym) => {
                    let value = self.current_frame().program.env.get(sym.as_ref()).unwrap_or(Atom::Nil);
                    self.push(value);
                },
                CONST(ref atom) => {
                    self.push(atom.clone())
                },
                DEFINE(sym) => {
                    let v = self.pop();
                    let a = try!(self.current_env().define(sym, v));
                    self.push(a);
                },
                DCALL(arity) => {
                    let func = self.pop();

                    match func {
                        Atom::Function(Function::Compiled(mut f)) => {
                            f.env.bind_mut(&f.params, &self.stack[self.sp-arity..]);
                            for _ in 0..arity {
                                self.pop();
                            }
                            self.frames.push(Frame::new(f.clone()));
                            self.fp += 1;
                            continue; // don't advance PC of the new frame
                        },
                        _ => {
                            println!("not a function: {:?}", func);
                            return Err(Error::NotAFunction)
                        }
                    }
                },
                CALL(func, arity) => {
                    match func {
                        Function::Native(n) => {
                            let len = self.stack.len();
                            let r = try!(eval_native_borrow(self, n, len - arity));
                            for _ in 1..arity {
                                self.pop();
                            }
                            self.push(r);
                        },
                        _ => {
                            return Err(Error::NotImplemented)
                        }
                    }
                },
            }
            self.current_frame().pc += 1;
        }
        Ok(self.pop())
    }

    fn push(&mut self, a: Atom)  {
        self.stack.push(a);
        self.sp += 1
    }

    fn pop(&mut self) -> Atom {
        self.sp -= 1;
        self.stack.pop().unwrap()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use parser::*;
    use atom::*;
    use env::*;
    use super::Frame;

    fn empty_frame() -> Frame {
        Frame::new(CompiledFunction{
            body: Vec::new(),
            params: empty_list(),
            env: Env::new(None)
        })
    }

    fn run_expr(s: &str) -> Atom {
        let p = tokenize(s).unwrap();
        let mut out = Vec::new();
        let env = Env::new(None);
        compile(p, &mut out, &env).unwrap();
        println!("output: {:?}", out);
        let mut vm = VirtualMachine::new();
        let mut f = empty_frame();
        f.program.env = env;
        f.program.body = out;
        vm.frames.push(f);
        vm.run().unwrap()
    }

    #[test]
    fn test() {
        assert_eq!(Atom::from(6), run_expr("(+ (+ 2 2) 2)"));
        assert_eq!(Atom::symbol("a"), run_expr("(def a 1)"));
        assert_eq!(Atom::from(3), run_expr("(do (def foo (fn (x) (+ x 2))) (foo 1))"));
    }

    #[test]
    fn test_invoke_fn_inline() {
        assert_eq!(Atom::from(1), run_expr("((fn (x) 1) 0)"));
    }

    #[test]
    fn test_closure() {
        assert_eq!(Atom::from(2), run_expr("(do (def foo (fn (x) (fn () (* 2 x)))) (def foo2 (foo 1)) (foo2))"))
    }

    #[test]
    fn let_binding() {
        assert_eq!(Atom::from(-1), run_expr("(let* ((x 2) (y 3)) (- x y))"));
    }

    #[test]
    fn let_doesnt_leak() {
        assert_eq!(Atom::from(0), run_expr("(do (let* ((x 1)) (let* ((x 0)) x)))"))
    }
}
<|endoftext|>


<commit_before>#![allow(unused_imports, unused_must_use, unused_variables, dead_code)]

extern crate combine;

use std::io::prelude::*;
use std::fs::File;
use std::path::Path;
use std::str;
use std::marker::PhantomData;
use combine::{spaces, parser, between, many, many1, digit, char, any, string, token, Parser, ParserExt, ParseError};
use combine::primitives::{State, Stream, ParseResult, Consumed};
use combine::combinator::{Between, Token, FnParser, satisfy};

//my own bencode stuff!
//
pub enum Bencode {
    Number(i64),
    String(Vec<u8>)
}

fn open_file <P: AsRef<Path>>(path: P) -> Vec<u8> {
    let mut fd = File::open(path).unwrap();
    let mut buffer:Vec<u8> = Vec::new();
    fd.read_to_end(&mut buffer);
    buffer
}

enum Expr {
    Integer(i64)
}

//TODO: this does not handle negatives, yo
fn bencode_integer<I>(input: State<I>) -> ParseResult<i64, I> where I:Stream<Item=char> {
    let (open, close) = (char('i'), char('e'));
    let mut int = between(open, close, many1::<String, _>(digit())).map(|x| {
        x.parse::<i64>().unwrap()
    });
    int.parse_state(input)
}

fn bencode_string<I>(input: State<I>) -> ParseResult<String, I> where I:Stream<Item=char> {
    let (len, input_) = try!(bencode_string_length_prefix(input));
    input_.combine(|input__| take(len).parse_state(input__))
}

fn bencode_string_length_prefix<I>(input: State<I>) -> ParseResult<i32, I> where I:Stream<Item=char> {
    let many_digit = many1::<String, _>(digit());
    let mut get_len = (many_digit, char(':')).map(|(length, _)| length.parse::<i32>().unwrap());
    get_len.parse_state(input)
}

fn take <I> (num: i32) -> SizedBuffer<I> where I: Stream<Item=char> {
    SizedBuffer(num, PhantomData)
}

#[derive(Clone)]
pub struct SizedBuffer <I>(i32, PhantomData<I>);
impl <I> Parser for SizedBuffer<I> where I: Stream<Item=char> {
    type Input = I;
    type Output = String;
    fn parse_lazy(&mut self, mut input: State<I>) -> ParseResult<String, I> {
        let start = input.position;
        let mut vec_buf:Vec<char> = Vec::new();
        for _ in 0..self.0 {
            match input.uncons() {
                Ok((other, rest)) => {
                    vec_buf.push(other);
                    input = rest.into_inner();
                }
                Err(error) => {
                    return error.combine(|mut error| {
                        error.position = start;
                        Err(Consumed::Consumed(error))
                    })
                }
            };
        }
        Ok((vec_buf.into_iter().collect(), Consumed::Consumed(input)))
    }
}

fn main () {
    let file_contents = open_file("t-test.torrent");

    match str::from_utf8(&file_contents) {
        Ok(v) => println!("file: {}", v),
        Err(e) => panic!("Not a UTF-8 String: {}", e)
    }

    //let result = parser(bencode_integer).parse("i5e");
    //let result = parser(bencode_string).parse("3:sss");
    let result = parser(bencode_string).parse("5:abcde");

    match result {
        Ok((a, _)) => println!("found {}", a),
        Err(e) => println!("{}", e)
    }
}
<commit_msg>adding 64 bit integer lists, and adding tests<commit_after>#![allow(unused_imports, unused_must_use, unused_variables, dead_code)]

extern crate combine;

use std::io::prelude::*;
use std::fs::File;
use std::path::Path;
use std::str;
use std::marker::PhantomData;
use combine::{spaces, parser, between, many, many1, digit, char, any, string, token, Parser, ParserExt, ParseError};
use combine::primitives::{State, Stream, ParseResult, Consumed};
use combine::combinator::{Between, Token, FnParser, satisfy};

//my own bencode stuff!
//
pub enum Bencode {
    Number(i64),
    String(Vec<u8>)
}

fn open_file <P: AsRef<Path>>(path: P) -> Vec<u8> {
    let mut fd = File::open(path).unwrap();
    let mut buffer:Vec<u8> = Vec::new();
    fd.read_to_end(&mut buffer);
    buffer
}

enum Expr {
    Integer(i64),
    String(String)
}

//TODO: this does not handle negatives, yo
fn bencode_integer<I>(input: State<I>) -> ParseResult<i64, I> where I:Stream<Item=char> {
    let (open, close) = (char('i'), char('e'));
    let mut int = between(open, close, many1::<String, _>(digit())).map(|x| {
        x.parse::<i64>().unwrap()
    });
    int.parse_state(input)
}

fn bencode_string<I>(input: State<I>) -> ParseResult<String, I> where I:Stream<Item=char> {
    let (len, input_) = try!(bencode_string_length_prefix(input));
    input_.combine(|input__| take(len).parse_state(input__))
}

fn bencode_string_length_prefix<I>(input: State<I>) -> ParseResult<i32, I> where I:Stream<Item=char> {
    let many_digit = many1::<String, _>(digit());
    let mut get_len = (many_digit, char(':')).map(|(length, _)| length.parse::<i32>().unwrap());
    get_len.parse_state(input)
}

//for now only handle i64s
fn bencode_list<I>(input: State<I>) -> ParseResult<Vec<i64>, I> where I:Stream<Item=char> {
    let (open, close) = (char('l'), char('e'));
    let list_contents = many(parser(bencode_integer));
    let mut list = between(open, close, list_contents).map(|x| {
        x
    });
    list.parse_state(input)
}

fn take <I> (num: i32) -> SizedBuffer<I> where I: Stream<Item=char> {
    SizedBuffer(num, PhantomData)
}

#[derive(Clone)]
pub struct SizedBuffer <I>(i32, PhantomData<I>);
impl <I> Parser for SizedBuffer<I> where I: Stream<Item=char> {
    type Input = I;
    type Output = String;
    fn parse_lazy(&mut self, mut input: State<I>) -> ParseResult<String, I> {
        let start = input.position;
        let mut vec_buf:Vec<char> = Vec::new();
        for _ in 0..self.0 {
            match input.uncons() {
                Ok((other, rest)) => {
                    vec_buf.push(other);
                    input = rest.into_inner();
                }
                Err(error) => {
                    return error.combine(|mut error| {
                        error.position = start;
                        Err(Consumed::Consumed(error))
                    })
                }
            };
        }
        Ok((vec_buf.into_iter().collect(), Consumed::Consumed(input)))
    }
}

fn main () {
    let file_contents = open_file("t-test.torrent");

    match str::from_utf8(&file_contents) {
        Ok(v) => println!("file: {}", v),
        Err(e) => panic!("Not a UTF-8 String: {}", e)
    }
}

#[test]
fn test_integer() {
    let result = parser(bencode_integer).parse("i57e");
    assert_eq!(result, Ok((57,"")));
}

#[test]
fn test_string() {
    let result = parser(bencode_string).parse("5:abcde");
    assert_eq!(result, Ok(("abcde".to_string(), "")));;
}
<|endoftext|>


<commit_before>//! Compact binary representations of nucleic acid kmers
pub type BitKmerSeq = u64;
pub type BitKmer = (BitKmerSeq, u8);

/// Takes a BitKmer and adds a new base on the end, optionally loping off the
/// first base if the resulting kmer is too long.
fn extend_kmer(kmer: &mut BitKmer, new_char: u8) -> bool {
    let new_char_int;
    match new_char {
        b'A' | b'a' => new_char_int = 0 as BitKmerSeq,
        b'C' | b'c' => new_char_int = 1 as BitKmerSeq,
        b'G' | b'g' => new_char_int = 2 as BitKmerSeq,
        b'T' | b't' => new_char_int = 3 as BitKmerSeq,
        _ => return false,
    };
    let new_kmer = (kmer.0 << 2) + new_char_int;

    // mask out any overflowed bits
    kmer.0 = new_kmer & (BitKmerSeq::pow(2, u32::from(2 * kmer.1)) - 1) as BitKmerSeq;
    true
}

/// Used for the BitNuclKmer iterator to handle skipping invalid bases.
fn update_position(
    start_pos: &mut usize,
    kmer: &mut BitKmer,
    buffer: &[u8],
    initial: bool,
) -> bool {
    // check if we have enough "physical" space for one more kmer
    if *start_pos + kmer.1 as usize > buffer.len() {
        return false;
    }

    let (mut kmer_len, stop_len) = if initial {
        (0, (kmer.1 - 1) as usize)
    } else {
        ((kmer.1 - 1) as usize, kmer.1 as usize)
    };

    let mut cur_kmer = kmer;
    while kmer_len < stop_len {
        if extend_kmer(&mut cur_kmer, buffer[*start_pos + kmer_len]) {
            kmer_len += 1;
        } else {
            kmer_len = 0;
            *cur_kmer = (0u64, cur_kmer.1);
            *start_pos += kmer_len + 1;
            if *start_pos + cur_kmer.1 as usize > buffer.len() {
                return false;
            }
        }
    }
    true
}

pub struct BitNuclKmer<'a> {
    start_pos: usize,
    cur_kmer: BitKmer,
    buffer: &'a [u8],
    canonical: bool,
}

impl<'a> BitNuclKmer<'a> {
    pub fn new(slice: &'a [u8], k: u8, canonical: bool) -> BitNuclKmer<'a> {
        let mut kmer = (0u64, k);
        let mut start_pos = 0;
        update_position(&mut start_pos, &mut kmer, &slice, true);

        BitNuclKmer {
            start_pos,
            cur_kmer: kmer,
            buffer: slice,
            canonical,
        }
    }
}

impl<'a> Iterator for BitNuclKmer<'a> {
    type Item = (usize, BitKmer, bool);

    fn next(&mut self) -> Option<(usize, BitKmer, bool)> {
        if !update_position(&mut self.start_pos, &mut self.cur_kmer, &self.buffer, false) {
            return None;
        }
        self.start_pos += 1;
        if self.canonical {
            let (kmer, was_rc) = canonical(self.cur_kmer);
            Some((self.start_pos - 1, kmer, was_rc))
        } else {
            Some((self.start_pos - 1, self.cur_kmer, false))
        }
    }
}

/// Reverse complement a BitKmer (reverses the sequence and swaps A<>T and G<>C)
pub fn reverse_complement(kmer: BitKmer) -> BitKmer {
    // FIXME: this is not going to work with BitKmers of u128 or u32
    // inspired from https://www.biostars.org/p/113640/
    let mut new_kmer = kmer.0;
    // reverse it
    new_kmer = (new_kmer >> 2 & 0x3333_3333_3333_3333) | (new_kmer & 0x3333_3333_3333_3333) << 2;
    new_kmer = (new_kmer >> 4 & 0x0F0F_0F0F_0F0F_0F0F) | (new_kmer & 0x0F0F_0F0F_0F0F_0F0F) << 4;
    new_kmer = (new_kmer >> 8 & 0x00FF_00FF_00FF_00FF) | (new_kmer & 0x00FF_00FF_00FF_00FF) << 8;
    new_kmer = (new_kmer >> 16 & 0x0000_FFFF_0000_FFFF) | (new_kmer & 0x0000_FFFF_0000_FFFF) << 16;
    new_kmer = (new_kmer >> 32 & 0x0000_0000_FFFF_FFFF) | (new_kmer & 0x0000_0000_FFFF_FFFF) << 32;
    // complement it
    new_kmer ^= 0xFFFF_FFFF_FFFF_FFFF;
    // shift it to the right size
    new_kmer >>= 2 * (32 - kmer.1);
    (new_kmer, kmer.1)
}

/// Return the lexigraphically lowest of the BitKmer and its reverse complement and
/// whether the returned kmer is the reverse_complement (true) or the original (false)
pub fn canonical(kmer: BitKmer) -> (BitKmer, bool) {
    let rc = reverse_complement(kmer);
    if kmer.0 > rc.0 {
        (rc, true)
    } else {
        (kmer, false)
    }
}

/// Find the lexigraphically lowest substring of a given length in the BitKmer
pub fn minimizer(kmer: BitKmer, minmer_size: u8) -> BitKmer {
    let mut new_kmer = kmer.0;
    let mut lowest = !(0 as BitKmerSeq);
    let bitmask = (BitKmerSeq::pow(2, u32::from(2 * minmer_size)) - 1) as BitKmerSeq;
    for _ in 0..=(kmer.1 - minmer_size) {
        let cur = bitmask & new_kmer;
        if cur < lowest {
            lowest = cur;
        }
        let cur_rev = reverse_complement((bitmask & new_kmer, kmer.1));
        if cur_rev.0 < lowest {
            lowest = cur_rev.0;
        }
        new_kmer >>= 2;
    }
    (lowest, kmer.1)
}

pub fn bitmer_to_bytes(kmer: BitKmer) -> Vec<u8> {
    let mut new_kmer = kmer.0;
    let mut new_kmer_str = Vec::new();
    // we're reading the bases off from the "high" end of the integer so we need to do some
    // math to figure out where they start (this helps us just pop the bases on the end
    // of the working buffer as we read them off "left to right")
    let offset = (kmer.1 - 1) * 2;
    let bitmask = BitKmerSeq::pow(2, u32::from(2 * kmer.1 - 1))
        + BitKmerSeq::pow(2, u32::from(2 * kmer.1 - 2));

    for _ in 0..kmer.1 {
        let new_char = (new_kmer & bitmask) >> offset;
        new_kmer <<= 2;
        new_kmer_str.push(match new_char {
            0 => b'A',
            1 => b'C',
            2 => b'G',
            3 => b'T',
            _ => panic!("Mathematical impossibility"),
        });
    }
    new_kmer_str
}

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

    #[test]
    fn can_kmerize() {
        // test general function
        let mut i = 0;
        for (_, k, _) in BitNuclKmer::new(b"AGCT", 1, false) {
            match i {
                0 => assert_eq!(k.0, 0b00 as BitKmerSeq),
                1 => assert_eq!(k.0, 0b10 as BitKmerSeq),
                2 => assert_eq!(k.0, 0b01 as BitKmerSeq),
                3 => assert_eq!(k.0, 0b11 as BitKmerSeq),
                _ => unreachable!("Too many kmers"),
            }
            i += 1;
        }

        // test that we skip over N's
        i = 0;
        for (_, k, _) in BitNuclKmer::new(b"ACNGT", 2, false) {
            match i {
                0 => assert_eq!(k.0, 0b0001 as BitKmerSeq),
                1 => assert_eq!(k.0, 0b1011 as BitKmerSeq),
                _ => unreachable!("Too many kmers"),
            }
            i += 1;
        }

        // test that we skip over N's and handle short kmers
        i = 0;
        for (_, k, _) in BitNuclKmer::new(b"ACNG", 2, false) {
            match i {
                0 => assert_eq!(k.0, 0x0001 as BitKmerSeq),
                _ => unreachable!("Too many kmers"),
            }
            i += 1;
        }

        // test that the minimum length works
        i = 0;
        for (_, k, _) in BitNuclKmer::new(b"AC", 2, false) {
            match i {
                0 => assert_eq!(k.0, 0x0001 as BitKmerSeq),
                _ => unreachable!("Too many kmers"),
            }
            i += 1;
        }
    }

    #[test]
    fn test_iterator() {
        let seq = b"ACGTA";
        let mut kmer_iter = BitNuclKmer::new(seq, 3, false);
        assert_eq!(kmer_iter.next(), Some((0, (6, 3), false)));
        assert_eq!(kmer_iter.next(), Some((1, (27, 3), false)));
        assert_eq!(kmer_iter.next(), Some((2, (44, 3), false)));
        assert_eq!(kmer_iter.next(), None);

        let seq = b"TA";
        let mut kmer_iter = BitNuclKmer::new(seq, 3, false);
        assert_eq!(kmer_iter.next(), None);
    }

    #[test]
    fn test_reverse_complement() {
        assert_eq!(reverse_complement((0b00_0000, 3)).0, 0b11_1111);
        assert_eq!(reverse_complement((0b11_1111, 3)).0, 0b00_0000);
        assert_eq!(reverse_complement((0b0000_0000, 4)).0, 0b1111_1111);
        assert_eq!(reverse_complement((0b0001_1011, 4)).0, 0b0001_1011);
    }

    #[test]
    fn test_minimizer() {
        assert_eq!(minimizer((0b00_1011, 3), 2).0, 0b0010);
        assert_eq!(minimizer((0b00_1011, 3), 1).0, 0b00);
        assert_eq!(minimizer((0b1100_0011, 4), 2).0, 0b0000);
        assert_eq!(minimizer((0b11_0001, 3), 2).0, 0b0001);
    }

    #[test]
    fn test_bytes_to_bitkmer() {
        let mut ikmer: BitKmer = bytes_to_bitmer(b"C");
        assert_eq!(ikmer.0, 1 as BitKmerSeq);

        ikmer = bytes_to_bitmer(b"TTA");
        assert_eq!(ikmer.0, 60 as BitKmerSeq);

        ikmer = bytes_to_bitmer(b"AAA");
        assert_eq!(ikmer.0, 0 as BitKmerSeq);
    }

    #[test]
    fn test_bitmer_to_bytes() {
        assert_eq!(bitmer_to_bytes((1 as BitKmerSeq, 1)), b"C".to_vec());
        assert_eq!(bitmer_to_bytes((60 as BitKmerSeq, 3)), b"TTA".to_vec());
        assert_eq!(bitmer_to_bytes((0 as BitKmerSeq, 3)), b"AAA".to_vec());
    }

    pub fn bytes_to_bitmer(kmer: &[u8]) -> BitKmer {
        let k = kmer.len() as u8;

        let mut bit_kmer = (0u64, k);
        for i in 0..k {
            extend_kmer(&mut bit_kmer, kmer[i as usize]);
        }
        bit_kmer
    }
}
<commit_msg>Uses a lookup table instead of match to convert nucleotides to 2 bits<commit_after>//! Compact binary representations of nucleic acid kmers
pub type BitKmerSeq = u64;
pub type BitKmer = (BitKmerSeq, u8);

static NUC2BIT_LOOKUP: [Option<u8>; 256] = [
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    Some(0),
    None,
    Some(1),
    None,
    None,
    None,
    Some(2),
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    Some(3),
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    Some(0),
    None,
    Some(1),
    None,
    None,
    None,
    Some(2),
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    Some(3),
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
    None,
];

fn nuc2bti_lookup_nocheck(nuc: u8) -> Option<u8> {
    unsafe { *NUC2BIT_LOOKUP.get_unchecked(nuc as usize) }
}

/// Takes a BitKmer and adds a new base on the end, optionally loping off the
/// first base if the resulting kmer is too long.
fn extend_kmer(kmer: &mut BitKmer, new_char: u8) -> bool {
    if let Some(new_char_int) = nuc2bti_lookup_nocheck(new_char) {
        let new_kmer = (kmer.0 << 2) + new_char_int as BitKmerSeq;

        // mask out any overflowed bits
        kmer.0 = new_kmer & (BitKmerSeq::pow(2, u32::from(2 * kmer.1)) - 1) as BitKmerSeq;
        true
    } else {
        false
    }
}

/// Used for the BitNuclKmer iterator to handle skipping invalid bases.
fn update_position(
    start_pos: &mut usize,
    kmer: &mut BitKmer,
    buffer: &[u8],
    initial: bool,
) -> bool {
    // check if we have enough "physical" space for one more kmer
    if *start_pos + kmer.1 as usize > buffer.len() {
        return false;
    }

    let (mut kmer_len, stop_len) = if initial {
        (0, (kmer.1 - 1) as usize)
    } else {
        ((kmer.1 - 1) as usize, kmer.1 as usize)
    };

    let mut cur_kmer = kmer;
    while kmer_len < stop_len {
        if extend_kmer(&mut cur_kmer, buffer[*start_pos + kmer_len]) {
            kmer_len += 1;
        } else {
            kmer_len = 0;
            *cur_kmer = (0u64, cur_kmer.1);
            *start_pos += kmer_len + 1;
            if *start_pos + cur_kmer.1 as usize > buffer.len() {
                return false;
            }
        }
    }
    true
}

pub struct BitNuclKmer<'a> {
    start_pos: usize,
    cur_kmer: BitKmer,
    buffer: &'a [u8],
    canonical: bool,
}

impl<'a> BitNuclKmer<'a> {
    pub fn new(slice: &'a [u8], k: u8, canonical: bool) -> BitNuclKmer<'a> {
        let mut kmer = (0u64, k);
        let mut start_pos = 0;
        update_position(&mut start_pos, &mut kmer, &slice, true);

        BitNuclKmer {
            start_pos,
            cur_kmer: kmer,
            buffer: slice,
            canonical,
        }
    }
}

impl<'a> Iterator for BitNuclKmer<'a> {
    type Item = (usize, BitKmer, bool);

    fn next(&mut self) -> Option<(usize, BitKmer, bool)> {
        if !update_position(&mut self.start_pos, &mut self.cur_kmer, &self.buffer, false) {
            return None;
        }
        self.start_pos += 1;
        if self.canonical {
            let (kmer, was_rc) = canonical(self.cur_kmer);
            Some((self.start_pos - 1, kmer, was_rc))
        } else {
            Some((self.start_pos - 1, self.cur_kmer, false))
        }
    }
}

/// Reverse complement a BitKmer (reverses the sequence and swaps A<>T and G<>C)
pub fn reverse_complement(kmer: BitKmer) -> BitKmer {
    // FIXME: this is not going to work with BitKmers of u128 or u32
    // inspired from https://www.biostars.org/p/113640/
    let mut new_kmer = kmer.0;
    // reverse it
    new_kmer = (new_kmer >> 2 & 0x3333_3333_3333_3333) | (new_kmer & 0x3333_3333_3333_3333) << 2;
    new_kmer = (new_kmer >> 4 & 0x0F0F_0F0F_0F0F_0F0F) | (new_kmer & 0x0F0F_0F0F_0F0F_0F0F) << 4;
    new_kmer = (new_kmer >> 8 & 0x00FF_00FF_00FF_00FF) | (new_kmer & 0x00FF_00FF_00FF_00FF) << 8;
    new_kmer = (new_kmer >> 16 & 0x0000_FFFF_0000_FFFF) | (new_kmer & 0x0000_FFFF_0000_FFFF) << 16;
    new_kmer = (new_kmer >> 32 & 0x0000_0000_FFFF_FFFF) | (new_kmer & 0x0000_0000_FFFF_FFFF) << 32;
    // complement it
    new_kmer ^= 0xFFFF_FFFF_FFFF_FFFF;
    // shift it to the right size
    new_kmer >>= 2 * (32 - kmer.1);
    (new_kmer, kmer.1)
}

/// Return the lexigraphically lowest of the BitKmer and its reverse complement and
/// whether the returned kmer is the reverse_complement (true) or the original (false)
pub fn canonical(kmer: BitKmer) -> (BitKmer, bool) {
    let rc = reverse_complement(kmer);
    if kmer.0 > rc.0 {
        (rc, true)
    } else {
        (kmer, false)
    }
}

/// Find the lexigraphically lowest substring of a given length in the BitKmer
pub fn minimizer(kmer: BitKmer, minmer_size: u8) -> BitKmer {
    let mut new_kmer = kmer.0;
    let mut lowest = !(0 as BitKmerSeq);
    let bitmask = (BitKmerSeq::pow(2, u32::from(2 * minmer_size)) - 1) as BitKmerSeq;
    for _ in 0..=(kmer.1 - minmer_size) {
        let cur = bitmask & new_kmer;
        if cur < lowest {
            lowest = cur;
        }
        let cur_rev = reverse_complement((bitmask & new_kmer, kmer.1));
        if cur_rev.0 < lowest {
            lowest = cur_rev.0;
        }
        new_kmer >>= 2;
    }
    (lowest, kmer.1)
}

pub fn bitmer_to_bytes(kmer: BitKmer) -> Vec<u8> {
    let mut new_kmer = kmer.0;
    let mut new_kmer_str = Vec::new();
    // we're reading the bases off from the "high" end of the integer so we need to do some
    // math to figure out where they start (this helps us just pop the bases on the end
    // of the working buffer as we read them off "left to right")
    let offset = (kmer.1 - 1) * 2;
    let bitmask = BitKmerSeq::pow(2, u32::from(2 * kmer.1 - 1))
        + BitKmerSeq::pow(2, u32::from(2 * kmer.1 - 2));

    for _ in 0..kmer.1 {
        let new_char = (new_kmer & bitmask) >> offset;
        new_kmer <<= 2;
        new_kmer_str.push(match new_char {
            0 => b'A',
            1 => b'C',
            2 => b'G',
            3 => b'T',
            _ => panic!("Mathematical impossibility"),
        });
    }
    new_kmer_str
}

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

    #[test]
    fn can_kmerize() {
        // test general function
        let mut i = 0;
        for (_, k, _) in BitNuclKmer::new(b"AGCT", 1, false) {
            match i {
                0 => assert_eq!(k.0, 0b00 as BitKmerSeq),
                1 => assert_eq!(k.0, 0b10 as BitKmerSeq),
                2 => assert_eq!(k.0, 0b01 as BitKmerSeq),
                3 => assert_eq!(k.0, 0b11 as BitKmerSeq),
                _ => unreachable!("Too many kmers"),
            }
            i += 1;
        }

        // test that we skip over N's
        i = 0;
        for (_, k, _) in BitNuclKmer::new(b"ACNGT", 2, false) {
            match i {
                0 => assert_eq!(k.0, 0b0001 as BitKmerSeq),
                1 => assert_eq!(k.0, 0b1011 as BitKmerSeq),
                _ => unreachable!("Too many kmers"),
            }
            i += 1;
        }

        // test that we skip over N's and handle short kmers
        i = 0;
        for (_, k, _) in BitNuclKmer::new(b"ACNG", 2, false) {
            match i {
                0 => assert_eq!(k.0, 0x0001 as BitKmerSeq),
                _ => unreachable!("Too many kmers"),
            }
            i += 1;
        }

        // test that the minimum length works
        i = 0;
        for (_, k, _) in BitNuclKmer::new(b"AC", 2, false) {
            match i {
                0 => assert_eq!(k.0, 0x0001 as BitKmerSeq),
                _ => unreachable!("Too many kmers"),
            }
            i += 1;
        }
    }

    #[test]
    fn test_iterator() {
        let seq = b"ACGTA";
        let mut kmer_iter = BitNuclKmer::new(seq, 3, false);
        assert_eq!(kmer_iter.next(), Some((0, (6, 3), false)));
        assert_eq!(kmer_iter.next(), Some((1, (27, 3), false)));
        assert_eq!(kmer_iter.next(), Some((2, (44, 3), false)));
        assert_eq!(kmer_iter.next(), None);

        let seq = b"TA";
        let mut kmer_iter = BitNuclKmer::new(seq, 3, false);
        assert_eq!(kmer_iter.next(), None);
    }

    #[test]
    fn test_reverse_complement() {
        assert_eq!(reverse_complement((0b00_0000, 3)).0, 0b11_1111);
        assert_eq!(reverse_complement((0b11_1111, 3)).0, 0b00_0000);
        assert_eq!(reverse_complement((0b0000_0000, 4)).0, 0b1111_1111);
        assert_eq!(reverse_complement((0b0001_1011, 4)).0, 0b0001_1011);
    }

    #[test]
    fn test_minimizer() {
        assert_eq!(minimizer((0b00_1011, 3), 2).0, 0b0010);
        assert_eq!(minimizer((0b00_1011, 3), 1).0, 0b00);
        assert_eq!(minimizer((0b1100_0011, 4), 2).0, 0b0000);
        assert_eq!(minimizer((0b11_0001, 3), 2).0, 0b0001);
    }

    #[test]
    fn test_bytes_to_bitkmer() {
        let mut ikmer: BitKmer = bytes_to_bitmer(b"C");
        assert_eq!(ikmer.0, 1 as BitKmerSeq);

        ikmer = bytes_to_bitmer(b"TTA");
        assert_eq!(ikmer.0, 60 as BitKmerSeq);

        ikmer = bytes_to_bitmer(b"AAA");
        assert_eq!(ikmer.0, 0 as BitKmerSeq);
    }

    #[test]
    fn test_bitmer_to_bytes() {
        assert_eq!(bitmer_to_bytes((1 as BitKmerSeq, 1)), b"C".to_vec());
        assert_eq!(bitmer_to_bytes((60 as BitKmerSeq, 3)), b"TTA".to_vec());
        assert_eq!(bitmer_to_bytes((0 as BitKmerSeq, 3)), b"AAA".to_vec());
    }

    pub fn bytes_to_bitmer(kmer: &[u8]) -> BitKmer {
        let k = kmer.len() as u8;

        let mut bit_kmer = (0u64, k);
        for i in 0..k {
            extend_kmer(&mut bit_kmer, kmer[i as usize]);
        }
        bit_kmer
    }
}
<|endoftext|>


<commit_before>use crate::lexical::algorithm::*;
use crate::lexical::num::Float;

#[test]
fn float_fast_path_test() {
    // valid
    let mantissa = (1 << f32::MANTISSA_SIZE) - 1;
    let (min_exp, max_exp) = f32::exponent_limit();
    for exp in min_exp..=max_exp {
        let f = fast_path::<f32>(mantissa, exp);
        assert!(f.is_some(), "should be valid {:?}.", (mantissa, exp));
    }

    // Check slightly above valid exponents
    let f = fast_path::<f32>(123, 15);
    assert_eq!(f, Some(1.23e+17));

    // Exponent is 1 too high, pushes over the mantissa.
    let f = fast_path::<f32>(123, 16);
    assert!(f.is_none());

    // Mantissa is too large, checked_mul should overflow.
    let f = fast_path::<f32>(mantissa, 11);
    assert!(f.is_none());

    // invalid mantissa
    #[cfg(feature = "radix")]
    {
        let (_, max_exp) = f64::exponent_limit(3);
        let f = fast_path::<f32>(1 << f32::MANTISSA_SIZE, 3, max_exp + 1);
        assert!(f.is_none(), "invalid mantissa");
    }

    // invalid exponents
    let (min_exp, max_exp) = f32::exponent_limit();
    let f = fast_path::<f32>(mantissa, min_exp - 1);
    assert!(f.is_none(), "exponent under min_exp");

    let f = fast_path::<f32>(mantissa, max_exp + 1);
    assert!(f.is_none(), "exponent above max_exp");
}

#[test]
fn double_fast_path_test() {
    // valid
    let mantissa = (1 << f64::MANTISSA_SIZE) - 1;
    let (min_exp, max_exp) = f64::exponent_limit();
    for exp in min_exp..=max_exp {
        let f = fast_path::<f64>(mantissa, exp);
        assert!(f.is_some(), "should be valid {:?}.", (mantissa, exp));
    }

    // invalid mantissa
    #[cfg(feature = "radix")]
    {
        let (_, max_exp) = f64::exponent_limit(3);
        let f = fast_path::<f64>(1 << f64::MANTISSA_SIZE, 3, max_exp + 1);
        assert!(f.is_none(), "invalid mantissa");
    }

    // invalid exponents
    let (min_exp, max_exp) = f64::exponent_limit();
    let f = fast_path::<f64>(mantissa, min_exp - 1);
    assert!(f.is_none(), "exponent under min_exp");

    let f = fast_path::<f64>(mantissa, max_exp + 1);
    assert!(f.is_none(), "exponent above max_exp");

    assert_eq!(
        Some(0.04628372940652459),
        fast_path::<f64>(4628372940652459, -17)
    );
    assert_eq!(None, fast_path::<f64>(26383446160308229, -272));
}

#[test]
fn moderate_path_test() {
    let (f, valid) = moderate_path::<f64>(1234567890, -1, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.0);

    let (f, valid) = moderate_path::<f64>(1234567891, -1, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.1);

    let (f, valid) = moderate_path::<f64>(12345678912, -2, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.12);

    let (f, valid) = moderate_path::<f64>(123456789123, -3, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.123);

    let (f, valid) = moderate_path::<f64>(1234567891234, -4, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.1234);

    let (f, valid) = moderate_path::<f64>(12345678912345, -5, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.12345);

    let (f, valid) = moderate_path::<f64>(123456789123456, -6, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.123456);

    let (f, valid) = moderate_path::<f64>(1234567891234567, -7, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.1234567);

    let (f, valid) = moderate_path::<f64>(12345678912345679, -8, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.12345679);

    let (f, valid) = moderate_path::<f64>(4628372940652459, -17, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 0.04628372940652459);

    let (f, valid) = moderate_path::<f64>(26383446160308229, -272, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 2.6383446160308229e-256);

    let (_, valid) = moderate_path::<f64>(26383446160308230, -272, false);
    assert!(!valid, "should be invalid");
}
<commit_msg>Remove dead lexical test code involving radix<commit_after>use crate::lexical::algorithm::*;
use crate::lexical::num::Float;

#[test]
fn float_fast_path_test() {
    // valid
    let mantissa = (1 << f32::MANTISSA_SIZE) - 1;
    let (min_exp, max_exp) = f32::exponent_limit();
    for exp in min_exp..=max_exp {
        let f = fast_path::<f32>(mantissa, exp);
        assert!(f.is_some(), "should be valid {:?}.", (mantissa, exp));
    }

    // Check slightly above valid exponents
    let f = fast_path::<f32>(123, 15);
    assert_eq!(f, Some(1.23e+17));

    // Exponent is 1 too high, pushes over the mantissa.
    let f = fast_path::<f32>(123, 16);
    assert!(f.is_none());

    // Mantissa is too large, checked_mul should overflow.
    let f = fast_path::<f32>(mantissa, 11);
    assert!(f.is_none());

    // invalid exponents
    let (min_exp, max_exp) = f32::exponent_limit();
    let f = fast_path::<f32>(mantissa, min_exp - 1);
    assert!(f.is_none(), "exponent under min_exp");

    let f = fast_path::<f32>(mantissa, max_exp + 1);
    assert!(f.is_none(), "exponent above max_exp");
}

#[test]
fn double_fast_path_test() {
    // valid
    let mantissa = (1 << f64::MANTISSA_SIZE) - 1;
    let (min_exp, max_exp) = f64::exponent_limit();
    for exp in min_exp..=max_exp {
        let f = fast_path::<f64>(mantissa, exp);
        assert!(f.is_some(), "should be valid {:?}.", (mantissa, exp));
    }

    // invalid exponents
    let (min_exp, max_exp) = f64::exponent_limit();
    let f = fast_path::<f64>(mantissa, min_exp - 1);
    assert!(f.is_none(), "exponent under min_exp");

    let f = fast_path::<f64>(mantissa, max_exp + 1);
    assert!(f.is_none(), "exponent above max_exp");

    assert_eq!(
        Some(0.04628372940652459),
        fast_path::<f64>(4628372940652459, -17)
    );
    assert_eq!(None, fast_path::<f64>(26383446160308229, -272));
}

#[test]
fn moderate_path_test() {
    let (f, valid) = moderate_path::<f64>(1234567890, -1, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.0);

    let (f, valid) = moderate_path::<f64>(1234567891, -1, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.1);

    let (f, valid) = moderate_path::<f64>(12345678912, -2, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.12);

    let (f, valid) = moderate_path::<f64>(123456789123, -3, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.123);

    let (f, valid) = moderate_path::<f64>(1234567891234, -4, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.1234);

    let (f, valid) = moderate_path::<f64>(12345678912345, -5, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.12345);

    let (f, valid) = moderate_path::<f64>(123456789123456, -6, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.123456);

    let (f, valid) = moderate_path::<f64>(1234567891234567, -7, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.1234567);

    let (f, valid) = moderate_path::<f64>(12345678912345679, -8, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 123456789.12345679);

    let (f, valid) = moderate_path::<f64>(4628372940652459, -17, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 0.04628372940652459);

    let (f, valid) = moderate_path::<f64>(26383446160308229, -272, false);
    assert!(valid, "should be valid");
    assert_eq!(f.into_float::<f64>(), 2.6383446160308229e-256);

    let (_, valid) = moderate_path::<f64>(26383446160308230, -272, false);
    assert!(!valid, "should be invalid");
}
<|endoftext|>