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

use bitflags::bitflags;

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

impl Msr {
    /// Create an instance from a register.
    #[inline]
    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;
        /// 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(feature = "instructions")]
mod x86_64 {
    use super::*;
    use crate::addr::VirtAddr;
    use crate::registers::rflags::RFlags;
    use crate::structures::gdt::SegmentSelector;
    use crate::PrivilegeLevel;
    use bit_field::BitField;
    use core::convert::TryInto;

    impl Msr {
        /// Read 64 bits msr register.
        ///
        /// ## Safety
        ///
        /// The caller must ensure that this read operation has no unsafe side
        /// effects.
        #[inline]
        pub unsafe fn read(&self) -> u64 {
            #[cfg(feature = "inline_asm")]
            {
                let (high, low): (u32, u32);
                asm!("rdmsr", out("eax") low, out("edx") high, in("ecx") self.0, options(nostack));
                ((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.
        ///
        /// ## Safety
        ///
        /// The caller must ensure that this write operation has no unsafe side
        /// effects.
        #[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!("rdmsr", in("ecx") self.0, in("eax") low, in("edx") high, options(nostack))
            }

            #[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.
        ///
        /// ## Safety
        ///
        /// Unsafe because it's possible to break memory
        /// safety with wrong flags, 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.
        ///
        /// ## Safety
        ///
        /// Unsafe because it's possible to
        /// break memory safety with wrong flags, e.g. by disabling long mode.
        #[inline]
        pub unsafe fn write_raw(flags: u64) {
            let mut msr = Self::MSR;
            msr.write(flags);
        }

        /// Update EFER flags.
        ///
        /// Preserves the value of reserved fields.
        ///
        /// ## Safety
        ///
        /// Unsafe because it's possible to break memory
        /// safety with wrong flags, 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) {
            let mut msr = Self::MSR;
            unsafe { 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) {
            let mut msr = Self::MSR;
            unsafe { 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) {
            let mut msr = Self::MSR;
            unsafe { 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),
                SegmentSelector(raw.0 + 8),
                SegmentSelector(raw.1),
                SegmentSelector(raw.1 + 8),
            )
        }

        /// 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.
        ///
        /// # Safety
        ///
        /// 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());
            let mut msr = Self::MSR;
            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.
        #[inline]
        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, cs_syscall.0) };

            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) {
            let mut msr = Self::MSR;
            unsafe { 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) {
            let mut msr = Self::MSR;
            unsafe { msr.write(value.bits()) };
        }
    }
}
<commit_msg>Fix a typo in MSR write<commit_after>//! Functions to read and write model specific registers.

use bitflags::bitflags;

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

impl Msr {
    /// Create an instance from a register.
    #[inline]
    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;
        /// 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(feature = "instructions")]
mod x86_64 {
    use super::*;
    use crate::addr::VirtAddr;
    use crate::registers::rflags::RFlags;
    use crate::structures::gdt::SegmentSelector;
    use crate::PrivilegeLevel;
    use bit_field::BitField;
    use core::convert::TryInto;

    impl Msr {
        /// Read 64 bits msr register.
        ///
        /// ## Safety
        ///
        /// The caller must ensure that this read operation has no unsafe side
        /// effects.
        #[inline]
        pub unsafe fn read(&self) -> u64 {
            #[cfg(feature = "inline_asm")]
            {
                let (high, low): (u32, u32);
                asm!("rdmsr", out("eax") low, out("edx") high, in("ecx") self.0, options(nostack));
                ((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.
        ///
        /// ## Safety
        ///
        /// The caller must ensure that this write operation has no unsafe side
        /// effects.
        #[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", in("ecx") self.0, in("eax") low, in("edx") high, options(nostack))
            }

            #[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.
        ///
        /// ## Safety
        ///
        /// Unsafe because it's possible to break memory
        /// safety with wrong flags, 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.
        ///
        /// ## Safety
        ///
        /// Unsafe because it's possible to
        /// break memory safety with wrong flags, e.g. by disabling long mode.
        #[inline]
        pub unsafe fn write_raw(flags: u64) {
            let mut msr = Self::MSR;
            msr.write(flags);
        }

        /// Update EFER flags.
        ///
        /// Preserves the value of reserved fields.
        ///
        /// ## Safety
        ///
        /// Unsafe because it's possible to break memory
        /// safety with wrong flags, 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) {
            let mut msr = Self::MSR;
            unsafe { 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) {
            let mut msr = Self::MSR;
            unsafe { 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) {
            let mut msr = Self::MSR;
            unsafe { 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),
                SegmentSelector(raw.0 + 8),
                SegmentSelector(raw.1),
                SegmentSelector(raw.1 + 8),
            )
        }

        /// 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.
        ///
        /// # Safety
        ///
        /// 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());
            let mut msr = Self::MSR;
            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.
        #[inline]
        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, cs_syscall.0) };

            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) {
            let mut msr = Self::MSR;
            unsafe { 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) {
            let mut msr = Self::MSR;
            unsafe { msr.write(value.bits()) };
        }
    }
}
<|endoftext|>


<commit_before>use gl;
use glutin;
use native::NativeTaskBuilder;
use std::mem;
use std::sync::atomic::{AtomicUint, Relaxed};
use std::sync::{Arc, Mutex};
use std::task::TaskBuilder;

enum Message {
    EndFrame,
    Execute(proc(&gl::Gl, &mut GLState, (u8, u8), &[String]):Send),
}

pub struct Context {
    commands: Mutex<Sender<Message>>,
    events: Mutex<Receiver<glutin::Event>>,

    /// Dimensions of the frame buffer.
    dimensions: Arc<(AtomicUint, AtomicUint)>,
}

/// Represents the current OpenGL state.
/// The current state is passed to each function and can be freely updated.
pub struct GLState {
    /// Whether GL_BLEND is enabled
    pub enabled_blend: bool,

    /// Whether GL_CULL_FACE is enabled
    pub enabled_cull_face: bool,

    /// Whether GL_DEPTH_TEST is enabled
    pub enabled_depth_test: bool,

    /// Whether GL_DITHER is enabled
    pub enabled_dither: bool,

    /// Whether GL_POLYGON_OFFSET_FILL is enabled
    pub enabled_polygon_offset_fill: bool,

    /// Whether GL_SAMPLE_ALPHA_TO_COVERAGE is enabled
    pub enabled_sample_alpha_to_coverage: bool,

    /// Whether GL_SAMPLE_COVERAGE is enabled
    pub enabled_sample_coverage: bool,

    /// Whether GL_SCISSOR_TEST is enabled
    pub enabled_scissor_test: bool,

    /// Whether GL_STENCIL_TEST is enabled
    pub enabled_stencil_test: bool,

    // The latest value passed to `glUseProgram`.
    pub program: gl::types::GLuint,

    // The latest value passed to `glClearColor`.
    pub clear_color: (gl::types::GLclampf, gl::types::GLclampf,
                      gl::types::GLclampf, gl::types::GLclampf),

    // The latest value passed to `glClearDepthf`.
    pub clear_depth: gl::types::GLclampf,

    // The latest value passed to `glClearStencil`.
    pub clear_stencil: gl::types::GLint,

    /// The latest buffer bound to `GL_ARRAY_BUFFER`.
    pub array_buffer_binding: Option<gl::types::GLuint>,

    /// The latest buffer bound to `GL_ELEMENT_ARRAY_BUFFER`.
    pub element_array_buffer_binding: Option<gl::types::GLuint>,

    /// The latest framebuffer bound to `GL_READ_FRAMEBUFFER`.
    pub read_framebuffer: Option<gl::types::GLuint>,

    /// The latest framebuffer bound to `GL_DRAW_FRAMEBUFFER`.
    pub draw_framebuffer: Option<gl::types::GLuint>,

    /// The latest values passed to `glBlendFunc`.
    pub blend_func: (gl::types::GLenum, gl::types::GLenum),

