<commit_before>//! Functions to read and write model specific registers.

use crate::registers::rflags::RFlags;
use crate::structures::gdt::SegmentSelector;
use crate::PrivilegeLevel;
use bit_field::BitField;
use bitflags::bitflags;
use core::convert::TryInto;

/// A model specific register.
#[derive(Debug)]
pub struct Msr(u32);

impl Msr {
    /// Create an instance from a register.
    pub const fn new(reg: u32) -> Msr {
        Msr(reg)
    }
}

/// The Extended Feature Enable Register.
#[derive(Debug)]
pub struct Efer;

/// FS.Base Model Specific Register.
#[derive(Debug)]
pub struct FsBase;

/// GS.Base Model Specific Register.
#[derive(Debug)]
pub struct GsBase;

/// KernelGsBase Model Specific Register.
#[derive(Debug)]
pub struct KernelGsBase;

/// Syscall Register: STAR
#[derive(Debug)]
pub struct Star;

/// Syscall Register: LSTAR
#[derive(Debug)]
pub struct LStar;

/// Syscall Register: SFMASK
#[derive(Debug)]
pub struct SFMask;

impl Efer {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0080);
}

impl FsBase {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0100);
}

impl GsBase {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0101);
}

impl KernelGsBase {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0102);
}

impl Star {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0081);
}

impl LStar {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0082);
}

impl SFMask {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0084);
}

bitflags! {
    /// Flags of the Extended Feature Enable Register.
    pub struct EferFlags: u64 {
        /// Enables the `syscall` and `sysret` instructions.
        const SYSTEM_CALL_EXTENSIONS = 1 << 0;
        /// Activates long mode, requires activating paging.
        const LONG_MODE_ENABLE = 1 << 8;
        /// Indicates that long mode is active.
        const LONG_MODE_ACTIVE = 1 << 10;
        /// Enables the no-execute page-protection feature.
        const NO_EXECUTE_ENABLE = 1 << 11;
        /// Enables SVM extensions.
        const SECURE_VIRTUAL_MACHINE_ENABLE = 1 << 12;
        /// Enable certain limit checks in 64-bit mode.
        const LONG_MODE_SEGMENT_LIMIT_ENABLE = 1 << 13;
        /// Enable the `fxsave` and `fxrstor` instructions to execute faster in 64-bit mode.
        const FAST_FXSAVE_FXRSTOR = 1 << 14;
        /// Changes how the `invlpg` instruction operates on TLB entries of upper-level entries.
        const TRANSLATION_CACHE_EXTENSION = 1 << 15;
    }
}

#[cfg(target_arch = "x86_64")]
mod x86_64 {
    use super::*;
    use crate::addr::VirtAddr;

    impl Msr {
        /// Read 64 bits msr register.
        #[inline]
        pub unsafe fn read(&self) -> u64 {
            #[cfg(feature = "inline_asm")]
            {
                let (high, low): (u32, u32);
                asm!("rdmsr" : "={eax}" (low), "={edx}" (high) : "{ecx}" (self.0) : "memory" : "volatile");
                ((high as u64) << 32) | (low as u64)
            }

            #[cfg(not(feature = "inline_asm"))]
            crate::asm::x86_64_asm_rdmsr(self.0)
        }

        /// Write 64 bits to msr register.
        #[inline]
        pub unsafe fn write(&mut self, value: u64) {
            #[cfg(feature = "inline_asm")]
            {
                let low = value as u32;
                let high = (value >> 32) as u32;
                asm!("wrmsr" :: "{ecx}" (self.0), "{eax}" (low), "{edx}" (high) : "memory" : "volatile" );
            }

            #[cfg(not(feature = "inline_asm"))]
            crate::asm::x86_64_asm_wrmsr(self.0, value);
        }
    }

    impl Efer {
        /// Read the current EFER flags.
        #[inline]
        pub fn read() -> EferFlags {
            EferFlags::from_bits_truncate(Self::read_raw())
        }

        /// Read the current raw EFER flags.
        #[inline]
        pub fn read_raw() -> u64 {
            unsafe { Self::MSR.read() }
        }

        /// Write the EFER flags, preserving reserved values.
        ///
        /// Preserves the value of reserved fields. Unsafe because it's possible to break memory
        /// safety, e.g. by disabling long mode.
        #[inline]
        pub unsafe fn write(flags: EferFlags) {
            let old_value = Self::read_raw();
            let reserved = old_value & !(EferFlags::all().bits());
            let new_value = reserved | flags.bits();

            Self::write_raw(new_value);
        }

        /// Write the EFER flags.
        ///
        /// Does not preserve any bits, including reserved fields. Unsafe because it's possible to
        /// break memory safety, e.g. by disabling long mode.
        #[inline]
        pub unsafe fn write_raw(flags: u64) {
            Self::MSR.write(flags);
        }

        /// Update EFER flags.
        ///
        /// Preserves the value of reserved fields. Unsafe because it's possible to break memory
        /// safety, e.g. by disabling long mode.
        #[inline]
        pub unsafe fn update<F>(f: F)
        where
            F: FnOnce(&mut EferFlags),
        {
            let mut flags = Self::read();
            f(&mut flags);
            Self::write(flags);
        }
    }

    impl FsBase {
        /// Read the current FsBase register.
        #[inline]
        pub fn read() -> VirtAddr {
            VirtAddr::new(unsafe { Self::MSR.read() })
        }

        /// Write a given virtual address to the FS.Base register.
        #[inline]
        pub fn write(address: VirtAddr) {
            unsafe { Self::MSR.write(address.as_u64()) };
        }
    }

    impl GsBase {
        /// Read the current GsBase register.
        #[inline]
        pub fn read() -> VirtAddr {
            VirtAddr::new(unsafe { Self::MSR.read() })
        }

        /// Write a given virtual address to the GS.Base register.
        #[inline]
        pub fn write(address: VirtAddr) {
            unsafe { Self::MSR.write(address.as_u64()) };
        }
    }

    impl KernelGsBase {
        /// Read the current KernelGsBase register.
        #[inline]
        pub fn read() -> VirtAddr {
            VirtAddr::new(unsafe { Self::MSR.read() })
        }

        /// Write a given virtual address to the KernelGsBase register.
        #[inline]
        pub fn write(address: VirtAddr) {
            unsafe { Self::MSR.write(address.as_u64()) };
        }
    }

    impl Star {
        /// Read the Ring 0 and Ring 3 segment bases.
        /// The remaining fields are ignored because they are
        /// not valid for long mode.
        ///
        /// # Returns
        /// - Field 1 (SYSRET): The CS selector is set to this field + 16. SS.Sel is set to
        /// this field + 8. Because SYSRET always returns to CPL 3, the
        /// RPL bits 1:0 should be initialized to 11b.
        /// - Field 2 (SYSCALL): This field is copied directly into CS.Sel. SS.Sel is set to
        ///  this field + 8. Because SYSCALL always switches to CPL 0, the RPL bits
        /// 33:32 should be initialized to 00b.
        #[inline]
        pub fn read_raw() -> (u16, u16) {
            let msr_value = unsafe { Self::MSR.read() };
            let sysret = msr_value.get_bits(48..64);
            let syscall = msr_value.get_bits(32..48);
            (sysret.try_into().unwrap(), syscall.try_into().unwrap())
        }

        /// Read the Ring 0 and Ring 3 segment bases.
        /// Returns
        /// - CS Selector SYSRET
        /// - SS Selector SYSRET
        /// - CS Selector SYSCALL
        /// - SS Selector SYSCALL
        #[inline]
        pub fn read() -> (
            SegmentSelector,
            SegmentSelector,
            SegmentSelector,
            SegmentSelector,
        ) {
            let raw = Self::read_raw();
            return (
                SegmentSelector((raw.0 + 16).try_into().unwrap()),
                SegmentSelector((raw.0 + 8).try_into().unwrap()),
                SegmentSelector((raw.1).try_into().unwrap()),
                SegmentSelector((raw.1 + 8).try_into().unwrap()),
            );
        }

        /// Write the Ring 0 and Ring 3 segment bases.
        /// The remaining fields are ignored because they are
        /// not valid for long mode.
        ///
        /// # Parameters
        /// - sysret: The CS selector is set to this field + 16. SS.Sel is set to
        /// this field + 8. Because SYSRET always returns to CPL 3, the
        /// RPL bits 1:0 should be initialized to 11b.
        /// - syscall: This field is copied directly into CS.Sel. SS.Sel is set to
        ///  this field + 8. Because SYSCALL always switches to CPL 0, the RPL bits
        /// 33:32 should be initialized to 00b.
        ///
        /// # Unsafety
        /// Unsafe because this can cause system instability if passed in the
        /// wrong values for the fields.
        #[inline]
        pub unsafe fn write_raw(sysret: u16, syscall: u16) {
            let mut msr_value = 0u64;
            msr_value.set_bits(48..64, sysret.into());
            msr_value.set_bits(32..48, syscall.into());
            Self::MSR.write(msr_value);
        }

        /// Write the Ring 0 and Ring 3 segment bases.
        /// The remaining fields are ignored because they are
        /// not valid for long mode.
        /// This function will fail if the segment selectors are
        /// not in the correct offset of each other or if the
        /// segment selectors do not have correct privileges.
        pub fn write(
            cs_sysret: SegmentSelector,
            ss_sysret: SegmentSelector,
            cs_syscall: SegmentSelector,
            ss_syscall: SegmentSelector,
        ) -> Result<(), &'static str> {
            if cs_sysret.0 - 16 != ss_sysret.0 - 8 {
                return Err("Sysret CS and SS is not offset by 8.");
            }

            if cs_syscall.0 != ss_syscall.0 - 8 {
                return Err("Syscall CS and SS is not offset by 8.");
            }

            if ss_sysret.rpl() != PrivilegeLevel::Ring3 {
                return Err("Sysret's segment must be a Ring3 segment.");
            }

            if ss_syscall.rpl() != PrivilegeLevel::Ring0 {
                return Err("Syscall's segment must be a Ring0 segment.");
            }

            unsafe { Self::write_raw((ss_sysret.0 - 8).into(), cs_syscall.0.into()) };

