<commit_before>use {
    AddColor,
    AddPolygons,
    Borrowed,
    Color,
    Field,
    Matrix2d,
    TweenColorContext,
    TweenPolygonsContext,
    Value,
};
use internal::{
    CanTransform,
    CanViewTransform,
    HasTransform,
    HasViewTransform,
};

/// An animation inbetweening context.
pub struct TweenContext<'a> {
    /// Base/origin transform.
    pub base: Field<'a, Matrix2d>,
    /// Current transform.
    pub transform: Field<'a, Matrix2d>,
    /// Animation inbetweening factor.
    pub tween_factor: Field<'a, f64>,
}

impl<'a> Clone for TweenContext<'a> {
    #[inline(always)]
    fn clone(&self) -> TweenContext<'static> {
        TweenContext {
            base: self.base.clone(),
            transform: self.transform.clone(),
            tween_factor: self.tween_factor.clone(),
        }
    }
}

impl<'a> AddColor<'a, TweenColorContext<'a>> for TweenContext<'a> {
    #[inline(always)]
    fn rgba(&'a self, r: f32, g: f32, b: f32, a: f32) -> TweenColorContext<'a> {
        TweenColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            tween_factor: Borrowed(self.tween_factor.get()),
            color: Value(Color([r, g, b, a])),
        }
    }
}

impl<'a, 'b> AddPolygons<'a, TweenPolygonsContext<'a, 'b>> for TweenContext<'a> {
    #[inline(always)]
    fn polygons(&'a self, polygons: &'b [&'b [f64]]) -> TweenPolygonsContext<'a, 'b> {
        TweenPolygonsContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            tween_factor: Borrowed(self.tween_factor.get()),
            polygons: Value(polygons),
        }
    }
}

impl<'a> HasTransform<'a, Matrix2d> for TweenContext<'a> {
    #[inline(always)]
    fn get_transform(&'a self) -> &'a Matrix2d {
        self.transform.get()
    }
}

impl<'a> CanTransform<'a, TweenContext<'a>, Matrix2d> for TweenContext<'a> {
    #[inline(always)]
    fn transform(&'a self, value: Matrix2d) -> TweenContext<'a> {
        TweenContext {
            base: Borrowed(self.base.get()),
            transform: Value(value),
            tween_factor: Borrowed(self.tween_factor.get()),
        }
    }
}

impl<'a> HasViewTransform<'a, Matrix2d> for TweenContext<'a> {
    #[inline(always)]
    fn get_view_transform(&'a self) -> &'a Matrix2d {
        self.base.get()
    }
}

impl<'a> CanViewTransform<'a, TweenContext<'a>, Matrix2d> for TweenContext<'a> {
    #[inline(always)]
    fn view_transform(&'a self, value: Matrix2d) -> TweenContext<'a> {
        TweenContext {
            base: Value(value),
            transform: Borrowed(self.transform.get()),
            tween_factor: Borrowed(self.tween_factor.get()),
        }
    }
}

<commit_msg>Using type alias for `Polygons` in `TweenContext`<commit_after>use {
    AddColor,
    AddPolygons,
    Borrowed,
    Color,
    Field,
    Matrix2d,
    Polygons,
    TweenColorContext,
    TweenPolygonsContext,
    Value,
};
use internal::{
    CanTransform,
    CanViewTransform,
    HasTransform,
    HasViewTransform,
};

/// An animation inbetweening context.
pub struct TweenContext<'a> {
    /// Base/origin transform.
    pub base: Field<'a, Matrix2d>,
    /// Current transform.
    pub transform: Field<'a, Matrix2d>,
    /// Animation inbetweening factor.
    pub tween_factor: Field<'a, f64>,
}

impl<'a> Clone for TweenContext<'a> {
    #[inline(always)]
    fn clone(&self) -> TweenContext<'static> {
        TweenContext {
            base: self.base.clone(),
            transform: self.transform.clone(),
            tween_factor: self.tween_factor.clone(),
        }
    }
}

impl<'a> AddColor<'a, TweenColorContext<'a>> for TweenContext<'a> {
    #[inline(always)]
    fn rgba(&'a self, r: f32, g: f32, b: f32, a: f32) -> TweenColorContext<'a> {
        TweenColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            tween_factor: Borrowed(self.tween_factor.get()),
            color: Value(Color([r, g, b, a])),
        }
    }
}

impl<'a, 'b> AddPolygons<'a, TweenPolygonsContext<'a, 'b>> for TweenContext<'a> {
    #[inline(always)]
    fn polygons(&'a self, polygons: Polygons<'b>) -> TweenPolygonsContext<'a, 'b> {
        TweenPolygonsContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.transform.get()),
            tween_factor: Borrowed(self.tween_factor.get()),
            polygons: Value(polygons),
        }
    }
}

impl<'a> HasTransform<'a, Matrix2d> for TweenContext<'a> {
    #[inline(always)]
    fn get_transform(&'a self) -> &'a Matrix2d {
        self.transform.get()
    }
}

