<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

#[cfg(target_arch = "x86")]
use std::arch::x86::{CpuidResult, __cpuid_count, __get_cpuid_max};
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::{CpuidResult, __cpuid_count, __get_cpuid_max};

#[cfg(target_arch = "x86_64")]
use kvm_bindings::CpuId;

use crate::bit_helper::BitHelper;
use crate::cpu_leaf::*;

/// Intel brand string.
pub const VENDOR_ID_INTEL: &[u8; 12] = b"GenuineIntel";
/// AMD brand string.
pub const VENDOR_ID_AMD: &[u8; 12] = b"AuthenticAMD";

/// cpuid related error.
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
    /// The function was called with invalid parameters.
    #[error("The function was called with invalid parameters.")]
    InvalidParameters(String),
    /// Function not supported on the current architecture.
    #[error("Function not supported on the current architecture.")]
    NotSupported,
}

/// Extract entry from the cpuid.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn get_cpuid(function: u32, count: u32) -> Result<CpuidResult, Error> {
    // TODO: replace with validation based on `has_cpuid()` when it becomes stable:
    //  https://doc.rust-lang.org/core/arch/x86/fn.has_cpuid.html
    #[cfg(target_env = "sgx")]
    {
        return Err(Error::NotSupported);
    }
    // For x86 the host supports the `cpuid` instruction if SSE is enabled. Otherwise it's hard to
    // check.
    #[cfg(target_arch = "x86")]
    {
        #[cfg(not(target_feature = "sse"))]
        {
            return Err(Error::NotSupported);
        }
    }

    // SAFETY: this is safe because the host supports the `cpuid` instruction
    let max_function = unsafe { __get_cpuid_max(function & leaf_0x80000000::LEAF_NUM).0 };
    if function > max_function {
        return Err(Error::InvalidParameters(format!(
            "Function not supported: 0x{:x}",
            function
        )));
    }

    // SAFETY: this is safe because the host supports the `cpuid` instruction
    let entry = unsafe { __cpuid_count(function, count) };
    if entry.eax == 0 && entry.ebx == 0 && entry.ecx == 0 && entry.edx == 0 {
        return Err(Error::InvalidParameters(format!(
            "Invalid count: {}",
            count
        )));
    }

    Ok(entry)
}

/// Extracts the CPU vendor id from leaf 0x0.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn get_vendor_id_from_host() -> Result<[u8; 12], Error> {
    // SAFETY: This is safe because the resulting type has a lower alignment requirement
    get_cpuid(0, 0).map(|vendor_entry| unsafe {
        std::mem::transmute::<[u32; 3], [u8; 12]>([
            vendor_entry.ebx,
            vendor_entry.edx,
            vendor_entry.ecx,
        ])
    })
}

/// Extracts the CPU vendor id from leaf 0x0.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn get_vendor_id_from_cpuid(cpuid: &CpuId) -> Result<[u8; 12], Error> {
    // Search for vendor id entry.
    for entry in cpuid.as_slice().iter() {
        if entry.function == 0 && entry.index == 0 {
            let cpu_vendor_id: [u8; 12] =
            // SAFETY: This is safe because the resulting type has a lower alignment requirement
                unsafe { std::mem::transmute([entry.ebx, entry.edx, entry.ecx]) };
            return Ok(cpu_vendor_id);
        }
    }

    Err(Error::NotSupported)
}

/// Validates that the provided CPUID belongs to a CPU of the same
/// model as the host's.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn is_same_model(cpuid: &CpuId) -> bool {
    // Try to get the vendor IDs from the host and the CPUID struct.
    if let (Ok(host_vendor_id), Ok(cpuid_vendor_id)) =
        (get_vendor_id_from_host(), get_vendor_id_from_cpuid(cpuid))
    {
        // If the vendor IDs aren't the same, the CPUs are not identical.
        if host_vendor_id != cpuid_vendor_id {
            return false;
        }
    } else {
        // This only fails when CPUID is not supported, in which case
        // we can't tell if the CPUs are identical.
        return false;
    }

    // Try to get the feature information leaf from the host CPUID.
    let host_feature_info_leaf = get_cpuid(leaf_0x1::LEAF_NUM, 0);

    // The relevant information for this comparison is in the EAX register.
    let host_feature_info_leaf_eax = match host_feature_info_leaf {
        Ok(leaf) => leaf.eax,
        Err(_) => {
            // If this fails, we can't tell if the CPUs are identical.
            return false;
        }
    };

    // Search for the entry for leaf0x1.
    let feature_info_leaf = cpuid
        .as_slice()
        .iter()
        .find(|entry| entry.function == leaf_0x1::LEAF_NUM);

    // The relevant information is in EAX.
    let feature_info_leaf_eax = match feature_info_leaf {
        Some(leaf) => leaf.eax,
        None => {
            // Fail fast if we can't retrieve the relevant
            // information from CPUID.
            return false;
        }
    };

    // Validate that all of these properties are the same.
    for elem in &[
        leaf_0x1::eax::EXTENDED_FAMILY_ID_BITRANGE,
        leaf_0x1::eax::EXTENDED_PROCESSOR_MODEL_BITRANGE,
        leaf_0x1::eax::PROCESSOR_FAMILY_BITRANGE,
        leaf_0x1::eax::PROCESSOR_MODEL_BITRANGE,
        leaf_0x1::eax::STEPPING_BITRANGE,
    ] {
        if feature_info_leaf_eax.read_bits_in_range(elem)
            != host_feature_info_leaf_eax.read_bits_in_range(elem)
        {
            return false;
        }
    }

    true
}

/// Scans through the CPUID and determines if a feature bit is set.
// TODO: This currently involves a linear search which would be improved
//       when we'll refactor the cpuid crate.
#[macro_export]
macro_rules! cpuid_is_feature_set {
    ($cpuid:ident, $leaf:expr, $index:expr, $reg:tt, $feature_bit:expr) => {{
        let mut res = false;
        for entry in $cpuid.as_slice().iter() {
            if entry.function == $leaf && entry.index == $index {
                if entry.$reg & (1 << $feature_bit) != 0 {
                    res = true;
                    break;
                }
            }
        }
        res
    }};
}

#[cfg(test)]
pub mod tests {
    use crate::common::*;

    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    pub fn get_topoext_fn() -> u32 {
        let vendor_id = get_vendor_id_from_host();
        assert!(vendor_id.is_ok());
        let function = match &vendor_id.ok().unwrap() {
            VENDOR_ID_INTEL => leaf_0x4::LEAF_NUM,
            VENDOR_ID_AMD => leaf_0x8000001d::LEAF_NUM,
            _ => 0,
        };
        assert!(function != 0);

        function
    }

    #[test]
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    fn test_get_cpu_id() {
        // get_cpu_id should work correctly here
        let topoext_fn = get_topoext_fn();

        // check that get_cpuid works for valid parameters
        match get_cpuid(topoext_fn, 0) {
            Ok(topoext_entry) => {
                assert!(topoext_entry.eax != 0);
            }
            _ => panic!("Wrong behavior"),
        }

        // check that get_cpuid returns correct error for invalid `function`
        match get_cpuid(0x9000_0000, 0) {
            Err(Error::InvalidParameters(s)) => {
                assert!(s == "Function not supported: 0x90000000");
            }
            _ => panic!("Wrong behavior"),
        }

        // check that get_cpuid returns correct error for invalid `count`
        match get_cpuid(topoext_fn, 100) {
            Err(Error::InvalidParameters(s)) => {
                assert!(s == "Invalid count: 100");
            }
            _ => panic!("Wrong behavior"),
        }
    }

    #[test]
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    fn test_get_vendor_id() {
        let vendor_id = get_vendor_id_from_host();
        assert!(vendor_id.is_ok());
        matches!(&vendor_id.ok().unwrap(), VENDOR_ID_INTEL | VENDOR_ID_AMD);
    }

    #[test]
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    fn test_is_same_model() {
        let mut curr_cpuid = CpuId::new(2).unwrap();

        // Add the vendor ID leaf.
        let vendor = get_cpuid(0x0, 0).unwrap();
        curr_cpuid.as_mut_slice()[0].function = 0x0;
        curr_cpuid.as_mut_slice()[0].index = 0;
        curr_cpuid.as_mut_slice()[0].ebx = vendor.ebx;
        curr_cpuid.as_mut_slice()[0].ecx = vendor.ecx;
        curr_cpuid.as_mut_slice()[0].edx = vendor.edx;

        // Add the feature info leaf.
        let feature_info = get_cpuid(0x1, 0).unwrap();
        curr_cpuid.as_mut_slice()[1].function = 0x1;
        curr_cpuid.as_mut_slice()[1].index = 0;
        curr_cpuid.as_mut_slice()[1].eax = feature_info.eax;

        assert!(is_same_model(&curr_cpuid));

        let mut diff_vendor_cpuid = curr_cpuid.clone();
        for mut entry in diff_vendor_cpuid.as_mut_slice() {
            if entry.function == 0x0 && entry.index == 0 {
                entry.ebx = 0xFFFF_FFFF;
            }
        }
        assert!(!is_same_model(&diff_vendor_cpuid));

        let mut diff_feature_cpuid = curr_cpuid;
        for mut entry in diff_feature_cpuid.as_mut_slice() {
            if entry.function == 0x1 && entry.index == 0 {
                entry.eax ^= 0x1;
            }
        }
        assert!(!is_same_model(&diff_feature_cpuid));
    }
}
<commit_msg>Linking `has_cpuid` todo<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

#[cfg(target_arch = "x86")]
use std::arch::x86::{CpuidResult, __cpuid_count, __get_cpuid_max};
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::{CpuidResult, __cpuid_count, __get_cpuid_max};

#[cfg(target_arch = "x86_64")]
use kvm_bindings::CpuId;

use crate::bit_helper::BitHelper;
use crate::cpu_leaf::*;

/// Intel brand string.
pub const VENDOR_ID_INTEL: &[u8; 12] = b"GenuineIntel";
/// AMD brand string.
pub const VENDOR_ID_AMD: &[u8; 12] = b"AuthenticAMD";

/// cpuid related error.
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
    /// The function was called with invalid parameters.
    #[error("The function was called with invalid parameters.")]
    InvalidParameters(String),
    /// Function not supported on the current architecture.
    #[error("Function not supported on the current architecture.")]
    NotSupported,
}

/// Extract entry from the cpuid.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn get_cpuid(function: u32, count: u32) -> Result<CpuidResult, Error> {
    // TODO: Use `core::arch::x86_64::has_cpuid`
    // (https://github.com/firecracker-microvm/firecracker/issues/3271)
    #[cfg(target_env = "sgx")]
    {
        return Err(Error::NotSupported);
    }
    // For x86 the host supports the `cpuid` instruction if SSE is enabled. Otherwise it's hard to
    // check.
    #[cfg(target_arch = "x86")]
    {
        #[cfg(not(target_feature = "sse"))]
        {
            return Err(Error::NotSupported);
        }
    }

    // SAFETY: this is safe because the host supports the `cpuid` instruction
    let max_function = unsafe { __get_cpuid_max(function & leaf_0x80000000::LEAF_NUM).0 };
    if function > max_function {
        return Err(Error::InvalidParameters(format!(
            "Function not supported: 0x{:x}",
            function
        )));
    }

    // SAFETY: this is safe because the host supports the `cpuid` instruction
    let entry = unsafe { __cpuid_count(function, count) };
    if entry.eax == 0 && entry.ebx == 0 && entry.ecx == 0 && entry.edx == 0 {
        return Err(Error::InvalidParameters(format!(
            "Invalid count: {}",
            count
        )));
    }

    Ok(entry)
}

