<commit_before>//! Low level functionality for writing DWARF debugging information.

use std::error;
use std::fmt;
use std::result;

mod endian_vec;
pub use self::endian_vec::*;

mod writer;
pub use self::writer::*;

#[macro_use]
mod section;
pub use self::section::*;

macro_rules! define_id {
    ($name:ident, $docs:expr) => {
        #[doc=$docs]
        #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
        pub struct $name {
            base_id: BaseId,
            index: usize,
        }

        impl $name {
            #[inline]
            fn new(base_id: BaseId, index: usize) -> Self {
                $name { base_id, index }
            }
        }
    };
}

macro_rules! define_offsets {
    ($offsets:ident: $id:ident => $offset:ident, $off_doc:expr) => {
        #[doc=$off_doc]
        #[derive(Debug)]
        pub struct $offsets {
            base_id: BaseId,
            // We know ids start at 0.
            offsets: Vec<$offset>,
        }

        impl $offsets {
            /// Return an empty list of offsets.
            #[inline]
            pub fn none() -> Self {
                $offsets {
                    base_id: BaseId::default(),
                    offsets: Vec::new(),
                }
            }

            /// Get the offset
            ///
            /// # Panics
            ///
            /// Panics if `id` is invalid.
            #[inline]
            pub fn get(&self, id: $id) -> $offset {
                debug_assert_eq!(self.base_id, id.base_id);
                self.offsets[id.index]
            }

            /// Return the number of offsets.
            #[inline]
            pub fn count(&self) -> usize {
                self.offsets.len()
            }
        }
    };
}

mod dwarf;
pub use self::dwarf::*;

mod abbrev;
pub use self::abbrev::*;

mod line;
pub use self::line::*;

mod str;
pub use self::str::*;

mod unit;
pub use self::unit::*;

mod range;
pub use self::range::*;

/// An error that occurred when writing.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Error {
    /// The given offset is out of bounds.
    OffsetOutOfBounds,
    /// The given length is out of bounds.
    LengthOutOfBounds,
    /// The attribute value is an invalid for writing.
    InvalidAttributeValue,
    /// The value is too large for the encoding form.
    ValueTooLarge,
    /// Unsupported word size.
    UnsupportedWordSize(u8),
    /// Unsupported DWARF version.
    UnsupportedVersion(u16),
    /// The unit length is too large for the requested DWARF format.
    InitialLengthOverflow,
    /// The address is invalid.
    InvalidAddress,
    /// A requested feature requires a different DWARF version.
    NeedVersion(u16),
    /// Strings in line number program have mismatched forms.
    LineStringFormMismatch,
    /// The range is empty or otherwise invalid.
    InvalidRange,
    /// The line number program encoding is incompatible with the unit encoding.
    IncompatibleLineProgramEncoding,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
        match *self {
            Error::OffsetOutOfBounds => write!(f, "The given offset is out of bounds."),
            Error::LengthOutOfBounds => write!(f, "The given length is out of bounds."),
            Error::InvalidAttributeValue => {
                write!(f, "The attribute value is an invalid for writing.")
            }
            Error::ValueTooLarge => write!(f, "The value is too large for the encoding form."),
            Error::UnsupportedWordSize(size) => write!(f, "Unsupported word size: {}", size),
            Error::UnsupportedVersion(version) => {
                write!(f, "Unsupported DWARF version: {}", version)
            }
            Error::InitialLengthOverflow => write!(
                f,
                "The unit length is too large for the requested DWARF format."
            ),
            Error::InvalidAddress => write!(f, "The address is invalid."),
            Error::NeedVersion(version) => write!(
                f,
                "A requested feature requires a DWARF version {}.",
                version
            ),
            Error::LineStringFormMismatch => {
                write!(f, "Strings in line number program have mismatched forms.")
            }
            Error::InvalidRange => write!(f, "The range is empty or otherwise invalid."),
            Error::IncompatibleLineProgramEncoding => write!(
                f,
                "The line number program encoding is incompatible with the unit encoding."
            ),
        }
    }
}

impl error::Error for Error {}

/// The result of a write.
pub type Result<T> = result::Result<T, Error>;

/// An address.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Address {
    /// An absolute address that does not require relocation.
    Absolute(u64),
    /// An address that is relative to a symbol which may be relocated.
    Relative {
        /// The symbol that the address is relative to.
        ///
        /// The meaning of this value is decided by the writer, but
        /// will typically be an index into a symbol table.
        symbol: usize,
        /// The offset of the address relative to the symbol.
        ///
        /// This will typically be used as the addend in a relocation.
        addend: i64,
    },
}

// This type is only used in debug assertions.
#[cfg(not(debug_assertions))]
type BaseId = ();

#[cfg(debug_assertions)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct BaseId(usize);

#[cfg(debug_assertions)]
impl Default for BaseId {
    fn default() -> Self {
        use std::sync::atomic;
        static BASE_ID: atomic::AtomicUsize = atomic::ATOMIC_USIZE_INIT;
        BaseId(BASE_ID.fetch_add(1, atomic::Ordering::Relaxed))
    }
}

#[cfg(feature = "read")]
mod convert {
    use super::*;
    use read;

    pub(crate) use super::unit::convert::*;

    /// An error that occurred when converting a read value into a write value.
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub enum ConvertError {
        /// An error occurred when reading.
        Read(read::Error),
        /// Writing of this attribute value is not implemented yet.
        UnsupportedAttributeValue,
        /// This attribute value is an invalid name/form combination.
        InvalidAttributeValue,
        /// A `.debug_info` reference does not refer to a valid entry.
        InvalidDebugInfoOffset,
        /// An address could not be converted.
        InvalidAddress,
        /// The line number program is missing the compilation directory.
        MissingCompilationDirectory,
        /// The line number program is missing the compilation file.
        MissingCompilationFile,
        /// Writing this line number instruction is not implemented yet.
        UnsupportedLineInstruction,
        /// Writing this form of line string is not implemented yet.
        UnsupportedLineStringForm,
        /// A `.debug_line` file index is invalid.
        InvalidFileIndex,
        /// A `.debug_line` directory index is invalid.
        InvalidDirectoryIndex,
        /// A `.debug_line` line base is invalid.
        InvalidLineBase,
        /// A `.debug_line` reference is invalid.
        InvalidLineRef,
        /// Invalid relative address in a range list.
        InvalidRangeRelativeAddress,
    }

    impl fmt::Display for ConvertError {
        fn fmt(&self, f: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
            use self::ConvertError::*;
            match *self {
                Read(ref e) => e.fmt(f),
                UnsupportedAttributeValue => {
                    write!(f, "Writing of this attribute value is not implemented yet.")
                }
                InvalidAttributeValue => write!(
                    f,
                    "This attribute value is an invalid name/form combination."
                ),
                InvalidDebugInfoOffset => write!(
                    f,
                    "A `.debug_info` reference does not refer to a valid entry."
                ),
                InvalidAddress => write!(f, "An address could not be converted."),
                MissingCompilationDirectory => write!(
                    f,
                    "The line number program is missing the compilation directory."
                ),
                MissingCompilationFile => write!(
                    f,
                    "The line number program is missing the compilation file."
                ),
                UnsupportedLineInstruction => write!(
                    f,
                    "Writing this line number instruction is not implemented yet."
                ),
                UnsupportedLineStringForm => write!(
                    f,
                    "Writing this form of line string is not implemented yet."
                ),
                InvalidFileIndex => write!(f, "A `.debug_line` file index is invalid."),
                InvalidDirectoryIndex => write!(f, "A `.debug_line` directory index is invalid."),
                InvalidLineBase => write!(f, "A `.debug_line` line base is invalid."),
                InvalidLineRef => write!(f, "A `.debug_line` reference is invalid."),
                InvalidRangeRelativeAddress => {
                    write!(f, "Invalid relative address in a range list.")
                }
            }
        }
    }

    impl error::Error for ConvertError {}

    impl From<read::Error> for ConvertError {
        fn from(e: read::Error) -> Self {
            ConvertError::Read(e)
        }
    }

    /// The result of a conversion.
    pub type ConvertResult<T> = result::Result<T, ConvertError>;
}
#[cfg(feature = "read")]
pub use self::convert::*;
<commit_msg>Add a little bit of write module doc<commit_after>//! Write DWARF debugging information.
//!
//! ## API Structure
//!
//! This module works by building up a representation of the debugging information
//! in memory, and then writing it all at once. It supports two major use cases:
//!
//! * Use the [`DwarfUnit`](./struct.DwarfUnit.html) type when writing DWARF
//! for a single compilation unit.
//!
//! * Use the [`Dwarf`](./struct.Dwarf.html) type when writing DWARF for multiple
//! compilation units.
//!
//! The module also supports reading in DWARF debugging information and writing it out
//! again, possibly after modifying it. Create a [`read::Dwarf`](../read/struct.Dwarf.html)
//! instance, and then use [`Dwarf::from`](./struct.Dwarf.html#method.from) to convert
//! it to a writable instance.

use std::error;
use std::fmt;
use std::result;

mod endian_vec;
pub use self::endian_vec::*;

mod writer;
pub use self::writer::*;

#[macro_use]
mod section;
pub use self::section::*;

macro_rules! define_id {
    ($name:ident, $docs:expr) => {
        #[doc=$docs]
        #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
        pub struct $name {
            base_id: BaseId,
            index: usize,
        }

        impl $name {
            #[inline]
            fn new(base_id: BaseId, index: usize) -> Self {
                $name { base_id, index }
            }
        }
    };
}

macro_rules! define_offsets {
    ($offsets:ident: $id:ident => $offset:ident, $off_doc:expr) => {
        #[doc=$off_doc]
        #[derive(Debug)]
        pub struct $offsets {
            base_id: BaseId,
            // We know ids start at 0.
            offsets: Vec<$offset>,
        }

        impl $offsets {
            /// Return an empty list of offsets.
            #[inline]
            pub fn none() -> Self {
                $offsets {
                    base_id: BaseId::default(),
                    offsets: Vec::new(),
                }
            }

            /// Get the offset
            ///
            /// # Panics
            ///
            /// Panics if `id` is invalid.
            #[inline]
            pub fn get(&self, id: $id) -> $offset {
                debug_assert_eq!(self.base_id, id.base_id);
                self.offsets[id.index]
            }

            /// Return the number of offsets.
            #[inline]
            pub fn count(&self) -> usize {
                self.offsets.len()
            }
        }
    };
}

mod dwarf;
pub use self::dwarf::*;

mod abbrev;
pub use self::abbrev::*;

mod line;
pub use self::line::*;

mod str;
pub use self::str::*;

mod unit;
pub use self::unit::*;

mod range;
pub use self::range::*;

