<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::*;
use rand::{Rng, SeedableRng};
use std::{
    io::{self, Write},
    num::NonZeroUsize,
    process,
};
use structopt::{clap::arg_enum, StructOpt};
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
// TODO going to remove random seed once cluster deployment supports re-run genesis
use rand::rngs::OsRng;

#[derive(Debug, StructOpt)]
#[structopt(about = "Forged in Fire")]
pub struct Options {
    /// The FILTER string is tested against the name of all tests, and only those tests whose names
    /// contain the filter are run.
    filter: Option<String>,
    #[structopt(long = "exact")]
    /// Exactly match filters rather than by substring
    filter_exact: bool,
    #[structopt(long, default_value = "1", env = "RUST_TEST_THREADS")]
    /// NO-OP: unsupported option, exists for compatibility with the default test harness
    /// Number of threads used for running tests in parallel
    test_threads: NonZeroUsize,
    #[structopt(short = "q", long)]
    /// Output minimal information
    quiet: bool,
    #[structopt(long)]
    /// NO-OP: unsupported option, exists for compatibility with the default test harness
    nocapture: bool,
    #[structopt(long)]
    /// List all tests
    list: bool,
    #[structopt(long)]
    /// List or run ignored tests
    ignored: bool,
    #[structopt(long)]
    /// Include ignored tests when listing or running tests
    include_ignored: bool,
    /// Configure formatting of output:
    ///   pretty = Print verbose output;
    ///   terse = Display one character per test;
    ///   (json is unsupported, exists for compatibility with the default test harness)
    #[structopt(long, possible_values = &Format::variants(), default_value, case_insensitive = true)]
    format: Format,
}

impl Options {
    pub fn from_args() -> Self {
        StructOpt::from_args()
    }
}

arg_enum! {
    #[derive(Debug, Eq, PartialEq)]
    pub enum Format {
        Pretty,
        Terse,
        Json,
    }
}

impl Default for Format {
    fn default() -> Self {
        Format::Pretty
    }
}

pub fn forge_main<F: Factory>(tests: ForgeConfig<'_>, factory: F, options: &Options) -> Result<()> {
    let forge = Forge::new(options, tests, factory);

    if options.list {
        forge.list()?;

        return Ok(());
    }

    match forge.run() {
        Ok(()) => Ok(()),
        Err(_) => process::exit(101), // Exit with a non-zero exit code if tests failed
    }
}

pub struct ForgeConfig<'cfg> {
    pub public_usage_tests: &'cfg [&'cfg dyn PublicUsageTest],
    pub admin_tests: &'cfg [&'cfg dyn AdminTest],
    pub network_tests: &'cfg [&'cfg dyn NetworkTest],
}

impl<'cfg> ForgeConfig<'cfg> {
    pub fn number_of_tests(&self) -> usize {
        self.public_usage_tests.len() + self.admin_tests.len() + self.network_tests.len()
    }

    pub fn all_tests(&self) -> impl Iterator<Item = &'cfg dyn Test> + 'cfg {
        self.public_usage_tests
            .iter()
            .map(|t| t as &dyn Test)
            .chain(self.admin_tests.iter().map(|t| t as &dyn Test))
            .chain(self.network_tests.iter().map(|t| t as &dyn Test))
    }
}

pub struct Forge<'cfg, F> {
    options: &'cfg Options,
    tests: ForgeConfig<'cfg>,
    factory: F,
}

impl<'cfg, F: Factory> Forge<'cfg, F> {
    pub fn new(options: &'cfg Options, tests: ForgeConfig<'cfg>, factory: F) -> Self {
        Self {
            options,
            tests,
            factory,
        }
    }

    pub fn list(&self) -> Result<()> {
        for test in self.filter_tests(self.tests.all_tests()) {
            println!("{}: test", test.name());
        }

        if self.options.format == Format::Pretty {
            println!();
            println!(
                "{} tests",
                self.filter_tests(self.tests.all_tests()).count()
            );
        }

        Ok(())
    }

    pub fn run(&self) -> Result<()> {
        let mut summary = TestSummary::new(self.filter_tests(self.tests.all_tests()).count(), 0);
        summary.write_starting_msg()?;

        let mut rng = ::rand::rngs::StdRng::from_seed(OsRng.gen());
        let mut swarm = self.factory.launch_swarm(1);

        // Run PublicUsageTests
        for test in self.filter_tests(self.tests.public_usage_tests.iter()) {
            let mut public_ctx = PublicUsageContext::new(
                CoreContext::from_rng(&mut rng),
                swarm.chain_info().into_public_info(),
            );
            let result = run_test(|| test.run(&mut public_ctx));
            summary.handle_result(test.name().to_owned(), result)?;
        }

        // Run AdminTests
        for test in self.filter_tests(self.tests.admin_tests.iter()) {
            let mut admin_ctx =
                AdminContext::new(CoreContext::from_rng(&mut rng), swarm.chain_info());
            let result = run_test(|| test.run(&mut admin_ctx));
            summary.handle_result(test.name().to_owned(), result)?;
        }

        for test in self.filter_tests(self.tests.network_tests.iter()) {
            let report = TestReport::new();
            let mut network_ctx =
                NetworkContext::new(CoreContext::from_rng(&mut rng), &mut *swarm, report);
            let result = run_test(|| test.run(&mut network_ctx));
            summary.handle_result(test.name().to_owned(), result)?;
        }

        summary.write_summary()?;

        if summary.success() {
            Ok(())
        } else {
            Err(anyhow::anyhow!("Tests Failed"))
        }
    }

    fn filter_tests<'a, T: Test, I: Iterator<Item = T> + 'a>(
        &'a self,
        tests: I,
    ) -> impl Iterator<Item = T> + 'a {
        tests
            // Filter by ignored
            .filter(
                move |test| match (self.options.include_ignored, self.options.ignored) {
                    (true, _) => true, // Don't filter anything
                    (false, true) => test.ignored(),
                    (false, false) => !test.ignored(),
                },
            )
            // Filter by test name
            .filter(move |test| {
                if let Some(filter) = &self.options.filter {
                    if self.options.filter_exact {
                        test.name() == &filter[..]
                    } else {
                        test.name().contains(&filter[..])
                    }
                } else {
                    true
                }
            })
    }
}

enum TestResult {
    Ok,
    Failed,
    FailedWithMsg(String),
}

fn run_test<F: FnOnce() -> Result<()>>(f: F) -> TestResult {
    match ::std::panic::catch_unwind(::std::panic::AssertUnwindSafe(f)) {
        Ok(Ok(())) => TestResult::Ok,
        Ok(Err(e)) => TestResult::FailedWithMsg(format!("{:?}", e)),
        Err(_) => TestResult::Failed,
    }
}

struct TestSummary {
    stdout: StandardStream,
    total: usize,
    filtered_out: usize,
    passed: usize,
    failed: Vec<String>,
}

impl TestSummary {
    fn new(total: usize, filtered_out: usize) -> Self {
        Self {
            stdout: StandardStream::stdout(ColorChoice::Auto),
            total,
            filtered_out,
            passed: 0,
            failed: Vec::new(),
        }
    }

    fn handle_result(&mut self, name: String, result: TestResult) -> io::Result<()> {
        write!(self.stdout, "test {} ... ", name)?;
        match result {
            TestResult::Ok => {
                self.passed += 1;
                self.write_ok()?;
            }
            TestResult::Failed => {
                self.failed.push(name);
                self.write_failed()?;
            }
            TestResult::FailedWithMsg(msg) => {
                self.failed.push(name);
                self.write_failed()?;
                writeln!(self.stdout)?;

                write!(self.stdout, "Error: {}", msg)?;
            }
        }
        writeln!(self.stdout)?;
        Ok(())
    }

    fn write_ok(&mut self) -> io::Result<()> {
        self.stdout
            .set_color(ColorSpec::new().set_fg(Some(Color::Green)))?;
        write!(self.stdout, "ok")?;
        self.stdout.reset()?;
        Ok(())
    }

    fn write_failed(&mut self) -> io::Result<()> {
        self.stdout
            .set_color(ColorSpec::new().set_fg(Some(Color::Red)))?;
        write!(self.stdout, "FAILED")?;
        self.stdout.reset()?;
        Ok(())
    }

    fn write_starting_msg(&mut self) -> io::Result<()> {
        writeln!(self.stdout)?;
        writeln!(
            self.stdout,
            "running {} tests",
            self.total - self.filtered_out
        )?;
        Ok(())
    }

    fn write_summary(&mut self) -> io::Result<()> {
        // Print out the failing tests
        if !self.failed.is_empty() {
            writeln!(self.stdout)?;
            writeln!(self.stdout, "failures:")?;
            for name in &self.failed {
                writeln!(self.stdout, "    {}", name)?;
            }
        }

        writeln!(self.stdout)?;
        write!(self.stdout, "test result: ")?;
        if self.failed.is_empty() {
            self.write_ok()?;
        } else {
            self.write_failed()?;
        }
        writeln!(
            self.stdout,
            ". {} passed; {} failed; {} filtered out",
            self.passed,
            self.failed.len(),
            self.filtered_out
        )?;
        writeln!(self.stdout)?;
        Ok(())
    }

    fn success(&self) -> bool {
        self.failed.is_empty()
    }
}
<commit_msg>forge: don't start a swarm if there are no tests to run<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::*;
use rand::{Rng, SeedableRng};
use std::{
    io::{self, Write},
    num::NonZeroUsize,
    process,
};
use structopt::{clap::arg_enum, StructOpt};
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
// TODO going to remove random seed once cluster deployment supports re-run genesis
use rand::rngs::OsRng;

#[derive(Debug, StructOpt)]
#[structopt(about = "Forged in Fire")]
pub struct Options {
    /// The FILTER string is tested against the name of all tests, and only those tests whose names
    /// contain the filter are run.
    filter: Option<String>,
    #[structopt(long = "exact")]
    /// Exactly match filters rather than by substring
    filter_exact: bool,
    #[structopt(long, default_value = "1", env = "RUST_TEST_THREADS")]
    /// NO-OP: unsupported option, exists for compatibility with the default test harness
    /// Number of threads used for running tests in parallel
    test_threads: NonZeroUsize,
    #[structopt(short = "q", long)]
    /// Output minimal information
    quiet: bool,
    #[structopt(long)]
    /// NO-OP: unsupported option, exists for compatibility with the default test harness
    nocapture: bool,
    #[structopt(long)]
    /// List all tests
    list: bool,
    #[structopt(long)]
    /// List or run ignored tests
    ignored: bool,
    #[structopt(long)]
    /// Include ignored tests when listing or running tests
    include_ignored: bool,
    /// Configure formatting of output:
    ///   pretty = Print verbose output;
    ///   terse = Display one character per test;
    ///   (json is unsupported, exists for compatibility with the default test harness)
    #[structopt(long, possible_values = &Format::variants(), default_value, case_insensitive = true)]
    format: Format,
}

impl Options {
    pub fn from_args() -> Self {
        StructOpt::from_args()
    }
}

arg_enum! {
    #[derive(Debug, Eq, PartialEq)]
    pub enum Format {
        Pretty,
        Terse,
        Json,
    }
}

impl Default for Format {
    fn default() -> Self {
        Format::Pretty
    }
}

pub fn forge_main<F: Factory>(tests: ForgeConfig<'_>, factory: F, options: &Options) -> Result<()> {
    let forge = Forge::new(options, tests, factory);

    if options.list {
        forge.list()?;

        return Ok(());
    }

    match forge.run() {
        Ok(()) => Ok(()),
        Err(_) => process::exit(101), // Exit with a non-zero exit code if tests failed
    }
}

