<commit_before>//! This crate implements the standard attribute dictionary.
//!
//! This dictionary is a singleton containing all information about the
//! DICOM attributes specified in the standard according to DICOM PS3.6 2019c,
//! and it will be used by default in most other abstractions available.
//!
//! When not using private tags, this dictionary should suffice.

pub mod tags;

use crate::tags::ENTRIES;
use dicom_core::dictionary::{DataDictionary, DictionaryEntryRef, TagRange::*};
use dicom_core::header::Tag;
use lazy_static::lazy_static;
use std::collections::{HashMap, HashSet};
use std::fmt;
use std::fmt::{Display, Formatter};

lazy_static! {
    static ref DICT: StandardDictionaryRegistry = init_dictionary();
}

/// Retrieve a singleton instance of the standard dictionary registry.
pub fn registry() -> &'static StandardDictionaryRegistry {
    &DICT
}

/// The data struct containing the standard dictionary.
#[derive(Debug)]
pub struct StandardDictionaryRegistry {
    /// mapping: name → tag
    by_name: HashMap<&'static str, &'static DictionaryEntryRef<'static>>,
    /// mapping: tag → name
    by_tag: HashMap<Tag, &'static DictionaryEntryRef<'static>>,
    /// repeating elements of the form (ggxx, eeee). The `xx` portion is zeroed.
    repeating_ggxx: HashSet<Tag>,
    /// repeating elements of the form (gggg, eexx). The `xx` portion is zeroed.
    repeating_eexx: HashSet<Tag>,
}

impl StandardDictionaryRegistry {
    fn new() -> StandardDictionaryRegistry {
        StandardDictionaryRegistry {
            by_name: HashMap::with_capacity(5000),
            by_tag: HashMap::with_capacity(5000),
            repeating_ggxx: HashSet::with_capacity(75),
            repeating_eexx: HashSet::new(),
        }
    }

    /// record the given dictionary entry reference
    fn index(&mut self, entry: &'static DictionaryEntryRef<'static>) -> &mut Self {
        self.by_name.insert(entry.alias, entry);
        self.by_tag.insert(entry.tag.inner(), entry);
        match entry.tag {
            Group100(tag) => {
                self.repeating_ggxx.insert(tag);
            }
            Element100(tag) => {
                self.repeating_eexx.insert(tag);
            }
            _ => {}
        }
        self
    }
}

/// A data dictionary which consults the library's global DICOM attribute registry.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct StandardDataDictionary;

impl StandardDataDictionary {
    fn indexed_tag(tag: Tag) -> Option<&'static DictionaryEntryRef<'static>> {
        let r = registry();

        r.by_tag
            .get(&tag)
            .or_else(|| {
                let group_trimmed = Tag(tag.0 & 0xFF00, tag.1);

                if r.repeating_ggxx.contains(&group_trimmed) {
                    r.by_tag.get(&group_trimmed)
                } else {
                    let elem_trimmed = Tag(tag.0, tag.1 & 0xFF00);
                    if r.repeating_eexx.contains(&elem_trimmed) {
                        r.by_tag.get(&elem_trimmed)
                    } else {
                        None
                    }
                }
            })
            .cloned()
    }
}

impl DataDictionary for StandardDataDictionary {
    type Entry = DictionaryEntryRef<'static>;

    fn by_name(&self, name: &str) -> Option<&Self::Entry> {
        registry().by_name.get(name).cloned()
    }

    fn by_tag(&self, tag: Tag) -> Option<&Self::Entry> {
        StandardDataDictionary::indexed_tag(tag)
    }
}

impl<'a> DataDictionary for &'a StandardDataDictionary {
    type Entry = DictionaryEntryRef<'static>;

    fn by_name(&self, name: &str) -> Option<&'static DictionaryEntryRef<'static>> {
        registry().by_name.get(name).cloned()
    }

    fn by_tag(&self, tag: Tag) -> Option<&'static DictionaryEntryRef<'static>> {
        StandardDataDictionary::indexed_tag(tag)
    }
}

impl Display for StandardDataDictionary {
    fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
        f.write_str("Standard DICOM Data Dictionary")
    }
}

fn init_dictionary() -> StandardDictionaryRegistry {
    let mut d = StandardDictionaryRegistry::new();
    for entry in ENTRIES {
        d.index(&entry);
    }
    d
}

#[cfg(test)]
mod tests {
    use super::StandardDataDictionary;
    use dicom_core::dictionary::{DataDictionary, DictionaryEntryRef, TagRange::*};
    use dicom_core::header::{Tag, VR};

    // tests for just a few attributes to make sure that the entries
    // were well installed into the crate
    #[test]
    fn smoke_test() {
        let dict = StandardDataDictionary::default();

        assert_eq!(
            dict.by_name("PatientName"),
            Some(&DictionaryEntryRef {
                tag: Single(Tag(0x0010, 0x0010)),
                alias: "PatientName",
                vr: VR::PN,
            })
        );

        assert_eq!(
            dict.by_name("Modality"),
            Some(&DictionaryEntryRef {
                tag: Single(Tag(0x0008, 0x0060)),
                alias: "Modality",
                vr: VR::CS,
            })
        );

        let pixel_data = dict
            .by_tag(Tag(0x7FE0, 0x0010))
            .expect("Pixel Data attribute should exist");
        eprintln!("{:X?}", pixel_data.tag);
        assert_eq!(pixel_data.tag, Single(Tag(0x7FE0, 0x0010)));
        assert_eq!(pixel_data.alias, "PixelData");
        assert!(pixel_data.vr == VR::OB || pixel_data.vr == VR::OW);

        let overlay_data = dict
            .by_tag(Tag(0x6000, 0x3000))
            .expect("Overlay Data attribute should exist");
        assert_eq!(overlay_data.tag, Group100(Tag(0x6000, 0x3000)));
        assert_eq!(overlay_data.alias, "OverlayData");
        assert!(overlay_data.vr == VR::OB || overlay_data.vr == VR::OW);

        // repeated overlay data
        let overlay_data = dict
            .by_tag(Tag(0x60EE, 0x3000))
            .expect("Repeated Overlay Data attribute should exist");
        assert_eq!(overlay_data.tag, Group100(Tag(0x6000, 0x3000)));
        assert_eq!(overlay_data.alias, "OverlayData");
        assert!(overlay_data.vr == VR::OB || overlay_data.vr == VR::OW);
    }

    // tests for just a few attributes to make sure that the tag constants
    // were well installed into the crate
    #[test]
    fn constants_available() {
        use crate::tags::*;
        assert_eq!(PATIENT_NAME, Tag(0x0010, 0x0010));
        assert_eq!(MODALITY, Tag(0x0008, 0x0060));
        assert_eq!(PIXEL_DATA, Tag(0x7FE0, 0x0010));
    }
}
<commit_msg>[dictionary-std] add a command group attribute<commit_after>//! This crate implements the standard attribute dictionary.
//!
//! This dictionary is a singleton containing all information about the
//! DICOM attributes specified in the standard according to DICOM PS3.6 2019c,
//! and it will be used by default in most other abstractions available.
//!
//! When not using private tags, this dictionary should suffice.

pub mod tags;

use crate::tags::ENTRIES;
use dicom_core::dictionary::{DataDictionary, DictionaryEntryRef, TagRange::*};
use dicom_core::header::Tag;
use lazy_static::lazy_static;
use std::collections::{HashMap, HashSet};
use std::fmt;
use std::fmt::{Display, Formatter};

lazy_static! {
    static ref DICT: StandardDictionaryRegistry = init_dictionary();
}

/// Retrieve a singleton instance of the standard dictionary registry.
pub fn registry() -> &'static StandardDictionaryRegistry {
    &DICT
}

/// The data struct containing the standard dictionary.
#[derive(Debug)]
pub struct StandardDictionaryRegistry {
    /// mapping: name → tag
    by_name: HashMap<&'static str, &'static DictionaryEntryRef<'static>>,
    /// mapping: tag → name
    by_tag: HashMap<Tag, &'static DictionaryEntryRef<'static>>,
    /// repeating elements of the form (ggxx, eeee). The `xx` portion is zeroed.
    repeating_ggxx: HashSet<Tag>,
    /// repeating elements of the form (gggg, eexx). The `xx` portion is zeroed.
    repeating_eexx: HashSet<Tag>,
}

impl StandardDictionaryRegistry {
    fn new() -> StandardDictionaryRegistry {
        StandardDictionaryRegistry {
            by_name: HashMap::with_capacity(5000),
            by_tag: HashMap::with_capacity(5000),
            repeating_ggxx: HashSet::with_capacity(75),
            repeating_eexx: HashSet::new(),
        }
    }

    /// record the given dictionary entry reference
    fn index(&mut self, entry: &'static DictionaryEntryRef<'static>) -> &mut Self {
        self.by_name.insert(entry.alias, entry);
        self.by_tag.insert(entry.tag.inner(), entry);
        match entry.tag {
            Group100(tag) => {
                self.repeating_ggxx.insert(tag);
            }
            Element100(tag) => {
                self.repeating_eexx.insert(tag);
            }
            _ => {}
        }
        self
    }
}

/// A data dictionary which consults the library's global DICOM attribute registry.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct StandardDataDictionary;

impl StandardDataDictionary {
    fn indexed_tag(tag: Tag) -> Option<&'static DictionaryEntryRef<'static>> {
        let r = registry();

        r.by_tag
            .get(&tag)
            .or_else(|| {
                let group_trimmed = Tag(tag.0 & 0xFF00, tag.1);

                if r.repeating_ggxx.contains(&group_trimmed) {
                    r.by_tag.get(&group_trimmed)
                } else {
                    let elem_trimmed = Tag(tag.0, tag.1 & 0xFF00);
                    if r.repeating_eexx.contains(&elem_trimmed) {
                        r.by_tag.get(&elem_trimmed)
                    } else {
                        None
                    }
                }
            })
            .cloned()
    }
}

impl DataDictionary for StandardDataDictionary {
    type Entry = DictionaryEntryRef<'static>;

    fn by_name(&self, name: &str) -> Option<&Self::Entry> {
        registry().by_name.get(name).cloned()
    }

    fn by_tag(&self, tag: Tag) -> Option<&Self::Entry> {
        StandardDataDictionary::indexed_tag(tag)
    }
}

impl<'a> DataDictionary for &'a StandardDataDictionary {
    type Entry = DictionaryEntryRef<'static>;

    fn by_name(&self, name: &str) -> Option<&'static DictionaryEntryRef<'static>> {
        registry().by_name.get(name).cloned()
    }

    fn by_tag(&self, tag: Tag) -> Option<&'static DictionaryEntryRef<'static>> {
        StandardDataDictionary::indexed_tag(tag)
    }
}

impl Display for StandardDataDictionary {
    fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
        f.write_str("Standard DICOM Data Dictionary")
    }
}

fn init_dictionary() -> StandardDictionaryRegistry {
    let mut d = StandardDictionaryRegistry::new();
    for entry in ENTRIES {
        d.index(&entry);
    }
    d
}

#[cfg(test)]
mod tests {
    use super::StandardDataDictionary;
    use dicom_core::dictionary::{DataDictionary, DictionaryEntryRef, TagRange::*};
    use dicom_core::header::{Tag, VR};

