<commit_before>extern crate bio;
extern crate argparse;
use std::io;
use std::io::Write;
use std::cmp;
use bio::io::fasta;
use std::fmt;
use std::fs::File;
use argparse::{ArgumentParser, StoreTrue, Store, Print, Collect};

const EXIT_FILE_IO_ERROR: i32 = 1;
const EXIT_COMMAND_LINE_ERROR: i32 = 2;
const EXIT_FASTA_PARSE_ERROR: i32 = 3;
static PROGRAM_NAME: &'static str = "biotool";

fn exit_with_error(status: i32, message: &String) -> () {
    writeln!(&mut std::io::stderr(),
             "{} ERROR: {}!",
             PROGRAM_NAME,
             message)
        .unwrap();
    std::process::exit(status);
}

#[derive(Debug, PartialEq)]
pub struct FastaStats {
    min_len: u64,
    average_len: u64,
    max_len: u64,
    total: u64,
    num_seqs: u64,
}

#[derive(Debug)]
pub enum StatsResult {
    StatsNone,
    StatsSome(FastaStats),
    StatsError(io::Error),
}

impl FastaStats {
    pub fn new<R: io::Read>(minlen: u64, reader: R) -> Result<Option<FastaStats>, io::Error> {
        let fasta_reader = fasta::Reader::new(reader);
        let mut num_seqs: u64 = 0;
        let mut total: u64 = 0;
        let mut max_len: u64 = 0;
        let mut min_len: u64 = 0;
        let mut this_len: u64;

        for next in fasta_reader.records() {
            match next {
                Ok(record) => {
                    this_len = record.seq().len() as u64;
                    if this_len >= minlen {
                        num_seqs += 1;
                        total += this_len;
                        if num_seqs == 1 {
                            max_len = this_len;
                            min_len = this_len;
                        } else {
                            max_len = cmp::max(max_len, this_len);
                            min_len = cmp::min(min_len, this_len);
                        }
                    }
                }
                // Err(error) => return StatsResult::StatsError(error)
                Err(error) => return Err(error),
            }
        }
        if num_seqs > 0 {
            let average_len = ((total as f64) / (num_seqs as f64)).floor() as u64;
            Ok(Some(FastaStats {
                min_len: min_len,
                average_len: average_len,
                max_len: max_len,
                total: total,
                num_seqs: num_seqs,
            }))
        } else {
            Ok(None)
        }
    }
}

impl fmt::Display for FastaStats {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "{}\t{}\t{}\t{}\t{}",
               self.num_seqs,
               self.total,
               self.min_len,
               self.average_len,
               self.max_len)
    }
}

struct Options {
    verbose: bool,
    minlen: u64,
    fasta_files: Vec<String>,
}

fn parse_options() -> Options {
    let mut options = Options {
        verbose: false,
        minlen: 0,
        fasta_files: Vec::new(),
    };
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Print fasta stats");
        ap.refer(&mut options.verbose)
            .add_option(&["-v", "--verbose"],
                        StoreTrue,
                        "Print more stuff about what's happening");
        ap.refer(&mut options.minlen)
            .add_option(&["--minlen"],
                        Store,
                        "Minimum length sequence to include in stats");
        ap.refer(&mut options.fasta_files)
            .add_argument(&"fasta_files", Collect, "Input FASTA files");
        ap.add_option(&["--version"],
                      Print(env!("CARGO_PKG_VERSION").to_string()),
                      "Show version");
        match ap.parse_args() {
            Ok(()) => (),
            Err(0) => {
                // Help and version commands
                std::process::exit(0);
            }
            Err(_error_code) => {
                // Parse error
                std::process::exit(EXIT_COMMAND_LINE_ERROR);
            }
        }
    }
    return options;
}

fn compute_print_stats<R: io::Read>(options: &Options, filename: &String, reader: R) -> () {
    match FastaStats::new(options.minlen, reader) {
        // StatsResult::StatsSome(stats) => {
        Ok(Some(stats)) => {
            println!("{}\t{}", filename, stats);
        }
        Ok(None) => {
            println!("{}\t0\t0\t-\t-\t-", filename);
        }
        Err(error) => exit_with_error(EXIT_FASTA_PARSE_ERROR, &format!("{}", error)),
    }
}

fn main() {
    let options = parse_options();
    println!("FILENAME\tTOTAL\tNUMSEQ\tMIN\tAVG\tMAX");
    if options.fasta_files.len() == 0 {
        compute_print_stats(&options, &String::from("stdin"), io::stdin());
    } else {
        for filename in &options.fasta_files {
            match File::open(filename) {
                Ok(file) => {
                    compute_print_stats(&options, filename, file);
                }
                Err(error) => exit_with_error(EXIT_FILE_IO_ERROR, &format!("{}", error)),
            }
        }
    }
}

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

    fn test_fastastats_ok(minlen: u64, input: &String, expected: Option<FastaStats>) -> () {
        match FastaStats::new(minlen, input.as_bytes()) {
            Ok(result) => {
                if result != expected {
                    panic!(format!("expected {:?} got {:?}", expected, result))
                }
            }
            Err(_error) => panic!(format!("expected {:?} got Err(..)", expected)),
        }
    }

    // io::Error does not currently implement PartialEq, so we resort to
    // just checking for the existence of an error, not its actual content.
    // Hopefully this will be fixed in Rust soon:
    // https://github.com/rust-lang/rust/issues/34158
    fn test_fastastats_err(minlen: u64, input: &String) -> () {
        let comp = FastaStats::new(minlen, input.as_bytes());
        match comp {
            Ok(result) => panic!(format!("expected Err(..) got {:?}", result)),
            Err(_error) => (),
        }
    }

    #[test]
    fn zero_byte_input() {
        test_fastastats_ok(0, &String::from(""), None)
    }

    #[test]
    fn single_newline_input() {
        test_fastastats_err(0, &String::from("\n"))
    }

    #[test]
    fn single_greater_than_input() {
        test_fastastats_ok(0,
                           &String::from(">"),
                           Some(FastaStats {
                               min_len: 0,
                               average_len: 0,
                               max_len: 0,
                               total: 0,
                               num_seqs: 1,
                           }))
    }

    #[test]
    fn one_sequence() {
        test_fastastats_ok(0,
                           &String::from(">header\nATGC\nA"),
                           Some(FastaStats {
                               min_len: 5,
                               average_len: 5,
                               max_len: 5,
                               total: 5,
                               num_seqs: 1,
                           }))
    }

    #[test]
    fn two_sequence() {
        test_fastastats_ok(0,
                           &String::from(">header1\nATGC\nAGG\n>header2\nTT\n"),
                           Some(FastaStats {
                               min_len: 2,
                               average_len: 4,
                               max_len: 7,
                               total: 9,
                               num_seqs: 2,
                           }))
    }

    #[test]
    fn no_header() {
        test_fastastats_err(0, &String::from("no header\n"))
    }

    #[test]
    fn minlen_less_than_all() {
        test_fastastats_ok(2,
                           &String::from(">header1\nATGC\nAGG\n>header2\nTT\n"),
                           Some(FastaStats {
                               min_len: 2,
                               average_len: 4,
                               max_len: 7,
                               total: 9,
                               num_seqs: 2,
                           }))
    }

    #[test]
    fn minlen_greater_than_one() {
        test_fastastats_ok(3,
                           &String::from(">header1\nATGC\nAGG\n>header2\nTT\n"),
                           Some(FastaStats {
                               min_len: 7,
                               average_len: 7,
                               max_len: 7,
                               total: 7,
                               num_seqs: 1,
                           }))
    }

    #[test]
    fn minlen_greater_than_all() {
        test_fastastats_ok(8,
                           &String::from(">header1\nATGC\nAGG\n>header2\nTT\n"),
                           None)
    }
}
<commit_msg>Add comments to Rust version<commit_after>/// Module      : Main
/// Description : The main entry point for the program.
/// Copyright   : (c) Bernie Pope, 2016
/// License     : MIT
/// Maintainer  : bjpope@unimelb.edu.au
/// Stability   : experimental
/// Portability : POSIX
///
/// Read a single FASTA file as input and compute:
///  * Num sequences.
///  * Total number of bases in all sequences.
///  * Minimum sequence length.
///  * Maximum sequence length.
///  * Average sequence length, as an integer, rounded towards zero.
///
/// Sequences whose length is less than a specified minimum length are
/// ignored (skipped) and not included in the calculation of the statistics.
///
/// The basic statistics cannot be computed for empty files, and those
/// with no reads that meet the mimimum length requirement. In such cases
/// the result is Nothing.
///
/// We do not make any assumptions about the type of data stored in the
/// FASTA file, it could be DNA, proteins or anything else. No checks are
/// made to ensure that the FASTA file makes sense biologically.
///
/// Whitespace within the sequences will be ignored.