/// Extracts the CPU vendor id from leaf 0x0.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn get_vendor_id_from_host() -> Result<[u8; 12], Error> {
    // SAFETY: This is safe because the resulting type has a lower alignment requirement
    get_cpuid(0, 0).map(|vendor_entry| unsafe {
        std::mem::transmute::<[u32; 3], [u8; 12]>([
            vendor_entry.ebx,
            vendor_entry.edx,
            vendor_entry.ecx,
        ])
    })
}

/// Extracts the CPU vendor id from leaf 0x0.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn get_vendor_id_from_cpuid(cpuid: &CpuId) -> Result<[u8; 12], Error> {
    // Search for vendor id entry.
    for entry in cpuid.as_slice().iter() {
        if entry.function == 0 && entry.index == 0 {
            let cpu_vendor_id: [u8; 12] =
            // SAFETY: This is safe because the resulting type has a lower alignment requirement
                unsafe { std::mem::transmute([entry.ebx, entry.edx, entry.ecx]) };
            return Ok(cpu_vendor_id);
        }
    }

    Err(Error::NotSupported)
}

/// Validates that the provided CPUID belongs to a CPU of the same
/// model as the host's.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn is_same_model(cpuid: &CpuId) -> bool {
    // Try to get the vendor IDs from the host and the CPUID struct.
    if let (Ok(host_vendor_id), Ok(cpuid_vendor_id)) =
        (get_vendor_id_from_host(), get_vendor_id_from_cpuid(cpuid))
    {
        // If the vendor IDs aren't the same, the CPUs are not identical.
        if host_vendor_id != cpuid_vendor_id {
            return false;
        }
    } else {
        // This only fails when CPUID is not supported, in which case
        // we can't tell if the CPUs are identical.
        return false;
    }

    // Try to get the feature information leaf from the host CPUID.
    let host_feature_info_leaf = get_cpuid(leaf_0x1::LEAF_NUM, 0);

    // The relevant information for this comparison is in the EAX register.
    let host_feature_info_leaf_eax = match host_feature_info_leaf {
        Ok(leaf) => leaf.eax,
        Err(_) => {
            // If this fails, we can't tell if the CPUs are identical.
            return false;
        }
    };

    // Search for the entry for leaf0x1.
    let feature_info_leaf = cpuid
        .as_slice()
        .iter()
        .find(|entry| entry.function == leaf_0x1::LEAF_NUM);

    // The relevant information is in EAX.
    let feature_info_leaf_eax = match feature_info_leaf {
        Some(leaf) => leaf.eax,
        None => {
            // Fail fast if we can't retrieve the relevant
            // information from CPUID.
            return false;
        }
    };

    // Validate that all of these properties are the same.
    for elem in &[
        leaf_0x1::eax::EXTENDED_FAMILY_ID_BITRANGE,
        leaf_0x1::eax::EXTENDED_PROCESSOR_MODEL_BITRANGE,
        leaf_0x1::eax::PROCESSOR_FAMILY_BITRANGE,
        leaf_0x1::eax::PROCESSOR_MODEL_BITRANGE,
        leaf_0x1::eax::STEPPING_BITRANGE,
    ] {
        if feature_info_leaf_eax.read_bits_in_range(elem)
            != host_feature_info_leaf_eax.read_bits_in_range(elem)
        {
            return false;
        }
    }

    true
}

/// Scans through the CPUID and determines if a feature bit is set.
// TODO: This currently involves a linear search which would be improved
//       when we'll refactor the cpuid crate.
#[macro_export]
macro_rules! cpuid_is_feature_set {
    ($cpuid:ident, $leaf:expr, $index:expr, $reg:tt, $feature_bit:expr) => {{
        let mut res = false;
        for entry in $cpuid.as_slice().iter() {
            if entry.function == $leaf && entry.index == $index {
                if entry.$reg & (1 << $feature_bit) != 0 {
                    res = true;
                    break;
                }
            }
        }
        res
    }};
}

#[cfg(test)]
pub mod tests {
    use crate::common::*;

    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    pub fn get_topoext_fn() -> u32 {
        let vendor_id = get_vendor_id_from_host();
        assert!(vendor_id.is_ok());
        let function = match &vendor_id.ok().unwrap() {
            VENDOR_ID_INTEL => leaf_0x4::LEAF_NUM,
            VENDOR_ID_AMD => leaf_0x8000001d::LEAF_NUM,
            _ => 0,
        };
        assert!(function != 0);

        function
    }

    #[test]
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    fn test_get_cpu_id() {
        // get_cpu_id should work correctly here
        let topoext_fn = get_topoext_fn();

        // check that get_cpuid works for valid parameters
        match get_cpuid(topoext_fn, 0) {
            Ok(topoext_entry) => {
                assert!(topoext_entry.eax != 0);
            }
            _ => panic!("Wrong behavior"),
        }

        // check that get_cpuid returns correct error for invalid `function`
        match get_cpuid(0x9000_0000, 0) {
            Err(Error::InvalidParameters(s)) => {
                assert!(s == "Function not supported: 0x90000000");
            }
            _ => panic!("Wrong behavior"),
        }

        // check that get_cpuid returns correct error for invalid `count`
        match get_cpuid(topoext_fn, 100) {
            Err(Error::InvalidParameters(s)) => {
                assert!(s == "Invalid count: 100");
            }
            _ => panic!("Wrong behavior"),
        }
    }

    #[test]
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    fn test_get_vendor_id() {
        let vendor_id = get_vendor_id_from_host();
        assert!(vendor_id.is_ok());
        matches!(&vendor_id.ok().unwrap(), VENDOR_ID_INTEL | VENDOR_ID_AMD);
    }

    #[test]
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    fn test_is_same_model() {
        let mut curr_cpuid = CpuId::new(2).unwrap();

        // Add the vendor ID leaf.
        let vendor = get_cpuid(0x0, 0).unwrap();
        curr_cpuid.as_mut_slice()[0].function = 0x0;
        curr_cpuid.as_mut_slice()[0].index = 0;
        curr_cpuid.as_mut_slice()[0].ebx = vendor.ebx;
        curr_cpuid.as_mut_slice()[0].ecx = vendor.ecx;
        curr_cpuid.as_mut_slice()[0].edx = vendor.edx;

        // Add the feature info leaf.
        let feature_info = get_cpuid(0x1, 0).unwrap();
        curr_cpuid.as_mut_slice()[1].function = 0x1;
        curr_cpuid.as_mut_slice()[1].index = 0;
        curr_cpuid.as_mut_slice()[1].eax = feature_info.eax;

        assert!(is_same_model(&curr_cpuid));

        let mut diff_vendor_cpuid = curr_cpuid.clone();
        for mut entry in diff_vendor_cpuid.as_mut_slice() {
            if entry.function == 0x0 && entry.index == 0 {
                entry.ebx = 0xFFFF_FFFF;
            }
        }
        assert!(!is_same_model(&diff_vendor_cpuid));

        let mut diff_feature_cpuid = curr_cpuid;
        for mut entry in diff_feature_cpuid.as_mut_slice() {
            if entry.function == 0x1 && entry.index == 0 {
                entry.eax ^= 0x1;
            }
        }
        assert!(!is_same_model(&diff_feature_cpuid));
    }
}
<|endoftext|>


<commit_before>use std::collections::VecDeque;
use std::fmt;
use term;
use time::precise_time_s;
use multirust::Notification;
use rust_install::Notification as In;
use rust_install::utils::Notification as Un;

/// Keep track of this many past download amounts
const DOWNLOAD_TRACK_COUNT: usize = 5;

/// Tracks download progress and displays information about it to a terminal.
pub struct DownloadTracker {
    /// Content-Length of the to-be downloaded object.
    content_len: Option<u64>,
    /// Total data downloaded in bytes.
    total_downloaded: usize,
    /// Data downloaded this second.
    downloaded_this_sec: usize,
    /// Keeps track of amount of data downloaded every last few secs.
    /// Used for averaging the download speed.
    downloaded_last_few_secs: VecDeque<usize>,
    /// Time stamp of the last second
    last_sec: Option<f64>,
    /// How many seconds have elapsed since the download started
    seconds_elapsed: u32,
    /// The terminal we write the information to.
    term: Box<term::StdoutTerminal>,
    /// Whether we displayed progress for the download or not.
    ///
    /// If the download is quick enough, we don't have time to
    /// display the progress info.
    /// In that case, we do not want to do some cleanup stuff we normally do.
    displayed_progress: bool,
}

impl DownloadTracker {
    /// Creates a new DownloadTracker.
    pub fn new() -> Self {
        DownloadTracker {
            content_len: None,
            total_downloaded: 0,
            downloaded_this_sec: 0,
            downloaded_last_few_secs: VecDeque::with_capacity(DOWNLOAD_TRACK_COUNT),
            seconds_elapsed: 0,
            last_sec: None,
            term: term::stdout().expect("Failed to open terminal"),
            displayed_progress: false,
        }
    }

    pub fn handle_notification(&mut self, n: &Notification) -> bool {
        match n {
            &Notification::Install(In::Utils(Un::DownloadContentLengthReceived(content_len))) => {
                self.content_length_received(content_len);
                true
            }
            &Notification::Install(In::Utils(Un::DownloadDataReceived(len))) => {
                if super::stderr_isatty() {
                    self.data_received(len);
                }
                true
            }
            &Notification::Install(In::Utils(Un::DownloadFinished)) => {
                self.download_finished();
                true
            }
            _ => false
        }
    }

    /// Notifies self that Content-Length information has been received.
    pub fn content_length_received(&mut self, content_len: u64) {
        self.content_len = Some(content_len);
    }
    /// Notifies self that data of size `len` has been received.
    pub fn data_received(&mut self, len: usize) {
        self.total_downloaded += len;
        self.downloaded_this_sec += len;

        match self.last_sec {
            None => self.last_sec = Some(precise_time_s()),
            Some(start) => {
                let elapsed = precise_time_s() - start;
                if elapsed >= 1.0 {
                    self.seconds_elapsed += 1;

                    self.display();
                    self.last_sec = Some(precise_time_s());
                    if self.downloaded_last_few_secs.len() == DOWNLOAD_TRACK_COUNT {
                        self.downloaded_last_few_secs.pop_back();
                    }
                    self.downloaded_last_few_secs.push_front(self.downloaded_this_sec);
                    self.downloaded_this_sec = 0;
                }
            }
        }
    }
    /// Notifies self that the download has finished.
    pub fn download_finished(&mut self) {
        if self.displayed_progress {
            // Display the finished state
            self.display();
            let _ = writeln!(&mut self.term, "");
        }
        self.prepare_for_new_download();
    }
    /// Resets the state to be ready for a new download.
    fn prepare_for_new_download(&mut self) {
        self.content_len = None;
        self.total_downloaded = 0;
        self.downloaded_this_sec = 0;
        self.downloaded_last_few_secs.clear();
        self.seconds_elapsed = 0;
        self.last_sec = None;
        self.displayed_progress = false;
    }
    /// Display the tracked download information to the terminal.
    fn display(&mut self) {
        let total_h = HumanReadable(self.total_downloaded as u64);
        let sum = self.downloaded_last_few_secs
                      .iter()
                      .fold(0usize, |a, &v| a + v);
        let len = self.downloaded_last_few_secs.len();
        let speed = if len > 0 {
            (sum / len) as u64
        } else {
            0
        };
        let speed_h = HumanReadable(speed);

        match self.content_len {
            Some(content_len) => {
                use std::borrow::Cow;

                let percent = (self.total_downloaded as f64 / content_len as f64) * 100.;
                let content_len_h = HumanReadable(content_len);
                let remaining = content_len - self.total_downloaded as u64;
                let eta = if speed > 0 {
                    Cow::Owned(format!("{}s", remaining / speed))
                } else {
                    Cow::Borrowed("Unknown")
                };
                let _ = write!(&mut self.term,
                               "{} / {} ({:.2}%) {}/s ETA: {}",
                               total_h,
                               content_len_h,
                               percent,
                               speed_h,
                               eta);
            }
            None => {
                let _ = write!(&mut self.term, "Total: {} Speed: {}/s", total_h, speed_h);
            }
        }
        // delete_line() doesn't seem to clear the line properly.
        // Instead, let's just print some whitespace to clear it.
        let _ = write!(&mut self.term, "                ");
        let _ = self.term.flush();
        let _ = self.term.carriage_return();
        self.displayed_progress = true;
    }
}