pub struct ForgeConfig<'cfg> {
    pub public_usage_tests: &'cfg [&'cfg dyn PublicUsageTest],
    pub admin_tests: &'cfg [&'cfg dyn AdminTest],
    pub network_tests: &'cfg [&'cfg dyn NetworkTest],
}

impl<'cfg> ForgeConfig<'cfg> {
    pub fn number_of_tests(&self) -> usize {
        self.public_usage_tests.len() + self.admin_tests.len() + self.network_tests.len()
    }

    pub fn all_tests(&self) -> impl Iterator<Item = &'cfg dyn Test> + 'cfg {
        self.public_usage_tests
            .iter()
            .map(|t| t as &dyn Test)
            .chain(self.admin_tests.iter().map(|t| t as &dyn Test))
            .chain(self.network_tests.iter().map(|t| t as &dyn Test))
    }
}

pub struct Forge<'cfg, F> {
    options: &'cfg Options,
    tests: ForgeConfig<'cfg>,
    factory: F,
}

impl<'cfg, F: Factory> Forge<'cfg, F> {
    pub fn new(options: &'cfg Options, tests: ForgeConfig<'cfg>, factory: F) -> Self {
        Self {
            options,
            tests,
            factory,
        }
    }

    pub fn list(&self) -> Result<()> {
        for test in self.filter_tests(self.tests.all_tests()) {
            println!("{}: test", test.name());
        }

        if self.options.format == Format::Pretty {
            println!();
            println!(
                "{} tests",
                self.filter_tests(self.tests.all_tests()).count()
            );
        }

        Ok(())
    }

    pub fn run(&self) -> Result<()> {
        let test_count = self.filter_tests(self.tests.all_tests()).count();
        let filtered_out = test_count.saturating_sub(self.tests.all_tests().count());

        let mut summary = TestSummary::new(test_count, filtered_out);
        summary.write_starting_msg()?;

        if test_count > 0 {
            let mut rng = ::rand::rngs::StdRng::from_seed(OsRng.gen());
            let mut swarm = self.factory.launch_swarm(1);

            // Run PublicUsageTests
            for test in self.filter_tests(self.tests.public_usage_tests.iter()) {
                let mut public_ctx = PublicUsageContext::new(
                    CoreContext::from_rng(&mut rng),
                    swarm.chain_info().into_public_info(),
                );
                let result = run_test(|| test.run(&mut public_ctx));
                summary.handle_result(test.name().to_owned(), result)?;
            }

            // Run AdminTests
            for test in self.filter_tests(self.tests.admin_tests.iter()) {
                let mut admin_ctx =
                    AdminContext::new(CoreContext::from_rng(&mut rng), swarm.chain_info());
                let result = run_test(|| test.run(&mut admin_ctx));
                summary.handle_result(test.name().to_owned(), result)?;
            }

            for test in self.filter_tests(self.tests.network_tests.iter()) {
                let report = TestReport::new();
                let mut network_ctx =
                    NetworkContext::new(CoreContext::from_rng(&mut rng), &mut *swarm, report);
                let result = run_test(|| test.run(&mut network_ctx));
                summary.handle_result(test.name().to_owned(), result)?;
            }
        }

        summary.write_summary()?;

        if summary.success() {
            Ok(())
        } else {
            Err(anyhow::anyhow!("Tests Failed"))
        }
    }

    fn filter_tests<'a, T: Test, I: Iterator<Item = T> + 'a>(
        &'a self,
        tests: I,
    ) -> impl Iterator<Item = T> + 'a {
        tests
            // Filter by ignored
            .filter(
                move |test| match (self.options.include_ignored, self.options.ignored) {
                    (true, _) => true, // Don't filter anything
                    (false, true) => test.ignored(),
                    (false, false) => !test.ignored(),
                },
            )
            // Filter by test name
            .filter(move |test| {
                if let Some(filter) = &self.options.filter {
                    if self.options.filter_exact {
                        test.name() == &filter[..]
                    } else {
                        test.name().contains(&filter[..])
                    }
                } else {
                    true
                }
            })
    }
}

enum TestResult {
    Ok,
    Failed,
    FailedWithMsg(String),
}

fn run_test<F: FnOnce() -> Result<()>>(f: F) -> TestResult {
    match ::std::panic::catch_unwind(::std::panic::AssertUnwindSafe(f)) {
        Ok(Ok(())) => TestResult::Ok,
        Ok(Err(e)) => TestResult::FailedWithMsg(format!("{:?}", e)),
        Err(_) => TestResult::Failed,
    }
}

struct TestSummary {
    stdout: StandardStream,
    total: usize,
    filtered_out: usize,
    passed: usize,
    failed: Vec<String>,
}

impl TestSummary {
    fn new(total: usize, filtered_out: usize) -> Self {
        Self {
            stdout: StandardStream::stdout(ColorChoice::Auto),
            total,
            filtered_out,
            passed: 0,
            failed: Vec::new(),
        }
    }

    fn handle_result(&mut self, name: String, result: TestResult) -> io::Result<()> {
        write!(self.stdout, "test {} ... ", name)?;
        match result {
            TestResult::Ok => {
                self.passed += 1;
                self.write_ok()?;
            }
            TestResult::Failed => {
                self.failed.push(name);
                self.write_failed()?;
            }
            TestResult::FailedWithMsg(msg) => {
                self.failed.push(name);
                self.write_failed()?;
                writeln!(self.stdout)?;

                write!(self.stdout, "Error: {}", msg)?;
            }
        }
        writeln!(self.stdout)?;
        Ok(())
    }

    fn write_ok(&mut self) -> io::Result<()> {
        self.stdout
            .set_color(ColorSpec::new().set_fg(Some(Color::Green)))?;
        write!(self.stdout, "ok")?;
        self.stdout.reset()?;
        Ok(())
    }

    fn write_failed(&mut self) -> io::Result<()> {
        self.stdout
            .set_color(ColorSpec::new().set_fg(Some(Color::Red)))?;
        write!(self.stdout, "FAILED")?;
        self.stdout.reset()?;
        Ok(())
    }

    fn write_starting_msg(&mut self) -> io::Result<()> {
        writeln!(self.stdout)?;
        writeln!(
            self.stdout,
            "running {} tests",
            self.total - self.filtered_out
        )?;
        Ok(())
    }

    fn write_summary(&mut self) -> io::Result<()> {
        // Print out the failing tests
        if !self.failed.is_empty() {
            writeln!(self.stdout)?;
            writeln!(self.stdout, "failures:")?;
            for name in &self.failed {
                writeln!(self.stdout, "    {}", name)?;
            }
        }

        writeln!(self.stdout)?;
        write!(self.stdout, "test result: ")?;
        if self.failed.is_empty() {
            self.write_ok()?;
        } else {
            self.write_failed()?;
        }
        writeln!(
            self.stdout,
            ". {} passed; {} failed; {} filtered out",
            self.passed,
            self.failed.len(),
            self.filtered_out
        )?;
        writeln!(self.stdout)?;
        Ok(())
    }

    fn success(&self) -> bool {
        self.failed.is_empty()
    }
}
<|endoftext|>


<commit_before>

fn add(x: uint, y: uint) -> uint { ret x + y; }

fn sub(x: uint, y: uint) -> uint { ret x - y; }

fn mul(x: uint, y: uint) -> uint { ret x * y; }

fn div(x: uint, y: uint) -> uint { ret x / y; }

fn rem(x: uint, y: uint) -> uint { ret x % y; }

pure fn lt(x: uint, y: uint) -> bool { ret x < y; }

pure fn le(x: uint, y: uint) -> bool { ret x <= y; }

pure fn eq(x: uint, y: uint) -> bool { ret x == y; }

pure fn ne(x: uint, y: uint) -> bool { ret x != y; }

pure fn ge(x: uint, y: uint) -> bool { ret x >= y; }

pure fn gt(x: uint, y: uint) -> bool { ret x > y; }

fn max(x: uint, y: uint) -> uint { if x > y { ret x; } ret y; }

fn min(x: uint, y: uint) -> uint { if x > y { ret y; } ret x; }

iter range(lo: uint, hi: uint) -> uint {
    let lo_ = lo;
    while lo_ < hi { put lo_; lo_ += 1u; }
}

fn next_power_of_two(n: uint) -> uint {
    let halfbits: uint = sys::size_of::<uint>() * 4u;
    let tmp: uint = n - 1u;
    let shift: uint = 1u;
    while shift <= halfbits { tmp |= tmp >> shift; shift <<= 1u; }
    ret tmp + 1u;
}

fn parse_buf(buf: [u8], radix: uint) -> uint {
    if vec::len::<u8>(buf) == 0u {
        log_err "parse_buf(): buf is empty";
        fail;
    }
    let i = vec::len::<u8>(buf) - 1u;
    let power = 1u;
    let n = 0u;
    while true {
        n += (buf[i] - ('0' as u8) as uint) * power;
        power *= radix;
        if i == 0u { ret n; }
        i -= 1u;
    }
    fail;
}

fn from_str(s: str) -> uint { parse_buf(str::bytes(s), 10u) }

fn to_str(num: uint, radix: uint) -> str {
    let n = num;
    assert (0u < radix && radix <= 16u);
    fn digit(n: uint) -> char {
        ret alt n {
              0u { '0' }
              1u { '1' }
              2u { '2' }
              3u { '3' }
              4u { '4' }
              5u { '5' }
              6u { '6' }
              7u { '7' }
              8u { '8' }
              9u { '9' }
              10u { 'a' }
              11u { 'b' }
              12u { 'c' }
              13u { 'd' }
              14u { 'e' }
              15u { 'f' }
              _ { fail }
            };
    }
    if n == 0u { ret "0"; }
    let s: str = "";
    while n != 0u {
        s += str::unsafe_from_byte(digit(n % radix) as u8);
        n /= radix;
    }
    let s1: str = "";
    let len: uint = str::byte_len(s);
    while len != 0u { len -= 1u; s1 += str::unsafe_from_byte(s[len]); }
    ret s1;
}
fn str(i: uint) -> str { ret to_str(i, 10u); }

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C $RBUILD 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
<commit_msg>[Lib] uint.rs: +min_value<commit_after>pure fn min_value() -> uint { ret 0u; }

fn add(x: uint, y: uint) -> uint { ret x + y; }

fn sub(x: uint, y: uint) -> uint { ret x - y; }

fn mul(x: uint, y: uint) -> uint { ret x * y; }

fn div(x: uint, y: uint) -> uint { ret x / y; }

fn rem(x: uint, y: uint) -> uint { ret x % y; }

pure fn lt(x: uint, y: uint) -> bool { ret x < y; }

pure fn le(x: uint, y: uint) -> bool { ret x <= y; }

pure fn eq(x: uint, y: uint) -> bool { ret x == y; }

pure fn ne(x: uint, y: uint) -> bool { ret x != y; }

pure fn ge(x: uint, y: uint) -> bool { ret x >= y; }

pure fn gt(x: uint, y: uint) -> bool { ret x > y; }

fn max(x: uint, y: uint) -> uint { if x > y { ret x; } ret y; }

fn min(x: uint, y: uint) -> uint { if x > y { ret y; } ret x; }

iter range(lo: uint, hi: uint) -> uint {
    let lo_ = lo;
    while lo_ < hi { put lo_; lo_ += 1u; }
}