            Ok(())
        }
    }

    impl LStar {
        /// Read the current LStar register.
        /// This holds the target RIP of a syscall.
        #[inline]
        pub fn read() -> VirtAddr {
            VirtAddr::new(unsafe { Self::MSR.read() })
        }

        /// Write a given virtual address to the LStar register.
        /// This holds the target RIP of a syscall.
        #[inline]
        pub fn write(address: VirtAddr) {
            unsafe { Self::MSR.write(address.as_u64()) };
        }
    }

    impl SFMask {
        /// Read to the SFMask register.
        /// The SFMASK register is used to specify which RFLAGS bits
        /// are cleared during a SYSCALL. In long mode, SFMASK is used
        /// to specify which RFLAGS bits are cleared when SYSCALL is
        /// executed. If a bit in SFMASK is set to 1, the corresponding
        /// bit in RFLAGS is cleared to 0. If a bit in SFMASK is cleared
        /// to 0, the corresponding rFLAGS bit is not modified.
        #[inline]
        pub fn read() -> RFlags {
            RFlags::from_bits(unsafe { Self::MSR.read() }).unwrap()
        }

        /// Write to the SFMask register.
        /// The SFMASK register is used to specify which RFLAGS bits
        /// are cleared during a SYSCALL. In long mode, SFMASK is used
        /// to specify which RFLAGS bits are cleared when SYSCALL is
        /// executed. If a bit in SFMASK is set to 1, the corresponding
        /// bit in RFLAGS is cleared to 0. If a bit in SFMASK is cleared
        /// to 0, the corresponding rFLAGS bit is not modified.
        #[inline]
        pub fn write(value: RFlags) {
            unsafe { Self::MSR.write(value.bits()) };
        }
    }
}
<commit_msg>unneeded `return` statement<commit_after>//! Functions to read and write model specific registers.

use crate::registers::rflags::RFlags;
use crate::structures::gdt::SegmentSelector;
use crate::PrivilegeLevel;
use bit_field::BitField;
use bitflags::bitflags;
use core::convert::TryInto;

/// A model specific register.
#[derive(Debug)]
pub struct Msr(u32);

impl Msr {
    /// Create an instance from a register.
    pub const fn new(reg: u32) -> Msr {
        Msr(reg)
    }
}

/// The Extended Feature Enable Register.
#[derive(Debug)]
pub struct Efer;

/// FS.Base Model Specific Register.
#[derive(Debug)]
pub struct FsBase;

/// GS.Base Model Specific Register.
#[derive(Debug)]
pub struct GsBase;

/// KernelGsBase Model Specific Register.
#[derive(Debug)]
pub struct KernelGsBase;

/// Syscall Register: STAR
#[derive(Debug)]
pub struct Star;

/// Syscall Register: LSTAR
#[derive(Debug)]
pub struct LStar;

/// Syscall Register: SFMASK
#[derive(Debug)]
pub struct SFMask;

impl Efer {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0080);
}

impl FsBase {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0100);
}

impl GsBase {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0101);
}

impl KernelGsBase {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0102);
}

impl Star {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0081);
}

impl LStar {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0082);
}

impl SFMask {
    /// The underlying model specific register.
    pub const MSR: Msr = Msr(0xC000_0084);
}

bitflags! {
    /// Flags of the Extended Feature Enable Register.
    pub struct EferFlags: u64 {
        /// Enables the `syscall` and `sysret` instructions.
        const SYSTEM_CALL_EXTENSIONS = 1 << 0;
        /// Activates long mode, requires activating paging.
        const LONG_MODE_ENABLE = 1 << 8;
        /// Indicates that long mode is active.
        const LONG_MODE_ACTIVE = 1 << 10;
        /// Enables the no-execute page-protection feature.
        const NO_EXECUTE_ENABLE = 1 << 11;
        /// Enables SVM extensions.
        const SECURE_VIRTUAL_MACHINE_ENABLE = 1 << 12;
        /// Enable certain limit checks in 64-bit mode.
        const LONG_MODE_SEGMENT_LIMIT_ENABLE = 1 << 13;
        /// Enable the `fxsave` and `fxrstor` instructions to execute faster in 64-bit mode.
        const FAST_FXSAVE_FXRSTOR = 1 << 14;
        /// Changes how the `invlpg` instruction operates on TLB entries of upper-level entries.
        const TRANSLATION_CACHE_EXTENSION = 1 << 15;
    }
}

#[cfg(target_arch = "x86_64")]
mod x86_64 {
    use super::*;
    use crate::addr::VirtAddr;

    impl Msr {
        /// Read 64 bits msr register.
        #[inline]
        pub unsafe fn read(&self) -> u64 {
            #[cfg(feature = "inline_asm")]
            {
                let (high, low): (u32, u32);
                asm!("rdmsr" : "={eax}" (low), "={edx}" (high) : "{ecx}" (self.0) : "memory" : "volatile");
                ((high as u64) << 32) | (low as u64)
            }

            #[cfg(not(feature = "inline_asm"))]
            crate::asm::x86_64_asm_rdmsr(self.0)
        }

        /// Write 64 bits to msr register.
        #[inline]
        pub unsafe fn write(&mut self, value: u64) {
            #[cfg(feature = "inline_asm")]
            {
                let low = value as u32;
                let high = (value >> 32) as u32;
                asm!("wrmsr" :: "{ecx}" (self.0), "{eax}" (low), "{edx}" (high) : "memory" : "volatile" );
            }

            #[cfg(not(feature = "inline_asm"))]
            crate::asm::x86_64_asm_wrmsr(self.0, value);
        }
    }

    impl Efer {
        /// Read the current EFER flags.
        #[inline]
        pub fn read() -> EferFlags {
            EferFlags::from_bits_truncate(Self::read_raw())
        }

        /// Read the current raw EFER flags.
        #[inline]
        pub fn read_raw() -> u64 {
            unsafe { Self::MSR.read() }
        }

        /// Write the EFER flags, preserving reserved values.
        ///
        /// Preserves the value of reserved fields. Unsafe because it's possible to break memory
        /// safety, e.g. by disabling long mode.
        #[inline]
        pub unsafe fn write(flags: EferFlags) {
            let old_value = Self::read_raw();
            let reserved = old_value & !(EferFlags::all().bits());
            let new_value = reserved | flags.bits();

            Self::write_raw(new_value);
        }

        /// Write the EFER flags.
        ///
        /// Does not preserve any bits, including reserved fields. Unsafe because it's possible to
        /// break memory safety, e.g. by disabling long mode.
        #[inline]
        pub unsafe fn write_raw(flags: u64) {
            Self::MSR.write(flags);
        }

        /// Update EFER flags.
        ///
        /// Preserves the value of reserved fields. Unsafe because it's possible to break memory
        /// safety, e.g. by disabling long mode.
        #[inline]
        pub unsafe fn update<F>(f: F)
        where
            F: FnOnce(&mut EferFlags),
        {
            let mut flags = Self::read();
            f(&mut flags);
            Self::write(flags);
        }
    }

    impl FsBase {
        /// Read the current FsBase register.
        #[inline]
        pub fn read() -> VirtAddr {
            VirtAddr::new(unsafe { Self::MSR.read() })
        }

        /// Write a given virtual address to the FS.Base register.
        #[inline]
        pub fn write(address: VirtAddr) {
            unsafe { Self::MSR.write(address.as_u64()) };
        }
    }

    impl GsBase {
        /// Read the current GsBase register.
        #[inline]
        pub fn read() -> VirtAddr {
            VirtAddr::new(unsafe { Self::MSR.read() })
        }

        /// Write a given virtual address to the GS.Base register.
        #[inline]
        pub fn write(address: VirtAddr) {
            unsafe { Self::MSR.write(address.as_u64()) };
        }
    }

    impl KernelGsBase {
        /// Read the current KernelGsBase register.
        #[inline]
        pub fn read() -> VirtAddr {
            VirtAddr::new(unsafe { Self::MSR.read() })
        }

        /// Write a given virtual address to the KernelGsBase register.
        #[inline]
        pub fn write(address: VirtAddr) {
            unsafe { Self::MSR.write(address.as_u64()) };
        }
    }

    impl Star {
        /// Read the Ring 0 and Ring 3 segment bases.
        /// The remaining fields are ignored because they are
        /// not valid for long mode.
        ///
        /// # Returns
        /// - Field 1 (SYSRET): The CS selector is set to this field + 16. SS.Sel is set to
        /// this field + 8. Because SYSRET always returns to CPL 3, the
        /// RPL bits 1:0 should be initialized to 11b.
        /// - Field 2 (SYSCALL): This field is copied directly into CS.Sel. SS.Sel is set to
        ///  this field + 8. Because SYSCALL always switches to CPL 0, the RPL bits
        /// 33:32 should be initialized to 00b.
        #[inline]
        pub fn read_raw() -> (u16, u16) {
            let msr_value = unsafe { Self::MSR.read() };
            let sysret = msr_value.get_bits(48..64);
            let syscall = msr_value.get_bits(32..48);
            (sysret.try_into().unwrap(), syscall.try_into().unwrap())
        }

        /// Read the Ring 0 and Ring 3 segment bases.
        /// Returns
        /// - CS Selector SYSRET
        /// - SS Selector SYSRET
        /// - CS Selector SYSCALL
        /// - SS Selector SYSCALL
        #[inline]
        pub fn read() -> (
            SegmentSelector,
            SegmentSelector,
            SegmentSelector,
            SegmentSelector,
        ) {
            let raw = Self::read_raw();
            (
                SegmentSelector((raw.0 + 16).try_into().unwrap()),
                SegmentSelector((raw.0 + 8).try_into().unwrap()),
                SegmentSelector((raw.1).try_into().unwrap()),
                SegmentSelector((raw.1 + 8).try_into().unwrap()),
            )
        }

        /// Write the Ring 0 and Ring 3 segment bases.
        /// The remaining fields are ignored because they are
        /// not valid for long mode.
        ///
        /// # Parameters
        /// - sysret: The CS selector is set to this field + 16. SS.Sel is set to
        /// this field + 8. Because SYSRET always returns to CPL 3, the
        /// RPL bits 1:0 should be initialized to 11b.
        /// - syscall: This field is copied directly into CS.Sel. SS.Sel is set to
        ///  this field + 8. Because SYSCALL always switches to CPL 0, the RPL bits
        /// 33:32 should be initialized to 00b.
        ///
        /// # Unsafety
        /// Unsafe because this can cause system instability if passed in the
        /// wrong values for the fields.
        #[inline]
        pub unsafe fn write_raw(sysret: u16, syscall: u16) {
            let mut msr_value = 0u64;
            msr_value.set_bits(48..64, sysret.into());
            msr_value.set_bits(32..48, syscall.into());
            Self::MSR.write(msr_value);
        }

        /// Write the Ring 0 and Ring 3 segment bases.
        /// The remaining fields are ignored because they are
        /// not valid for long mode.
        /// This function will fail if the segment selectors are
        /// not in the correct offset of each other or if the
        /// segment selectors do not have correct privileges.
        pub fn write(
            cs_sysret: SegmentSelector,
            ss_sysret: SegmentSelector,
            cs_syscall: SegmentSelector,
            ss_syscall: SegmentSelector,
        ) -> Result<(), &'static str> {
            if cs_sysret.0 - 16 != ss_sysret.0 - 8 {
                return Err("Sysret CS and SS is not offset by 8.");
            }