    /// The latest value passed to `glDepthFunc`.
    pub depth_func: gl::types::GLenum,

    /// The latest values passed to `glViewport`.
    pub viewport: (gl::types::GLint, gl::types::GLint, gl::types::GLsizei, gl::types::GLsizei),

    /// The latest value passed to `glLineWidth`.
    pub line_width: gl::types::GLfloat,

    /// The latest value passed to `glCullFace`.
    pub cull_face: gl::types::GLenum,
}

impl GLState {
    /// Builds the `GLState` corresponding to the default GL values.
    fn new_defaults(viewport: (gl::types::GLint, gl::types::GLint, gl::types::GLsizei, gl::types::GLsizei)) -> GLState {
        GLState {
            enabled_blend: false,
            enabled_cull_face: false,
            enabled_depth_test: false,
            enabled_dither: false,
            enabled_polygon_offset_fill: false,
            enabled_sample_alpha_to_coverage: false,
            enabled_sample_coverage: false,
            enabled_scissor_test: false,
            enabled_stencil_test: false,

            program: 0,
            clear_color: (0.0, 0.0, 0.0, 0.0),
            clear_depth: 1.0,
            clear_stencil: 0,
            array_buffer_binding: None,
            element_array_buffer_binding: None,
            read_framebuffer: None,
            draw_framebuffer: None,
            depth_func: gl::LESS,
            blend_func: (0, 0),     // no default specified
            viewport: viewport,
            line_width: 1.0,
            cull_face: gl::BACK,
        }
    }
}

impl Context {
    pub fn new_from_window(window: glutin::Window) -> Context {
        let (tx_events, rx_events) = channel();
        let (tx_commands, rx_commands) = channel();

        let dimensions = Arc::new((AtomicUint::new(800), AtomicUint::new(600)));

        let context = Context {
            commands: Mutex::new(tx_commands),
            events: Mutex::new(rx_events),
            dimensions: dimensions.clone(),
        };

        TaskBuilder::new().native().spawn(proc() {
            unsafe { window.make_current(); }

            let gl = gl::Gl::load_with(|symbol| window.get_proc_address(symbol));

            // building the GLState and modifying to GL state to match it
            let mut gl_state = {
                let viewport = {
                    let dim = window.get_inner_size().unwrap();
                    dimensions.0.store(dim.0, Relaxed);
                    dimensions.1.store(dim.1, Relaxed);
                    (0, 0, dim.0 as gl::types::GLsizei, dim.1 as gl::types::GLsizei)
                };

                GLState::new_defaults(viewport)
            };

            // getting the GL version and extensions
            let version = get_gl_version(&gl);
            let extensions = get_extensions(&gl);

            // main loop
            'main: loop {
                // processing commands
                loop {
                    match rx_commands.recv_opt() {
                        Ok(EndFrame) => break,
                        Ok(Execute(cmd)) => cmd(&gl, &mut gl_state, version, extensions.as_slice()),
                        Err(_) => break 'main
                    }
                }

                // this is necessary on Windows 8, or nothing is being displayed
                gl.Flush();

                // swapping
                window.swap_buffers();

                // getting events
                for event in window.poll_events() {
                    // calling `glViewport` if the window has been resized
                    if let &glutin::Resized(width, height) = &event {
                        dimensions.0.store(width, Relaxed);
                        dimensions.1.store(height, Relaxed);

                        if gl_state.viewport != (0, 0, width as gl::types::GLsizei,
                                                 height as gl::types::GLsizei)
                        {
                            gl.Viewport(0, 0, width as gl::types::GLsizei,
                                height as gl::types::GLsizei);
                            gl_state.viewport = (0, 0, width as gl::types::GLsizei,
                                height as gl::types::GLsizei);
                        }
                    }

                    // sending the event outside
                    if tx_events.send_opt(event.clone()).is_err() {
                        break 'main;
                    }
                }
            }
        });

        context
    }

    pub fn new_from_headless(window: glutin::HeadlessContext) -> Context {
        let (_, rx_events) = channel();
        let (tx_commands, rx_commands) = channel();

        // TODO: fixme
        let dimensions = Arc::new((AtomicUint::new(800), AtomicUint::new(600)));

        let context = Context {
            commands: Mutex::new(tx_commands),
            events: Mutex::new(rx_events),
            dimensions: dimensions,
        };

        TaskBuilder::new().native().spawn(proc() {
            unsafe { window.make_current(); }

            let gl = gl::Gl::load_with(|symbol| window.get_proc_address(symbol));

            // building the GLState, version, and extensions
            let mut gl_state = GLState::new_defaults((0, 0, 0, 0));    // FIXME:
            let version = get_gl_version(&gl);
            let extensions = get_extensions(&gl);

            loop {
                match rx_commands.recv_opt() {
                    Ok(Execute(cmd)) => cmd(&gl, &mut gl_state, version, extensions.as_slice()),
                    Ok(EndFrame) => (),     // ignoring buffer swapping
                    Err(_) => break
                }
            }
        });

        context
    }

    pub fn get_framebuffer_dimensions(&self) -> (uint, uint) {
        (
            self.dimensions.0.load(Relaxed),
            self.dimensions.1.load(Relaxed),
        )
    }

    pub fn exec(&self, f: proc(&gl::Gl, &mut GLState, (u8, u8), &[String]): Send) {
        self.commands.lock().send(Execute(f));
    }

    pub fn swap_buffers(&self) {
        self.commands.lock().send(EndFrame);
    }

    pub fn recv(&self) -> Vec<glutin::Event> {
        let events = self.events.lock();

        let mut result = Vec::new();
        loop {
            match events.try_recv() {
                Ok(ev) => result.push(ev),
                Err(_) => break
            }
        }
        result
    }
}

fn get_gl_version(gl: &gl::Gl) -> (u8, u8) {
    use std::c_str::CString;

    unsafe {
        let version = gl.GetString(gl::VERSION);
        let version = CString::new(version as *const i8, false);
        let version = version.as_str().expect("OpenGL version contains non-utf8 characters");

        let version = version.words().next().expect("glGetString(GL_VERSION) returned an empty \
                                                     string");

        let mut iter = version.split(|c: char| c == '.');
        let major = iter.next().unwrap();
        let minor = iter.next().expect("glGetString(GL_VERSION) did not return a correct version");

        (
            from_str(major).expect("failed to parse GL major version"),
            from_str(minor).expect("failed to parse GL minor version")
        )
    }
}

fn get_extensions(gl: &gl::Gl) -> Vec<String> {
    use std::c_str::CString;

    unsafe {
        let list = gl.GetString(gl::EXTENSIONS);

        if list.is_null() {
            let mut num_extensions = 0;
            gl.GetIntegerv(gl::NUM_EXTENSIONS, &mut num_extensions);

            range(0, num_extensions).map(|num| {
                let ext = gl.GetStringi(gl::EXTENSIONS, num as gl::types::GLuint);
                let ext = CString::new(ext as *const i8, false);
                ext.as_str().expect("OpenGL extension contains non-utf8 characters").to_string()
            }).collect()

        } else {
            let list = CString::new(list as *const i8, false);
            let list = list.as_str()
                .expect("List of OpenGL extensions contains non-utf8 characters");
            list.words().map(|e| e.to_string()).collect()
        }
    }
}
<commit_msg>Now passing an `ExtensionsList` object instead of a `&[String]`<commit_after>use gl;
use glutin;
use native::NativeTaskBuilder;
use std::mem;
use std::sync::atomic::{AtomicUint, Relaxed};
use std::sync::{Arc, Mutex};
use std::task::TaskBuilder;

enum Message {
    EndFrame,
    Execute(proc(&gl::Gl, &mut GLState, (u8, u8), &ExtensionsList):Send),
}

pub struct Context {
    commands: Mutex<Sender<Message>>,
    events: Mutex<Receiver<glutin::Event>>,