    // tests for just a few attributes to make sure that the entries
    // were well installed into the crate
    #[test]
    fn smoke_test() {
        let dict = StandardDataDictionary::default();

        assert_eq!(
            dict.by_name("PatientName"),
            Some(&DictionaryEntryRef {
                tag: Single(Tag(0x0010, 0x0010)),
                alias: "PatientName",
                vr: VR::PN,
            })
        );

        assert_eq!(
            dict.by_name("Modality"),
            Some(&DictionaryEntryRef {
                tag: Single(Tag(0x0008, 0x0060)),
                alias: "Modality",
                vr: VR::CS,
            })
        );

        let pixel_data = dict
            .by_tag(Tag(0x7FE0, 0x0010))
            .expect("Pixel Data attribute should exist");
        eprintln!("{:X?}", pixel_data.tag);
        assert_eq!(pixel_data.tag, Single(Tag(0x7FE0, 0x0010)));
        assert_eq!(pixel_data.alias, "PixelData");
        assert!(pixel_data.vr == VR::OB || pixel_data.vr == VR::OW);

        let overlay_data = dict
            .by_tag(Tag(0x6000, 0x3000))
            .expect("Overlay Data attribute should exist");
        assert_eq!(overlay_data.tag, Group100(Tag(0x6000, 0x3000)));
        assert_eq!(overlay_data.alias, "OverlayData");
        assert!(overlay_data.vr == VR::OB || overlay_data.vr == VR::OW);

        // repeated overlay data
        let overlay_data = dict
            .by_tag(Tag(0x60EE, 0x3000))
            .expect("Repeated Overlay Data attribute should exist");
        assert_eq!(overlay_data.tag, Group100(Tag(0x6000, 0x3000)));
        assert_eq!(overlay_data.alias, "OverlayData");
        assert!(overlay_data.vr == VR::OB || overlay_data.vr == VR::OW);
    }

    // tests for just a few attributes to make sure that the tag constants
    // were well installed into the crate
    #[test]
    fn constants_available() {
        use crate::tags::*;
        assert_eq!(PATIENT_NAME, Tag(0x0010, 0x0010));
        assert_eq!(MODALITY, Tag(0x0008, 0x0060));
        assert_eq!(PIXEL_DATA, Tag(0x7FE0, 0x0010));
        assert_eq!(STATUS, Tag(0x0000, 0x0900));
    }
}
<|endoftext|>


<commit_before>use super::*;

use rustc_middle::mir::coverage::*;
use rustc_middle::mir::visit::Visitor;
use rustc_middle::mir::{self, Coverage, CoverageInfo, Location};
use rustc_middle::ty::query::Providers;
use rustc_middle::ty::TyCtxt;
use rustc_span::def_id::DefId;

/// The `query` provider for `CoverageInfo`, requested by `codegen_coverage()` (to inject each
/// counter) and `FunctionCoverage::new()` (to extract the coverage map metadata from the MIR).
pub(crate) fn provide(providers: &mut Providers) {
    providers.coverageinfo = |tcx, def_id| coverageinfo_from_mir(tcx, def_id);
    providers.covered_file_name = |tcx, def_id| covered_file_name(tcx, def_id);
    providers.covered_code_regions = |tcx, def_id| covered_code_regions(tcx, def_id);
}

/// The `num_counters` argument to `llvm.instrprof.increment` is the max counter_id + 1, or in
/// other words, the number of counter value references injected into the MIR (plus 1 for the
/// reserved `ZERO` counter, which uses counter ID `0` when included in an expression). Injected
/// counters have a counter ID from `1..num_counters-1`.
///
/// `num_expressions` is the number of counter expressions added to the MIR body.
///
/// Both `num_counters` and `num_expressions` are used to initialize new vectors, during backend
/// code generate, to lookup counters and expressions by simple u32 indexes.
///
/// MIR optimization may split and duplicate some BasicBlock sequences, or optimize out some code
/// including injected counters. (It is OK if some counters are optimized out, but those counters
/// are still included in the total `num_counters` or `num_expressions`.) Simply counting the
/// calls may not work; but computing the number of counters or expressions by adding `1` to the
/// highest ID (for a given instrumented function) is valid.
///
/// This visitor runs twice, first with `add_missing_operands` set to `false`, to find the maximum
/// counter ID and maximum expression ID based on their enum variant `id` fields; then, as a
/// safeguard, with `add_missing_operands` set to `true`, to find any other counter or expression
/// IDs referenced by expression operands, if not already seen.
///
/// Ideally, each operand ID in a MIR `CoverageKind::Expression` will have a separate MIR `Coverage`
/// statement for the `Counter` or `Expression` with the referenced ID. but since current or future
/// MIR optimizations can theoretically optimize out segments of a MIR, it may not be possible to
/// guarantee this, so the second pass ensures the `CoverageInfo` counts include all referenced IDs.
struct CoverageVisitor {
    info: CoverageInfo,
    add_missing_operands: bool,
}

impl CoverageVisitor {
    /// Updates `num_counters` to the maximum encountered zero-based counter_id plus 1. Note the
    /// final computed number of counters should be the number of all `CoverageKind::Counter`
    /// statements in the MIR *plus one* for the implicit `ZERO` counter.
    #[inline(always)]
    fn update_num_counters(&mut self, counter_id: u32) {
        self.info.num_counters = std::cmp::max(self.info.num_counters, counter_id + 1);
    }

    /// Computes an expression index for each expression ID, and updates `num_expressions` to the
    /// maximum encountered index plus 1.
    #[inline(always)]
    fn update_num_expressions(&mut self, expression_id: u32) {
        let expression_index = u32::MAX - expression_id;
        self.info.num_expressions = std::cmp::max(self.info.num_expressions, expression_index + 1);
    }

    fn update_from_expression_operand(&mut self, operand_id: u32) {
        if operand_id >= self.info.num_counters {
            let operand_as_expression_index = u32::MAX - operand_id;
            if operand_as_expression_index >= self.info.num_expressions {
                // The operand ID is outside the known range of counter IDs and also outside the
                // known range of expression IDs. In either case, the result of a missing operand
                // (if and when used in an expression) will be zero, so from a computation
                // perspective, it doesn't matter whether it is interepretted as a counter or an
                // expression.
                //
                // However, the `num_counters` and `num_expressions` query results are used to
                // allocate arrays when generating the coverage map (during codegen), so choose
                // the type that grows either `num_counters` or `num_expressions` the least.
                if operand_id - self.info.num_counters
                    < operand_as_expression_index - self.info.num_expressions
                {
                    self.update_num_counters(operand_id)
                } else {
                    self.update_num_expressions(operand_id)
                }
            }
        }
    }
}

impl Visitor<'_> for CoverageVisitor {
    fn visit_coverage(&mut self, coverage: &Coverage, _location: Location) {
        if self.add_missing_operands {
            match coverage.kind {
                CoverageKind::Expression { lhs, rhs, .. } => {
                    self.update_from_expression_operand(u32::from(lhs));
                    self.update_from_expression_operand(u32::from(rhs));
                }
                _ => {}
            }
        } else {
            match coverage.kind {
                CoverageKind::Counter { id, .. } => {
                    self.update_num_counters(u32::from(id));
                }
                CoverageKind::Expression { id, .. } => {
                    self.update_num_expressions(u32::from(id));
                }
                _ => {}
            }
        }
    }
}

fn coverageinfo_from_mir<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> CoverageInfo {
    let mir_body = tcx.optimized_mir(def_id);

    let mut coverage_visitor = CoverageVisitor {
        // num_counters always has at least the `ZERO` counter.
        info: CoverageInfo { num_counters: 1, num_expressions: 0 },
        add_missing_operands: false,
    };

    coverage_visitor.visit_body(mir_body);

    coverage_visitor.add_missing_operands = true;
    coverage_visitor.visit_body(mir_body);

    coverage_visitor.info
}

fn covered_file_name<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<Symbol> {
    let mir_body = tcx.optimized_mir(def_id);
    for bb_data in mir_body.basic_blocks().iter() {
        for statement in bb_data.statements.iter() {
            if let StatementKind::Coverage(box ref coverage) = statement.kind {
                if let Some(code_region) = coverage.code_region.as_ref() {
                    return Some(code_region.file_name);
                }
            }
        }
    }
    None
}

fn covered_code_regions<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Vec<&'tcx CodeRegion> {
    let mir_body: &'tcx mir::Body<'tcx> = tcx.optimized_mir(def_id);
    mir_body
        .basic_blocks()
        .iter()
        .map(|data| {
            data.statements.iter().filter_map(|statement| match statement.kind {
                StatementKind::Coverage(box ref coverage) => {
                    coverage.code_region.as_ref() // may be None
                }
                _ => None,
            })
        })
        .flatten()
        .collect()
}
<commit_msg>Coverage computation needs access to the MIR, too<commit_after>use super::*;

use rustc_middle::mir::coverage::*;
use rustc_middle::mir::visit::Visitor;
use rustc_middle::mir::{self, Coverage, CoverageInfo, Location};
use rustc_middle::ty::query::Providers;
use rustc_middle::ty::{self, TyCtxt};
use rustc_span::def_id::DefId;

/// The `query` provider for `CoverageInfo`, requested by `codegen_coverage()` (to inject each
/// counter) and `FunctionCoverage::new()` (to extract the coverage map metadata from the MIR).
pub(crate) fn provide(providers: &mut Providers) {
    providers.coverageinfo = |tcx, def_id| coverageinfo_from_mir(tcx, def_id);
    providers.covered_file_name = |tcx, def_id| covered_file_name(tcx, def_id);
    providers.covered_code_regions = |tcx, def_id| covered_code_regions(tcx, def_id);
}

/// The `num_counters` argument to `llvm.instrprof.increment` is the max counter_id + 1, or in
/// other words, the number of counter value references injected into the MIR (plus 1 for the
/// reserved `ZERO` counter, which uses counter ID `0` when included in an expression). Injected
/// counters have a counter ID from `1..num_counters-1`.
///
/// `num_expressions` is the number of counter expressions added to the MIR body.
///
/// Both `num_counters` and `num_expressions` are used to initialize new vectors, during backend
/// code generate, to lookup counters and expressions by simple u32 indexes.
///
/// MIR optimization may split and duplicate some BasicBlock sequences, or optimize out some code
/// including injected counters. (It is OK if some counters are optimized out, but those counters
/// are still included in the total `num_counters` or `num_expressions`.) Simply counting the
/// calls may not work; but computing the number of counters or expressions by adding `1` to the
/// highest ID (for a given instrumented function) is valid.
///
/// This visitor runs twice, first with `add_missing_operands` set to `false`, to find the maximum
/// counter ID and maximum expression ID based on their enum variant `id` fields; then, as a
/// safeguard, with `add_missing_operands` set to `true`, to find any other counter or expression
/// IDs referenced by expression operands, if not already seen.
///
/// Ideally, each operand ID in a MIR `CoverageKind::Expression` will have a separate MIR `Coverage`
/// statement for the `Counter` or `Expression` with the referenced ID. but since current or future
/// MIR optimizations can theoretically optimize out segments of a MIR, it may not be possible to
/// guarantee this, so the second pass ensures the `CoverageInfo` counts include all referenced IDs.
struct CoverageVisitor {
    info: CoverageInfo,
    add_missing_operands: bool,
}