            if cs_syscall.0 != ss_syscall.0 - 8 {
                return Err("Syscall CS and SS is not offset by 8.");
            }

            if ss_sysret.rpl() != PrivilegeLevel::Ring3 {
                return Err("Sysret's segment must be a Ring3 segment.");
            }

            if ss_syscall.rpl() != PrivilegeLevel::Ring0 {
                return Err("Syscall's segment must be a Ring0 segment.");
            }

            unsafe { Self::write_raw((ss_sysret.0 - 8).into(), cs_syscall.0.into()) };

            Ok(())
        }
    }

    impl LStar {
        /// Read the current LStar register.
        /// This holds the target RIP of a syscall.
        #[inline]
        pub fn read() -> VirtAddr {
            VirtAddr::new(unsafe { Self::MSR.read() })
        }

        /// Write a given virtual address to the LStar register.
        /// This holds the target RIP of a syscall.
        #[inline]
        pub fn write(address: VirtAddr) {
            unsafe { Self::MSR.write(address.as_u64()) };
        }
    }

    impl SFMask {
        /// Read to the SFMask register.
        /// The SFMASK register is used to specify which RFLAGS bits
        /// are cleared during a SYSCALL. In long mode, SFMASK is used
        /// to specify which RFLAGS bits are cleared when SYSCALL is
        /// executed. If a bit in SFMASK is set to 1, the corresponding
        /// bit in RFLAGS is cleared to 0. If a bit in SFMASK is cleared
        /// to 0, the corresponding rFLAGS bit is not modified.
        #[inline]
        pub fn read() -> RFlags {
            RFlags::from_bits(unsafe { Self::MSR.read() }).unwrap()
        }

        /// Write to the SFMask register.
        /// The SFMASK register is used to specify which RFLAGS bits
        /// are cleared during a SYSCALL. In long mode, SFMASK is used
        /// to specify which RFLAGS bits are cleared when SYSCALL is
        /// executed. If a bit in SFMASK is set to 1, the corresponding
        /// bit in RFLAGS is cleared to 0. If a bit in SFMASK is cleared
        /// to 0, the corresponding rFLAGS bit is not modified.
        #[inline]
        pub fn write(value: RFlags) {
            unsafe { Self::MSR.write(value.bits()) };
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Mozilla
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use
// this file except in compliance with the License. You may obtain a copy of the
// License at http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

#![allow(dead_code)]

use db::TypedSQLValue;
use edn;
use errors::{ErrorKind, Result};
use edn::symbols;
use mentat_core::{
    attribute,
    Attribute,
    Entid,
    EntidMap,
    IdentMap,
    Schema,
    SchemaMap,
    TypedValue,
    ValueType,
};
use metadata;
use metadata::{
    AttributeAlteration,
};

/// Return `Ok(())` if `schema_map` defines a valid Mentat schema.
fn validate_schema_map(entid_map: &EntidMap, schema_map: &SchemaMap) -> Result<()> {
    for (entid, attribute) in schema_map {
        let ident = || entid_map.get(entid).map(|ident| ident.to_string()).unwrap_or(entid.to_string());
        if attribute.unique == Some(attribute::Unique::Value) && !attribute.index {
            bail!(ErrorKind::BadSchemaAssertion(format!(":db/unique :db/unique_value without :db/index true for entid: {}", ident())))
        }
        if attribute.unique == Some(attribute::Unique::Identity) && !attribute.index {
            bail!(ErrorKind::BadSchemaAssertion(format!(":db/unique :db/unique_identity without :db/index true for entid: {}", ident())))
        }
        if attribute.fulltext && attribute.value_type != ValueType::String {
            bail!(ErrorKind::BadSchemaAssertion(format!(":db/fulltext true without :db/valueType :db.type/string for entid: {}", ident())))
        }
        if attribute.fulltext && !attribute.index {
            bail!(ErrorKind::BadSchemaAssertion(format!(":db/fulltext true without :db/index true for entid: {}", ident())))
        }
        if attribute.component && attribute.value_type != ValueType::Ref {
            bail!(ErrorKind::BadSchemaAssertion(format!(":db/isComponent true without :db/valueType :db.type/ref for entid: {}", ident())))
        }
        // TODO: consider warning if we have :db/index true for :db/valueType :db.type/string,
        // since this may be inefficient.  More generally, we should try to drive complex
        // :db/valueType (string, uri, json in the future) users to opt-in to some hash-indexing
        // scheme, as discussed in https://github.com/mozilla/mentat/issues/69.
    }
    Ok(())
}

#[derive(Clone,Debug,Default,Eq,Hash,Ord,PartialOrd,PartialEq)]
pub struct AttributeBuilder {
    value_type: Option<ValueType>,
    multival: Option<bool>,
    unique: Option<Option<attribute::Unique>>,
    index: Option<bool>,
    fulltext: Option<bool>,
    component: Option<bool>,
}

impl AttributeBuilder {
    pub fn value_type<'a>(&'a mut self, value_type: ValueType) -> &'a mut Self {
        self.value_type = Some(value_type);
        self
    }

    pub fn multival<'a>(&'a mut self, multival: bool) -> &'a mut Self {
        self.multival = Some(multival);
        self
    }

    pub fn unique<'a>(&'a mut self, unique: Option<attribute::Unique>) -> &'a mut Self {
        self.unique = Some(unique);
        self
    }

    pub fn index<'a>(&'a mut self, index: bool) -> &'a mut Self {
        self.index = Some(index);
        self
    }

    pub fn fulltext<'a>(&'a mut self, fulltext: bool) -> &'a mut Self {
        self.fulltext = Some(fulltext);
        self
    }

    pub fn component<'a>(&'a mut self, component: bool) -> &'a mut Self {
        self.component = Some(component);
        self
    }

    pub fn validate_install_attribute(&self) -> Result<()> {
        if self.value_type.is_none() {
            bail!(ErrorKind::BadSchemaAssertion("Schema attribute for new attribute does not set :db/valueType".into()));
        }
        Ok(())
    }

    pub fn validate_alter_attribute(&self) -> Result<()> {
        if self.value_type.is_some() {
            bail!(ErrorKind::BadSchemaAssertion("Schema alteration must not set :db/valueType".into()));
        }
        if self.fulltext.is_some() {
            bail!(ErrorKind::BadSchemaAssertion("Schema alteration must not set :db/fulltext".into()));
        }
        Ok(())
    }

    pub fn build(&self) -> Attribute {
        let mut attribute = Attribute::default();
        if let Some(value_type) = self.value_type {
            attribute.value_type = value_type;
        }
        if let Some(fulltext) = self.fulltext {
            attribute.fulltext = fulltext;
        }
        if let Some(multival) = self.multival {
            attribute.multival = multival;
        }
        if let Some(ref unique) = self.unique {
            attribute.unique = unique.clone();
        }
        if let Some(index) = self.index {
            attribute.index = index;
        }
        if let Some(component) = self.component {
            attribute.component = component;
        }

        attribute
    }

    pub fn mutate(&self, attribute: &mut Attribute) -> Vec<AttributeAlteration> {
        let mut mutations = Vec::new();
        if let Some(multival) = self.multival {
            if multival != attribute.multival {
                attribute.multival = multival;
                mutations.push(AttributeAlteration::Cardinality);
            }
        }
        if let Some(ref unique) = self.unique {
            if *unique != attribute.unique {
                attribute.unique = unique.clone();
                mutations.push(AttributeAlteration::Unique);
            }
        }
        if let Some(index) = self.index {
            if index != attribute.index {
                attribute.index = index;
                mutations.push(AttributeAlteration::Index);
            }
        }
        if let Some(component) = self.component {
            if component != attribute.component {
                attribute.component = component;
                mutations.push(AttributeAlteration::IsComponent);
            }
        }

        mutations
    }
}

pub trait SchemaBuilding {
    fn require_ident(&self, entid: Entid) -> Result<&symbols::NamespacedKeyword>;
    fn require_entid(&self, ident: &symbols::NamespacedKeyword) -> Result<Entid>;
    fn require_attribute_for_entid(&self, entid: Entid) -> Result<&Attribute>;
    fn from_ident_map_and_schema_map(ident_map: IdentMap, schema_map: SchemaMap) -> Result<Schema>;
    fn from_ident_map_and_triples<U>(ident_map: IdentMap, assertions: U) -> Result<Schema>
        where U: IntoIterator<Item=(symbols::NamespacedKeyword, symbols::NamespacedKeyword, TypedValue)>;
}

impl SchemaBuilding for Schema {
    fn require_ident(&self, entid: Entid) -> Result<&symbols::NamespacedKeyword> {
        self.get_ident(entid).ok_or(ErrorKind::UnrecognizedEntid(entid).into())
    }

    fn require_entid(&self, ident: &symbols::NamespacedKeyword) -> Result<Entid> {
        self.get_entid(&ident).ok_or(ErrorKind::UnrecognizedIdent(ident.to_string()).into())
    }

    fn require_attribute_for_entid(&self, entid: Entid) -> Result<&Attribute> {
        self.attribute_for_entid(entid).ok_or(ErrorKind::UnrecognizedEntid(entid).into())
    }

    /// Create a valid `Schema` from the constituent maps.
    fn from_ident_map_and_schema_map(ident_map: IdentMap, schema_map: SchemaMap) -> Result<Schema> {
        let entid_map: EntidMap = ident_map.iter().map(|(k, v)| (v.clone(), k.clone())).collect();

        validate_schema_map(&entid_map, &schema_map)?;

        Ok(Schema {
            ident_map: ident_map,
            entid_map: entid_map,
            schema_map: schema_map,
        })
    }

    /// Turn vec![(NamespacedKeyword(:ident), NamespacedKeyword(:key), TypedValue(:value)), ...] into a Mentat `Schema`.
    fn from_ident_map_and_triples<U>(ident_map: IdentMap, assertions: U) -> Result<Schema>
        where U: IntoIterator<Item=(symbols::NamespacedKeyword, symbols::NamespacedKeyword, TypedValue)>{

        let entid_assertions: Result<Vec<(Entid, Entid, TypedValue)>> = assertions.into_iter().map(|(symbolic_ident, symbolic_attr, value)| {
            let ident: i64 = *ident_map.get(&symbolic_ident).ok_or(ErrorKind::UnrecognizedIdent(symbolic_ident.to_string()))?;
            let attr: i64 = *ident_map.get(&symbolic_attr).ok_or(ErrorKind::UnrecognizedIdent(symbolic_attr.to_string()))?;
            Ok((ident, attr, value))
        }).collect();
        let mut schema = Schema::from_ident_map_and_schema_map(ident_map, SchemaMap::default())?;
        metadata::update_schema_map_from_entid_triples(&mut schema.schema_map, entid_assertions?)?;
        Ok(schema)
    }
}