impl<'a> CanTransform<'a, TweenContext<'a>, Matrix2d> for TweenContext<'a> {
    #[inline(always)]
    fn transform(&'a self, value: Matrix2d) -> TweenContext<'a> {
        TweenContext {
            base: Borrowed(self.base.get()),
            transform: Value(value),
            tween_factor: Borrowed(self.tween_factor.get()),
        }
    }
}

impl<'a> HasViewTransform<'a, Matrix2d> for TweenContext<'a> {
    #[inline(always)]
    fn get_view_transform(&'a self) -> &'a Matrix2d {
        self.base.get()
    }
}

impl<'a> CanViewTransform<'a, TweenContext<'a>, Matrix2d> for TweenContext<'a> {
    #[inline(always)]
    fn view_transform(&'a self, value: Matrix2d) -> TweenContext<'a> {
        TweenContext {
            base: Value(value),
            transform: Borrowed(self.transform.get()),
            tween_factor: Borrowed(self.tween_factor.get()),
        }
    }
}

<|endoftext|>


<commit_before>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

use crate::DiskType;

use crate::sys::disk::{new_disk, Disk};
use crate::sys::processor::{self, Processor, Query};

use std::collections::HashMap;
use std::ffi::OsStr;
use std::mem::{size_of, zeroed};

use winapi::{ctypes::c_void, um::winbase::DRIVE_REMOVABLE};

use winapi::shared::minwindef::{BYTE, DWORD, MAX_PATH, TRUE};
use winapi::um::fileapi::{
    CreateFileW, GetDriveTypeW, GetLogicalDrives, GetVolumeInformationW, OPEN_EXISTING,
};
use winapi::um::handleapi::CloseHandle;
use winapi::um::handleapi::INVALID_HANDLE_VALUE;
use winapi::um::ioapiset::DeviceIoControl;
use winapi::um::sysinfoapi::{GetSystemInfo, SYSTEM_INFO};
use winapi::um::winbase::DRIVE_FIXED;
use winapi::um::winioctl::{
    DISK_GEOMETRY, IOCTL_DISK_GET_DRIVE_GEOMETRY, IOCTL_STORAGE_QUERY_PROPERTY,
};
use winapi::um::winnt::{BOOLEAN, FILE_SHARE_READ, FILE_SHARE_WRITE, HANDLE};

pub struct KeyHandler {
    pub unique_id: String,
    pub win_key: Vec<u16>,
}

impl KeyHandler {
    pub fn new(unique_id: String, win_key: Vec<u16>) -> KeyHandler {
        KeyHandler { unique_id, win_key }
    }
}

pub fn init_processors() -> (Vec<Processor>, String, String) {
    unsafe {
        let mut sys_info: SYSTEM_INFO = zeroed();
        GetSystemInfo(&mut sys_info);
        let (vendor_id, brand) = processor::get_vendor_id_and_brand(&sys_info);
        let frequencies = processor::get_frequencies(sys_info.dwNumberOfProcessors as usize);
        let mut ret = Vec::with_capacity(sys_info.dwNumberOfProcessors as usize + 1);
        for nb in 0..sys_info.dwNumberOfProcessors {
            ret.push(Processor::new_with_values(
                &format!("CPU {}", nb + 1),
                vendor_id.clone(),
                brand.clone(),
                frequencies[nb as usize],
            ));
        }
        (ret, vendor_id, brand)
    }
}

pub unsafe fn open_drive(drive_name: &[u16], open_rights: DWORD) -> HANDLE {
    CreateFileW(
        drive_name.as_ptr(),
        open_rights,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        std::ptr::null_mut(),
        OPEN_EXISTING,
        0,
        std::ptr::null_mut(),
    )
}

pub unsafe fn get_drive_size(handle: HANDLE) -> u64 {
    let mut pdg: DISK_GEOMETRY = std::mem::zeroed();
    let mut junk = 0;
    let result = DeviceIoControl(
        handle,
        IOCTL_DISK_GET_DRIVE_GEOMETRY,
        std::ptr::null_mut(),
        0,
        &mut pdg as *mut DISK_GEOMETRY as *mut c_void,
        size_of::<DISK_GEOMETRY>() as DWORD,
        &mut junk,
        std::ptr::null_mut(),
    );
    if result == TRUE {
        *pdg.Cylinders.QuadPart() as u64
            * pdg.TracksPerCylinder as u64
            * pdg.SectorsPerTrack as u64
            * pdg.BytesPerSector as u64
    } else {
        0
    }
}

pub unsafe fn get_disks() -> Vec<Disk> {
    let drives = GetLogicalDrives();
    if drives == 0 {
        return Vec::new();
    }

    #[cfg(feature = "multithread")]
    use rayon::iter::ParallelIterator;

    crate::utils::into_iter(0..size_of::<DWORD>() * 8)
        .filter_map(|x| {
            if (drives >> x) & 1 == 0 {
                return None;
            }
            let mount_point = [b'A' as u16 + x as u16, b':' as u16, b'\\' as u16, 0];

            let drive_type = GetDriveTypeW(mount_point.as_ptr());

            let is_removable = drive_type == DRIVE_REMOVABLE;

            if drive_type != DRIVE_FIXED && drive_type != DRIVE_REMOVABLE {
                return None;
            }
            let mut name = [0u16; MAX_PATH + 1];
            let mut file_system = [0u16; 32];
            if GetVolumeInformationW(
                mount_point.as_ptr(),
                name.as_mut_ptr(),
                name.len() as DWORD,
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                file_system.as_mut_ptr(),
                file_system.len() as DWORD,
            ) == 0
            {
                return None;
            }
            let mut pos = 0;
            for x in name.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let name = String::from_utf16_lossy(&name[..pos]);
            let name = OsStr::new(&name);

            pos = 0;
            for x in file_system.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let file_system: Vec<u8> = file_system[..pos].iter().map(|x| *x as u8).collect();

            let drive_name = [
                b'\\' as u16,
                b'\\' as u16,
                b'.' as u16,
                b'\\' as u16,
                b'A' as u16 + x as u16,
                b':' as u16,
                0,
            ];
            let handle = open_drive(&drive_name, 0);
            if handle == INVALID_HANDLE_VALUE {
                CloseHandle(handle);
                return new_disk(
                    name,
                    &mount_point,
                    &file_system,
                    if is_removable {
                        DiskType::Removable
                    } else {
                        DiskType::Unknown(-1)
                    },
                    0,
                );
            }
            let disk_size = get_drive_size(handle);
            /*let mut spq_trim: ffi::STORAGE_PROPERTY_QUERY = std::mem::zeroed();
            spq_trim.PropertyId = ffi::StorageDeviceTrimProperty;
            spq_trim.QueryType = ffi::PropertyStandardQuery;
            let mut dtd: ffi::DEVICE_TRIM_DESCRIPTOR = std::mem::zeroed();*/
            #[allow(non_snake_case)]
            #[repr(C)]
            struct STORAGE_PROPERTY_QUERY {
                PropertyId: i32,
                QueryType: i32,
                AdditionalParameters: [BYTE; 1],
            }
            #[allow(non_snake_case)]
            #[repr(C)]
            struct DEVICE_TRIM_DESCRIPTOR {
                Version: DWORD,
                Size: DWORD,
                TrimEnabled: BOOLEAN,
            }
            let mut spq_trim = STORAGE_PROPERTY_QUERY {
                PropertyId: 8i32,
                QueryType: 0i32,
                AdditionalParameters: [0],
            };
            let mut dtd: DEVICE_TRIM_DESCRIPTOR = std::mem::zeroed();

            let mut dw_size = 0;
            if DeviceIoControl(
                handle,
                IOCTL_STORAGE_QUERY_PROPERTY,
                &mut spq_trim as *mut STORAGE_PROPERTY_QUERY as *mut c_void,
                size_of::<STORAGE_PROPERTY_QUERY>() as DWORD,
                &mut dtd as *mut DEVICE_TRIM_DESCRIPTOR as *mut c_void,
                size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD,
                &mut dw_size,
                std::ptr::null_mut(),
            ) == 0
                || dw_size != size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD
            {
                CloseHandle(handle);
                return new_disk(
                    name,
                    &mount_point,
                    &file_system,
                    if is_removable {
                        DiskType::Removable
                    } else {
                        DiskType::Unknown(-1)
                    },
                    disk_size,
                );
            }
            let is_ssd = dtd.TrimEnabled != 0;
            CloseHandle(handle);
            new_disk(
                name,
                &mount_point,
                &file_system,
                match (is_removable, is_ssd) {
                    (true, _) => DiskType::Removable,
                    (false, true) => DiskType::SSD,
                    (false, false) => DiskType::HDD,
                },
                disk_size,
            )
        })
        .collect::<Vec<_>>()
}

#[allow(non_snake_case)]
pub unsafe fn load_symbols() -> HashMap<String, u32> {
    use winapi::um::winreg::{RegQueryValueExA, HKEY_PERFORMANCE_DATA};

    let mut cbCounters = 0;
    let mut dwType = 0;
    let mut ret = HashMap::new();

    let _dwStatus = RegQueryValueExA(
        HKEY_PERFORMANCE_DATA,
        b"Counter 009\0".as_ptr() as *const _,
        std::ptr::null_mut(),
        &mut dwType as *mut i32 as *mut _,
        std::ptr::null_mut(),
        &mut cbCounters as *mut i32 as *mut _,
    );

    let mut lpmszCounters = Vec::with_capacity(cbCounters as usize);
    lpmszCounters.set_len(cbCounters as usize);
    let _dwStatus = RegQueryValueExA(
        HKEY_PERFORMANCE_DATA,
        b"Counter 009\0".as_ptr() as *const _,
        std::ptr::null_mut(),
        &mut dwType as *mut i32 as *mut _,
        lpmszCounters.as_mut_ptr(),
        &mut cbCounters as *mut i32 as *mut _,
    );
    for (pos, s) in lpmszCounters
        .split(|x| *x == 0)
        .filter(|x| !x.is_empty())
        .collect::<Vec<_>>()
        .chunks(2)
        .filter(|&x| x.len() == 2)
        .filter_map(
            |x| match (std::str::from_utf8(x[0]), String::from_utf8(x[1].to_vec())) {
                (Ok(n), Ok(s)) => {
                    if let Ok(n) = u32::from_str_radix(n, 10) {
                        Some((n, s))
                    } else {
                        None
                    }
                }
                _ => None,
            },
        )
    {
        ret.insert(s, pos as u32);
    }
    ret
}

#[allow(clippy::ptr_arg)]
pub fn get_translation(s: &String, map: &HashMap<String, u32>) -> Option<String> {
    use winapi::um::pdh::PdhLookupPerfNameByIndexW;

    if let Some(index) = map.get(s) {
        let mut size: usize = 0;
        unsafe {
            let _res = PdhLookupPerfNameByIndexW(
                std::ptr::null(),
                *index,
                std::ptr::null_mut(),
                &mut size as *mut usize as *mut _,
            );
            if size == 0 {
                return Some(String::new());
            } else {
                let mut v = Vec::with_capacity(size);
                v.set_len(size);
                let _res = PdhLookupPerfNameByIndexW(
                    std::ptr::null(),
                    *index,
                    v.as_mut_ptr() as *mut _,
                    &mut size as *mut usize as *mut _,
                );
                return Some(String::from_utf16(&v[..size - 1]).expect("invalid utf16"));
            }
        }
    }
    None
}

pub fn add_counter(
    s: String,
    query: &mut Query,
    keys: &mut Option<KeyHandler>,
    counter_name: String,
) {
    let mut full = s.encode_utf16().collect::<Vec<_>>();
    full.push(0);
    if query.add_counter(&counter_name, full.clone()) {
        *keys = Some(KeyHandler::new(counter_name, full));
    }
}
<commit_msg>Windows use partition size instaid of disk size<commit_after>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

use crate::DiskType;

use crate::sys::disk::{new_disk, Disk};
use crate::sys::processor::{self, Processor, Query};

use std::collections::HashMap;
use std::ffi::OsStr;
use std::mem::{size_of, zeroed};

use winapi::{ctypes::c_void, um::winbase::DRIVE_REMOVABLE};

use winapi::shared::minwindef::{BYTE, DWORD, MAX_PATH, TRUE};
use winapi::um::fileapi::{
    CreateFileW, GetDriveTypeW, GetLogicalDrives, GetVolumeInformationW, OPEN_EXISTING,
};
use winapi::um::handleapi::CloseHandle;
use winapi::um::handleapi::INVALID_HANDLE_VALUE;
use winapi::um::ioapiset::DeviceIoControl;
use winapi::um::sysinfoapi::{GetSystemInfo, SYSTEM_INFO};
use winapi::um::winbase::DRIVE_FIXED;
use winapi::um::winioctl::{
    IOCTL_DISK_GET_PARTITION_INFO_EX, IOCTL_STORAGE_QUERY_PROPERTY, PARTITION_INFORMATION_EX,
};
use winapi::um::winnt::{BOOLEAN, FILE_SHARE_READ, FILE_SHARE_WRITE, HANDLE};

pub struct KeyHandler {
    pub unique_id: String,
    pub win_key: Vec<u16>,
}

impl KeyHandler {
    pub fn new(unique_id: String, win_key: Vec<u16>) -> KeyHandler {
        KeyHandler { unique_id, win_key }
    }
}

pub fn init_processors() -> (Vec<Processor>, String, String) {
    unsafe {
        let mut sys_info: SYSTEM_INFO = zeroed();
        GetSystemInfo(&mut sys_info);
        let (vendor_id, brand) = processor::get_vendor_id_and_brand(&sys_info);
        let frequencies = processor::get_frequencies(sys_info.dwNumberOfProcessors as usize);
        let mut ret = Vec::with_capacity(sys_info.dwNumberOfProcessors as usize + 1);
        for nb in 0..sys_info.dwNumberOfProcessors {
            ret.push(Processor::new_with_values(
                &format!("CPU {}", nb + 1),
                vendor_id.clone(),
                brand.clone(),
                frequencies[nb as usize],
            ));
        }
        (ret, vendor_id, brand)
    }
}

pub unsafe fn open_drive(drive_name: &[u16], open_rights: DWORD) -> HANDLE {
    CreateFileW(
        drive_name.as_ptr(),
        open_rights,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        std::ptr::null_mut(),
        OPEN_EXISTING,
        0,
        std::ptr::null_mut(),
    )
}

pub unsafe fn get_drive_size(handle: HANDLE) -> u64 {
    let mut pdg: PARTITION_INFORMATION_EX = std::mem::zeroed();
    let mut junk = 0;
    let result = DeviceIoControl(
        handle,
        IOCTL_DISK_GET_PARTITION_INFO_EX,
        std::ptr::null_mut(),
        0,
        &mut pdg as *mut PARTITION_INFORMATION_EX as *mut c_void,
        size_of::<PARTITION_INFORMATION_EX>() as DWORD,
        &mut junk,
        std::ptr::null_mut(),
    );
    if result == TRUE {
        *pdg.PartitionLength.QuadPart() as u64
    } else {
        0
    }
}

pub unsafe fn get_disks() -> Vec<Disk> {
    let drives = GetLogicalDrives();
    if drives == 0 {
        return Vec::new();
    }

    #[cfg(feature = "multithread")]
    use rayon::iter::ParallelIterator;

    crate::utils::into_iter(0..size_of::<DWORD>() * 8)
        .filter_map(|x| {
            if (drives >> x) & 1 == 0 {
                return None;
            }
            let mount_point = [b'A' as u16 + x as u16, b':' as u16, b'\\' as u16, 0];

            let drive_type = GetDriveTypeW(mount_point.as_ptr());

            let is_removable = drive_type == DRIVE_REMOVABLE;

            if drive_type != DRIVE_FIXED && drive_type != DRIVE_REMOVABLE {
                return None;
            }
            let mut name = [0u16; MAX_PATH + 1];
            let mut file_system = [0u16; 32];
            if GetVolumeInformationW(
                mount_point.as_ptr(),
                name.as_mut_ptr(),
                name.len() as DWORD,
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                file_system.as_mut_ptr(),
                file_system.len() as DWORD,
            ) == 0
            {
                return None;
            }
            let mut pos = 0;
            for x in name.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let name = String::from_utf16_lossy(&name[..pos]);
            let name = OsStr::new(&name);

            pos = 0;
            for x in file_system.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let file_system: Vec<u8> = file_system[..pos].iter().map(|x| *x as u8).collect();

            let drive_name = [
                b'\\' as u16,
                b'\\' as u16,
                b'.' as u16,
                b'\\' as u16,
                b'A' as u16 + x as u16,
                b':' as u16,
                0,
            ];
            let handle = open_drive(&drive_name, 0);
            if handle == INVALID_HANDLE_VALUE {
                CloseHandle(handle);
                return new_disk(
                    name,
                    &mount_point,
                    &file_system,
                    if is_removable {
                        DiskType::Removable
                    } else {
                        DiskType::Unknown(-1)
                    },
                    0,
                );
            }
            let disk_size = get_drive_size(handle);
            /*let mut spq_trim: ffi::STORAGE_PROPERTY_QUERY = std::mem::zeroed();
            spq_trim.PropertyId = ffi::StorageDeviceTrimProperty;
            spq_trim.QueryType = ffi::PropertyStandardQuery;
            let mut dtd: ffi::DEVICE_TRIM_DESCRIPTOR = std::mem::zeroed();*/
            #[allow(non_snake_case)]
            #[repr(C)]
            struct STORAGE_PROPERTY_QUERY {
                PropertyId: i32,
                QueryType: i32,
                AdditionalParameters: [BYTE; 1],
            }
            #[allow(non_snake_case)]
            #[repr(C)]
            struct DEVICE_TRIM_DESCRIPTOR {
                Version: DWORD,
                Size: DWORD,
                TrimEnabled: BOOLEAN,
            }
            let mut spq_trim = STORAGE_PROPERTY_QUERY {
                PropertyId: 8i32,
                QueryType: 0i32,
                AdditionalParameters: [0],
            };
            let mut dtd: DEVICE_TRIM_DESCRIPTOR = std::mem::zeroed();

            let mut dw_size = 0;
            if DeviceIoControl(
                handle,
                IOCTL_STORAGE_QUERY_PROPERTY,
                &mut spq_trim as *mut STORAGE_PROPERTY_QUERY as *mut c_void,
                size_of::<STORAGE_PROPERTY_QUERY>() as DWORD,
                &mut dtd as *mut DEVICE_TRIM_DESCRIPTOR as *mut c_void,
                size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD,
                &mut dw_size,
                std::ptr::null_mut(),
            ) == 0
                || dw_size != size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD
            {
                CloseHandle(handle);
                return new_disk(
                    name,
                    &mount_point,
                    &file_system,
                    if is_removable {
                        DiskType::Removable
                    } else {
                        DiskType::Unknown(-1)
                    },
                    disk_size,
                );
            }
            let is_ssd = dtd.TrimEnabled != 0;
            CloseHandle(handle);
            new_disk(
                name,
                &mount_point,
                &file_system,
                match (is_removable, is_ssd) {
                    (true, _) => DiskType::Removable,
                    (false, true) => DiskType::SSD,
                    (false, false) => DiskType::HDD,
                },
                disk_size,
            )
        })
        .collect::<Vec<_>>()
}

#[allow(non_snake_case)]
pub unsafe fn load_symbols() -> HashMap<String, u32> {
    use winapi::um::winreg::{RegQueryValueExA, HKEY_PERFORMANCE_DATA};

    let mut cbCounters = 0;
    let mut dwType = 0;
    let mut ret = HashMap::new();

    let _dwStatus = RegQueryValueExA(
        HKEY_PERFORMANCE_DATA,
        b"Counter 009\0".as_ptr() as *const _,
        std::ptr::null_mut(),
        &mut dwType as *mut i32 as *mut _,
        std::ptr::null_mut(),
        &mut cbCounters as *mut i32 as *mut _,
    );

    let mut lpmszCounters = Vec::with_capacity(cbCounters as usize);
    lpmszCounters.set_len(cbCounters as usize);
    let _dwStatus = RegQueryValueExA(
        HKEY_PERFORMANCE_DATA,
        b"Counter 009\0".as_ptr() as *const _,
        std::ptr::null_mut(),
        &mut dwType as *mut i32 as *mut _,
        lpmszCounters.as_mut_ptr(),
        &mut cbCounters as *mut i32 as *mut _,
    );
    for (pos, s) in lpmszCounters
        .split(|x| *x == 0)
        .filter(|x| !x.is_empty())
        .collect::<Vec<_>>()
        .chunks(2)
        .filter(|&x| x.len() == 2)
        .filter_map(
            |x| match (std::str::from_utf8(x[0]), String::from_utf8(x[1].to_vec())) {
                (Ok(n), Ok(s)) => {
                    if let Ok(n) = u32::from_str_radix(n, 10) {
                        Some((n, s))
                    } else {
                        None
                    }
                }
                _ => None,
            },
        )
    {
        ret.insert(s, pos as u32);
    }
    ret
}

#[allow(clippy::ptr_arg)]
pub fn get_translation(s: &String, map: &HashMap<String, u32>) -> Option<String> {
    use winapi::um::pdh::PdhLookupPerfNameByIndexW;

    if let Some(index) = map.get(s) {
        let mut size: usize = 0;
        unsafe {
            let _res = PdhLookupPerfNameByIndexW(
                std::ptr::null(),
                *index,
                std::ptr::null_mut(),
                &mut size as *mut usize as *mut _,
            );
            if size == 0 {
                return Some(String::new());
            } else {
                let mut v = Vec::with_capacity(size);
                v.set_len(size);
                let _res = PdhLookupPerfNameByIndexW(
                    std::ptr::null(),
                    *index,
                    v.as_mut_ptr() as *mut _,
                    &mut size as *mut usize as *mut _,
                );
                return Some(String::from_utf16(&v[..size - 1]).expect("invalid utf16"));
            }
        }
    }
    None
}

pub fn add_counter(
    s: String,
    query: &mut Query,
    keys: &mut Option<KeyHandler>,
    counter_name: String,
) {
    let mut full = s.encode_utf16().collect::<Vec<_>>();
    full.push(0);
    if query.add_counter(&counter_name, full.clone()) {
        *keys = Some(KeyHandler::new(counter_name, full));
    }
}
<|endoftext|>


<commit_before>use regex::Regex;
use std::path::PathBuf;
use std::result::Result;
use subprocess::{Exec, Redirection};

const LOOP_BACKING_DEVICE_BASENAME: &'static str = "flaglock6-home";
const LOOP_DEVICE_PARTITION_LABEL: &'static str = "fl6";
const RSYNC_DOT_NET_REMOTE_DIR: &'static str = "usw-s007.rsync.net:./";

#[derive(Debug)]
pub enum AuricError {
    Invocation(String),
    Io(String),
    Subprocess(String),
    NotFound,
}

enum AuricOperationMode {
    Mount,
    Unmount,
}

impl From<subprocess::PopenError> for AuricError {
    fn from(err: subprocess::PopenError) -> AuricError {
        AuricError::Subprocess(err.to_string())
    }
}

struct SshfsManager {
    mountpoint: PathBuf,
}

struct LoopManager {
    target: Regex,
    backing: PathBuf,
}

struct LuksManager {
    device: PathBuf,
}

struct AuricImpl {
    sshfs_manager: SshfsManager,
    loop_manager: LoopManager,
    luks_manager: LuksManager,
}

fn exec_with_stdout(command: &str, args: &[&str], error_hint: &str) -> Result<String, AuricError> {
    let result = Exec::cmd(command)
        .args(args)
        .stdout(Redirection::Pipe)
        .stderr(Redirection::Pipe)
        .capture()?;
    if !result.success() {
        return Err(AuricError::Subprocess(format!(
            "{} failed: ``{}''",
            error_hint,
            result.stderr_str(),
        )));
    }
    Ok(result.stdout_str())
}

fn exec(command: &str, args: &[&str], error_hint: &str) -> Result<(), AuricError> {
    exec_with_stdout(command, args, error_hint).map(|_stdout_str| ())
}

impl SshfsManager {
    pub fn new() -> Result<SshfsManager, AuricError> {
        let xdg_runtime_dir = match std::env::var("XDG_RUNTIME_DIR") {
            Ok(val) => val,
            Err(e) => {
                return Err(AuricError::Invocation(format!(
                    "fatal error getting ${{XDG_RUNTIME_DIR}}: ``{}''",
                    e
                )))
            }
        };
        let mut mountpoint = PathBuf::new();
        mountpoint.push(xdg_runtime_dir);
        mountpoint.push("rsync.net-mountpoint");

        Ok(SshfsManager {
            mountpoint: mountpoint,
        })
    }