    /// Dimensions of the frame buffer.
    dimensions: Arc<(AtomicUint, AtomicUint)>,
}

/// Represents the current OpenGL state.
/// The current state is passed to each function and can be freely updated.
pub struct GLState {
    /// Whether GL_BLEND is enabled
    pub enabled_blend: bool,

    /// Whether GL_CULL_FACE is enabled
    pub enabled_cull_face: bool,

    /// Whether GL_DEPTH_TEST is enabled
    pub enabled_depth_test: bool,

    /// Whether GL_DITHER is enabled
    pub enabled_dither: bool,

    /// Whether GL_POLYGON_OFFSET_FILL is enabled
    pub enabled_polygon_offset_fill: bool,

    /// Whether GL_SAMPLE_ALPHA_TO_COVERAGE is enabled
    pub enabled_sample_alpha_to_coverage: bool,

    /// Whether GL_SAMPLE_COVERAGE is enabled
    pub enabled_sample_coverage: bool,

    /// Whether GL_SCISSOR_TEST is enabled
    pub enabled_scissor_test: bool,

    /// Whether GL_STENCIL_TEST is enabled
    pub enabled_stencil_test: bool,

    // The latest value passed to `glUseProgram`.
    pub program: gl::types::GLuint,

    // The latest value passed to `glClearColor`.
    pub clear_color: (gl::types::GLclampf, gl::types::GLclampf,
                      gl::types::GLclampf, gl::types::GLclampf),

    // The latest value passed to `glClearDepthf`.
    pub clear_depth: gl::types::GLclampf,

    // The latest value passed to `glClearStencil`.
    pub clear_stencil: gl::types::GLint,

    /// The latest buffer bound to `GL_ARRAY_BUFFER`.
    pub array_buffer_binding: Option<gl::types::GLuint>,

    /// The latest buffer bound to `GL_ELEMENT_ARRAY_BUFFER`.
    pub element_array_buffer_binding: Option<gl::types::GLuint>,

    /// The latest framebuffer bound to `GL_READ_FRAMEBUFFER`.
    pub read_framebuffer: Option<gl::types::GLuint>,

    /// The latest framebuffer bound to `GL_DRAW_FRAMEBUFFER`.
    pub draw_framebuffer: Option<gl::types::GLuint>,

    /// The latest values passed to `glBlendFunc`.
    pub blend_func: (gl::types::GLenum, gl::types::GLenum),

    /// The latest value passed to `glDepthFunc`.
    pub depth_func: gl::types::GLenum,

    /// The latest values passed to `glViewport`.
    pub viewport: (gl::types::GLint, gl::types::GLint, gl::types::GLsizei, gl::types::GLsizei),

    /// The latest value passed to `glLineWidth`.
    pub line_width: gl::types::GLfloat,

    /// The latest value passed to `glCullFace`.
    pub cull_face: gl::types::GLenum,
}

impl GLState {
    /// Builds the `GLState` corresponding to the default GL values.
    fn new_defaults(viewport: (gl::types::GLint, gl::types::GLint, gl::types::GLsizei, gl::types::GLsizei)) -> GLState {
        GLState {
            enabled_blend: false,
            enabled_cull_face: false,
            enabled_depth_test: false,
            enabled_dither: false,
            enabled_polygon_offset_fill: false,
            enabled_sample_alpha_to_coverage: false,
            enabled_sample_coverage: false,
            enabled_scissor_test: false,
            enabled_stencil_test: false,

            program: 0,
            clear_color: (0.0, 0.0, 0.0, 0.0),
            clear_depth: 1.0,
            clear_stencil: 0,
            array_buffer_binding: None,
            element_array_buffer_binding: None,
            read_framebuffer: None,
            draw_framebuffer: None,
            depth_func: gl::LESS,
            blend_func: (0, 0),     // no default specified
            viewport: viewport,
            line_width: 1.0,
            cull_face: gl::BACK,
        }
    }
}

/// Contains data about the list of extensions
pub struct ExtensionsList {
    /// GL_EXT_direct_state_access
    pub gl_ext_direct_state_access: bool,
    /// GL_EXT_framebuffer_object
    pub gl_ext_framebuffer_object: bool,
}

impl Context {
    pub fn new_from_window(window: glutin::Window) -> Context {
        let (tx_events, rx_events) = channel();
        let (tx_commands, rx_commands) = channel();

        let dimensions = Arc::new((AtomicUint::new(800), AtomicUint::new(600)));

        let context = Context {
            commands: Mutex::new(tx_commands),
            events: Mutex::new(rx_events),
            dimensions: dimensions.clone(),
        };

        TaskBuilder::new().native().spawn(proc() {
            unsafe { window.make_current(); }

            let gl = gl::Gl::load_with(|symbol| window.get_proc_address(symbol));

            // building the GLState and modifying to GL state to match it
            let mut gl_state = {
                let viewport = {
                    let dim = window.get_inner_size().unwrap();
                    dimensions.0.store(dim.0, Relaxed);
                    dimensions.1.store(dim.1, Relaxed);
                    (0, 0, dim.0 as gl::types::GLsizei, dim.1 as gl::types::GLsizei)
                };

                GLState::new_defaults(viewport)
            };

            // getting the GL version and extensions
            let version = get_gl_version(&gl);
            let extensions = get_extensions(&gl);

            // main loop
            'main: loop {
                // processing commands
                loop {
                    match rx_commands.recv_opt() {
                        Ok(EndFrame) => break,
                        Ok(Execute(cmd)) => cmd(&gl, &mut gl_state, version, &extensions),
                        Err(_) => break 'main
                    }
                }

                // this is necessary on Windows 8, or nothing is being displayed
                gl.Flush();

                // swapping
                window.swap_buffers();

                // getting events
                for event in window.poll_events() {
                    // calling `glViewport` if the window has been resized
                    if let &glutin::Resized(width, height) = &event {
                        dimensions.0.store(width, Relaxed);
                        dimensions.1.store(height, Relaxed);

                        if gl_state.viewport != (0, 0, width as gl::types::GLsizei,
                                                 height as gl::types::GLsizei)
                        {
                            gl.Viewport(0, 0, width as gl::types::GLsizei,
                                height as gl::types::GLsizei);
                            gl_state.viewport = (0, 0, width as gl::types::GLsizei,
                                height as gl::types::GLsizei);
                        }
                    }

                    // sending the event outside
                    if tx_events.send_opt(event.clone()).is_err() {
                        break 'main;
                    }
                }
            }
        });

        context
    }

    pub fn new_from_headless(window: glutin::HeadlessContext) -> Context {
        let (_, rx_events) = channel();
        let (tx_commands, rx_commands) = channel();

        // TODO: fixme
        let dimensions = Arc::new((AtomicUint::new(800), AtomicUint::new(600)));

        let context = Context {
            commands: Mutex::new(tx_commands),
            events: Mutex::new(rx_events),
            dimensions: dimensions,
        };

        TaskBuilder::new().native().spawn(proc() {
            unsafe { window.make_current(); }

            let gl = gl::Gl::load_with(|symbol| window.get_proc_address(symbol));

            // building the GLState, version, and extensions
            let mut gl_state = GLState::new_defaults((0, 0, 0, 0));    // FIXME:
            let version = get_gl_version(&gl);
            let extensions = get_extensions(&gl);

            loop {
                match rx_commands.recv_opt() {
                    Ok(Execute(cmd)) => cmd(&gl, &mut gl_state, version, &extensions),
                    Ok(EndFrame) => (),     // ignoring buffer swapping
                    Err(_) => break
                }
            }
        });

        context
    }

    pub fn get_framebuffer_dimensions(&self) -> (uint, uint) {
        (
            self.dimensions.0.load(Relaxed),
            self.dimensions.1.load(Relaxed),
        )
    }

    pub fn exec(&self, f: proc(&gl::Gl, &mut GLState, (u8, u8), &ExtensionsList): Send) {
        self.commands.lock().send(Execute(f));
    }