impl CoverageVisitor {
    /// Updates `num_counters` to the maximum encountered zero-based counter_id plus 1. Note the
    /// final computed number of counters should be the number of all `CoverageKind::Counter`
    /// statements in the MIR *plus one* for the implicit `ZERO` counter.
    #[inline(always)]
    fn update_num_counters(&mut self, counter_id: u32) {
        self.info.num_counters = std::cmp::max(self.info.num_counters, counter_id + 1);
    }

    /// Computes an expression index for each expression ID, and updates `num_expressions` to the
    /// maximum encountered index plus 1.
    #[inline(always)]
    fn update_num_expressions(&mut self, expression_id: u32) {
        let expression_index = u32::MAX - expression_id;
        self.info.num_expressions = std::cmp::max(self.info.num_expressions, expression_index + 1);
    }

    fn update_from_expression_operand(&mut self, operand_id: u32) {
        if operand_id >= self.info.num_counters {
            let operand_as_expression_index = u32::MAX - operand_id;
            if operand_as_expression_index >= self.info.num_expressions {
                // The operand ID is outside the known range of counter IDs and also outside the
                // known range of expression IDs. In either case, the result of a missing operand
                // (if and when used in an expression) will be zero, so from a computation
                // perspective, it doesn't matter whether it is interepretted as a counter or an
                // expression.
                //
                // However, the `num_counters` and `num_expressions` query results are used to
                // allocate arrays when generating the coverage map (during codegen), so choose
                // the type that grows either `num_counters` or `num_expressions` the least.
                if operand_id - self.info.num_counters
                    < operand_as_expression_index - self.info.num_expressions
                {
                    self.update_num_counters(operand_id)
                } else {
                    self.update_num_expressions(operand_id)
                }
            }
        }
    }
}

impl Visitor<'_> for CoverageVisitor {
    fn visit_coverage(&mut self, coverage: &Coverage, _location: Location) {
        if self.add_missing_operands {
            match coverage.kind {
                CoverageKind::Expression { lhs, rhs, .. } => {
                    self.update_from_expression_operand(u32::from(lhs));
                    self.update_from_expression_operand(u32::from(rhs));
                }
                _ => {}
            }
        } else {
            match coverage.kind {
                CoverageKind::Counter { id, .. } => {
                    self.update_num_counters(u32::from(id));
                }
                CoverageKind::Expression { id, .. } => {
                    self.update_num_expressions(u32::from(id));
                }
                _ => {}
            }
        }
    }
}

fn coverageinfo_from_mir<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> CoverageInfo {
    let mir_body = mir_body(tcx, def_id);

    let mut coverage_visitor = CoverageVisitor {
        // num_counters always has at least the `ZERO` counter.
        info: CoverageInfo { num_counters: 1, num_expressions: 0 },
        add_missing_operands: false,
    };

    coverage_visitor.visit_body(mir_body);

    coverage_visitor.add_missing_operands = true;
    coverage_visitor.visit_body(mir_body);

    coverage_visitor.info
}

fn covered_file_name<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<Symbol> {
    for bb_data in mir_body(tcx, def_id).basic_blocks().iter() {
        for statement in bb_data.statements.iter() {
            if let StatementKind::Coverage(box ref coverage) = statement.kind {
                if let Some(code_region) = coverage.code_region.as_ref() {
                    return Some(code_region.file_name);
                }
            }
        }
    }
    None
}

/// This function ensures we obtain the correct MIR for the given item irrespective of
/// whether that means const mir or runtime mir. For `const fn` this opts for runtime
/// mir.
fn mir_body<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> &'tcx mir::Body<'tcx> {
    let id = ty::WithOptConstParam::unknown(def_id);
    let def = ty::InstanceDef::Item(id);
    tcx.instance_mir(def)
}

fn covered_code_regions<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Vec<&'tcx CodeRegion> {
    mir_body(tcx, def_id)
        .basic_blocks()
        .iter()
        .map(|data| {
            data.statements.iter().filter_map(|statement| match statement.kind {
                StatementKind::Coverage(box ref coverage) => {
                    coverage.code_region.as_ref() // may be None
                }
                _ => None,
            })
        })
        .flatten()
        .collect()
}
<|endoftext|>


<commit_before>//! Types for the [`m.room.guest_access`] event.
//!
//! [`m.room.guest_access`]: https://spec.matrix.org/v1.2/client-server-api/#mroomguest_access

use ruma_macros::EventContent;
use serde::{Deserialize, Serialize};

use crate::{serde::StringEnum, PrivOwnedStr};

/// The content of an `m.room.guest_access` event.
///
/// Controls whether guest users are allowed to join rooms.
///
/// This event controls whether guest users are allowed to join rooms. If this event is absent,
/// servers should act as if it is present and has the value `GuestAccess::Forbidden`.
#[derive(Clone, Debug, Deserialize, Serialize, EventContent)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
#[ruma_event(type = "m.room.guest_access", kind = State)]
pub struct RoomGuestAccessEventContent {
    /// A policy for guest user access to a room.
    pub guest_access: GuestAccess,
}

impl RoomGuestAccessEventContent {
    /// Creates a new `RoomGuestAccessEventContent` with the given policy.
    pub fn new(guest_access: GuestAccess) -> Self {
        Self { guest_access }
    }
}

/// A policy for guest user access to a room.
#[doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/doc/string_enum.md"))]
#[derive(Clone, Debug, PartialEq, Eq, StringEnum)]
#[ruma_enum(rename_all = "snake_case")]
#[non_exhaustive]
pub enum GuestAccess {
    /// Guests are allowed to join the room.
    CanJoin,

    /// Guests are not allowed to join the room.
    Forbidden,

    #[doc(hidden)]
    _Custom(PrivOwnedStr),
}

impl GuestAccess {
    /// Creates a string slice from this `GuestAccess`.
    pub fn as_str(&self) -> &str {
        self.as_ref()
    }
}
<commit_msg>events: Add {RoomGuestAccessEvent, SyncRoomGuestAccessEvent}::guest_access<commit_after>//! Types for the [`m.room.guest_access`] event.
//!
//! [`m.room.guest_access`]: https://spec.matrix.org/v1.2/client-server-api/#mroomguest_access

use ruma_macros::EventContent;
use serde::{Deserialize, Serialize};

use crate::{serde::StringEnum, PrivOwnedStr};

/// The content of an `m.room.guest_access` event.
///
/// Controls whether guest users are allowed to join rooms.
///
/// This event controls whether guest users are allowed to join rooms. If this event is absent,
/// servers should act as if it is present and has the value `GuestAccess::Forbidden`.
#[derive(Clone, Debug, Deserialize, Serialize, EventContent)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
#[ruma_event(type = "m.room.guest_access", kind = State)]
pub struct RoomGuestAccessEventContent {
    /// A policy for guest user access to a room.
    pub guest_access: GuestAccess,
}

impl RoomGuestAccessEventContent {
    /// Creates a new `RoomGuestAccessEventContent` with the given policy.
    pub fn new(guest_access: GuestAccess) -> Self {
        Self { guest_access }
    }
}

impl RoomGuestAccessEvent {
    /// Obtain the guest access policy, regardless of whether this event is redacted.
    pub fn guest_access(&self) -> &GuestAccess {
        match self {
            Self::Original(ev) => &ev.content.guest_access,
            Self::Redacted(_) => &GuestAccess::Forbidden,
        }
    }
}

impl SyncRoomGuestAccessEvent {
    /// Obtain the guest access policy, regardless of whether this event is redacted.
    pub fn guest_access(&self) -> &GuestAccess {
        match self {
            Self::Original(ev) => &ev.content.guest_access,
            Self::Redacted(_) => &GuestAccess::Forbidden,
        }
    }
}

/// A policy for guest user access to a room.
#[doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/doc/string_enum.md"))]
#[derive(Clone, Debug, PartialEq, Eq, StringEnum)]
#[ruma_enum(rename_all = "snake_case")]
#[non_exhaustive]
pub enum GuestAccess {
    /// Guests are allowed to join the room.
    CanJoin,

    /// Guests are not allowed to join the room.
    Forbidden,

    #[doc(hidden)]
    _Custom(PrivOwnedStr),
}

impl GuestAccess {
    /// Creates a string slice from this `GuestAccess`.
    pub fn as_str(&self) -> &str {
        self.as_ref()
    }
}
<|endoftext|>


<commit_before>use alloc::boxed::Box;

use arch::memory;

use collections::slice;
use collections::string::ToString;
use collections::vec::Vec;
use collections::vec_deque::VecDeque;

use core::ptr;

use common::debug;

use drivers::pci::config::PciConfig;
use drivers::io::{Io, Pio};

use network::common::*;
use network::scheme::*;

use fs::{KScheme, Resource, Url};

use syscall::Result;

use sync::Intex;

const RTL8139_TSR_OWN: u32 = 1 << 13;

const RTL8139_CR_RST: u8 = 1 << 4;
const RTL8139_CR_RE: u8 = 1 << 3;
const RTL8139_CR_TE: u8 = 1 << 2;
const RTL8139_CR_BUFE: u8 = 1 << 0;

const RTL8139_ISR_SERR: u16 = 1 << 15;
const RTL8139_ISR_TIMEOUT: u16 = 1 << 14;
const RTL8139_ISR_LENCHG: u16 = 1 << 13;
const RTL8139_ISR_FOVW: u16 = 1 << 6;
const RTL8139_ISR_PUN_LINKCHG: u16 = 1 << 5;
const RTL8139_ISR_RXOVW: u16 = 1 << 4;
const RTL8139_ISR_TER: u16 = 1 << 3;
const RTL8139_ISR_TOK: u16 = 1 << 2;
const RTL8139_ISR_RER: u16 = 1 << 1;
const RTL8139_ISR_ROK: u16 = 1 << 0;

const RTL8139_TCR_IFG: u32 = 0b11 << 24;

const RTL8139_RCR_WRAP: u32 = 1 << 7;
const RTL8139_RCR_AR: u32 = 1 << 4;
const RTL8139_RCR_AB: u32 = 1 << 3;
const RTL8139_RCR_AM: u32 = 1 << 2;
const RTL8139_RCR_APM: u32 = 1 << 1;

#[repr(packed)]
struct Txd {
    pub address_port: Pio<u32>,
    pub status_port: Pio<u32>,
    pub buffer: usize,
}

pub struct Rtl8139Port {
    pub idr: [Pio<u8>; 6],
    pub rbstart: Pio<u32>,
    pub cr: Pio<u8>,
    pub capr: Pio<u16>,
    pub cbr: Pio<u16>,
    pub imr: Pio<u16>,
    pub isr: Pio<u16>,
    pub tcr: Pio<u32>,
    pub rcr: Pio<u32>,
    pub config1: Pio<u8>,
}

impl Rtl8139Port {
    pub fn new(base: u16) -> Self {
        return Rtl8139Port {
            idr: [Pio::<u8>::new(base + 0x00),
                  Pio::<u8>::new(base + 0x01),
                  Pio::<u8>::new(base + 0x02),
                  Pio::<u8>::new(base + 0x03),
                  Pio::<u8>::new(base + 0x04),
                  Pio::<u8>::new(base + 0x05)],
            rbstart: Pio::<u32>::new(base + 0x30),
            cr: Pio::<u8>::new(base + 0x37),
            capr: Pio::<u16>::new(base + 0x38),
            cbr: Pio::<u16>::new(base + 0x3A),
            imr: Pio::<u16>::new(base + 0x3C),
            isr: Pio::<u16>::new(base + 0x3E),
            tcr: Pio::<u32>::new(base + 0x40),
            rcr: Pio::<u32>::new(base + 0x44),
            config1: Pio::<u8>::new(base + 0x52),
        };
    }
}