    fn ensure_mountpoint_is_dir(&self) -> Result<(), AuricError> {
        if self.mountpoint.is_dir() {
            return Ok(());
        }
        std::fs::create_dir(self.mountpoint.as_path())
            .map_err(|e| AuricError::Io(format!("bad mountpoint: ``{}''", e)))
    }

    // Does _not_ ensure that backing device exists.
    pub fn loop_backing_device_path(&self) -> PathBuf {
        let mut result: PathBuf = self.mountpoint.clone();
        result.push(LOOP_BACKING_DEVICE_BASENAME);
        result
    }

    fn loop_backing_device_present(&self) -> bool {
        self.loop_backing_device_path().is_file()
    }

    pub fn mount(&self) -> Result<(), AuricError> {
        self.ensure_mountpoint_is_dir()?;
        if self.loop_backing_device_present() {
            return Ok(());
        }

        // Invoke `sshfs`
        // *    with compression and
        // *    allowing other users to access the mounted filesystem
        //      (necessary to allow loop-mounts).
        exec(
            "sshfs",
            &[
                "-C",
                "-o",
                "allow_other",
                RSYNC_DOT_NET_REMOTE_DIR,
                self.mountpoint.to_str().unwrap(),
            ],
            "sshfs mount",
        )
    }

    pub fn unmount(&self) -> Result<(), AuricError> {
        if !self.loop_backing_device_present() {
            return Ok(());
        }
        exec(
            "fusermount",
            &["-u", self.mountpoint.to_str().unwrap()],
            "fusermount -u",
        )
    }
}

impl LoopManager {
    pub fn new(backing: &PathBuf) -> LoopManager {
        let target = format!(
            r"^(/dev/loop\d)(\s+\d+){{4}}\s+({})\s+",
            backing.to_str().unwrap()
        );
        LoopManager {
            target: Regex::new(&target).unwrap(),
            backing: backing.clone(),
        }
    }

    pub fn find(&self) -> Result<PathBuf, AuricError> {
        let losetup_stdout = exec_with_stdout("losetup", &["-l"], "loop device enumeration")?;
        for line in losetup_stdout.lines() {
            if let Some(captures) = self.target.captures(line) {
                return Ok(PathBuf::from(captures.get(1).unwrap().as_str()));
            }
        }
        Err(AuricError::NotFound)
    }

    pub fn attach(&self) -> Result<(), AuricError> {
        if let Ok(_attached_loop) = self.find() {
            return Ok(());
        }

        exec(
            "udisksctl",
            &["loop-setup", "-f", self.backing.to_str().unwrap()],
            "loop device attach",
        )
    }

    // Note that the loop device will linger if the volume remains
    // unlocked. Prefer not to call this before locking.
    pub fn detach(&self) -> Result<(), AuricError> {
        let attached_loop = match self.find() {
            Ok(path) => path,
            Err(e) => {
                if let AuricError::NotFound = e {
                    return Ok(());
                }
                return Err(e);
            }
        };
        exec(
            "udisksctl",
            &["loop-delete", "-b", attached_loop.to_str().unwrap()],
            "loop device detach",
        )
    }
}

impl LuksManager {
    pub fn new() -> LuksManager {
        let mut device = PathBuf::from("/dev/disk/by-label");
        device.push(LOOP_DEVICE_PARTITION_LABEL);
        LuksManager { device: device }
    }

    pub fn unlock(&self, locked_device: &PathBuf) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["unlock", "-b", locked_device.to_str().unwrap()],
            format!("unlocking {}", locked_device.to_str().unwrap()).as_str(),
        )
    }

    pub fn lock(&self, underlying_device: &PathBuf) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["lock", "-b", underlying_device.to_str().unwrap()],
            format!("locking {}", underlying_device.to_str().unwrap()).as_str(),
        )
    }

    pub fn mount(&self) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["mount", "-b", self.device.to_str().unwrap()],
            format!("mounting {}", self.device.to_str().unwrap()).as_str(),
        )
    }

    pub fn unmount(&self) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["unmount", "-b", self.device.to_str().unwrap()],
            format!("unmounting {}", self.device.to_str().unwrap()).as_str(),
        )
    }
}

fn get_action() -> Result<AuricOperationMode, AuricError> {
    let mut args_iter = std::env::args();
    args_iter.next();
    if let Some(action) = args_iter.next() {
        match action.as_str() {
            "mount" => return Ok(AuricOperationMode::Mount),
            "unmount" => return Ok(AuricOperationMode::Unmount),
            _ => {
                return Err(AuricError::Invocation(format!(
                    "unknown action: ``{}''",
                    action
                )))
            }
        }
    }
    Err(AuricError::Invocation("no action given".to_string()))
}