extern crate bio;
extern crate argparse;
use std::io;
use std::io::Write;
use std::cmp;
use bio::io::fasta;
use std::fmt;
use std::fs::File;
use argparse::{ArgumentParser, StoreTrue, Store, Print, Collect};

// File I/O error. This can occur if at least one of the input FASTA
// files cannot be opened for reading. This can occur because the file
// does not exist at the specified path, or biotool does not have
// permission to read from the file.
const EXIT_FILE_IO_ERROR: i32 = 1;

// A command line error occurred. This can happen if the user specifies
// an incorrect command line argument. In this circumstance biotool will
// also print a usage message to the standard error device (stderr).
const EXIT_COMMAND_LINE_ERROR: i32 = 2;

// Input FASTA file is invalid. This can occur if biotool can read an
// input file but the file format is invalid.
const EXIT_FASTA_PARSE_ERROR: i32 = 3;

// Name of the program, to be used in diagnostic messages.
static PROGRAM_NAME: &'static str = "biotool";

/// Exit the program, printing an error message on stderr, and returning
/// a specific error code. The program name is prefixed onto the front of
/// the error message.
fn exit_with_error(status: i32, message: &String) -> () {
    writeln!(&mut std::io::stderr(),
             "{} ERROR: {}!",
             PROGRAM_NAME,
             message)
        .unwrap();
    std::process::exit(status);
}

/// Basic statistics computed for a FASTA file.
/// Note that all values are with respect to only those reads whose
/// length is at least as long as the minimum.
#[derive(Debug, PartialEq)]
pub struct FastaStats {
    /// Minimum length of all sequences in the file.
    min_len: u64,
    /// Average length of all sequences in the file rounded towards zero.
    average_len: u64,
    /// Maximum length of all sequences in the file.
    max_len: u64,
    /// Total number of bases from all the sequences in the file.
    total: u64,
    /// Total number of sequences in the file.
    num_seqs: u64,
}

/// Construct a FastaStats value from an input FASTA file. Sequences
/// of length < `minlen` are ignored.
/// The input FASTA file is represented as an `io::Read` reader.
/// The result is of type `Result<Option<FastaStats>, io::Error>`
///  * `Err(error)` if an error occurred while reading the FASTA file.
///  * `Ok(Nothing)` if the FASTA file had no sequences of length >= `minlen`.
///  * `Ok(Some(stats))` if there were no errors reading the file and there was
///     at least one sequence in the file whose length was >= `minlen`.
impl FastaStats {
    pub fn new<R: io::Read>(minlen: u64, reader: R) -> Result<Option<FastaStats>, io::Error> {
        let fasta_reader = fasta::Reader::new(reader);
        let mut num_seqs: u64 = 0;
        let mut total: u64 = 0;
        let mut max_len: u64 = 0;
        let mut min_len: u64 = 0;
        let mut this_len: u64;

        // Read each sequence in the input FASTA file.
        for next in fasta_reader.records() {
            match next {
                // The sequence was parsed correctly.
                Ok(record) => {
                    // Filter out sequences whose length is too short.
                    this_len = record.seq().len() as u64;
                    if this_len >= minlen {
                        num_seqs += 1;
                        total += this_len;
                        if num_seqs == 1 {
                            // This is the first sequence we have
                            // encountered, set max and min to the
                            // length of this sequence.
                            max_len = this_len;
                            min_len = this_len;
                        } else {
                            // Update max and min.
                            max_len = cmp::max(max_len, this_len);
                            min_len = cmp::min(min_len, this_len);
                        }
                    }
                }
                // There was an error parsing the sequence.
                Err(error) => return Err(error),
            }
        }
        if num_seqs > 0 {
            // We encountered at least one sequence, so we can
            // compute statistics for this file.
            let average_len = ((total as f64) / (num_seqs as f64)).floor() as u64;
            Ok(Some(FastaStats {
                min_len: min_len,
                average_len: average_len,
                max_len: max_len,
                total: total,
                num_seqs: num_seqs,
            }))
        } else {
            // We did not encounter any sequences (satisfying the length
            // requirement). So we cannot compute statistics for this file.
            Ok(None)
        }
    }
}

/// Format the FastaStats for output.
/// Use a tab as a delimiter.
impl fmt::Display for FastaStats {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "{}\t{}\t{}\t{}\t{}",
               self.num_seqs,
               self.total,
               self.min_len,
               self.average_len,
               self.max_len)
    }
}

/// Command line arguments supplied to the program.
struct Options {
    /// If True, make the program produce more detailed output
    /// about its progress.
    verbose: bool,
    /// Minimum length sequence considered by the program. Sequences
    /// shorter than this length are ignored.
    minlen: u64,
    /// Possibly empty vector of input FASTA file paths.
    fasta_files: Vec<String>,
}

/// Parse the command line options.
/// Program exits with status 2 on parsing failure.
fn parse_options() -> Options {
    let mut options = Options {
        verbose: false,
        minlen: 0,
        fasta_files: Vec::new(),
    };
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Print fasta stats");
        ap.refer(&mut options.verbose)
            .add_option(&["-v", "--verbose"],
                        StoreTrue,
                        "Print more stuff about what's happening");
        ap.refer(&mut options.minlen)
            .add_option(&["--minlen"],
                        Store,
                        "Minimum length sequence to include in stats");
        ap.refer(&mut options.fasta_files)
            .add_argument(&"fasta_files", Collect, "Input FASTA files");
        ap.add_option(&["--version"],
                      Print(env!("CARGO_PKG_VERSION").to_string()),
                      "Show version");
        match ap.parse_args() {
            Ok(()) => (),
            Err(0) => {
                // Help and version commands
                std::process::exit(0);
            }
            Err(_error_code) => {
                // Parse error
                std::process::exit(EXIT_COMMAND_LINE_ERROR);
            }
        }
    }
    return options;
}

/// Compute basic statistics for the input FASTA files,
/// and pretty print the results to standard output.
/// Program exits if a parse error occurs when reading an
/// input FASTA file.
fn compute_print_stats<R: io::Read>(options: &Options, filename: &String, reader: R) -> () {
    match FastaStats::new(options.minlen, reader) {
        Ok(Some(stats)) => {
            // Prefix the FASTA filename onto the front of the statistics
            println!("{}\t{}", filename, stats);
        }
        Ok(None) => {
            // We could not compute any statistics for the file because
            // it contained no sequences at least as long as the minimum
            // length. In this case we just print dashes.
            println!("{}\t0\t0\t-\t-\t-", filename);
        }
        // There was a parse error when reading a FASTA file.
        // Exit the program.
        Err(error) => exit_with_error(EXIT_FASTA_PARSE_ERROR, &format!("{}", error)),
    }
}

/// The entry point for the program.
///  * Parse the command line arguments.
///  * Print the output header.
///  * Process each input FASTA file, compute stats and display
///    results.
fn main() {
    let options = parse_options();
    // Display the output header.
    println!("FILENAME\tTOTAL\tNUMSEQ\tMIN\tAVG\tMAX");
    if options.fasta_files.len() == 0 {
        // No FASTA files were specified on the command line, so
        // read from stdin instead.
        compute_print_stats(&options, &String::from("stdin"), io::stdin());
    } else {
        // Process each FASTA file specified on the command line.
        // Exit the program if a file I/O error occurs.
        for filename in &options.fasta_files {
            match File::open(filename) {
                Ok(file) => {
                    compute_print_stats(&options, filename, file);
                }
                Err(error) => exit_with_error(EXIT_FILE_IO_ERROR, &format!("{}", error)),
            }
        }
    }
}

/// Unit testing.
#[cfg(test)]
mod tests {
    use super::*;

    /// Helper function for running test cases.
    /// Arguments are:
    ///  * `minlen`: minimum length of sequences in the FASTA file to be
    ///     considered by the program.
    ///  * `input`: a string containing the contents of the input FASTA file.
    ///  * `expected`: the expected result of computing statistics for the input
    ///     FASTA file. `None` means that we do not expect any statistics to be
    ///     computed because there are no sequences in the input which meet the
    ///     minimum length requirements. `Some(stats)` means that we expect some
    ///     particular statistics to be computed.
    /// We compute the `FastaStats` for the input FASTA file and then compare against
    /// the expected result. If they are different we `panic` with a error message.
    fn test_fastastats_ok(minlen: u64, input: &String, expected: Option<FastaStats>) -> () {
        match FastaStats::new(minlen, input.as_bytes()) {
            Ok(result) => {
                if result != expected {
                    panic!(format!("expected {:?} got {:?}", expected, result))
                }
            }
            Err(_error) => panic!(format!("expected {:?} got Err(..)", expected)),
        }
    }