pub struct Rtl8139 {
    pci: PciConfig,
    base: usize,
    memory_mapped: bool,
    irq: u8,
    resources: Intex<Vec<*mut NetworkResource>>,
    inbound: VecDeque<Vec<u8>>,
    outbound: VecDeque<Vec<u8>>,
    txds: Vec<Txd>,
    txd_i: usize,
    port: Rtl8139Port,
}

impl Rtl8139 {
    pub fn new(mut pci: PciConfig) -> Box<Self> {
        let pci_id = unsafe { pci.read(0x00) };
        let revision = (unsafe { pci.read(0x08) } & 0xFF) as u8;
        if pci_id == 0x813910EC && revision < 0x20 {
            debug::d("Not an 8139C+ compatible chip")
        }

        let base = unsafe { pci.read(0x10) as usize };
        let irq = unsafe { pci.read(0x3C) as u8 & 0xF };

        let mut module = box Rtl8139 {
            pci: pci,
            base: base & 0xFFFFFFF0,
            memory_mapped: base & 1 == 0,
            irq: irq,
            resources: Intex::new(Vec::new()),
            inbound: VecDeque::new(),
            outbound: VecDeque::new(),
            txds: Vec::new(),
            txd_i: 0,
            port: Rtl8139Port::new((base & 0xFFFFFFF0) as u16),
        };

        unsafe { module.init() };

        module
    }

    unsafe fn init(&mut self) {
        debug::d("RTL8139 on: ");
        debug::dh(self.base);
        if self.memory_mapped {
            debug::d(" memory mapped");
        } else {
            debug::d(" port mapped");
        }
        debug::d(" IRQ: ");
        debug::dbh(self.irq);

        self.pci.flag(4, 4, true); // Bus mastering

        let base = self.base as u16;

        self.port.config1.write(0);
        self.port.cr.write(RTL8139_CR_RST);
        while self.port.cr.read() & RTL8139_CR_RST != 0 {}

        debug::d(" MAC: ");
        MAC_ADDR = MacAddr {
            bytes: [self.port.idr[0].read(),
                    self.port.idr[1].read(),
                    self.port.idr[2].read(),
                    self.port.idr[3].read(),
                    self.port.idr[4].read(),
                    self.port.idr[5].read()],
        };
        debug::d(&MAC_ADDR.to_string());

        let receive_buffer = memory::alloc(10240);
        self.port.rbstart.write(receive_buffer as u32);

        for i in 0..4 {
            self.txds.push(Txd {
                address_port: Pio::<u32>::new(base + 0x20 + (i as u16) * 4),
                status_port: Pio::<u32>::new(base + 0x10 + (i as u16) * 4),
                buffer: memory::alloc(4096),
            });
        }

        self.port.imr.write(RTL8139_ISR_TOK | RTL8139_ISR_ROK);
        debug::d(" IMR: ");
        debug::dh(self.port.imr.read() as usize);

        self.port.cr.write(RTL8139_CR_RE | RTL8139_CR_TE);
        debug::d(" CMD: ");
        debug::dbh(self.port.cr.read());

        self.port.rcr.write(RTL8139_RCR_WRAP | RTL8139_RCR_AR | RTL8139_RCR_AB | RTL8139_RCR_AM |
                            RTL8139_RCR_APM);
        debug::d(" RCR: ");
        debug::dh(self.port.rcr.read() as usize);

        self.port.tcr.writef(RTL8139_TCR_IFG, true);
        debug::d(" TCR: ");
        debug::dh(self.port.tcr.read() as usize);

        debug::dl();
    }

    unsafe fn receive_inbound(&mut self) {
        let receive_buffer = self.port.rbstart.read() as usize;
        let mut capr = (self.port.capr.read() + 16) as usize;
        let cbr = self.port.cbr.read() as usize;

        while capr != cbr {
            let frame_addr = receive_buffer + capr + 4;
            let frame_status = ptr::read((receive_buffer + capr) as *const u16) as usize;
            let frame_len = ptr::read((receive_buffer + capr + 2) as *const u16) as usize;

            debug::d("Recv ");
            debug::dh(capr as usize);
            debug::d(" ");
            debug::dh(frame_status);
            debug::d(" ");
            debug::dh(frame_addr);
            debug::d(" ");
            debug::dh(frame_len);
            debug::dl();

            self.inbound.push_back(Vec::from(slice::from_raw_parts(frame_addr as *const u8, frame_len - 4)));

            capr = capr + frame_len + 4;
            capr = (capr + 3) & (0xFFFFFFFF - 3);
            if capr >= 8192 {
                capr -= 8192
            }

            self.port.capr.write((capr as u16) - 16);
        }
    }

    unsafe fn send_outbound(&mut self) {
        while let Some(bytes) = self.outbound.pop_front() {
            if let Some(ref mut txd) = self.txds.get_mut(self.txd_i) {
                if bytes.len() < 4096 {
                    while !txd.status_port.readf(RTL8139_TSR_OWN) {}

                    debug::d("Send ");
                    debug::dh(self.txd_i as usize);
                    debug::d(" ");
                    debug::dh(txd.status_port.read() as usize);
                    debug::d(" ");
                    debug::dh(txd.buffer);
                    debug::d(" ");
                    debug::dh(bytes.len() & 0xFFF);
                    debug::dl();

                    ::memcpy(txd.buffer as *mut u8, bytes.as_ptr(), bytes.len());

                    txd.address_port.write(txd.buffer as u32);
                    txd.status_port.write(bytes.len() as u32 & 0xFFF);

                    self.txd_i = (self.txd_i + 1) % 4;
                } else {
                    debug::dl();
                    debug::d("RTL8139: Frame too long for transmit: ");
                    debug::dd(bytes.len());
                    debug::dl();
                }
            } else {
                debug::d("RTL8139: TXD Overflow!\n");
                self.txd_i = 0;
            }
        }
    }
}

impl KScheme for Rtl8139 {
    fn scheme(&self) -> &str {
        "network"
    }

    fn open(&mut self, _: &Url, _: usize) -> Result<Box<Resource>> {
        Ok(NetworkResource::new(self))
    }

    fn on_irq(&mut self, irq: u8) {
        if irq == self.irq {
            let isr = self.port.isr.read();
            self.port.isr.write(isr);

            // dh(isr as usize);
            // dl();

            self.sync();
        }
    }

    fn on_poll(&mut self) {
        self.sync();
    }
}

impl NetworkScheme for Rtl8139 {
    fn add(&mut self, resource: *mut NetworkResource) {
        self.resources.lock().push(resource);
    }

    fn remove(&mut self, resource: *mut NetworkResource) {
        let mut resources = self.resources.lock();

        let mut i = 0;
        while i < resources.len() {
            let mut remove = false;

            match resources.get(i) {
                Some(ptr) => if *ptr == resource {
                    remove = true;
                } else {
                    i += 1;
                },
                None => break,
            }

            if remove {
                resources.remove(i);
            }
        }
    }

    fn sync(&mut self) {
        unsafe {
            {
                let resources = self.resources.lock();

                for resource in resources.iter() {
                    while let Some(bytes) = (**resource).outbound.lock().pop_front() {
                        self.outbound.push_back(bytes);
                    }
                }
            }

            self.send_outbound();

            self.receive_inbound();

            {
                let resources = self.resources.lock();

                while let Some(bytes) = self.inbound.pop_front() {
                    for resource in resources.iter() {
                        (**resource).inbound.lock().push_back(bytes.clone());
                    }
                }
            }
        }
    }
}
<commit_msg>Newline in rtl8139 driver<commit_after>use alloc::boxed::Box;

use arch::memory;

use collections::slice;
use collections::string::ToString;
use collections::vec::Vec;
use collections::vec_deque::VecDeque;

use core::ptr;

use common::debug;

use drivers::pci::config::PciConfig;
use drivers::io::{Io, Pio};

use network::common::*;
use network::scheme::*;

use fs::{KScheme, Resource, Url};

use syscall::Result;

use sync::Intex;

const RTL8139_TSR_OWN: u32 = 1 << 13;

const RTL8139_CR_RST: u8 = 1 << 4;
const RTL8139_CR_RE: u8 = 1 << 3;
const RTL8139_CR_TE: u8 = 1 << 2;
const RTL8139_CR_BUFE: u8 = 1 << 0;

const RTL8139_ISR_SERR: u16 = 1 << 15;
const RTL8139_ISR_TIMEOUT: u16 = 1 << 14;
const RTL8139_ISR_LENCHG: u16 = 1 << 13;
const RTL8139_ISR_FOVW: u16 = 1 << 6;
const RTL8139_ISR_PUN_LINKCHG: u16 = 1 << 5;
const RTL8139_ISR_RXOVW: u16 = 1 << 4;
const RTL8139_ISR_TER: u16 = 1 << 3;
const RTL8139_ISR_TOK: u16 = 1 << 2;
const RTL8139_ISR_RER: u16 = 1 << 1;
const RTL8139_ISR_ROK: u16 = 1 << 0;

const RTL8139_TCR_IFG: u32 = 0b11 << 24;

const RTL8139_RCR_WRAP: u32 = 1 << 7;
const RTL8139_RCR_AR: u32 = 1 << 4;
const RTL8139_RCR_AB: u32 = 1 << 3;
const RTL8139_RCR_AM: u32 = 1 << 2;
const RTL8139_RCR_APM: u32 = 1 << 1;

#[repr(packed)]
struct Txd {
    pub address_port: Pio<u32>,
    pub status_port: Pio<u32>,
    pub buffer: usize,
}

pub struct Rtl8139Port {
    pub idr: [Pio<u8>; 6],
    pub rbstart: Pio<u32>,
    pub cr: Pio<u8>,
    pub capr: Pio<u16>,
    pub cbr: Pio<u16>,
    pub imr: Pio<u16>,
    pub isr: Pio<u16>,
    pub tcr: Pio<u32>,
    pub rcr: Pio<u32>,
    pub config1: Pio<u8>,
}

impl Rtl8139Port {
    pub fn new(base: u16) -> Self {
        return Rtl8139Port {
            idr: [Pio::<u8>::new(base + 0x00),
                  Pio::<u8>::new(base + 0x01),
                  Pio::<u8>::new(base + 0x02),
                  Pio::<u8>::new(base + 0x03),
                  Pio::<u8>::new(base + 0x04),
                  Pio::<u8>::new(base + 0x05)],
            rbstart: Pio::<u32>::new(base + 0x30),
            cr: Pio::<u8>::new(base + 0x37),
            capr: Pio::<u16>::new(base + 0x38),
            cbr: Pio::<u16>::new(base + 0x3A),
            imr: Pio::<u16>::new(base + 0x3C),
            isr: Pio::<u16>::new(base + 0x3E),
            tcr: Pio::<u32>::new(base + 0x40),
            rcr: Pio::<u32>::new(base + 0x44),
            config1: Pio::<u8>::new(base + 0x52),
        };
    }
}