impl AuricImpl {
    pub fn new() -> Result<AuricImpl, AuricError> {
        let sshfs_manager = SshfsManager::new()?;
        let loop_backing_device_path = sshfs_manager.loop_backing_device_path();
        Ok(AuricImpl {
            sshfs_manager: sshfs_manager,
            loop_manager: LoopManager::new(&loop_backing_device_path),
            luks_manager: LuksManager::new(),
        })
    }

    fn mount(&self) -> Result<(), AuricError> {
        self.sshfs_manager.mount()?;
        self.loop_manager.attach()?;
        let loop_device = self.loop_manager.find()?;
        self.luks_manager.unlock(&loop_device)?;
        self.luks_manager.mount()
    }

    fn unmount(&self) -> Result<(), AuricError> {
        self.luks_manager.unmount()?;
        let loop_device = self.loop_manager.find()?;
        self.luks_manager.lock(&loop_device)?;
        self.loop_manager.detach()?;
        self.sshfs_manager.unmount()
    }

    pub fn act(&self, mode: AuricOperationMode) -> Result<(), AuricError> {
        match mode {
            AuricOperationMode::Mount => return self.mount(),
            AuricOperationMode::Unmount => return self.unmount(),
        }
    }
}

fn main_impl() -> Result<(), AuricError> {
    let action = get_action()?;
    let auric_impl = AuricImpl::new()?;
    auric_impl.act(action)
}

fn main() {
    std::process::exit(match main_impl() {
        Ok(_) => 0,
        Err(e) => {
            eprintln!("{:?}", e);
            1
        }
    });
}
<commit_msg>apricot-auric: Add ``sshfs -o reconnect''<commit_after>use regex::Regex;
use std::path::PathBuf;
use std::result::Result;
use subprocess::{Exec, Redirection};

const LOOP_BACKING_DEVICE_BASENAME: &'static str = "flaglock6-home";
const LOOP_DEVICE_PARTITION_LABEL: &'static str = "fl6";
const RSYNC_DOT_NET_REMOTE_DIR: &'static str = "usw-s007.rsync.net:./";

#[derive(Debug)]
pub enum AuricError {
    Invocation(String),
    Io(String),
    Subprocess(String),
    NotFound,
}

enum AuricOperationMode {
    Mount,
    Unmount,
}

impl From<subprocess::PopenError> for AuricError {
    fn from(err: subprocess::PopenError) -> AuricError {
        AuricError::Subprocess(err.to_string())
    }
}

struct SshfsManager {
    mountpoint: PathBuf,
}

struct LoopManager {
    target: Regex,
    backing: PathBuf,
}

struct LuksManager {
    device: PathBuf,
}

struct AuricImpl {
    sshfs_manager: SshfsManager,
    loop_manager: LoopManager,
    luks_manager: LuksManager,
}

fn exec_with_stdout(command: &str, args: &[&str], error_hint: &str) -> Result<String, AuricError> {
    let result = Exec::cmd(command)
        .args(args)
        .stdout(Redirection::Pipe)
        .stderr(Redirection::Pipe)
        .capture()?;
    if !result.success() {
        return Err(AuricError::Subprocess(format!(
            "{} failed: ``{}''",
            error_hint,
            result.stderr_str(),
        )));
    }
    Ok(result.stdout_str())
}

fn exec(command: &str, args: &[&str], error_hint: &str) -> Result<(), AuricError> {
    exec_with_stdout(command, args, error_hint).map(|_stdout_str| ())
}

impl SshfsManager {
    pub fn new() -> Result<SshfsManager, AuricError> {
        let xdg_runtime_dir = match std::env::var("XDG_RUNTIME_DIR") {
            Ok(val) => val,
            Err(e) => {
                return Err(AuricError::Invocation(format!(
                    "fatal error getting ${{XDG_RUNTIME_DIR}}: ``{}''",
                    e
                )))
            }
        };
        let mut mountpoint = PathBuf::new();
        mountpoint.push(xdg_runtime_dir);
        mountpoint.push("rsync.net-mountpoint");

        Ok(SshfsManager {
            mountpoint: mountpoint,
        })
    }

    fn ensure_mountpoint_is_dir(&self) -> Result<(), AuricError> {
        if self.mountpoint.is_dir() {
            return Ok(());
        }
        std::fs::create_dir(self.mountpoint.as_path())
            .map_err(|e| AuricError::Io(format!("bad mountpoint: ``{}''", e)))
    }

    // Does _not_ ensure that backing device exists.
    pub fn loop_backing_device_path(&self) -> PathBuf {
        let mut result: PathBuf = self.mountpoint.clone();
        result.push(LOOP_BACKING_DEVICE_BASENAME);
        result
    }

    fn loop_backing_device_present(&self) -> bool {
        self.loop_backing_device_path().is_file()
    }

    pub fn mount(&self) -> Result<(), AuricError> {
        self.ensure_mountpoint_is_dir()?;
        if self.loop_backing_device_present() {
            return Ok(());
        }

        // Invoke `sshfs`
        // *    with compression,
        // *    allowing other users to access the mounted filesystem
        //      (necessary to allow loop-mounts), and
        // *    attempting to reconnect if we lose the connection.
        exec(
            "sshfs",
            &[
                "-C",
                "-o",
                "allow_other,reconnect",
                RSYNC_DOT_NET_REMOTE_DIR,
                self.mountpoint.to_str().unwrap(),
            ],
            "sshfs mount",
        )
    }

    pub fn unmount(&self) -> Result<(), AuricError> {
        if !self.loop_backing_device_present() {
            return Ok(());
        }
        exec(
            "fusermount",
            &["-u", self.mountpoint.to_str().unwrap()],
            "fusermount -u",
        )
    }
}

impl LoopManager {
    pub fn new(backing: &PathBuf) -> LoopManager {
        let target = format!(
            r"^(/dev/loop\d)(\s+\d+){{4}}\s+({})\s+",
            backing.to_str().unwrap()
        );
        LoopManager {
            target: Regex::new(&target).unwrap(),
            backing: backing.clone(),
        }
    }

    pub fn find(&self) -> Result<PathBuf, AuricError> {
        let losetup_stdout = exec_with_stdout("losetup", &["-l"], "loop device enumeration")?;
        for line in losetup_stdout.lines() {
            if let Some(captures) = self.target.captures(line) {
                return Ok(PathBuf::from(captures.get(1).unwrap().as_str()));
            }
        }
        Err(AuricError::NotFound)
    }

    pub fn attach(&self) -> Result<(), AuricError> {
        if let Ok(_attached_loop) = self.find() {
            return Ok(());
        }

        exec(
            "udisksctl",
            &["loop-setup", "-f", self.backing.to_str().unwrap()],
            "loop device attach",
        )
    }

    // Note that the loop device will linger if the volume remains
    // unlocked. Prefer not to call this before locking.
    pub fn detach(&self) -> Result<(), AuricError> {
        let attached_loop = match self.find() {
            Ok(path) => path,
            Err(e) => {
                if let AuricError::NotFound = e {
                    return Ok(());
                }
                return Err(e);
            }
        };
        exec(
            "udisksctl",
            &["loop-delete", "-b", attached_loop.to_str().unwrap()],
            "loop device detach",
        )
    }
}

impl LuksManager {
    pub fn new() -> LuksManager {
        let mut device = PathBuf::from("/dev/disk/by-label");
        device.push(LOOP_DEVICE_PARTITION_LABEL);
        LuksManager { device: device }
    }

    pub fn unlock(&self, locked_device: &PathBuf) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["unlock", "-b", locked_device.to_str().unwrap()],
            format!("unlocking {}", locked_device.to_str().unwrap()).as_str(),
        )
    }

    pub fn lock(&self, underlying_device: &PathBuf) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["lock", "-b", underlying_device.to_str().unwrap()],
            format!("locking {}", underlying_device.to_str().unwrap()).as_str(),
        )
    }

    pub fn mount(&self) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["mount", "-b", self.device.to_str().unwrap()],
            format!("mounting {}", self.device.to_str().unwrap()).as_str(),
        )
    }

    pub fn unmount(&self) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["unmount", "-b", self.device.to_str().unwrap()],
            format!("unmounting {}", self.device.to_str().unwrap()).as_str(),
        )
    }
}

fn get_action() -> Result<AuricOperationMode, AuricError> {
    let mut args_iter = std::env::args();
    args_iter.next();
    if let Some(action) = args_iter.next() {
        match action.as_str() {
            "mount" => return Ok(AuricOperationMode::Mount),
            "unmount" => return Ok(AuricOperationMode::Unmount),
            _ => {
                return Err(AuricError::Invocation(format!(
                    "unknown action: ``{}''",
                    action
                )))
            }
        }
    }
    Err(AuricError::Invocation("no action given".to_string()))
}

impl AuricImpl {
    pub fn new() -> Result<AuricImpl, AuricError> {
        let sshfs_manager = SshfsManager::new()?;
        let loop_backing_device_path = sshfs_manager.loop_backing_device_path();
        Ok(AuricImpl {
            sshfs_manager: sshfs_manager,
            loop_manager: LoopManager::new(&loop_backing_device_path),
            luks_manager: LuksManager::new(),
        })
    }

    fn mount(&self) -> Result<(), AuricError> {
        self.sshfs_manager.mount()?;
        self.loop_manager.attach()?;
        let loop_device = self.loop_manager.find()?;
        self.luks_manager.unlock(&loop_device)?;
        self.luks_manager.mount()
    }

    fn unmount(&self) -> Result<(), AuricError> {
        self.luks_manager.unmount()?;
        let loop_device = self.loop_manager.find()?;
        self.luks_manager.lock(&loop_device)?;
        self.loop_manager.detach()?;
        self.sshfs_manager.unmount()
    }

    pub fn act(&self, mode: AuricOperationMode) -> Result<(), AuricError> {
        match mode {
            AuricOperationMode::Mount => return self.mount(),
            AuricOperationMode::Unmount => return self.unmount(),
        }
    }
}

fn main_impl() -> Result<(), AuricError> {
    let action = get_action()?;
    let auric_impl = AuricImpl::new()?;
    auric_impl.act(action)
}

fn main() {
    std::process::exit(match main_impl() {
        Ok(_) => 0,
        Err(e) => {
            eprintln!("{:?}", e);
            1
        }
    });
}
<|endoftext|>


<commit_before>use support::{project, execs};
use hamcrest::assert_that;
use cargo;

fn setup() {}

test!(simple {
    let p = project("foo");

    assert_that(p.cargo_process("version"),
                execs().with_status(0).with_stdout(format!("{}\n",
        cargo::version()).as_slice()));
})
<commit_msg>add test: cargo --version<commit_after>use support::{project, execs};
use hamcrest::assert_that;
use cargo;

fn setup() {}

test!(simple {
    let p = project("foo");

    assert_that(p.cargo_process("version"),
                execs().with_status(0).with_stdout(format!("{}\n",
        cargo::version()).as_slice()));

    assert_that(p.cargo_process("--version"),
                execs().with_status(0).with_stdout(format!("{}\n",
        cargo::version()).as_slice()));

})
<|endoftext|>