pub trait SchemaTypeChecking {
    /// Do schema-aware typechecking and coercion.
    ///
    /// Either assert that the given value is in the attribute's value set, or (in limited cases)
    /// coerce the given value into the attribute's value set.
    fn to_typed_value(&self, value: &edn::Value, attribute: &Attribute) -> Result<TypedValue>;
}

impl SchemaTypeChecking for Schema {
    fn to_typed_value(&self, value: &edn::Value, attribute: &Attribute) -> Result<TypedValue> {
        // TODO: encapsulate entid-ident-attribute for better error messages.
        match TypedValue::from_edn_value(value) {
            // We don't recognize this EDN at all.  Get out!
            None => bail!(ErrorKind::BadEDNValuePair(value.clone(), attribute.value_type.clone())),
            Some(typed_value) => match (&attribute.value_type, typed_value) {
                // Most types don't coerce at all.
                (&ValueType::Boolean, tv @ TypedValue::Boolean(_)) => Ok(tv),
                (&ValueType::Long, tv @ TypedValue::Long(_)) => Ok(tv),
                (&ValueType::Double, tv @ TypedValue::Double(_)) => Ok(tv),
                (&ValueType::String, tv @ TypedValue::String(_)) => Ok(tv),
                (&ValueType::Keyword, tv @ TypedValue::Keyword(_)) => Ok(tv),
                // Ref coerces a little: we interpret some things depending on the schema as a Ref.
                (&ValueType::Ref, TypedValue::Long(x)) => Ok(TypedValue::Ref(x)),
                (&ValueType::Ref, TypedValue::Keyword(ref x)) => self.require_entid(&x).map(|entid| TypedValue::Ref(entid)),
                // Otherwise, we have a type mismatch.
                (value_type, _) => bail!(ErrorKind::BadEDNValuePair(value.clone(), value_type.clone())),
            }
        }
    }
}
<commit_msg>Add tests for `validate_schema_map` (#391) r=rnewman<commit_after>// Copyright 2016 Mozilla
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use
// this file except in compliance with the License. You may obtain a copy of the
// License at http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

#![allow(dead_code)]

use db::TypedSQLValue;
use edn;
use errors::{ErrorKind, Result};
use edn::symbols;
use mentat_core::{
    attribute,
    Attribute,
    Entid,
    EntidMap,
    IdentMap,
    Schema,
    SchemaMap,
    TypedValue,
    ValueType,
};
use metadata;
use metadata::{
    AttributeAlteration,
};

/// Return `Ok(())` if `schema_map` defines a valid Mentat schema.
fn validate_schema_map(entid_map: &EntidMap, schema_map: &SchemaMap) -> Result<()> {
    for (entid, attribute) in schema_map {
        let ident = || entid_map.get(entid).map(|ident| ident.to_string()).unwrap_or(entid.to_string());
        if attribute.unique == Some(attribute::Unique::Value) && !attribute.index {
            bail!(ErrorKind::BadSchemaAssertion(format!(":db/unique :db/unique_value without :db/index true for entid: {}", ident())))
        }
        if attribute.unique == Some(attribute::Unique::Identity) && !attribute.index {
            bail!(ErrorKind::BadSchemaAssertion(format!(":db/unique :db/unique_identity without :db/index true for entid: {}", ident())))
        }
        if attribute.fulltext && attribute.value_type != ValueType::String {
            bail!(ErrorKind::BadSchemaAssertion(format!(":db/fulltext true without :db/valueType :db.type/string for entid: {}", ident())))
        }
        if attribute.fulltext && !attribute.index {
            bail!(ErrorKind::BadSchemaAssertion(format!(":db/fulltext true without :db/index true for entid: {}", ident())))
        }
        if attribute.component && attribute.value_type != ValueType::Ref {
            bail!(ErrorKind::BadSchemaAssertion(format!(":db/isComponent true without :db/valueType :db.type/ref for entid: {}", ident())))
        }
        // TODO: consider warning if we have :db/index true for :db/valueType :db.type/string,
        // since this may be inefficient.  More generally, we should try to drive complex
        // :db/valueType (string, uri, json in the future) users to opt-in to some hash-indexing
        // scheme, as discussed in https://github.com/mozilla/mentat/issues/69.
    }
    Ok(())
}

#[derive(Clone,Debug,Default,Eq,Hash,Ord,PartialOrd,PartialEq)]
pub struct AttributeBuilder {
    value_type: Option<ValueType>,
    multival: Option<bool>,
    unique: Option<Option<attribute::Unique>>,
    index: Option<bool>,
    fulltext: Option<bool>,
    component: Option<bool>,
}

impl AttributeBuilder {
    pub fn value_type<'a>(&'a mut self, value_type: ValueType) -> &'a mut Self {
        self.value_type = Some(value_type);
        self
    }

    pub fn multival<'a>(&'a mut self, multival: bool) -> &'a mut Self {
        self.multival = Some(multival);
        self
    }

    pub fn unique<'a>(&'a mut self, unique: Option<attribute::Unique>) -> &'a mut Self {
        self.unique = Some(unique);
        self
    }

    pub fn index<'a>(&'a mut self, index: bool) -> &'a mut Self {
        self.index = Some(index);
        self
    }

    pub fn fulltext<'a>(&'a mut self, fulltext: bool) -> &'a mut Self {
        self.fulltext = Some(fulltext);
        self
    }

    pub fn component<'a>(&'a mut self, component: bool) -> &'a mut Self {
        self.component = Some(component);
        self
    }

    pub fn validate_install_attribute(&self) -> Result<()> {
        if self.value_type.is_none() {
            bail!(ErrorKind::BadSchemaAssertion("Schema attribute for new attribute does not set :db/valueType".into()));
        }
        Ok(())
    }

    pub fn validate_alter_attribute(&self) -> Result<()> {
        if self.value_type.is_some() {
            bail!(ErrorKind::BadSchemaAssertion("Schema alteration must not set :db/valueType".into()));
        }
        if self.fulltext.is_some() {
            bail!(ErrorKind::BadSchemaAssertion("Schema alteration must not set :db/fulltext".into()));
        }
        Ok(())
    }

    pub fn build(&self) -> Attribute {
        let mut attribute = Attribute::default();
        if let Some(value_type) = self.value_type {
            attribute.value_type = value_type;
        }
        if let Some(fulltext) = self.fulltext {
            attribute.fulltext = fulltext;
        }
        if let Some(multival) = self.multival {
            attribute.multival = multival;
        }
        if let Some(ref unique) = self.unique {
            attribute.unique = unique.clone();
        }
        if let Some(index) = self.index {
            attribute.index = index;
        }
        if let Some(component) = self.component {
            attribute.component = component;
        }

        attribute
    }

    pub fn mutate(&self, attribute: &mut Attribute) -> Vec<AttributeAlteration> {
        let mut mutations = Vec::new();
        if let Some(multival) = self.multival {
            if multival != attribute.multival {
                attribute.multival = multival;
                mutations.push(AttributeAlteration::Cardinality);
            }
        }
        if let Some(ref unique) = self.unique {
            if *unique != attribute.unique {
                attribute.unique = unique.clone();
                mutations.push(AttributeAlteration::Unique);
            }
        }
        if let Some(index) = self.index {
            if index != attribute.index {
                attribute.index = index;
                mutations.push(AttributeAlteration::Index);
            }
        }
        if let Some(component) = self.component {
            if component != attribute.component {
                attribute.component = component;
                mutations.push(AttributeAlteration::IsComponent);
            }
        }

        mutations
    }
}

pub trait SchemaBuilding {
    fn require_ident(&self, entid: Entid) -> Result<&symbols::NamespacedKeyword>;
    fn require_entid(&self, ident: &symbols::NamespacedKeyword) -> Result<Entid>;
    fn require_attribute_for_entid(&self, entid: Entid) -> Result<&Attribute>;
    fn from_ident_map_and_schema_map(ident_map: IdentMap, schema_map: SchemaMap) -> Result<Schema>;
    fn from_ident_map_and_triples<U>(ident_map: IdentMap, assertions: U) -> Result<Schema>
        where U: IntoIterator<Item=(symbols::NamespacedKeyword, symbols::NamespacedKeyword, TypedValue)>;
}

impl SchemaBuilding for Schema {
    fn require_ident(&self, entid: Entid) -> Result<&symbols::NamespacedKeyword> {
        self.get_ident(entid).ok_or(ErrorKind::UnrecognizedEntid(entid).into())
    }

    fn require_entid(&self, ident: &symbols::NamespacedKeyword) -> Result<Entid> {
        self.get_entid(&ident).ok_or(ErrorKind::UnrecognizedIdent(ident.to_string()).into())
    }

    fn require_attribute_for_entid(&self, entid: Entid) -> Result<&Attribute> {
        self.attribute_for_entid(entid).ok_or(ErrorKind::UnrecognizedEntid(entid).into())
    }

    /// Create a valid `Schema` from the constituent maps.
    fn from_ident_map_and_schema_map(ident_map: IdentMap, schema_map: SchemaMap) -> Result<Schema> {
        let entid_map: EntidMap = ident_map.iter().map(|(k, v)| (v.clone(), k.clone())).collect();

        validate_schema_map(&entid_map, &schema_map)?;

        Ok(Schema {
            ident_map: ident_map,
            entid_map: entid_map,
            schema_map: schema_map,
        })
    }

    /// Turn vec![(NamespacedKeyword(:ident), NamespacedKeyword(:key), TypedValue(:value)), ...] into a Mentat `Schema`.
    fn from_ident_map_and_triples<U>(ident_map: IdentMap, assertions: U) -> Result<Schema>
        where U: IntoIterator<Item=(symbols::NamespacedKeyword, symbols::NamespacedKeyword, TypedValue)>{

        let entid_assertions: Result<Vec<(Entid, Entid, TypedValue)>> = assertions.into_iter().map(|(symbolic_ident, symbolic_attr, value)| {
            let ident: i64 = *ident_map.get(&symbolic_ident).ok_or(ErrorKind::UnrecognizedIdent(symbolic_ident.to_string()))?;
            let attr: i64 = *ident_map.get(&symbolic_attr).ok_or(ErrorKind::UnrecognizedIdent(symbolic_attr.to_string()))?;
            Ok((ident, attr, value))
        }).collect();
        let mut schema = Schema::from_ident_map_and_schema_map(ident_map, SchemaMap::default())?;
        metadata::update_schema_map_from_entid_triples(&mut schema.schema_map, entid_assertions?)?;
        Ok(schema)
    }
}

pub trait SchemaTypeChecking {
    /// Do schema-aware typechecking and coercion.
    ///
    /// Either assert that the given value is in the attribute's value set, or (in limited cases)
    /// coerce the given value into the attribute's value set.
    fn to_typed_value(&self, value: &edn::Value, attribute: &Attribute) -> Result<TypedValue>;
}