/// Human readable representation of data size in bytes
struct HumanReadable(u64);

impl fmt::Display for HumanReadable {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        const KIB: f64 = 1024.0;
        const MIB: f64 = 1048576.0;
        let size = self.0 as f64;

        if size >= MIB {
            write!(f, "{:.2} MiB", size / MIB)
        } else if size >= KIB {
            write!(f, "{:.2} KiB", size / KIB)
        } else {
            write!(f, "{} B", size)
        }
    }
}
<commit_msg>Only get current time once<commit_after>use std::collections::VecDeque;
use std::fmt;
use term;
use time::precise_time_s;
use multirust::Notification;
use rust_install::Notification as In;
use rust_install::utils::Notification as Un;

/// Keep track of this many past download amounts
const DOWNLOAD_TRACK_COUNT: usize = 5;

/// Tracks download progress and displays information about it to a terminal.
pub struct DownloadTracker {
    /// Content-Length of the to-be downloaded object.
    content_len: Option<u64>,
    /// Total data downloaded in bytes.
    total_downloaded: usize,
    /// Data downloaded this second.
    downloaded_this_sec: usize,
    /// Keeps track of amount of data downloaded every last few secs.
    /// Used for averaging the download speed.
    downloaded_last_few_secs: VecDeque<usize>,
    /// Time stamp of the last second
    last_sec: Option<f64>,
    /// How many seconds have elapsed since the download started
    seconds_elapsed: u32,
    /// The terminal we write the information to.
    term: Box<term::StdoutTerminal>,
    /// Whether we displayed progress for the download or not.
    ///
    /// If the download is quick enough, we don't have time to
    /// display the progress info.
    /// In that case, we do not want to do some cleanup stuff we normally do.
    displayed_progress: bool,
}

impl DownloadTracker {
    /// Creates a new DownloadTracker.
    pub fn new() -> Self {
        DownloadTracker {
            content_len: None,
            total_downloaded: 0,
            downloaded_this_sec: 0,
            downloaded_last_few_secs: VecDeque::with_capacity(DOWNLOAD_TRACK_COUNT),
            seconds_elapsed: 0,
            last_sec: None,
            term: term::stdout().expect("Failed to open terminal"),
            displayed_progress: false,
        }
    }

    pub fn handle_notification(&mut self, n: &Notification) -> bool {
        match n {
            &Notification::Install(In::Utils(Un::DownloadContentLengthReceived(content_len))) => {
                self.content_length_received(content_len);
                true
            }
            &Notification::Install(In::Utils(Un::DownloadDataReceived(len))) => {
                if super::stderr_isatty() {
                    self.data_received(len);
                }
                true
            }
            &Notification::Install(In::Utils(Un::DownloadFinished)) => {
                self.download_finished();
                true
            }
            _ => false
        }
    }

    /// Notifies self that Content-Length information has been received.
    pub fn content_length_received(&mut self, content_len: u64) {
        self.content_len = Some(content_len);
    }
    /// Notifies self that data of size `len` has been received.
    pub fn data_received(&mut self, len: usize) {
        self.total_downloaded += len;
        self.downloaded_this_sec += len;

        let current_time = precise_time_s();

        match self.last_sec {
            None => self.last_sec = Some(current_time),
            Some(start) => {
                let elapsed = current_time - start;
                if elapsed >= 1.0 {
                    self.seconds_elapsed += 1;

                    self.display();
                    self.last_sec = Some(current_time);
                    if self.downloaded_last_few_secs.len() == DOWNLOAD_TRACK_COUNT {
                        self.downloaded_last_few_secs.pop_back();
                    }
                    self.downloaded_last_few_secs.push_front(self.downloaded_this_sec);
                    self.downloaded_this_sec = 0;
                }
            }
        }
    }
    /// Notifies self that the download has finished.
    pub fn download_finished(&mut self) {
        if self.displayed_progress {
            // Display the finished state
            self.display();
            let _ = writeln!(&mut self.term, "");
        }
        self.prepare_for_new_download();
    }
    /// Resets the state to be ready for a new download.
    fn prepare_for_new_download(&mut self) {
        self.content_len = None;
        self.total_downloaded = 0;
        self.downloaded_this_sec = 0;
        self.downloaded_last_few_secs.clear();
        self.seconds_elapsed = 0;
        self.last_sec = None;
        self.displayed_progress = false;
    }
    /// Display the tracked download information to the terminal.
    fn display(&mut self) {
        let total_h = HumanReadable(self.total_downloaded as u64);
        let sum = self.downloaded_last_few_secs
                      .iter()
                      .fold(0usize, |a, &v| a + v);
        let len = self.downloaded_last_few_secs.len();
        let speed = if len > 0 {
            (sum / len) as u64
        } else {
            0
        };
        let speed_h = HumanReadable(speed);

        match self.content_len {
            Some(content_len) => {
                use std::borrow::Cow;

                let percent = (self.total_downloaded as f64 / content_len as f64) * 100.;
                let content_len_h = HumanReadable(content_len);
                let remaining = content_len - self.total_downloaded as u64;
                let eta = if speed > 0 {
                    Cow::Owned(format!("{}s", remaining / speed))
                } else {
                    Cow::Borrowed("Unknown")
                };
                let _ = write!(&mut self.term,
                               "{} / {} ({:.2}%) {}/s ETA: {}",
                               total_h,
                               content_len_h,
                               percent,
                               speed_h,
                               eta);
            }
            None => {
                let _ = write!(&mut self.term, "Total: {} Speed: {}/s", total_h, speed_h);
            }
        }
        // delete_line() doesn't seem to clear the line properly.
        // Instead, let's just print some whitespace to clear it.
        let _ = write!(&mut self.term, "                ");
        let _ = self.term.flush();
        let _ = self.term.carriage_return();
        self.displayed_progress = true;
    }
}

/// Human readable representation of data size in bytes
struct HumanReadable(u64);

impl fmt::Display for HumanReadable {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        const KIB: f64 = 1024.0;
        const MIB: f64 = 1048576.0;
        let size = self.0 as f64;

        if size >= MIB {
            write!(f, "{:.2} MiB", size / MIB)
        } else if size >= KIB {
            write!(f, "{:.2} KiB", size / KIB)
        } else {
            write!(f, "{} B", size)
        }
    }
}
<|endoftext|>


<commit_before>use llvm_sys::core::LLVMDisposeMessage;
use llvm_sys::execution_engine::{LLVMGetExecutionEngineTargetData, LLVMExecutionEngineRef, LLVMRunFunction, LLVMRunFunctionAsMain, LLVMDisposeExecutionEngine, LLVMGetFunctionAddress, LLVMAddModule, LLVMFindFunction, LLVMLinkInMCJIT, LLVMLinkInInterpreter, LLVMRemoveModule, LLVMGenericValueRef, LLVMFreeMachineCodeForFunction, LLVMGetPointerToGlobal};

use module::Module;
use targets::TargetData;
use values::{AsValueRef, FunctionValue, GenericValue, GlobalValue};

use std::ffi::{CStr, CString};
use std::mem::{forget, uninitialized, zeroed};

#[derive(Debug, PartialEq, Eq)]
pub enum FunctionLookupError {
    JITNotEnabled,
    FunctionNotFound, // 404!
}

#[derive(Debug)]
pub struct ExecutionEngine {
    execution_engine: LLVMExecutionEngineRef,
    pub(crate) modules: Vec<Module>,
    target_data: Option<TargetData>,
    jit_mode: bool,
}

impl ExecutionEngine {
    pub(crate) fn new(execution_engine: LLVMExecutionEngineRef, jit_mode: bool) -> ExecutionEngine {
        assert!(!execution_engine.is_null());

        // REVIEW: Will we have to do this for LLVMGetExecutionEngineTargetMachine too?
        let target_data = unsafe {
            LLVMGetExecutionEngineTargetData(execution_engine)
        };

        ExecutionEngine {
            execution_engine: execution_engine,
            modules: vec![],
            target_data: Some(TargetData::new(target_data)),
            jit_mode: jit_mode,
        }
    }

    pub fn link_in_mc_jit() {
        unsafe {
            LLVMLinkInMCJIT()
        }
    }

    pub fn link_in_interpreter() {
        unsafe {
            LLVMLinkInInterpreter();
        }
    }

    // pub fn add_global_mapping(&mut self, global: &AnyValue, addr: *const ()) {
    //     unsafe {
    //         LLVMAddGlobalMapping(self.execution_engine, global.as_value_ref(), addr as *mut ::libc::c_void)
    //     }
    // }

    // TODOC: EE must *own* modules and deal out references
    pub fn add_module(&mut self, module: Module) -> &Module {
        unsafe {
            LLVMAddModule(self.execution_engine, module.module)
        }

        self.modules.push(module);

        &self.modules[self.modules.len() - 1]
    }

    pub fn remove_module(&mut self, module: &Module) -> Result<Module, String> {
        let mut new_module = unsafe { uninitialized() };
        let mut err_str = unsafe { zeroed() };

        let code = unsafe {
            LLVMRemoveModule(self.execution_engine, module.module, &mut new_module, &mut err_str)
        };

        if code == 1 {
            let rust_str = unsafe {
                let rust_str = CStr::from_ptr(err_str).to_string_lossy().into_owned();

                LLVMDisposeMessage(err_str);

                rust_str
            };

            return Err(rust_str);
        }

        // REVIEW: This might end up a hashtable for better performance
        let mut index = None;

        for (i, owned_module) in self.modules.iter().enumerate() {
            if module == owned_module {
                index = Some(i);
            }
        }

        match index {
            Some(idx) => self.modules.remove(idx),
            None => return Err("Module does not belong to this ExecutionEngine".into()),
        };

        Ok(Module::new(new_module))
    }

    // FIXME: Workaround until we can think of a better API
    // Ideally, we'd provide this as a hashmap from module name to module, but this may not be
    // possible since you can create modules, ie via create_module_from_ir, which either do not
    // have a name, or at least LLVM provides no way to obtain its name. Also, possible collisions
    // depending on how loose, if at all, LLVM is with module names belonging to execution engines
    pub fn get_module_at(&self, index: usize) -> &Module {
        &self.modules[index]
    }