    pub fn swap_buffers(&self) {
        self.commands.lock().send(EndFrame);
    }

    pub fn recv(&self) -> Vec<glutin::Event> {
        let events = self.events.lock();

        let mut result = Vec::new();
        loop {
            match events.try_recv() {
                Ok(ev) => result.push(ev),
                Err(_) => break
            }
        }
        result
    }
}

fn get_gl_version(gl: &gl::Gl) -> (u8, u8) {
    use std::c_str::CString;

    unsafe {
        let version = gl.GetString(gl::VERSION);
        let version = CString::new(version as *const i8, false);
        let version = version.as_str().expect("OpenGL version contains non-utf8 characters");

        let version = version.words().next().expect("glGetString(GL_VERSION) returned an empty \
                                                     string");

        let mut iter = version.split(|c: char| c == '.');
        let major = iter.next().unwrap();
        let minor = iter.next().expect("glGetString(GL_VERSION) did not return a correct version");

        (
            from_str(major).expect("failed to parse GL major version"),
            from_str(minor).expect("failed to parse GL minor version")
        )
    }
}

fn get_extensions_strings(gl: &gl::Gl) -> Vec<String> {
    use std::c_str::CString;

    unsafe {
        let list = gl.GetString(gl::EXTENSIONS);

        if list.is_null() {
            let mut num_extensions = 0;
            gl.GetIntegerv(gl::NUM_EXTENSIONS, &mut num_extensions);

            range(0, num_extensions).map(|num| {
                let ext = gl.GetStringi(gl::EXTENSIONS, num as gl::types::GLuint);
                let ext = CString::new(ext as *const i8, false);
                ext.as_str().expect("OpenGL extension contains non-utf8 characters").to_string()
            }).collect()

        } else {
            let list = CString::new(list as *const i8, false);
            let list = list.as_str()
                .expect("List of OpenGL extensions contains non-utf8 characters");
            list.words().map(|e| e.to_string()).collect()
        }
    }
}

fn get_extensions(gl: &gl::Gl) -> ExtensionsList {
    let strings = get_extensions_strings(gl);

    let mut extensions = ExtensionsList {
        gl_ext_direct_state_access: false,
        gl_ext_framebuffer_object: false,
    };

    for extension in strings.into_iter() {
        match extension.as_slice() {
            "GL_EXT_direct_state_access" => extensions.gl_ext_direct_state_access = true,
            "GL_EXT_framebuffer_object" => extensions.gl_ext_framebuffer_object = true,
            _ => ()
        }
    }

    extensions
}
<|endoftext|>


<commit_before>extern crate libc;

use std::io::Read;
use std::io::Result;
use std::ptr;
use super::liblz4::*;

use self::libc::size_t;

const BUFFER_SIZE: usize = 32 * 1024;

struct DecoderContext {
	c: LZ4FDecompressionContext,
}

pub struct Decoder<R> {
	c: DecoderContext,
	r: R,
	buf: [u8; BUFFER_SIZE],
	pos: usize,
	len: usize,
	eof: bool,
}

impl<R: Read> Decoder<R> {
	/// Creates a new encoder which will have its output written to the given
	/// output stream. The output stream can be re-acquired by calling
	/// `finish()`
	pub fn new(r: R) -> Result<Decoder<R>> {
		Ok (Decoder {
			r: r,
			c: try! (DecoderContext::new()),
			buf: [0; BUFFER_SIZE],
			pos: BUFFER_SIZE,
			len: BUFFER_SIZE,
			eof: false,
		})
	}
}

impl<R: Read> Read for Decoder<R> {
	fn read(&mut self, buf: &mut [u8]) -> Result<usize>
	{
		if self.eof
		{
			return Ok(0);
		}
		let mut dst_offset: usize = 0;
		while dst_offset < buf.len()
		{
			if self.pos >= self.len
			{
				self.pos = 0;
				self.len = try! (self.r.read(&mut self.buf));
				if self.len <= 0
				{
					break;
				}
			}
			while (dst_offset < buf.len()) && (self.pos < self.len)
			{
				let mut src_size = (self.len - self.pos) as size_t;
				let mut dst_size = (buf.len() - dst_offset) as size_t;
				let len = try! (check_error(unsafe {LZ4F_decompress(self.c.c, buf[dst_offset..].as_mut_ptr(), &mut dst_size, self.buf[self.pos..].as_ptr(), &mut src_size, ptr::null())}));
				self.pos += src_size as usize;
				dst_offset += dst_size as usize;
				if len == 0 {
					self.eof = true;
					break;
				}
			}
		}
		Ok(dst_offset)
	}
}

impl DecoderContext {
	fn new() -> Result<DecoderContext>
	{
		let mut context: LZ4FDecompressionContext = ptr::null_mut();
		try! (check_error(unsafe {
			LZ4F_createDecompressionContext(&mut context, LZ4F_VERSION)
		}));
		Ok(DecoderContext {
			c: context
		})
	}
}

impl Drop for DecoderContext {
	fn drop(&mut self) {
		unsafe
		{
			LZ4F_freeDecompressionContext(self.c)
		};
	}
}

#[cfg(test)]
mod test {
	use std::io::{Cursor, Read, Write};
	use super::super::encoder::EncoderBuilder;
	use super::Decoder;

	const BUFFER_SIZE: usize = 64 * 1024;

	#[test]
	fn test_decoder_smoke() {
		let mut encoder = EncoderBuilder::new().level(1).build(Vec::new()).unwrap();
		let expected = b"Some data";
		encoder.write(&expected[..4]).unwrap();
		encoder.write(&expected[4..]).unwrap();
		let (buffer, result) = encoder.finish();
		result.unwrap();

		let mut decoder = Decoder::new(Cursor::new(buffer)).unwrap();
		let mut actual = [0; BUFFER_SIZE];

		let size = decoder.read(&mut actual).unwrap();
		assert_eq!(expected, &actual[0..size]);
	}

	#[test]
	fn test_decoder_random() {
		let mut encoder = EncoderBuilder::new().level(1).build(Vec::new()).unwrap();
		let mut expected = Vec::new();
		let mut rnd: u32 = 42;
		for _ in 0..1027 * 1023 * 7 {
			expected.push((rnd & 0xFF) as u8);
			rnd = ((1664525 as u64) * (rnd as u64) + (1013904223 as u64)) as u32;
		}
		encoder.write(&expected).unwrap();
		let (encoded, result) = encoder.finish();
		result.unwrap();

		let mut decoder = Decoder::new(Cursor::new(encoded)).unwrap();
		let mut actual = Vec::new();
		loop {
			let mut buffer = [0; BUFFER_SIZE];
			let size = decoder.read(&mut buffer).unwrap();
			if size == 0 {
				break;
			}
			actual.write(&buffer[0..size]).unwrap();
		}
		assert_eq!(expected, actual);
	}
}
<commit_msg>Fix incorrect break that could cause reading after a frame ends<commit_after>extern crate libc;

use std::io::Read;
use std::io::Result;
use std::ptr;
use super::liblz4::*;

use self::libc::size_t;

const BUFFER_SIZE: usize = 32 * 1024;

struct DecoderContext {
	c: LZ4FDecompressionContext,
}

pub struct Decoder<R> {
	c: DecoderContext,
	r: R,
	buf: [u8; BUFFER_SIZE],
	pos: usize,
	len: usize,
	eof: bool,
}

impl<R: Read> Decoder<R> {
	/// Creates a new encoder which will have its output written to the given
	/// output stream. The output stream can be re-acquired by calling
	/// `finish()`
	pub fn new(r: R) -> Result<Decoder<R>> {
		Ok (Decoder {
			r: r,
			c: try! (DecoderContext::new()),
			buf: [0; BUFFER_SIZE],
			pos: BUFFER_SIZE,
			len: BUFFER_SIZE,
			eof: false,
		})
	}
}

impl<R: Read> Read for Decoder<R> {
	fn read(&mut self, buf: &mut [u8]) -> Result<usize>
	{
		if self.eof
		{
			return Ok(0);
		}
		let mut dst_offset: usize = 0;
		while dst_offset < buf.len()
		{
			if self.pos >= self.len
			{
				self.pos = 0;
				self.len = try! (self.r.read(&mut self.buf));
				if self.len <= 0
				{
					break;
				}
			}
			while (dst_offset < buf.len()) && (self.pos < self.len)
			{
				let mut src_size = (self.len - self.pos) as size_t;
				let mut dst_size = (buf.len() - dst_offset) as size_t;
				let len = try! (check_error(unsafe {LZ4F_decompress(self.c.c, buf[dst_offset..].as_mut_ptr(), &mut dst_size, self.buf[self.pos..].as_ptr(), &mut src_size, ptr::null())}));
				self.pos += src_size as usize;
				dst_offset += dst_size as usize;
				if len == 0 {
					self.eof = true;
					return Ok(dst_offset);
				}
			}
		}
		Ok(dst_offset)
	}
}

impl DecoderContext {
	fn new() -> Result<DecoderContext>
	{
		let mut context: LZ4FDecompressionContext = ptr::null_mut();
		try! (check_error(unsafe {
			LZ4F_createDecompressionContext(&mut context, LZ4F_VERSION)
		}));
		Ok(DecoderContext {
			c: context
		})
	}
}

impl Drop for DecoderContext {
	fn drop(&mut self) {
		unsafe
		{
			LZ4F_freeDecompressionContext(self.c)
		};
	}
}

#[cfg(test)]
mod test {
	use std::io::{Cursor, Read, Write};
	use super::super::encoder::EncoderBuilder;
	use super::Decoder;