/// An error that occurred when writing.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Error {
    /// The given offset is out of bounds.
    OffsetOutOfBounds,
    /// The given length is out of bounds.
    LengthOutOfBounds,
    /// The attribute value is an invalid for writing.
    InvalidAttributeValue,
    /// The value is too large for the encoding form.
    ValueTooLarge,
    /// Unsupported word size.
    UnsupportedWordSize(u8),
    /// Unsupported DWARF version.
    UnsupportedVersion(u16),
    /// The unit length is too large for the requested DWARF format.
    InitialLengthOverflow,
    /// The address is invalid.
    InvalidAddress,
    /// A requested feature requires a different DWARF version.
    NeedVersion(u16),
    /// Strings in line number program have mismatched forms.
    LineStringFormMismatch,
    /// The range is empty or otherwise invalid.
    InvalidRange,
    /// The line number program encoding is incompatible with the unit encoding.
    IncompatibleLineProgramEncoding,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
        match *self {
            Error::OffsetOutOfBounds => write!(f, "The given offset is out of bounds."),
            Error::LengthOutOfBounds => write!(f, "The given length is out of bounds."),
            Error::InvalidAttributeValue => {
                write!(f, "The attribute value is an invalid for writing.")
            }
            Error::ValueTooLarge => write!(f, "The value is too large for the encoding form."),
            Error::UnsupportedWordSize(size) => write!(f, "Unsupported word size: {}", size),
            Error::UnsupportedVersion(version) => {
                write!(f, "Unsupported DWARF version: {}", version)
            }
            Error::InitialLengthOverflow => write!(
                f,
                "The unit length is too large for the requested DWARF format."
            ),
            Error::InvalidAddress => write!(f, "The address is invalid."),
            Error::NeedVersion(version) => write!(
                f,
                "A requested feature requires a DWARF version {}.",
                version
            ),
            Error::LineStringFormMismatch => {
                write!(f, "Strings in line number program have mismatched forms.")
            }
            Error::InvalidRange => write!(f, "The range is empty or otherwise invalid."),
            Error::IncompatibleLineProgramEncoding => write!(
                f,
                "The line number program encoding is incompatible with the unit encoding."
            ),
        }
    }
}

impl error::Error for Error {}

/// The result of a write.
pub type Result<T> = result::Result<T, Error>;

/// An address.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Address {
    /// An absolute address that does not require relocation.
    Absolute(u64),
    /// An address that is relative to a symbol which may be relocated.
    Relative {
        /// The symbol that the address is relative to.
        ///
        /// The meaning of this value is decided by the writer, but
        /// will typically be an index into a symbol table.
        symbol: usize,
        /// The offset of the address relative to the symbol.
        ///
        /// This will typically be used as the addend in a relocation.
        addend: i64,
    },
}

// This type is only used in debug assertions.
#[cfg(not(debug_assertions))]
type BaseId = ();

#[cfg(debug_assertions)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct BaseId(usize);

#[cfg(debug_assertions)]
impl Default for BaseId {
    fn default() -> Self {
        use std::sync::atomic;
        static BASE_ID: atomic::AtomicUsize = atomic::ATOMIC_USIZE_INIT;
        BaseId(BASE_ID.fetch_add(1, atomic::Ordering::Relaxed))
    }
}

#[cfg(feature = "read")]
mod convert {
    use super::*;
    use read;

    pub(crate) use super::unit::convert::*;

    /// An error that occurred when converting a read value into a write value.
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub enum ConvertError {
        /// An error occurred when reading.
        Read(read::Error),
        /// Writing of this attribute value is not implemented yet.
        UnsupportedAttributeValue,
        /// This attribute value is an invalid name/form combination.
        InvalidAttributeValue,
        /// A `.debug_info` reference does not refer to a valid entry.
        InvalidDebugInfoOffset,
        /// An address could not be converted.
        InvalidAddress,
        /// The line number program is missing the compilation directory.
        MissingCompilationDirectory,
        /// The line number program is missing the compilation file.
        MissingCompilationFile,
        /// Writing this line number instruction is not implemented yet.
        UnsupportedLineInstruction,
        /// Writing this form of line string is not implemented yet.
        UnsupportedLineStringForm,
        /// A `.debug_line` file index is invalid.
        InvalidFileIndex,
        /// A `.debug_line` directory index is invalid.
        InvalidDirectoryIndex,
        /// A `.debug_line` line base is invalid.
        InvalidLineBase,
        /// A `.debug_line` reference is invalid.
        InvalidLineRef,
        /// Invalid relative address in a range list.
        InvalidRangeRelativeAddress,
    }

    impl fmt::Display for ConvertError {
        fn fmt(&self, f: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
            use self::ConvertError::*;
            match *self {
                Read(ref e) => e.fmt(f),
                UnsupportedAttributeValue => {
                    write!(f, "Writing of this attribute value is not implemented yet.")
                }
                InvalidAttributeValue => write!(
                    f,
                    "This attribute value is an invalid name/form combination."
                ),
                InvalidDebugInfoOffset => write!(
                    f,
                    "A `.debug_info` reference does not refer to a valid entry."
                ),
                InvalidAddress => write!(f, "An address could not be converted."),
                MissingCompilationDirectory => write!(
                    f,
                    "The line number program is missing the compilation directory."
                ),
                MissingCompilationFile => write!(
                    f,
                    "The line number program is missing the compilation file."
                ),
                UnsupportedLineInstruction => write!(
                    f,
                    "Writing this line number instruction is not implemented yet."
                ),
                UnsupportedLineStringForm => write!(
                    f,
                    "Writing this form of line string is not implemented yet."
                ),
                InvalidFileIndex => write!(f, "A `.debug_line` file index is invalid."),
                InvalidDirectoryIndex => write!(f, "A `.debug_line` directory index is invalid."),
                InvalidLineBase => write!(f, "A `.debug_line` line base is invalid."),
                InvalidLineRef => write!(f, "A `.debug_line` reference is invalid."),
                InvalidRangeRelativeAddress => {
                    write!(f, "Invalid relative address in a range list.")
                }
            }
        }
    }

    impl error::Error for ConvertError {}

    impl From<read::Error> for ConvertError {
        fn from(e: read::Error) -> Self {
            ConvertError::Read(e)
        }
    }

    /// The result of a conversion.
    pub type ConvertResult<T> = result::Result<T, ConvertError>;
}
#[cfg(feature = "read")]
pub use self::convert::*;
<|endoftext|>


<commit_before>#![cfg(test)]
#![allow(unused)]

use cairo;
use gio;
use glib;
use glib::translate::*;
use libc;
use std::env;
use std::ffi::CString;
use std::sync::Once;

use librsvg::{
    surface_utils::shared_surface::{SharedImageSurface, SurfaceType},
    CairoRenderer, Loader, LoadingError, RenderingError, SvgHandle,
};

pub fn load_svg(input: &'static [u8]) -> Result<SvgHandle, LoadingError> {
    let bytes = glib::Bytes::from_static(input);
    let stream = gio::MemoryInputStream::from_bytes(&bytes);

    Loader::new().read_stream(&stream, None::<&gio::File>, None::<&gio::Cancellable>)
}

#[derive(Copy, Clone)]
pub struct SurfaceSize(pub i32, pub i32);

pub fn render_document<F: FnOnce(&cairo::Context)>(
    svg: &SvgHandle,
    surface_size: SurfaceSize,
    cr_transform: F,
    viewport: cairo::Rectangle,
) -> Result<SharedImageSurface, RenderingError> {
    let renderer = CairoRenderer::new(svg);

    let SurfaceSize(width, height) = surface_size;

    let output = cairo::ImageSurface::create(cairo::Format::ARgb32, width, height).unwrap();

    let res = {
        let cr = cairo::Context::new(&output).expect("Failed to create a cairo context");
        cr_transform(&cr);
        Ok(renderer.render_document(&cr, &viewport)?)
    };

    res.and_then(|_| Ok(SharedImageSurface::wrap(output, SurfaceType::SRgb)?))
}

#[cfg(system_deps_have_pangoft2)]
mod pango_ft2 {
    use super::*;
    use fontconfig_sys::fontconfig;
    use glib::prelude::*;
    use pangocairo::FontMap;

    extern "C" {
        // pango_fc_font_map_set_config (PangoFcFontMap *fcfontmap,
        //                               FcConfig       *fcconfig);
        // This is not bound in gtk-rs, and PangoFcFontMap is not even exposed, so we'll bind it by hand.
        fn pango_fc_font_map_set_config(
            font_map: *mut libc::c_void,
            config: *mut fontconfig::FcConfig,
        );
    }

    pub unsafe fn load_test_fonts() {
        let font_paths = [
            "tests/resources/Ahem.ttf",
            "tests/resources/Roboto-Regular.ttf",
            "tests/resources/Roboto-Italic.ttf",
            "tests/resources/Roboto-Bold.ttf",
            "tests/resources/Roboto-BoldItalic.ttf",
        ];

        let config = fontconfig::FcConfigCreate();

        for path in &font_paths {
            let path_cstring = CString::new(*path).unwrap();

            if fontconfig::FcConfigAppFontAddFile(config, path_cstring.as_ptr() as *const _) == 0 {
                panic!("Could not load font file {} for tests; aborting", path,);
            }
        }

        let font_map = FontMap::for_font_type(cairo::FontType::FontTypeFt).unwrap();
        let raw_font_map: *mut pango::ffi::PangoFontMap = font_map.to_glib_none().0;

        pango_fc_font_map_set_config(raw_font_map as *mut _, config);
        fontconfig::FcConfigDestroy(config);

        FontMap::set_default(Some(font_map.downcast::<pangocairo::FontMap>().unwrap()));
    }
}

#[cfg(system_deps_have_pangoft2)]
pub fn setup_font_map() {
    unsafe {
        self::pango_ft2::load_test_fonts();
    }
}

#[cfg(not(system_deps_have_pangoft2))]
pub fn setup_font_map() {}

pub fn setup_language() {
    static ONCE: Once = Once::new();

    ONCE.call_once(|| {
        // For systemLanguage attribute tests.
        // The trailing ":" is intentional to test gitlab#425.
        env::set_var("LANGUAGE", "de:en_US:en:");
        env::set_var("LC_ALL", "de:en_US:en:");
    });
}
<commit_msg>For the test suite, ensure only the testing FcConfig is being used<commit_after>#![cfg(test)]
#![allow(unused)]

use cairo;
use gio;
use glib;
use glib::translate::*;
use libc;
use std::env;
use std::ffi::CString;
use std::sync::Once;

use librsvg::{
    surface_utils::shared_surface::{SharedImageSurface, SurfaceType},
    CairoRenderer, Loader, LoadingError, RenderingError, SvgHandle,
};