<commit_before>use crate::case_result::CaseResult;
use crate::cases::*;
use crate::doc_header::DocHeader;
use crate::eth_specs::MinimalEthSpec;
use crate::yaml_decode::{extract_yaml_by_key, YamlDecode};
use crate::EfTest;
use serde_derive::Deserialize;
use std::{fs::File, io::prelude::*, path::PathBuf};
use types::{EthSpec, FoundationEthSpec};

#[derive(Debug, Deserialize)]
pub struct Doc {
    pub yaml: String,
    pub path: PathBuf,
}

impl Doc {
    fn from_path(path: PathBuf) -> Self {
        let mut file = File::open(path.clone()).unwrap();

        let mut yaml = String::new();
        file.read_to_string(&mut yaml).unwrap();

        Self { yaml, path }
    }

    pub fn test_results(&self) -> Vec<CaseResult> {
        let header: DocHeader = serde_yaml::from_str(&self.yaml.as_str()).unwrap();

        match (
            header.runner.as_ref(),
            header.handler.as_ref(),
            header.config.as_ref(),
        ) {
            ("ssz", "uint", _) => run_test::<SszGeneric, FoundationEthSpec>(&self.yaml),
            ("ssz", "static", "minimal") => run_test::<SszStatic, MinimalEthSpec>(&self.yaml),
            (runner, handler, config) => panic!(
                "No implementation for runner: \"{}\", handler: \"{}\", config: \"{}\"",
                runner, handler, config
            ),
        }
    }

    pub fn assert_tests_pass(path: PathBuf) {
        let doc = Self::from_path(path);
        let results = doc.test_results();

        if results.iter().any(|r| r.result.is_err()) {
            print_failures(&doc, &results);
            panic!("Tests failed (see above)");
        }
    }
}

pub fn run_test<T, E: EthSpec>(test_doc_yaml: &String) -> Vec<CaseResult>
where
    Cases<T>: EfTest + YamlDecode,
{
    // Extract only the "test_cases" YAML as a stand-alone string.
    let test_cases_yaml = extract_yaml_by_key(test_doc_yaml, "test_cases");

    // Pass only the "test_cases" YAML string to `yaml_decode`.
    let test_cases: Cases<T> = Cases::yaml_decode(&test_cases_yaml.to_string()).unwrap();

    test_cases.test_results::<E>()
}

pub fn print_failures(doc: &Doc, results: &[CaseResult]) {
    let header: DocHeader = serde_yaml::from_str(&doc.yaml).unwrap();
    let failures: Vec<&CaseResult> = results.iter().filter(|r| r.result.is_err()).collect();

    println!("--------------------------------------------------");
    println!("Test Failure");
    println!("Title: {}", header.title);
    println!("File: {:?}", doc.path);
    println!("");
    println!(
        "{} tests, {} failures, {} passes.",
        results.len(),
        failures.len(),
        results.len() - failures.len()
    );
    println!("");

    for failure in failures {
        println!("-------");
        println!(
            "case[{}] failed with: {:#?}",
            failure.case_index, failure.result
        );
    }
    println!("");
}
<commit_msg>Update `ef_test` error messages<commit_after>use crate::case_result::CaseResult;
use crate::cases::*;
use crate::doc_header::DocHeader;
use crate::eth_specs::MinimalEthSpec;
use crate::yaml_decode::{extract_yaml_by_key, YamlDecode};
use crate::EfTest;
use serde_derive::Deserialize;
use std::{fs::File, io::prelude::*, path::PathBuf};
use types::{EthSpec, FoundationEthSpec};

#[derive(Debug, Deserialize)]
pub struct Doc {
    pub yaml: String,
    pub path: PathBuf,
}

impl Doc {
    fn from_path(path: PathBuf) -> Self {
        let mut file = File::open(path.clone()).unwrap();

        let mut yaml = String::new();
        file.read_to_string(&mut yaml).unwrap();

        Self { yaml, path }
    }

    pub fn test_results(&self) -> Vec<CaseResult> {
        let header: DocHeader = serde_yaml::from_str(&self.yaml.as_str()).unwrap();

        match (
            header.runner.as_ref(),
            header.handler.as_ref(),
            header.config.as_ref(),
        ) {
            ("ssz", "uint", _) => run_test::<SszGeneric, FoundationEthSpec>(&self.yaml),
            ("ssz", "static", "minimal") => run_test::<SszStatic, MinimalEthSpec>(&self.yaml),
            (runner, handler, config) => panic!(
                "No implementation for runner: \"{}\", handler: \"{}\", config: \"{}\"",
                runner, handler, config
            ),
        }
    }

    pub fn assert_tests_pass(path: PathBuf) {
        let doc = Self::from_path(path);
        let results = doc.test_results();

        if results.iter().any(|r| r.result.is_err()) {
            print_failures(&doc, &results);
            panic!("Tests failed (see above)");
        }
    }
}

pub fn run_test<T, E: EthSpec>(test_doc_yaml: &String) -> Vec<CaseResult>
where
    Cases<T>: EfTest + YamlDecode,
{
    // Extract only the "test_cases" YAML as a stand-alone string.
    let test_cases_yaml = extract_yaml_by_key(test_doc_yaml, "test_cases");

    // Pass only the "test_cases" YAML string to `yaml_decode`.
    let test_cases: Cases<T> = Cases::yaml_decode(&test_cases_yaml.to_string()).unwrap();

    test_cases.test_results::<E>()
}

pub fn print_failures(doc: &Doc, results: &[CaseResult]) {
    let header: DocHeader = serde_yaml::from_str(&doc.yaml).unwrap();
    let failures: Vec<&CaseResult> = results.iter().filter(|r| r.result.is_err()).collect();

    println!("--------------------------------------------------");
    println!("Test Failure");
    println!("Title: {}", header.title);
    println!("File: {:?}", doc.path);
    println!("");
    println!(
        "{} tests, {} failures, {} passes.",
        results.len(),
        failures.len(),
        results.len() - failures.len()
    );
    println!("");

    for failure in failures {
        println!("-------");
        println!("case[{}].result:", failure.case_index);
        println!("{:#?}", failure.result);
    }
    println!("");
}
<|endoftext|>


<commit_before>//! Basic syntax highlighting functionality.
//!
//! This module uses librustc_ast's lexer to provide token-based highlighting for
//! the HTML documentation generated by rustdoc.
//!
//! Use the `render_with_highlighting` to highlight some rust code.

use crate::html::escape::Escape;

use std::fmt::{Display, Write};
use std::iter::Peekable;

use rustc_lexer::{LiteralKind, TokenKind};
use rustc_span::symbol::Ident;
use rustc_span::with_default_session_globals;

/// Highlights `src`, returning the HTML output.
pub fn render_with_highlighting(
    src: String,
    class: Option<&str>,
    playground_button: Option<&str>,
    tooltip: Option<(&str, &str)>,
) -> String {
    debug!("highlighting: ================\n{}\n==============", src);
    let mut out = String::with_capacity(src.len());
    if let Some((tooltip, class)) = tooltip {
        write!(
            out,
            "

ⓘ{}

",
            class, tooltip
        )
        .unwrap();
    }

    write_header(&mut out, class);
    write_code(&mut out, &src);
    write_footer(&mut out, playground_button);

    out
}

fn write_header(out: &mut String, class: Option<&str>) {
    write!(out, "

\n", class.unwrap_or_default())
        .unwrap()
}

fn write_code(out: &mut String, src: &str) {
    Classifier::new(src).highlight(&mut |highlight| {
        match highlight {
            Highlight::Token { text, class } => string(out, Escape(text), class),
            Highlight::EnterSpan { class } => enter_span(out, class),
            Highlight::ExitSpan => exit_span(out),
        };
    });
}

fn write_footer(out: &mut String, playground_button: Option<&str>) {
    write!(out, "
{}
\n", playground_button.unwrap_or_default()).unwrap()
}

/// How a span of text is classified. Mostly corresponds to token kinds.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Class {
    None,
    Comment,
    DocComment,
    Attribute,
    KeyWord,
    // Keywords that do pointer/reference stuff.
    RefKeyWord,
    Self_,
    Op,
    Macro,
    MacroNonTerminal,
    String,
    Number,
    Bool,
    Ident,
    Lifetime,
    PreludeTy,
    PreludeVal,
    QuestionMark,
}

impl Class {
    /// Returns the css class expected by rustdoc for each `Class`.
    fn as_html(self) -> &'static str {
        match self {
            Class::None => "",
            Class::Comment => "comment",
            Class::DocComment => "doccomment",
            Class::Attribute => "attribute",
            Class::KeyWord => "kw",
            Class::RefKeyWord => "kw-2",
            Class::Self_ => "self",
            Class::Op => "op",
            Class::Macro => "macro",
            Class::MacroNonTerminal => "macro-nonterminal",
            Class::String => "string",
            Class::Number => "number",
            Class::Bool => "bool-val",
            Class::Ident => "ident",
            Class::Lifetime => "lifetime",
            Class::PreludeTy => "prelude-ty",
            Class::PreludeVal => "prelude-val",
            Class::QuestionMark => "question-mark",
        }
    }
}

enum Highlight<'a> {
    Token { text: &'a str, class: Class },
    EnterSpan { class: Class },
    ExitSpan,
}

struct TokenIter<'a> {
    src: &'a str,
}

impl Iterator for TokenIter<'a> {
    type Item = (TokenKind, &'a str);
    fn next(&mut self) -> Option<(TokenKind, &'a str)> {
        if self.src.is_empty() {
            return None;
        }
        let token = rustc_lexer::first_token(self.src);
        let (text, rest) = self.src.split_at(token.len);
        self.src = rest;
        Some((token.kind, text))
    }
}

/// Processes program tokens, classifying strings of text by highlighting
/// category (`Class`).
struct Classifier<'a> {
    tokens: Peekable<TokenIter<'a>>,
    in_attribute: bool,
    in_macro: bool,
    in_macro_nonterminal: bool,
}