    /// Test for computations which are expected to fail.
    /// NB io::Error does not currently implement PartialEq, so we resort to
    /// just checking for the existence of an error, not its actual content.
    /// Hopefully this will be fixed in Rust soon:
    /// https://github.com/rust-lang/rust/issues/34158
    fn test_fastastats_err(minlen: u64, input: &String) -> () {
        let comp = FastaStats::new(minlen, input.as_bytes());
        match comp {
            Ok(result) => panic!(format!("expected Err(..) got {:?}", result)),
            Err(_error) => (),
        }
    }

    /// Empty input FASTA file. Should not produce any statistics.
    #[test]
    fn zero_byte_input() {
        test_fastastats_ok(0, &String::from(""), None)
    }

    /// Input FASTA file consisting of just a newline character.
    /// Result is expected to be an error, as the input is not considered
    /// a valid FASTA file.
    #[test]
    fn single_newline_input() {
        test_fastastats_err(0, &String::from("\n"))
    }

    /// Input FASTA file consisting of a single greater than sign,
    /// which is the minimal requirement for the FASTA header.
    #[test]
    fn single_greater_than_input() {
        test_fastastats_ok(0,
                           &String::from(">"),
                           Some(FastaStats {
                               min_len: 0,
                               average_len: 0,
                               max_len: 0,
                               total: 0,
                               num_seqs: 1,
                           }))
    }

    /// Input FASTA file consisting of a single sequence.
    /// The sequence is split over two lines.
    #[test]
    fn one_sequence() {
        test_fastastats_ok(0,
                           &String::from(">header\nATGC\nA"),
                           Some(FastaStats {
                               min_len: 5,
                               average_len: 5,
                               max_len: 5,
                               total: 5,
                               num_seqs: 1,
                           }))
    }

    /// Input FASTA file consisting of two sequences.
    /// The sequences are split over multiple lines.
    #[test]
    fn two_sequence() {
        test_fastastats_ok(0,
                           &String::from(">header1\nATGC\nAGG\n>header2\nTT\n"),
                           Some(FastaStats {
                               min_len: 2,
                               average_len: 4,
                               max_len: 7,
                               total: 9,
                               num_seqs: 2,
                           }))
    }

    /// Input FASTA file without a sequence header. This is considered an
    /// error, because it is not a valid FASTA file.
    #[test]
    fn no_header() {
        test_fastastats_err(0, &String::from("no header\n"))
    }

    /// Input FASTA file with 2 sequences, and minlen less than the lengths
    /// of all the sequences in the file. None of the sequences should be
    /// filtered out.
    #[test]
    fn minlen_less_than_all() {
        test_fastastats_ok(2,
                           &String::from(">header1\nATGC\nAGG\n>header2\nTT\n"),
                           Some(FastaStats {
                               min_len: 2,
                               average_len: 4,
                               max_len: 7,
                               total: 9,
                               num_seqs: 2,
                           }))
    }

    /// Input FASTA file with 2 sequences, and minlen greater than
    /// the length of one of the sequences. This sequence should be
    /// filtered out, and not considered in the calculation of the
    /// statistics.
    #[test]
    fn minlen_greater_than_one() {
        test_fastastats_ok(3,
                           &String::from(">header1\nATGC\nAGG\n>header2\nTT\n"),
                           Some(FastaStats {
                               min_len: 7,
                               average_len: 7,
                               max_len: 7,
                               total: 7,
                               num_seqs: 1,
                           }))
    }

    /// Input FASTA file with 2 sequences, and minlen greater than
    /// the length of all the sequences. All sequences should be filtered
    /// out, and thus there are no statistics to compute.
    #[test]
    fn minlen_greater_than_all() {
        test_fastastats_ok(8,
                           &String::from(">header1\nATGC\nAGG\n>header2\nTT\n"),
                           None)
    }
}
<|endoftext|>


<commit_before>/// Detect NIF version to build against

use std::env;
use std::process::Command;

extern crate lazy_static;
use lazy_static::lazy_static;

extern crate which;
use which::which;

lazy_static! {
    // keep this sorted by version number
    static ref ERTS_VERSIONS: Vec<&'static str> = vec![
        "2.7", "2.8", "2.9", "2.10", "2.11", "2.12", "2.13", "2.14"
    ];
}

fn main() {
    let version = match env::var("RUSTLER_NIF_VERSION") {
        Ok(version) => version,
        Err(_) => get_version_from_erl()
    };

    activate_versions(&version);
}

fn get_version_from_erl() -> String {
    let erl = which("erl").expect("expected to find 'erl' executable");
    let args = vec![
        "-noshell",
        "-eval",
        r#"io:format("~s~n", [erlang:system_info(nif_version)]), init:stop()."#
    ];

    let version = Command::new(erl)
        .args(&args)
        .output()
        .expect("failed to execute 'erl'")
        .stdout;

    let version = String::from_utf8(version).expect("convert version to String");
    dbg!(version.clone());

    version.trim().into()
}

fn activate_versions(version: &str) {
    let index = ERTS_VERSIONS
        .iter()
        .position(|&v| v == version)
        .expect(&format!("Erlang version {} not handled, please file a a bug report.", version));

    for i in 0..=index {
        println!("cargo:rustc-cfg=nif_version_{}", version_feature(ERTS_VERSIONS[i]));
    }
}

fn version_feature(version: &str) -> String {
    version.replace(".", "_")
}
<commit_msg>Remove debug output<commit_after>/// Detect NIF version to build against

use std::env;
use std::process::Command;

extern crate lazy_static;
use lazy_static::lazy_static;

extern crate which;
use which::which;

lazy_static! {
    // keep this sorted by version number
    static ref ERTS_VERSIONS: Vec<&'static str> = vec![
        "2.7", "2.8", "2.9", "2.10", "2.11", "2.12", "2.13", "2.14"
    ];
}

fn main() {
    let version = match env::var("RUSTLER_NIF_VERSION") {
        Ok(version) => version,
        Err(_) => get_version_from_erl()
    };

    activate_versions(&version);
}

fn get_version_from_erl() -> String {
    let erl = which("erl").expect("expected to find 'erl' executable");
    let args = vec![
        "-noshell",
        "-eval",
        r#"io:format("~s~n", [erlang:system_info(nif_version)]), init:stop()."#
    ];

    let version = Command::new(erl)
        .args(&args)
        .output()
        .expect("failed to execute 'erl'")
        .stdout;

    let version = String::from_utf8(version).expect("convert version to String");

    version.trim().into()
}

fn activate_versions(version: &str) {
    let index = ERTS_VERSIONS
        .iter()
        .position(|&v| v == version)
        .expect(&format!("Erlang version {} not handled, please file a a bug report.", version));

    for i in 0..=index {
        println!("cargo:rustc-cfg=nif_version_{}", version_feature(ERTS_VERSIONS[i]));
    }
}

fn version_feature(version: &str) -> String {
    version.replace(".", "_")
}
<|endoftext|>


<commit_before>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A vector type intended to be used for collecting from iterators onto the stack.
//!
//! Space for up to N elements is provided on the stack.  If more elements are collected, Vec is
//! used to store the values on the heap. SmallVec is similar to AccumulateVec, but adds
//! the ability to push elements.
//!
//! The N above is determined by Array's implementor, by way of an associated constant.

use std::ops::{Deref, DerefMut};
use std::iter::{IntoIterator, FromIterator};
use std::fmt::{self, Debug};
use std::mem;
use std::ptr;

use rustc_serialize::{Encodable, Encoder, Decodable, Decoder};

use accumulate_vec::{IntoIter, AccumulateVec};
use array_vec::Array;

pub struct SmallVec(AccumulateVec);

impl Clone for SmallVec
    where A: Array,
          A::Element: Clone {
    fn clone(&self) -> Self {
        SmallVec(self.0.clone())
    }
}

impl Debug for SmallVec
    where A: Array + Debug,
          A::Element: Debug {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_tuple("SmallVec").field(&self.0).finish()
    }
}

impl SmallVec {
    pub fn new() -> Self {
        SmallVec(AccumulateVec::new())
    }

    pub fn is_array(&self) -> bool {
        self.0.is_array()
    }

    pub fn with_capacity(cap: usize) -> Self {
        let mut vec = SmallVec::new();
        vec.reserve(cap);
        vec
    }

    pub fn one(el: A::Element) -> Self {
        SmallVec(AccumulateVec::one(el))
    }