    /// WARNING: The returned address *will* be invalid if the EE drops first
    /// Do not attempt to transmute it to a function if the ExecutionEngine is gone
    // TODOC: Initializing a target MUST occur before creating the EE or else it will not count
    // TODOC: Can still add functions after EE has been created
    pub fn get_function_address(&self, fn_name: &str) -> Result<u64, FunctionLookupError> {
        if !self.jit_mode {
            return Err(FunctionLookupError::JITNotEnabled);
        }

        let c_string = CString::new(fn_name).expect("Conversion to CString failed unexpectedly");

        let address = unsafe {
            LLVMGetFunctionAddress(self.execution_engine, c_string.as_ptr())
        };

        // REVIEW: Can also return 0 if no targets are initialized.
        // One option might be to set a global to true if any at all of the targets have been
        // initialized (maybe we could figure out which config in particular is the trigger)
        // and if not return an "NoTargetsInitialized" error, instead of not found.
        if address == 0 {
            return Err(FunctionLookupError::FunctionNotFound);
        }

        Ok(address)
    }

    // REVIEW: Not sure if an EE's target data can change.. if so we might want to update the value
    // when making this call
    pub fn get_target_data(&self) -> &TargetData {
        self.target_data.as_ref().expect("TargetData should always exist until Drop")
    }

    // REVIEW: Can also find nothing if no targeting is initialized. Maybe best to
    // do have a global flag for anything initialized. Catch is that it must be initialized
    // before EE is created
    pub fn get_function_value(&self, fn_name: &str) -> Result<FunctionValue, FunctionLookupError> {
        if !self.jit_mode {
            return Err(FunctionLookupError::JITNotEnabled);
        }

        let c_string = CString::new(fn_name).expect("Conversion to CString failed unexpectedly");
        let mut function = unsafe { zeroed() };

        let code = unsafe {
            LLVMFindFunction(self.execution_engine, c_string.as_ptr(), &mut function)
        };

        if code == 0 {
            return FunctionValue::new(function).ok_or(FunctionLookupError::FunctionNotFound)
        };

        Err(FunctionLookupError::FunctionNotFound)
    }

    // TODOC: Marked as unsafe because input function could very well do something unsafe. It's up to the caller
    // to ensure that doesn't happen by defining their function correctly.
    pub unsafe fn run_function(&self, function: &FunctionValue, args: &[&GenericValue]) -> GenericValue {
        let mut args: Vec<LLVMGenericValueRef> = args.iter()
                                                     .map(|val| val.generic_value)
                                                     .collect();

        let value = LLVMRunFunction(self.execution_engine, function.as_value_ref(), args.len() as u32, args.as_mut_ptr()); // REVIEW: usize to u32 ok??

        GenericValue::new(value)
    }

    pub fn run_function_as_main(&self, function: &FunctionValue) {
        let args = vec![]; // TODO: Support argc, argv
        let env_p = vec![]; // REVIEW: No clue what this is

        unsafe {
            LLVMRunFunctionAsMain(self.execution_engine, function.as_value_ref(), args.len() as u32, args.as_ptr(), env_p.as_ptr()); // REVIEW: usize to u32 cast ok??
        }
    }

    pub fn free_fn_machine_code(&self, function: &FunctionValue) {
        unsafe {
            LLVMFreeMachineCodeForFunction(self.execution_engine, function.as_value_ref())
        }
    }
}

// Modules owned by the EE will be discarded by the EE so we don't
// want owned modules to drop.
impl Drop for ExecutionEngine {
    fn drop(&mut self) {
        for module in self.modules.drain(..) {
            forget(module);
        }

        forget(self.target_data.take().expect("TargetData should always exist until Drop"));

        unsafe {
            LLVMDisposeExecutionEngine(self.execution_engine);
        }
    }
}
<commit_msg>Added add_global_mapping + docs<commit_after>use llvm_sys::core::LLVMDisposeMessage;
use llvm_sys::execution_engine::{LLVMGetExecutionEngineTargetData, LLVMExecutionEngineRef, LLVMRunFunction, LLVMRunFunctionAsMain, LLVMDisposeExecutionEngine, LLVMGetFunctionAddress, LLVMAddModule, LLVMFindFunction, LLVMLinkInMCJIT, LLVMLinkInInterpreter, LLVMRemoveModule, LLVMGenericValueRef, LLVMFreeMachineCodeForFunction, LLVMGetPointerToGlobal, LLVMAddGlobalMapping};

use module::Module;
use targets::TargetData;
use values::{AnyValue, AsValueRef, FunctionValue, GenericValue, GlobalValue};

use std::ffi::{CStr, CString};
use std::mem::{forget, uninitialized, zeroed};

#[derive(Debug, PartialEq, Eq)]
pub enum FunctionLookupError {
    JITNotEnabled,
    FunctionNotFound, // 404!
}

#[derive(Debug)]
pub struct ExecutionEngine {
    execution_engine: LLVMExecutionEngineRef,
    pub(crate) modules: Vec<Module>,
    target_data: Option<TargetData>,
    jit_mode: bool,
}

impl ExecutionEngine {
    pub(crate) fn new(execution_engine: LLVMExecutionEngineRef, jit_mode: bool) -> ExecutionEngine {
        assert!(!execution_engine.is_null());

        // REVIEW: Will we have to do this for LLVMGetExecutionEngineTargetMachine too?
        let target_data = unsafe {
            LLVMGetExecutionEngineTargetData(execution_engine)
        };

        ExecutionEngine {
            execution_engine: execution_engine,
            modules: vec![],
            target_data: Some(TargetData::new(target_data)),
            jit_mode: jit_mode,
        }
    }

    pub fn link_in_mc_jit() {
        unsafe {
            LLVMLinkInMCJIT()
        }
    }

    pub fn link_in_interpreter() {
        unsafe {
            LLVMLinkInInterpreter();
        }
    }

    /// Maps the specified value to an address.
    ///
    /// # Example
    /// ```
    /// use inkwell::targets::{InitializationConfig, Target};
    /// use inkwell::context::Context;
    /// use inkwell::OptimizationLevel;
    ///
    /// Target::initialize_native(&InitializationConfig::default()).unwrap();
    ///
    /// extern fn sumf(a: f64, b: f64) -> f64 {
    ///     a + b
    /// }
    ///
    /// let context = Context::create();
    /// let module = context.create_module("test");
    /// let builder = context.create_builder();
    ///
    /// let ft = context.f64_type();
    /// let fnt = ft.fn_type(&[], false);
    ///
    /// let f = module.add_function("test_fn", &fnt, None);
    /// let b = context.append_basic_block(&f, "entry");
    /// builder.position_at_end(&b);
    ///
    /// let extf = module.add_function("sumf", &ft.fn_type(&[ &ft, &ft ], false), None);
    ///
    /// let argf = ft.const_float(64.);
    /// let retv = builder.build_call(&extf, &[ &argf, &argf ], "retv", false).left().unwrap().into_float_value();
    ///
    /// builder.build_return(Some(&retv));
    ///
    /// let mut ee = module.create_jit_execution_engine(OptimizationLevel::None).unwrap();
    /// ee.add_global_mapping(&extf, sumf as *mut _);
    ///
    /// let result = unsafe { ee.run_function(&f, &[]) }.as_float(&ft);
    ///
    /// assert_eq!(result, 128.);
    /// ```
    pub fn add_global_mapping(&mut self, value: &AnyValue, addr: *const ()) {
        unsafe {
            LLVMAddGlobalMapping(self.execution_engine, value.as_value_ref(), addr as *mut _)
        }
    }

    // TODOC: EE must *own* modules and deal out references
    pub fn add_module(&mut self, module: Module) -> &Module {
        unsafe {
            LLVMAddModule(self.execution_engine, module.module)
        }

        self.modules.push(module);

        &self.modules[self.modules.len() - 1]
    }

    pub fn remove_module(&mut self, module: &Module) -> Result<Module, String> {
        let mut new_module = unsafe { uninitialized() };
        let mut err_str = unsafe { zeroed() };

        let code = unsafe {
            LLVMRemoveModule(self.execution_engine, module.module, &mut new_module, &mut err_str)
        };

        if code == 1 {
            let rust_str = unsafe {
                let rust_str = CStr::from_ptr(err_str).to_string_lossy().into_owned();

                LLVMDisposeMessage(err_str);

                rust_str
            };

            return Err(rust_str);
        }

        // REVIEW: This might end up a hashtable for better performance
        let mut index = None;

        for (i, owned_module) in self.modules.iter().enumerate() {
            if module == owned_module {
                index = Some(i);
            }
        }

        match index {
            Some(idx) => self.modules.remove(idx),
            None => return Err("Module does not belong to this ExecutionEngine".into()),
        };

        Ok(Module::new(new_module))
    }

    // FIXME: Workaround until we can think of a better API
    // Ideally, we'd provide this as a hashmap from module name to module, but this may not be
    // possible since you can create modules, ie via create_module_from_ir, which either do not
    // have a name, or at least LLVM provides no way to obtain its name. Also, possible collisions
    // depending on how loose, if at all, LLVM is with module names belonging to execution engines
    pub fn get_module_at(&self, index: usize) -> &Module {
        &self.modules[index]
    }

    /// WARNING: The returned address *will* be invalid if the EE drops first
    /// Do not attempt to transmute it to a function if the ExecutionEngine is gone
    // TODOC: Initializing a target MUST occur before creating the EE or else it will not count
    // TODOC: Can still add functions after EE has been created
    pub fn get_function_address(&self, fn_name: &str) -> Result<u64, FunctionLookupError> {
        if !self.jit_mode {
            return Err(FunctionLookupError::JITNotEnabled);
        }

        let c_string = CString::new(fn_name).expect("Conversion to CString failed unexpectedly");

        let address = unsafe {
            LLVMGetFunctionAddress(self.execution_engine, c_string.as_ptr())
        };

        // REVIEW: Can also return 0 if no targets are initialized.
        // One option might be to set a global to true if any at all of the targets have been
        // initialized (maybe we could figure out which config in particular is the trigger)
        // and if not return an "NoTargetsInitialized" error, instead of not found.
        if address == 0 {
            return Err(FunctionLookupError::FunctionNotFound);
        }

        Ok(address)
    }

    // REVIEW: Not sure if an EE's target data can change.. if so we might want to update the value
    // when making this call
    pub fn get_target_data(&self) -> &TargetData {
        self.target_data.as_ref().expect("TargetData should always exist until Drop")
    }

    // REVIEW: Can also find nothing if no targeting is initialized. Maybe best to
    // do have a global flag for anything initialized. Catch is that it must be initialized
    // before EE is created
    pub fn get_function_value(&self, fn_name: &str) -> Result<FunctionValue, FunctionLookupError> {
        if !self.jit_mode {
            return Err(FunctionLookupError::JITNotEnabled);
        }

        let c_string = CString::new(fn_name).expect("Conversion to CString failed unexpectedly");
        let mut function = unsafe { zeroed() };

        let code = unsafe {
            LLVMFindFunction(self.execution_engine, c_string.as_ptr(), &mut function)
        };

        if code == 0 {
            return FunctionValue::new(function).ok_or(FunctionLookupError::FunctionNotFound)
        };

        Err(FunctionLookupError::FunctionNotFound)
    }

    // TODOC: Marked as unsafe because input function could very well do something unsafe. It's up to the caller
    // to ensure that doesn't happen by defining their function correctly.
    pub unsafe fn run_function(&self, function: &FunctionValue, args: &[&GenericValue]) -> GenericValue {
        let mut args: Vec<LLVMGenericValueRef> = args.iter()
                                                     .map(|val| val.generic_value)
                                                     .collect();

        let value = LLVMRunFunction(self.execution_engine, function.as_value_ref(), args.len() as u32, args.as_mut_ptr()); // REVIEW: usize to u32 ok??

        GenericValue::new(value)
    }