	const BUFFER_SIZE: usize = 64 * 1024;

	#[test]
	fn test_decoder_smoke() {
		let mut encoder = EncoderBuilder::new().level(1).build(Vec::new()).unwrap();
		let expected = b"Some data";
		encoder.write(&expected[..4]).unwrap();
		encoder.write(&expected[4..]).unwrap();
		let (buffer, result) = encoder.finish();
		result.unwrap();

		let mut decoder = Decoder::new(Cursor::new(buffer)).unwrap();
		let mut actual = [0; BUFFER_SIZE];

		let size = decoder.read(&mut actual).unwrap();
		assert_eq!(expected, &actual[0..size]);
	}

	#[test]
	fn test_decoder_random() {
		let mut encoder = EncoderBuilder::new().level(1).build(Vec::new()).unwrap();
		let mut expected = Vec::new();
		let mut rnd: u32 = 42;
		for _ in 0..1027 * 1023 * 7 {
			expected.push((rnd & 0xFF) as u8);
			rnd = ((1664525 as u64) * (rnd as u64) + (1013904223 as u64)) as u32;
		}
		encoder.write(&expected).unwrap();
		let (encoded, result) = encoder.finish();
		result.unwrap();

		let mut decoder = Decoder::new(Cursor::new(encoded)).unwrap();
		let mut actual = Vec::new();
		loop {
			let mut buffer = [0; BUFFER_SIZE];
			let size = decoder.read(&mut buffer).unwrap();
			if size == 0 {
				break;
			}
			actual.write(&buffer[0..size]).unwrap();
		}
		assert_eq!(expected, actual);
	}
}
<|endoftext|>


<commit_before>use interpreter::*;

pub enum Program {
    Integer(u8),
    Binary {op: u8, left: Box<Program>, right: Box<Program>},
    Print(Box<Program>)
}

pub fn emit(program: &Program) -> Vec<u8> {
    let mut code = vec![];
    emit_into(&mut code, program, 0);
    code.push(OP_RETURN);
    code.push(0);
    code
}

fn emit_into(code: &mut Vec<u8>, program: &Program, target: u8) {
    match *program {
        Program::Integer(n) => {
            load_u8(code, n, target)
        },
        Program::Binary {op, ref left, ref right} => {
            emit_into(code, &*left, target);
            emit_into(code, &*right, target + 1);
            code.push(op);
            code.push(target);
            code.push(target + 1);
            code.push(target);
        },
        Program::Print(ref program) => {
            emit_into(code, &*program, target);
            code.push(OP_PRINT);
        }
    }
}

fn load_u8(code: &mut Vec<u8>, value: u8, dest: u8) {
    code.push(OP_LOAD_U8);
    code.push(value);
    code.push(dest);
}

fn int(n: u8) -> Box<Program> {
    Box::new(Program::Integer(n))
}

fn add(left: Box<Program>, right: Box<Program>) -> Box<Program> {
    let program = Program::Binary {
        op: OP_ADD,
        left: left,
        right: right
    };
    Box::new(program)
}

#[test]
fn test_emit_integer() {
    let code = emit(&Program::Integer(37));
    assert_eq!(37, interpret(&code));

    let code1 = emit(&int(56));
    assert_eq!(56, interpret(&code1));

    assert_eq!(73, interpret(&emit(&int(73))));
}

#[test]
fn test_emit_simple_binary_op() {
    let program = Program::Binary {
        op: OP_ADD,
        left: int(73),
        right: int(68)
    };

    let code = emit(&program);
    assert_eq!(141, interpret(&code));
}

#[test]
fn test_emit_nested_binary_op() {
    // (73 + 68) + 50 = 191
    let program0 = Program::Binary {
        op: OP_ADD,
        left: add(int(73), int(68)),
        right: int(50)
    };

    let code0 = emit(&program0);
    assert_eq!(191, interpret(&code0));

    // 1 + (2 + 3) = 6
    let program1 = Program::Binary {
        op: OP_ADD,
        left: int(1),
        right: add(int(2), int(3))
    };

    let code1 = emit(&program1);
    assert_eq!(6, interpret(&code1));
}

#[test]
fn test_emit_deep_nested_binary_op() {
    // (1 + (2 + (3 + (4 + 5))))
    let program = add(int(1), add(int(2), add(int(3), add(int(4), int(5)))));
    let code = emit(&program);
    assert_eq!(15, interpret(&code));
}

#[test]
fn test_emit_print() {
    assert_eq!(7, interpret(&emit(&Program::Print(add(int(3), int(4))))));
    assert_eq!(7, interpret(&emit(&Program::Integer(7))));
}
<commit_msg>actually test the print functionality<commit_after>use interpreter::*;

pub enum Program {
    Integer(u8),
    Binary {op: u8, left: Box<Program>, right: Box<Program>},
    Print(Box<Program>)
}

pub fn emit(program: &Program) -> Vec<u8> {
    let mut code = vec![];
    emit_into(&mut code, program, 0);
    code.push(OP_RETURN);
    code.push(0);
    code
}

fn emit_into(code: &mut Vec<u8>, program: &Program, target: u8) {
    match *program {
        Program::Integer(n) => {
            load_u8(code, n, target)
        },
        Program::Binary {op, ref left, ref right} => {
            emit_into(code, &*left, target);
            emit_into(code, &*right, target + 1);
            code.push(op);
            code.push(target);
            code.push(target + 1);
            code.push(target);
        },
        Program::Print(ref program) => {
            emit_into(code, &*program, target);
            code.push(OP_PRINT);
        }
    }
}

fn load_u8(code: &mut Vec<u8>, value: u8, dest: u8) {
    code.push(OP_LOAD_U8);
    code.push(value);
    code.push(dest);
}

fn int(n: u8) -> Box<Program> {
    Box::new(Program::Integer(n))
}

fn add(left: Box<Program>, right: Box<Program>) -> Box<Program> {
    let program = Program::Binary {
        op: OP_ADD,
        left: left,
        right: right
    };
    Box::new(program)
}

#[test]
fn test_emit_integer() {
    let code = emit(&Program::Integer(37));
    assert_eq!(37, interpret(&code));

    let code1 = emit(&int(56));
    assert_eq!(56, interpret(&code1));

    assert_eq!(73, interpret(&emit(&int(73))));
}

#[test]
fn test_emit_simple_binary_op() {
    let program = Program::Binary {
        op: OP_ADD,
        left: int(73),
        right: int(68)
    };

    let code = emit(&program);
    assert_eq!(141, interpret(&code));
}

#[test]
fn test_emit_nested_binary_op() {
    // (73 + 68) + 50 = 191
    let program0 = Program::Binary {
        op: OP_ADD,
        left: add(int(73), int(68)),
        right: int(50)
    };

    let code0 = emit(&program0);
    assert_eq!(191, interpret(&code0));

    // 1 + (2 + 3) = 6
    let program1 = Program::Binary {
        op: OP_ADD,
        left: int(1),
        right: add(int(2), int(3))
    };

    let code1 = emit(&program1);
    assert_eq!(6, interpret(&code1));
}

#[test]
fn test_emit_deep_nested_binary_op() {
    // (1 + (2 + (3 + (4 + 5))))
    let program = add(int(1), add(int(2), add(int(3), add(int(4), int(5)))));
    let code = emit(&program);
    assert_eq!(15, interpret(&code));
}

#[test]
fn test_emit_print() {
    assert_eq!(7, interpret(&emit(&Program::Print(add(int(3), int(4))))));
    assert_eq!(7, interpret(&emit(&Program::Print(int(7)))));
}
<|endoftext|>


<commit_before>use {HasLua, CopyRead, ConsumeRead, LoadedVariable, Push, Index};
use ffi;

#[unstable]
pub struct LuaTable<'var, L> {
    variable: LoadedVariable<'var, L>
}

impl<'var, 'lua, L: HasLua> HasLua for LuaTable<'var, L> {
    fn use_lua(&mut self) -> *mut ffi::lua_State {
        self.variable.use_lua()
    }
}