pub struct Rtl8139 {
    pci: PciConfig,
    base: usize,
    memory_mapped: bool,
    irq: u8,
    resources: Intex<Vec<*mut NetworkResource>>,
    inbound: VecDeque<Vec<u8>>,
    outbound: VecDeque<Vec<u8>>,
    txds: Vec<Txd>,
    txd_i: usize,
    port: Rtl8139Port,
}

impl Rtl8139 {
    pub fn new(mut pci: PciConfig) -> Box<Self> {
        let pci_id = unsafe { pci.read(0x00) };
        let revision = (unsafe { pci.read(0x08) } & 0xFF) as u8;
        if pci_id == 0x813910EC && revision < 0x20 {
            debugln!("Not an 8139C+ compatible chip")
        }

        let base = unsafe { pci.read(0x10) as usize };
        let irq = unsafe { pci.read(0x3C) as u8 & 0xF };

        let mut module = box Rtl8139 {
            pci: pci,
            base: base & 0xFFFFFFF0,
            memory_mapped: base & 1 == 0,
            irq: irq,
            resources: Intex::new(Vec::new()),
            inbound: VecDeque::new(),
            outbound: VecDeque::new(),
            txds: Vec::new(),
            txd_i: 0,
            port: Rtl8139Port::new((base & 0xFFFFFFF0) as u16),
        };

        unsafe { module.init() };

        module
    }

    unsafe fn init(&mut self) {
        debug::d("RTL8139 on: ");
        debug::dh(self.base);
        if self.memory_mapped {
            debug::d(" memory mapped");
        } else {
            debug::d(" port mapped");
        }
        debug::d(" IRQ: ");
        debug::dbh(self.irq);

        self.pci.flag(4, 4, true); // Bus mastering

        let base = self.base as u16;

        self.port.config1.write(0);
        self.port.cr.write(RTL8139_CR_RST);
        while self.port.cr.read() & RTL8139_CR_RST != 0 {}

        debug::d(" MAC: ");
        MAC_ADDR = MacAddr {
            bytes: [self.port.idr[0].read(),
                    self.port.idr[1].read(),
                    self.port.idr[2].read(),
                    self.port.idr[3].read(),
                    self.port.idr[4].read(),
                    self.port.idr[5].read()],
        };
        debug::d(&MAC_ADDR.to_string());

        let receive_buffer = memory::alloc(10240);
        self.port.rbstart.write(receive_buffer as u32);

        for i in 0..4 {
            self.txds.push(Txd {
                address_port: Pio::<u32>::new(base + 0x20 + (i as u16) * 4),
                status_port: Pio::<u32>::new(base + 0x10 + (i as u16) * 4),
                buffer: memory::alloc(4096),
            });
        }

        self.port.imr.write(RTL8139_ISR_TOK | RTL8139_ISR_ROK);
        debug::d(" IMR: ");
        debug::dh(self.port.imr.read() as usize);

        self.port.cr.write(RTL8139_CR_RE | RTL8139_CR_TE);
        debug::d(" CMD: ");
        debug::dbh(self.port.cr.read());

        self.port.rcr.write(RTL8139_RCR_WRAP | RTL8139_RCR_AR | RTL8139_RCR_AB | RTL8139_RCR_AM |
                            RTL8139_RCR_APM);
        debug::d(" RCR: ");
        debug::dh(self.port.rcr.read() as usize);

        self.port.tcr.writef(RTL8139_TCR_IFG, true);
        debug::d(" TCR: ");
        debug::dh(self.port.tcr.read() as usize);

        debug::dl();
    }

    unsafe fn receive_inbound(&mut self) {
        let receive_buffer = self.port.rbstart.read() as usize;
        let mut capr = (self.port.capr.read() + 16) as usize;
        let cbr = self.port.cbr.read() as usize;

        while capr != cbr {
            let frame_addr = receive_buffer + capr + 4;
            let frame_status = ptr::read((receive_buffer + capr) as *const u16) as usize;
            let frame_len = ptr::read((receive_buffer + capr + 2) as *const u16) as usize;

            debug::d("Recv ");
            debug::dh(capr as usize);
            debug::d(" ");
            debug::dh(frame_status);
            debug::d(" ");
            debug::dh(frame_addr);
            debug::d(" ");
            debug::dh(frame_len);
            debug::dl();

            self.inbound.push_back(Vec::from(slice::from_raw_parts(frame_addr as *const u8, frame_len - 4)));

            capr = capr + frame_len + 4;
            capr = (capr + 3) & (0xFFFFFFFF - 3);
            if capr >= 8192 {
                capr -= 8192
            }

            self.port.capr.write((capr as u16) - 16);
        }
    }

    unsafe fn send_outbound(&mut self) {
        while let Some(bytes) = self.outbound.pop_front() {
            if let Some(ref mut txd) = self.txds.get_mut(self.txd_i) {
                if bytes.len() < 4096 {
                    while !txd.status_port.readf(RTL8139_TSR_OWN) {}

                    debug::d("Send ");
                    debug::dh(self.txd_i as usize);
                    debug::d(" ");
                    debug::dh(txd.status_port.read() as usize);
                    debug::d(" ");
                    debug::dh(txd.buffer);
                    debug::d(" ");
                    debug::dh(bytes.len() & 0xFFF);
                    debug::dl();

                    ::memcpy(txd.buffer as *mut u8, bytes.as_ptr(), bytes.len());

                    txd.address_port.write(txd.buffer as u32);
                    txd.status_port.write(bytes.len() as u32 & 0xFFF);

                    self.txd_i = (self.txd_i + 1) % 4;
                } else {
                    debug::dl();
                    debug::d("RTL8139: Frame too long for transmit: ");
                    debug::dd(bytes.len());
                    debug::dl();
                }
            } else {
                debug::d("RTL8139: TXD Overflow!\n");
                self.txd_i = 0;
            }
        }
    }
}

impl KScheme for Rtl8139 {
    fn scheme(&self) -> &str {
        "network"
    }

    fn open(&mut self, _: &Url, _: usize) -> Result<Box<Resource>> {
        Ok(NetworkResource::new(self))
    }

    fn on_irq(&mut self, irq: u8) {
        if irq == self.irq {
            let isr = self.port.isr.read();
            self.port.isr.write(isr);

            // dh(isr as usize);
            // dl();

            self.sync();
        }
    }

    fn on_poll(&mut self) {
        self.sync();
    }
}

impl NetworkScheme for Rtl8139 {
    fn add(&mut self, resource: *mut NetworkResource) {
        self.resources.lock().push(resource);
    }

    fn remove(&mut self, resource: *mut NetworkResource) {
        let mut resources = self.resources.lock();

        let mut i = 0;
        while i < resources.len() {
            let mut remove = false;

            match resources.get(i) {
                Some(ptr) => if *ptr == resource {
                    remove = true;
                } else {
                    i += 1;
                },
                None => break,
            }

            if remove {
                resources.remove(i);
            }
        }
    }

    fn sync(&mut self) {
        unsafe {
            {
                let resources = self.resources.lock();

                for resource in resources.iter() {
                    while let Some(bytes) = (**resource).outbound.lock().pop_front() {
                        self.outbound.push_back(bytes);
                    }
                }
            }

            self.send_outbound();

            self.receive_inbound();

            {
                let resources = self.resources.lock();

                while let Some(bytes) = self.inbound.pop_front() {
                    for resource in resources.iter() {
                        (**resource).inbound.lock().push_back(bytes.clone());
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2013 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.

//! The Gamma and derived distributions.

use std::num::Float;
use {Rng, Open01};
use super::normal::StandardNormal;
use super::{IndependentSample, Sample, Exp};

/// The Gamma distribution `Gamma(shape, scale)` distribution.
///
/// The density function of this distribution is
///
/// ```notrust
/// f(x) =  x^(k - 1) * exp(-x / θ) / (Γ(k) * θ^k)
/// ```
///
/// where `Γ` is the Gamma function, `k` is the shape and `θ` is the
/// scale and both `k` and `θ` are strictly positive.
///
/// The algorithm used is that described by Marsaglia & Tsang 2000[1],
/// falling back to directly sampling from an Exponential for `shape
/// == 1`, and using the boosting technique described in [1] for
/// `shape < 1`.
///
/// # Example
///
/// ```rust
/// use rand::distributions::{IndependentSample, Gamma};
///
/// let gamma = Gamma::new(2.0, 5.0);
/// let v = gamma.ind_sample(&mut rand::task_rng());
/// println!("{} is from a Gamma(2, 5) distribution", v);
/// ```
///
/// [1]: George Marsaglia and Wai Wan Tsang. 2000. "A Simple Method
/// for Generating Gamma Variables" *ACM Trans. Math. Softw.* 26, 3
/// (September 2000),
/// 363-372. DOI:[10.1145/358407.358414](http://doi.acm.org/10.1145/358407.358414)
pub struct Gamma {
    repr: GammaRepr,
}

enum GammaRepr {
    Large(GammaLargeShape),
    One(Exp),
    Small(GammaSmallShape)
}

// These two helpers could be made public, but saving the
// match-on-Gamma-enum branch from using them directly (e.g. if one
// knows that the shape is always > 1) doesn't appear to be much
// faster.

/// Gamma distribution where the shape parameter is less than 1.
///
/// Note, samples from this require a compulsory floating-point `pow`
/// call, which makes it significantly slower than sampling from a
/// gamma distribution where the shape parameter is greater than or
/// equal to 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaSmallShape {
    inv_shape: f64,
    large_shape: GammaLargeShape
}

/// Gamma distribution where the shape parameter is larger than 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaLargeShape {
    shape: f64,
    scale: f64,
    c: f64,
    d: f64
}

impl Gamma {
    /// Construct an object representing the `Gamma(shape, scale)`
    /// distribution.
    ///
    /// Fails if `shape <= 0` or `scale <= 0`.
    pub fn new(shape: f64, scale: f64) -> Gamma {
        assert!(shape > 0.0, "Gamma::new called with shape <= 0");
        assert!(scale > 0.0, "Gamma::new called with scale <= 0");

        let repr = match shape {
            1.0        => One(Exp::new(1.0 / scale)),
            0.0 .. 1.0 => Small(GammaSmallShape::new_raw(shape, scale)),
            _          => Large(GammaLargeShape::new_raw(shape, scale))
        };
        Gamma { repr: repr }
    }
}

impl GammaSmallShape {
    fn new_raw(shape: f64, scale: f64) -> GammaSmallShape {
        GammaSmallShape {
            inv_shape: 1. / shape,
            large_shape: GammaLargeShape::new_raw(shape + 1.0, scale)
        }
    }
}

impl GammaLargeShape {
    fn new_raw(shape: f64, scale: f64) -> GammaLargeShape {
        let d = shape - 1. / 3.;
        GammaLargeShape {
            shape: shape,
            scale: scale,
            c: 1. / (9. * d).sqrt(),
            d: d
        }
    }
}

impl Sample<f64> for Gamma {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaSmallShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaLargeShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}

impl IndependentSample<f64> for Gamma {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match self.repr {
            Small(ref g) => g.ind_sample(rng),
            One(ref g) => g.ind_sample(rng),
            Large(ref g) => g.ind_sample(rng),
        }
    }
}
impl IndependentSample<f64> for GammaSmallShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let Open01(u) = rng.gen::<Open01<f64>>();

        self.large_shape.ind_sample(rng) * u.powf(&self.inv_shape)
    }
}
impl IndependentSample<f64> for GammaLargeShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        loop {
            let StandardNormal(x) = rng.gen::<StandardNormal>();
            let v_cbrt = 1.0 + self.c * x;
            if v_cbrt <= 0.0 { // a^3 <= 0 iff a <= 0
                continue
            }