    pub fn many<I: IntoIterator<Item=A::Element>>(els: I) -> Self {
        SmallVec(AccumulateVec::many(els))
    }

    pub fn expect_one(self, err: &'static str) -> A::Element {
        assert!(self.len() == 1, err);
        match self.0 {
            AccumulateVec::Array(arr) => arr.into_iter().next().unwrap(),
            AccumulateVec::Heap(vec) => vec.into_iter().next().unwrap(),
        }
    }

    /// Will reallocate onto the heap if needed.
    pub fn push(&mut self, el: A::Element) {
        self.reserve(1);
        match self.0 {
            AccumulateVec::Array(ref mut array) => array.push(el),
            AccumulateVec::Heap(ref mut vec) => vec.push(el),
        }
    }

    pub fn reserve(&mut self, n: usize) {
        match self.0 {
            AccumulateVec::Array(_) => {
                if self.len() + n > A::LEN {
                    let len = self.len();
                    let array = mem::replace(&mut self.0,
                            AccumulateVec::Heap(Vec::with_capacity(len + n)));
                    if let AccumulateVec::Array(array) = array {
                        match self.0 {
                            AccumulateVec::Heap(ref mut vec) => vec.extend(array),
                            _ => unreachable!()
                        }
                    }
                }
            }
            AccumulateVec::Heap(ref mut vec) => vec.reserve(n)
        }
    }

    pub unsafe fn set_len(&mut self, len: usize) {
        match self.0 {
            AccumulateVec::Array(ref mut arr) => arr.set_len(len),
            AccumulateVec::Heap(ref mut vec) => vec.set_len(len),
        }
    }

    pub fn insert(&mut self, index: usize, element: A::Element) {
        let len = self.len();

        // Reserve space for shifting elements to the right
        self.reserve(1);

        assert!(index <= len);

        unsafe {
            // infallible
            // The spot to put the new value
            {
                let p = self.as_mut_ptr().offset(index as isize);
                // Shift everything over to make space. (Duplicating the
                // `index`th element into two consecutive places.)
                ptr::copy(p, p.offset(1), len - index);
                // Write it in, overwriting the first copy of the `index`th
                // element.
                ptr::write(p, element);
            }
            self.set_len(len + 1);
        }
    }

    pub fn truncate(&mut self, len: usize) {
        unsafe {
            while len < self.len() {
                // Decrement len before the drop_in_place(), so a panic on Drop
                // doesn't re-drop the just-failed value.
                let newlen = self.len() - 1;
                self.set_len(newlen);
                ::std::ptr::drop_in_place(self.get_unchecked_mut(newlen));
            }
        }
    }
}

impl Deref for SmallVec {
    type Target = AccumulateVec;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for SmallVec {
    fn deref_mut(&mut self) -> &mut AccumulateVec {
        &mut self.0
    }
}

impl FromIterator for SmallVec {
    fn from_iter(iter: I) -> Self where I: IntoIterator<Item=A::Element> {
        SmallVec(iter.into_iter().collect())
    }
}

impl Extend for SmallVec {
    fn extend<I: IntoIterator<Item=A::Element>>(&mut self, iter: I) {
        let iter = iter.into_iter();
        self.reserve(iter.size_hint().0);
        for el in iter {
            self.push(el);
        }
    }
}

impl IntoIterator for SmallVec {
    type Item = A::Element;
    type IntoIter = IntoIter;
    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl Default for SmallVec {
    fn default() -> SmallVec {
        SmallVec::new()
    }
}

impl Encodable for SmallVec
    where A: Array,
          A::Element: Encodable {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        s.emit_seq(self.len(), |s| {
            for (i, e) in self.iter().enumerate() {
                try!(s.emit_seq_elt(i, |s| e.encode(s)));
            }
            Ok(())
        })
    }
}

impl Decodable for SmallVec
    where A: Array,
          A::Element: Decodable {
    fn decode<D: Decoder>(d: &mut D) -> Result<SmallVec, D::Error> {
        d.read_seq(|d, len| {
            let mut vec = SmallVec::with_capacity(len);
            for i in 0..len {
                vec.push(try!(d.read_seq_elt(i, |d| Decodable::decode(d))));
            }
            Ok(vec)
        })
    }
}
<commit_msg>Use Vec::extend in SmallVec::extend when applicable<commit_after>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A vector type intended to be used for collecting from iterators onto the stack.
//!
//! Space for up to N elements is provided on the stack.  If more elements are collected, Vec is
//! used to store the values on the heap. SmallVec is similar to AccumulateVec, but adds
//! the ability to push elements.
//!
//! The N above is determined by Array's implementor, by way of an associated constant.

use std::ops::{Deref, DerefMut};
use std::iter::{IntoIterator, FromIterator};
use std::fmt::{self, Debug};
use std::mem;
use std::ptr;

use rustc_serialize::{Encodable, Encoder, Decodable, Decoder};

use accumulate_vec::{IntoIter, AccumulateVec};
use array_vec::Array;

pub struct SmallVec(AccumulateVec);

impl Clone for SmallVec
    where A: Array,
          A::Element: Clone {
    fn clone(&self) -> Self {
        SmallVec(self.0.clone())
    }
}

impl Debug for SmallVec
    where A: Array + Debug,
          A::Element: Debug {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_tuple("SmallVec").field(&self.0).finish()
    }
}

impl SmallVec {
    pub fn new() -> Self {
        SmallVec(AccumulateVec::new())
    }

    pub fn is_array(&self) -> bool {
        self.0.is_array()
    }

    pub fn with_capacity(cap: usize) -> Self {
        let mut vec = SmallVec::new();
        vec.reserve(cap);
        vec
    }

    pub fn one(el: A::Element) -> Self {
        SmallVec(AccumulateVec::one(el))
    }

    pub fn many<I: IntoIterator<Item=A::Element>>(els: I) -> Self {
        SmallVec(AccumulateVec::many(els))
    }

    pub fn expect_one(self, err: &'static str) -> A::Element {
        assert!(self.len() == 1, err);
        match self.0 {
            AccumulateVec::Array(arr) => arr.into_iter().next().unwrap(),
            AccumulateVec::Heap(vec) => vec.into_iter().next().unwrap(),
        }
    }

    /// Will reallocate onto the heap if needed.
    pub fn push(&mut self, el: A::Element) {
        self.reserve(1);
        match self.0 {
            AccumulateVec::Array(ref mut array) => array.push(el),
            AccumulateVec::Heap(ref mut vec) => vec.push(el),
        }
    }

    pub fn reserve(&mut self, n: usize) {
        match self.0 {
            AccumulateVec::Array(_) => {
                if self.len() + n > A::LEN {
                    let len = self.len();
                    let array = mem::replace(&mut self.0,
                            AccumulateVec::Heap(Vec::with_capacity(len + n)));
                    if let AccumulateVec::Array(array) = array {
                        match self.0 {
                            AccumulateVec::Heap(ref mut vec) => vec.extend(array),
                            _ => unreachable!()
                        }
                    }
                }
            }
            AccumulateVec::Heap(ref mut vec) => vec.reserve(n)
        }
    }

    pub unsafe fn set_len(&mut self, len: usize) {
        match self.0 {
            AccumulateVec::Array(ref mut arr) => arr.set_len(len),
            AccumulateVec::Heap(ref mut vec) => vec.set_len(len),
        }
    }

    pub fn insert(&mut self, index: usize, element: A::Element) {
        let len = self.len();

        // Reserve space for shifting elements to the right
        self.reserve(1);

        assert!(index <= len);

        unsafe {
            // infallible
            // The spot to put the new value
            {
                let p = self.as_mut_ptr().offset(index as isize);
                // Shift everything over to make space. (Duplicating the
                // `index`th element into two consecutive places.)
                ptr::copy(p, p.offset(1), len - index);
                // Write it in, overwriting the first copy of the `index`th
                // element.
                ptr::write(p, element);
            }
            self.set_len(len + 1);
        }
    }

    pub fn truncate(&mut self, len: usize) {
        unsafe {
            while len < self.len() {
                // Decrement len before the drop_in_place(), so a panic on Drop
                // doesn't re-drop the just-failed value.
                let newlen = self.len() - 1;
                self.set_len(newlen);
                ::std::ptr::drop_in_place(self.get_unchecked_mut(newlen));
            }
        }
    }
}