    pub fn run_function_as_main(&self, function: &FunctionValue) {
        let args = vec![]; // TODO: Support argc, argv
        let env_p = vec![]; // REVIEW: No clue what this is

        unsafe {
            LLVMRunFunctionAsMain(self.execution_engine, function.as_value_ref(), args.len() as u32, args.as_ptr(), env_p.as_ptr()); // REVIEW: usize to u32 cast ok??
        }
    }

    pub fn free_fn_machine_code(&self, function: &FunctionValue) {
        unsafe {
            LLVMFreeMachineCodeForFunction(self.execution_engine, function.as_value_ref())
        }
    }
}

// Modules owned by the EE will be discarded by the EE so we don't
// want owned modules to drop.
impl Drop for ExecutionEngine {
    fn drop(&mut self) {
        for module in self.modules.drain(..) {
            forget(module);
        }

        forget(self.target_data.take().expect("TargetData should always exist until Drop"));

        unsafe {
            LLVMDisposeExecutionEngine(self.execution_engine);
        }
    }
}
<|endoftext|>


<commit_before>//! This module contains type aliases for C's fixed-width integer types .
//!
//! These aliases are deprecated: use the Rust types instead.

pub type int8_t = i8;
pub type int16_t = i16;
pub type int32_t = i32;
pub type int64_t = i64;
pub type uint8_t = u8;
pub type uint16_t = u16;
pub type uint32_t = u32;
pub type uint64_t = u64;
<commit_msg>Deprecates fixed-width integer type aliases<commit_after>//! This module contains type aliases for C's fixed-width integer types .
//!
//! These aliases are deprecated: use the Rust types instead.

#[deprecated(
    since = "0.2.55",
    note = "Use i8 instead."
)]
pub type int8_t = i8;
#[deprecated(
    since = "0.2.55",
    note = "Use i16 instead."
)]
pub type int16_t = i16;
#[deprecated(
    since = "0.2.55",
    note = "Use i32 instead."
)]
pub type int32_t = i32;
#[deprecated(
    since = "0.2.55",
    note = "Use i64 instead."
)]
pub type int64_t = i64;
#[deprecated(
    since = "0.2.55",
    note = "Use u8 instead."
)]
pub type uint8_t = u8;
#[deprecated(
    since = "0.2.55",
    note = "Use u16 instead."
)]
pub type uint16_t = u16;
#[deprecated(
    since = "0.2.55",
    note = "Use u32 instead."
)]
pub type uint32_t = u32;
#[deprecated(
    since = "0.2.55",
    note = "Use u64 instead."
)]
pub type uint64_t = u64;
<|endoftext|>


<commit_before>extern crate frame_stack;

use self::frame_stack::Frame;

extern {
    static kernel_end_symbol_table_entry: ();
    //static kernel_start_symbol_table_entry: ();
}

const PAGE_SIZE: u64 = 4096;

pub unsafe fn init_frame_stack(multiboot: *const ::multiboot::Information) {

    let kernel_end = &kernel_end_symbol_table_entry as *const () as u32;

    let stack_start_frame = Frame{number: kernel_end >> 12};

    // map frame stack to 2mb behind kernel
    map_p1_entries(0, 0, stack_start_frame);
    // map p1 entries
    for i in 0..512 {
        map_p1_entries(0, i, Frame{
            number: stack_start_frame.number + i,
        });
    }

    frame_stack::init();

    let mut areas = (*multiboot).memory_areas().unwrap();
    let number_of_frames = areas.clone().fold(0, |n, area| {n + area.length / PAGE_SIZE}) as u32;

    // for memory > 1GB we need more p1 tables
    let p1_tables_required = (frame_stack::max_size(number_of_frames) / 4096 / 512) as u32;
    let mut last_mapped_p1_table = 0;

    // the first free physical frame (the 512 frames after kernel are used)
    let first_free_frame = Frame{
        number: stack_start_frame.number + 512,
    };

    for area in areas {
        for frame in ((area.base_addr >> 12)..((area.base_addr+area.length) >> 12)).map(
            |frame_number| Frame{number: frame_number as u32}) {
            if frame >= first_free_frame {
                if last_mapped_p1_table < p1_tables_required {
                    // add as p1 table
                    last_mapped_p1_table += 1;
                    map_to_p1_table(last_mapped_p1_table, frame);
                    // map p1 entries
                    for i in 0..512 {
                        map_p1_entries(last_mapped_p1_table, i, Frame{
                            number: stack_start_frame.number + last_mapped_p1_table*512 + i,
                        });
                    }
                } else {
                    // add as free
                    frame_stack::deallocate_frame(frame);
                }
            }
        }
    }

    unsafe fn map_to_p1_table(p2_index: u32, frame: Frame) {
        let p1_field: *mut u64 = (0o177777_777_777_000_001_0000 + (p2_index as u64) * 8)
            as *mut u64;
        *p1_field = ((frame.number as u64) << 12) | 1;
    }

    unsafe fn map_p1_entries(p2_index: u32, p1_index: u32, to: Frame) {
        let entry: *mut u64 = (0o177777_777_000_001_000_0000 | ((p2_index as u64) << 12)
            | (p1_index as u64 * 8)) as *mut u64;
        *entry = ((to.number as u64) << 12) | 3;
    }
}
<commit_msg>specify type in range notation<commit_after>extern crate frame_stack;

use self::frame_stack::Frame;

extern {
    static kernel_end_symbol_table_entry: ();
    //static kernel_start_symbol_table_entry: ();
}

const PAGE_SIZE: u64 = 4096;

pub unsafe fn init_frame_stack(multiboot: *const ::multiboot::Information) {

    let kernel_end = &kernel_end_symbol_table_entry as *const () as u32;

    let stack_start_frame = Frame{number: kernel_end >> 12};

    // map frame stack to 2mb behind kernel
    map_p1_entries(0, 0, stack_start_frame);
    // map p1 entries
    for i in 0..512u32 {
        map_p1_entries(0, i, Frame{
            number: stack_start_frame.number + i,
        });
    }

    frame_stack::init();

    let mut areas = (*multiboot).memory_areas().unwrap();
    let number_of_frames = areas.clone().fold(0, |n, area| {n + area.length / PAGE_SIZE}) as u32;

    // for memory > 1GB we need more p1 tables
    let p1_tables_required = (frame_stack::max_size(number_of_frames) / 4096 / 512) as u32;
    let mut last_mapped_p1_table = 0;

    // the first free physical frame (the 512 frames after kernel are used)
    let first_free_frame = Frame{
        number: stack_start_frame.number + 512,
    };

    for area in areas {
        for frame in ((area.base_addr >> 12)..((area.base_addr+area.length) >> 12)).map(
            |frame_number| Frame{number: frame_number as u32}) {
            if frame >= first_free_frame {
                if last_mapped_p1_table < p1_tables_required {
                    // add as p1 table
                    last_mapped_p1_table += 1;
                    map_to_p1_table(last_mapped_p1_table, frame);
                    // map p1 entries
                    for i in 0..512u32 {
                        map_p1_entries(last_mapped_p1_table, i, Frame{
                            number: stack_start_frame.number + last_mapped_p1_table*512 + i,
                        });
                    }
                } else {
                    // add as free
                    frame_stack::deallocate_frame(frame);
                }
            }
        }
    }

    unsafe fn map_to_p1_table(p2_index: u32, frame: Frame) {
        let p1_field: *mut u64 = (0o177777_777_777_000_001_0000 + (p2_index as u64) * 8)
            as *mut u64;
        *p1_field = ((frame.number as u64) << 12) | 1;
    }

    unsafe fn map_p1_entries(p2_index: u32, p1_index: u32, to: Frame) {
        let entry: *mut u64 = (0o177777_777_000_001_000_0000 | ((p2_index as u64) << 12)
            | (p1_index as u64 * 8)) as *mut u64;
        *entry = ((to.number as u64) << 12) | 3;
    }
}
<|endoftext|>


<commit_before>//! Special x86_64 instructions.

pub mod interrupts;
pub mod port;
pub mod segmentation;
pub mod tables;
pub mod tlb;

/// Cause a breakpoint exception by invoking the `int3` instruction.
pub fn int3() {
    unsafe {
        asm!("int3");
    }
}

/// Halts the CPU until the next interrupt arrives.
#[inline(always)]
pub unsafe fn hlt() {
    asm!("hlt" :::: "volatile");
}
<commit_msg>Make hlt function safe<commit_after>//! Special x86_64 instructions.

pub mod interrupts;
pub mod port;
pub mod segmentation;
pub mod tables;
pub mod tlb;

/// Cause a breakpoint exception by invoking the `int3` instruction.
pub fn int3() {
    unsafe {
        asm!("int3");
    }
}

/// Halts the CPU until the next interrupt arrives.
#[inline(always)]
pub fn hlt() {
    unsafe {
        asm!("hlt" :::: "volatile");
    }
}
<|endoftext|>


<commit_before>//! Converts the from ESIL/CFG representation to SSA.

use std::collections::HashMap;
use petgraph::graph::NodeIndex;
use frontend::structs::LRegInfo;
use middle::cfg::NodeData as CFGNodeData;
use middle::cfg::EdgeType as CFGEdgeType;
use middle::cfg::{CFG, BasicBlock};
use middle::ssa::{BBInfo, SSA, SSAMod, ValueType};
use middle::ssa::verifier::Verify;
use middle::ssa::verifier::VerifiedAdd;
use middle::ir::{MVal, MInst, MOpcode, MValType};
use middle::phiplacement::PhiPlacer;
use middle::regfile::SubRegisterFile;
use middle::dce;

pub type VarId = usize;

pub struct SSAConstruction<'a, T> where
	T: Clone + Verify + SSAMod<BBInfo=BBInfo> + 'a
{
	pub phiplacer: PhiPlacer<'a, T>,
	pub regfile:   SubRegisterFile,
	pub temps:     HashMap<String, T::ValueRef>,
}