impl SchemaTypeChecking for Schema {
    fn to_typed_value(&self, value: &edn::Value, attribute: &Attribute) -> Result<TypedValue> {
        // TODO: encapsulate entid-ident-attribute for better error messages.
        match TypedValue::from_edn_value(value) {
            // We don't recognize this EDN at all.  Get out!
            None => bail!(ErrorKind::BadEDNValuePair(value.clone(), attribute.value_type.clone())),
            Some(typed_value) => match (&attribute.value_type, typed_value) {
                // Most types don't coerce at all.
                (&ValueType::Boolean, tv @ TypedValue::Boolean(_)) => Ok(tv),
                (&ValueType::Long, tv @ TypedValue::Long(_)) => Ok(tv),
                (&ValueType::Double, tv @ TypedValue::Double(_)) => Ok(tv),
                (&ValueType::String, tv @ TypedValue::String(_)) => Ok(tv),
                (&ValueType::Keyword, tv @ TypedValue::Keyword(_)) => Ok(tv),
                // Ref coerces a little: we interpret some things depending on the schema as a Ref.
                (&ValueType::Ref, TypedValue::Long(x)) => Ok(TypedValue::Ref(x)),
                (&ValueType::Ref, TypedValue::Keyword(ref x)) => self.require_entid(&x).map(|entid| TypedValue::Ref(entid)),
                // Otherwise, we have a type mismatch.
                (value_type, _) => bail!(ErrorKind::BadEDNValuePair(value.clone(), value_type.clone())),
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use self::edn::NamespacedKeyword;
    use errors::Error;

    fn add_attribute(schema: &mut Schema,
            ident: NamespacedKeyword,
            entid: Entid,
            attribute: Attribute) {

        schema.entid_map.insert(entid, ident.clone());
        schema.ident_map.insert(ident.clone(), entid);

        schema.schema_map.insert(entid, attribute);
    }

    #[test]
    fn validate_schema_map_success() {
        let mut schema = Schema::default();
        // attribute that is not an index has no uniqueness
        add_attribute(&mut schema, NamespacedKeyword::new("foo", "bar"), 97, Attribute {
            index: false,
            value_type: ValueType::Boolean,
            fulltext: false,
            unique: None,
            multival: false,
            component: false,
        });
        // attribute is unique by value and an index
        add_attribute(&mut schema, NamespacedKeyword::new("foo", "baz"), 98, Attribute {
            index: true,
            value_type: ValueType::Long,
            fulltext: false,
            unique: Some(attribute::Unique::Value),
            multival: false,
            component: false,
        });
        // attribue is unique by identity and an index
        add_attribute(&mut schema, NamespacedKeyword::new("foo", "bat"), 99, Attribute {
            index: true,
            value_type: ValueType::Ref,
            fulltext: false,
            unique: Some(attribute::Unique::Identity),
            multival: false,
            component: false,
        });
        // attribute is a components and a `Ref`
        add_attribute(&mut schema, NamespacedKeyword::new("foo", "bak"), 100, Attribute {
            index: false,
            value_type: ValueType::Ref,
            fulltext: false,
            unique: None,
            multival: false,
            component: true,
        });
        // fulltext attribute is a string and an index
        add_attribute(&mut schema, NamespacedKeyword::new("foo", "bap"), 101, Attribute {
            index: true,
            value_type: ValueType::String,
            fulltext: true,
            unique: None,
            multival: false,
            component: false,
        });

        assert!(validate_schema_map(&schema.entid_map, &schema.schema_map).is_ok());
    }

    #[test]
    fn invalid_schema_unique_value_not_index() {
        let mut schema = Schema::default();
        // attribute unique by value but not index
        let ident = NamespacedKeyword::new("foo", "bar");
        add_attribute(&mut schema, ident , 99, Attribute {
            index: false,
            value_type: ValueType::Boolean,
            fulltext: false,
            unique: Some(attribute::Unique::Value),
            multival: false,
            component: false,
        });

        let err = validate_schema_map(&schema.entid_map, &schema.schema_map).err();
        assert!(err.is_some());

        match err.unwrap() {
            Error(ErrorKind::BadSchemaAssertion(message), _) => { assert_eq!(message, ":db/unique :db/unique_value without :db/index true for entid: :foo/bar"); },
            x => panic!("expected Bad Schema Assertion error, got {:?}", x),
        }
    }

    #[test]
    fn invalid_schema_unique_identity_not_index() {
        let mut schema = Schema::default();
        // attribute is unique by identity but not index
        add_attribute(&mut schema, NamespacedKeyword::new("foo", "bar"), 99, Attribute {
            index: false,
            value_type: ValueType::Long,
            fulltext: false,
            unique: Some(attribute::Unique::Identity),
            multival: false,
            component: false,
        });

        let err = validate_schema_map(&schema.entid_map, &schema.schema_map).err();
        assert!(err.is_some());

        match err.unwrap() {
            Error(ErrorKind::BadSchemaAssertion(message), _) => { assert_eq!(message, ":db/unique :db/unique_identity without :db/index true for entid: :foo/bar"); },
            x => panic!("expected Bad Schema Assertion error, got {:?}", x),
        }
    }

    #[test]
    fn invalid_schema_component_not_ref() {
        let mut schema = Schema::default();
        // attribute that is a component is not a `Ref`
        add_attribute(&mut schema, NamespacedKeyword::new("foo", "bar"), 99, Attribute {
            index: false,
            value_type: ValueType::Boolean,
            fulltext: false,
            unique: None,
            multival: false,
            component: true,
        });

        let err = validate_schema_map(&schema.entid_map, &schema.schema_map).err();
        assert!(err.is_some());

        match err.unwrap() {
            Error(ErrorKind::BadSchemaAssertion(message), _) => { assert_eq!(message, ":db/isComponent true without :db/valueType :db.type/ref for entid: :foo/bar"); },
            x => panic!("expected Bad Schema Assertion error, got {:?}", x),
        }
    }

    #[test]
    fn invalid_schema_fulltext_not_index() {
        let mut schema = Schema::default();
        // attribute that is fulltext is not an index
        add_attribute(&mut schema, NamespacedKeyword::new("foo", "bar"), 99, Attribute {
            index: false,
            value_type: ValueType::String,
            fulltext: true,
            unique: None,
            multival: false,
            component: false,
        });

        let err = validate_schema_map(&schema.entid_map, &schema.schema_map).err();
        assert!(err.is_some());

        match err.unwrap() {
            Error(ErrorKind::BadSchemaAssertion(message), _) => { assert_eq!(message, ":db/fulltext true without :db/index true for entid: :foo/bar"); },
            x => panic!("expected Bad Schema Assertion error, got {:?}", x),
        }
    }

    fn invalid_schema_fulltext_index_not_string() {
        let mut schema = Schema::default();
        // attribute that is fulltext and not a `String`
        add_attribute(&mut schema, NamespacedKeyword::new("foo", "bar"), 99, Attribute {
            index: true,
            value_type: ValueType::Long,
            fulltext: true,
            unique: None,
            multival: false,
            component: false,
        });

        let err = validate_schema_map(&schema.entid_map, &schema.schema_map).err();
        assert!(err.is_some());

        match err.unwrap() {
            Error(ErrorKind::BadSchemaAssertion(message), _) => { assert_eq!(message, ":db/fulltext true without :db/valueType :db.type/string for entid: :foo/bar"); },
            x => panic!("expected Bad Schema Assertion error, got {:?}", x),
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Memory profiling functions.

use libc::{c_char,c_int,c_void,size_t};
use std::borrow::ToOwned;
use std::ffi::CString;
use std::iter::AdditiveIterator;
use std::old_io::timer::sleep;
#[cfg(target_os="linux")]
use std::old_io::File;
use std::mem::size_of;
use std::ptr::null_mut;
use std::sync::mpsc::{Sender, channel, Receiver};
use std::time::duration::Duration;
use task::spawn_named;
#[cfg(target_os="macos")]
use task_info::task_basic_info::{virtual_size,resident_size};

pub struct MemoryProfilerChan(pub Sender<MemoryProfilerMsg>);

impl MemoryProfilerChan {
    pub fn send(&self, msg: MemoryProfilerMsg) {
        let MemoryProfilerChan(ref c) = *self;
        c.send(msg).unwrap();
    }
}

pub enum MemoryProfilerMsg {
    /// Message used to force print the memory profiling metrics.
    Print,
    /// Tells the memory profiler to shut down.
    Exit,
}

pub struct MemoryProfiler {
    pub port: Receiver<MemoryProfilerMsg>,
}

impl MemoryProfiler {
    pub fn create(period: Option<f64>) -> MemoryProfilerChan {
        let (chan, port) = channel();
        match period {
            Some(period) => {
                let period = Duration::milliseconds((period * 1000f64) as i64);
                let chan = chan.clone();
                spawn_named("Memory profiler timer".to_owned(), move || {
                    loop {
                        sleep(period);
                        if chan.send(MemoryProfilerMsg::Print).is_err() {
                            break;
                        }
                    }
                });
                // Spawn the memory profiler.
                spawn_named("Memory profiler".to_owned(), move || {
                    let memory_profiler = MemoryProfiler::new(port);
                    memory_profiler.start();
                });
            }
            None => {
                // No-op to handle messages when the memory profiler is
                // inactive.
                spawn_named("Memory profiler".to_owned(), move || {
                    loop {
                        match port.recv() {
                            Err(_) | Ok(MemoryProfilerMsg::Exit) => break,
                            _ => {}
                        }
                    }
                });
            }
        }

        MemoryProfilerChan(chan)
    }

    pub fn new(port: Receiver<MemoryProfilerMsg>) -> MemoryProfiler {
        MemoryProfiler {
            port: port
        }
    }

    pub fn start(&self) {
        loop {
            match self.port.recv() {
               Ok(msg) => {
                   if !self.handle_msg(msg) {
                       break
                   }
               }
               _ => break
            }
        }
    }

    fn handle_msg(&self, msg: MemoryProfilerMsg) -> bool {
        match msg {
            MemoryProfilerMsg::Print => {
                self.handle_print_msg();
                true
            },
            MemoryProfilerMsg::Exit => false
        }
    }

    fn print_measurement(path: &str, nbytes: Option<u64>) {
        match nbytes {
            Some(nbytes) => {
                let mebi = 1024f64 * 1024f64;
                println!("{:12.2}: {}", (nbytes as f64) / mebi, path);
            }
            None => {
                println!("{:>12}: {}", "???", path);
            }
        }
    }

    fn handle_print_msg(&self) {
        println!("{:12}: {}", "_size (MiB)_", "_category_");

        // Virtual and physical memory usage, as reported by the OS.
        MemoryProfiler::print_measurement("vsize", get_vsize());
        MemoryProfiler::print_measurement("resident", get_resident());

        for seg in get_resident_segments().iter() {
            MemoryProfiler::print_measurement(seg.0.as_slice(), Some(seg.1));
        }

        // Total number of bytes allocated by the application on the system
        // heap.
        MemoryProfiler::print_measurement("system-heap-allocated",
                                          get_system_heap_allocated());

        // The descriptions of the following jemalloc measurements are taken
        // directly from the jemalloc documentation.

        // "Total number of bytes allocated by the application."
        MemoryProfiler::print_measurement("jemalloc-heap-allocated",
                                          get_jemalloc_stat("stats.allocated"));

        // "Total number of bytes in active pages allocated by the application.
        // This is a multiple of the page size, and greater than or equal to
        // |stats.allocated|."
        MemoryProfiler::print_measurement("jemalloc-heap-active",
                                          get_jemalloc_stat("stats.active"));

        // "Total number of bytes in chunks mapped on behalf of the application.
        // This is a multiple of the chunk size, and is at least as large as
        // |stats.active|. This does not include inactive chunks."
        MemoryProfiler::print_measurement("jemalloc-heap-mapped",
                                          get_jemalloc_stat("stats.mapped"));

        println!("");
    }
}

#[cfg(target_os="linux")]
extern {
    fn mallinfo() -> struct_mallinfo;
}

#[cfg(target_os="linux")]
#[repr(C)]
pub struct struct_mallinfo {
    arena:    c_int,
    ordblks:  c_int,
    smblks:   c_int,
    hblks:    c_int,
    hblkhd:   c_int,
    usmblks:  c_int,
    fsmblks:  c_int,
    uordblks: c_int,
    fordblks: c_int,
    keepcost: c_int,
}

#[cfg(target_os="linux")]
fn get_system_heap_allocated() -> Option<u64> {
    let mut info: struct_mallinfo;
    unsafe {
        info = mallinfo();
    }
    // The documentation in the glibc man page makes it sound like |uordblks|
    // would suffice, but that only gets the small allocations that are put in
    // the brk heap. We need |hblkhd| as well to get the larger allocations
    // that are mmapped.
    Some((info.hblkhd + info.uordblks) as u64)
}

#[cfg(not(target_os="linux"))]
fn get_system_heap_allocated() -> Option<u64> {
    None
}

extern {
    fn je_mallctl(name: *const c_char, oldp: *mut c_void, oldlenp: *mut size_t,
                  newp: *mut c_void, newlen: size_t) -> c_int;
}

fn get_jemalloc_stat(value_name: &str) -> Option<u64> {
    // Before we request the measurement of interest, we first send an "epoch"
    // request. Without that jemalloc gives cached statistics(!) which can be
    // highly inaccurate.
    let epoch_name = "epoch";
    let epoch_c_name = CString::from_slice(epoch_name.as_bytes());
    let mut epoch: u64 = 0;
    let epoch_ptr = &mut epoch as *mut _ as *mut c_void;
    let mut epoch_len = size_of::<u64>() as size_t;

    let value_c_name = CString::from_slice(value_name.as_bytes());
    let mut value: size_t = 0;
    let value_ptr = &mut value as *mut _ as *mut c_void;
    let mut value_len = size_of::<size_t>() as size_t;

    // Using the same values for the `old` and `new` parameters is enough
    // to get the statistics updated.
    let rv = unsafe {
        je_mallctl(epoch_c_name.as_ptr(), epoch_ptr, &mut epoch_len, epoch_ptr,
                   epoch_len)
    };
    if rv != 0 {
        return None;
    }

    let rv = unsafe {
        je_mallctl(value_c_name.as_ptr(), value_ptr, &mut value_len,
                   null_mut(), 0)
    };
    if rv != 0 {
        return None;
    }

    Some(value as u64)
}

// Like std::macros::try!, but for Option<>.
macro_rules! option_try(
    ($e:expr) => (match $e { Some(e) => e, None => return None })
);

#[cfg(target_os="linux")]
fn get_proc_self_statm_field(field: usize) -> Option<u64> {
    let mut f = File::open(&Path::new("/proc/self/statm"));
    match f.read_to_string() {
        Ok(contents) => {
            let s = option_try!(contents.as_slice().words().nth(field));
            let npages = option_try!(s.parse::<u64>().ok());
            Some(npages * (::std::env::page_size() as u64))
        }
        Err(_) => None
    }
}

#[cfg(target_os="linux")]
fn get_vsize() -> Option<u64> {
    get_proc_self_statm_field(0)
}

#[cfg(target_os="linux")]
fn get_resident() -> Option<u64> {
    get_proc_self_statm_field(1)
}

#[cfg(target_os="macos")]
fn get_vsize() -> Option<u64> {
    virtual_size()
}

#[cfg(target_os="macos")]
fn get_resident() -> Option<u64> {
    resident_size()
}

#[cfg(not(any(target_os="linux", target_os = "macos")))]
fn get_vsize() -> Option<u64> {
    None
}

#[cfg(not(any(target_os="linux", target_os = "macos")))]
fn get_resident() -> Option<u64> {
    None
}

#[cfg(target_os="linux")]
fn get_resident_segments() -> Vec<(String, u64)> {
    use regex::Regex;
    use std::collections::HashMap;
    use std::collections::hash_map::Entry;

    // The first line of an entry in /proc/<pid>/smaps looks just like an entry
    // in /proc/<pid>/maps:
    //
    //   address           perms offset  dev   inode  pathname
    //   02366000-025d8000 rw-p 00000000 00:00 0      [heap]
    //
    // Each of the following lines contains a key and a value, separated
    // by ": ", where the key does not contain either of those characters.
    // For example:
    //
    //   Rss:           132 kB

    let path = Path::new("/proc/self/smaps");
    let mut f = ::std::old_io::BufferedReader::new(File::open(&path));

    let seg_re = Regex::new(
        r"^[:xdigit:]+-[:xdigit:]+ (....) [:xdigit:]+ [:xdigit:]+:[:xdigit:]+ \d+ +(.*)").unwrap();
    let rss_re = Regex::new(r"^Rss: +(\d+) kB").unwrap();

    // We record each segment's resident size.
    let mut seg_map: HashMap<String, u64> = HashMap::new();

    #[derive(PartialEq)]
    enum LookingFor { Segment, Rss }
    let mut looking_for = LookingFor::Segment;

    let mut curr_seg_name = String::new();

    // Parse the file.
    for line in f.lines() {
        let line = match line {
            Ok(line) => line,
            Err(_) => continue,
        };
        if looking_for == LookingFor::Segment {
            // Look for a segment info line.
            let cap = match seg_re.captures(line.as_slice()) {
                Some(cap) => cap,
                None => continue,
            };
            let perms = cap.at(1).unwrap();
            let pathname = cap.at(2).unwrap();

            // Construct the segment name from its pathname and permissions.
            curr_seg_name.clear();
            curr_seg_name.push_str("- ");
            if pathname == "" || pathname.starts_with("[stack:") {
                // Anonymous memory. Entries marked with "[stack:nnn]"
                // look like thread stacks but they may include other
                // anonymous mappings, so we can't trust them and just
                // treat them as entirely anonymous.
                curr_seg_name.push_str("anonymous");
            } else {
                curr_seg_name.push_str(pathname);
            }
            curr_seg_name.push_str(" (");
            curr_seg_name.push_str(perms);
            curr_seg_name.push_str(")");

            looking_for = LookingFor::Rss;
        } else {
            // Look for an "Rss:" line.
            let cap = match rss_re.captures(line.as_slice()) {
                Some(cap) => cap,
                None => continue,
            };
            let rss = cap.at(1).unwrap().parse::<u64>().unwrap() * 1024;

            if rss > 0 {
                // Aggregate small segments into "- other".
                let seg_name = if rss < 512 * 1024 {
                    "- other".to_owned()
                } else {
                    curr_seg_name.clone()
                };
                match seg_map.entry(seg_name) {
                    Entry::Vacant(entry) => { entry.insert(rss); },
                    Entry::Occupied(mut entry) => *entry.get_mut() += rss,
                }
            }

            looking_for = LookingFor::Segment;
        }
    }

    let mut segs: Vec<(String, u64)> = seg_map.into_iter().collect();

    // Get the total and add it to the vector. Note that this total differs
    // from the "resident" measurement obtained via /proc/<pid>/statm in
    // get_resident(). It's unclear why this difference occurs; for some
    // processes the measurements match, but for Servo they do not.
    let total = segs.iter().map(|&(_, size)| size).sum();
    segs.push(("resident-according-to-smaps".to_owned(), total));

    // Sort by size; the total will be first.
    segs.sort_by(|&(_, rss1), &(_, rss2)| rss2.cmp(&rss1));

    segs
}

#[cfg(not(target_os="linux"))]
fn get_resident_segments() -> Vec<(String, u64)> {
    vec![]
}

<commit_msg>Cleanup compilation warning in OSX (should still work on Linux).<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Memory profiling functions.

use libc::{c_char,c_int,c_void,size_t};
use std::borrow::ToOwned;
use std::ffi::CString;
#[cfg(target_os = "linux")]
use std::iter::AdditiveIterator;
use std::old_io::timer::sleep;
#[cfg(target_os="linux")]
use std::old_io::File;
use std::mem::size_of;
use std::ptr::null_mut;
use std::sync::mpsc::{Sender, channel, Receiver};
use std::time::duration::Duration;
use task::spawn_named;
#[cfg(target_os="macos")]
use task_info::task_basic_info::{virtual_size,resident_size};

pub struct MemoryProfilerChan(pub Sender<MemoryProfilerMsg>);

impl MemoryProfilerChan {
    pub fn send(&self, msg: MemoryProfilerMsg) {
        let MemoryProfilerChan(ref c) = *self;
        c.send(msg).unwrap();
    }
}

pub enum MemoryProfilerMsg {
    /// Message used to force print the memory profiling metrics.
    Print,
    /// Tells the memory profiler to shut down.
    Exit,
}

pub struct MemoryProfiler {
    pub port: Receiver<MemoryProfilerMsg>,
}

impl MemoryProfiler {
    pub fn create(period: Option<f64>) -> MemoryProfilerChan {
        let (chan, port) = channel();
        match period {
            Some(period) => {
                let period = Duration::milliseconds((period * 1000f64) as i64);
                let chan = chan.clone();
                spawn_named("Memory profiler timer".to_owned(), move || {
                    loop {
                        sleep(period);
                        if chan.send(MemoryProfilerMsg::Print).is_err() {
                            break;
                        }
                    }
                });
                // Spawn the memory profiler.
                spawn_named("Memory profiler".to_owned(), move || {
                    let memory_profiler = MemoryProfiler::new(port);
                    memory_profiler.start();
                });
            }
            None => {
                // No-op to handle messages when the memory profiler is
                // inactive.
                spawn_named("Memory profiler".to_owned(), move || {
                    loop {
                        match port.recv() {
                            Err(_) | Ok(MemoryProfilerMsg::Exit) => break,
                            _ => {}
                        }
                    }
                });
            }
        }

        MemoryProfilerChan(chan)
    }

    pub fn new(port: Receiver<MemoryProfilerMsg>) -> MemoryProfiler {
        MemoryProfiler {
            port: port
        }
    }

    pub fn start(&self) {
        loop {
            match self.port.recv() {
               Ok(msg) => {
                   if !self.handle_msg(msg) {
                       break
                   }
               }
               _ => break
            }
        }
    }

    fn handle_msg(&self, msg: MemoryProfilerMsg) -> bool {
        match msg {
            MemoryProfilerMsg::Print => {
                self.handle_print_msg();
                true
            },
            MemoryProfilerMsg::Exit => false
        }
    }

    fn print_measurement(path: &str, nbytes: Option<u64>) {
        match nbytes {
            Some(nbytes) => {
                let mebi = 1024f64 * 1024f64;
                println!("{:12.2}: {}", (nbytes as f64) / mebi, path);
            }
            None => {
                println!("{:>12}: {}", "???", path);
            }
        }
    }

    fn handle_print_msg(&self) {
        println!("{:12}: {}", "_size (MiB)_", "_category_");

        // Virtual and physical memory usage, as reported by the OS.
        MemoryProfiler::print_measurement("vsize", get_vsize());
        MemoryProfiler::print_measurement("resident", get_resident());

        for seg in get_resident_segments().iter() {
            MemoryProfiler::print_measurement(seg.0.as_slice(), Some(seg.1));
        }

        // Total number of bytes allocated by the application on the system
        // heap.
        MemoryProfiler::print_measurement("system-heap-allocated",
                                          get_system_heap_allocated());

        // The descriptions of the following jemalloc measurements are taken
        // directly from the jemalloc documentation.

        // "Total number of bytes allocated by the application."
        MemoryProfiler::print_measurement("jemalloc-heap-allocated",
                                          get_jemalloc_stat("stats.allocated"));

        // "Total number of bytes in active pages allocated by the application.
        // This is a multiple of the page size, and greater than or equal to
        // |stats.allocated|."
        MemoryProfiler::print_measurement("jemalloc-heap-active",
                                          get_jemalloc_stat("stats.active"));

        // "Total number of bytes in chunks mapped on behalf of the application.
        // This is a multiple of the chunk size, and is at least as large as
        // |stats.active|. This does not include inactive chunks."
        MemoryProfiler::print_measurement("jemalloc-heap-mapped",
                                          get_jemalloc_stat("stats.mapped"));

        println!("");
    }
}

#[cfg(target_os="linux")]
extern {
    fn mallinfo() -> struct_mallinfo;
}

#[cfg(target_os="linux")]
#[repr(C)]
pub struct struct_mallinfo {
    arena:    c_int,
    ordblks:  c_int,
    smblks:   c_int,
    hblks:    c_int,
    hblkhd:   c_int,
    usmblks:  c_int,
    fsmblks:  c_int,
    uordblks: c_int,
    fordblks: c_int,
    keepcost: c_int,
}

#[cfg(target_os="linux")]
fn get_system_heap_allocated() -> Option<u64> {
    let mut info: struct_mallinfo;
    unsafe {
        info = mallinfo();
    }
    // The documentation in the glibc man page makes it sound like |uordblks|
    // would suffice, but that only gets the small allocations that are put in
    // the brk heap. We need |hblkhd| as well to get the larger allocations
    // that are mmapped.
    Some((info.hblkhd + info.uordblks) as u64)
}

#[cfg(not(target_os="linux"))]
fn get_system_heap_allocated() -> Option<u64> {
    None
}

extern {
    fn je_mallctl(name: *const c_char, oldp: *mut c_void, oldlenp: *mut size_t,
                  newp: *mut c_void, newlen: size_t) -> c_int;
}

fn get_jemalloc_stat(value_name: &str) -> Option<u64> {
    // Before we request the measurement of interest, we first send an "epoch"
    // request. Without that jemalloc gives cached statistics(!) which can be
    // highly inaccurate.
    let epoch_name = "epoch";
    let epoch_c_name = CString::from_slice(epoch_name.as_bytes());
    let mut epoch: u64 = 0;
    let epoch_ptr = &mut epoch as *mut _ as *mut c_void;
    let mut epoch_len = size_of::<u64>() as size_t;

    let value_c_name = CString::from_slice(value_name.as_bytes());
    let mut value: size_t = 0;
    let value_ptr = &mut value as *mut _ as *mut c_void;
    let mut value_len = size_of::<size_t>() as size_t;

    // Using the same values for the `old` and `new` parameters is enough
    // to get the statistics updated.
    let rv = unsafe {
        je_mallctl(epoch_c_name.as_ptr(), epoch_ptr, &mut epoch_len, epoch_ptr,
                   epoch_len)
    };
    if rv != 0 {
        return None;
    }

    let rv = unsafe {
        je_mallctl(value_c_name.as_ptr(), value_ptr, &mut value_len,
                   null_mut(), 0)
    };
    if rv != 0 {
        return None;
    }

    Some(value as u64)
}

// Like std::macros::try!, but for Option<>.
macro_rules! option_try(
    ($e:expr) => (match $e { Some(e) => e, None => return None })
);

#[cfg(target_os="linux")]
fn get_proc_self_statm_field(field: usize) -> Option<u64> {
    let mut f = File::open(&Path::new("/proc/self/statm"));
    match f.read_to_string() {
        Ok(contents) => {
            let s = option_try!(contents.as_slice().words().nth(field));
            let npages = option_try!(s.parse::<u64>().ok());
            Some(npages * (::std::env::page_size() as u64))
        }
        Err(_) => None
    }
}

#[cfg(target_os="linux")]
fn get_vsize() -> Option<u64> {
    get_proc_self_statm_field(0)
}

#[cfg(target_os="linux")]
fn get_resident() -> Option<u64> {
    get_proc_self_statm_field(1)
}

#[cfg(target_os="macos")]
fn get_vsize() -> Option<u64> {
    virtual_size()
}

#[cfg(target_os="macos")]
fn get_resident() -> Option<u64> {
    resident_size()
}

#[cfg(not(any(target_os="linux", target_os = "macos")))]
fn get_vsize() -> Option<u64> {
    None
}

#[cfg(not(any(target_os="linux", target_os = "macos")))]
fn get_resident() -> Option<u64> {
    None
}

#[cfg(target_os="linux")]
fn get_resident_segments() -> Vec<(String, u64)> {
    use regex::Regex;
    use std::collections::HashMap;
    use std::collections::hash_map::Entry;

    // The first line of an entry in /proc/<pid>/smaps looks just like an entry
    // in /proc/<pid>/maps:
    //
    //   address           perms offset  dev   inode  pathname
    //   02366000-025d8000 rw-p 00000000 00:00 0      [heap]
    //
    // Each of the following lines contains a key and a value, separated
    // by ": ", where the key does not contain either of those characters.
    // For example:
    //
    //   Rss:           132 kB

    let path = Path::new("/proc/self/smaps");
    let mut f = ::std::old_io::BufferedReader::new(File::open(&path));

    let seg_re = Regex::new(
        r"^[:xdigit:]+-[:xdigit:]+ (....) [:xdigit:]+ [:xdigit:]+:[:xdigit:]+ \d+ +(.*)").unwrap();
    let rss_re = Regex::new(r"^Rss: +(\d+) kB").unwrap();

    // We record each segment's resident size.
    let mut seg_map: HashMap<String, u64> = HashMap::new();

    #[derive(PartialEq)]
    enum LookingFor { Segment, Rss }
    let mut looking_for = LookingFor::Segment;

    let mut curr_seg_name = String::new();

    // Parse the file.
    for line in f.lines() {
        let line = match line {
            Ok(line) => line,
            Err(_) => continue,
        };
        if looking_for == LookingFor::Segment {
            // Look for a segment info line.
            let cap = match seg_re.captures(line.as_slice()) {
                Some(cap) => cap,
                None => continue,
            };
            let perms = cap.at(1).unwrap();
            let pathname = cap.at(2).unwrap();

            // Construct the segment name from its pathname and permissions.
            curr_seg_name.clear();
            curr_seg_name.push_str("- ");
            if pathname == "" || pathname.starts_with("[stack:") {
                // Anonymous memory. Entries marked with "[stack:nnn]"
                // look like thread stacks but they may include other
                // anonymous mappings, so we can't trust them and just
                // treat them as entirely anonymous.
                curr_seg_name.push_str("anonymous");
            } else {
                curr_seg_name.push_str(pathname);
            }
            curr_seg_name.push_str(" (");
            curr_seg_name.push_str(perms);
            curr_seg_name.push_str(")");

            looking_for = LookingFor::Rss;
        } else {
            // Look for an "Rss:" line.
            let cap = match rss_re.captures(line.as_slice()) {
                Some(cap) => cap,
                None => continue,
            };
            let rss = cap.at(1).unwrap().parse::<u64>().unwrap() * 1024;

            if rss > 0 {
                // Aggregate small segments into "- other".
                let seg_name = if rss < 512 * 1024 {
                    "- other".to_owned()
                } else {
                    curr_seg_name.clone()
                };
                match seg_map.entry(seg_name) {
                    Entry::Vacant(entry) => { entry.insert(rss); },
                    Entry::Occupied(mut entry) => *entry.get_mut() += rss,
                }
            }

            looking_for = LookingFor::Segment;
        }
    }

    let mut segs: Vec<(String, u64)> = seg_map.into_iter().collect();

    // Get the total and add it to the vector. Note that this total differs
    // from the "resident" measurement obtained via /proc/<pid>/statm in
    // get_resident(). It's unclear why this difference occurs; for some
    // processes the measurements match, but for Servo they do not.
    let total = segs.iter().map(|&(_, size)| size).sum();
    segs.push(("resident-according-to-smaps".to_owned(), total));

    // Sort by size; the total will be first.
    segs.sort_by(|&(_, rss1), &(_, rss2)| rss2.cmp(&rss1));

    segs
}

#[cfg(not(target_os="linux"))]
fn get_resident_segments() -> Vec<(String, u64)> {
    vec![]
}

<|endoftext|>


<commit_before>//! Procedural macro implementing `#[derive(Enum)]`
//!
//! This is supposed to used with `enum-map` crate, which provides the
//! actual usage documentation.

mod derive_enum;
mod derive_struct;

use proc_macro2::TokenStream;
use quote::quote;
use syn::{Data, DeriveInput, Type};

/// Derive macro generating an implementation of trait `Enum`.
///
/// When using a derive, enum maps are maintained in the order in which
/// enum variants are declared. This is reflected in the iterators of
/// enum map as well as its `as_slice` method.
///
/// # Examples
///
/// ## Enums Without Payload
/// ```
/// use enum_map::Enum;
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum A {
///     B,
///     C,
///     D,
/// }
///
/// assert_eq!(A::B.into_usize(), 0);
/// assert_eq!(A::C.into_usize(), 1);
/// assert_eq!(A::D.into_usize(), 2);
///
/// assert_eq!(A::from_usize(0), A::B);
/// assert_eq!(A::from_usize(1), A::C);
/// assert_eq!(A::from_usize(2), A::D);
/// ```
///
/// ## Enums With Payload
///
/// ```
/// use enum_map::Enum;
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum A {
///     B,
///     C,
///     D,
/// }
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum X {
///     Y,
///     Z,
/// }
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum Foo {
///     Bar(bool, A),
///     Empty,
///     Baz { fa: A, fx: X },
/// }
///
/// assert_eq!(Foo::Bar(false, A::B).into_usize(), 0);
/// assert_eq!(Foo::Bar(false, A::D).into_usize(), 4);
/// assert_eq!(Foo::Bar(true, A::B).into_usize(), 1);
/// assert_eq!(Foo::Bar(true, A::C).into_usize(), 3);
/// assert_eq!(Foo::Empty.into_usize(), 6);
/// assert_eq!(Foo::Baz { fa: A::B, fx: X::Y }.into_usize(), 7);
/// assert_eq!(Foo::Baz { fa: A::B, fx: X::Z }.into_usize(), 10);
/// assert_eq!(Foo::Baz { fa: A::D, fx: X::Y }.into_usize(), 9);
///
/// assert_eq!(Foo::from_usize(0), Foo::Bar(false, A::B));
/// assert_eq!(Foo::from_usize(4), Foo::Bar(false, A::D));
/// assert_eq!(Foo::from_usize(1), Foo::Bar(true, A::B));
/// assert_eq!(Foo::from_usize(3), Foo::Bar(true, A::C));
/// assert_eq!(Foo::from_usize(6), Foo::Empty);
/// assert_eq!(Foo::from_usize(7), Foo::Baz { fa: A::B, fx: X::Y });
/// assert_eq!(Foo::from_usize(10), Foo::Baz { fa: A::B, fx: X::Z });
/// assert_eq!(Foo::from_usize(9), Foo::Baz { fa: A::D, fx: X::Y });
///
/// ```
///
/// ## Structs
///
/// ```
/// use enum_map::Enum;
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum A {
///     B,
///     C,
///     D,
/// }
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum X {
///     Y,
///     Z,
/// }
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// struct Foo {
///     bar: bool,
///     baz: A,
///     end: X,
/// }
///
/// assert_eq!(Foo { bar: false, baz: A::B, end: X::Y }.into_usize(), 0);
/// assert_eq!(Foo { bar: true, baz: A::B, end: X::Y }.into_usize(), 1);
/// assert_eq!(Foo { bar: false, baz: A::D, end: X::Y }.into_usize(), 4);
/// assert_eq!(Foo { bar: true, baz: A::C, end: X::Z }.into_usize(), 9);
///
/// assert_eq!(Foo::from_usize(0), Foo { bar: false, baz: A::B, end: X::Y });
/// assert_eq!(Foo::from_usize(1), Foo { bar: true, baz: A::B, end: X::Y });
/// assert_eq!(Foo::from_usize(4), Foo { bar: false, baz: A::D, end: X::Y });
/// assert_eq!(Foo::from_usize(9), Foo { bar: true, baz: A::C, end: X::Z });
/// ```
///
/// ## Tuple Structs
///
/// ```
/// use enum_map::Enum;
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum A {
///     B,
///     C,
///     D,
/// }
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum X {
///     Y,
///     Z,
/// }
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// struct Foo(bool, A, X);
///
/// assert_eq!(Foo(false, A::B, X::Y ).into_usize(), 0);
/// assert_eq!(Foo(true, A::B, X::Y ).into_usize(), 1);
/// assert_eq!(Foo(false, A::D, X::Y ).into_usize(), 4);
/// assert_eq!(Foo(true, A::C, X::Z ).into_usize(), 9);
///
/// assert_eq!(Foo::from_usize(0), Foo(false, A::B, X::Y));
/// assert_eq!(Foo::from_usize(1), Foo(true, A::B, X::Y));
/// assert_eq!(Foo::from_usize(4), Foo(false, A::D, X::Y));
/// assert_eq!(Foo::from_usize(9), Foo(true, A::C, X::Z));
#[proc_macro_derive(Enum)]
pub fn derive_enum_map(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input: DeriveInput = syn::parse(input).unwrap();

    let result = match input.data {
        Data::Enum(data_enum) => derive_enum::generate(input.ident, data_enum),
        Data::Struct(data_struct) => derive_struct::generate(input.ident, data_struct),
        _ => quote! { compile_error! {"#[derive(Enum)] is only defined for enums and structs"} },
    };

    result.into()
}

fn type_length(ty: &Type) -> TokenStream {
    quote! {
        <#ty as ::enum_map::Enum>::LENGTH
    }
}
<commit_msg>Document that into_usize generated by derive is deterministic<commit_after>//! Procedural macro implementing `#[derive(Enum)]`
//!
//! This is supposed to used with `enum-map` crate, which provides the
//! actual usage documentation.

mod derive_enum;
mod derive_struct;

use proc_macro2::TokenStream;
use quote::quote;
use syn::{Data, DeriveInput, Type};

/// Derive macro generating an implementation of trait `Enum`.
///
/// When using a derive, enum maps are maintained in the order in which
/// enum variants are declared. This is reflected in the value returned
/// by `Enum::into_usize`, iterators of enum map as well as
/// `EnumMap::as_slice` method.
///
/// # Examples
///
/// ## Enums Without Payload
/// ```
/// use enum_map::Enum;
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum A {
///     B,
///     C,
///     D,
/// }
///
/// assert_eq!(A::B.into_usize(), 0);
/// assert_eq!(A::C.into_usize(), 1);
/// assert_eq!(A::D.into_usize(), 2);
///
/// assert_eq!(A::from_usize(0), A::B);
/// assert_eq!(A::from_usize(1), A::C);
/// assert_eq!(A::from_usize(2), A::D);
/// ```
///
/// ## Enums With Payload
///
/// ```
/// use enum_map::Enum;
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum A {
///     B,
///     C,
///     D,
/// }
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum X {
///     Y,
///     Z,
/// }
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum Foo {
///     Bar(bool, A),
///     Empty,
///     Baz { fa: A, fx: X },
/// }
///
/// assert_eq!(Foo::Bar(false, A::B).into_usize(), 0);
/// assert_eq!(Foo::Bar(false, A::D).into_usize(), 4);
/// assert_eq!(Foo::Bar(true, A::B).into_usize(), 1);
/// assert_eq!(Foo::Bar(true, A::C).into_usize(), 3);
/// assert_eq!(Foo::Empty.into_usize(), 6);
/// assert_eq!(Foo::Baz { fa: A::B, fx: X::Y }.into_usize(), 7);
/// assert_eq!(Foo::Baz { fa: A::B, fx: X::Z }.into_usize(), 10);
/// assert_eq!(Foo::Baz { fa: A::D, fx: X::Y }.into_usize(), 9);
///
/// assert_eq!(Foo::from_usize(0), Foo::Bar(false, A::B));
/// assert_eq!(Foo::from_usize(4), Foo::Bar(false, A::D));
/// assert_eq!(Foo::from_usize(1), Foo::Bar(true, A::B));
/// assert_eq!(Foo::from_usize(3), Foo::Bar(true, A::C));
/// assert_eq!(Foo::from_usize(6), Foo::Empty);
/// assert_eq!(Foo::from_usize(7), Foo::Baz { fa: A::B, fx: X::Y });
/// assert_eq!(Foo::from_usize(10), Foo::Baz { fa: A::B, fx: X::Z });
/// assert_eq!(Foo::from_usize(9), Foo::Baz { fa: A::D, fx: X::Y });
///
/// ```
///
/// ## Structs
///
/// ```
/// use enum_map::Enum;
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum A {
///     B,
///     C,
///     D,
/// }
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum X {
///     Y,
///     Z,
/// }
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// struct Foo {
///     bar: bool,
///     baz: A,
///     end: X,
/// }
///
/// assert_eq!(Foo { bar: false, baz: A::B, end: X::Y }.into_usize(), 0);
/// assert_eq!(Foo { bar: true, baz: A::B, end: X::Y }.into_usize(), 1);
/// assert_eq!(Foo { bar: false, baz: A::D, end: X::Y }.into_usize(), 4);
/// assert_eq!(Foo { bar: true, baz: A::C, end: X::Z }.into_usize(), 9);
///
/// assert_eq!(Foo::from_usize(0), Foo { bar: false, baz: A::B, end: X::Y });
/// assert_eq!(Foo::from_usize(1), Foo { bar: true, baz: A::B, end: X::Y });
/// assert_eq!(Foo::from_usize(4), Foo { bar: false, baz: A::D, end: X::Y });
/// assert_eq!(Foo::from_usize(9), Foo { bar: true, baz: A::C, end: X::Z });
/// ```
///
/// ## Tuple Structs
///
/// ```
/// use enum_map::Enum;
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum A {
///     B,
///     C,
///     D,
/// }
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// enum X {
///     Y,
///     Z,
/// }
///
/// #[derive(Enum, Debug, PartialEq, Eq)]
/// struct Foo(bool, A, X);
///
/// assert_eq!(Foo(false, A::B, X::Y ).into_usize(), 0);
/// assert_eq!(Foo(true, A::B, X::Y ).into_usize(), 1);
/// assert_eq!(Foo(false, A::D, X::Y ).into_usize(), 4);
/// assert_eq!(Foo(true, A::C, X::Z ).into_usize(), 9);
///
/// assert_eq!(Foo::from_usize(0), Foo(false, A::B, X::Y));
/// assert_eq!(Foo::from_usize(1), Foo(true, A::B, X::Y));
/// assert_eq!(Foo::from_usize(4), Foo(false, A::D, X::Y));
/// assert_eq!(Foo::from_usize(9), Foo(true, A::C, X::Z));
#[proc_macro_derive(Enum)]
pub fn derive_enum_map(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input: DeriveInput = syn::parse(input).unwrap();

    let result = match input.data {
        Data::Enum(data_enum) => derive_enum::generate(input.ident, data_enum),
        Data::Struct(data_struct) => derive_struct::generate(input.ident, data_struct),
        _ => quote! { compile_error! {"#[derive(Enum)] is only defined for enums and structs"} },
    };

    result.into()
}

fn type_length(ty: &Type) -> TokenStream {
    quote! {
        <#ty as ::enum_map::Enum>::LENGTH
    }
}
<|endoftext|>