impl Deref for SmallVec {
    type Target = AccumulateVec;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for SmallVec {
    fn deref_mut(&mut self) -> &mut AccumulateVec {
        &mut self.0
    }
}

impl FromIterator for SmallVec {
    fn from_iter(iter: I) -> Self where I: IntoIterator<Item=A::Element> {
        SmallVec(iter.into_iter().collect())
    }
}

impl Extend for SmallVec {
    fn extend<I: IntoIterator<Item=A::Element>>(&mut self, iter: I) {
        if self.is_array() {
            let iter = iter.into_iter();
            self.reserve(iter.size_hint().0);

            for el in iter {
                self.push(el);
            }
        } else {
            match self.0 {
                AccumulateVec::Heap(ref mut vec) => vec.extend(iter),
                _ => unreachable!()
            }
        }
    }
}

impl IntoIterator for SmallVec {
    type Item = A::Element;
    type IntoIter = IntoIter;
    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl Default for SmallVec {
    fn default() -> SmallVec {
        SmallVec::new()
    }
}

impl Encodable for SmallVec
    where A: Array,
          A::Element: Encodable {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        s.emit_seq(self.len(), |s| {
            for (i, e) in self.iter().enumerate() {
                try!(s.emit_seq_elt(i, |s| e.encode(s)));
            }
            Ok(())
        })
    }
}

impl Decodable for SmallVec
    where A: Array,
          A::Element: Decodable {
    fn decode<D: Decoder>(d: &mut D) -> Result<SmallVec, D::Error> {
        d.read_seq(|d, len| {
            let mut vec = SmallVec::with_capacity(len);
            for i in 0..len {
                vec.push(try!(d.read_seq_elt(i, |d| Decodable::decode(d))));
            }
            Ok(vec)
        })
    }
}
<|endoftext|>


<commit_before>use gemlab::mesh::{At, Mesh};
use gemlab::StrError;
use russell_chk::assert_vec_approx_eq;
use russell_lab::Vector;

#[test]
fn column_distorted_tris_quads() -> Result<(), StrError> {
    // read mesh
    let mesh = Mesh::from_text_file("./data/meshes/column_distorted_tris_quads.msh")?;
    println!("{}", mesh);

    // check sizes
    assert_eq!(mesh.points.len(), 13);
    assert_eq!(mesh.cells.len(), 7);
    assert_eq!(mesh.boundary_points.len(), 13);
    assert_eq!(mesh.boundary_edges.len(), 13);
    assert_eq!(mesh.boundary_faces.len(), 0);

    // check cells
    assert_eq!(mesh.cells[0].points, &[0, 7, 8, 1]);
    assert_eq!(mesh.cells[1].points, &[1, 8, 9, 2]);
    assert_eq!(mesh.cells[2].points, &[2, 9, 10, 3]);
    assert_eq!(mesh.cells[3].points, &[3, 10, 4]);
    assert_eq!(mesh.cells[4].points, &[4, 10, 11]);
    assert_eq!(mesh.cells[5].points, &[4, 11, 5]);
    assert_eq!(mesh.cells[6].points, &[5, 11, 12, 6]);

    // check edges
    let edge = mesh.boundary_edges.get(&(0, 1)).unwrap();
    assert_eq!(edge.points, &[0, 1]);
    let edge = mesh.boundary_edges.get(&(9, 10)).unwrap();
    assert_eq!(edge.points, &[10, 9]);
    let edge = mesh.boundary_edges.get(&(3, 4)).unwrap();
    assert_eq!(edge.points, &[3, 4]);

    // the magnitude of the normal vector should be equal to edge_length / 2.0
    // for both tris or quas where 2.0 corresponds to the edge_length in the reference system
    // Note that the edge is mapped to [-1,+1] in both Lin or Qua

    // edge keys and correct normal vectors (solutions)
    let edge_keys_and_solutions = [
        // left
        (vec![(0, 1)], [-0.5 / 2.0, 0.0]),
        (vec![(1, 2)], [-0.5 / 2.0, 0.0]),
        (vec![(2, 3)], [-1.0 / 2.0, 0.0]),
        (vec![(3, 4)], [-0.5 / 2.0, 0.0]),
        (vec![(4, 5)], [-0.5 / 2.0, 0.0]),
        (vec![(5, 6)], [-0.1 / 2.0, 0.0]),
        // right
        (vec![(7, 8)], [0.2 / 2.0, 0.0]),
        (vec![(8, 9)], [0.8 / 2.0, 0.0]),
        (vec![(9, 10)], [0.8 / 2.0, 0.0]),
        (vec![(10, 11)], [1.2 / 2.0, 0.0]),
        (vec![(11, 12)], [0.1 / 2.0, 0.0]),
        // bottom
        (vec![(0, 7)], [0.0, -1.0 / 2.0]),
        // top
        (vec![(6, 12)], [0.0, 1.0 / 2.0]),
    ];

    // check if the normal vectors at boundary are outward
    let mut normal = Vector::new(mesh.space_ndim);
    let ksi = &[0.0, 0.0, 0.0];
    for (edge_keys, solution) in &edge_keys_and_solutions {
        for edge_key in edge_keys {
            let mut edge_shape = mesh.alloc_shape_boundary_edge(edge_key)?;
            assert_eq!(edge_shape.nnode, 2);
            edge_shape.calc_boundary_normal(&mut normal, ksi)?;
            assert_vec_approx_eq!(normal.as_data(), solution, 1e-14);
        }
    }

    // find points
    let points = mesh.find_boundary_points(At::X(0.0))?;
    assert_eq!(points, &[0, 1, 2, 3, 4, 5, 6]);
    let points = mesh.find_boundary_points(At::X(1.0))?;
    assert_eq!(points, &[7, 8, 9, 10, 11, 12]);

    // find edges
    let edges = mesh.find_boundary_edges(At::X(0.0))?;
    assert_eq!(edges, &[(0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6)]);
    let edges = mesh.find_boundary_edges(At::X(1.0))?;
    assert_eq!(edges, &[(7, 8), (8, 9), (9, 10), (10, 11), (11, 12)]);

    // find faces
    let faces = mesh.find_boundary_faces(At::X(0.0))?;
    assert_eq!(faces.len(), 0);
    Ok(())
}

#[test]
fn rectangle_tris_quads() -> Result<(), StrError> {
    // read mesh
    let mesh = Mesh::from_text_file("./data/meshes/rectangle_tris_quads.msh")?;
    println!("{}", mesh);
    Ok(())
}
<commit_msg>Improve test<commit_after>use gemlab::mesh::{At, Mesh};
use gemlab::StrError;
use russell_chk::assert_vec_approx_eq;
use russell_lab::Vector;

#[test]
fn column_distorted_tris_quads() -> Result<(), StrError> {
    // read mesh
    let mesh = Mesh::from_text_file("./data/meshes/column_distorted_tris_quads.msh")?;

    // check sizes
    assert_eq!(mesh.points.len(), 13);
    assert_eq!(mesh.cells.len(), 7);
    assert_eq!(mesh.boundary_points.len(), 13);
    assert_eq!(mesh.boundary_edges.len(), 13);
    assert_eq!(mesh.boundary_faces.len(), 0);

    // check cells
    assert_eq!(mesh.cells[0].points, &[0, 7, 8, 1]);
    assert_eq!(mesh.cells[1].points, &[1, 8, 9, 2]);
    assert_eq!(mesh.cells[2].points, &[2, 9, 10, 3]);
    assert_eq!(mesh.cells[3].points, &[3, 10, 4]);
    assert_eq!(mesh.cells[4].points, &[4, 10, 11]);
    assert_eq!(mesh.cells[5].points, &[4, 11, 5]);
    assert_eq!(mesh.cells[6].points, &[5, 11, 12, 6]);

    // check edges
    let edge = mesh.boundary_edges.get(&(0, 1)).unwrap();
    assert_eq!(edge.points, &[0, 1]);
    let edge = mesh.boundary_edges.get(&(9, 10)).unwrap();
    assert_eq!(edge.points, &[10, 9]);
    let edge = mesh.boundary_edges.get(&(3, 4)).unwrap();
    assert_eq!(edge.points, &[3, 4]);

    // the magnitude of the normal vector should be equal to edge_length / 2.0
    // for both tris or quas where 2.0 corresponds to the edge_length in the reference system
    // Note that the edge is mapped to [-1,+1] in both Lin or Qua

    // edge keys and correct normal vectors (solutions)
    let edge_keys_and_solutions = [
        // left
        (vec![(0, 1)], [-0.5 / 2.0, 0.0]),
        (vec![(1, 2)], [-0.5 / 2.0, 0.0]),
        (vec![(2, 3)], [-1.0 / 2.0, 0.0]),
        (vec![(3, 4)], [-0.5 / 2.0, 0.0]),
        (vec![(4, 5)], [-0.5 / 2.0, 0.0]),
        (vec![(5, 6)], [-0.1 / 2.0, 0.0]),
        // right
        (vec![(7, 8)], [0.2 / 2.0, 0.0]),
        (vec![(8, 9)], [0.8 / 2.0, 0.0]),
        (vec![(9, 10)], [0.8 / 2.0, 0.0]),
        (vec![(10, 11)], [1.2 / 2.0, 0.0]),
        (vec![(11, 12)], [0.1 / 2.0, 0.0]),
        // bottom
        (vec![(0, 7)], [0.0, -1.0 / 2.0]),
        // top
        (vec![(6, 12)], [0.0, 1.0 / 2.0]),
    ];

    // check if the normal vectors at boundary are outward
    let mut normal = Vector::new(mesh.space_ndim);
    let ksi = &[0.0, 0.0, 0.0];
    for (edge_keys, solution) in &edge_keys_and_solutions {
        for edge_key in edge_keys {
            let mut edge_shape = mesh.alloc_shape_boundary_edge(edge_key)?;
            assert_eq!(edge_shape.nnode, 2);
            edge_shape.calc_boundary_normal(&mut normal, ksi)?;
            assert_vec_approx_eq!(normal.as_data(), solution, 1e-14);
        }
    }

    // find points
    let points = mesh.find_boundary_points(At::X(0.0))?;
    assert_eq!(points, &[0, 1, 2, 3, 4, 5, 6]);
    let points = mesh.find_boundary_points(At::X(1.0))?;
    assert_eq!(points, &[7, 8, 9, 10, 11, 12]);

    // find edges
    let edges = mesh.find_boundary_edges(At::X(0.0))?;
    assert_eq!(edges, &[(0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6)]);
    let edges = mesh.find_boundary_edges(At::X(1.0))?;
    assert_eq!(edges, &[(7, 8), (8, 9), (9, 10), (10, 11), (11, 12)]);

    // find faces
    let faces = mesh.find_boundary_faces(At::X(0.0))?;
    assert_eq!(faces.len(), 0);
    Ok(())
}

#[test]
fn rectangle_tris_quads() -> Result<(), StrError> {
    // read mesh
    let mesh = Mesh::from_text_file("./data/meshes/rectangle_tris_quads.msh")?;

    // the magnitude of the normal vector should be equal to edge_length / 2.0
    // for both tris or quas where 2.0 corresponds to the edge_length in the reference system
    // Note that the edge is mapped to [-1,+1] in both Lin or Qua

    // edge keys and correct normal vectors (solutions)
    let edge_keys_and_solutions = [
        // left
        (vec![(0, 3), (3, 7), (7, 10)], [-1.0 / 2.0, 0.0]),
        // right
        (vec![(2, 6), (6, 9), (9, 13)], [1.0 / 2.0, 0.0]),
        // bottom
        (vec![(0, 1), (1, 2)], [0.0, -2.0 / 2.0]),
    ];

    // check if the normal vectors at boundary are outward
    let mut normal = Vector::new(mesh.space_ndim);
    let ksi = &[0.0, 0.0, 0.0];
    for (edge_keys, solution) in &edge_keys_and_solutions {
        for edge_key in edge_keys {
            let mut edge_shape = mesh.alloc_shape_boundary_edge(edge_key)?;
            assert_eq!(edge_shape.nnode, 2);
            edge_shape.calc_boundary_normal(&mut normal, ksi)?;
            assert_vec_approx_eq!(normal.as_data(), solution, 1e-14);
        }
    }

    // find edges
    let edges = mesh.find_boundary_edges(At::X(0.0))?;
    assert_eq!(edges, &[(0, 3), (3, 7), (7, 10), (10, 14)]);
    let edges = mesh.find_boundary_edges(At::X(4.0))?;
    assert_eq!(edges, &[(2, 6), (6, 9), (9, 13)]);
    Ok(())
}
<|endoftext|>


<commit_before>use core::ops::{Add, AddAssign, Sub, SubAssign};
use core::iter::Step;
use core::{fmt, mem};
use core::convert::{Into, TryFrom, TryInto};

use bit_field::BitField;
use ux::*;

/// A canonical 64-bit virtual memory address.
///
/// This is a wrapper type around an `u64`, so it is always 8 bytes, even when compiled
/// on non 64-bit systems. The `UsizeConversions` trait can be used for performing conversions
/// between `u64` and `usize`.
///
/// On `x86_64`, only the 48 lower bits of a virtual address can be used. The top 16 bits need
/// to be copies of bit 47, i.e. the most significant bit. Addresses that fulfil this criterium
/// are called “canonical”. This type guarantees that it always represents a canonical address.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct VirtAddr(u64);

/// A 64-bit physical memory address.
///
/// This is a wrapper type around an `u64`, so it is always 8 bytes, even when compiled
/// on non 64-bit systems. The `UsizeConversions` trait can be used for performing conversions
/// between `u64` and `usize`.
///
/// On `x86_64`, only the 52 lower bits of a physical address can be used. The top 12 bits need
/// to be zero. This type guarantees that it always represents a valid physical address.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct PhysAddr(u64);

/// A passed `u64` was not a valid virtual address.
///
/// This means that bits 48 to 64 are not
/// a valid sign extension and are not null either. So automatic sign extension would have
/// overwritten possibly meaningful bits. This likely indicates a bug, for example an invalid
/// address calculation.
#[derive(Debug)]
pub struct VirtAddrNotValid(u64);

impl VirtAddr {
    /// Creates a new canonical virtual address.
    ///
    /// This function performs sign extension of bit 47 to make the address canonical. Panics
    /// if the bits in the range 48 to 64 contain data (i.e. are not null and no sign extension).
    pub fn new(addr: u64) -> VirtAddr {
        Self::try_new(addr).expect("address passed to VirtAddr::new must not contain any data \
        in bits 48 to 64")
    }