impl<'a> Classifier<'a> {
    fn new(src: &str) -> Classifier<'_> {
        let tokens = TokenIter { src }.peekable();
        Classifier { tokens, in_attribute: false, in_macro: false, in_macro_nonterminal: false }
    }

    /// Exhausts the `Classifier` writing the output into `sink`.
    ///
    /// The general structure for this method is to iterate over each token,
    /// possibly giving it an HTML span with a class specifying what flavor of
    /// token is used.
    fn highlight(mut self, sink: &mut dyn FnMut(Highlight<'a>)) {
        with_default_session_globals(|| {
            while let Some((token, text)) = self.tokens.next() {
                self.advance(token, text, sink);
            }
        })
    }

    /// Single step of highlighting. This will classify `token`, but maybe also
    /// a couple of following ones as well.
    fn advance(&mut self, token: TokenKind, text: &'a str, sink: &mut dyn FnMut(Highlight<'a>)) {
        let lookahead = self.peek();
        let class = match token {
            TokenKind::Whitespace => Class::None,
            TokenKind::LineComment { doc_style } | TokenKind::BlockComment { doc_style, .. } => {
                if doc_style.is_some() {
                    Class::DocComment
                } else {
                    Class::Comment
                }
            }
            // Consider this as part of a macro invocation if there was a
            // leading identifier.
            TokenKind::Bang if self.in_macro => {
                self.in_macro = false;
                Class::Macro
            }

            // Assume that '&' or '*' is the reference or dereference operator
            // or a reference or pointer type. Unless, of course, it looks like
            // a logical and or a multiplication operator: `&&` or `* `.
            TokenKind::Star => match lookahead {
                Some(TokenKind::Whitespace) => Class::Op,
                _ => Class::RefKeyWord,
            },
            TokenKind::And => match lookahead {
                Some(TokenKind::And) => {
                    let _and = self.tokens.next();
                    sink(Highlight::Token { text: "&&", class: Class::Op });
                    return;
                }
                Some(TokenKind::Eq) => {
                    let _eq = self.tokens.next();
                    sink(Highlight::Token { text: "&=", class: Class::Op });
                    return;
                }
                Some(TokenKind::Whitespace) => Class::Op,
                _ => Class::RefKeyWord,
            },

            // Operators.
            TokenKind::Minus
            | TokenKind::Plus
            | TokenKind::Or
            | TokenKind::Slash
            | TokenKind::Caret
            | TokenKind::Percent
            | TokenKind::Bang
            | TokenKind::Eq
            | TokenKind::Lt
            | TokenKind::Gt => Class::Op,

            // Miscellaneous, no highlighting.
            TokenKind::Dot
            | TokenKind::Semi
            | TokenKind::Comma
            | TokenKind::OpenParen
            | TokenKind::CloseParen
            | TokenKind::OpenBrace
            | TokenKind::CloseBrace
            | TokenKind::OpenBracket
            | TokenKind::At
            | TokenKind::Tilde
            | TokenKind::Colon
            | TokenKind::Unknown => Class::None,

            TokenKind::Question => Class::QuestionMark,

            TokenKind::Dollar => match lookahead {
                Some(TokenKind::Ident) => {
                    self.in_macro_nonterminal = true;
                    Class::MacroNonTerminal
                }
                _ => Class::None,
            },

            // This might be the start of an attribute. We're going to want to
            // continue highlighting it as an attribute until the ending ']' is
            // seen, so skip out early. Down below we terminate the attribute
            // span when we see the ']'.
            TokenKind::Pound => {
                match lookahead {
                    // Case 1: #![inner_attribute]
                    Some(TokenKind::Bang) => {
                        let _not = self.tokens.next().unwrap();
                        if let Some(TokenKind::OpenBracket) = self.peek() {
                            self.in_attribute = true;
                            sink(Highlight::EnterSpan { class: Class::Attribute });
                        }
                        sink(Highlight::Token { text: "#", class: Class::None });
                        sink(Highlight::Token { text: "!", class: Class::None });
                        return;
                    }
                    // Case 2: #[outer_attribute]
                    Some(TokenKind::OpenBracket) => {
                        self.in_attribute = true;
                        sink(Highlight::EnterSpan { class: Class::Attribute });
                    }
                    _ => (),
                }
                Class::None
            }
            TokenKind::CloseBracket => {
                if self.in_attribute {
                    self.in_attribute = false;
                    sink(Highlight::Token { text: "]", class: Class::None });
                    sink(Highlight::ExitSpan);
                    return;
                }
                Class::None
            }
            TokenKind::Literal { kind, .. } => match kind {
                // Text literals.
                LiteralKind::Byte { .. }
                | LiteralKind::Char { .. }
                | LiteralKind::Str { .. }
                | LiteralKind::ByteStr { .. }
                | LiteralKind::RawStr { .. }
                | LiteralKind::RawByteStr { .. } => Class::String,
                // Number literals.
                LiteralKind::Float { .. } | LiteralKind::Int { .. } => Class::Number,
            },
            TokenKind::Ident | TokenKind::RawIdent if lookahead == Some(TokenKind::Bang) => {
                self.in_macro = true;
                Class::Macro
            }
            TokenKind::Ident => match text {
                "ref" | "mut" => Class::RefKeyWord,
                "self" | "Self" => Class::Self_,
                "false" | "true" => Class::Bool,
                "Option" | "Result" => Class::PreludeTy,
                "Some" | "None" | "Ok" | "Err" => Class::PreludeVal,
                // Keywords are also included in the identifier set.
                _ if Ident::from_str(text).is_reserved() => Class::KeyWord,
                _ if self.in_macro_nonterminal => {
                    self.in_macro_nonterminal = false;
                    Class::MacroNonTerminal
                }
                _ => Class::Ident,
            },
            TokenKind::RawIdent => Class::Ident,
            TokenKind::Lifetime { .. } => Class::Lifetime,
        };
        // Anything that didn't return above is the simple case where we the
        // class just spans a single token, so we can use the `string` method.
        sink(Highlight::Token { text, class });
    }

    fn peek(&mut self) -> Option<TokenKind> {
        self.tokens.peek().map(|(toke_kind, _text)| *toke_kind)
    }
}

/// Called when we start processing a span of text that should be highlighted.
/// The `Class` argument specifies how it should be highlighted.
fn enter_span(out: &mut String, klass: Class) {
    write!(out, "", klass.as_html()).unwrap()
}

/// Called at the end of a span of highlighted text.
fn exit_span(out: &mut String) {
    write!(out, "").unwrap()
}

/// Called for a span of text. If the text should be highlighted differently
/// from the surrounding text, then the `Class` argument will be a value other
/// than `None`.
///
/// The following sequences of callbacks are equivalent:
/// ```plain
///     enter_span(Foo), string("text", None), exit_span()
///     string("text", Foo)
/// ```
/// The latter can be thought of as a shorthand for the former, which is more
/// flexible.
fn string<T: Display>(out: &mut String, text: T, klass: Class) {
    match klass {
        Class::None => write!(out, "{}", text).unwrap(),
        klass => write!(out, "{}", klass.as_html(), text).unwrap(),
    }
}

#[cfg(test)]
mod tests;
<commit_msg>Get rid of `Class::None`<commit_after>//! Basic syntax highlighting functionality.
//!
//! This module uses librustc_ast's lexer to provide token-based highlighting for
//! the HTML documentation generated by rustdoc.
//!
//! Use the `render_with_highlighting` to highlight some rust code.

use crate::html::escape::Escape;

use std::fmt::{Display, Write};
use std::iter::Peekable;

use rustc_lexer::{LiteralKind, TokenKind};
use rustc_span::symbol::Ident;
use rustc_span::with_default_session_globals;

/// Highlights `src`, returning the HTML output.
pub fn render_with_highlighting(
    src: String,
    class: Option<&str>,
    playground_button: Option<&str>,
    tooltip: Option<(&str, &str)>,
) -> String {
    debug!("highlighting: ================\n{}\n==============", src);
    let mut out = String::with_capacity(src.len());
    if let Some((tooltip, class)) = tooltip {
        write!(
            out,
            "

ⓘ{}

",
            class, tooltip
        )
        .unwrap();
    }

    write_header(&mut out, class);
    write_code(&mut out, &src);
    write_footer(&mut out, playground_button);

    out
}

fn write_header(out: &mut String, class: Option<&str>) {
    write!(out, "

\n", class.unwrap_or_default())
        .unwrap()
}

fn write_code(out: &mut String, src: &str) {
    Classifier::new(src).highlight(&mut |highlight| {
        match highlight {
            Highlight::Token { text, class } => string(out, Escape(text), class),
            Highlight::EnterSpan { class } => enter_span(out, class),
            Highlight::ExitSpan => exit_span(out),
        };
    });
}

fn write_footer(out: &mut String, playground_button: Option<&str>) {
    write!(out, "
{}
\n", playground_button.unwrap_or_default()).unwrap()
}

/// How a span of text is classified. Mostly corresponds to token kinds.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Class {
    Comment,
    DocComment,
    Attribute,
    KeyWord,
    // Keywords that do pointer/reference stuff.
    RefKeyWord,
    Self_,
    Op,
    Macro,
    MacroNonTerminal,
    String,
    Number,
    Bool,
    Ident,
    Lifetime,
    PreludeTy,
    PreludeVal,
    QuestionMark,
}

impl Class {
    /// Returns the css class expected by rustdoc for each `Class`.
    fn as_html(self) -> &'static str {
        match self {
            Class::Comment => "comment",
            Class::DocComment => "doccomment",
            Class::Attribute => "attribute",
            Class::KeyWord => "kw",
            Class::RefKeyWord => "kw-2",
            Class::Self_ => "self",
            Class::Op => "op",
            Class::Macro => "macro",
            Class::MacroNonTerminal => "macro-nonterminal",
            Class::String => "string",
            Class::Number => "number",
            Class::Bool => "bool-val",
            Class::Ident => "ident",
            Class::Lifetime => "lifetime",
            Class::PreludeTy => "prelude-ty",
            Class::PreludeVal => "prelude-val",
            Class::QuestionMark => "question-mark",
        }
    }
}

enum Highlight<'a> {
    Token { text: &'a str, class: Option<Class> },
    EnterSpan { class: Class },
    ExitSpan,
}

struct TokenIter<'a> {
    src: &'a str,
}

impl Iterator for TokenIter<'a> {
    type Item = (TokenKind, &'a str);
    fn next(&mut self) -> Option<(TokenKind, &'a str)> {
        if self.src.is_empty() {
            return None;
        }
        let token = rustc_lexer::first_token(self.src);
        let (text, rest) = self.src.split_at(token.len);
        self.src = rest;
        Some((token.kind, text))
    }
}

/// Processes program tokens, classifying strings of text by highlighting
/// category (`Class`).
struct Classifier<'a> {
    tokens: Peekable<TokenIter<'a>>,
    in_attribute: bool,
    in_macro: bool,
    in_macro_nonterminal: bool,
}