            let v = v_cbrt * v_cbrt * v_cbrt;
            let Open01(u) = rng.gen::<Open01<f64>>();

            let x_sqr = x * x;
            if u < 1.0 - 0.0331 * x_sqr * x_sqr ||
                u.ln() < 0.5 * x_sqr + self.d * (1.0 - v + v.ln()) {
                return self.d * v * self.scale
            }
        }
    }
}

/// The chi-squared distribution `χ²(k)`, where `k` is the degrees of
/// freedom.
///
/// For `k > 0` integral, this distribution is the sum of the squares
/// of `k` independent standard normal random variables. For other
/// `k`, this uses the equivalent characterisation `χ²(k) = Gamma(k/2,
/// 2)`.
///
/// # Example
///
/// ```rust
/// use rand::distributions::{ChiSquared, IndependentSample};
///
/// let chi = ChiSquared::new(11.0);
/// let v = chi.ind_sample(&mut rand::task_rng());
/// println!("{} is from a χ²(11) distribution", v)
/// ```
pub struct ChiSquared {
    repr: ChiSquaredRepr,
}

enum ChiSquaredRepr {
    // k == 1, Gamma(alpha, ..) is particularly slow for alpha < 1,
    // e.g. when alpha = 1/2 as it would be for this case, so special-
    // casing and using the definition of N(0,1)^2 is faster.
    DoFExactlyOne,
    DoFAnythingElse(Gamma),
}

impl ChiSquared {
    /// Create a new chi-squared distribution with degrees-of-freedom
    /// `k`. Fails if `k < 0`.
    pub fn new(k: f64) -> ChiSquared {
        let repr = if k == 1.0 {
            DoFExactlyOne
        } else {
            assert!(k > 0.0, "ChiSquared::new called with `k` < 0");
            DoFAnythingElse(Gamma::new(0.5 * k, 2.0))
        };
        ChiSquared { repr: repr }
    }
}
impl Sample<f64> for ChiSquared {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for ChiSquared {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match self.repr {
            DoFExactlyOne => {
                // k == 1 => N(0,1)^2
                let StandardNormal(norm) = rng.gen::<StandardNormal>();
                norm * norm
            }
            DoFAnythingElse(ref g) => g.ind_sample(rng)
        }
    }
}

/// The Fisher F distribution `F(m, n)`.
///
/// This distribution is equivalent to the ratio of two normalised
/// chi-squared distributions, that is, `F(m,n) = (χ²(m)/m) /
/// (χ²(n)/n)`.
///
/// # Example
///
/// ```rust
/// use rand::distributions::{FisherF, IndependentSample};
///
/// let f = FisherF::new(2.0, 32.0);
/// let v = f.ind_sample(&mut rand::task_rng());
/// println!("{} is from an F(2, 32) distribution", v)
/// ```
pub struct FisherF {
    numer: ChiSquared,
    denom: ChiSquared,
    // denom_dof / numer_dof so that this can just be a straight
    // multiplication, rather than a division.
    dof_ratio: f64,
}

impl FisherF {
    /// Create a new `FisherF` distribution, with the given
    /// parameter. Fails if either `m` or `n` are not positive.
    pub fn new(m: f64, n: f64) -> FisherF {
        assert!(m > 0.0, "FisherF::new called with `m < 0`");
        assert!(n > 0.0, "FisherF::new called with `n < 0`");

        FisherF {
            numer: ChiSquared::new(m),
            denom: ChiSquared::new(n),
            dof_ratio: n / m
        }
    }
}
impl Sample<f64> for FisherF {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for FisherF {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        self.numer.ind_sample(rng) / self.denom.ind_sample(rng) * self.dof_ratio
    }
}

/// The Student t distribution, `t(nu)`, where `nu` is the degrees of
/// freedom.
///
/// # Example
///
/// ```rust
/// use rand::distributions::{StudentT, IndependentSample};
///
/// let t = StudentT::new(11.0);
/// let v = t.ind_sample(&mut rand::task_rng());
/// println!("{} is from a t(11) distribution", v)
/// ```
pub struct StudentT {
    chi: ChiSquared,
    dof: f64
}

impl StudentT {
    /// Create a new Student t distribution with `n` degrees of
    /// freedom. Fails if `n <= 0`.
    pub fn new(n: f64) -> StudentT {
        assert!(n > 0.0, "StudentT::new called with `n <= 0`");
        StudentT {
            chi: ChiSquared::new(n),
            dof: n
        }
    }
}
impl Sample<f64> for StudentT {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for StudentT {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let StandardNormal(norm) = rng.gen::<StandardNormal>();
        norm * (self.dof / self.chi.ind_sample(rng)).sqrt()
    }
}

#[cfg(test)]
mod test {
    use distributions::{Sample, IndependentSample};
    use {Rng, task_rng};
    use super::{ChiSquared, StudentT, FisherF};