fn next_power_of_two(n: uint) -> uint {
    let halfbits: uint = sys::size_of::<uint>() * 4u;
    let tmp: uint = n - 1u;
    let shift: uint = 1u;
    while shift <= halfbits { tmp |= tmp >> shift; shift <<= 1u; }
    ret tmp + 1u;
}

fn parse_buf(buf: [u8], radix: uint) -> uint {
    if vec::len::<u8>(buf) == 0u {
        log_err "parse_buf(): buf is empty";
        fail;
    }
    let i = vec::len::<u8>(buf) - 1u;
    let power = 1u;
    let n = 0u;
    while true {
        n += (buf[i] - ('0' as u8) as uint) * power;
        power *= radix;
        if i == 0u { ret n; }
        i -= 1u;
    }
    fail;
}

fn from_str(s: str) -> uint { parse_buf(str::bytes(s), 10u) }

fn to_str(num: uint, radix: uint) -> str {
    let n = num;
    assert (0u < radix && radix <= 16u);
    fn digit(n: uint) -> char {
        ret alt n {
              0u { '0' }
              1u { '1' }
              2u { '2' }
              3u { '3' }
              4u { '4' }
              5u { '5' }
              6u { '6' }
              7u { '7' }
              8u { '8' }
              9u { '9' }
              10u { 'a' }
              11u { 'b' }
              12u { 'c' }
              13u { 'd' }
              14u { 'e' }
              15u { 'f' }
              _ { fail }
            };
    }
    if n == 0u { ret "0"; }
    let s: str = "";
    while n != 0u {
        s += str::unsafe_from_byte(digit(n % radix) as u8);
        n /= radix;
    }
    let s1: str = "";
    let len: uint = str::byte_len(s);
    while len != 0u { len -= 1u; s1 += str::unsafe_from_byte(s[len]); }
    ret s1;
}
fn str(i: uint) -> str { ret to_str(i, 10u); }

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C $RBUILD 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
<|endoftext|>


<commit_before>//!
//! Generic addresses
//!

use std::{fmt, mem};
use std::ops::{Not, BitAnd, BitOr, BitXor};

/// A trait for all addresses
pub trait Address: Copy + Ord + Eq + Not<Output=Self> + BitAnd<Output=Self> + BitOr<Output=Self> + BitXor<Output=Self> + fmt::UpperHex {
    /// Type of address offset
    type Offset;

    /// The zero address
    /// TODO: Use std::num::Zero instead when it becomes stable
    fn zero () -> Self;

    /// The succeeding address (wrapping)
    fn next (self) -> Self;

    /// Calculate new address with given offset (wrapping)
    fn offset (self, offset: Self::Offset) -> Self;

    /// Calculate new address with given offset (wrapping) while protecting the masked part of
    /// the address
    fn offset_masked (&self, offset: Self::Offset, mask: Self) -> Self {
        (*self & mask) | (self.offset(offset) & !mask)
    }

    /// Return an iterator for getting successive addresses (wrapping)
    fn successive (self) -> Iter<Self> {
        Iter { addr: self }
    }

    /// Return an iterator for getting successive addresses until the given last address
    fn upto (self, last_addr: Self) -> UpTo<Self> {
        UpTo { it: self.successive(), last_addr: last_addr, flag: false }
    }

    /// Return an object for displaying the address
    fn display (&self) -> Display<Self> {
        Display { addr: self }
    }
}

macro_rules! impl_address {
    ($addr_type:ty, $offset_type:ty) => (
        impl Address for $addr_type {
            type Offset = $offset_type;

            fn zero () -> $addr_type { 0 }

            fn next (self) -> $addr_type { self.wrapping_add(1) }

            fn offset (self, offset: $offset_type) -> $addr_type {
                if offset < 0 {
                    self.wrapping_sub((-offset as $addr_type))
                } else {
                    self.wrapping_add(offset as $addr_type)
                }
            }
        }
    );
}

impl_address!( u8,  i8);
impl_address!(u16, i16);
impl_address!(u32, i32);

/// Iterator for getting successive addresses
pub struct Iter<A> {
    addr: A,
}

impl<A: Address> Iterator for Iter<A> {
    type Item = A;

    fn next (&mut self) -> Option<A> {
        let addr = self.addr;
        self.addr = self.addr.next();
        Some(addr)
    }
}

/// Iterator for getting successive addresses until a given last address
pub struct UpTo<A> {
    it: Iter<A>,
    last_addr: A,
    flag: bool,
}

impl<A: Address> Iterator for UpTo<A> {
    type Item = A;

    fn next (&mut self) -> Option<A> {
        if self.flag {
            None
        } else {
            self.it.next().and_then(|x| {
                if x == self.last_addr { self.flag = true; }
                Some(x)
            })
        }
    }
}

/// Helper struct for displaying an address
pub struct Display<'a, A: 'a> {
    addr: &'a A,
}

impl<'a, A: Address> fmt::Display for Display<'a, A> {
    fn fmt (&self, f: &mut fmt::Formatter) -> fmt::Result {
        match mem::size_of::<A>() {
            1 => write!(f, "${:02X}", self.addr),
            2 => write!(f, "${:04X}", self.addr),
            4 => write!(f, "${:08X}", self.addr),
            _ => write!(f, "${:X}", self.addr),
        }
    }
}

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

    #[test]
    fn zero () {
        assert_eq!(u16::zero(), 0x0000);
    }

    #[test]
    fn next () {
        assert_eq!(0x1234_u16.next(), 0x1235);
        assert_eq!(0xffff_u16.next(), 0x0000);
    }

    #[test]
    fn offset () {
        assert_eq!(0x1234_u16.offset( 5), 0x1239);
        assert_eq!(0x1234_u16.offset(-3), 0x1231);
    }

    #[test]
    fn offset_wrapping () {
        assert_eq!(0xffff_u16.offset( 1), 0x0000);
        assert_eq!(0x0000_u16.offset(-1), 0xffff);
    }

    #[test]
    fn offset_masked () {
        assert_eq!(0x12ff_u16.offset_masked( 1, 0x0000), 0x1300);
        assert_eq!(0x12ff_u16.offset_masked( 1, 0xff00), 0x1200);
        assert_eq!(0x1300_u16.offset_masked(-1, 0x0000), 0x12ff);
        assert_eq!(0x1300_u16.offset_masked(-1, 0xff00), 0x13ff);
    }

    #[test]
    fn iterating () {
        let mut it = 0xfffe_u16.successive();
        assert_eq!(it.next(), Some(0xfffe));
        assert_eq!(it.next(), Some(0xffff));
        assert_eq!(it.next(), Some(0x0000));
        assert_eq!(it.next(), Some(0x0001));
    }

    #[test]
    fn bounded_iterating () {
        let mut it = 0xfffe_u16.upto(0xffff);
        assert_eq!(it.next(), Some(0xfffe));
        assert_eq!(it.next(), Some(0xffff));
        assert_eq!(it.next(), None);
        assert_eq!(it.next(), None);
    }

    #[test]
    fn displaying () {
        assert_eq!(format!("{}", 0x0f_u8.display()), "$0F");
        assert_eq!(format!("{}", 0x01ff_u16.display()), "$01FF");
    }
}
<commit_msg>Allow conversion of addresses from and to system integers<commit_after>//!
//! Generic addresses
//!

use std::{fmt, mem};
use std::ops::{Not, BitAnd, BitOr, BitXor};

/// A trait for all addresses
pub trait Address: Copy + Ord + Eq + Not<Output=Self> + BitAnd<Output=Self> + BitOr<Output=Self> + BitXor<Output=Self> + fmt::UpperHex {
    /// Type of address offset
    type Offset;

    /// The zero address
    /// TODO: Use std::num::Zero instead when it becomes stable
    fn zero () -> Self;

    /// The succeeding address (wrapping)
    fn next (self) -> Self;

    /// Calculate new address with given offset (wrapping)
    fn offset (self, offset: Self::Offset) -> Self;

    /// Calculate new address with given offset (wrapping) while protecting the masked part of
    /// the address
    fn offset_masked (&self, offset: Self::Offset, mask: Self) -> Self {
        (*self & mask) | (self.offset(offset) & !mask)
    }

    /// Return an iterator for getting successive addresses (wrapping)
    fn successive (self) -> Iter<Self> {
        Iter { addr: self }
    }

    /// Return an iterator for getting successive addresses until the given last address
    fn upto (self, last_addr: Self) -> UpTo<Self> {
        UpTo { it: self.successive(), last_addr: last_addr, flag: false }
    }

    /// Return an object for displaying the address
    fn display (&self) -> Display<Self> {
        Display { addr: self }
    }

    /// Return the address for a host system integer. This should normally not be used and only
    /// exists to provide a way for memory implementations to easily map from system integer
    /// indeces. This method is therefore marked as unsafe to prevent usage in normal contexts.
    unsafe fn from_usize (i: usize) -> Self;

    /// Return the address as a host system integer. This should normally not be used and only
    /// exists to provide a way for memory implementations to easily map to system integer
    /// indeces. This method is therefore marked as unsafe to prevent usage in normal contexts.
    unsafe fn to_usize (self) -> usize;
}

macro_rules! impl_address {
    ($addr_type:ty, $offset_type:ty) => (
        impl Address for $addr_type {
            type Offset = $offset_type;

            fn zero () -> $addr_type { 0 }

            fn next (self) -> $addr_type { self.wrapping_add(1) }

            fn offset (self, offset: $offset_type) -> $addr_type {
                if offset < 0 {
                    self.wrapping_sub((-offset as $addr_type))
                } else {
                    self.wrapping_add(offset as $addr_type)
                }
            }

            unsafe fn from_usize (i: usize) -> $addr_type { i as $addr_type }

            unsafe fn to_usize (self) -> usize { self as usize }
        }
    );
}

impl_address!( u8,  i8);
impl_address!(u16, i16);
impl_address!(u32, i32);

/// Iterator for getting successive addresses
pub struct Iter<A> {
    addr: A,
}

impl<A: Address> Iterator for Iter<A> {
    type Item = A;

    fn next (&mut self) -> Option<A> {
        let addr = self.addr;
        self.addr = self.addr.next();
        Some(addr)
    }
}

/// Iterator for getting successive addresses until a given last address
pub struct UpTo<A> {
    it: Iter<A>,
    last_addr: A,
    flag: bool,
}

impl<A: Address> Iterator for UpTo<A> {
    type Item = A;

    fn next (&mut self) -> Option<A> {
        if self.flag {
            None
        } else {
            self.it.next().and_then(|x| {
                if x == self.last_addr { self.flag = true; }
                Some(x)
            })
        }
    }
}

/// Helper struct for displaying an address
pub struct Display<'a, A: 'a> {
    addr: &'a A,
}