// while the LuaTableIterator is active, the current key is constantly pushed over the table
#[unstable]
pub struct LuaTableIterator<'var, 'table, L> {
    table: &'table mut LuaTable<'var, L>
}

impl<'var, 'lua, 'table, L: HasLua> HasLua for LuaTableIterator<'var, 'table, L> {
    fn use_lua(&mut self) -> *mut ffi::lua_State {
        self.table.use_lua()
    }
}

impl<'var, 'lua, L: HasLua> ConsumeRead<'var, L> for LuaTable<'var, L> {
    fn read_from_variable(mut var: LoadedVariable<'var, L>)
        -> Result<LuaTable<'var, L>, LoadedVariable<'var, L>>
    {
        if unsafe { ffi::lua_istable(var.use_lua(), -1) } {
            Ok(LuaTable{ variable: var })
        } else {
            Err(var)
        }
    }
}

impl<'var, L: HasLua> LuaTable<'var, L> {
    pub fn iter<'me>(&'me mut self)
        -> LuaTableIterator<'var, 'me, L>
    {
        unsafe { ffi::lua_pushnil(self.variable.use_lua()) };
        LuaTableIterator { table: self }
    }

    pub fn load<'a, R: ConsumeRead<'a, LuaTable<'var, L>>, I: Index<LuaTable<'var, L>>>(&'a mut self, index: I) -> Option<R> {
        index.push_to_lua(self);
        unsafe { ffi::lua_gettable(self.use_lua(), -2); }
        let var = LoadedVariable{lua: self, size: 1};
        ConsumeRead::read_from_variable(var).ok()
    }

    pub fn load_table<'a, I: Index<LuaTable<'var, L>>>(&'a mut self, index: I) -> Option<LuaTable<'a, LuaTable<'var, L>>> {
        self.load(index)
    }

    pub fn get<R: CopyRead<LuaTable<'var, L>>, I: Index<LuaTable<'var, L>>>(&mut self, index: I) -> Option<R> {
        index.push_to_lua(self);
        unsafe { ffi::lua_gettable(self.use_lua(), -2); }
        let value = CopyRead::read_from_lua(self, -1);
        unsafe { ffi::lua_pop(self.use_lua(), 1); }
        value
    }

    pub fn set<I: Index<LuaTable<'var, L>>, V: Push<LuaTable<'var, L>>>(&mut self, index: I, value: V) {
        index.push_to_lua(self);
        value.push_to_lua(self);
        unsafe { ffi::lua_settable(self.use_lua(), -3); }
    }

    // Obtains or create the metatable of the table
    pub fn get_or_create_metatable(mut self) -> LuaTable<'var, L> {
        let result = unsafe { ffi::lua_getmetatable(self.variable.use_lua(), -1) };

        if result == 0 {
            unsafe {
                ffi::lua_newtable(self.variable.use_lua());
                ffi::lua_setmetatable(self.variable.use_lua(), -2);
                let r = ffi::lua_getmetatable(self.variable.use_lua(), -1);
                assert!(r != 0);
            }
        }

        // note: it would be cleaner to create another table, but cannot manage to make it compile
        self.variable.size += 1;
        self
    }
}

impl<'a, 'b, 'lua, L: HasLua, K: CopyRead<LuaTableIterator<'a, 'b, L>>, V: CopyRead<LuaTableIterator<'a, 'b, L>>>
    Iterator<Option<(K, V)>> for LuaTableIterator<'a, 'b, L>
{
    fn next(&mut self)
        -> Option<Option<(K,V)>>
    {
        // this call pushes the next key and value on the stack
        if unsafe { ffi::lua_next(self.table.use_lua(), -2) } == 0 {
            return None
        }

        let key = CopyRead::read_from_lua(self, -2);
        let value = CopyRead::read_from_lua(self, -1);

        // removing the value, leaving only the key on the top of the stack
        unsafe { ffi::lua_pop(self.table.use_lua(), 1) };

        //
        if key.is_none() || value.is_none() {
            Some(None)
        } else {
            Some(Some((key.unwrap(), value.unwrap())))
        }
    }
}

// TODO: this destructor crashes the compiler
/*impl<'a, 'b> Drop for LuaTableIterator<'a, 'b> {
    fn drop(&mut self) {
        unsafe { ffi::lua_pop(self.table.variable.lua.lua, 1) }
    }
}*/
<commit_msg>Lifetimes cleanup in lua_tables<commit_after>use {HasLua, CopyRead, ConsumeRead, LoadedVariable, Push, Index};
use ffi;

#[unstable]
pub struct LuaTable<'var, L> {
    variable: LoadedVariable<'var, L>
}

impl<'var, L: HasLua> HasLua for LuaTable<'var, L> {
    fn use_lua(&mut self) -> *mut ffi::lua_State {
        self.variable.use_lua()
    }
}

// while the LuaTableIterator is active, the current key is constantly pushed over the table
#[unstable]
pub struct LuaTableIterator<'var, 'table, L> {
    table: &'table mut LuaTable<'var, L>
}

impl<'var, 'table, L: HasLua> HasLua for LuaTableIterator<'var, 'table, L> {
    fn use_lua(&mut self) -> *mut ffi::lua_State {
        self.table.use_lua()
    }
}

impl<'var, L: HasLua> ConsumeRead<'var, L> for LuaTable<'var, L> {
    fn read_from_variable(mut var: LoadedVariable<'var, L>)
        -> Result<LuaTable<'var, L>, LoadedVariable<'var, L>>
    {
        if unsafe { ffi::lua_istable(var.use_lua(), -1) } {
            Ok(LuaTable{ variable: var })
        } else {
            Err(var)
        }
    }
}