impl<'a, T> SSAConstruction<'a, T> where
	T: 'a + Clone + Verify +
       SSAMod<BBInfo=BBInfo, ValueRef=NodeIndex, ActionRef=NodeIndex>
{
	pub fn new(ssa: &'a mut T, reg_info: &LRegInfo) -> SSAConstruction<'a, T> {
		let mut sc = SSAConstruction {
			phiplacer: PhiPlacer::new(ssa),
			regfile:   SubRegisterFile::new(reg_info),
			temps:     HashMap::new(),
		};
		sc.phiplacer.add_variables(vec![
								   ValueType::Integer{width: 64}, // cur
								   ValueType::Integer{width: 64},  // old
								   ValueType::Integer{width: 64}  // lastsz
		]);
		// make the following a method of regfile?
		sc.phiplacer.add_variables(sc.regfile.whole_registers.clone());
		sc
	}

	pub fn run(&mut self, cfg: &CFG) {
		let mut blocks = Vec::<T::ActionRef>::with_capacity(cfg.g.node_count());
		let bb_iter = cfg.bbs.iter();

		for _ in 0..cfg.g.node_count() {
			blocks.push(self.phiplacer.ssa.invalid_action());
		}

		{
			// Insert the entry and exit blocks for the ssa.
			let block = self.phiplacer.add_block(BBInfo { addr: 0 });
			self.phiplacer.ssa.mark_start_node(&block);
			let zero = self.phiplacer.ssa.add_const(block, 0);
			self.phiplacer.write_variable(block, 0, zero); // cur = 0
			self.phiplacer.write_variable(block, 1, zero); // old = 0

			for (i, name) in self.regfile.whole_names.iter().enumerate() {
				let argnode = self.phiplacer.ssa.add_comment(block, self.regfile.whole_registers[i], name.clone());
				self.phiplacer.write_variable(block, i+2, argnode);
			}

			blocks[cfg.entry.index()] = block;

			let block = self.phiplacer.add_dynamic();
			self.phiplacer.ssa.mark_exit_node(&block);
			blocks[cfg.exit.index()] = block;
		}

		for (addr, i) in bb_iter {
			let block = self.phiplacer.add_block(BBInfo { addr: *addr });
			blocks[i.index()] = block;
			match cfg.g[*i] {
				CFGNodeData::Block(ref srcbb) => {
					self.process_block(block, srcbb);
				}
				_ => unreachable!(),
			}
		}

		for edge in cfg.g.raw_edges() {
			let i = match edge.weight.edge_type {
				CFGEdgeType::False => 0,
				CFGEdgeType::True => 1,
				CFGEdgeType::Unconditional => 2,
			};
			self.phiplacer.ssa.add_control_edge(
				blocks[edge.source().index()],
				blocks[edge.target().index()],
				i);
		}
		for &block in &blocks {
			self.phiplacer.seal_block(block);
		}

		//self.phiplacer.ssa.stable_indexing = false;
		//self.phiplacer.ssa.cleanup();

		{
			dce::collect(self.phiplacer.ssa);
		}
		//dce::collect(self.phiplacer.ssa, &[exit_regstate]);
		//self.phiplacer.ssa.cleanup();
	}

	fn process_in_flag(&mut self, block: T::ActionRef, _mval: &MVal) -> T::ValueRef {
		// this would be neccesary if parser didn't use EsilCur/EsilOld
		let _cur = self.phiplacer.read_variable(block, 0);
		let _old = self.phiplacer.read_variable(block, 1);
		let _lastsz = self.phiplacer.read_variable(block, 2);
		//unimplemented!();
		self.phiplacer.ssa.invalid_value()
	}

	fn process_in(&mut self, block: T::ActionRef, mval: &MVal) -> T::ValueRef {
		match mval.val_type {
			MValType::Register  => self.regfile.read_register(&mut self.phiplacer, 2, block, &mval.name),
			MValType::Temporary => self.temps[&mval.name],
			MValType::Internal  => self.process_in_flag(block, mval),
			MValType::EsilCur   => self.phiplacer.read_variable(block, 0),
			MValType::EsilOld   => self.phiplacer.read_variable(block, 1),
			MValType::Lastsz   => self.phiplacer.read_variable(block, 2),
			MValType::Unknown   => self.phiplacer.ssa.invalid_value(),
			                       //self.phiplacer.ssa.add_comment(block, &"Unknown".to_string()), // unimplemented!()
			MValType::Null      => self.phiplacer.ssa.invalid_value(),
		}
	}

	fn process_out(&mut self, block: T::ActionRef, mval: &MVal, value: T::ValueRef) {
		match mval.val_type {
			MValType::Register  => self.regfile.write_register(&mut self.phiplacer, 2, block, &mval.name, value),
			MValType::Temporary => {self.temps.insert(mval.name.clone(), value);},
			MValType::Null      => {},
			_                   => panic!(),
		}
	}

	fn process_op(&mut self, block: T::ActionRef, inst: &MInst, n0: T::ValueRef, n1: T::ValueRef) -> T::ValueRef {
		if inst.opcode == MOpcode::OpEq {
			return n0
		}

		let dsttype = match inst.dst.val_type {
			MValType::Null => ValueType::Integer{width: 0}, // there is no ValueType::None?
			_              => ValueType::Integer{width: inst.dst.size},
		};

		let nn = {
			// let ref mut ssa = self.phiplacer.ssa;
			// let nn = ssa.add_op(block, inst.opcode, dsttype);
			// ssa.op_use(nn, 0, n0);
			// ssa.op_use(nn, 1, n1);
			// nn
			(*self.phiplacer.ssa).verified_add_op(block, inst.opcode, dsttype, &[n0, n1])
		};

		if inst.update_flags {
			let mut nn64 = nn;

			if inst.dst.size != 64 {
				let ref mut ssa = self.phiplacer.ssa;
				nn64 = ssa.add_op(block,
					MOpcode::OpWiden(64),
					ValueType::Integer{width: 64});
				ssa.op_use(nn64, 0, nn);
			}

			let old = self.phiplacer.read_variable(block, 0);
			self.phiplacer.write_variable(block, 1, old);
			self.phiplacer.write_variable(block, 0, nn64);
		}

		return nn
	}

	fn process_block(&mut self, block: T::ActionRef, source: &BasicBlock) {
		let mut machinestate = self.phiplacer.ssa.to_value(block);
		for ref instruction in &source.instructions {
			let n0 = self.process_in(block, &instruction.operand_1);
			let n1 = self.process_in(block, &instruction.operand_2);

			if let MOpcode::OpNarrow(i) = instruction.opcode {
				println!("Found narrow. Narrow to : {}, n0: {:?}, n1: {:?},
				op1: {}, op2: {}", i, n0, n1, instruction.operand_1, instruction.operand_2);
			}

			if instruction.opcode == MOpcode::OpJmp {
				// TODO: In case of static jumps, this is trivial and does not need a selector.
				// In case of dynamic jump, the jump targets have to be determined.
				//self.ssa.g.add_edge(block, n0, SSAEdgeData::DynamicControl(0));
				break;
			}

			if instruction.opcode == MOpcode::OpCJmp {
				println!("Identified a selector: {:?} for block {:?}", n0, block);
				self.phiplacer.ssa.mark_selector(n0, block);
				continue;
			}

			let nn = self.process_op(block, instruction, n0, n1);

			if instruction.opcode == MOpcode::OpLoad {
				self.phiplacer.ssa.op_use(nn, 3, machinestate);
			}
			if instruction.opcode == MOpcode::OpStore {
				self.phiplacer.ssa.op_use(nn, 3, machinestate);
				machinestate = nn;
			}

			self.process_out(block, &instruction.dst, nn);
		}
	}
}
<commit_msg>esilssa.rs: Since all flags are handled now, this code no longer needs to tolerate leftover MValType::Internals.<commit_after>//! Converts the from ESIL/CFG representation to SSA.

use std::collections::HashMap;
use petgraph::graph::NodeIndex;
use frontend::structs::LRegInfo;
use middle::cfg::NodeData as CFGNodeData;
use middle::cfg::EdgeType as CFGEdgeType;
use middle::cfg::{CFG, BasicBlock};
use middle::ssa::{BBInfo, SSA, SSAMod, ValueType};
use middle::ssa::verifier::Verify;
use middle::ssa::verifier::VerifiedAdd;
use middle::ir::{MVal, MInst, MOpcode, MValType};
use middle::phiplacement::PhiPlacer;
use middle::regfile::SubRegisterFile;
use middle::dce;

pub type VarId = usize;

pub struct SSAConstruction<'a, T> where
	T: Clone + Verify + SSAMod<BBInfo=BBInfo> + 'a
{
	pub phiplacer: PhiPlacer<'a, T>,
	pub regfile:   SubRegisterFile,
	pub temps:     HashMap<String, T::ValueRef>,
}

impl<'a, T> SSAConstruction<'a, T> where
	T: 'a + Clone + Verify +
       SSAMod<BBInfo=BBInfo, ValueRef=NodeIndex, ActionRef=NodeIndex>
{
	pub fn new(ssa: &'a mut T, reg_info: &LRegInfo) -> SSAConstruction<'a, T> {
		let mut sc = SSAConstruction {
			phiplacer: PhiPlacer::new(ssa),
			regfile:   SubRegisterFile::new(reg_info),
			temps:     HashMap::new(),
		};
		sc.phiplacer.add_variables(vec![
								   ValueType::Integer{width: 64}, // cur
								   ValueType::Integer{width: 64},  // old
								   ValueType::Integer{width: 64}  // lastsz
		]);
		// make the following a method of regfile?
		sc.phiplacer.add_variables(sc.regfile.whole_registers.clone());
		sc
	}

	pub fn run(&mut self, cfg: &CFG) {
		let mut blocks = Vec::<T::ActionRef>::with_capacity(cfg.g.node_count());
		let bb_iter = cfg.bbs.iter();

		for _ in 0..cfg.g.node_count() {
			blocks.push(self.phiplacer.ssa.invalid_action());
		}

		{
			// Insert the entry and exit blocks for the ssa.
			let block = self.phiplacer.add_block(BBInfo { addr: 0 });
			self.phiplacer.ssa.mark_start_node(&block);
			let zero = self.phiplacer.ssa.add_const(block, 0);
			self.phiplacer.write_variable(block, 0, zero); // cur = 0
			self.phiplacer.write_variable(block, 1, zero); // old = 0

			for (i, name) in self.regfile.whole_names.iter().enumerate() {
				let argnode = self.phiplacer.ssa.add_comment(block, self.regfile.whole_registers[i], name.clone());
				self.phiplacer.write_variable(block, i+2, argnode);
			}

			blocks[cfg.entry.index()] = block;

			let block = self.phiplacer.add_dynamic();
			self.phiplacer.ssa.mark_exit_node(&block);
			blocks[cfg.exit.index()] = block;
		}

		for (addr, i) in bb_iter {
			let block = self.phiplacer.add_block(BBInfo { addr: *addr });
			blocks[i.index()] = block;
			match cfg.g[*i] {
				CFGNodeData::Block(ref srcbb) => {
					self.process_block(block, srcbb);
				}
				_ => unreachable!(),
			}
		}

		for edge in cfg.g.raw_edges() {
			let i = match edge.weight.edge_type {
				CFGEdgeType::False => 0,
				CFGEdgeType::True => 1,
				CFGEdgeType::Unconditional => 2,
			};
			self.phiplacer.ssa.add_control_edge(
				blocks[edge.source().index()],
				blocks[edge.target().index()],
				i);
		}
		for &block in &blocks {
			self.phiplacer.seal_block(block);
		}

		//self.phiplacer.ssa.stable_indexing = false;
		//self.phiplacer.ssa.cleanup();

		{
			dce::collect(self.phiplacer.ssa);
		}
		//dce::collect(self.phiplacer.ssa, &[exit_regstate]);
		//self.phiplacer.ssa.cleanup();
	}

	fn process_in(&mut self, block: T::ActionRef, mval: &MVal) -> T::ValueRef {
		match mval.val_type {
			MValType::Register  => self.regfile.read_register(&mut self.phiplacer, 2, block, &mval.name),
			MValType::Temporary => self.temps[&mval.name],
			MValType::Internal  => panic!("This value type should be eliminated during parsing"),
			MValType::EsilCur   => self.phiplacer.read_variable(block, 0),
			MValType::EsilOld   => self.phiplacer.read_variable(block, 1),
			MValType::Lastsz    => self.phiplacer.read_variable(block, 2),
			MValType::Unknown   => self.phiplacer.ssa.invalid_value(),
			                       //self.phiplacer.ssa.add_comment(block, &"Unknown".to_string()), // unimplemented!()
			MValType::Null      => self.phiplacer.ssa.invalid_value(),
		}
	}

	fn process_out(&mut self, block: T::ActionRef, mval: &MVal, value: T::ValueRef) {
		match mval.val_type {
			MValType::Register  => self.regfile.write_register(&mut self.phiplacer, 2, block, &mval.name, value),
			MValType::Temporary => {self.temps.insert(mval.name.clone(), value);},
			MValType::Null      => {},
			_                   => panic!(),
		}
	}

	fn process_op(&mut self, block: T::ActionRef, inst: &MInst, n0: T::ValueRef, n1: T::ValueRef) -> T::ValueRef {
		if inst.opcode == MOpcode::OpEq {
			return n0
		}

		let dsttype = match inst.dst.val_type {
			MValType::Null => ValueType::Integer{width: 0}, // there is no ValueType::None?
			_              => ValueType::Integer{width: inst.dst.size},
		};

		let nn = {
			// let ref mut ssa = self.phiplacer.ssa;
			// let nn = ssa.add_op(block, inst.opcode, dsttype);
			// ssa.op_use(nn, 0, n0);
			// ssa.op_use(nn, 1, n1);
			// nn
			(*self.phiplacer.ssa).verified_add_op(block, inst.opcode, dsttype, &[n0, n1])
		};

		if inst.update_flags {
			let mut nn64 = nn;

			if inst.dst.size != 64 {
				let ref mut ssa = self.phiplacer.ssa;
				nn64 = ssa.add_op(block,
					MOpcode::OpWiden(64),
					ValueType::Integer{width: 64});
				ssa.op_use(nn64, 0, nn);
			}

			let old = self.phiplacer.read_variable(block, 0);
			self.phiplacer.write_variable(block, 1, old);
			self.phiplacer.write_variable(block, 0, nn64);
		}

		return nn
	}

	fn process_block(&mut self, block: T::ActionRef, source: &BasicBlock) {
		let mut machinestate = self.phiplacer.ssa.to_value(block);
		for ref instruction in &source.instructions {
			let n0 = self.process_in(block, &instruction.operand_1);
			let n1 = self.process_in(block, &instruction.operand_2);

			if let MOpcode::OpNarrow(i) = instruction.opcode {
				println!("Found narrow. Narrow to : {}, n0: {:?}, n1: {:?},
				op1: {}, op2: {}", i, n0, n1, instruction.operand_1, instruction.operand_2);
			}

			if instruction.opcode == MOpcode::OpJmp {
				// TODO: In case of static jumps, this is trivial and does not need a selector.
				// In case of dynamic jump, the jump targets have to be determined.
				//self.ssa.g.add_edge(block, n0, SSAEdgeData::DynamicControl(0));
				break;
			}

			if instruction.opcode == MOpcode::OpCJmp {
				println!("Identified a selector: {:?} for block {:?}", n0, block);
				self.phiplacer.ssa.mark_selector(n0, block);
				continue;
			}

			let nn = self.process_op(block, instruction, n0, n1);

			if instruction.opcode == MOpcode::OpLoad {
				self.phiplacer.ssa.op_use(nn, 3, machinestate);
			}
			if instruction.opcode == MOpcode::OpStore {
				self.phiplacer.ssa.op_use(nn, 3, machinestate);
				machinestate = nn;
			}

			self.process_out(block, &instruction.dst, nn);
		}
	}
}
<|endoftext|>


<commit_before>
/// An `Allocator` implements a particular memory allocation strategy.
pub trait Allocator {

    /// Allocate a new block of size `size` on alignment `align`.
    ///
    /// # Arguments:
    ///   - `size`: the amount of memory to allocate (in bytes)
    ///   - `align`: the alignment for the allocation request
    ///
    /// # Returns:
    ///   - `Some(*mut u8)` if the request was allocated successfully
    ///   - `None` if the allocator is out of memory or if the request was
    ///     invalid.
    unsafe fn allocate( &mut self
                      , size: usize
                      , align: usize)
                      -> Option<*mut u8>;

    /// Release an allocated block of memory.
    ///
    /// The `size` and `align` parameters _must_ be the same as the original
    /// size and alignment of the frame being deallocated, otherwise our
    /// heap will become corrupted.
    ///
    /// # Arguments:
    ///   - `frame`: a pointer to the block of memory to deallocate
    ///   - `size`: the size of the block being deallocated
    ///   - `align`: the alignment of the block being deallocated
    unsafe fn deallocate( &mut self
                        , frame: *mut u8
                        , size: usize, align: usize);

    /// Reallocate `old_frame` from `old_size` bytes to `new_size` bytes
    ///
    /// The `old_size` and `align` parameters _must_ be the same as the
    /// original size and alignment of the frame being reallocated, otherwise
    /// our heap will become corrupted.
    ///
    /// # Arguments:
    ///   - `old_frame`: a pointer to the frame to be reallocated
    ///   - `old_size`: the size (in bytes) of the frame being reallocated
    ///   - `new_size`: the size to reallocate the frame to.
    ///   - `align`: the alignment for the allocation request
    ///
    /// # Returns:
    ///   - `Some(*mut u8)` if the frame was reallocated successfully
    ///   - `None` if the allocator is out of memory or if the request was
    ///     invalid.
    // TODO: Optimization: check if the reallocation request fits in
    // the old frame and return immediately if it does
    unsafe fn reallocate( &mut self
                        , old_frame: *mut u8
                        , old_size: usize
                        , new_size: usize
                        , align: usize )
                        -> Option<*mut u8> {
        // First, attempt to allocate a new frame...
        self.allocate(new_size, align)
            .map(|new_frame| {
                // If a new frame was allocated, copy all the data from the
                // old frame into the new frame.
                ptr::copy(new_frame, old_frame, min(old_size, new_size));
                // Then we can deallocate the old frame
                self.deallocate(old_frame, old_size, align);
                new_frame
            })
    }

    unsafe fn zero_alloc( &mut self
                        , _size: usize
                        , _align: usize)
                        -> Option<*mut u8> {
        unimplemented!()
    }
}
<commit_msg>[kernel] start on making allocator API smarter<commit_after>use super::paging::{Page, PageRange};
use core::ops;

/// A borrowed handle on a frame with a specified lifetime.
///
/// This automatically deallocates the frame when the borrow's lifetime
/// ends. It also ensures that the borrow only lives as long as the allocator
/// that provided it, and that the borrow is dropped if the allocator is
/// dropped.
pub struct BorrowedFrame<'a, F, A>
where F: Page
    , A: FrameAllocator<F>
    , A: 'a {
    frame: F
  , allocator: &'a A
}

impl<'a, F, A> ops::Deref for BorrowedFrame<'a, F, A>
where F: Page
    , A: 'a {
    type Target = F;
    fn deref(&self) -> &F { &self.frame }
}

impl<'a, F, A> ops::DerefMut for BorrowedFrame<'a, F, A>
where F: Page
    , A: 'a {
    fn deref_mut(&mut self) -> &mut F { &mut self.frame }
}

impl<'a, F, A> Drop for BorrowedFrame<'a, F, A>
where F: Page
    , A: FrameAllocator<F>
    , A: 'a {
    fn drop(&mut self) {
        unsafe { self.allocator.deallocate(self.frame) }
    }
}