impl<'a, A: Address> fmt::Display for Display<'a, A> {
    fn fmt (&self, f: &mut fmt::Formatter) -> fmt::Result {
        match mem::size_of::<A>() {
            1 => write!(f, "${:02X}", self.addr),
            2 => write!(f, "${:04X}", self.addr),
            4 => write!(f, "${:08X}", self.addr),
            _ => write!(f, "${:X}", self.addr),
        }
    }
}

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

    #[test]
    fn zero () {
        assert_eq!(u16::zero(), 0x0000);
    }

    #[test]
    fn next () {
        assert_eq!(0x1234_u16.next(), 0x1235);
        assert_eq!(0xffff_u16.next(), 0x0000);
    }

    #[test]
    fn offset () {
        assert_eq!(0x1234_u16.offset( 5), 0x1239);
        assert_eq!(0x1234_u16.offset(-3), 0x1231);
    }

    #[test]
    fn offset_wrapping () {
        assert_eq!(0xffff_u16.offset( 1), 0x0000);
        assert_eq!(0x0000_u16.offset(-1), 0xffff);
    }

    #[test]
    fn offset_masked () {
        assert_eq!(0x12ff_u16.offset_masked( 1, 0x0000), 0x1300);
        assert_eq!(0x12ff_u16.offset_masked( 1, 0xff00), 0x1200);
        assert_eq!(0x1300_u16.offset_masked(-1, 0x0000), 0x12ff);
        assert_eq!(0x1300_u16.offset_masked(-1, 0xff00), 0x13ff);
    }

    #[test]
    fn iterating () {
        let mut it = 0xfffe_u16.successive();
        assert_eq!(it.next(), Some(0xfffe));
        assert_eq!(it.next(), Some(0xffff));
        assert_eq!(it.next(), Some(0x0000));
        assert_eq!(it.next(), Some(0x0001));
    }

    #[test]
    fn bounded_iterating () {
        let mut it = 0xfffe_u16.upto(0xffff);
        assert_eq!(it.next(), Some(0xfffe));
        assert_eq!(it.next(), Some(0xffff));
        assert_eq!(it.next(), None);
        assert_eq!(it.next(), None);
    }

    #[test]
    fn displaying () {
        assert_eq!(format!("{}", 0x0f_u8.display()), "$0F");
        assert_eq!(format!("{}", 0x01ff_u16.display()), "$01FF");
    }

    #[test]
    fn usize_conversion () {
        assert_eq!(0x1234_usize, unsafe { 0x1234_u16.to_usize() });
        assert_eq!(0x1234_u16,   unsafe { Address::from_usize(0x1234) });
    }
}
<|endoftext|>


<commit_before>use core::ptr;
use spin::Mutex;

use paging::PhysicalAddress;

// This MUST be kept in sync with the one from src/asm_routines and plan.md
const BOOT_TMP_MMAP_BUFFER:     usize   = 0x2000;

pub const MEM_PAGE_SIZE_BYTES:      usize   = 0x1_000; // 4096
pub const MEM_PAGE_MAP_SIZE_BYTES:  usize   = 0x10_000;
pub const MEM_PAGE_MAP1_ADDRESS:    usize   = 0x70_000;
pub const MEM_PAGE_MAP2_ADDRESS:    usize   = 0x80_000;
pub const KERNEL_LOCATION:          usize   = 0x100_000;
pub const KERNEL_SIZE_LIMIT:        usize   = 200*0x200;
pub const MEMORY_RESERVED_BELOW:    usize   = KERNEL_LOCATION+KERNEL_SIZE_LIMIT;

// Memory frame (single allocation unit)
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Frame {
    pub index: usize
}
impl Frame {
    /// Create new Frame from memory address. Rounds down.
    pub fn containing_address(address: usize) -> Frame {
        Frame {index: address / MEM_PAGE_SIZE_BYTES}
    }
    /// Create new Frame from memory map index.
    pub fn from_index(index: usize) -> Frame {
        Frame {index: index}
    }
    /// Start address of the frame
    pub fn start_address(&self) -> PhysicalAddress {
        self.index * MEM_PAGE_SIZE_BYTES
    }
    /// Clone frame
    pub fn clone(&self) -> Frame {
        Frame { index: self.index }
    }
    /// Inclusive range between frames
    pub fn range_inclusive(start: Frame, end: Frame) -> FrameIter {
        FrameIter {
            start: start,
            end: end,
        }
    }
}

/// Iterator for frames
pub struct FrameIter {
    start: Frame,
    end: Frame,
}
impl Iterator for FrameIter {
    type Item = Frame;

    fn next(&mut self) -> Option<Frame> {
        if self.start <= self.end {
            let frame = self.start.clone();
            self.start.index += 1;
            Some(frame)
        } else {
            None
        }
    }
}

// Frame allocators all work this way
pub trait FrameAllocator {
    fn allocate_frame(&mut self) -> Option<Frame>;
    fn deallocate_frame(&mut self, frame: Frame);
}

// A simple first-fit frame allocator
// Currently we can only get one frame at a time
pub struct BitmapAllocator {
    index: usize
}

impl BitmapAllocator {
    pub fn new() -> BitmapAllocator {
        BitmapAllocator {index: 0}
    }
    fn is_free(&self, index: usize) -> bool {
        let free    = unsafe {ptr::read_volatile((MEM_PAGE_MAP1_ADDRESS + index/8) as *mut u8)} & (1 << (index%8)) != 0; // 1: free, 0: reserved
        let usable  = unsafe {ptr::read_volatile((MEM_PAGE_MAP2_ADDRESS + index/8) as *mut u8)} & (1 << (index%8)) != 0; // 1: usable, 0: unusable

        if free && !usable { // error in free pages table
            panic!("PHYS_MEM: Page {} is incorrectly marked as free.", index);
        }
        free
    }
}
impl FrameAllocator for BitmapAllocator {
    fn allocate_frame(&mut self) -> Option<Frame> {
        // Find first free block
        for i in 0..(MEM_PAGE_MAP_SIZE_BYTES*8) {
            if self.is_free(i) {
                unsafe {
                    // set frame reserved
                    ptr::write_volatile(
                        (MEM_PAGE_MAP1_ADDRESS + i/8) as *mut u8,
                        ptr::read_volatile((MEM_PAGE_MAP1_ADDRESS + i/8) as *mut u8) ^ (1 << (i%8))
                    );
                }
                return Some(Frame::from_index(i));
            }
        }
        // We could not find a free memory block
        None
    }
    fn deallocate_frame(&mut self, frame: Frame) {
        if self.is_free(frame.index) {
            panic!("PHYS_MEM: deallocate_frame: Page {} is already free.", frame.index);
        }
        unsafe {
            // set the correct bit
            ptr::write_volatile(
                (MEM_PAGE_MAP1_ADDRESS + frame.index/8) as *mut u8,
                ptr::read_volatile((MEM_PAGE_MAP1_ADDRESS + frame.index/8) as *mut u8) | (1 << (frame.index%8))
            );
        }
    }
}

fn mt_align_address(address: usize, upwards: bool) -> usize {
    if address % MEM_PAGE_SIZE_BYTES == 0 {
        address
    }
    else if upwards {
        address + MEM_PAGE_SIZE_BYTES - address % MEM_PAGE_SIZE_BYTES
    }
    else {
        address - address % MEM_PAGE_SIZE_BYTES
    }
}