    /// Tries to create a new canonical virtual address.
    ///
    /// This function tries to performs sign
    /// extension of bit 47 to make the address canonical. It succeeds if bits 48 to 64 are
    /// either a correct sign extension (i.e. copies of bit 47) or all null. Else, an error
    /// is returned.
    pub fn try_new(addr: u64) -> Result<VirtAddr, VirtAddrNotValid> {
        match addr.get_bits(47..64) {
            0 | 0x1ffff => Ok(VirtAddr(addr)), // address is canonical
            1 => Ok(VirtAddr::new_unchecked(addr)), // address needs sign extension
            other => Err(VirtAddrNotValid(other))
        }
    }

    /// Creates a new canonical virtual address without checks.
    ///
    /// This function performs sign extension of bit 47 to make the address canonical, so
    /// bits 48 to 64 are overwritten. If you want to check that these bits contain no data,
    /// use `new` or `try_new`.
    pub fn new_unchecked(mut addr: u64) -> VirtAddr {
        if addr.get_bit(47) {
            addr.set_bits(48..64, 0xffff);
        } else {
            addr.set_bits(48..64, 0);
        }
        VirtAddr(addr)
    }

    /// Converts the address to an `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }

    /// Aligns the virtual address upwards to the given alignment.
    ///
    /// See the `align_up` function for more information.
    pub fn align_up<U>(self, align: U) -> Self
        where U: Into<u64>
    {
        VirtAddr(align_up(self.0, align.into()))
    }

    /// Aligns the virtual address downwards to the given alignment.
    ///
    /// See the `align_down` function for more information.
    pub fn align_down<U>(self, align: U) -> Self
        where U: Into<u64>
    {
        VirtAddr(align_down(self.0, align.into()))
    }

    /// Returns the 12-bit page offset of this virtual address.
    pub fn page_offset(&self) -> u12 {
        u12::new((self.0 & 0xfff).try_into().unwrap())
    }

    /// Returns the 9-bit level 1 page table index.
    pub fn p1_index(&self) -> u9 {
        u9::new(((self.0 >> 12) & 0o777).try_into().unwrap())
    }

    /// Returns the 9-bit level 2 page table index.
    pub fn p2_index(&self) -> u9 {
        u9::new(((self.0 >> 12 >> 9) & 0o777).try_into().unwrap())
    }

    /// Returns the 9-bit level 3 page table index.
    pub fn p3_index(&self) -> u9 {
        u9::new(((self.0 >> 12 >> 9 >> 9) & 0o777).try_into().unwrap())
    }

    /// Returns the 9-bit level 4 page table index.
    pub fn p4_index(&self) -> u9 {
        u9::new(((self.0 >> 12 >> 9 >> 9 >> 9) & 0o777).try_into().unwrap())
    }
}

impl fmt::Debug for VirtAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "VirtAddr({:#x})", self.0)
    }
}