impl<'var, L: HasLua> LuaTable<'var, L> {
    pub fn iter<'me>(&'me mut self)
        -> LuaTableIterator<'var, 'me, L>
    {
        unsafe { ffi::lua_pushnil(self.variable.use_lua()) };
        LuaTableIterator { table: self }
    }

    pub fn load<'a, R: ConsumeRead<'a, LuaTable<'var, L>>, I: Index<LuaTable<'var, L>>>(&'a mut self, index: I) -> Option<R> {
        index.push_to_lua(self);
        unsafe { ffi::lua_gettable(self.use_lua(), -2); }
        let var = LoadedVariable{lua: self, size: 1};
        ConsumeRead::read_from_variable(var).ok()
    }

    pub fn load_table<'a, I: Index<LuaTable<'var, L>>>(&'a mut self, index: I) -> Option<LuaTable<'a, LuaTable<'var, L>>> {
        self.load(index)
    }

    pub fn get<R: CopyRead<LuaTable<'var, L>>, I: Index<LuaTable<'var, L>>>(&mut self, index: I) -> Option<R> {
        index.push_to_lua(self);
        unsafe { ffi::lua_gettable(self.use_lua(), -2); }
        let value = CopyRead::read_from_lua(self, -1);
        unsafe { ffi::lua_pop(self.use_lua(), 1); }
        value
    }

    pub fn set<I: Index<LuaTable<'var, L>>, V: Push<LuaTable<'var, L>>>(&mut self, index: I, value: V) {
        index.push_to_lua(self);
        value.push_to_lua(self);
        unsafe { ffi::lua_settable(self.use_lua(), -3); }
    }

    // Obtains or create the metatable of the table
    pub fn get_or_create_metatable(mut self) -> LuaTable<'var, L> {
        let result = unsafe { ffi::lua_getmetatable(self.variable.use_lua(), -1) };

        if result == 0 {
            unsafe {
                ffi::lua_newtable(self.variable.use_lua());
                ffi::lua_setmetatable(self.variable.use_lua(), -2);
                let r = ffi::lua_getmetatable(self.variable.use_lua(), -1);
                assert!(r != 0);
            }
        }

        // note: it would be cleaner to create another table, but cannot manage to make it compile
        self.variable.size += 1;
        self
    }
}

impl<'a, 'b, L: HasLua, K: CopyRead<LuaTableIterator<'a, 'b, L>>, V: CopyRead<LuaTableIterator<'a, 'b, L>>>
    Iterator<Option<(K, V)>> for LuaTableIterator<'a, 'b, L>
{
    fn next(&mut self)
        -> Option<Option<(K,V)>>
    {
        // this call pushes the next key and value on the stack
        if unsafe { ffi::lua_next(self.table.use_lua(), -2) } == 0 {
            return None
        }

        let key = CopyRead::read_from_lua(self, -2);
        let value = CopyRead::read_from_lua(self, -1);

        // removing the value, leaving only the key on the top of the stack
        unsafe { ffi::lua_pop(self.table.use_lua(), 1) };

        //
        if key.is_none() || value.is_none() {
            Some(None)
        } else {
            Some(Some((key.unwrap(), value.unwrap())))
        }
    }
}

// TODO: this destructor crashes the compiler
/*impl<'a, 'b> Drop for LuaTableIterator<'a, 'b> {
    fn drop(&mut self) {
        unsafe { ffi::lua_pop(self.table.variable.lua.lua, 1) }
    }
}*/
<|endoftext|>


<commit_before>use chrono::{DateTime, UTC};

pub fn format_for_send(metric: String, namespace: &String, tags: Vec<String>) -> String {
    let mut result = metric;

    if namespace != "" {
        result = format!("{}.{}", namespace, result)
    }

    let joined_tags = tags.join(",");
    if joined_tags != "" {
        result = format!("{}|{}", result, joined_tags)
    }

    result
}

pub trait Metric {
    fn metric_type_format(&self) -> String;
}

pub enum CountMetric {
    Incr(String),
    Decr(String),
}

impl Metric for CountMetric {
    // my_count:1|c
    // my_count:-1|c
    fn metric_type_format(&self) -> String {
        match *self {
            CountMetric::Incr(ref stat) => {
                format!("{}:1|c", stat)
            },
            CountMetric::Decr(ref stat) => {
                format!("{}:-1|c", stat)
            },
        }
    }
}

pub struct TimeMetric {
    start_time: DateTime<UTC>,
    end_time: DateTime<UTC>,
    stat: String,
}

impl Metric for TimeMetric {
    // my_stat:500|ms
    fn metric_type_format(&self) -> String {
        let dur = self.end_time - self.start_time;
        format!("{}:{}|ms", self.stat, dur.num_milliseconds())
    }
}

impl TimeMetric {
    pub fn new(stat: String, start_time: DateTime<UTC>, end_time: DateTime<UTC>) -> Self {
        TimeMetric {
            start_time: start_time,
            end_time: end_time,
            stat: stat,
        }
    }
}

pub struct TimingMetric {
    ms: i64,
    stat: String,
}

impl Metric for TimingMetric {
    // my_stat:500|ms
    fn metric_type_format(&self) -> String {
        format!("{}:{}|ms", self.stat, self.ms)
    }
}

impl TimingMetric {
    pub fn new(stat: String, ms: i64) -> Self {
        TimingMetric {
            ms: ms,
            stat: stat,
        }
    }
}

pub struct GaugeMetric {
    stat: String,
    val: String,
}

impl Metric for GaugeMetric {
    // my_gauge:1000|g
    fn metric_type_format(&self) -> String {
        format!("{}:{}|g", self.stat, self.val)
    }
}

impl GaugeMetric {
    pub fn new(stat: String, val: String) -> Self {
        GaugeMetric {
            stat: stat,
            val: val,
        }
    }
}

pub struct HistogramMetric {
    stat: String,
    val: String,
}

impl Metric for HistogramMetric {
    // my_histogram:1000|h
    fn metric_type_format(&self) -> String {
        format!("{}:{}|h", self.stat, self.val)
    }
}

impl HistogramMetric {
    pub fn new(stat: String, val: String) -> Self {
        HistogramMetric {
            stat: stat,
            val: val,
        }
    }
}

pub struct SetMetric {
    stat: String,
    val: String,
}

impl Metric for SetMetric {
    // my_set:45|s
    fn metric_type_format(&self) -> String {
        format!("{}:{}|s", self.stat, self.val)
    }
}

impl SetMetric {
    pub fn new(stat: String, val: String) -> Self {
        SetMetric {
            stat: stat,
            val: val,
        }
    }
}

pub struct Event {
    title: String,
    text: String,
}

impl Metric for Event {
    fn metric_type_format(&self) -> String {
        format!("_e{{{title_len},{text_len}}}:{title}|{text}",
                title_len = self.title.len(),
                text_len = self.text.len(),
                title = self.title,
                text = self.text)
    }
}

impl Event {
    pub fn new(title: String, text: String) -> Self {
        Event {
            title: title,
            text: text,
        }
    }
}

#[cfg(test)]
mod tests {
    use chrono::{TimeZone, UTC};
    use super::*;

    #[test]
    fn test_format_for_send_no_tags() {
        assert_eq!(
            "namespace.metric:val|v".to_string(),
            format_for_send("metric:val|v".to_string(), &"namespace".to_string(), vec![])
        )
    }

    #[test]
    fn test_format_for_send_no_namespace() {
        assert_eq!(
            "metric:val|v|tag:1,tag:2".to_string(),
            format_for_send("metric:val|v".to_string(), &String::new(), vec!["tag:1".into(), "tag:2".into()])
        )
    }

    #[test]
    fn test_format_for_send_everything() {
        assert_eq!(
            "namespace.metric:val|v|tag:1,tag:2".to_string(),
            format_for_send("metric:val|v".to_string(), &"namespace".to_string(), vec!["tag:1".into(), "tag:2".into()])
        )
    }

    #[test]
    fn test_count_incr_metric() {
        let metric = CountMetric::Incr("incr".into());

        assert_eq!("incr:1|c", metric.metric_type_format())
    }