pub fn load_svg(input: &'static [u8]) -> Result<SvgHandle, LoadingError> {
    let bytes = glib::Bytes::from_static(input);
    let stream = gio::MemoryInputStream::from_bytes(&bytes);

    Loader::new().read_stream(&stream, None::<&gio::File>, None::<&gio::Cancellable>)
}

#[derive(Copy, Clone)]
pub struct SurfaceSize(pub i32, pub i32);

pub fn render_document<F: FnOnce(&cairo::Context)>(
    svg: &SvgHandle,
    surface_size: SurfaceSize,
    cr_transform: F,
    viewport: cairo::Rectangle,
) -> Result<SharedImageSurface, RenderingError> {
    let renderer = CairoRenderer::new(svg);

    let SurfaceSize(width, height) = surface_size;

    let output = cairo::ImageSurface::create(cairo::Format::ARgb32, width, height).unwrap();

    let res = {
        let cr = cairo::Context::new(&output).expect("Failed to create a cairo context");
        cr_transform(&cr);
        Ok(renderer.render_document(&cr, &viewport)?)
    };

    res.and_then(|_| Ok(SharedImageSurface::wrap(output, SurfaceType::SRgb)?))
}

#[cfg(system_deps_have_pangoft2)]
mod pango_ft2 {
    use super::*;
    use fontconfig_sys::fontconfig;
    use glib::prelude::*;
    use pangocairo::FontMap;

    extern "C" {
        // pango_fc_font_map_set_config (PangoFcFontMap *fcfontmap,
        //                               FcConfig       *fcconfig);
        // This is not bound in gtk-rs, and PangoFcFontMap is not even exposed, so we'll bind it by hand.
        fn pango_fc_font_map_set_config(
            font_map: *mut libc::c_void,
            config: *mut fontconfig::FcConfig,
        );
    }

    pub unsafe fn load_test_fonts() {
        let font_paths = [
            "tests/resources/Ahem.ttf",
            "tests/resources/Roboto-Regular.ttf",
            "tests/resources/Roboto-Italic.ttf",
            "tests/resources/Roboto-Bold.ttf",
            "tests/resources/Roboto-BoldItalic.ttf",
        ];

        let config = fontconfig::FcConfigCreate();
        if fontconfig::FcConfigSetCurrent(config) == 0 {
            panic!("Could not set a fontconfig configuration");
        }

        for path in &font_paths {
            let path_cstring = CString::new(*path).unwrap();

            if fontconfig::FcConfigAppFontAddFile(config, path_cstring.as_ptr() as *const _) == 0 {
                panic!("Could not load font file {} for tests; aborting", path,);
            }
        }

        let font_map = FontMap::for_font_type(cairo::FontType::FontTypeFt).unwrap();
        let raw_font_map: *mut pango::ffi::PangoFontMap = font_map.to_glib_none().0;

        pango_fc_font_map_set_config(raw_font_map as *mut _, config);
        fontconfig::FcConfigDestroy(config);

        FontMap::set_default(Some(font_map.downcast::<pangocairo::FontMap>().unwrap()));
    }
}

#[cfg(system_deps_have_pangoft2)]
pub fn setup_font_map() {
    unsafe {
        self::pango_ft2::load_test_fonts();
    }
}

#[cfg(not(system_deps_have_pangoft2))]
pub fn setup_font_map() {}

pub fn setup_language() {
    static ONCE: Once = Once::new();

    ONCE.call_once(|| {
        // For systemLanguage attribute tests.
        // The trailing ":" is intentional to test gitlab#425.
        env::set_var("LANGUAGE", "de:en_US:en:");
        env::set_var("LC_ALL", "de:en_US:en:");
    });
}
<|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.

//! A radix trie for storing integers in sorted order

use prelude::*;

// FIXME: #5244: need to manually update the TrieNode constructor
const SHIFT: uint = 4;
const SIZE: uint = 1 << SHIFT;
const MASK: uint = SIZE - 1;

enum Child<T> {
    Internal(~TrieNode<T>),
    External(uint, T),
    Nothing
}

pub struct TrieMap<T> {
    priv root: TrieNode<T>,
    priv length: uint
}