impl Add<u64> for VirtAddr {
    type Output = Self;
    fn add(self, rhs: u64) -> Self::Output {
        VirtAddr::new(self.0 + rhs)
    }
}

impl AddAssign<u64> for VirtAddr {
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl Sub<u64> for VirtAddr {
    type Output = Self;
    fn sub(self, rhs: u64) -> Self::Output {
        VirtAddr::new(self.0 - rhs)
    }
}

impl SubAssign<u64> for VirtAddr {
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl Sub<VirtAddr> for VirtAddr {
    type Output = u64;
    fn sub(self, rhs: VirtAddr) -> Self::Output {
        self.as_u64() - rhs.as_u64()
    }
}

impl Step for VirtAddr {
    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
        if *start < *end {
            usize::try_from(end.0.get_bits(0..48) - start.0.get_bits(0..48)).ok()
        } else {
            Some(0)
        }
    }

    fn replace_one(&mut self) -> Self {
        mem::replace(self, VirtAddr(1))
    }

    fn replace_zero(&mut self) -> Self {
        mem::replace(self, VirtAddr(0))
    }

    fn add_one(&self) -> Self {
        *self + 1
    }

    fn sub_one(&self) -> Self {
        *self - 1
    }

    fn add_usize(&self, n: usize) -> Option<Self> {
        u64::try_from(n).ok().and_then(|n| self.0.checked_add(n)).map(VirtAddr)
    }
}

/// A passed `u64` was not a valid physical address.
///
/// This means that bits 52 to 64 are not were not all null.
#[derive(Debug)]
pub struct PhysAddrNotValid(u64);

impl PhysAddr {
    /// Creates a new physical address.
    ///
    /// Panics if a bit in the range 52 to 64 is set.
    pub fn new(addr: u64) -> PhysAddr {
        assert_eq!(addr.get_bits(52..64), 0,
            "physical addresses must not have any bits in the range 52 to 64 set");
        PhysAddr(addr)
    }

    /// Tries to create a new physical address.
    ///
    /// Fails if any bits in the range 52 to 64 are set.
    pub fn try_new(addr: u64) -> Result<PhysAddr, PhysAddrNotValid> {
        match addr.get_bits(52..64) {
            0 => Ok(PhysAddr(addr)), // address is valid
            other => Err(PhysAddrNotValid(other))
        }
    }

    /// Converts the address to an `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }

    /// Convenience method for checking if a physical address is null.
    pub fn is_null(&self) -> bool {
        self.0 == 0
    }

    /// Aligns the physical address upwards to the given alignment.
    ///
    /// See the `align_up` function for more information.
    pub fn align_up<U>(self, align: U) -> Self
        where U: Into<u64>
    {
        PhysAddr(align_up(self.0, align.into()))
    }

    /// Aligns the physical address downwards to the given alignment.
    ///
    /// See the `align_down` function for more information.
    pub fn align_down<U>(self, align: U) -> Self
        where U: Into<u64>
    {
        PhysAddr(align_down(self.0, align.into()))
    }
}

impl fmt::Debug for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "PhysAddr({:#x})", self.0)
    }
}

impl fmt::Binary for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl fmt::LowerHex for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl fmt::Octal for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl fmt::UpperHex for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl Add<u64> for PhysAddr {
    type Output = Self;
    fn add(self, rhs: u64) -> Self::Output {
        PhysAddr::new(self.0 + rhs)
    }
}

impl AddAssign<u64> for PhysAddr {
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl Sub<u64> for PhysAddr {
    type Output = Self;
    fn sub(self, rhs: u64) -> Self::Output {
        PhysAddr::new(self.0 - rhs)
    }
}

impl SubAssign<u64> for PhysAddr {
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl Sub<PhysAddr> for PhysAddr {
    type Output = u64;
    fn sub(self, rhs: PhysAddr) -> Self::Output {
        self.as_u64() - rhs.as_u64()
    }
}

impl Step for PhysAddr {
    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
        if *start < *end {
            usize::try_from(end.0 - start.0).ok()
        } else {
            Some(0)
        }
    }

    fn replace_one(&mut self) -> Self {
        mem::replace(self, PhysAddr(1))
    }

    fn replace_zero(&mut self) -> Self {
        mem::replace(self, PhysAddr(0))
    }

    fn add_one(&self) -> Self {
        *self + 1
    }

    fn sub_one(&self) -> Self {
        *self - 1
    }

    fn add_usize(&self, n: usize) -> Option<Self> {
        u64::try_from(n).ok().and_then(|n| self.0.checked_add(n)).map(PhysAddr)
    }
}

/// Align address downwards.
///
/// Returns the greatest x with alignment `align` so that x <= addr. The alignment must be
///  a power of 2.
pub fn align_down(addr: u64, align: u64) -> u64 {
    assert!(align.is_power_of_two(), "`align` must be a power of two");
    addr & !(align - 1)
}

/// Align address upwards.
///
/// Returns the smallest x with alignment `align` so that x >= addr. The alignment must be
/// a power of 2.
pub fn align_up(addr: u64, align: u64) -> u64 {
    assert!(align.is_power_of_two(), "`align` must be a power of two");
    let align_mask = align - 1;
    if addr & align_mask == 0 {
        addr // already aligned
    } else {
        (addr | align_mask) + 1
    }
}

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

    #[test]
    pub fn test_align_up() {
        // align 1
        assert_eq!(align_up(0, 1), 0);
        assert_eq!(align_up(1234, 1), 1234);
        assert_eq!(align_up(0xffffffffffffffff, 1), 0xffffffffffffffff);
        // align 2
        assert_eq!(align_up(0, 2), 0);
        assert_eq!(align_up(1233, 2), 1234);
        assert_eq!(align_up(0xfffffffffffffffe, 2), 0xfffffffffffffffe);
        // address 0
        assert_eq!(align_up(0, 128), 0);
        assert_eq!(align_up(0, 1), 0);
        assert_eq!(align_up(0, 2), 0);
        assert_eq!(align_up(0, 0x8000000000000000), 0);
    }
}
<commit_msg>Add VirtAddr::{as_ptr, as_mut_ptr} methods<commit_after>use core::ops::{Add, AddAssign, Sub, SubAssign};
use core::iter::Step;
use core::{fmt, mem};
use core::convert::{Into, TryFrom, TryInto};

use usize_conversions::usize_from;
use bit_field::BitField;
use ux::*;

/// A canonical 64-bit virtual memory address.
///
/// This is a wrapper type around an `u64`, so it is always 8 bytes, even when compiled
/// on non 64-bit systems. The `UsizeConversions` trait can be used for performing conversions
/// between `u64` and `usize`.
///
/// On `x86_64`, only the 48 lower bits of a virtual address can be used. The top 16 bits need
/// to be copies of bit 47, i.e. the most significant bit. Addresses that fulfil this criterium
/// are called “canonical”. This type guarantees that it always represents a canonical address.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct VirtAddr(u64);

/// A 64-bit physical memory address.
///
/// This is a wrapper type around an `u64`, so it is always 8 bytes, even when compiled
/// on non 64-bit systems. The `UsizeConversions` trait can be used for performing conversions
/// between `u64` and `usize`.
///
/// On `x86_64`, only the 52 lower bits of a physical address can be used. The top 12 bits need
/// to be zero. This type guarantees that it always represents a valid physical address.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct PhysAddr(u64);

/// A passed `u64` was not a valid virtual address.
///
/// This means that bits 48 to 64 are not
/// a valid sign extension and are not null either. So automatic sign extension would have
/// overwritten possibly meaningful bits. This likely indicates a bug, for example an invalid
/// address calculation.
#[derive(Debug)]
pub struct VirtAddrNotValid(u64);

impl VirtAddr {
    /// Creates a new canonical virtual address.
    ///
    /// This function performs sign extension of bit 47 to make the address canonical. Panics
    /// if the bits in the range 48 to 64 contain data (i.e. are not null and no sign extension).
    pub fn new(addr: u64) -> VirtAddr {
        Self::try_new(addr).expect("address passed to VirtAddr::new must not contain any data \
        in bits 48 to 64")
    }

    /// Tries to create a new canonical virtual address.
    ///
    /// This function tries to performs sign
    /// extension of bit 47 to make the address canonical. It succeeds if bits 48 to 64 are
    /// either a correct sign extension (i.e. copies of bit 47) or all null. Else, an error
    /// is returned.
    pub fn try_new(addr: u64) -> Result<VirtAddr, VirtAddrNotValid> {
        match addr.get_bits(47..64) {
            0 | 0x1ffff => Ok(VirtAddr(addr)), // address is canonical
            1 => Ok(VirtAddr::new_unchecked(addr)), // address needs sign extension
            other => Err(VirtAddrNotValid(other))
        }
    }