/// Identical to a `BorrowedFrame` but borrowing a range of `Frame`s.
pub struct BorrowedFrameRange<'a, F, A>
where F: Page
    , A: FrameAllocator<F>
    , A: 'a {
    range: PageRange<F>
  , allocator: &'a A
}

impl<'a, F, A> ops::Deref for BorrowedFrame<'a, F, A>
where F: Page
    , A: 'a {
    type Target = PageRange<F>;
    fn deref(&self) -> &Self::Target { &self.range }
}

impl<'a, F, A> ops::DerefMut for BorrowedFrame<'a, F, A>
where F: Page
    , A: 'a {
    fn deref_mut(&mut self) -> &mut Self::Target { &mut self.range }
}

impl<'a, F, A> Drop for BorrowedFrame<'a, F, A>
where F: Page
    , A: FrameAllocator<F>
    , A: 'a {
    fn drop(&mut self) {
        unsafe { self.allocator.deallocate_range(self.range) }
    }
}

pub trait FrameAllocator<Frame>
where Frame: Page {

    unsafe fn allocate(&self) -> Option<Frame>;
    unsafe fn deallocate(&self, frame: Frame);
    fn borrow(&self) -> Option<BorrowedFrame<Frame, Self>>;

    unsafe fn allocate_range(&self, num: usize) -> Option<PageRange<Frame>>;
    unsafe fn deallocate_range(&self, range: PageRange<Frame>);
    fn borrow_range(&self, num: usize)
                    -> Option<BorrowedFrameRange<Frame, Self>>;

}

/// An `Allocator` implements a particular memory allocation strategy.
pub trait Allocator {

    /// Allocate a new block of size `size` on alignment `align`.
    ///
    /// # Arguments:
    ///   - `size`: the amount of memory to allocate (in bytes)
    ///   - `align`: the alignment for the allocation request
    ///
    /// # Returns:
    ///   - `Some(*mut u8)` if the request was allocated successfully
    ///   - `None` if the allocator is out of memory or if the request was
    ///     invalid.
    unsafe fn allocate( &mut self
                      , size: usize
                      , align: usize)
                      -> Option<*mut u8>;

    /// Release an allocated block of memory.
    ///
    /// The `size` and `align` parameters _must_ be the same as the original
    /// size and alignment of the frame being deallocated, otherwise our
    /// heap will become corrupted.
    ///
    /// # Arguments:
    ///   - `frame`: a pointer to the block of memory to deallocate
    ///   - `size`: the size of the block being deallocated
    ///   - `align`: the alignment of the block being deallocated
    unsafe fn deallocate( &mut self
                        , frame: *mut u8
                        , size: usize, align: usize);

    /// Reallocate `old_frame` from `old_size` bytes to `new_size` bytes
    ///
    /// The `old_size` and `align` parameters _must_ be the same as the
    /// original size and alignment of the frame being reallocated, otherwise
    /// our heap will become corrupted.
    ///
    /// # Arguments:
    ///   - `old_frame`: a pointer to the frame to be reallocated
    ///   - `old_size`: the size (in bytes) of the frame being reallocated
    ///   - `new_size`: the size to reallocate the frame to.
    ///   - `align`: the alignment for the allocation request
    ///
    /// # Returns:
    ///   - `Some(*mut u8)` if the frame was reallocated successfully
    ///   - `None` if the allocator is out of memory or if the request was
    ///     invalid.
    // TODO: Optimization: check if the reallocation request fits in
    // the old frame and return immediately if it does
    unsafe fn reallocate( &mut self
                        , old_frame: *mut u8
                        , old_size: usize
                        , new_size: usize
                        , align: usize )
                        -> Option<*mut u8> {
        // First, attempt to allocate a new frame...
        self.allocate(new_size, align)
            .map(|new_frame| {
                // If a new frame was allocated, copy all the data from the
                // old frame into the new frame.
                ptr::copy(new_frame, old_frame, min(old_size, new_size));
                // Then we can deallocate the old frame
                self.deallocate(old_frame, old_size, align);
                new_frame
            })
    }

    unsafe fn zero_alloc( &mut self
                        , _size: usize
                        , _align: usize)
                        -> Option<*mut u8> {
        unimplemented!()
    }
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;

use libimagstore::storeid::StoreId;
use libimagstore::store::Entry;
use libimagstore::store::EntryHeader;
use libimagstore::store::Result as StoreResult;

use error::{LinkError, LinkErrorKind};
use result::Result;

use toml::Value;
use itertools::Itertools;

pub type Link = StoreId;

pub trait InternalLinker {