    #[test]
    fn test_count_decr_metric() {
        let metric = CountMetric::Decr("decr".into());

        assert_eq!("decr:-1|c", metric.metric_type_format())
    }

    #[test]
    fn test_time_metric() {
        let start_time = UTC.ymd(2016, 4, 24).and_hms_milli(0, 0, 0, 0);
        let end_time = UTC.ymd(2016, 4, 24).and_hms_milli(0, 0, 0, 900);
        let metric = TimeMetric::new("time".into(), start_time, end_time);

        assert_eq!("time:900|ms", metric.metric_type_format())
    }

    #[test]
    fn test_timing_metric() {
        let metric = TimingMetric::new("timing".into(), 720);

        assert_eq!("timing:720|ms", metric.metric_type_format())
    }

    #[test]
    fn test_gauge_metric() {
        let metric = GaugeMetric::new("gauge".into(), "12345".into());

        assert_eq!("gauge:12345|g", metric.metric_type_format())
    }

    #[test]
    fn test_histogram_metric() {
        let metric = HistogramMetric::new("histogram".into(), "67890".into());

        assert_eq!("histogram:67890|h", metric.metric_type_format())
    }

    #[test]
    fn test_set_metric() {
        let metric = SetMetric::new("set".into(), "13579".into());

        assert_eq!("set:13579|s", metric.metric_type_format())
    }

    #[test]
    fn test_event() {
        let metric = Event::new("Event Title".into(), "Event Body - Something Happened".into());

        assert_eq!("_e{11,31}:Event Title|Event Body - Something Happened",
                   metric.metric_type_format())
    }
}
<commit_msg>Fix dogstatsD tag formatting<commit_after>use chrono::{DateTime, UTC};

pub fn format_for_send(metric: String, namespace: &String, tags: Vec<String>) -> String {
    let mut result = metric;

    if namespace != "" {
        result = format!("{}.{}", namespace, result)
    }

    let joined_tags = tags.join(",");
    if joined_tags != "" {
        result = format!("{}|#{}", result, joined_tags)
    }

    result
}

pub trait Metric {
    fn metric_type_format(&self) -> String;
}

pub enum CountMetric {
    Incr(String),
    Decr(String),
}

impl Metric for CountMetric {
    // my_count:1|c
    // my_count:-1|c
    fn metric_type_format(&self) -> String {
        match *self {
            CountMetric::Incr(ref stat) => {
                format!("{}:1|c", stat)
            },
            CountMetric::Decr(ref stat) => {
                format!("{}:-1|c", stat)
            },
        }
    }
}

pub struct TimeMetric {
    start_time: DateTime<UTC>,
    end_time: DateTime<UTC>,
    stat: String,
}

impl Metric for TimeMetric {
    // my_stat:500|ms
    fn metric_type_format(&self) -> String {
        let dur = self.end_time - self.start_time;
        format!("{}:{}|ms", self.stat, dur.num_milliseconds())
    }
}

impl TimeMetric {
    pub fn new(stat: String, start_time: DateTime<UTC>, end_time: DateTime<UTC>) -> Self {
        TimeMetric {
            start_time: start_time,
            end_time: end_time,
            stat: stat,
        }
    }
}

pub struct TimingMetric {
    ms: i64,
    stat: String,
}

impl Metric for TimingMetric {
    // my_stat:500|ms
    fn metric_type_format(&self) -> String {
        format!("{}:{}|ms", self.stat, self.ms)
    }
}

impl TimingMetric {
    pub fn new(stat: String, ms: i64) -> Self {
        TimingMetric {
            ms: ms,
            stat: stat,
        }
    }
}

pub struct GaugeMetric {
    stat: String,
    val: String,
}

impl Metric for GaugeMetric {
    // my_gauge:1000|g
    fn metric_type_format(&self) -> String {
        format!("{}:{}|g", self.stat, self.val)
    }
}

impl GaugeMetric {
    pub fn new(stat: String, val: String) -> Self {
        GaugeMetric {
            stat: stat,
            val: val,
        }
    }
}

pub struct HistogramMetric {
    stat: String,
    val: String,
}

impl Metric for HistogramMetric {
    // my_histogram:1000|h
    fn metric_type_format(&self) -> String {
        format!("{}:{}|h", self.stat, self.val)
    }
}

impl HistogramMetric {
    pub fn new(stat: String, val: String) -> Self {
        HistogramMetric {
            stat: stat,
            val: val,
        }
    }
}

pub struct SetMetric {
    stat: String,
    val: String,
}

impl Metric for SetMetric {
    // my_set:45|s
    fn metric_type_format(&self) -> String {
        format!("{}:{}|s", self.stat, self.val)
    }
}

impl SetMetric {
    pub fn new(stat: String, val: String) -> Self {
        SetMetric {
            stat: stat,
            val: val,
        }
    }
}

pub struct Event {
    title: String,
    text: String,
}

impl Metric for Event {
    fn metric_type_format(&self) -> String {
        format!("_e{{{title_len},{text_len}}}:{title}|{text}",
                title_len = self.title.len(),
                text_len = self.text.len(),
                title = self.title,
                text = self.text)
    }
}

impl Event {
    pub fn new(title: String, text: String) -> Self {
        Event {
            title: title,
            text: text,
        }
    }
}

#[cfg(test)]
mod tests {
    use chrono::{TimeZone, UTC};
    use super::*;

    #[test]
    fn test_format_for_send_no_tags() {
        assert_eq!(
            "namespace.metric:val|v".to_string(),
            format_for_send("metric:val|v".to_string(), &"namespace".to_string(), vec![])
        )
    }

    #[test]
    fn test_format_for_send_no_namespace() {
        assert_eq!(
            "metric:val|v|#tag:1,tag:2".to_string(),
            format_for_send("metric:val|v".to_string(), &String::new(), vec!["tag:1".into(), "tag:2".into()])
        )
    }

    #[test]
    fn test_format_for_send_everything() {
        assert_eq!(
            "namespace.metric:val|v|#tag:1,tag:2".to_string(),
            format_for_send("metric:val|v".to_string(), &"namespace".to_string(), vec!["tag:1".into(), "tag:2".into()])
        )
    }

    #[test]
    fn test_count_incr_metric() {
        let metric = CountMetric::Incr("incr".into());

        assert_eq!("incr:1|c", metric.metric_type_format())
    }

    #[test]
    fn test_count_decr_metric() {
        let metric = CountMetric::Decr("decr".into());

        assert_eq!("decr:-1|c", metric.metric_type_format())
    }

    #[test]
    fn test_time_metric() {
        let start_time = UTC.ymd(2016, 4, 24).and_hms_milli(0, 0, 0, 0);
        let end_time = UTC.ymd(2016, 4, 24).and_hms_milli(0, 0, 0, 900);
        let metric = TimeMetric::new("time".into(), start_time, end_time);

        assert_eq!("time:900|ms", metric.metric_type_format())
    }

    #[test]
    fn test_timing_metric() {
        let metric = TimingMetric::new("timing".into(), 720);

        assert_eq!("timing:720|ms", metric.metric_type_format())
    }

    #[test]
    fn test_gauge_metric() {
        let metric = GaugeMetric::new("gauge".into(), "12345".into());

        assert_eq!("gauge:12345|g", metric.metric_type_format())
    }

    #[test]
    fn test_histogram_metric() {
        let metric = HistogramMetric::new("histogram".into(), "67890".into());

        assert_eq!("histogram:67890|h", metric.metric_type_format())
    }

    #[test]
    fn test_set_metric() {
        let metric = SetMetric::new("set".into(), "13579".into());

        assert_eq!("set:13579|s", metric.metric_type_format())
    }

    #[test]
    fn test_event() {
        let metric = Event::new("Event Title".into(), "Event Body - Something Happened".into());

        assert_eq!("_e{11,31}:Event Title|Event Body - Something Happened",
                   metric.metric_type_format())
    }
}
<|endoftext|>