    /// Creates a new canonical virtual address without checks.
    ///
    /// This function performs sign extension of bit 47 to make the address canonical, so
    /// bits 48 to 64 are overwritten. If you want to check that these bits contain no data,
    /// use `new` or `try_new`.
    pub fn new_unchecked(mut addr: u64) -> VirtAddr {
        if addr.get_bit(47) {
            addr.set_bits(48..64, 0xffff);
        } else {
            addr.set_bits(48..64, 0);
        }
        VirtAddr(addr)
    }

    /// Converts the address to an `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }

    pub fn as_ptr<T>(self) -> *const T {
        usize_from(self.as_u64()) as *const T
    }

    pub fn as_mut_ptr<T>(self) -> *mut T {
        self.as_ptr::<T>() as *mut T
    }

    /// Aligns the virtual address upwards to the given alignment.
    ///
    /// See the `align_up` function for more information.
    pub fn align_up<U>(self, align: U) -> Self
        where U: Into<u64>
    {
        VirtAddr(align_up(self.0, align.into()))
    }

    /// Aligns the virtual address downwards to the given alignment.
    ///
    /// See the `align_down` function for more information.
    pub fn align_down<U>(self, align: U) -> Self
        where U: Into<u64>
    {
        VirtAddr(align_down(self.0, align.into()))
    }

    /// Returns the 12-bit page offset of this virtual address.
    pub fn page_offset(&self) -> u12 {
        u12::new((self.0 & 0xfff).try_into().unwrap())
    }

    /// Returns the 9-bit level 1 page table index.
    pub fn p1_index(&self) -> u9 {
        u9::new(((self.0 >> 12) & 0o777).try_into().unwrap())
    }

    /// Returns the 9-bit level 2 page table index.
    pub fn p2_index(&self) -> u9 {
        u9::new(((self.0 >> 12 >> 9) & 0o777).try_into().unwrap())
    }

    /// Returns the 9-bit level 3 page table index.
    pub fn p3_index(&self) -> u9 {
        u9::new(((self.0 >> 12 >> 9 >> 9) & 0o777).try_into().unwrap())
    }

    /// Returns the 9-bit level 4 page table index.
    pub fn p4_index(&self) -> u9 {
        u9::new(((self.0 >> 12 >> 9 >> 9 >> 9) & 0o777).try_into().unwrap())
    }
}

impl fmt::Debug for VirtAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "VirtAddr({:#x})", self.0)
    }
}

impl Add<u64> for VirtAddr {
    type Output = Self;
    fn add(self, rhs: u64) -> Self::Output {
        VirtAddr::new(self.0 + rhs)
    }
}

impl AddAssign<u64> for VirtAddr {
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl Sub<u64> for VirtAddr {
    type Output = Self;
    fn sub(self, rhs: u64) -> Self::Output {
        VirtAddr::new(self.0 - rhs)
    }
}

impl SubAssign<u64> for VirtAddr {
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl Sub<VirtAddr> for VirtAddr {
    type Output = u64;
    fn sub(self, rhs: VirtAddr) -> Self::Output {
        self.as_u64() - rhs.as_u64()
    }
}

impl Step for VirtAddr {
    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
        if *start < *end {
            usize::try_from(end.0.get_bits(0..48) - start.0.get_bits(0..48)).ok()
        } else {
            Some(0)
        }
    }

    fn replace_one(&mut self) -> Self {
        mem::replace(self, VirtAddr(1))
    }

    fn replace_zero(&mut self) -> Self {
        mem::replace(self, VirtAddr(0))
    }

    fn add_one(&self) -> Self {
        *self + 1
    }

    fn sub_one(&self) -> Self {
        *self - 1
    }

    fn add_usize(&self, n: usize) -> Option<Self> {
        u64::try_from(n).ok().and_then(|n| self.0.checked_add(n)).map(VirtAddr)
    }
}

/// A passed `u64` was not a valid physical address.
///
/// This means that bits 52 to 64 are not were not all null.
#[derive(Debug)]
pub struct PhysAddrNotValid(u64);

impl PhysAddr {
    /// Creates a new physical address.
    ///
    /// Panics if a bit in the range 52 to 64 is set.
    pub fn new(addr: u64) -> PhysAddr {
        assert_eq!(addr.get_bits(52..64), 0,
            "physical addresses must not have any bits in the range 52 to 64 set");
        PhysAddr(addr)
    }

    /// Tries to create a new physical address.
    ///
    /// Fails if any bits in the range 52 to 64 are set.
    pub fn try_new(addr: u64) -> Result<PhysAddr, PhysAddrNotValid> {
        match addr.get_bits(52..64) {
            0 => Ok(PhysAddr(addr)), // address is valid
            other => Err(PhysAddrNotValid(other))
        }
    }

    /// Converts the address to an `u64`.
    pub fn as_u64(self) -> u64 {
        self.0
    }

    /// Convenience method for checking if a physical address is null.
    pub fn is_null(&self) -> bool {
        self.0 == 0
    }

    /// Aligns the physical address upwards to the given alignment.
    ///
    /// See the `align_up` function for more information.
    pub fn align_up<U>(self, align: U) -> Self
        where U: Into<u64>
    {
        PhysAddr(align_up(self.0, align.into()))
    }

    /// Aligns the physical address downwards to the given alignment.
    ///
    /// See the `align_down` function for more information.
    pub fn align_down<U>(self, align: U) -> Self
        where U: Into<u64>
    {
        PhysAddr(align_down(self.0, align.into()))
    }
}

impl fmt::Debug for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "PhysAddr({:#x})", self.0)
    }
}

impl fmt::Binary for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl fmt::LowerHex for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl fmt::Octal for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl fmt::UpperHex for PhysAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl Add<u64> for PhysAddr {
    type Output = Self;
    fn add(self, rhs: u64) -> Self::Output {
        PhysAddr::new(self.0 + rhs)
    }
}

impl AddAssign<u64> for PhysAddr {
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl Sub<u64> for PhysAddr {
    type Output = Self;
    fn sub(self, rhs: u64) -> Self::Output {
        PhysAddr::new(self.0 - rhs)
    }
}

impl SubAssign<u64> for PhysAddr {
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl Sub<PhysAddr> for PhysAddr {
    type Output = u64;
    fn sub(self, rhs: PhysAddr) -> Self::Output {
        self.as_u64() - rhs.as_u64()
    }
}

impl Step for PhysAddr {
    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
        if *start < *end {
            usize::try_from(end.0 - start.0).ok()
        } else {
            Some(0)
        }
    }

    fn replace_one(&mut self) -> Self {
        mem::replace(self, PhysAddr(1))
    }

    fn replace_zero(&mut self) -> Self {
        mem::replace(self, PhysAddr(0))
    }

    fn add_one(&self) -> Self {
        *self + 1
    }

    fn sub_one(&self) -> Self {
        *self - 1
    }

    fn add_usize(&self, n: usize) -> Option<Self> {
        u64::try_from(n).ok().and_then(|n| self.0.checked_add(n)).map(PhysAddr)
    }
}

/// Align address downwards.
///
/// Returns the greatest x with alignment `align` so that x <= addr. The alignment must be
///  a power of 2.
pub fn align_down(addr: u64, align: u64) -> u64 {
    assert!(align.is_power_of_two(), "`align` must be a power of two");
    addr & !(align - 1)
}

/// Align address upwards.
///
/// Returns the smallest x with alignment `align` so that x >= addr. The alignment must be
/// a power of 2.
pub fn align_up(addr: u64, align: u64) -> u64 {
    assert!(align.is_power_of_two(), "`align` must be a power of two");
    let align_mask = align - 1;
    if addr & align_mask == 0 {
        addr // already aligned
    } else {
        (addr | align_mask) + 1
    }
}

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

    #[test]
    pub fn test_align_up() {
        // align 1
        assert_eq!(align_up(0, 1), 0);
        assert_eq!(align_up(1234, 1), 1234);
        assert_eq!(align_up(0xffffffffffffffff, 1), 0xffffffffffffffff);
        // align 2
        assert_eq!(align_up(0, 2), 0);
        assert_eq!(align_up(1233, 2), 1234);
        assert_eq!(align_up(0xfffffffffffffffe, 2), 0xfffffffffffffffe);
        // address 0
        assert_eq!(align_up(0, 128), 0);
        assert_eq!(align_up(0, 1), 0);
        assert_eq!(align_up(0, 2), 0);
        assert_eq!(align_up(0, 0x8000000000000000), 0);
    }
}
<|endoftext|>