impl<T> BaseIter<(uint, &'self T)> for TrieMap<T> {
    /// Visit all key-value pairs in order
    #[inline(always)]
    fn each(&self, f: &fn(&(uint, &'self T)) -> bool) {
        self.root.each(f);
    }
    #[inline(always)]
    fn size_hint(&self) -> Option<uint> { Some(self.len()) }
}

impl<T> ReverseIter<(uint, &'self T)> for TrieMap<T> {
    /// Visit all key-value pairs in reverse order
    #[inline(always)]
    fn each_reverse(&self, f: &fn(&(uint, &'self T)) -> bool) {
        self.root.each_reverse(f);
    }
}

impl<T> Container for TrieMap<T> {
    /// Return the number of elements in the map
    #[inline(always)]
    fn len(&const self) -> uint { self.length }

    /// Return true if the map contains no elements
    #[inline(always)]
    fn is_empty(&const self) -> bool { self.len() == 0 }
}

impl<T> Mutable for TrieMap<T> {
    /// Clear the map, removing all values.
    #[inline(always)]
    fn clear(&mut self) {
        self.root = TrieNode::new();
        self.length = 0;
    }
}

impl<T> Map<uint, T> for TrieMap<T> {
    /// Return true if the map contains a value for the specified key
    #[inline(always)]
    fn contains_key(&self, key: &uint) -> bool {
        self.find(key).is_some()
    }

    /// Visit all keys in order
    #[inline(always)]
    fn each_key(&self, f: &fn(&uint) -> bool) {
        self.each(|&(k, _)| f(&k))
    }

    /// Visit all values in order
    #[inline(always)]
    fn each_value(&self, f: &fn(&T) -> bool) {
        self.each(|&(_, v)| f(v))
    }

    /// Iterate over the map and mutate the contained values
    #[inline(always)]
    fn mutate_values(&mut self, f: &fn(&uint, &mut T) -> bool) {
        self.root.mutate_values(f);
    }

    /// Return the value corresponding to the key in the map
    #[inline(hint)]
    fn find(&self, key: &uint) -> Option<&'self T> {
        let mut node: &'self TrieNode<T> = &self.root;
        let mut idx = 0;
        loop {
            match node.children[chunk(*key, idx)] {
              Internal(ref x) => node = &**x,
              External(stored, ref value) => {
                if stored == *key {
                    return Some(value)
                } else {
                    return None
                }
              }
              Nothing => return None
            }
            idx += 1;
        }
    }

    /// Insert a key-value pair into the map. An existing value for a
    /// key is replaced by the new value. Return true if the key did
    /// not already exist in the map.
    #[inline(always)]
    fn insert(&mut self, key: uint, value: T) -> bool {
        let ret = insert(&mut self.root.count,
                         &mut self.root.children[chunk(key, 0)],
                         key, value, 1);
        if ret { self.length += 1 }
        ret
    }

    /// Remove a key-value pair from the map. Return true if the key
    /// was present in the map, otherwise false.
    #[inline(always)]
    fn remove(&mut self, key: &uint) -> bool {
        let ret = remove(&mut self.root.count,
                         &mut self.root.children[chunk(*key, 0)],
                         *key, 1);
        if ret { self.length -= 1 }
        ret
    }
}

pub impl<T> TrieMap<T> {
    /// Create an empty TrieMap
    #[inline(always)]
    fn new() -> TrieMap<T> {
        TrieMap{root: TrieNode::new(), length: 0}
    }

    /// Visit all keys in reverse order
    #[inline(always)]
    fn each_key_reverse(&self, f: &fn(&uint) -> bool) {
        self.each_reverse(|&(k, _)| f(&k))
    }

    /// Visit all values in reverse order
    #[inline(always)]
    fn each_value_reverse(&self, f: &fn(&T) -> bool) {
        self.each_reverse(|&(_, v)| f(v))
    }
}

pub struct TrieSet {
    priv map: TrieMap<()>
}

impl BaseIter<uint> for TrieSet {
    /// Visit all values in order
    fn each(&self, f: &fn(&uint) -> bool) { self.map.each_key(f) }
    fn size_hint(&self) -> Option<uint> { Some(self.len()) }
}

impl ReverseIter<uint> for TrieSet {
    /// Visit all values in reverse order
    fn each_reverse(&self, f: &fn(&uint) -> bool) {
        self.map.each_key_reverse(f)
    }
}

impl Container for TrieSet {
    /// Return the number of elements in the set
    #[inline(always)]
    fn len(&const self) -> uint { self.map.len() }

    /// Return true if the set contains no elements
    #[inline(always)]
    fn is_empty(&const self) -> bool { self.map.is_empty() }
}

impl Mutable for TrieSet {
    /// Clear the set, removing all values.
    #[inline(always)]
    fn clear(&mut self) { self.map.clear() }
}

impl TrieSet {
    /// Create an empty TrieSet
    #[inline(always)]
    fn new() -> TrieSet {
        TrieSet{map: TrieMap::new()}
    }

    /// Return true if the set contains a value
    #[inline(always)]
    fn contains(&self, value: &uint) -> bool {
        self.map.contains_key(value)
    }

    /// Add a value to the set. Return true if the value was not already
    /// present in the set.
    #[inline(always)]
    fn insert(&mut self, value: uint) -> bool { self.map.insert(value, ()) }

    /// Remove a value from the set. Return true if the value was
    /// present in the set.
    #[inline(always)]
    fn remove(&mut self, value: &uint) -> bool { self.map.remove(value) }
}

struct TrieNode<T> {
    count: uint,
    children: [Child<T> * SIZE]
}

impl<T> TrieNode<T> {
    #[inline(always)]
    fn new() -> TrieNode<T> {
        // FIXME: #5244: [Nothing, ..SIZE] should be possible without Copy
        TrieNode{count: 0,
                 children: [Nothing, Nothing, Nothing, Nothing,
                            Nothing, Nothing, Nothing, Nothing,
                            Nothing, Nothing, Nothing, Nothing,
                            Nothing, Nothing, Nothing, Nothing]}
    }
}

impl<T> TrieNode<T> {
    fn each(&self, f: &fn(&(uint, &'self T)) -> bool) -> bool {
        for uint::range(0, self.children.len()) |idx| {
            match self.children[idx] {
                Internal(ref x) => if !x.each(f) { return false },
                External(k, ref v) => if !f(&(k, v)) { return false },
                Nothing => ()
            }
        }
        true
    }

    fn each_reverse(&self, f: &fn(&(uint, &'self T)) -> bool) -> bool {
        for uint::range_rev(self.children.len(), 0) |idx| {
            match self.children[idx - 1] {
                Internal(ref x) => if !x.each_reverse(f) { return false },
                External(k, ref v) => if !f(&(k, v)) { return false },
                Nothing => ()
            }
        }
        true
    }

    fn mutate_values(&mut self, f: &fn(&uint, &mut T) -> bool) -> bool {
        for vec::each_mut(self.children) |child| {
            match *child {
                Internal(ref mut x) => if !x.mutate_values(f) {
                    return false
                },
                External(k, ref mut v) => if !f(&k, v) { return false },
                Nothing => ()
            }
        }
        true
    }
}

// if this was done via a trait, the key could be generic
#[inline(always)]
fn chunk(n: uint, idx: uint) -> uint {
    let sh = uint::bits - (SHIFT * (idx + 1));
    (n >> sh) & MASK
}

fn insert<T>(count: &mut uint, child: &mut Child<T>, key: uint, value: T,
             idx: uint) -> bool {
    let mut tmp = Nothing;
    tmp <-> *child;
    let mut added = false;

    *child = match tmp {
      External(stored_key, stored_value) => {
          if stored_key == key {
              External(stored_key, value)
          } else {
              // conflict - split the node
              let mut new = ~TrieNode::new();
              insert(&mut new.count,
                     &mut new.children[chunk(stored_key, idx)],
                     stored_key, stored_value, idx + 1);
              insert(&mut new.count, &mut new.children[chunk(key, idx)], key,
                     value, idx + 1);
              added = true;
              Internal(new)
          }
      }
      Internal(x) => {
        let mut x = x;
        added = insert(&mut x.count, &mut x.children[chunk(key, idx)], key,
                       value, idx + 1);
        Internal(x)

      }
      Nothing => {
        *count += 1;
        added = true;
        External(key, value)
      }
    };
    added
}

fn remove<T>(count: &mut uint, child: &mut Child<T>, key: uint,
             idx: uint) -> bool {
    let (ret, this) = match *child {
      External(stored, _) => {
          if stored == key { (true, true) } else { (false, false) }
      }
      Internal(ref mut x) => {
          let ret = remove(&mut x.count, &mut x.children[chunk(key, idx)],
                           key, idx + 1);
          (ret, x.count == 0)
      }
      Nothing => (false, false)
    };

    if this {
        *child = Nothing;
        *count -= 1;
    }
    ret
}

#[cfg(test)]
pub fn check_integrity<T>(trie: &TrieNode<T>) {
    fail_unless!(trie.count != 0);

    let mut sum = 0;

    for trie.children.each |x| {
        match *x {
          Nothing => (),
          Internal(ref y) => {
              check_integrity(&**y);
              sum += 1
          }
          External(_, _) => { sum += 1 }
        }
    }

    fail_unless!(sum == trie.count);
}

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

    #[test]
    fn test_step() {
        let mut trie = TrieMap::new();
        let n = 300;

        for uint::range_step(1, n, 2) |x| {
            fail_unless!(trie.insert(x, x + 1));
            fail_unless!(trie.contains_key(&x));
            check_integrity(&trie.root);
        }

        for uint::range_step(0, n, 2) |x| {
            fail_unless!(!trie.contains_key(&x));
            fail_unless!(trie.insert(x, x + 1));
            check_integrity(&trie.root);
        }

        for uint::range(0, n) |x| {
            fail_unless!(trie.contains_key(&x));
            fail_unless!(!trie.insert(x, x + 1));
            check_integrity(&trie.root);
        }

        for uint::range_step(1, n, 2) |x| {
            fail_unless!(trie.remove(&x));
            fail_unless!(!trie.contains_key(&x));
            check_integrity(&trie.root);
        }

        for uint::range_step(0, n, 2) |x| {
            fail_unless!(trie.contains_key(&x));
            fail_unless!(!trie.insert(x, x + 1));
            check_integrity(&trie.root);
        }
    }

    #[test]
    fn test_each() {
        let mut m = TrieMap::new();

        fail_unless!(m.insert(3, 6));
        fail_unless!(m.insert(0, 0));
        fail_unless!(m.insert(4, 8));
        fail_unless!(m.insert(2, 4));
        fail_unless!(m.insert(1, 2));

        let mut n = 0;
        for m.each |&(k, v)| {
            fail_unless!(k == n);
            fail_unless!(*v == n * 2);
            n += 1;
        }
    }

    #[test]
    fn test_each_break() {
        let mut m = TrieMap::new();

        for uint::range_rev(uint::max_value, uint::max_value - 10000) |x| {
            m.insert(x, x / 2);
        }

        let mut n = uint::max_value - 9999;
        for m.each |&(k, v)| {
            if n == uint::max_value - 5000 { break }
            fail_unless!(n < uint::max_value - 5000);

            fail_unless!(k == n);
            fail_unless!(*v == n / 2);
            n += 1;
        }
    }

    #[test]
    fn test_each_reverse() {
        let mut m = TrieMap::new();

        fail_unless!(m.insert(3, 6));
        fail_unless!(m.insert(0, 0));
        fail_unless!(m.insert(4, 8));
        fail_unless!(m.insert(2, 4));
        fail_unless!(m.insert(1, 2));

        let mut n = 4;
        for m.each_reverse |&(k, v)| {
            fail_unless!(k == n);
            fail_unless!(*v == n * 2);
            n -= 1;
        }
    }

    #[test]
    fn test_each_reverse_break() {
        let mut m = TrieMap::new();

        for uint::range_rev(uint::max_value, uint::max_value - 10000) |x| {
            m.insert(x, x / 2);
        }

        let mut n = uint::max_value;
        for m.each_reverse |&(k, v)| {
            if n == uint::max_value - 5000 { break }
            fail_unless!(n > uint::max_value - 5000);

            fail_unless!(k == n);
            fail_unless!(*v == n / 2);
            n -= 1;
        }
    }

    #[test]
    fn test_sane_chunk() {
        let x = 1;
        let y = 1 << (uint::bits - 1);

        let mut trie = TrieSet::new();

        fail_unless!(trie.insert(x));
        fail_unless!(trie.insert(y));

        fail_unless!(trie.len() == 2);

        let expected = [x, y];

        let mut i = 0;

        for trie.each |x| {
            fail_unless!(expected[i] == *x);
            i += 1;
        }
    }
}
<commit_msg>trie: inline the other TrieSet wrapper methods<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.

//! A radix trie for storing integers in sorted order

use prelude::*;

// FIXME: #5244: need to manually update the TrieNode constructor
const SHIFT: uint = 4;
const SIZE: uint = 1 << SHIFT;
const MASK: uint = SIZE - 1;

enum Child<T> {
    Internal(~TrieNode<T>),
    External(uint, T),
    Nothing
}

pub struct TrieMap<T> {
    priv root: TrieNode<T>,
    priv length: uint
}

impl<T> BaseIter<(uint, &'self T)> for TrieMap<T> {
    /// Visit all key-value pairs in order
    #[inline(always)]
    fn each(&self, f: &fn(&(uint, &'self T)) -> bool) {
        self.root.each(f);
    }
    #[inline(always)]
    fn size_hint(&self) -> Option<uint> { Some(self.len()) }
}

impl<T> ReverseIter<(uint, &'self T)> for TrieMap<T> {
    /// Visit all key-value pairs in reverse order
    #[inline(always)]
    fn each_reverse(&self, f: &fn(&(uint, &'self T)) -> bool) {
        self.root.each_reverse(f);
    }
}

impl<T> Container for TrieMap<T> {
    /// Return the number of elements in the map
    #[inline(always)]
    fn len(&const self) -> uint { self.length }

    /// Return true if the map contains no elements
    #[inline(always)]
    fn is_empty(&const self) -> bool { self.len() == 0 }
}

impl<T> Mutable for TrieMap<T> {
    /// Clear the map, removing all values.
    #[inline(always)]
    fn clear(&mut self) {
        self.root = TrieNode::new();
        self.length = 0;
    }
}

impl<T> Map<uint, T> for TrieMap<T> {
    /// Return true if the map contains a value for the specified key
    #[inline(always)]
    fn contains_key(&self, key: &uint) -> bool {
        self.find(key).is_some()
    }

    /// Visit all keys in order
    #[inline(always)]
    fn each_key(&self, f: &fn(&uint) -> bool) {
        self.each(|&(k, _)| f(&k))
    }

    /// Visit all values in order
    #[inline(always)]
    fn each_value(&self, f: &fn(&T) -> bool) {
        self.each(|&(_, v)| f(v))
    }

    /// Iterate over the map and mutate the contained values
    #[inline(always)]
    fn mutate_values(&mut self, f: &fn(&uint, &mut T) -> bool) {
        self.root.mutate_values(f);
    }

    /// Return the value corresponding to the key in the map
    #[inline(hint)]
    fn find(&self, key: &uint) -> Option<&'self T> {
        let mut node: &'self TrieNode<T> = &self.root;
        let mut idx = 0;
        loop {
            match node.children[chunk(*key, idx)] {
              Internal(ref x) => node = &**x,
              External(stored, ref value) => {
                if stored == *key {
                    return Some(value)
                } else {
                    return None
                }
              }
              Nothing => return None
            }
            idx += 1;
        }
    }

    /// Insert a key-value pair into the map. An existing value for a
    /// key is replaced by the new value. Return true if the key did
    /// not already exist in the map.
    #[inline(always)]
    fn insert(&mut self, key: uint, value: T) -> bool {
        let ret = insert(&mut self.root.count,
                         &mut self.root.children[chunk(key, 0)],
                         key, value, 1);
        if ret { self.length += 1 }
        ret
    }

    /// Remove a key-value pair from the map. Return true if the key
    /// was present in the map, otherwise false.
    #[inline(always)]
    fn remove(&mut self, key: &uint) -> bool {
        let ret = remove(&mut self.root.count,
                         &mut self.root.children[chunk(*key, 0)],
                         *key, 1);
        if ret { self.length -= 1 }
        ret
    }
}

pub impl<T> TrieMap<T> {
    /// Create an empty TrieMap
    #[inline(always)]
    fn new() -> TrieMap<T> {
        TrieMap{root: TrieNode::new(), length: 0}
    }

    /// Visit all keys in reverse order
    #[inline(always)]
    fn each_key_reverse(&self, f: &fn(&uint) -> bool) {
        self.each_reverse(|&(k, _)| f(&k))
    }

    /// Visit all values in reverse order
    #[inline(always)]
    fn each_value_reverse(&self, f: &fn(&T) -> bool) {
        self.each_reverse(|&(_, v)| f(v))
    }
}

pub struct TrieSet {
    priv map: TrieMap<()>
}

impl BaseIter<uint> for TrieSet {
    /// Visit all values in order
    #[inline(always)]
    fn each(&self, f: &fn(&uint) -> bool) { self.map.each_key(f) }
    #[inline(always)]
    fn size_hint(&self) -> Option<uint> { Some(self.len()) }
}

impl ReverseIter<uint> for TrieSet {
    /// Visit all values in reverse order
    #[inline(always)]
    fn each_reverse(&self, f: &fn(&uint) -> bool) {
        self.map.each_key_reverse(f)
    }
}

impl Container for TrieSet {
    /// Return the number of elements in the set
    #[inline(always)]
    fn len(&const self) -> uint { self.map.len() }

    /// Return true if the set contains no elements
    #[inline(always)]
    fn is_empty(&const self) -> bool { self.map.is_empty() }
}

impl Mutable for TrieSet {
    /// Clear the set, removing all values.
    #[inline(always)]
    fn clear(&mut self) { self.map.clear() }
}

impl TrieSet {
    /// Create an empty TrieSet
    #[inline(always)]
    fn new() -> TrieSet {
        TrieSet{map: TrieMap::new()}
    }

    /// Return true if the set contains a value
    #[inline(always)]
    fn contains(&self, value: &uint) -> bool {
        self.map.contains_key(value)
    }

    /// Add a value to the set. Return true if the value was not already
    /// present in the set.
    #[inline(always)]
    fn insert(&mut self, value: uint) -> bool { self.map.insert(value, ()) }

    /// Remove a value from the set. Return true if the value was
    /// present in the set.
    #[inline(always)]
    fn remove(&mut self, value: &uint) -> bool { self.map.remove(value) }
}

struct TrieNode<T> {
    count: uint,
    children: [Child<T> * SIZE]
}

impl<T> TrieNode<T> {
    #[inline(always)]
    fn new() -> TrieNode<T> {
        // FIXME: #5244: [Nothing, ..SIZE] should be possible without Copy
        TrieNode{count: 0,
                 children: [Nothing, Nothing, Nothing, Nothing,
                            Nothing, Nothing, Nothing, Nothing,
                            Nothing, Nothing, Nothing, Nothing,
                            Nothing, Nothing, Nothing, Nothing]}
    }
}

impl<T> TrieNode<T> {
    fn each(&self, f: &fn(&(uint, &'self T)) -> bool) -> bool {
        for uint::range(0, self.children.len()) |idx| {
            match self.children[idx] {
                Internal(ref x) => if !x.each(f) { return false },
                External(k, ref v) => if !f(&(k, v)) { return false },
                Nothing => ()
            }
        }
        true
    }

    fn each_reverse(&self, f: &fn(&(uint, &'self T)) -> bool) -> bool {
        for uint::range_rev(self.children.len(), 0) |idx| {
            match self.children[idx - 1] {
                Internal(ref x) => if !x.each_reverse(f) { return false },
                External(k, ref v) => if !f(&(k, v)) { return false },
                Nothing => ()
            }
        }
        true
    }

    fn mutate_values(&mut self, f: &fn(&uint, &mut T) -> bool) -> bool {
        for vec::each_mut(self.children) |child| {
            match *child {
                Internal(ref mut x) => if !x.mutate_values(f) {
                    return false
                },
                External(k, ref mut v) => if !f(&k, v) { return false },
                Nothing => ()
            }
        }
        true
    }
}

// if this was done via a trait, the key could be generic
#[inline(always)]
fn chunk(n: uint, idx: uint) -> uint {
    let sh = uint::bits - (SHIFT * (idx + 1));
    (n >> sh) & MASK
}

fn insert<T>(count: &mut uint, child: &mut Child<T>, key: uint, value: T,
             idx: uint) -> bool {
    let mut tmp = Nothing;
    tmp <-> *child;
    let mut added = false;

    *child = match tmp {
      External(stored_key, stored_value) => {
          if stored_key == key {
              External(stored_key, value)
          } else {
              // conflict - split the node
              let mut new = ~TrieNode::new();
              insert(&mut new.count,
                     &mut new.children[chunk(stored_key, idx)],
                     stored_key, stored_value, idx + 1);
              insert(&mut new.count, &mut new.children[chunk(key, idx)], key,
                     value, idx + 1);
              added = true;
              Internal(new)
          }
      }
      Internal(x) => {
        let mut x = x;
        added = insert(&mut x.count, &mut x.children[chunk(key, idx)], key,
                       value, idx + 1);
        Internal(x)

      }
      Nothing => {
        *count += 1;
        added = true;
        External(key, value)
      }
    };
    added
}

fn remove<T>(count: &mut uint, child: &mut Child<T>, key: uint,
             idx: uint) -> bool {
    let (ret, this) = match *child {
      External(stored, _) => {
          if stored == key { (true, true) } else { (false, false) }
      }
      Internal(ref mut x) => {
          let ret = remove(&mut x.count, &mut x.children[chunk(key, idx)],
                           key, idx + 1);
          (ret, x.count == 0)
      }
      Nothing => (false, false)
    };

    if this {
        *child = Nothing;
        *count -= 1;
    }
    ret
}

#[cfg(test)]
pub fn check_integrity<T>(trie: &TrieNode<T>) {
    fail_unless!(trie.count != 0);

    let mut sum = 0;

    for trie.children.each |x| {
        match *x {
          Nothing => (),
          Internal(ref y) => {
              check_integrity(&**y);
              sum += 1
          }
          External(_, _) => { sum += 1 }
        }
    }

    fail_unless!(sum == trie.count);
}

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

    #[test]
    fn test_step() {
        let mut trie = TrieMap::new();
        let n = 300;

        for uint::range_step(1, n, 2) |x| {
            fail_unless!(trie.insert(x, x + 1));
            fail_unless!(trie.contains_key(&x));
            check_integrity(&trie.root);
        }

        for uint::range_step(0, n, 2) |x| {
            fail_unless!(!trie.contains_key(&x));
            fail_unless!(trie.insert(x, x + 1));
            check_integrity(&trie.root);
        }

        for uint::range(0, n) |x| {
            fail_unless!(trie.contains_key(&x));
            fail_unless!(!trie.insert(x, x + 1));
            check_integrity(&trie.root);
        }

        for uint::range_step(1, n, 2) |x| {
            fail_unless!(trie.remove(&x));
            fail_unless!(!trie.contains_key(&x));
            check_integrity(&trie.root);
        }

        for uint::range_step(0, n, 2) |x| {
            fail_unless!(trie.contains_key(&x));
            fail_unless!(!trie.insert(x, x + 1));
            check_integrity(&trie.root);
        }
    }

    #[test]
    fn test_each() {
        let mut m = TrieMap::new();

        fail_unless!(m.insert(3, 6));
        fail_unless!(m.insert(0, 0));
        fail_unless!(m.insert(4, 8));
        fail_unless!(m.insert(2, 4));
        fail_unless!(m.insert(1, 2));

        let mut n = 0;
        for m.each |&(k, v)| {
            fail_unless!(k == n);
            fail_unless!(*v == n * 2);
            n += 1;
        }
    }

    #[test]
    fn test_each_break() {
        let mut m = TrieMap::new();

        for uint::range_rev(uint::max_value, uint::max_value - 10000) |x| {
            m.insert(x, x / 2);
        }

        let mut n = uint::max_value - 9999;
        for m.each |&(k, v)| {
            if n == uint::max_value - 5000 { break }
            fail_unless!(n < uint::max_value - 5000);

            fail_unless!(k == n);
            fail_unless!(*v == n / 2);
            n += 1;
        }
    }

    #[test]
    fn test_each_reverse() {
        let mut m = TrieMap::new();

        fail_unless!(m.insert(3, 6));
        fail_unless!(m.insert(0, 0));
        fail_unless!(m.insert(4, 8));
        fail_unless!(m.insert(2, 4));
        fail_unless!(m.insert(1, 2));

        let mut n = 4;
        for m.each_reverse |&(k, v)| {
            fail_unless!(k == n);
            fail_unless!(*v == n * 2);
            n -= 1;
        }
    }

    #[test]
    fn test_each_reverse_break() {
        let mut m = TrieMap::new();

        for uint::range_rev(uint::max_value, uint::max_value - 10000) |x| {
            m.insert(x, x / 2);
        }

        let mut n = uint::max_value;
        for m.each_reverse |&(k, v)| {
            if n == uint::max_value - 5000 { break }
            fail_unless!(n > uint::max_value - 5000);

            fail_unless!(k == n);
            fail_unless!(*v == n / 2);
            n -= 1;
        }
    }

    #[test]
    fn test_sane_chunk() {
        let x = 1;
        let y = 1 << (uint::bits - 1);

        let mut trie = TrieSet::new();

        fail_unless!(trie.insert(x));
        fail_unless!(trie.insert(y));

        fail_unless!(trie.len() == 2);

        let expected = [x, y];

        let mut i = 0;

        for trie.each |x| {
            fail_unless!(expected[i] == *x);
            i += 1;
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use]
mod macros;

use proc_macro2::{Delimiter, Group, Ident, Punct, Spacing, Span, TokenStream, TokenTree};
use quote::quote;
use std::iter::FromIterator;
use syn::{Expr, Type};

#[test]
fn parse_interpolated_leading_component() {
    // mimics the token stream corresponding to `$mod::rest`
    let tokens = TokenStream::from_iter(vec![
        TokenTree::Group(Group::new(Delimiter::None, quote! { first })),
        TokenTree::Punct(Punct::new(':', Spacing::Joint)),
        TokenTree::Punct(Punct::new(':', Spacing::Alone)),
        TokenTree::Ident(Ident::new("rest", Span::call_site())),
    ]);

    snapshot!(tokens.clone() as Expr, @r###"
    Expr::Path {
        path: Path {
            segments: [
                PathSegment {
                    ident: "first",
                    arguments: None,
                },
                PathSegment {
                    ident: "rest",
                    arguments: None,
                },
            ],
        },
    }
    "###);

    snapshot!(tokens as Type, @r###"
    Type::Path {
        path: Path {
            segments: [
                PathSegment {
                    ident: "first",
                    arguments: None,
                },
                PathSegment {
                    ident: "rest",
                    arguments: None,
                },
            ],
        },
    }
    "###);
}

#[test]
fn print_incomplete_qpath() {
    use quote::ToTokens;
    use syn::TypePath;

    // qpath with `as` token
    let mut ty: TypePath = syn::parse_quote!(<Self as A>::Q);
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`< Self as A > :: Q`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`< Self as A > ::`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`< Self >`)
    "###);
    assert!(ty.path.segments.pop().is_none());

    // qpath without `as` token
    let mut ty: TypePath = syn::parse_quote!(<Self>::A::B);
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`< Self > :: A :: B`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`< Self > :: A ::`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`< Self > ::`)
    "###);
    assert!(ty.path.segments.pop().is_none());

    // normal path
    let mut ty: TypePath = syn::parse_quote!(Self::A::B);
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`Self :: A :: B`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`Self :: A ::`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`Self ::`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(``)
    "###);
    assert!(ty.path.segments.pop().is_none());
}
<commit_msg>Touch up PR 891<commit_after>#[macro_use]
mod macros;

use proc_macro2::{Delimiter, Group, Ident, Punct, Spacing, Span, TokenStream, TokenTree};
use quote::{quote, ToTokens};
use std::iter::FromIterator;
use syn::{parse_quote, Expr, Type, TypePath};

#[test]
fn parse_interpolated_leading_component() {
    // mimics the token stream corresponding to `$mod::rest`
    let tokens = TokenStream::from_iter(vec![
        TokenTree::Group(Group::new(Delimiter::None, quote! { first })),
        TokenTree::Punct(Punct::new(':', Spacing::Joint)),
        TokenTree::Punct(Punct::new(':', Spacing::Alone)),
        TokenTree::Ident(Ident::new("rest", Span::call_site())),
    ]);

    snapshot!(tokens.clone() as Expr, @r###"
    Expr::Path {
        path: Path {
            segments: [
                PathSegment {
                    ident: "first",
                    arguments: None,
                },
                PathSegment {
                    ident: "rest",
                    arguments: None,
                },
            ],
        },
    }
    "###);

    snapshot!(tokens as Type, @r###"
    Type::Path {
        path: Path {
            segments: [
                PathSegment {
                    ident: "first",
                    arguments: None,
                },
                PathSegment {
                    ident: "rest",
                    arguments: None,
                },
            ],
        },
    }
    "###);
}

#[test]
fn print_incomplete_qpath() {
    // qpath with `as` token
    let mut ty: TypePath = parse_quote!(<Self as A>::Q);
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`< Self as A > :: Q`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`< Self as A > ::`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`< Self >`)
    "###);
    assert!(ty.path.segments.pop().is_none());

    // qpath without `as` token
    let mut ty: TypePath = parse_quote!(<Self>::A::B);
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`< Self > :: A :: B`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`< Self > :: A ::`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`< Self > ::`)
    "###);
    assert!(ty.path.segments.pop().is_none());

    // normal path
    let mut ty: TypePath = parse_quote!(Self::A::B);
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`Self :: A :: B`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`Self :: A ::`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(`Self ::`)
    "###);
    assert!(ty.path.segments.pop().is_some());
    snapshot!(ty.to_token_stream(), @r###"
    TokenStream(``)
    "###);
    assert!(ty.path.segments.pop().is_none());
}
<|endoftext|>


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

use crate::{
    application::storage::PeerMetadataStorage,
    peer_manager::{ConnectionNotification, PeerManagerNotification, PeerManagerRequest},
    protocols::{direct_send::Message, rpc::InboundRpcRequest},
    transport::ConnectionMetadata,
    DisconnectReason, ProtocolId,
};
use async_trait::async_trait;
use diem_config::{
    config::{PeerRole, RoleType},
    network_id::{NetworkContext, NetworkId, PeerNetworkId},
};
use diem_types::PeerId;
use futures::StreamExt;
use netcore::transport::ConnectionOrigin;
use std::{collections::HashMap, sync::Arc};

/// A sender to a node to mock an inbound network message from [`PeerManager`]
pub type InboundMessageSender =
    channel::diem_channel::Sender<(PeerId, ProtocolId), PeerManagerNotification>;

/// A sender to a node to mock an inbound connection from [`PeerManager`]
pub type ConnectionUpdateSender = crate::peer_manager::conn_notifs_channel::Sender;

/// A receiver to get outbound network messages to [`PeerManager`]
pub type OutboundMessageReceiver =
    channel::diem_channel::Receiver<(PeerId, ProtocolId), PeerManagerRequest>;

/// A connection handle describing the network for a node.
///
/// Use this to interact with the node
#[derive(Clone)]
pub struct InboundNetworkHandle {
    /// To send new incoming network messages
    pub inbound_message_sender: InboundMessageSender,
    /// To send new incoming connections or disconnections
    pub connection_update_sender: ConnectionUpdateSender,
    /// To update the local state (normally done by peer manager)
    pub peer_metadata_storage: Arc<PeerMetadataStorage>,
}

impl InboundNetworkHandle {
    /// Push connection update, and update the local storage
    pub fn connect(
        &self,
        role: RoleType,
        self_peer_network_id: PeerNetworkId,
        conn_metadata: ConnectionMetadata,
    ) {
        let self_peer_id = self_peer_network_id.peer_id();
        let network_id = self_peer_network_id.network_id();

        // PeerManager pushes this data before it's received by events
        self.peer_metadata_storage
            .insert_connection(network_id, conn_metadata.clone());
        self.connection_update_sender
            .push(
                conn_metadata.remote_peer_id,
                ConnectionNotification::NewPeer(
                    conn_metadata,
                    NetworkContext::new(role, network_id, self_peer_id),
                ),
            )
            .unwrap();
    }

    /// Push disconnect update, and update the local storage
    pub fn disconnect(
        &self,
        role: RoleType,
        self_peer_network_id: PeerNetworkId,
        conn_metadata: ConnectionMetadata,
    ) {
        let self_peer_id = self_peer_network_id.peer_id();
        let network_id = self_peer_network_id.network_id();

        // PeerManager pushes this data before it's received by events
        self.peer_metadata_storage.remove(&PeerNetworkId::new(
            network_id,
            conn_metadata.remote_peer_id,
        ));
        self.connection_update_sender
            .push(
                conn_metadata.remote_peer_id,
                ConnectionNotification::LostPeer(
                    conn_metadata,
                    NetworkContext::new(role, network_id, self_peer_id),
                    DisconnectReason::ConnectionLost,
                ),
            )
            .unwrap();
    }
}

/// An application specific network handle
pub type ApplicationNetworkHandle<Sender, Events> = (NetworkId, Sender, Events);

/// A unique identifier of a node across the entire network
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct NodeId {
    pub owner: u32,
    pub node_type: NodeType,
}

impl NodeId {
    pub fn validator(owner: u32) -> Self {
        Self {
            owner,
            node_type: NodeType::Validator,
        }
    }

    pub fn vfn(owner: u32) -> Self {
        Self {
            owner,
            node_type: NodeType::ValidatorFullNode,
        }
    }

    pub fn pfn(owner: u32) -> Self {
        Self {
            owner,
            node_type: NodeType::PublicFullNode,
        }
    }

    pub fn role(&self) -> RoleType {
        match self.node_type {
            NodeType::Validator => RoleType::Validator,
            _ => RoleType::FullNode,
        }
    }

    pub fn peer_role(&self) -> PeerRole {
        match self.node_type {
            NodeType::Validator => PeerRole::Validator,
            NodeType::ValidatorFullNode => PeerRole::ValidatorFullNode,
            NodeType::PublicFullNode => PeerRole::Unknown,
        }
    }
}

impl std::fmt::Display for NodeId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}-{:?}", self.owner, self.node_type)
    }
}

/// An enum defining the type of node
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum NodeType {
    Validator,
    ValidatorFullNode,
    PublicFullNode,
}

/// A trait defining an application specific node with networking abstracted
///
/// This is built as an abstract implementation of networking around a node
pub trait ApplicationNode {
    fn node_id(&self) -> NodeId;

    /// Default ['ProtocolId`]s to connect with
    fn default_protocols(&self) -> &[ProtocolId];

    /// For sending to this node. Generally should not be used after setup
    fn get_inbound_handle(&self, network_id: NetworkId) -> InboundNetworkHandle;

    /// For adding handles to other peers
    fn add_inbound_handle_for_peer(
        &mut self,
        peer_network_id: PeerNetworkId,
        handle: InboundNetworkHandle,
    );

    /// For sending to other nodes
    fn get_inbound_handle_for_peer(&self, peer_network_id: PeerNetworkId) -> InboundNetworkHandle;

    /// For receiving messages from other nodes
    fn get_outbound_handle(&mut self, network_id: NetworkId) -> &mut OutboundMessageReceiver;

    fn get_peer_metadata_storage(&self) -> &PeerMetadataStorage;

    fn peer_network_ids(&self) -> &HashMap<NetworkId, PeerNetworkId>;
}

/// An extension trait for an `ApplicationNode` to run tests on.
///
/// Handles common implementation and helper functions
#[async_trait]
pub trait TestNode: ApplicationNode + Sync {
    /// Retrieve the [`PeerNetworkId`] for a specific [`NetworkId`].
    ///
    /// There can only be one per network.
    fn peer_network_id(&self, network_id: NetworkId) -> PeerNetworkId {
        *self.peer_network_ids().get(&network_id).unwrap_or_else(|| {
            panic!(
                "Expected network {} to exist on node {}",
                network_id,
                self.node_id()
            )
        })
    }

    /// Retrieve all [`NetworkId`] for the node
    fn network_ids(&self) -> Vec<NetworkId> {
        self.peer_network_ids().keys().copied().collect()
    }

    /// Connects a node to another node.  The other's inbound handle must already be added.
    fn connect(&self, network_id: NetworkId, metadata: ConnectionMetadata) {
        assert_eq!(ConnectionOrigin::Outbound, metadata.origin);
        let self_metadata = self.conn_metadata(network_id, ConnectionOrigin::Inbound, &[]);
        let remote_peer_id = metadata.remote_peer_id;

        // Tell the other node it's good to send to the connected peer now
        let remote_peer_network_id = PeerNetworkId::new(network_id, remote_peer_id);
        self.get_inbound_handle_for_peer(remote_peer_network_id)
            .connect(self.node_id().role(), remote_peer_network_id, self_metadata);

        // Then connect us
        self.connect_self(network_id, metadata);
    }

    /// Connects only the local side, useful for mocking the other node
    fn connect_self(&self, network_id: NetworkId, metadata: ConnectionMetadata) {
        self.get_inbound_handle(network_id).connect(
            self.node_id().role(),
            self.peer_network_id(network_id),
            metadata,
        );
    }

    /// Disconnects a node from another node
    fn disconnect(&self, network_id: NetworkId, metadata: ConnectionMetadata) {
        let self_metadata = self.conn_metadata(network_id, ConnectionOrigin::Inbound, &[]);
        let remote_peer_id = metadata.remote_peer_id;

        // Tell the other node it's disconnected
        let remote_peer_network_id = PeerNetworkId::new(network_id, remote_peer_id);
        self.get_inbound_handle_for_peer(remote_peer_network_id)
            .disconnect(self.node_id().role(), remote_peer_network_id, self_metadata);

        // Then disconnect us
        self.disconnect_self(network_id, metadata);
    }

    /// Disconnects only the local side, useful for mocking the other node
    fn disconnect_self(&self, network_id: NetworkId, metadata: ConnectionMetadata) {
        self.get_inbound_handle(network_id).disconnect(
            self.node_id().role(),
            self.peer_network_id(network_id),
            metadata,
        );
    }

    /// Find a common [`NetworkId`] between nodes based on [`NodeType`]
    fn find_common_network(&self, other: &Self) -> Option<NetworkId> {
        let self_node_type = self.node_id().node_type;
        let other_node_type = other.node_id().node_type;
        match self_node_type {
            NodeType::Validator => match other_node_type {
                NodeType::Validator => Some(NetworkId::Validator),
                NodeType::ValidatorFullNode => Some(NetworkId::Vfn),
                NodeType::PublicFullNode => None,
            },
            NodeType::ValidatorFullNode => match other_node_type {
                NodeType::Validator => Some(NetworkId::Vfn),
                _ => Some(NetworkId::Public),
            },
            NodeType::PublicFullNode => match other_node_type {
                NodeType::Validator => None,
                _ => Some(NetworkId::Public),
            },
        }
    }

    /// Build `ConnectionMetadata` for a connection on another node
    fn conn_metadata(
        &self,
        network_id: NetworkId,
        origin: ConnectionOrigin,
        protocol_ids: &[ProtocolId],
    ) -> ConnectionMetadata {
        mock_conn_metadata(
            self.peer_network_id(network_id),
            self.node_id().peer_role(),
            origin,
            if protocol_ids.is_empty() {
                self.default_protocols()
            } else {
                protocol_ids
            },
        )
    }

    /// Gets the next queued network message on `Node`'s network (`NetworkId`).  Doesn't propagate
    /// to downstream node
    async fn get_next_network_msg(&mut self, network_id: NetworkId) -> PeerManagerRequest {
        self.get_outbound_handle(network_id)
            .next()
            .await
            .expect("Expecting a message")
    }

    /// Sends the next queued network message on `Node`'s network (`NetworkId`)
    async fn send_next_network_msg(&mut self, network_id: NetworkId) {
        let request = self.get_next_network_msg(network_id).await;

        let (remote_peer_id, protocol_id, data, maybe_rpc_info) = match request {
            PeerManagerRequest::SendRpc(peer_id, msg) => (
                peer_id,
                msg.protocol_id,
                msg.data,
                Some((msg.timeout, msg.res_tx)),
            ),
            PeerManagerRequest::SendDirectSend(peer_id, msg) => {
                (peer_id, msg.protocol_id, msg.mdata, None)
            }
        };

        let sender_peer_network_id = self.peer_network_id(network_id);
        let receiver_peer_network_id = PeerNetworkId::new(network_id, remote_peer_id);
        let receiver_handle = self.get_inbound_handle_for_peer(receiver_peer_network_id);
        let sender_peer_id = sender_peer_network_id.peer_id();

        // TODO: Add timeout functionality
        let peer_manager_notif = if let Some((_timeout, res_tx)) = maybe_rpc_info {
            PeerManagerNotification::RecvRpc(
                sender_peer_id,
                InboundRpcRequest {
                    protocol_id,
                    data,
                    res_tx,
                },
            )
        } else {
            PeerManagerNotification::RecvMessage(
                sender_peer_id,
                Message {
                    protocol_id,
                    mdata: data,
                },
            )
        };
        receiver_handle
            .inbound_message_sender
            .push((sender_peer_id, protocol_id), peer_manager_notif)
            .unwrap();
    }
}

/// Creates a [`ConnectionMetadata`].
pub fn mock_conn_metadata(
    peer_network_id: PeerNetworkId,
    peer_role: PeerRole,
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> ConnectionMetadata {
    let mut metadata =
        ConnectionMetadata::mock_with_role_and_origin(peer_network_id.peer_id(), peer_role, origin);
    for protocol_id in protocol_ids {
        metadata.application_protocols.insert(*protocol_id);
    }
    metadata
}

/// Creates a mock connection based on the `Validator` to `Validator` connection
pub fn validator_mock_connection(
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> (PeerNetworkId, ConnectionMetadata) {
    mock_connection(
        NetworkId::Validator,
        PeerRole::Validator,
        origin,
        protocol_ids,
    )
}

/// Creates a mock connection based on the `Vfn` to `Validator` connection
pub fn vfn_validator_mock_connection(
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> (PeerNetworkId, ConnectionMetadata) {
    let peer_role = match origin {
        ConnectionOrigin::Inbound => PeerRole::ValidatorFullNode,
        ConnectionOrigin::Outbound => PeerRole::Validator,
    };
    mock_connection(NetworkId::Vfn, peer_role, origin, protocol_ids)
}

/// Creates a mock connection based on the `Pfn` to `Vfn` connection
pub fn pfn_vfn_mock_connection(
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> (PeerNetworkId, ConnectionMetadata) {
    let peer_role = match origin {
        ConnectionOrigin::Inbound => PeerRole::Unknown,
        ConnectionOrigin::Outbound => PeerRole::ValidatorFullNode,
    };
    mock_connection(NetworkId::Public, peer_role, origin, protocol_ids)
}

/// Creates a mock connection based on the `Vfn` to `Vfn` connection
pub fn vfn_vfn_mock_connection(
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> (PeerNetworkId, ConnectionMetadata) {
    mock_connection(
        NetworkId::Public,
        PeerRole::ValidatorFullNode,
        origin,
        protocol_ids,
    )
}

/// Creates a mock connection based on the `Pfn` to `Pfn` connection
pub fn pfn_pfn_mock_connection(
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> (PeerNetworkId, ConnectionMetadata) {
    mock_connection(NetworkId::Public, PeerRole::Known, origin, protocol_ids)
}

fn mock_connection(
    network_id: NetworkId,
    peer_role: PeerRole,
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> (PeerNetworkId, ConnectionMetadata) {
    let peer = PeerNetworkId::new(network_id, PeerId::random());
    let metadata = mock_conn_metadata(peer, peer_role, origin, protocol_ids);
    (peer, metadata)
}
<commit_msg>[network] Add ability to drop network messages in framework<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    application::storage::PeerMetadataStorage,
    peer_manager::{ConnectionNotification, PeerManagerNotification, PeerManagerRequest},
    protocols::{
        direct_send::Message,
        rpc::{InboundRpcRequest, OutboundRpcRequest},
    },
    transport::ConnectionMetadata,
    DisconnectReason, ProtocolId,
};
use async_trait::async_trait;
use diem_config::{
    config::{PeerRole, RoleType},
    network_id::{NetworkContext, NetworkId, PeerNetworkId},
};
use diem_types::PeerId;
use futures::StreamExt;
use netcore::transport::ConnectionOrigin;
use std::{collections::HashMap, sync::Arc};

/// A sender to a node to mock an inbound network message from [`PeerManager`]
pub type InboundMessageSender =
    channel::diem_channel::Sender<(PeerId, ProtocolId), PeerManagerNotification>;

/// A sender to a node to mock an inbound connection from [`PeerManager`]
pub type ConnectionUpdateSender = crate::peer_manager::conn_notifs_channel::Sender;

/// A receiver to get outbound network messages to [`PeerManager`]
pub type OutboundMessageReceiver =
    channel::diem_channel::Receiver<(PeerId, ProtocolId), PeerManagerRequest>;

/// A connection handle describing the network for a node.
///
/// Use this to interact with the node
#[derive(Clone)]
pub struct InboundNetworkHandle {
    /// To send new incoming network messages
    pub inbound_message_sender: InboundMessageSender,
    /// To send new incoming connections or disconnections
    pub connection_update_sender: ConnectionUpdateSender,
    /// To update the local state (normally done by peer manager)
    pub peer_metadata_storage: Arc<PeerMetadataStorage>,
}

impl InboundNetworkHandle {
    /// Push connection update, and update the local storage
    pub fn connect(
        &self,
        role: RoleType,
        self_peer_network_id: PeerNetworkId,
        conn_metadata: ConnectionMetadata,
    ) {
        let self_peer_id = self_peer_network_id.peer_id();
        let network_id = self_peer_network_id.network_id();

        // PeerManager pushes this data before it's received by events
        self.peer_metadata_storage
            .insert_connection(network_id, conn_metadata.clone());
        self.connection_update_sender
            .push(
                conn_metadata.remote_peer_id,
                ConnectionNotification::NewPeer(
                    conn_metadata,
                    NetworkContext::new(role, network_id, self_peer_id),
                ),
            )
            .unwrap();
    }

    /// Push disconnect update, and update the local storage
    pub fn disconnect(
        &self,
        role: RoleType,
        self_peer_network_id: PeerNetworkId,
        conn_metadata: ConnectionMetadata,
    ) {
        let self_peer_id = self_peer_network_id.peer_id();
        let network_id = self_peer_network_id.network_id();

        // PeerManager pushes this data before it's received by events
        self.peer_metadata_storage.remove(&PeerNetworkId::new(
            network_id,
            conn_metadata.remote_peer_id,
        ));
        self.connection_update_sender
            .push(
                conn_metadata.remote_peer_id,
                ConnectionNotification::LostPeer(
                    conn_metadata,
                    NetworkContext::new(role, network_id, self_peer_id),
                    DisconnectReason::ConnectionLost,
                ),
            )
            .unwrap();
    }
}

/// An application specific network handle
pub type ApplicationNetworkHandle<Sender, Events> = (NetworkId, Sender, Events);

/// A unique identifier of a node across the entire network
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct NodeId {
    pub owner: u32,
    pub node_type: NodeType,
}

impl NodeId {
    pub fn validator(owner: u32) -> Self {
        Self {
            owner,
            node_type: NodeType::Validator,
        }
    }

    pub fn vfn(owner: u32) -> Self {
        Self {
            owner,
            node_type: NodeType::ValidatorFullNode,
        }
    }

    pub fn pfn(owner: u32) -> Self {
        Self {
            owner,
            node_type: NodeType::PublicFullNode,
        }
    }

    pub fn role(&self) -> RoleType {
        match self.node_type {
            NodeType::Validator => RoleType::Validator,
            _ => RoleType::FullNode,
        }
    }

    pub fn peer_role(&self) -> PeerRole {
        match self.node_type {
            NodeType::Validator => PeerRole::Validator,
            NodeType::ValidatorFullNode => PeerRole::ValidatorFullNode,
            NodeType::PublicFullNode => PeerRole::Unknown,
        }
    }
}

impl std::fmt::Display for NodeId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}-{:?}", self.owner, self.node_type)
    }
}

/// An enum defining the type of node
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum NodeType {
    Validator,
    ValidatorFullNode,
    PublicFullNode,
}

/// A trait defining an application specific node with networking abstracted
///
/// This is built as an abstract implementation of networking around a node
pub trait ApplicationNode {
    fn node_id(&self) -> NodeId;

    /// Default ['ProtocolId`]s to connect with
    fn default_protocols(&self) -> &[ProtocolId];

    /// For sending to this node. Generally should not be used after setup
    fn get_inbound_handle(&self, network_id: NetworkId) -> InboundNetworkHandle;

    /// For adding handles to other peers
    fn add_inbound_handle_for_peer(
        &mut self,
        peer_network_id: PeerNetworkId,
        handle: InboundNetworkHandle,
    );

    /// For sending to other nodes
    fn get_inbound_handle_for_peer(&self, peer_network_id: PeerNetworkId) -> InboundNetworkHandle;

    /// For receiving messages from other nodes
    fn get_outbound_handle(&mut self, network_id: NetworkId) -> &mut OutboundMessageReceiver;

    fn get_peer_metadata_storage(&self) -> &PeerMetadataStorage;

    fn peer_network_ids(&self) -> &HashMap<NetworkId, PeerNetworkId>;
}

/// An extension trait for an `ApplicationNode` to run tests on.
///
/// Handles common implementation and helper functions
#[async_trait]
pub trait TestNode: ApplicationNode + Sync {
    /// Retrieve the [`PeerNetworkId`] for a specific [`NetworkId`].
    ///
    /// There can only be one per network.
    fn peer_network_id(&self, network_id: NetworkId) -> PeerNetworkId {
        *self.peer_network_ids().get(&network_id).unwrap_or_else(|| {
            panic!(
                "Expected network {} to exist on node {}",
                network_id,
                self.node_id()
            )
        })
    }

    /// Retrieve all [`NetworkId`] for the node
    fn network_ids(&self) -> Vec<NetworkId> {
        self.peer_network_ids().keys().copied().collect()
    }

    /// Connects a node to another node.  The other's inbound handle must already be added.
    fn connect(&self, network_id: NetworkId, metadata: ConnectionMetadata) {
        assert_eq!(ConnectionOrigin::Outbound, metadata.origin);
        let self_metadata = self.conn_metadata(network_id, ConnectionOrigin::Inbound, &[]);
        let remote_peer_id = metadata.remote_peer_id;

        // Tell the other node it's good to send to the connected peer now
        let remote_peer_network_id = PeerNetworkId::new(network_id, remote_peer_id);
        self.get_inbound_handle_for_peer(remote_peer_network_id)
            .connect(self.node_id().role(), remote_peer_network_id, self_metadata);

        // Then connect us
        self.connect_self(network_id, metadata);
    }

    /// Connects only the local side, useful for mocking the other node
    fn connect_self(&self, network_id: NetworkId, metadata: ConnectionMetadata) {
        self.get_inbound_handle(network_id).connect(
            self.node_id().role(),
            self.peer_network_id(network_id),
            metadata,
        );
    }

    /// Disconnects a node from another node
    fn disconnect(&self, network_id: NetworkId, metadata: ConnectionMetadata) {
        let self_metadata = self.conn_metadata(network_id, ConnectionOrigin::Inbound, &[]);
        let remote_peer_id = metadata.remote_peer_id;

        // Tell the other node it's disconnected
        let remote_peer_network_id = PeerNetworkId::new(network_id, remote_peer_id);
        self.get_inbound_handle_for_peer(remote_peer_network_id)
            .disconnect(self.node_id().role(), remote_peer_network_id, self_metadata);

        // Then disconnect us
        self.disconnect_self(network_id, metadata);
    }

    /// Disconnects only the local side, useful for mocking the other node
    fn disconnect_self(&self, network_id: NetworkId, metadata: ConnectionMetadata) {
        self.get_inbound_handle(network_id).disconnect(
            self.node_id().role(),
            self.peer_network_id(network_id),
            metadata,
        );
    }

    /// Find a common [`NetworkId`] between nodes based on [`NodeType`]
    fn find_common_network(&self, other: &Self) -> Option<NetworkId> {
        let self_node_type = self.node_id().node_type;
        let other_node_type = other.node_id().node_type;
        match self_node_type {
            NodeType::Validator => match other_node_type {
                NodeType::Validator => Some(NetworkId::Validator),
                NodeType::ValidatorFullNode => Some(NetworkId::Vfn),
                NodeType::PublicFullNode => None,
            },
            NodeType::ValidatorFullNode => match other_node_type {
                NodeType::Validator => Some(NetworkId::Vfn),
                _ => Some(NetworkId::Public),
            },
            NodeType::PublicFullNode => match other_node_type {
                NodeType::Validator => None,
                _ => Some(NetworkId::Public),
            },
        }
    }

    /// Build `ConnectionMetadata` for a connection on another node
    fn conn_metadata(
        &self,
        network_id: NetworkId,
        origin: ConnectionOrigin,
        protocol_ids: &[ProtocolId],
    ) -> ConnectionMetadata {
        mock_conn_metadata(
            self.peer_network_id(network_id),
            self.node_id().peer_role(),
            origin,
            if protocol_ids.is_empty() {
                self.default_protocols()
            } else {
                protocol_ids
            },
        )
    }

    /// Gets the next queued network message on `Node`'s network [`NetworkId`].  Doesn't propagate
    /// to downstream node.  If dropping a message use [`TestNode::drop_next_network_msg`]
    async fn get_next_network_msg(&mut self, network_id: NetworkId) -> PeerManagerRequest {
        self.get_outbound_handle(network_id)
            .next()
            .await
            .expect("Expecting a message")
    }

    /// Drop a network message.  This is required over [`TestNode::get_network_msg`] because the
    /// oneshot channel must be dropped.
    async fn drop_next_network_msg(
        &mut self,
        network_id: NetworkId,
    ) -> (PeerId, ProtocolId, bytes::Bytes) {
        let message = self.get_next_network_msg(network_id).await;
        match message {
            PeerManagerRequest::SendRpc(
                peer_id,
                OutboundRpcRequest {
                    protocol_id,
                    res_tx,
                    data,
                    ..
                },
            ) => {
                // Forcefully close the oneshot channel, otherwise listening task will hang forever.
                drop(res_tx);
                (peer_id, protocol_id, data)
            }
            PeerManagerRequest::SendDirectSend(peer_id, message) => {
                (peer_id, message.protocol_id, message.mdata)
            }
        }
    }

    /// Sends the next queued network message on `Node`'s network (`NetworkId`)
    async fn send_next_network_msg(&mut self, network_id: NetworkId) {
        let request = self.get_next_network_msg(network_id).await;

        let (remote_peer_id, protocol_id, data, maybe_rpc_info) = match request {
            PeerManagerRequest::SendRpc(peer_id, msg) => (
                peer_id,
                msg.protocol_id,
                msg.data,
                Some((msg.timeout, msg.res_tx)),
            ),
            PeerManagerRequest::SendDirectSend(peer_id, msg) => {
                (peer_id, msg.protocol_id, msg.mdata, None)
            }
        };

        let sender_peer_network_id = self.peer_network_id(network_id);
        let receiver_peer_network_id = PeerNetworkId::new(network_id, remote_peer_id);
        let receiver_handle = self.get_inbound_handle_for_peer(receiver_peer_network_id);
        let sender_peer_id = sender_peer_network_id.peer_id();

        // TODO: Add timeout functionality
        let peer_manager_notif = if let Some((_timeout, res_tx)) = maybe_rpc_info {
            PeerManagerNotification::RecvRpc(
                sender_peer_id,
                InboundRpcRequest {
                    protocol_id,
                    data,
                    res_tx,
                },
            )
        } else {
            PeerManagerNotification::RecvMessage(
                sender_peer_id,
                Message {
                    protocol_id,
                    mdata: data,
                },
            )
        };
        receiver_handle
            .inbound_message_sender
            .push((sender_peer_id, protocol_id), peer_manager_notif)
            .unwrap();
    }
}

/// Creates a [`ConnectionMetadata`].
pub fn mock_conn_metadata(
    peer_network_id: PeerNetworkId,
    peer_role: PeerRole,
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> ConnectionMetadata {
    let mut metadata =
        ConnectionMetadata::mock_with_role_and_origin(peer_network_id.peer_id(), peer_role, origin);
    for protocol_id in protocol_ids {
        metadata.application_protocols.insert(*protocol_id);
    }
    metadata
}

/// Creates a mock connection based on the `Validator` to `Validator` connection
pub fn validator_mock_connection(
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> (PeerNetworkId, ConnectionMetadata) {
    mock_connection(
        NetworkId::Validator,
        PeerRole::Validator,
        origin,
        protocol_ids,
    )
}

/// Creates a mock connection based on the `Vfn` to `Validator` connection
pub fn vfn_validator_mock_connection(
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> (PeerNetworkId, ConnectionMetadata) {
    let peer_role = match origin {
        ConnectionOrigin::Inbound => PeerRole::ValidatorFullNode,
        ConnectionOrigin::Outbound => PeerRole::Validator,
    };
    mock_connection(NetworkId::Vfn, peer_role, origin, protocol_ids)
}

/// Creates a mock connection based on the `Pfn` to `Vfn` connection
pub fn pfn_vfn_mock_connection(
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> (PeerNetworkId, ConnectionMetadata) {
    let peer_role = match origin {
        ConnectionOrigin::Inbound => PeerRole::Unknown,
        ConnectionOrigin::Outbound => PeerRole::ValidatorFullNode,
    };
    mock_connection(NetworkId::Public, peer_role, origin, protocol_ids)
}

/// Creates a mock connection based on the `Vfn` to `Vfn` connection
pub fn vfn_vfn_mock_connection(
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> (PeerNetworkId, ConnectionMetadata) {
    mock_connection(
        NetworkId::Public,
        PeerRole::ValidatorFullNode,
        origin,
        protocol_ids,
    )
}

/// Creates a mock connection based on the `Pfn` to `Pfn` connection
pub fn pfn_pfn_mock_connection(
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> (PeerNetworkId, ConnectionMetadata) {
    mock_connection(NetworkId::Public, PeerRole::Known, origin, protocol_ids)
}

fn mock_connection(
    network_id: NetworkId,
    peer_role: PeerRole,
    origin: ConnectionOrigin,
    protocol_ids: &[ProtocolId],
) -> (PeerNetworkId, ConnectionMetadata) {
    let peer = PeerNetworkId::new(network_id, PeerId::random());
    let metadata = mock_conn_metadata(peer, peer_role, origin, protocol_ids);
    (peer, metadata)
}
<|endoftext|>