pub fn create_memory_bitmap() {
    // load memory map from where out bootloader left it
    // http://wiki.osdev.org/Detecting_Memory_(x86)#BIOS_Function:_INT_0x15.2C_EAX_.3D_0xE820

    rprintln!("BITM: 1");

    // zero out the bitmap sections
    for address in (MEM_PAGE_MAP1_ADDRESS..MEM_PAGE_MAP2_ADDRESS+MEM_PAGE_SIZE_BYTES).step_by(8) {
        rprintln!("{:#x}", address);
        unsafe {
            ptr::write_volatile(address as *mut u8, 0); // default to (reserved, unusable)
        }
    }

    rprintln!("BITM: 2");

    let entry_count: u8 = unsafe {ptr::read_volatile(BOOT_TMP_MMAP_BUFFER as *mut u8)};
    let base = (BOOT_TMP_MMAP_BUFFER+2) as *mut u8;
    let mut memory_amount_counter_kib = 0;
    for index in 0..(entry_count as isize) {
        let entry_start:    usize   = unsafe { ptr::read_volatile(base.offset(24*index+ 0) as *mut u64) } as usize;
        let entry_size:     usize   = unsafe { ptr::read_volatile(base.offset(24*index+ 8) as *mut u64) } as usize;
        let entry_type:     u32     = unsafe { ptr::read_volatile(base.offset(24*index+16) as *mut u32) };
        let acpi_data:      u32     = unsafe { ptr::read_volatile(base.offset(24*index+20) as *mut u32) };
        // rprintln!("Section {}: {:#016x}-{:#016x}: type: {:#x}, acpi: {:#x}", index, entry_start, entry_start+entry_size, entry_type, acpi_data);

        // is this usable?
        // Types 1, 4 ok to use and acpi_data bit 0 must be set
        if (entry_type == 1 || entry_type == 4) && (acpi_data & 1) == 1 {
            // set frame data. accept only full frames
            for address in (mt_align_address(entry_start, true)..mt_align_address(entry_start+entry_size, false)).step_by(MEM_PAGE_SIZE_BYTES) {
                memory_amount_counter_kib += 1;
                if address/MEM_PAGE_SIZE_BYTES > MEM_PAGE_MAP_SIZE_BYTES*8 {
                    // Page table is full.
                    break;
                }
                if address < MEMORY_RESERVED_BELOW {
                    // these are permanently reserved for the kernel
                    continue;
                }
                unsafe {
                    // set free
                    ptr::write_volatile(
                        (MEM_PAGE_MAP1_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8,
                        ptr::read_volatile((MEM_PAGE_MAP1_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8) | (1 << (address%8))
                    );
                    // set usable
                    ptr::write_volatile(
                        (MEM_PAGE_MAP2_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8,
                        ptr::read_volatile((MEM_PAGE_MAP2_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8) | (1 << (address%8))
                    );
                }
            }
        }
    }
    rprintln!("Memory size {} MiB", memory_amount_counter_kib/1024);
}

// Create static pointer mutex with spinlock to make ALLOCATOR thread-safe
// pub static ALLOCATOR: Mutex<BitmapAllocator> = Mutex::new(BitmapAllocator {});

// XXX / FIXME / TODO: We should be using mutexed allocator.
macro_rules! ALLOCATOR {
    () => ({
        $crate::mem_map::BitmapAllocator::new()
    });
}
<commit_msg>Cleanup<commit_after>use core::ptr;
use spin::Mutex;

use paging::PhysicalAddress;

// This MUST be kept in sync with the one from src/asm_routines and plan.md
const BOOT_TMP_MMAP_BUFFER:     usize   = 0x2000;

pub const MEM_PAGE_SIZE_BYTES:      usize   = 0x1_000; // 4096
pub const MEM_PAGE_MAP_SIZE_BYTES:  usize   = 0x10_000;
pub const MEM_PAGE_MAP1_ADDRESS:    usize   = 0x70_000;
pub const MEM_PAGE_MAP2_ADDRESS:    usize   = 0x80_000;
pub const KERNEL_LOCATION:          usize   = 0x100_000;
pub const KERNEL_SIZE_LIMIT:        usize   = 200*0x200;
pub const MEMORY_RESERVED_BELOW:    usize   = KERNEL_LOCATION+KERNEL_SIZE_LIMIT;

// Memory frame (single allocation unit)
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Frame {
    pub index: usize
}
impl Frame {
    /// Create new Frame from memory address. Rounds down.
    pub fn containing_address(address: usize) -> Frame {
        Frame {index: address / MEM_PAGE_SIZE_BYTES}
    }
    /// Create new Frame from memory map index.
    pub fn from_index(index: usize) -> Frame {
        Frame {index: index}
    }
    /// Start address of the frame
    pub fn start_address(&self) -> PhysicalAddress {
        self.index * MEM_PAGE_SIZE_BYTES
    }
    /// Clone frame
    pub fn clone(&self) -> Frame {
        Frame { index: self.index }
    }
    /// Inclusive range between frames
    pub fn range_inclusive(start: Frame, end: Frame) -> FrameIter {
        FrameIter {
            start: start,
            end: end,
        }
    }
}

/// Iterator for frames
pub struct FrameIter {
    start: Frame,
    end: Frame,
}
impl Iterator for FrameIter {
    type Item = Frame;

    fn next(&mut self) -> Option<Frame> {
        if self.start <= self.end {
            let frame = self.start.clone();
            self.start.index += 1;
            Some(frame)
        } else {
            None
        }
    }
}

// Frame allocators all work this way
pub trait FrameAllocator {
    fn allocate_frame(&mut self) -> Option<Frame>;
    fn deallocate_frame(&mut self, frame: Frame);
}

// A simple first-fit frame allocator
// Currently we can only get one frame at a time
pub struct BitmapAllocator {
    index: usize
}

impl BitmapAllocator {
    pub fn new() -> BitmapAllocator {
        BitmapAllocator {index: 0}
    }
    fn is_free(&self, index: usize) -> bool {
        let free    = unsafe {ptr::read_volatile((MEM_PAGE_MAP1_ADDRESS + index/8) as *mut u8)} & (1 << (index%8)) != 0; // 1: free, 0: reserved
        let usable  = unsafe {ptr::read_volatile((MEM_PAGE_MAP2_ADDRESS + index/8) as *mut u8)} & (1 << (index%8)) != 0; // 1: usable, 0: unusable

        if free && !usable { // error in free pages table
            panic!("PHYS_MEM: Page {} is incorrectly marked as free.", index);
        }
        free
    }
}
impl FrameAllocator for BitmapAllocator {
    fn allocate_frame(&mut self) -> Option<Frame> {
        // Find first free block
        for i in 0..(MEM_PAGE_MAP_SIZE_BYTES*8) {
            if self.is_free(i) {
                unsafe {
                    // set frame reserved
                    ptr::write_volatile(
                        (MEM_PAGE_MAP1_ADDRESS + i/8) as *mut u8,
                        ptr::read_volatile((MEM_PAGE_MAP1_ADDRESS + i/8) as *mut u8) ^ (1 << (i%8))
                    );
                }
                return Some(Frame::from_index(i));
            }
        }
        // We could not find a free memory block
        None
    }
    fn deallocate_frame(&mut self, frame: Frame) {
        if self.is_free(frame.index) {
            panic!("PHYS_MEM: deallocate_frame: Page {} is already free.", frame.index);
        }
        unsafe {
            // set the correct bit
            ptr::write_volatile(
                (MEM_PAGE_MAP1_ADDRESS + frame.index/8) as *mut u8,
                ptr::read_volatile((MEM_PAGE_MAP1_ADDRESS + frame.index/8) as *mut u8) | (1 << (frame.index%8))
            );
        }
    }
}

fn mt_align_address(address: usize, upwards: bool) -> usize {
    if address % MEM_PAGE_SIZE_BYTES == 0 {
        address
    }
    else if upwards {
        address + MEM_PAGE_SIZE_BYTES - address % MEM_PAGE_SIZE_BYTES
    }
    else {
        address - address % MEM_PAGE_SIZE_BYTES
    }
}

pub fn create_memory_bitmap() {
    // load memory map from where out bootloader left it
    // http://wiki.osdev.org/Detecting_Memory_(x86)#BIOS_Function:_INT_0x15.2C_EAX_.3D_0xE820

    // zero out the bitmap sections
    for address in (MEM_PAGE_MAP1_ADDRESS..MEM_PAGE_MAP2_ADDRESS+MEM_PAGE_SIZE_BYTES).step_by(8) {
        unsafe {
            ptr::write_volatile(address as *mut u8, 0); // default to (reserved, unusable)
        }
    }

    let entry_count: u8 = unsafe {ptr::read_volatile(BOOT_TMP_MMAP_BUFFER as *mut u8)};
    let base = (BOOT_TMP_MMAP_BUFFER+2) as *mut u8;
    let mut memory_amount_counter_kib = 0;
    for index in 0..(entry_count as isize) {
        let entry_start:    usize   = unsafe { ptr::read_volatile(base.offset(24*index+ 0) as *mut u64) } as usize;
        let entry_size:     usize   = unsafe { ptr::read_volatile(base.offset(24*index+ 8) as *mut u64) } as usize;
        let entry_type:     u32     = unsafe { ptr::read_volatile(base.offset(24*index+16) as *mut u32) };
        let acpi_data:      u32     = unsafe { ptr::read_volatile(base.offset(24*index+20) as *mut u32) };
        // rprintln!("Section {}: {:#016x}-{:#016x}: type: {:#x}, acpi: {:#x}", index, entry_start, entry_start+entry_size, entry_type, acpi_data);

        // is this usable?
        // Types 1, 4 ok to use and acpi_data bit 0 must be set
        if (entry_type == 1 || entry_type == 4) && (acpi_data & 1) == 1 {
            // set frame data. accept only full frames
            for address in (mt_align_address(entry_start, true)..mt_align_address(entry_start+entry_size, false)).step_by(MEM_PAGE_SIZE_BYTES) {
                memory_amount_counter_kib += 1;
                if address/MEM_PAGE_SIZE_BYTES > MEM_PAGE_MAP_SIZE_BYTES*8 {
                    // Page table is full.
                    break;
                }
                if address < MEMORY_RESERVED_BELOW {
                    // these are permanently reserved for the kernel
                    continue;
                }
                unsafe {
                    // set free
                    ptr::write_volatile(
                        (MEM_PAGE_MAP1_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8,
                        ptr::read_volatile((MEM_PAGE_MAP1_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8) | (1 << (address%8))
                    );
                    // set usable
                    ptr::write_volatile(
                        (MEM_PAGE_MAP2_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8,
                        ptr::read_volatile((MEM_PAGE_MAP2_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8) | (1 << (address%8))
                    );
                }
            }
        }
    }
    rprintln!("Memory size {} MiB", memory_amount_counter_kib/1024);
}

// Create static pointer mutex with spinlock to make ALLOCATOR thread-safe
// pub static ALLOCATOR: Mutex<BitmapAllocator> = Mutex::new(BitmapAllocator {});

// XXX / FIXME / TODO: We should be using mutexed allocator.
macro_rules! ALLOCATOR {
    () => ({
        $crate::mem_map::BitmapAllocator::new()
    });
}
<|endoftext|>


<commit_before>use spin::Mutex;

use paging::PhysicalAddress;

// This MUST be kept in sync with the one from src/asm_routines and plan.md
const BOOT_TMP_MMAP_BUFFER:     usize   = 0x2000;

pub const MEM_PAGE_SIZE_BYTES:      usize   = 0x1_000; // 4096
pub const MEM_PAGE_MAP_SIZE_BYTES:  usize   = 0x10_000;
pub const MEM_PAGE_MAP1_ADDRESS:    usize   = 0x60_000;
pub const MEM_PAGE_MAP2_ADDRESS:    usize   = 0x70_000;
pub const KERNEL_LOCATION:          usize   = 0x100_000;
pub const KERNEL_SIZE_LIMIT:        usize   = 200*0x200;
pub const MEMORY_RESERVED_BELOW:    usize   = KERNEL_LOCATION+KERNEL_SIZE_LIMIT;

// Memory frame (single allocation unit)
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Frame {
    pub index: usize
}
impl Frame {
    /// Create new Frame from memory address. Rounds down.
    pub fn containing_address(address: usize) -> Frame {
        Frame {index: address / MEM_PAGE_SIZE_BYTES}
    }
    /// Create new Frame from memory map index.
    pub fn from_index(index: usize) -> Frame {
        Frame {index: index}
    }
    /// Start address of the frame
    pub fn start_address(&self) -> PhysicalAddress {
        self.index * MEM_PAGE_SIZE_BYTES
    }
    /// Clone frame
    pub fn clone(&self) -> Frame {
        Frame { index: self.index }
    }
    /// Inclusive range between frames
    pub fn range_inclusive(start: Frame, end: Frame) -> FrameIter {
        FrameIter {
            start: start,
            end: end,
        }
    }
}

/// Iterator for frames
pub struct FrameIter {
    start: Frame,
    end: Frame,
}
impl Iterator for FrameIter {
    type Item = Frame;

    fn next(&mut self) -> Option<Frame> {
        if self.start <= self.end {
            let frame = self.start.clone();
            self.start.index += 1;
            Some(frame)
        } else {
            None
        }
    }
}

// Frame allocators all work this way
pub trait FrameAllocator {
    fn allocate_frame(&mut self) -> Option<Frame>;
    fn deallocate_frame(&mut self, frame: Frame);
}

// A simple first-fit frame allocator
// Currently we can only get one frame at a time
pub struct BitmapAllocator {
    index: usize
}

impl BitmapAllocator {
    pub fn new() -> BitmapAllocator {
        BitmapAllocator {index: 0}
    }
    fn is_free(&self, index: usize) -> bool {
        let free    = unsafe { *((MEM_PAGE_MAP1_ADDRESS + index/8) as *mut u8) } & (1 << (index%8)) != 0; // 1: free, 0: reserved
        let usable  = unsafe { *((MEM_PAGE_MAP2_ADDRESS + index/8) as *mut u8) } & (1 << (index%8)) != 0; // 1: usable, 0: unusable
        if free && !usable { // error in free pages table
            panic!("PHYS_MEM: Page {} is incorrectly marked as free.", index);
        }
        free
    }
}
impl FrameAllocator for BitmapAllocator {
    fn allocate_frame(&mut self) -> Option<Frame> {
        // Find first free block
        for i in 0..(MEM_PAGE_MAP_SIZE_BYTES*8) {
            if self.is_free(i) {
                unsafe {
                    *((MEM_PAGE_MAP1_ADDRESS + i/8) as *mut u8) ^= 1 << (i%8); // set frame reserved
                }
                return Some(Frame::from_index(i));
            }
        }
        // We could not find a free memory block
        None
    }
    fn deallocate_frame(&mut self, frame: Frame) {
        if self.is_free(frame.index) {
            panic!("PHYS_MEM: deallocate_frame: Page {} is already free.", frame.index);
        }
        unsafe {
            *((MEM_PAGE_MAP1_ADDRESS + frame.index/8) as *mut u8) |= 1 << (frame.index%8); // set the correct bit
        }
    }
}

fn mt_align_address(address: usize, upwards: bool) -> usize {
    if address % MEM_PAGE_SIZE_BYTES == 0 {
        address
    }
    else if upwards {
        address + MEM_PAGE_SIZE_BYTES - address % MEM_PAGE_SIZE_BYTES
    }
    else {
        address - address % MEM_PAGE_SIZE_BYTES
    }
}

pub fn create_memory_bitmap() {
    // load memory map from where out bootloader left it
    // http://wiki.osdev.org/Detecting_Memory_(x86)#BIOS_Function:_INT_0x15.2C_EAX_.3D_0xE820

    // zero out the bitmap sections
    for address in (MEM_PAGE_MAP1_ADDRESS..MEM_PAGE_MAP2_ADDRESS+MEM_PAGE_SIZE_BYTES).step_by(8) {
        unsafe {
            *(address as *mut u8) = 0; // default to: reserved, unusable
        }
    }

    let entry_count: u8 = unsafe {*((BOOT_TMP_MMAP_BUFFER) as *mut u8)};
    let base = (BOOT_TMP_MMAP_BUFFER+2) as *mut u8;
    let mut memory_amount_counter_kib = 0;
    for index in 0..(entry_count as isize) {
        let entry_start:    usize   = unsafe { *(base.offset(24*index+ 0) as *mut u64) } as usize;
        let entry_size:     usize   = unsafe { *(base.offset(24*index+ 8) as *mut u64) } as usize;
        let entry_type:     u32     = unsafe { *(base.offset(24*index+16) as *mut u32) };
        let acpi_data:      u32     = unsafe { *(base.offset(24*index+20) as *mut u32) };
        // rprintln!("Section {}: {:#016x}-{:#016x}: type: {:#x}, acpi: {:#x}", index, entry_start, entry_start+entry_size, entry_type, acpi_data);

        // is this usable?
        // Types 1, 4 ok to use and acpi_data bit 0 must be set
        if (entry_type == 1 || entry_type == 4) && (acpi_data & 1) == 1 {
            // set frame data. accept only full frames
            for address in (mt_align_address(entry_start, true)..mt_align_address(entry_start+entry_size, false)).step_by(MEM_PAGE_SIZE_BYTES) {
                memory_amount_counter_kib += 1;
                if address/MEM_PAGE_SIZE_BYTES > MEM_PAGE_MAP_SIZE_BYTES*8 {
                    // Page table is full.
                    break;
                }
                if address < MEMORY_RESERVED_BELOW {
                    // these are permanently reserved for the kernel
                    continue;
                }
                unsafe {
                    *((MEM_PAGE_MAP1_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8) |= 1 << (address%8); // set free
                    *((MEM_PAGE_MAP2_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8) |= 1 << (address%8); // set usable
                }
            }
        }
    }
    rprintln!("Memory size {} MiB", memory_amount_counter_kib/1024);
}

// Create static pointer mutex with spinlock to make ALLOCATOR thread-safe
// pub static ALLOCATOR: Mutex<BitmapAllocator> = Mutex::new(BitmapAllocator {});

macro_rules! ALLOCATOR {
    () => ({
        $crate::mem_map::BitmapAllocator::new()
    });
}
<commit_msg>Use volatile reads/writes<commit_after>use core::ptr;
use spin::Mutex;

use paging::PhysicalAddress;

// This MUST be kept in sync with the one from src/asm_routines and plan.md
const BOOT_TMP_MMAP_BUFFER:     usize   = 0x2000;

pub const MEM_PAGE_SIZE_BYTES:      usize   = 0x1_000; // 4096
pub const MEM_PAGE_MAP_SIZE_BYTES:  usize   = 0x10_000;
pub const MEM_PAGE_MAP1_ADDRESS:    usize   = 0x60_000;
pub const MEM_PAGE_MAP2_ADDRESS:    usize   = 0x70_000;
pub const KERNEL_LOCATION:          usize   = 0x100_000;
pub const KERNEL_SIZE_LIMIT:        usize   = 200*0x200;
pub const MEMORY_RESERVED_BELOW:    usize   = KERNEL_LOCATION+KERNEL_SIZE_LIMIT;

// Memory frame (single allocation unit)
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Frame {
    pub index: usize
}
impl Frame {
    /// Create new Frame from memory address. Rounds down.
    pub fn containing_address(address: usize) -> Frame {
        Frame {index: address / MEM_PAGE_SIZE_BYTES}
    }
    /// Create new Frame from memory map index.
    pub fn from_index(index: usize) -> Frame {
        Frame {index: index}
    }
    /// Start address of the frame
    pub fn start_address(&self) -> PhysicalAddress {
        self.index * MEM_PAGE_SIZE_BYTES
    }
    /// Clone frame
    pub fn clone(&self) -> Frame {
        Frame { index: self.index }
    }
    /// Inclusive range between frames
    pub fn range_inclusive(start: Frame, end: Frame) -> FrameIter {
        FrameIter {
            start: start,
            end: end,
        }
    }
}

/// Iterator for frames
pub struct FrameIter {
    start: Frame,
    end: Frame,
}
impl Iterator for FrameIter {
    type Item = Frame;

    fn next(&mut self) -> Option<Frame> {
        if self.start <= self.end {
            let frame = self.start.clone();
            self.start.index += 1;
            Some(frame)
        } else {
            None
        }
    }
}

// Frame allocators all work this way
pub trait FrameAllocator {
    fn allocate_frame(&mut self) -> Option<Frame>;
    fn deallocate_frame(&mut self, frame: Frame);
}

// A simple first-fit frame allocator
// Currently we can only get one frame at a time
pub struct BitmapAllocator {
    index: usize
}

impl BitmapAllocator {
    pub fn new() -> BitmapAllocator {
        BitmapAllocator {index: 0}
    }
    fn is_free(&self, index: usize) -> bool {
        let free    = unsafe {ptr::read_volatile((MEM_PAGE_MAP1_ADDRESS + index/8) as *mut u8)} & (1 << (index%8)) != 0; // 1: free, 0: reserved
        let usable  = unsafe {ptr::read_volatile((MEM_PAGE_MAP2_ADDRESS + index/8) as *mut u8)} & (1 << (index%8)) != 0; // 1: usable, 0: unusable

        if free && !usable { // error in free pages table
            panic!("PHYS_MEM: Page {} is incorrectly marked as free.", index);
        }
        free
    }
}
impl FrameAllocator for BitmapAllocator {
    fn allocate_frame(&mut self) -> Option<Frame> {
        // Find first free block
        for i in 0..(MEM_PAGE_MAP_SIZE_BYTES*8) {
            if self.is_free(i) {
                unsafe {
                    // set frame reserved
                    ptr::write_volatile(
                        (MEM_PAGE_MAP1_ADDRESS + i/8) as *mut u8,
                        ptr::read_volatile((MEM_PAGE_MAP1_ADDRESS + i/8) as *mut u8) ^ (1 << (i%8))
                    );
                }
                return Some(Frame::from_index(i));
            }
        }
        // We could not find a free memory block
        None
    }
    fn deallocate_frame(&mut self, frame: Frame) {
        if self.is_free(frame.index) {
            panic!("PHYS_MEM: deallocate_frame: Page {} is already free.", frame.index);
        }
        unsafe {
            // set the correct bit
            ptr::write_volatile(
                (MEM_PAGE_MAP1_ADDRESS + frame.index/8) as *mut u8,
                ptr::read_volatile((MEM_PAGE_MAP1_ADDRESS + frame.index/8) as *mut u8) | (1 << (frame.index%8))
            );
        }
    }
}

fn mt_align_address(address: usize, upwards: bool) -> usize {
    if address % MEM_PAGE_SIZE_BYTES == 0 {
        address
    }
    else if upwards {
        address + MEM_PAGE_SIZE_BYTES - address % MEM_PAGE_SIZE_BYTES
    }
    else {
        address - address % MEM_PAGE_SIZE_BYTES
    }
}

pub fn create_memory_bitmap() {
    // load memory map from where out bootloader left it
    // http://wiki.osdev.org/Detecting_Memory_(x86)#BIOS_Function:_INT_0x15.2C_EAX_.3D_0xE820

    // zero out the bitmap sections
    for address in (MEM_PAGE_MAP1_ADDRESS..MEM_PAGE_MAP2_ADDRESS+MEM_PAGE_SIZE_BYTES).step_by(8) {
        unsafe {
            ptr::write_volatile(address as *mut u8, 0); // default to (reserved, unusable)
        }
    }

    let entry_count: u8 = unsafe {ptr::read_volatile(BOOT_TMP_MMAP_BUFFER as *mut u8)};
    let base = (BOOT_TMP_MMAP_BUFFER+2) as *mut u8;
    let mut memory_amount_counter_kib = 0;
    for index in 0..(entry_count as isize) {
        let entry_start:    usize   = unsafe { ptr::read_volatile(base.offset(24*index+ 0) as *mut u64) } as usize;
        let entry_size:     usize   = unsafe { ptr::read_volatile(base.offset(24*index+ 8) as *mut u64) } as usize;
        let entry_type:     u32     = unsafe { ptr::read_volatile(base.offset(24*index+16) as *mut u32) };
        let acpi_data:      u32     = unsafe { ptr::read_volatile(base.offset(24*index+20) as *mut u32) };
        // rprintln!("Section {}: {:#016x}-{:#016x}: type: {:#x}, acpi: {:#x}", index, entry_start, entry_start+entry_size, entry_type, acpi_data);

        // is this usable?
        // Types 1, 4 ok to use and acpi_data bit 0 must be set
        if (entry_type == 1 || entry_type == 4) && (acpi_data & 1) == 1 {
            // set frame data. accept only full frames
            for address in (mt_align_address(entry_start, true)..mt_align_address(entry_start+entry_size, false)).step_by(MEM_PAGE_SIZE_BYTES) {
                memory_amount_counter_kib += 1;
                if address/MEM_PAGE_SIZE_BYTES > MEM_PAGE_MAP_SIZE_BYTES*8 {
                    // Page table is full.
                    break;
                }
                if address < MEMORY_RESERVED_BELOW {
                    // these are permanently reserved for the kernel
                    continue;
                }
                unsafe {
                    // set free
                    ptr::write_volatile(
                        (MEM_PAGE_MAP1_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8,
                        ptr::read_volatile((MEM_PAGE_MAP1_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8) | (1 << (address%8))
                    );
                    // set usable
                    ptr::write_volatile(
                        (MEM_PAGE_MAP2_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8,
                        ptr::read_volatile((MEM_PAGE_MAP2_ADDRESS + (address/8)/MEM_PAGE_SIZE_BYTES) as *mut u8) | (1 << (address%8))
                    );
                }
            }
        }
    }
    rprintln!("Memory size {} MiB", memory_amount_counter_kib/1024);
}

// Create static pointer mutex with spinlock to make ALLOCATOR thread-safe
// pub static ALLOCATOR: Mutex<BitmapAllocator> = Mutex::new(BitmapAllocator {});

macro_rules! ALLOCATOR {
    () => ({
        $crate::mem_map::BitmapAllocator::new()
    });
}
<|endoftext|>


<commit_before>//! This module implements the in-memory database
//!
//! A support folder with SDK debug symbols can be processed into a
//! in-memory database format which is a flat file on the file system
//! that gets mmaped into the process.
use std::io::{Write, Seek, SeekFrom};
use std::fs::File;
use std::mem;
use std::slice;
use std::path::Path;
use std::cell::RefCell;
use std::collections::HashMap;

use uuid::Uuid;
use xz2::write::XzEncoder;
use tempfile::tempfile;

use super::types::{IndexItem, StoredSlice, MemDbHeader, IndexedUuid};
use super::super::Result;
use super::super::sdk::{SdkInfo, DumpOptions, Objects};
use super::super::dsym::{Object, Variant};
use super::super::utils::{file_size_format, ProgressIndicator, copy_with_progress};

struct MemDbBuilder<W> {
    writer: RefCell<W>,
    tempfile: Option<RefCell<File>>,
    info: SdkInfo,
    symbols: Vec<String>,
    symbols_map: HashMap<String, u32>,
    object_names: Vec<String>,
    object_names_map: HashMap<String, u16>,
    object_uuid_mapping: Vec<(String, Uuid)>,
    variant_uuids: Vec<IndexedUuid>,
    variants: Vec<Vec<IndexItem>>,
    symbol_count: usize,
    progress: ProgressIndicator,
    options: DumpOptions,
}

trait WriteSeek : Write + Seek {}
impl<T: Write+Seek> WriteSeek for T {}

impl<W: Write + Seek> MemDbBuilder<W> {

    pub fn new(writer: W, info: &SdkInfo, opts: DumpOptions, object_count: usize)
        -> Result<MemDbBuilder<W>>
    {
        let rv = MemDbBuilder {
            writer: RefCell::new(writer),
            tempfile: if opts.compress {
                Some(RefCell::new(tempfile()?))
            } else {
                None
            },
            info: info.clone(),
            symbols: vec![],
            symbols_map: HashMap::new(),
            object_names: vec![],
            object_names_map: HashMap::new(),
            object_uuid_mapping: vec![],
            variant_uuids: vec![],
            variants: vec![],
            symbol_count: 0,
            progress: if opts.show_progress_bar {
                ProgressIndicator::new(object_count)
            } else {
                ProgressIndicator::disabled()
            },
            options: opts,
        };
        let header = MemDbHeader { ..Default::default() };
        rv.write(&header)?;
        rv.progress.set_message("Initializing");
        Ok(rv)
    }

    fn with_file<T, F: FnOnce(&mut WriteSeek) -> T>(&self, f: F) -> T {
        if self.options.compress {
            f(&mut *self.tempfile.as_ref().unwrap().borrow_mut() as &mut WriteSeek)
        } else {
            f(&mut *self.writer.borrow_mut() as &mut WriteSeek)
        }
    }

    fn write_bytes(&self, x: &[u8]) -> Result<usize> {
        self.with_file(|mut w| {
            w.write_all(x)?;
            Ok(x.len())
        })
    }

    fn write<T>(&self, x: &T) -> Result<usize> {
        unsafe {
            let bytes : *const u8 = mem::transmute(x);
            let size = mem::size_of_val(x);
            self.with_file(|mut w| {
                w.write_all(slice::from_raw_parts(bytes, size))?;
                Ok(size)
            })
        }
    }

    fn seek(&self, new_pos: usize) -> Result<()> {
        self.with_file(|mut w| {
            w.seek(SeekFrom::Start(new_pos as u64))?;
            Ok(())
        })
    }

    fn tell(&self) -> Result<usize> {
        self.with_file(|mut w| {
            Ok(w.seek(SeekFrom::Current(0))? as usize)
        })
    }

    fn add_symbol(&mut self, sym: &str) -> usize {
        if let Some(&sym_id) = self.symbols_map.get(sym) {
            return sym_id as usize;
        }
        let symbol_count = self.symbols.len();
        self.symbols.push(sym.to_string());
        self.symbols_map.insert(sym.to_string(), symbol_count as u32);
        symbol_count
    }

    fn add_object_name(&mut self, src: &str) -> usize {
        if let Some(&src_id) = self.object_names_map.get(src) {
            return src_id as usize;
        }
        let object_count = self.object_names.len();
        self.object_names.push(src.to_string());
        self.object_names_map.insert(src.to_string(), object_count as u16);
        object_count
    }

    pub fn write_object(&mut self, obj: &Object, filename: Option<&str>,
                        offset: usize) -> Result<()> {
        self.progress.inc(offset);
        for variant in obj.variants() {
            let src = variant.name().or(filename).unwrap();
            if let Some(fname) = Path::new(src).file_name().and_then(|x| x.to_str()) {
                self.progress.set_message(fname);
            }
            if variant.uuid().is_some() {
                self.write_object_variant(&obj, &variant, src)?;
            }
        }
        Ok(())
    }

    pub fn write_object_variant(&mut self, obj: &Object, var: &Variant,
                                src: &str) -> Result<()> {
        let mut symbols = obj.symbols(var.arch())?;
        let src_id = self.add_object_name(src);

        // build symbol index
        let mut index = vec![];
        for (idx, (addr, sym)) in symbols.iter().enumerate() {
            let sym_id = self.add_symbol(sym);
            index.push(IndexItem::new(addr, src_id, sym_id));
            if idx % 100 == 0 {
                self.progress.tick();
            }
            self.symbol_count += 1;
        }
        index.sort_by_key(|item| item.addr());

        // register variant and uuid
        self.variant_uuids.push(IndexedUuid::new(&var.uuid().unwrap(), self.variants.len()));
        self.object_uuid_mapping.push((
            format!("{}:{}", src, var.arch()),
            var.uuid().unwrap()
        ));
        self.variants.push(index);

        Ok(())
    }

    fn make_string_slices(&self, strings: &[String], _try_compress: bool) -> Result<Vec<StoredSlice>> {
        let mut slices = vec![];
        for (idx, string) in strings.iter().enumerate() {
            let offset = self.tell()?;
            let len = self.write_bytes(string.as_bytes())?;
            slices.push(StoredSlice::new(offset, len, false));
            if idx % 500 == 0 {
                self.progress.tick();
            }
        }
        Ok(slices)
    }

    fn write_slices(&self, slices: &[StoredSlice], start: &mut u32, len: &mut u32) -> Result<()> {
        *start = self.tell()? as u32;
        for (idx, item) in slices.iter().enumerate() {
            self.write(item)?;
            if idx % 500 == 0 {
                self.progress.tick();
            }
        }
        *len = slices.len() as u32;
        Ok(())
    }

    pub fn flush(&mut self) -> Result<()> {
        self.progress.finish(&format!(
            "Processed {} symbols", self.symbol_count));
        self.progress.add_bar(3);

        let mut header = MemDbHeader { ..Default::default() };
        header.version = 1;
        header.sdk_info.set_from_sdk_info(&self.info);

        self.progress.inc(1);
        self.progress.set_message("Writing metadata");

        // start by writing out the index of the variants and record the slices.
        let mut slices = vec![];
        let mut ii_idx = 0;
        for variant in self.variants.iter() {
            let offset = self.tell()?;
            for index_item in variant {
                self.write(index_item)?;
                ii_idx += 1;
                if ii_idx % 100 == 0 {
                    self.progress.tick();
                }
            }
            slices.push(StoredSlice::new(offset, (self.tell()? - offset), false));
        }
        self.write_slices(&slices[..], &mut header.variants_start,
                          &mut header.variants_count)?;

        // next write out the UUIDs.  Since these are fixed length we do not
        // need to use slices here.
        header.uuids_start = self.tell()? as u32;
        header.uuids_count = self.variant_uuids.len() as u32;
        self.variant_uuids.sort_by_key(|x| x.uuid);
        for (idx, indexed_uuid) in self.variant_uuids.iter().enumerate() {
            self.write(indexed_uuid)?;
            if idx % 200 == 0 {
                self.progress.tick();
            }
        }

        // next we write out the name + arch -> uuid index mapping.  We also sort
        // this by uuid so that the index matches up.
        header.tagged_object_names_start = self.tell()? as u32;
        self.object_uuid_mapping.sort_by_key(|&(_, b)| b);
        for (idx, &(ref tagged_object, _)) in self.object_uuid_mapping.iter().enumerate() {
            self.write_bytes(format!("{}\x00", tagged_object).as_bytes())?;
            if idx % 100 == 0 {
                self.progress.tick();
            }
        }
        header.tagged_object_names_end = self.tell()? as u32;

        // now write out all the object name sources
        let slices = self.make_string_slices(&self.object_names[..], true)?;
        self.write_slices(&slices[..], &mut header.object_names_start,
                          &mut header.object_names_count)?;

        self.progress.inc(1);
        self.progress.set_message("Writing symbols");

        // now write out all the symbols
        let slices = self.make_string_slices(&self.symbols[..], true)?;
        self.write_slices(&slices[..], &mut header.symbols_start,
                          &mut header.symbols_count)?;

        self.progress.inc(1);
        self.progress.set_message("Writing headers");

        let file_size = self.tell()?;

        // write the updated header
        self.seek(0)?;
        self.write(&header)?;

        self.progress.finish(&format!(
            "Indexed {} variants", self.variant_uuids.len()));

        // compress if necessary
        if self.options.compress {
            self.progress.add_bar(file_size);
            self.progress.set_message("Compressing");
            self.seek(0)?;
            let mut reader = self.tempfile.as_ref().unwrap().borrow_mut();
            let mut writer = self.writer.borrow_mut();
            {
                let mut zwriter = XzEncoder::new(&mut *writer, 9);
                copy_with_progress(&self.progress, &mut *reader, &mut zwriter)?;
            }
            let compressed_file_size = writer.seek(SeekFrom::Current(0))? as usize;
            let pct = (compressed_file_size * 100) / file_size;
            self.progress.finish(&format!(
                "Compressed from {} to {} ({}% of original size)",
                file_size_format(file_size),
                file_size_format(compressed_file_size),
                pct));
        }

        Ok(())
    }
}

/// Dumps objects into a writer
pub fn dump_memdb<W: Write + Seek>(writer: W, info: &SdkInfo,
                                   opts: DumpOptions, objects: Objects)
    -> Result<()>
{
    let mut builder = MemDbBuilder::new(writer, info, opts, objects.file_count())?;
    for obj_res in objects {
        let (offset, filename, obj) = obj_res?;
        builder.write_object(&obj, Some(&filename), offset)?;
    }
    builder.flush()?;
    Ok(())
}
<commit_msg>Subtract vmaddr from addr before storing<commit_after>//! This module implements the in-memory database
//!
//! A support folder with SDK debug symbols can be processed into a
//! in-memory database format which is a flat file on the file system
//! that gets mmaped into the process.
use std::io::{Write, Seek, SeekFrom};
use std::fs::File;
use std::mem;
use std::slice;
use std::path::Path;
use std::cell::RefCell;
use std::collections::HashMap;

use uuid::Uuid;
use xz2::write::XzEncoder;
use tempfile::tempfile;

use super::types::{IndexItem, StoredSlice, MemDbHeader, IndexedUuid};
use super::super::Result;
use super::super::sdk::{SdkInfo, DumpOptions, Objects};
use super::super::dsym::{Object, Variant};
use super::super::utils::{file_size_format, ProgressIndicator, copy_with_progress};

struct MemDbBuilder<W> {
    writer: RefCell<W>,
    tempfile: Option<RefCell<File>>,
    info: SdkInfo,
    symbols: Vec<String>,
    symbols_map: HashMap<String, u32>,
    object_names: Vec<String>,
    object_names_map: HashMap<String, u16>,
    object_uuid_mapping: Vec<(String, Uuid)>,
    variant_uuids: Vec<IndexedUuid>,
    variants: Vec<Vec<IndexItem>>,
    symbol_count: usize,
    progress: ProgressIndicator,
    options: DumpOptions,
}

trait WriteSeek : Write + Seek {}
impl<T: Write+Seek> WriteSeek for T {}

impl<W: Write + Seek> MemDbBuilder<W> {

    pub fn new(writer: W, info: &SdkInfo, opts: DumpOptions, object_count: usize)
        -> Result<MemDbBuilder<W>>
    {
        let rv = MemDbBuilder {
            writer: RefCell::new(writer),
            tempfile: if opts.compress {
                Some(RefCell::new(tempfile()?))
            } else {
                None
            },
            info: info.clone(),
            symbols: vec![],
            symbols_map: HashMap::new(),
            object_names: vec![],
            object_names_map: HashMap::new(),
            object_uuid_mapping: vec![],
            variant_uuids: vec![],
            variants: vec![],
            symbol_count: 0,
            progress: if opts.show_progress_bar {
                ProgressIndicator::new(object_count)
            } else {
                ProgressIndicator::disabled()
            },
            options: opts,
        };
        let header = MemDbHeader { ..Default::default() };
        rv.write(&header)?;
        rv.progress.set_message("Initializing");
        Ok(rv)
    }

    fn with_file<T, F: FnOnce(&mut WriteSeek) -> T>(&self, f: F) -> T {
        if self.options.compress {
            f(&mut *self.tempfile.as_ref().unwrap().borrow_mut() as &mut WriteSeek)
        } else {
            f(&mut *self.writer.borrow_mut() as &mut WriteSeek)
        }
    }

    fn write_bytes(&self, x: &[u8]) -> Result<usize> {
        self.with_file(|mut w| {
            w.write_all(x)?;
            Ok(x.len())
        })
    }

    fn write<T>(&self, x: &T) -> Result<usize> {
        unsafe {
            let bytes : *const u8 = mem::transmute(x);
            let size = mem::size_of_val(x);
            self.with_file(|mut w| {
                w.write_all(slice::from_raw_parts(bytes, size))?;
                Ok(size)
            })
        }
    }

    fn seek(&self, new_pos: usize) -> Result<()> {
        self.with_file(|mut w| {
            w.seek(SeekFrom::Start(new_pos as u64))?;
            Ok(())
        })
    }

    fn tell(&self) -> Result<usize> {
        self.with_file(|mut w| {
            Ok(w.seek(SeekFrom::Current(0))? as usize)
        })
    }

    fn add_symbol(&mut self, sym: &str) -> usize {
        if let Some(&sym_id) = self.symbols_map.get(sym) {
            return sym_id as usize;
        }
        let symbol_count = self.symbols.len();
        self.symbols.push(sym.to_string());
        self.symbols_map.insert(sym.to_string(), symbol_count as u32);
        symbol_count
    }

    fn add_object_name(&mut self, src: &str) -> usize {
        if let Some(&src_id) = self.object_names_map.get(src) {
            return src_id as usize;
        }
        let object_count = self.object_names.len();
        self.object_names.push(src.to_string());
        self.object_names_map.insert(src.to_string(), object_count as u16);
        object_count
    }

    pub fn write_object(&mut self, obj: &Object, filename: Option<&str>,
                        offset: usize) -> Result<()> {
        self.progress.inc(offset);
        for variant in obj.variants() {
            let src = variant.name().or(filename).unwrap();
            if let Some(fname) = Path::new(src).file_name().and_then(|x| x.to_str()) {
                self.progress.set_message(fname);
            }
            if variant.uuid().is_some() {
                self.write_object_variant(&obj, &variant, src)?;
            }
        }
        Ok(())
    }

    pub fn write_object_variant(&mut self, obj: &Object, var: &Variant,
                                src: &str) -> Result<()> {
        let mut symbols = obj.symbols(var.arch())?;
        let src_id = self.add_object_name(src);

        // build symbol index
        let mut index = vec![];
        for (idx, (addr, sym)) in symbols.iter().enumerate() {
            let sym_id = self.add_symbol(sym);
            index.push(IndexItem::new(addr - var.vmaddr(), src_id, sym_id));
            if idx % 100 == 0 {
                self.progress.tick();
            }
            self.symbol_count += 1;
        }
        index.sort_by_key(|item| item.addr());

        // register variant and uuid
        self.variant_uuids.push(IndexedUuid::new(&var.uuid().unwrap(), self.variants.len()));
        self.object_uuid_mapping.push((
            format!("{}:{}", src, var.arch()),
            var.uuid().unwrap()
        ));
        self.variants.push(index);

        Ok(())
    }

    fn make_string_slices(&self, strings: &[String], _try_compress: bool) -> Result<Vec<StoredSlice>> {
        let mut slices = vec![];
        for (idx, string) in strings.iter().enumerate() {
            let offset = self.tell()?;
            let len = self.write_bytes(string.as_bytes())?;
            slices.push(StoredSlice::new(offset, len, false));
            if idx % 500 == 0 {
                self.progress.tick();
            }
        }
        Ok(slices)
    }

    fn write_slices(&self, slices: &[StoredSlice], start: &mut u32, len: &mut u32) -> Result<()> {
        *start = self.tell()? as u32;
        for (idx, item) in slices.iter().enumerate() {
            self.write(item)?;
            if idx % 500 == 0 {
                self.progress.tick();
            }
        }
        *len = slices.len() as u32;
        Ok(())
    }

    pub fn flush(&mut self) -> Result<()> {
        self.progress.finish(&format!(
            "Processed {} symbols", self.symbol_count));
        self.progress.add_bar(3);

        let mut header = MemDbHeader { ..Default::default() };
        header.version = 1;
        header.sdk_info.set_from_sdk_info(&self.info);

        self.progress.inc(1);
        self.progress.set_message("Writing metadata");

        // start by writing out the index of the variants and record the slices.
        let mut slices = vec![];
        let mut ii_idx = 0;
        for variant in self.variants.iter() {
            let offset = self.tell()?;
            for index_item in variant {
                self.write(index_item)?;
                ii_idx += 1;
                if ii_idx % 100 == 0 {
                    self.progress.tick();
                }
            }
            slices.push(StoredSlice::new(offset, (self.tell()? - offset), false));
        }
        self.write_slices(&slices[..], &mut header.variants_start,
                          &mut header.variants_count)?;

        // next write out the UUIDs.  Since these are fixed length we do not
        // need to use slices here.
        header.uuids_start = self.tell()? as u32;
        header.uuids_count = self.variant_uuids.len() as u32;
        self.variant_uuids.sort_by_key(|x| x.uuid);
        for (idx, indexed_uuid) in self.variant_uuids.iter().enumerate() {
            self.write(indexed_uuid)?;
            if idx % 200 == 0 {
                self.progress.tick();
            }
        }

        // next we write out the name + arch -> uuid index mapping.  We also sort
        // this by uuid so that the index matches up.
        header.tagged_object_names_start = self.tell()? as u32;
        self.object_uuid_mapping.sort_by_key(|&(_, b)| b);
        for (idx, &(ref tagged_object, _)) in self.object_uuid_mapping.iter().enumerate() {
            self.write_bytes(format!("{}\x00", tagged_object).as_bytes())?;
            if idx % 100 == 0 {
                self.progress.tick();
            }
        }
        header.tagged_object_names_end = self.tell()? as u32;

        // now write out all the object name sources
        let slices = self.make_string_slices(&self.object_names[..], true)?;
        self.write_slices(&slices[..], &mut header.object_names_start,
                          &mut header.object_names_count)?;

        self.progress.inc(1);
        self.progress.set_message("Writing symbols");

        // now write out all the symbols
        let slices = self.make_string_slices(&self.symbols[..], true)?;
        self.write_slices(&slices[..], &mut header.symbols_start,
                          &mut header.symbols_count)?;

        self.progress.inc(1);
        self.progress.set_message("Writing headers");

        let file_size = self.tell()?;

        // write the updated header
        self.seek(0)?;
        self.write(&header)?;

        self.progress.finish(&format!(
            "Indexed {} variants", self.variant_uuids.len()));

        // compress if necessary
        if self.options.compress {
            self.progress.add_bar(file_size);
            self.progress.set_message("Compressing");
            self.seek(0)?;
            let mut reader = self.tempfile.as_ref().unwrap().borrow_mut();
            let mut writer = self.writer.borrow_mut();
            {
                let mut zwriter = XzEncoder::new(&mut *writer, 9);
                copy_with_progress(&self.progress, &mut *reader, &mut zwriter)?;
            }
            let compressed_file_size = writer.seek(SeekFrom::Current(0))? as usize;
            let pct = (compressed_file_size * 100) / file_size;
            self.progress.finish(&format!(
                "Compressed from {} to {} ({}% of original size)",
                file_size_format(file_size),
                file_size_format(compressed_file_size),
                pct));
        }

        Ok(())
    }
}

/// Dumps objects into a writer
pub fn dump_memdb<W: Write + Seek>(writer: W, info: &SdkInfo,
                                   opts: DumpOptions, objects: Objects)
    -> Result<()>
{
    let mut builder = MemDbBuilder::new(writer, info, opts, objects.file_count())?;
    for obj_res in objects {
        let (offset, filename, obj) = obj_res?;
        builder.write_object(&obj, Some(&filename), offset)?;
    }
    builder.flush()?;
    Ok(())
}
<|endoftext|>


<commit_before>use std::io::net::ip::{SocketAddr, IpAddr};
use std::sync::Arc;

use middleware::Application;
use request::Request;

use hyper::server::{Server, Handler, Incoming};
use hyper::server::Request as HyperRequest;
use hyper::server::Response as HyperResponse;
use hyper::header::common::ContentLength;
use hyper::net::{HttpStream, HttpAcceptor, Fresh};

pub trait ConcurrentHandler: Send + Sync {
    fn handle(&self, req: HyperRequest, res: HyperResponse<Fresh>);
}

pub struct Concurrent<H: ConcurrentHandler> { pub handler: Arc<H> }

impl<H: ConcurrentHandler> Handler<HttpAcceptor, HttpStream> for Concurrent<H> {
    fn handle(self, mut incoming: Incoming) {
        for (mut req, mut res) in incoming {
            let clone = self.handler.clone();
            spawn(proc() { clone.handle(req, res) })
        }
    }
}

macro_rules! try_abort(
    ($e:expr) => {{
        match $e {
            Ok(v) => v,
            Err(..) => return
        }
    }}
)

impl ConcurrentHandler for Application {
    fn handle(&self, mut req: HyperRequest, mut res: HyperResponse<Fresh>) {

        let mut request = Request::wrap(req).unwrap();
        let response = self.call(&mut request);

        match response {
            Ok(response) => {
                *res.status_mut() = response.status;
                *res.headers_mut() = response.headers;

                match response.body {
                    Some(mut reader) => {
                        let content = try_abort!(reader.read_to_end());
                        res.headers_mut().set(ContentLength(content.len()));
                        let mut res = try_abort!(res.start());
                        try_abort!(res.write(content.as_slice()));
                        try_abort!(res.end());
                    },
                    _ => ()
                }

            },

            Err(_) => println!("No response")
        }
    }
}<commit_msg>refactor: naming fixes<commit_after>use std::io::net::ip::{SocketAddr, IpAddr};
use std::sync::Arc;

use middleware::Application;
use request::Request;

use hyper::server::{Server, Handler, Incoming};
use hyper::server::Request as HyperRequest;
use hyper::server::Response as HyperResponse;
use hyper::header::common::ContentLength;
use hyper::net::{HttpStream, HttpAcceptor, Fresh};

pub trait ConcurrentHandler: Send + Sync {
    fn handle(&self, req: HyperRequest, res: HyperResponse<Fresh>);
}

pub struct Concurrent<H: ConcurrentHandler> { pub handler: Arc<H> }

impl<H: ConcurrentHandler> Handler<HttpAcceptor, HttpStream> for Concurrent<H> {
    fn handle(self, mut incoming: Incoming) {
        for (mut req, mut res) in incoming {
            let clone = self.handler.clone();
            spawn(proc() { clone.handle(req, res) })
        }
    }
}

macro_rules! try_abort(
    ($e:expr) => {{
        match $e {
            Ok(v) => v,
            Err(..) => return
        }
    }}
)

impl ConcurrentHandler for Application {
    fn handle(&self, mut req: HyperRequest, mut res: HyperResponse<Fresh>) {

        let mut request = Request::wrap(req).unwrap();
        let maybe_response = self.call(&mut request);

        match maybe_response {
            Ok(response) => {
                *res.status_mut() = response.status;
                *res.headers_mut() = response.headers;

                match response.body {
                    Some(mut reader) => {
                        let content = try_abort!(reader.read_to_end());
                        res.headers_mut().set(ContentLength(content.len()));
                        let mut res = try_abort!(res.start());
                        try_abort!(res.write(content.as_slice()));
                        try_abort!(res.end());
                    },
                    _ => ()
                }
            },

            Err(_) => println!("No response")
        }
    }
}<|endoftext|>