impl<'a> Classifier<'a> {
    fn new(src: &str) -> Classifier<'_> {
        let tokens = TokenIter { src }.peekable();
        Classifier { tokens, in_attribute: false, in_macro: false, in_macro_nonterminal: false }
    }

    /// Exhausts the `Classifier` writing the output into `sink`.
    ///
    /// The general structure for this method is to iterate over each token,
    /// possibly giving it an HTML span with a class specifying what flavor of
    /// token is used.
    fn highlight(mut self, sink: &mut dyn FnMut(Highlight<'a>)) {
        with_default_session_globals(|| {
            while let Some((token, text)) = self.tokens.next() {
                self.advance(token, text, sink);
            }
        })
    }

    /// Single step of highlighting. This will classify `token`, but maybe also
    /// a couple of following ones as well.
    fn advance(&mut self, token: TokenKind, text: &'a str, sink: &mut dyn FnMut(Highlight<'a>)) {
        let lookahead = self.peek();
        let no_highlight = |sink: &mut dyn FnMut(_)| sink(Highlight::Token { text, class: None });
        let class = match token {
            TokenKind::Whitespace => return no_highlight(sink),
            TokenKind::LineComment { doc_style } | TokenKind::BlockComment { doc_style, .. } => {
                if doc_style.is_some() {
                    Class::DocComment
                } else {
                    Class::Comment
                }
            }
            // Consider this as part of a macro invocation if there was a
            // leading identifier.
            TokenKind::Bang if self.in_macro => {
                self.in_macro = false;
                Class::Macro
            }

            // Assume that '&' or '*' is the reference or dereference operator
            // or a reference or pointer type. Unless, of course, it looks like
            // a logical and or a multiplication operator: `&&` or `* `.
            TokenKind::Star => match lookahead {
                Some(TokenKind::Whitespace) => Class::Op,
                _ => Class::RefKeyWord,
            },
            TokenKind::And => match lookahead {
                Some(TokenKind::And) => {
                    let _and = self.tokens.next();
                    sink(Highlight::Token { text: "&&", class: Some(Class::Op) });
                    return;
                }
                Some(TokenKind::Eq) => {
                    let _eq = self.tokens.next();
                    sink(Highlight::Token { text: "&=", class: Some(Class::Op) });
                    return;
                }
                Some(TokenKind::Whitespace) => Class::Op,
                _ => Class::RefKeyWord,
            },

            // Operators.
            TokenKind::Minus
            | TokenKind::Plus
            | TokenKind::Or
            | TokenKind::Slash
            | TokenKind::Caret
            | TokenKind::Percent
            | TokenKind::Bang
            | TokenKind::Eq
            | TokenKind::Lt
            | TokenKind::Gt => Class::Op,

            // Miscellaneous, no highlighting.
            TokenKind::Dot
            | TokenKind::Semi
            | TokenKind::Comma
            | TokenKind::OpenParen
            | TokenKind::CloseParen
            | TokenKind::OpenBrace
            | TokenKind::CloseBrace
            | TokenKind::OpenBracket
            | TokenKind::At
            | TokenKind::Tilde
            | TokenKind::Colon
            | TokenKind::Unknown => return no_highlight(sink),

            TokenKind::Question => Class::QuestionMark,

            TokenKind::Dollar => match lookahead {
                Some(TokenKind::Ident) => {
                    self.in_macro_nonterminal = true;
                    Class::MacroNonTerminal
                }
                _ => return no_highlight(sink),
            },

            // This might be the start of an attribute. We're going to want to
            // continue highlighting it as an attribute until the ending ']' is
            // seen, so skip out early. Down below we terminate the attribute
            // span when we see the ']'.
            TokenKind::Pound => {
                match lookahead {
                    // Case 1: #![inner_attribute]
                    Some(TokenKind::Bang) => {
                        let _not = self.tokens.next().unwrap();
                        if let Some(TokenKind::OpenBracket) = self.peek() {
                            self.in_attribute = true;
                            sink(Highlight::EnterSpan { class: Class::Attribute });
                        }
                        sink(Highlight::Token { text: "#", class: None });
                        sink(Highlight::Token { text: "!", class: None });
                        return;
                    }
                    // Case 2: #[outer_attribute]
                    Some(TokenKind::OpenBracket) => {
                        self.in_attribute = true;
                        sink(Highlight::EnterSpan { class: Class::Attribute });
                    }
                    _ => (),
                }
                return no_highlight(sink);
            }
            TokenKind::CloseBracket => {
                if self.in_attribute {
                    self.in_attribute = false;
                    sink(Highlight::Token { text: "]", class: None });
                    sink(Highlight::ExitSpan);
                    return;
                }
                return no_highlight(sink);
            }
            TokenKind::Literal { kind, .. } => match kind {
                // Text literals.
                LiteralKind::Byte { .. }
                | LiteralKind::Char { .. }
                | LiteralKind::Str { .. }
                | LiteralKind::ByteStr { .. }
                | LiteralKind::RawStr { .. }
                | LiteralKind::RawByteStr { .. } => Class::String,
                // Number literals.
                LiteralKind::Float { .. } | LiteralKind::Int { .. } => Class::Number,
            },
            TokenKind::Ident | TokenKind::RawIdent if lookahead == Some(TokenKind::Bang) => {
                self.in_macro = true;
                Class::Macro
            }
            TokenKind::Ident => match text {
                "ref" | "mut" => Class::RefKeyWord,
                "self" | "Self" => Class::Self_,
                "false" | "true" => Class::Bool,
                "Option" | "Result" => Class::PreludeTy,
                "Some" | "None" | "Ok" | "Err" => Class::PreludeVal,
                // Keywords are also included in the identifier set.
                _ if Ident::from_str(text).is_reserved() => Class::KeyWord,
                _ if self.in_macro_nonterminal => {
                    self.in_macro_nonterminal = false;
                    Class::MacroNonTerminal
                }
                _ => Class::Ident,
            },
            TokenKind::RawIdent => Class::Ident,
            TokenKind::Lifetime { .. } => Class::Lifetime,
        };
        // Anything that didn't return above is the simple case where we the
        // class just spans a single token, so we can use the `string` method.
        sink(Highlight::Token { text, class: Some(class) });
    }

    fn peek(&mut self) -> Option<TokenKind> {
        self.tokens.peek().map(|(toke_kind, _text)| *toke_kind)
    }
}

/// Called when we start processing a span of text that should be highlighted.
/// The `Class` argument specifies how it should be highlighted.
fn enter_span(out: &mut String, klass: Class) {
    write!(out, "", klass.as_html()).unwrap()
}

/// Called at the end of a span of highlighted text.
fn exit_span(out: &mut String) {
    write!(out, "").unwrap()
}

/// Called for a span of text. If the text should be highlighted differently
/// from the surrounding text, then the `Class` argument will be a value other
/// than `None`.
///
/// The following sequences of callbacks are equivalent:
/// ```plain
///     enter_span(Foo), string("text", None), exit_span()
///     string("text", Foo)
/// ```
/// The latter can be thought of as a shorthand for the former, which is more
/// flexible.
fn string<T: Display>(out: &mut String, text: T, klass: Option<Class>) {
    match klass {
        None => write!(out, "{}", text).unwrap(),
        Some(klass) => write!(out, "{}", klass.as_html(), text).unwrap(),
    }
}

#[cfg(test)]
mod tests;
<|endoftext|>


<commit_before>// Copyright 2018 Serde Developers
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#[macro_use]
extern crate serde_derive;

#[derive(Serialize)] //~ ERROR: proc-macro derive panicked
//~^ HELP: unknown serde variant attribute `flatten`
enum Foo {
    #[serde(flatten)]
    Foo {
        x: u32,
    }
}

fn main() {}
<commit_msg>Remove test that fails to parse flatten attribute<commit_after><|endoftext|>


<commit_before>//! A TCP RelayServer that intercepts messages between servers inside an
//! integration test. Right now it just allows us to simulate nodes going offline
//! but we could extend it to expose information about which messages have been sent
//! for testing asserts.
//!
//! #Examples
//!
//! ```
//! let relay_server = RelayServer::new_with_random_addresses(NUM_ADDRESSES);
//! let addrs = relay_server.get_bound_addresses();
//! for addr in addrs {
//!     // start up a new Raft server (can pass in all bound addrs as the addrs for the other
//!     // servers)
//!     # let raft_server_addr = "127.0.0.1:0";
//!     relay_server.relay_address(addr, raft_server_addr);
//! }
//! // relay server is now set up and forwarding messages between raft servers.
//! ```
//!

use std::collections::HashMap;
use std::collections::hash_map::Keys;
use std::net::{SocketAddr, TcpListener, TcpStream, Shutdown};
use std::thread;
use std::sync::{Arc, Mutex};
use std::io::{Read, Write};

pub struct RelayServer {
    addrs: Arc<Mutex<HashMap<SocketAddr, SocketInfo>>>
}

struct SocketInfo {
    to_addr: Option<SocketAddr>,
    online: bool
}

impl RelayServer {
    pub fn new() -> RelayServer {
        RelayServer {addrs: Arc::new(Mutex::new(HashMap::new()))}
    }

    pub fn new_with_random_addresses(num_addresses: u64) -> RelayServer {
        let mut server = RelayServer::new();
        server.bind_random_addresses(num_addresses);
        server
    }

    /// Binds to num_addresses addresses. Ports are assigned by the OS.
    /// Does not initially route their messages anywhere
    ///
    /// #Panics
    /// Panics if it is unable to acquire enough bound TCP ports.
    /// Or if one of the server's background threads have panicked.
    pub fn bind_random_addresses(&mut self, num_addresses: u64) {
        let mut addrs = self.addrs.lock().unwrap();
        for i in 0..num_addresses {
            let listener = TcpListener::bind("127.0.0.1:0").unwrap();
            let socket_info = SocketInfo {to_addr: None, online: false};
            addrs.insert(listener.local_addr().unwrap(), socket_info);

            let addrs_clone = self.addrs.clone();
            // spawn a background thread to listen on this port
            thread::spawn (move || RelayServer::relay_messages(listener, addrs_clone));
        }
    }

    /// Returns a vector containing the bound socket addrs
    pub fn get_bound_addresses(&self) -> Vec<SocketAddr> {
        let addrs = self.addrs.lock().unwrap();
        addrs.keys()
             .cloned()
             .take(addrs.len())
             .collect::<Vec<SocketAddr>>()
    }

    /// Sets up a new mapping from from_address to to_address.
    /// Initially the mapping is online. Use set_address_active to change this
    ///
    /// #Panics
    /// Panics if the from_address is not an existing mapping
    /// Also panics if any of the server's background threads have panicked.
    pub fn relay_address(&mut self, from_address: SocketAddr, to_address: SocketAddr) {
        let mut addrs = self.addrs.lock().unwrap();
        let socket_info: &mut SocketInfo = addrs.get_mut(&from_address).unwrap();

        socket_info.to_addr = Some(to_address);
        socket_info.online = true;
    }

    /// Toggles a mapping online or offline.
    ///
    /// #Panics
    /// Panics if the from_address is not an existing mapping
    /// Also panics if any of the server's background threads have panicked.
    pub fn set_address_active(&mut self, from_address: SocketAddr, online: bool) {
        let mut addrs = self.addrs.lock().unwrap();
        addrs.get_mut(&from_address).unwrap().online = true;
    }

    /// Runs in a background thread, and relays messages according to the hash map.
    fn relay_messages(listener: TcpListener, addrs: Arc<Mutex<HashMap<SocketAddr, SocketInfo>>>) {
        let local_addr = listener.local_addr().unwrap();
        for stream in listener.incoming() {
            match stream {
                Ok (mut stream) => {
                    let addrs = addrs.lock().unwrap();
                    match addrs.get(&local_addr) {
                        Some(info) => {
                            if info.online {
                                RelayServer::relay_stream_to_addr(stream, info.to_addr.unwrap());
                            }
                        },
                        None => {}
                    }
                },
                Err (e) => {
                    // conection errors are dropped
                }
            }
        }
    }

    /// Relays the given stream to the to_addr and sends the result back to the stream.
    fn relay_stream_to_addr(mut stream: TcpStream, to_addr: SocketAddr) {
        // Open a connection to to_addr
        let mut outgoing = TcpStream::connect(to_addr).unwrap();

        let mut outgoing_clone = outgoing.try_clone().unwrap();
        let mut stream_clone = stream.try_clone().unwrap();
        thread::spawn (move || RelayServer::relay_stream_to_stream(outgoing_clone, stream_clone));

        RelayServer::relay_stream_to_stream(stream, outgoing);
    }