    /// Get the internal links from the implementor object
    fn get_internal_links(&self) -> Result<Vec<Link>>;

    /// Set the internal links for the implementor object
    fn set_internal_links(&mut self, links: Vec<&mut Entry>) -> Result<Vec<Link>>;

    /// Add an internal link to the implementor object
    fn add_internal_link(&mut self, link: &mut Entry) -> Result<()>;

    /// Remove an internal link from the implementor object
    fn remove_internal_link(&mut self, link: &mut Entry) -> Result<()>;

}

impl InternalLinker for Entry {

    fn get_internal_links(&self) -> Result<Vec<Link>> {
        process_rw_result(self.get_header().read("imag.links"))
    }

    /// Set the links in a header and return the old links, if any.
    fn set_internal_links(&mut self, links: Vec<&mut Entry>) -> Result<Vec<Link>> {
        let self_location = self.get_location().clone();
        let mut new_links = vec![];

        for link in links {
            if let Err(e) = add_foreign_link(link, self_location.clone()) {
                return Err(e);
            }
            let link = link.get_location().clone();
            new_links.push(link);
        }

        let new_links = links_into_values(new_links);
        if new_links.iter().any(|o| o.is_none()) {
            return Err(LinkError::new(LinkErrorKind::InternalConversionError, None));
        }
        let new_links = new_links.into_iter().map(|o| o.unwrap()).collect();
        process_rw_result(self.get_header_mut().set("imag.links", Value::Array(new_links)))
    }

    fn add_internal_link(&mut self, link: &mut Entry) -> Result<()> {
        let new_link = link.get_location().clone();

        add_foreign_link(link, self.get_location().clone())
            .and_then(|_| {
                self.get_internal_links()
                    .and_then(|mut links| {
                        links.push(new_link);
                        rewrite_links(self.get_header_mut(), links)
                    })
            })
    }

    fn remove_internal_link(&mut self, link: &mut Entry) -> Result<()> {
        let own_loc   = link.get_location().clone();
        let other_loc = link.get_location().clone();

        link.get_internal_links()
            .and_then(|links| {
                let links = links.into_iter().filter(|l| l.clone() != own_loc).collect();
                rewrite_links(self.get_header_mut(), links)
            })
            .and_then(|_| {
                self.get_internal_links()
                    .and_then(|links| {
                        let links = links.into_iter().filter(|l| l.clone() != other_loc).collect();
                        rewrite_links(link.get_header_mut(), links)
                    })
            })
    }

}

fn links_into_values(links: Vec<StoreId>) -> Vec<Option<Value>> {
    links
        .into_iter()
        .map(|s| s.to_str().map(String::from))
        .unique()
        .map(|elem| elem.map(Value::String))
        .sorted_by(|a, b| {
            match (a, b) {
                (&Some(Value::String(ref a)), &Some(Value::String(ref b))) => Ord::cmp(a, b),
                (&None, _) | (_, &None) => Ordering::Equal,
                _                                              => unreachable!()
            }
        })
}

fn rewrite_links(header: &mut EntryHeader, links: Vec<StoreId>) -> Result<()> {
    let links = links_into_values(links);

    if links.iter().any(|o| o.is_none()) {
        // if any type convert failed we fail as well
        Err(LinkError::new(LinkErrorKind::InternalConversionError, None))
    } else {
        // I know it is ugly
        let links = links.into_iter().map(|opt| opt.unwrap()).collect();
        let process = header.set("imag.links", Value::Array(links));
        process_rw_result(process).map(|_| ())
    }
}

/// When Linking A -> B, the specification wants us to link back B -> A.
/// This is a helper function which does this.
fn add_foreign_link(target: &mut Entry, from: StoreId) -> Result<()> {
    target.get_internal_links()
        .and_then(|mut links| {
            links.push(from);
            let links = links_into_values(links);
            if links.iter().any(|o| o.is_none()) {
                Err(LinkError::new(LinkErrorKind::InternalConversionError, None))
            } else {
                let links = links.into_iter().map(|opt| opt.unwrap()).collect();
                process_rw_result(target.get_header_mut().set("imag.links", Value::Array(links)))
                    .map(|_| ())
            }
        })
}

fn process_rw_result(links: StoreResult<Option<Value>>) -> Result<Vec<Link>> {
    if links.is_err() {
        debug!("RW action on store failed. Generating LinkError");
        let lerr  = LinkError::new(LinkErrorKind::EntryHeaderReadError,
                                   Some(Box::new(links.unwrap_err())));
        return Err(lerr);
    }
    let links = links.unwrap();

    if links.is_none() {
        debug!("We got no value from the header!");
        return Ok(vec![])
    }
    let links = links.unwrap();

    let links = {
        match links {
            Value::Array(a) => a,
            _ => {
                debug!("We expected an Array for the links, but there was a non-Array!");
                return Err(LinkError::new(LinkErrorKind::ExistingLinkTypeWrong, None));
            },
        }
    };

    if !links.iter().all(|l| match *l { Value::String(_) => true, _ => false }) {
        debug!("At least one of the Values which were expected in the Array of links is a non-String!");
        debug!("Generating LinkError");
        return Err(LinkError::new(LinkErrorKind::ExistingLinkTypeWrong, None));
    }

    let links : Vec<Link> = links.into_iter()
        .map(|link| {
            match link {
                Value::String(s) => StoreId::from(s),
                _ => unreachable!(),
            }
        })
        .collect();

    debug!("Ok, the RW action was successful, returning link vector now!");
    Ok(links)
}

<commit_msg>libimagentrylink: Refactor code with is_match!() macro<commit_after>use std::cmp::Ordering;

use libimagstore::storeid::StoreId;
use libimagstore::store::Entry;
use libimagstore::store::EntryHeader;
use libimagstore::store::Result as StoreResult;

use error::{LinkError, LinkErrorKind};
use result::Result;

use toml::Value;
use itertools::Itertools;

pub type Link = StoreId;

pub trait InternalLinker {

    /// Get the internal links from the implementor object
    fn get_internal_links(&self) -> Result<Vec<Link>>;

    /// Set the internal links for the implementor object
    fn set_internal_links(&mut self, links: Vec<&mut Entry>) -> Result<Vec<Link>>;

    /// Add an internal link to the implementor object
    fn add_internal_link(&mut self, link: &mut Entry) -> Result<()>;

    /// Remove an internal link from the implementor object
    fn remove_internal_link(&mut self, link: &mut Entry) -> Result<()>;

}

impl InternalLinker for Entry {

    fn get_internal_links(&self) -> Result<Vec<Link>> {
        process_rw_result(self.get_header().read("imag.links"))
    }

    /// Set the links in a header and return the old links, if any.
    fn set_internal_links(&mut self, links: Vec<&mut Entry>) -> Result<Vec<Link>> {
        let self_location = self.get_location().clone();
        let mut new_links = vec![];

        for link in links {
            if let Err(e) = add_foreign_link(link, self_location.clone()) {
                return Err(e);
            }
            let link = link.get_location().clone();
            new_links.push(link);
        }

        let new_links = links_into_values(new_links);
        if new_links.iter().any(|o| o.is_none()) {
            return Err(LinkError::new(LinkErrorKind::InternalConversionError, None));
        }
        let new_links = new_links.into_iter().map(|o| o.unwrap()).collect();
        process_rw_result(self.get_header_mut().set("imag.links", Value::Array(new_links)))
    }

    fn add_internal_link(&mut self, link: &mut Entry) -> Result<()> {
        let new_link = link.get_location().clone();

        add_foreign_link(link, self.get_location().clone())
            .and_then(|_| {
                self.get_internal_links()
                    .and_then(|mut links| {
                        links.push(new_link);
                        rewrite_links(self.get_header_mut(), links)
                    })
            })
    }

    fn remove_internal_link(&mut self, link: &mut Entry) -> Result<()> {
        let own_loc   = link.get_location().clone();
        let other_loc = link.get_location().clone();

        link.get_internal_links()
            .and_then(|links| {
                let links = links.into_iter().filter(|l| l.clone() != own_loc).collect();
                rewrite_links(self.get_header_mut(), links)
            })
            .and_then(|_| {
                self.get_internal_links()
                    .and_then(|links| {
                        let links = links.into_iter().filter(|l| l.clone() != other_loc).collect();
                        rewrite_links(link.get_header_mut(), links)
                    })
            })
    }

}

fn links_into_values(links: Vec<StoreId>) -> Vec<Option<Value>> {
    links
        .into_iter()
        .map(|s| s.to_str().map(String::from))
        .unique()
        .map(|elem| elem.map(Value::String))
        .sorted_by(|a, b| {
            match (a, b) {
                (&Some(Value::String(ref a)), &Some(Value::String(ref b))) => Ord::cmp(a, b),
                (&None, _) | (_, &None) => Ordering::Equal,
                _                                              => unreachable!()
            }
        })
}

fn rewrite_links(header: &mut EntryHeader, links: Vec<StoreId>) -> Result<()> {
    let links = links_into_values(links);

    if links.iter().any(|o| o.is_none()) {
        // if any type convert failed we fail as well
        Err(LinkError::new(LinkErrorKind::InternalConversionError, None))
    } else {
        // I know it is ugly
        let links = links.into_iter().map(|opt| opt.unwrap()).collect();
        let process = header.set("imag.links", Value::Array(links));
        process_rw_result(process).map(|_| ())
    }
}

/// When Linking A -> B, the specification wants us to link back B -> A.
/// This is a helper function which does this.
fn add_foreign_link(target: &mut Entry, from: StoreId) -> Result<()> {
    target.get_internal_links()
        .and_then(|mut links| {
            links.push(from);
            let links = links_into_values(links);
            if links.iter().any(|o| o.is_none()) {
                Err(LinkError::new(LinkErrorKind::InternalConversionError, None))
            } else {
                let links = links.into_iter().map(|opt| opt.unwrap()).collect();
                process_rw_result(target.get_header_mut().set("imag.links", Value::Array(links)))
                    .map(|_| ())
            }
        })
}

fn process_rw_result(links: StoreResult<Option<Value>>) -> Result<Vec<Link>> {
    if links.is_err() {
        debug!("RW action on store failed. Generating LinkError");
        let lerr  = LinkError::new(LinkErrorKind::EntryHeaderReadError,
                                   Some(Box::new(links.unwrap_err())));
        return Err(lerr);
    }
    let links = links.unwrap();

    if links.is_none() {
        debug!("We got no value from the header!");
        return Ok(vec![])
    }
    let links = links.unwrap();

    let links = {
        match links {
            Value::Array(a) => a,
            _ => {
                debug!("We expected an Array for the links, but there was a non-Array!");
                return Err(LinkError::new(LinkErrorKind::ExistingLinkTypeWrong, None));
            },
        }
    };

    if !links.iter().all(|l| is_match!(*l, Value::String(_))) {
        debug!("At least one of the Values which were expected in the Array of links is a non-String!");
        debug!("Generating LinkError");
        return Err(LinkError::new(LinkErrorKind::ExistingLinkTypeWrong, None));
    }

    let links : Vec<Link> = links.into_iter()
        .map(|link| {
            match link {
                Value::String(s) => StoreId::from(s),
                _ => unreachable!(),
            }
        })
        .collect();

    debug!("Ok, the RW action was successful, returning link vector now!");
    Ok(links)
}

<|endoftext|>