    #[test]
    fn test_chi_squared_one() {
        let mut chi = ChiSquared::new(1.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_small() {
        let mut chi = ChiSquared::new(0.5);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_large() {
        let mut chi = ChiSquared::new(30.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    #[should_fail]
    fn test_chi_squared_invalid_dof() {
        ChiSquared::new(-1.0);
    }

    #[test]
    fn test_f() {
        let mut f = FisherF::new(2.0, 32.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            f.sample(&mut rng);
            f.ind_sample(&mut rng);
        }
    }

    #[test]
    fn test_t() {
        let mut t = StudentT::new(11.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            t.sample(&mut rng);
            t.ind_sample(&mut rng);
        }
    }
}

#[cfg(test)]
mod bench {
    extern crate test;
    use self::test::Bencher;
    use std::mem::size_of;
    use distributions::IndependentSample;
    use {XorShiftRng, RAND_BENCH_N};
    use super::Gamma;

    #[bench]
    fn bench_gamma_large_shape(b: &mut Bencher) {
        let gamma = Gamma::new(10., 1.0);
        let mut rng = XorShiftRng::new().unwrap();

        b.iter(|| {
            for _ in range(0, RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * RAND_BENCH_N;
    }

    #[bench]
    fn bench_gamma_small_shape(b: &mut Bencher) {
        let gamma = Gamma::new(0.1, 1.0);
        let mut rng = XorShiftRng::new().unwrap();

        b.iter(|| {
            for _ in range(0, RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * RAND_BENCH_N;
    }
}
<commit_msg>Have floating point functions take their parameters by value.<commit_after>// Copyright 2013 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.

//! The Gamma and derived distributions.

use std::num::Float;
use {Rng, Open01};
use super::normal::StandardNormal;
use super::{IndependentSample, Sample, Exp};

/// The Gamma distribution `Gamma(shape, scale)` distribution.
///
/// The density function of this distribution is
///
/// ```notrust
/// f(x) =  x^(k - 1) * exp(-x / θ) / (Γ(k) * θ^k)
/// ```
///
/// where `Γ` is the Gamma function, `k` is the shape and `θ` is the
/// scale and both `k` and `θ` are strictly positive.
///
/// The algorithm used is that described by Marsaglia & Tsang 2000[1],
/// falling back to directly sampling from an Exponential for `shape
/// == 1`, and using the boosting technique described in [1] for
/// `shape < 1`.
///
/// # Example
///
/// ```rust
/// use rand::distributions::{IndependentSample, Gamma};
///
/// let gamma = Gamma::new(2.0, 5.0);
/// let v = gamma.ind_sample(&mut rand::task_rng());
/// println!("{} is from a Gamma(2, 5) distribution", v);
/// ```
///
/// [1]: George Marsaglia and Wai Wan Tsang. 2000. "A Simple Method
/// for Generating Gamma Variables" *ACM Trans. Math. Softw.* 26, 3
/// (September 2000),
/// 363-372. DOI:[10.1145/358407.358414](http://doi.acm.org/10.1145/358407.358414)
pub struct Gamma {
    repr: GammaRepr,
}

enum GammaRepr {
    Large(GammaLargeShape),
    One(Exp),
    Small(GammaSmallShape)
}

// These two helpers could be made public, but saving the
// match-on-Gamma-enum branch from using them directly (e.g. if one
// knows that the shape is always > 1) doesn't appear to be much
// faster.

/// Gamma distribution where the shape parameter is less than 1.
///
/// Note, samples from this require a compulsory floating-point `pow`
/// call, which makes it significantly slower than sampling from a
/// gamma distribution where the shape parameter is greater than or
/// equal to 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaSmallShape {
    inv_shape: f64,
    large_shape: GammaLargeShape
}

/// Gamma distribution where the shape parameter is larger than 1.
///
/// See `Gamma` for sampling from a Gamma distribution with general
/// shape parameters.
struct GammaLargeShape {
    shape: f64,
    scale: f64,
    c: f64,
    d: f64
}

impl Gamma {
    /// Construct an object representing the `Gamma(shape, scale)`
    /// distribution.
    ///
    /// Fails if `shape <= 0` or `scale <= 0`.
    pub fn new(shape: f64, scale: f64) -> Gamma {
        assert!(shape > 0.0, "Gamma::new called with shape <= 0");
        assert!(scale > 0.0, "Gamma::new called with scale <= 0");

        let repr = match shape {
            1.0        => One(Exp::new(1.0 / scale)),
            0.0 .. 1.0 => Small(GammaSmallShape::new_raw(shape, scale)),
            _          => Large(GammaLargeShape::new_raw(shape, scale))
        };
        Gamma { repr: repr }
    }
}

impl GammaSmallShape {
    fn new_raw(shape: f64, scale: f64) -> GammaSmallShape {
        GammaSmallShape {
            inv_shape: 1. / shape,
            large_shape: GammaLargeShape::new_raw(shape + 1.0, scale)
        }
    }
}

impl GammaLargeShape {
    fn new_raw(shape: f64, scale: f64) -> GammaLargeShape {
        let d = shape - 1. / 3.;
        GammaLargeShape {
            shape: shape,
            scale: scale,
            c: 1. / (9. * d).sqrt(),
            d: d
        }
    }
}

impl Sample<f64> for Gamma {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaSmallShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl Sample<f64> for GammaLargeShape {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}

impl IndependentSample<f64> for Gamma {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match self.repr {
            Small(ref g) => g.ind_sample(rng),
            One(ref g) => g.ind_sample(rng),
            Large(ref g) => g.ind_sample(rng),
        }
    }
}
impl IndependentSample<f64> for GammaSmallShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let Open01(u) = rng.gen::<Open01<f64>>();

        self.large_shape.ind_sample(rng) * u.powf(self.inv_shape)
    }
}
impl IndependentSample<f64> for GammaLargeShape {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        loop {
            let StandardNormal(x) = rng.gen::<StandardNormal>();
            let v_cbrt = 1.0 + self.c * x;
            if v_cbrt <= 0.0 { // a^3 <= 0 iff a <= 0
                continue
            }

            let v = v_cbrt * v_cbrt * v_cbrt;
            let Open01(u) = rng.gen::<Open01<f64>>();

            let x_sqr = x * x;
            if u < 1.0 - 0.0331 * x_sqr * x_sqr ||
                u.ln() < 0.5 * x_sqr + self.d * (1.0 - v + v.ln()) {
                return self.d * v * self.scale
            }
        }
    }
}

/// The chi-squared distribution `χ²(k)`, where `k` is the degrees of
/// freedom.
///
/// For `k > 0` integral, this distribution is the sum of the squares
/// of `k` independent standard normal random variables. For other
/// `k`, this uses the equivalent characterisation `χ²(k) = Gamma(k/2,
/// 2)`.
///
/// # Example
///
/// ```rust
/// use rand::distributions::{ChiSquared, IndependentSample};
///
/// let chi = ChiSquared::new(11.0);
/// let v = chi.ind_sample(&mut rand::task_rng());
/// println!("{} is from a χ²(11) distribution", v)
/// ```
pub struct ChiSquared {
    repr: ChiSquaredRepr,
}

enum ChiSquaredRepr {
    // k == 1, Gamma(alpha, ..) is particularly slow for alpha < 1,
    // e.g. when alpha = 1/2 as it would be for this case, so special-
    // casing and using the definition of N(0,1)^2 is faster.
    DoFExactlyOne,
    DoFAnythingElse(Gamma),
}

impl ChiSquared {
    /// Create a new chi-squared distribution with degrees-of-freedom
    /// `k`. Fails if `k < 0`.
    pub fn new(k: f64) -> ChiSquared {
        let repr = if k == 1.0 {
            DoFExactlyOne
        } else {
            assert!(k > 0.0, "ChiSquared::new called with `k` < 0");
            DoFAnythingElse(Gamma::new(0.5 * k, 2.0))
        };
        ChiSquared { repr: repr }
    }
}
impl Sample<f64> for ChiSquared {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for ChiSquared {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        match self.repr {
            DoFExactlyOne => {
                // k == 1 => N(0,1)^2
                let StandardNormal(norm) = rng.gen::<StandardNormal>();
                norm * norm
            }
            DoFAnythingElse(ref g) => g.ind_sample(rng)
        }
    }
}

/// The Fisher F distribution `F(m, n)`.
///
/// This distribution is equivalent to the ratio of two normalised
/// chi-squared distributions, that is, `F(m,n) = (χ²(m)/m) /
/// (χ²(n)/n)`.
///
/// # Example
///
/// ```rust
/// use rand::distributions::{FisherF, IndependentSample};
///
/// let f = FisherF::new(2.0, 32.0);
/// let v = f.ind_sample(&mut rand::task_rng());
/// println!("{} is from an F(2, 32) distribution", v)
/// ```
pub struct FisherF {
    numer: ChiSquared,
    denom: ChiSquared,
    // denom_dof / numer_dof so that this can just be a straight
    // multiplication, rather than a division.
    dof_ratio: f64,
}

impl FisherF {
    /// Create a new `FisherF` distribution, with the given
    /// parameter. Fails if either `m` or `n` are not positive.
    pub fn new(m: f64, n: f64) -> FisherF {
        assert!(m > 0.0, "FisherF::new called with `m < 0`");
        assert!(n > 0.0, "FisherF::new called with `n < 0`");

        FisherF {
            numer: ChiSquared::new(m),
            denom: ChiSquared::new(n),
            dof_ratio: n / m
        }
    }
}
impl Sample<f64> for FisherF {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for FisherF {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        self.numer.ind_sample(rng) / self.denom.ind_sample(rng) * self.dof_ratio
    }
}

/// The Student t distribution, `t(nu)`, where `nu` is the degrees of
/// freedom.
///
/// # Example
///
/// ```rust
/// use rand::distributions::{StudentT, IndependentSample};
///
/// let t = StudentT::new(11.0);
/// let v = t.ind_sample(&mut rand::task_rng());
/// println!("{} is from a t(11) distribution", v)
/// ```
pub struct StudentT {
    chi: ChiSquared,
    dof: f64
}

impl StudentT {
    /// Create a new Student t distribution with `n` degrees of
    /// freedom. Fails if `n <= 0`.
    pub fn new(n: f64) -> StudentT {
        assert!(n > 0.0, "StudentT::new called with `n <= 0`");
        StudentT {
            chi: ChiSquared::new(n),
            dof: n
        }
    }
}
impl Sample<f64> for StudentT {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
}
impl IndependentSample<f64> for StudentT {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
        let StandardNormal(norm) = rng.gen::<StandardNormal>();
        norm * (self.dof / self.chi.ind_sample(rng)).sqrt()
    }
}

#[cfg(test)]
mod test {
    use distributions::{Sample, IndependentSample};
    use {Rng, task_rng};
    use super::{ChiSquared, StudentT, FisherF};

    #[test]
    fn test_chi_squared_one() {
        let mut chi = ChiSquared::new(1.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_small() {
        let mut chi = ChiSquared::new(0.5);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    fn test_chi_squared_large() {
        let mut chi = ChiSquared::new(30.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            chi.sample(&mut rng);
            chi.ind_sample(&mut rng);
        }
    }
    #[test]
    #[should_fail]
    fn test_chi_squared_invalid_dof() {
        ChiSquared::new(-1.0);
    }

    #[test]
    fn test_f() {
        let mut f = FisherF::new(2.0, 32.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            f.sample(&mut rng);
            f.ind_sample(&mut rng);
        }
    }

    #[test]
    fn test_t() {
        let mut t = StudentT::new(11.0);
        let mut rng = task_rng();
        for _ in range(0, 1000) {
            t.sample(&mut rng);
            t.ind_sample(&mut rng);
        }
    }
}

#[cfg(test)]
mod bench {
    extern crate test;
    use self::test::Bencher;
    use std::mem::size_of;
    use distributions::IndependentSample;
    use {XorShiftRng, RAND_BENCH_N};
    use super::Gamma;

    #[bench]
    fn bench_gamma_large_shape(b: &mut Bencher) {
        let gamma = Gamma::new(10., 1.0);
        let mut rng = XorShiftRng::new().unwrap();

        b.iter(|| {
            for _ in range(0, RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * RAND_BENCH_N;
    }

    #[bench]
    fn bench_gamma_small_shape(b: &mut Bencher) {
        let gamma = Gamma::new(0.1, 1.0);
        let mut rng = XorShiftRng::new().unwrap();

        b.iter(|| {
            for _ in range(0, RAND_BENCH_N) {
                gamma.ind_sample(&mut rng);
            }
        });
        b.bytes = size_of::<f64>() as u64 * RAND_BENCH_N;
    }
}
<|endoftext|>


<commit_before>use rustc_hir::Mutability;
use rustc_index::bit_set::HybridBitSet;
use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
use rustc_middle::mir::{self, BasicBlock, Local, Location};
use rustc_middle::ty::TyCtxt;

use crate::transform::{MirPass, MirSource};

/// This pass looks for MIR that always copies the same local into the return place and eliminates
/// the copy by renaming all uses of that local to `_0`.
///
/// This allows LLVM to perform an optimization similar to the named return value optimization
/// (NRVO) that is guaranteed in C++. This avoids a stack allocation and `memcpy` for the
/// relatively common pattern of allocating a buffer on the stack, mutating it, and returning it by
/// value like so:
///
/// ```rust
/// fn foo(init: fn(&mut [u8; 1024])) -> [u8; 1024] {
///     let mut buf = [0; 1024];
///     init(&mut buf);
///     buf
/// }
/// ```
///
/// For now, this pass is very simple and only capable of eliminating a single copy. A more general
/// version of copy propagation, such as the one based on non-overlapping live ranges in [#47954] and
/// [#71003], could yield even more benefits.
///
/// [#47954]: https://github.com/rust-lang/rust/pull/47954
/// [#71003]: https://github.com/rust-lang/rust/pull/71003
pub struct RenameReturnPlace;

impl<'tcx> MirPass<'tcx> for RenameReturnPlace {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, body: &mut mir::Body<'tcx>) {
        if tcx.sess.opts.debugging_opts.mir_opt_level == 0 {
            return;
        }

        let returned_local = match local_eligible_for_nrvo(body) {
            Some(l) => l,
            None => {
                debug!("`{:?}` was ineligible for NRVO", src.def_id());
                return;
            }
        };

        // Sometimes, the return place is assigned a local of a different but coercable type, for
        // example `&T` instead of `&mut T`. Overwriting the `LocalInfo` for the return place would
        // result in it having an incorrect type. Although this doesn't seem to cause a problem in
        // codegen, bail out anyways since it happens so rarely.
        let ret_ty = body.local_decls[mir::RETURN_PLACE].ty;
        let assigned_ty = body.local_decls[returned_local].ty;
        if ret_ty != assigned_ty {
            debug!("`{:?}` was eligible for NRVO but for type mismatch", src.def_id());
            debug!("typeof(_0) != typeof({:?}); {:?} != {:?}", returned_local, ret_ty, assigned_ty);
            return;
        }

        debug!(
            "`{:?}` was eligible for NRVO, making {:?} the return place",
            src.def_id(),
            returned_local
        );

        RenameToReturnPlace { tcx, to_rename: returned_local }.visit_body(body);

        // Clean up the `NOP`s we inserted for statements made useless by our renaming.
        for block_data in body.basic_blocks_mut() {
            block_data.statements.retain(|stmt| stmt.kind != mir::StatementKind::Nop);
        }

        // Overwrite the debuginfo of `_0` with that of the renamed local.
        let (renamed_decl, ret_decl) =
            body.local_decls.pick2_mut(returned_local, mir::RETURN_PLACE);
        ret_decl.clone_from(renamed_decl);

        // The return place is always mutable.
        ret_decl.mutability = Mutability::Mut;
    }
}

/// MIR that is eligible for the NRVO must fulfill two conditions:
///   1. The return place must not be read prior to the `Return` terminator.
///   2. A simple assignment of a whole local to the return place (e.g., `_0 = _1`) must be the
///      only definition of the return place reaching the `Return` terminator.
///
/// If the MIR fulfills both these conditions, this function returns the `Local` that is assigned
/// to the return place along all possible paths through the control-flow graph.
fn local_eligible_for_nrvo(body: &mut mir::Body<'_>) -> Option<Local> {
    if IsReturnPlaceRead::run(body) {
        return None;
    }

    let mut copied_to_return_place = None;
    for block in body.basic_blocks().indices() {
        // Look for blocks with a `Return` terminator.
        if !matches!(body[block].terminator().kind, mir::TerminatorKind::Return) {
            continue;
        }

        // Look for an assignment of a single local to the return place prior to the `Return`.
        let returned_local = find_local_assigned_to_return_place(block, body)?;
        match body.local_kind(returned_local) {
            // FIXME: Can we do this for arguments as well?
            mir::LocalKind::Arg => return None,

            mir::LocalKind::ReturnPointer => bug!("Return place was assigned to itself?"),
            mir::LocalKind::Var | mir::LocalKind::Temp => {}
        }

        // If multiple different locals are copied to the return place. We can't pick a
        // single one to rename.
        if copied_to_return_place.map_or(false, |old| old != returned_local) {
            return None;
        }

        copied_to_return_place = Some(returned_local);
    }

    return copied_to_return_place;
}

fn find_local_assigned_to_return_place(
    start: BasicBlock,
    body: &mut mir::Body<'_>,
) -> Option<Local> {
    let mut block = start;
    let mut seen = HybridBitSet::new_empty(body.basic_blocks().len());

    // Iterate as long as `block` has exactly one predecessor that we have not yet visited.
    while seen.insert(block) {
        trace!("Looking for assignments to `_0` in {:?}", block);

        let local = body[block].statements.iter().rev().find_map(as_local_assigned_to_return_place);
        if local.is_some() {
            return local;
        }

        match body.predecessors()[block].as_slice() {
            &[pred] => block = pred,
            _ => return None,
        }
    }

    return None;
}

// If this statement is an assignment of an unprojected local to the return place,
// return that local.
fn as_local_assigned_to_return_place(stmt: &mir::Statement<'_>) -> Option<Local> {
    if let mir::StatementKind::Assign(box (lhs, rhs)) = &stmt.kind {
        if lhs.as_local() == Some(mir::RETURN_PLACE) {
            if let mir::Rvalue::Use(mir::Operand::Copy(rhs) | mir::Operand::Move(rhs)) = rhs {
                return rhs.as_local();
            }
        }
    }

    None
}

struct RenameToReturnPlace<'tcx> {
    to_rename: Local,
    tcx: TyCtxt<'tcx>,
}

/// Replaces all uses of `self.to_rename` with `_0`.
impl MutVisitor<'tcx> for RenameToReturnPlace<'tcx> {
    fn tcx(&self) -> TyCtxt<'tcx> {
        self.tcx
    }

    fn visit_statement(&mut self, stmt: &mut mir::Statement<'tcx>, loc: Location) {
        // Remove assignments of the local being replaced to the return place, since it is now the
        // return place:
        //     _0 = _1
        if as_local_assigned_to_return_place(stmt) == Some(self.to_rename) {
            stmt.kind = mir::StatementKind::Nop;
            return;
        }

        // Remove storage annotations for the local being replaced:
        //     StorageLive(_1)
        if let mir::StatementKind::StorageLive(local) | mir::StatementKind::StorageDead(local) =
            stmt.kind
        {
            if local == self.to_rename {
                stmt.kind = mir::StatementKind::Nop;
                return;
            }
        }

        self.super_statement(stmt, loc)
    }

    fn visit_terminator(&mut self, terminator: &mut mir::Terminator<'tcx>, loc: Location) {
        // Ignore the implicit "use" of the return place in a `Return` statement.
        if let mir::TerminatorKind::Return = terminator.kind {
            return;
        }

        self.super_terminator(terminator, loc);
    }

    fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
        assert_ne!(*l, mir::RETURN_PLACE);
        if *l == self.to_rename {
            *l = mir::RETURN_PLACE;
        }
    }
}

struct IsReturnPlaceRead(bool);

impl IsReturnPlaceRead {
    fn run(body: &mir::Body<'_>) -> bool {
        let mut vis = IsReturnPlaceRead(false);
        vis.visit_body(body);
        vis.0
    }
}

impl Visitor<'tcx> for IsReturnPlaceRead {
    fn visit_local(&mut self, &l: &Local, ctxt: PlaceContext, _: Location) {
        if l == mir::RETURN_PLACE && ctxt.is_use() && !ctxt.is_place_assignment() {
            self.0 = true;
        }
    }

    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, loc: Location) {
        // Ignore the implicit "use" of the return place in a `Return` statement.
        if let mir::TerminatorKind::Return = terminator.kind {
            return;
        }

        self.super_terminator(terminator, loc);
    }
}
<commit_msg>Rollup merge of #72451 - ecstatic-morse:nrvo-type-mismatch, r=matthewjasper<commit_after>use rustc_hir::Mutability;
use rustc_index::bit_set::HybridBitSet;
use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
use rustc_middle::mir::{self, BasicBlock, Local, Location};
use rustc_middle::ty::TyCtxt;

use crate::transform::{MirPass, MirSource};

/// This pass looks for MIR that always copies the same local into the return place and eliminates
/// the copy by renaming all uses of that local to `_0`.
///
/// This allows LLVM to perform an optimization similar to the named return value optimization
/// (NRVO) that is guaranteed in C++. This avoids a stack allocation and `memcpy` for the
/// relatively common pattern of allocating a buffer on the stack, mutating it, and returning it by
/// value like so:
///
/// ```rust
/// fn foo(init: fn(&mut [u8; 1024])) -> [u8; 1024] {
///     let mut buf = [0; 1024];
///     init(&mut buf);
///     buf
/// }
/// ```
///
/// For now, this pass is very simple and only capable of eliminating a single copy. A more general
/// version of copy propagation, such as the one based on non-overlapping live ranges in [#47954] and
/// [#71003], could yield even more benefits.
///
/// [#47954]: https://github.com/rust-lang/rust/pull/47954
/// [#71003]: https://github.com/rust-lang/rust/pull/71003
pub struct RenameReturnPlace;

impl<'tcx> MirPass<'tcx> for RenameReturnPlace {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, body: &mut mir::Body<'tcx>) {
        if tcx.sess.opts.debugging_opts.mir_opt_level == 0 {
            return;
        }

        let returned_local = match local_eligible_for_nrvo(body) {
            Some(l) => l,
            None => {
                debug!("`{:?}` was ineligible for NRVO", src.def_id());
                return;
            }
        };

        debug!(
            "`{:?}` was eligible for NRVO, making {:?} the return place",
            src.def_id(),
            returned_local
        );

        RenameToReturnPlace { tcx, to_rename: returned_local }.visit_body(body);

        // Clean up the `NOP`s we inserted for statements made useless by our renaming.
        for block_data in body.basic_blocks_mut() {
            block_data.statements.retain(|stmt| stmt.kind != mir::StatementKind::Nop);
        }

        // Overwrite the debuginfo of `_0` with that of the renamed local.
        let (renamed_decl, ret_decl) =
            body.local_decls.pick2_mut(returned_local, mir::RETURN_PLACE);

        // Sometimes, the return place is assigned a local of a different but coercable type, for
        // example `&mut T` instead of `&T`. Overwriting the `LocalInfo` for the return place means
        // its type may no longer match the return type of its function. This doesn't cause a
        // problem in codegen because these two types are layout-compatible, but may be unexpected.
        debug!("_0: {:?} = {:?}: {:?}", ret_decl.ty, returned_local, renamed_decl.ty);
        ret_decl.clone_from(renamed_decl);

        // The return place is always mutable.
        ret_decl.mutability = Mutability::Mut;
    }
}

/// MIR that is eligible for the NRVO must fulfill two conditions:
///   1. The return place must not be read prior to the `Return` terminator.
///   2. A simple assignment of a whole local to the return place (e.g., `_0 = _1`) must be the
///      only definition of the return place reaching the `Return` terminator.
///
/// If the MIR fulfills both these conditions, this function returns the `Local` that is assigned
/// to the return place along all possible paths through the control-flow graph.
fn local_eligible_for_nrvo(body: &mut mir::Body<'_>) -> Option<Local> {
    if IsReturnPlaceRead::run(body) {
        return None;
    }

    let mut copied_to_return_place = None;
    for block in body.basic_blocks().indices() {
        // Look for blocks with a `Return` terminator.
        if !matches!(body[block].terminator().kind, mir::TerminatorKind::Return) {
            continue;
        }

        // Look for an assignment of a single local to the return place prior to the `Return`.
        let returned_local = find_local_assigned_to_return_place(block, body)?;
        match body.local_kind(returned_local) {
            // FIXME: Can we do this for arguments as well?
            mir::LocalKind::Arg => return None,

            mir::LocalKind::ReturnPointer => bug!("Return place was assigned to itself?"),
            mir::LocalKind::Var | mir::LocalKind::Temp => {}
        }

        // If multiple different locals are copied to the return place. We can't pick a
        // single one to rename.
        if copied_to_return_place.map_or(false, |old| old != returned_local) {
            return None;
        }

        copied_to_return_place = Some(returned_local);
    }

    return copied_to_return_place;
}

fn find_local_assigned_to_return_place(
    start: BasicBlock,
    body: &mut mir::Body<'_>,
) -> Option<Local> {
    let mut block = start;
    let mut seen = HybridBitSet::new_empty(body.basic_blocks().len());

    // Iterate as long as `block` has exactly one predecessor that we have not yet visited.
    while seen.insert(block) {
        trace!("Looking for assignments to `_0` in {:?}", block);

        let local = body[block].statements.iter().rev().find_map(as_local_assigned_to_return_place);
        if local.is_some() {
            return local;
        }

        match body.predecessors()[block].as_slice() {
            &[pred] => block = pred,
            _ => return None,
        }
    }

    return None;
}

// If this statement is an assignment of an unprojected local to the return place,
// return that local.
fn as_local_assigned_to_return_place(stmt: &mir::Statement<'_>) -> Option<Local> {
    if let mir::StatementKind::Assign(box (lhs, rhs)) = &stmt.kind {
        if lhs.as_local() == Some(mir::RETURN_PLACE) {
            if let mir::Rvalue::Use(mir::Operand::Copy(rhs) | mir::Operand::Move(rhs)) = rhs {
                return rhs.as_local();
            }
        }
    }

    None
}

struct RenameToReturnPlace<'tcx> {
    to_rename: Local,
    tcx: TyCtxt<'tcx>,
}

/// Replaces all uses of `self.to_rename` with `_0`.
impl MutVisitor<'tcx> for RenameToReturnPlace<'tcx> {
    fn tcx(&self) -> TyCtxt<'tcx> {
        self.tcx
    }

    fn visit_statement(&mut self, stmt: &mut mir::Statement<'tcx>, loc: Location) {
        // Remove assignments of the local being replaced to the return place, since it is now the
        // return place:
        //     _0 = _1
        if as_local_assigned_to_return_place(stmt) == Some(self.to_rename) {
            stmt.kind = mir::StatementKind::Nop;
            return;
        }

        // Remove storage annotations for the local being replaced:
        //     StorageLive(_1)
        if let mir::StatementKind::StorageLive(local) | mir::StatementKind::StorageDead(local) =
            stmt.kind
        {
            if local == self.to_rename {
                stmt.kind = mir::StatementKind::Nop;
                return;
            }
        }

        self.super_statement(stmt, loc)
    }

    fn visit_terminator(&mut self, terminator: &mut mir::Terminator<'tcx>, loc: Location) {
        // Ignore the implicit "use" of the return place in a `Return` statement.
        if let mir::TerminatorKind::Return = terminator.kind {
            return;
        }

        self.super_terminator(terminator, loc);
    }

    fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
        assert_ne!(*l, mir::RETURN_PLACE);
        if *l == self.to_rename {
            *l = mir::RETURN_PLACE;
        }
    }
}

struct IsReturnPlaceRead(bool);

impl IsReturnPlaceRead {
    fn run(body: &mir::Body<'_>) -> bool {
        let mut vis = IsReturnPlaceRead(false);
        vis.visit_body(body);
        vis.0
    }
}

impl Visitor<'tcx> for IsReturnPlaceRead {
    fn visit_local(&mut self, &l: &Local, ctxt: PlaceContext, _: Location) {
        if l == mir::RETURN_PLACE && ctxt.is_use() && !ctxt.is_place_assignment() {
            self.0 = true;
        }
    }

    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, loc: Location) {
        // Ignore the implicit "use" of the return place in a `Return` statement.
        if let mir::TerminatorKind::Return = terminator.kind {
            return;
        }

        self.super_terminator(terminator, loc);
    }
}
<|endoftext|>