    /// Relays messages from from_stream to to_stream until
    /// EOF is reached on from_stream. At this point we send
    /// an EOF on to_stream
    ///
    /// #Panics
    /// Panics if the from_stream cannot be read from or the to_stream can
    /// not be written to
    fn relay_stream_to_stream(mut from_stream: TcpStream, mut to_stream: TcpStream) {
        let mut buf = [0; 100];

        loop {
            let bytes_read = from_stream.read(&mut buf[..]).unwrap();
            if bytes_read == 0 {
                to_stream.shutdown(Shutdown::Write);
                break;
            }
            to_stream.write(&buf[0..bytes_read]);
            to_stream.flush();
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate rand;
    use super::*;
    use std::thread;
    use std::sync::mpsc::{channel, Receiver};
    use std::net::{SocketAddr, TcpListener, TcpStream};
    use std::vec;
    use std::io::{Read, Write};
    use self::rand::{thread_rng, Rng};
    use std::time::Duration;

    // Timeout value, in ms, when waiting for local network messages
    const TIMEOUT: u64 = 5000;

    #[test]
    /// Tests that the relay server will relay from A->B
    fn it_relays_in_one_direction() {
        const MESSAGE_LENGTH: usize = 1057;

        let mut relay_server = RelayServer::new_with_random_addresses(1);
        let addresses = relay_server.get_bound_addresses();
        assert_eq!(addresses.len(), 1);

        let (addr, rx) = start_tcp_server();
        relay_server.relay_address(addresses[0], addr);

        let msg: String = thread_rng().gen_ascii_chars().take(MESSAGE_LENGTH).collect();
        {
            let mut client = TcpStream::connect(addresses[0]).unwrap();
            // generate a random message string
            assert_eq!(msg.as_bytes().len(), MESSAGE_LENGTH);
            client.write_all(msg.as_bytes()).unwrap();
            client.flush().unwrap();
        }

        let result = rx.recv_timeout(Duration::from_millis(TIMEOUT)).unwrap();
        assert_eq!(msg.as_bytes(), &result[..]);
    }

    /// Starts a simple TCP server on the returned address (OS assigned) in a background thread
    /// that sends every message it recieves down through the returned channel
    fn start_tcp_server() -> (SocketAddr, Receiver<Vec<u8>>) {
        let (tx, rx) = channel();
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        thread::spawn(move || {
            for stream in listener.incoming() {
                match stream {
                    Ok(mut s) => {
                        let mut v = Vec::new();
                        s.read_to_end(&mut v);
                        tx.send(v).unwrap();
                    },
                    Err(e) => {}
                }
            }
        });
        (addr, rx)
    }
}
<commit_msg>test(RelayServer): Add relay server echoing tests<commit_after>//! A TCP RelayServer that intercepts messages between servers inside an
//! integration test. Right now it just allows us to simulate nodes going offline
//! but we could extend it to expose information about which messages have been sent
//! for testing asserts.
//!
//! #Examples
//!
//! ```
//! let relay_server = RelayServer::new_with_random_addresses(NUM_ADDRESSES);
//! let addrs = relay_server.get_bound_addresses();
//! for addr in addrs {
//!     // start up a new Raft server (can pass in all bound addrs as the addrs for the other
//!     // servers)
//!     # let raft_server_addr = "127.0.0.1:0";
//!     relay_server.relay_address(addr, raft_server_addr);
//! }
//! // relay server is now set up and forwarding messages between raft servers.
//! ```
//!

use std::collections::HashMap;
use std::net::{SocketAddr, TcpListener, TcpStream, Shutdown};
use std::thread;
use std::sync::{Arc, Mutex};
use std::io::{Read, Write};

pub struct RelayServer {
    addrs: Arc<Mutex<HashMap<SocketAddr, SocketInfo>>>
}

struct SocketInfo {
    to_addr: Option<SocketAddr>,
    online: bool
}

impl RelayServer {
    pub fn new() -> RelayServer {
        RelayServer {addrs: Arc::new(Mutex::new(HashMap::new()))}
    }

    pub fn new_with_random_addresses(num_addresses: u64) -> RelayServer {
        let mut server = RelayServer::new();
        server.bind_random_addresses(num_addresses);
        server
    }

    /// Binds to num_addresses addresses. Ports are assigned by the OS.
    /// Does not initially route their messages anywhere
    ///
    /// #Panics
    /// Panics if it is unable to acquire enough bound TCP ports.
    /// Or if one of the server's background threads have panicked.
    pub fn bind_random_addresses(&mut self, num_addresses: u64) {
        let mut addrs = self.addrs.lock().unwrap();
        for _ in 0..num_addresses {
            let listener = TcpListener::bind("127.0.0.1:0").unwrap();
            let socket_info = SocketInfo {to_addr: None, online: false};
            addrs.insert(listener.local_addr().unwrap(), socket_info);

            let addrs_clone = self.addrs.clone();
            // spawn a background thread to listen on this port
            thread::spawn (move || RelayServer::relay_messages(listener, addrs_clone));
        }
    }

    /// Returns a vector containing the bound socket addrs
    pub fn get_bound_addresses(&self) -> Vec<SocketAddr> {
        let addrs = self.addrs.lock().unwrap();
        addrs.keys()
             .cloned()
             .take(addrs.len())
             .collect::<Vec<SocketAddr>>()
    }

    /// Sets up a new mapping from from_address to to_address.
    /// Initially the mapping is online. Use set_address_active to change this
    ///
    /// #Panics
    /// Panics if the from_address is not an existing mapping
    /// Also panics if any of the server's background threads have panicked.
    pub fn relay_address(&mut self, from_address: SocketAddr, to_address: SocketAddr) {
        let mut addrs = self.addrs.lock().unwrap();
        let socket_info: &mut SocketInfo = addrs.get_mut(&from_address).unwrap();

        socket_info.to_addr = Some(to_address);
        socket_info.online = true;
    }

    /// Toggles a mapping online or offline.
    ///
    /// #Panics
    /// Panics if the from_address is not an existing mapping
    /// Also panics if any of the server's background threads have panicked.
    pub fn set_address_active(&mut self, from_address: SocketAddr, online: bool) {
        let mut addrs = self.addrs.lock().unwrap();
        addrs.get_mut(&from_address).unwrap().online = online;
    }

    /// Runs in a background thread, and relays messages according to the hash map.
    fn relay_messages(listener: TcpListener, addrs: Arc<Mutex<HashMap<SocketAddr, SocketInfo>>>) {
        let local_addr = listener.local_addr().unwrap();
        for stream in listener.incoming() {
            match stream {
                Ok (stream) => {
                    let addrs = addrs.lock().unwrap();
                    match addrs.get(&local_addr) {
                        Some(info) => {
                            if info.online {
                                RelayServer::relay_stream_to_addr(stream, info.to_addr.unwrap());
                            }
                        },
                        None => {}
                    }
                },
                Err (_) => {
                    // conection errors are dropped. We might want to panic instead?
                }
            }
        }
    }

    /// Relays the given stream to the to_addr and sends the result back to the stream.
    fn relay_stream_to_addr(stream: TcpStream, to_addr: SocketAddr) {
        // Open a connection to to_addr
        let outgoing = TcpStream::connect(to_addr).unwrap();

        let outgoing_clone = outgoing.try_clone().unwrap();
        let stream_clone = stream.try_clone().unwrap();
        thread::spawn (move || RelayServer::relay_stream_to_stream(outgoing_clone, stream_clone));

        RelayServer::relay_stream_to_stream(stream, outgoing);
    }

    /// Relays messages from from_stream to to_stream until
    /// EOF is reached on from_stream. At this point we send
    /// an EOF on to_stream
    ///
    /// #Panics
    /// Panics if the from_stream cannot be read from or the to_stream can
    /// not be written to
    fn relay_stream_to_stream(mut from_stream: TcpStream, mut to_stream: TcpStream) {
        let mut buf = [0; 100];

        loop {
            let bytes_read = from_stream.read(&mut buf[..]).unwrap();
            if bytes_read == 0 {
                to_stream.shutdown(Shutdown::Write).unwrap();
                break;
            }
            to_stream.write(&buf[0..bytes_read]).unwrap();
            to_stream.flush().unwrap();
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate rand;
    use super::*;
    use std::thread;
    use std::sync::mpsc::{channel, Receiver};
    use std::net::{SocketAddr, TcpListener, TcpStream, Shutdown};
    use std::io::{Read, Write};
    use self::rand::{thread_rng, Rng};
    use std::time::Duration;

    // Timeout value, in ms, when waiting for local network messages
    const TIMEOUT: u64 = 5000;

    #[test]
    /// Tests that the relay server will relay from A->B
    fn it_relays_from_client() {
        const MESSAGE_LENGTH: usize = 1057;

        let mut relay_server = RelayServer::new_with_random_addresses(1);
        let addresses = relay_server.get_bound_addresses();
        assert_eq!(addresses.len(), 1);

        let (addr, rx) = start_tcp_listening_server();
        relay_server.relay_address(addresses[0], addr);

        let msg: String = thread_rng().gen_ascii_chars().take(MESSAGE_LENGTH).collect();
        {
            let mut client = TcpStream::connect(addresses[0]).unwrap();
            // generate a random message string
            assert_eq!(msg.as_bytes().len(), MESSAGE_LENGTH);
            client.write_all(msg.as_bytes()).unwrap();
            client.flush().unwrap();
        }

        let result = rx.recv_timeout(Duration::from_millis(TIMEOUT)).unwrap();
        assert_eq!(msg.as_bytes(), &result[..]);
    }

    #[test]
    fn it_relays_from_server() {
        const MESSAGE_LENGTH: usize = 904;

        let mut relay_server = RelayServer::new_with_random_addresses(1);
        let addresses = relay_server.get_bound_addresses();
        assert_eq!(addresses.len(), 1);

        let msg: String = thread_rng().gen_ascii_chars().take(MESSAGE_LENGTH).collect();
        let addr = start_tcp_sending_server(msg.clone());
        relay_server.relay_address(addresses[0], addr);

        let mut client = TcpStream::connect(addresses[0]).unwrap();
        let mut v = Vec::new();
        client.read_to_end(&mut v).unwrap();
        assert_eq!(&v[..], msg.as_bytes());
    }

    #[test]
    fn it_relays_from_client_and_server() {
        const MESSAGE_LENGTH: usize = 904;

        let mut relay_server = RelayServer::new_with_random_addresses(1);
        let addresses = relay_server.get_bound_addresses();
        assert_eq!(addresses.len(), 1);

        let msg: String = thread_rng().gen_ascii_chars().take(MESSAGE_LENGTH).collect();
        let addr = start_tcp_echo_server();
        relay_server.relay_address(addresses[0], addr);

        let mut client = TcpStream::connect(addresses[0]).unwrap();
        client.write_all(msg.as_bytes()).unwrap();
        client.flush().unwrap();
        client.shutdown(Shutdown::Write).unwrap();

        let mut v = Vec::new();
        client.read_to_end(&mut v).unwrap();

        assert_eq!(&v[..], msg.as_bytes());
    }

    /// Starts a simple TCP server on the returned address (OS assigned) in a background thread
    /// that sends every message it recieves down through the returned channel
    fn start_tcp_listening_server() -> (SocketAddr, Receiver<Vec<u8>>) {
        let (tx, rx) = channel();
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        thread::spawn(move || {
            for stream in listener.incoming() {
                let mut s = stream.unwrap();
                let mut v = Vec::new();
                s.read_to_end(&mut v).unwrap();
                tx.send(v).unwrap();
            }
        });
        (addr, rx)
    }

    /// Starts a TCP server that sends msg to each client when they connect
    /// and then closes the connection
    fn start_tcp_sending_server(msg: String) -> SocketAddr {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        thread::spawn(move || {
            for stream in listener.incoming() {
                let mut s = stream.unwrap();
                s.write_all(msg.as_bytes()).unwrap();
                s.flush().unwrap();
            }
        });
        addr
    }

    /// Starts a TCP echo server and returns the address assigned to the server by the OS
    fn start_tcp_echo_server() -> SocketAddr {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        thread::spawn(move || {
            for stream in listener.incoming() {
                let mut s = stream.unwrap();
                let mut v = Vec::new();
                s.read_to_end(&mut v).unwrap();
                s.write_all(&v[..]).unwrap();
                s.flush().unwrap();
            }
        });
        addr
    }
}
<|endoftext|>