<commit_before>/// Macro that allows you to quickly define a handlebars helper by passing a
/// name and a closure. The closure arguments are mapped to helper parameters
/// one by one. Named argument with default value is also supported and mapped
/// to helper hash.
///
/// # Examples
///
/// ```rust
/// #[macro_use] extern crate handlebars;
/// #[macro_use] extern crate serde_json;
///
/// handlebars_helper!(is_above_10: |x: u64| x > 10);
/// handlebars_helper!(is_above: |x: u64, { compare: u64 = 10 }| x > compare);
///
/// # fn main() {
/// #
/// let mut handlebars = handlebars::Handlebars::new();
/// handlebars.register_helper("is-above-10", Box::new(is_above_10));
/// handlebars.register_helper("is-above", Box::new(is_above));
///
/// let result = handlebars
///     .render_template("{{#if (is-above-10 12)}}great!{{else}}okay{{/if}}", &json!({}))
///     .unwrap();
///  assert_eq!(&result, "great!");
/// let result2 = handlebars
///     .render_template("{{#if (is-above 12 compare=10)}}great!{{else}}okay{{/if}}", &json!({}))
///     .unwrap();
///  assert_eq!(&result2, "great!");
/// # }
/// ```

#[macro_export]
macro_rules! handlebars_helper {
    ($struct_name:ident: |$($name:ident: $tpe:tt),*
     $($(,)?{$($hash_name:ident: $hash_tpe:tt=$dft_val:literal),*})?
     $($(,)?*$args:ident)?
     $($(,)?**$kwargs:ident)?|
     $body:expr ) => {
        #[allow(non_camel_case_types)]
        pub struct $struct_name;

        impl $crate::HelperDef for $struct_name {
            #[allow(unused_assignments)]
            fn call_inner<'reg: 'rc, 'rc>(
                &self,
                h: &$crate::Helper<'reg, 'rc>,
                _: &'reg $crate::Handlebars<'reg>,
                _: &'rc $crate::Context,
                _: &mut $crate::RenderContext<'reg, 'rc>,
            ) -> Result<Option<$crate::ScopedJson<'reg, 'rc>>, $crate::RenderError> {
                let mut param_idx = 0;

                $(
                    let $name = h.param(param_idx)
                        .map(|x| x.value())
                        .ok_or_else(|| $crate::RenderError::new(&format!(
                            "`{}` helper: Couldn't read parameter {}",
                            stringify!($struct_name), stringify!($name),
                        )))
                        .and_then(|x|
                                  handlebars_helper!(@as_json_value x, $tpe)
                                  .ok_or_else(|| $crate::RenderError::new(&format!(
                                      "`{}` helper: Couldn't convert parameter {} to type `{}`. \
                                       It's {:?} as JSON. Got these params: {:?}",
                                      stringify!($struct_name), stringify!($name), stringify!($tpe),
                                      x, h.params(),
                                  )))
                        )?;
                    param_idx += 1;
                )*

                    $(
                        $(
                            let $hash_name = h.hash_get(stringify!($hash_name))
                                .map(|x| x.value())
                                .map(|x|
                                     handlebars_helper!(@as_json_value x, $hash_tpe)
                                     .ok_or_else(|| $crate::RenderError::new(&format!(
                                         "`{}` helper: Couldn't convert hash {} to type `{}`. \
                                          It's {:?} as JSON. Got these hash: {:?}",
                                         stringify!($struct_name), stringify!($hash_name), stringify!($hash_tpe),
                                         x, h.hash(),
                                     )))
                                )
                                .unwrap_or_else(|| Ok($dft_val))?;
                        )*
                    )?

                    $(let $args = h.params().iter().map(|x| x.value()).collect::<Vec<&serde_json::Value>>();)?
                    $(let $kwargs = h.hash().iter().map(|(k, v)| (k.to_owned(), v.value())).collect::<std::collections::BTreeMap<&str, &serde_json::Value>>();)?

                let result = $body;
                Ok(Some($crate::ScopedJson::Derived($crate::JsonValue::from(result))))
            }
        }
    };

    (@as_json_value $x:ident, object) => { $x.as_object() };
    (@as_json_value $x:ident, array) => { $x.as_array() };
    (@as_json_value $x:ident, str) => { $x.as_str() };
    (@as_json_value $x:ident, i64) => { $x.as_i64() };
    (@as_json_value $x:ident, u64) => { $x.as_u64() };
    (@as_json_value $x:ident, f64) => { $x.as_f64() };
    (@as_json_value $x:ident, bool) => { $x.as_bool() };
    (@as_json_value $x:ident, null) => { $x.as_null() };
    (@as_json_value $x:ident, Json) => { Some($x) };
}

/// This macro is defined if the `logging` feature is set.
///
/// It ignores all logging calls inside the library.
#[cfg(feature = "no_logging")]
#[macro_export]
macro_rules! debug {
    (target: $target:expr, $($arg:tt)*) => {};
    ($($arg:tt)*) => {};
}

/// This macro is defined if the `logging` feature is not set.
///
/// It ignores all logging calls inside the library.
#[cfg(feature = "no_logging")]
#[macro_export]
macro_rules! error {
    (target: $target:expr, $($arg:tt)*) => {};
    ($($arg:tt)*) => {};
}

/// This macro is defined if the `logging` feature is not set.
///
/// It ignores all logging calls inside the library.
#[cfg(feature = "no_logging")]
#[macro_export]
macro_rules! info {
    (target: $target:expr, $($arg:tt)*) => {};
    ($($arg:tt)*) => {};
}

/// This macro is defined if the `logging` feature is not set.
///
/// It ignores all logging calls inside the library.
#[cfg(feature = "no_logging")]
#[macro_export]
macro_rules! log {
    (target: $target:expr, $($arg:tt)*) => {};
    ($($arg:tt)*) => {};
}

/// This macro is defined if the `logging` feature is not set.
///
/// It ignores all logging calls inside the library.
#[cfg(feature = "no_logging")]
#[macro_export]
macro_rules! trace {
    (target: $target:expr, $($arg:tt)*) => {};
    ($($arg:tt)*) => {};
}

/// This macro is defined if the `logging` feature is not set.
///
/// It ignores all logging calls inside the library.
#[cfg(feature = "no_logging")]
#[macro_export]
macro_rules! warn {
    (target: $target:expr, $($arg:tt)*) => {};
    ($($arg:tt)*) => {};
}
<commit_msg>(doc) update docstring for macro<commit_after>/// Macro that allows you to quickly define a handlebars helper by passing a
/// name and a closure.
///
/// There are several types of arguments available to closure:
///
/// * Parameters are mapped to closure arguments one by one. Any declared
/// parameters are required
/// * Hash are mapped as named arguments and declared in a bracket block.
/// All named arguments are optional so default value is required.
/// * An optional `*args` provides a vector of all helper parameters.
/// * An optional `**kwargs` provides a map of all helper hash.
///
/// # Examples
///
/// ```rust
/// #[macro_use] extern crate handlebars;
/// #[macro_use] extern crate serde_json;
///
/// handlebars_helper!(is_above_10: |x: u64| x > 10);
/// handlebars_helper!(is_above: |x: u64, { compare: u64 = 10 }| x > compare);
///
/// # fn main() {
/// #
/// let mut handlebars = handlebars::Handlebars::new();
/// handlebars.register_helper("is-above-10", Box::new(is_above_10));
/// handlebars.register_helper("is-above", Box::new(is_above));
///
/// let result = handlebars
///     .render_template("{{#if (is-above-10 12)}}great!{{else}}okay{{/if}}", &json!({}))
///     .unwrap();
///  assert_eq!(&result, "great!");
/// let result2 = handlebars
///     .render_template("{{#if (is-above 12 compare=10)}}great!{{else}}okay{{/if}}", &json!({}))
///     .unwrap();
///  assert_eq!(&result2, "great!");
/// # }
/// ```

#[macro_export]
macro_rules! handlebars_helper {
    ($struct_name:ident: |$($name:ident: $tpe:tt),*
     $($(,)?{$($hash_name:ident: $hash_tpe:tt=$dft_val:literal),*})?
     $($(,)?*$args:ident)?
     $($(,)?**$kwargs:ident)?|
     $body:expr ) => {
        #[allow(non_camel_case_types)]
        pub struct $struct_name;

        impl $crate::HelperDef for $struct_name {
            #[allow(unused_assignments)]
            fn call_inner<'reg: 'rc, 'rc>(
                &self,
                h: &$crate::Helper<'reg, 'rc>,
                _: &'reg $crate::Handlebars<'reg>,
                _: &'rc $crate::Context,
                _: &mut $crate::RenderContext<'reg, 'rc>,
            ) -> Result<Option<$crate::ScopedJson<'reg, 'rc>>, $crate::RenderError> {
                let mut param_idx = 0;

                $(
                    let $name = h.param(param_idx)
                        .map(|x| x.value())
                        .ok_or_else(|| $crate::RenderError::new(&format!(
                            "`{}` helper: Couldn't read parameter {}",
                            stringify!($struct_name), stringify!($name),
                        )))
                        .and_then(|x|
                                  handlebars_helper!(@as_json_value x, $tpe)
                                  .ok_or_else(|| $crate::RenderError::new(&format!(
                                      "`{}` helper: Couldn't convert parameter {} to type `{}`. \
                                       It's {:?} as JSON. Got these params: {:?}",
                                      stringify!($struct_name), stringify!($name), stringify!($tpe),
                                      x, h.params(),
                                  )))
                        )?;
                    param_idx += 1;
                )*

                    $(
                        $(
                            let $hash_name = h.hash_get(stringify!($hash_name))
                                .map(|x| x.value())
                                .map(|x|
                                     handlebars_helper!(@as_json_value x, $hash_tpe)
                                     .ok_or_else(|| $crate::RenderError::new(&format!(
                                         "`{}` helper: Couldn't convert hash {} to type `{}`. \
                                          It's {:?} as JSON. Got these hash: {:?}",
                                         stringify!($struct_name), stringify!($hash_name), stringify!($hash_tpe),
                                         x, h.hash(),
                                     )))
                                )
                                .unwrap_or_else(|| Ok($dft_val))?;
                        )*
                    )?

                    $(let $args = h.params().iter().map(|x| x.value()).collect::<Vec<&serde_json::Value>>();)?
                    $(let $kwargs = h.hash().iter().map(|(k, v)| (k.to_owned(), v.value())).collect::<std::collections::BTreeMap<&str, &serde_json::Value>>();)?

                let result = $body;
                Ok(Some($crate::ScopedJson::Derived($crate::JsonValue::from(result))))
            }
        }
    };

    (@as_json_value $x:ident, object) => { $x.as_object() };
    (@as_json_value $x:ident, array) => { $x.as_array() };
    (@as_json_value $x:ident, str) => { $x.as_str() };
    (@as_json_value $x:ident, i64) => { $x.as_i64() };
    (@as_json_value $x:ident, u64) => { $x.as_u64() };
    (@as_json_value $x:ident, f64) => { $x.as_f64() };
    (@as_json_value $x:ident, bool) => { $x.as_bool() };
    (@as_json_value $x:ident, null) => { $x.as_null() };
    (@as_json_value $x:ident, Json) => { Some($x) };
}

/// This macro is defined if the `logging` feature is set.
///
/// It ignores all logging calls inside the library.
#[cfg(feature = "no_logging")]
#[macro_export]
macro_rules! debug {
    (target: $target:expr, $($arg:tt)*) => {};
    ($($arg:tt)*) => {};
}

/// This macro is defined if the `logging` feature is not set.
///
/// It ignores all logging calls inside the library.
#[cfg(feature = "no_logging")]
#[macro_export]
macro_rules! error {
    (target: $target:expr, $($arg:tt)*) => {};
    ($($arg:tt)*) => {};
}

/// This macro is defined if the `logging` feature is not set.
///
/// It ignores all logging calls inside the library.
#[cfg(feature = "no_logging")]
#[macro_export]
macro_rules! info {
    (target: $target:expr, $($arg:tt)*) => {};
    ($($arg:tt)*) => {};
}

/// This macro is defined if the `logging` feature is not set.
///
/// It ignores all logging calls inside the library.
#[cfg(feature = "no_logging")]
#[macro_export]
macro_rules! log {
    (target: $target:expr, $($arg:tt)*) => {};
    ($($arg:tt)*) => {};
}

/// This macro is defined if the `logging` feature is not set.
///
/// It ignores all logging calls inside the library.
#[cfg(feature = "no_logging")]
#[macro_export]
macro_rules! trace {
    (target: $target:expr, $($arg:tt)*) => {};
    ($($arg:tt)*) => {};
}

/// This macro is defined if the `logging` feature is not set.
///
/// It ignores all logging calls inside the library.
#[cfg(feature = "no_logging")]
#[macro_export]
macro_rules! warn {
    (target: $target:expr, $($arg:tt)*) => {};
    ($($arg:tt)*) => {};
}
<|endoftext|>


<commit_before>use super::Result;

#[derive(Debug,PartialEq)]
pub struct Memory {
    total: u64,
    free: u64,
    used: u64,
    buffers: u64,
    cached: u64,
    swap_total: u64,
    swap_free: u64,
    swap_used: u64,
}

const PROC_MEMORY_NUMBER_OF_FIELDS: usize = 6;
const SYS_MEMORY_NUMBER_OF_FIELDS: usize = 1;

impl Memory {
    /// Total amount of physical memory in Kb.
    pub fn total(&self) -> u64 {
        self.total
    }

    pub fn free(&self) -> u64 {
        self.free
    }

    /// Total amount of used physical memory in Kb.
    pub fn used(&self) -> u64 {
        self.used
    }

    /// Total amount of swap space in Kb.
    pub fn swap_total(&self) -> u64 {
        self.swap_total
    }

    /// Total amount of free swap space in Kb.
    pub fn swap_free(&self) -> u64 {
        self.swap_free
    }

    /// Total amount of used swap space in Kb.
    pub fn swap_used(&self) -> u64 {
        self.swap_used
    }
}

/// Read the current memory status of the system.
#[cfg(target_os = "linux")]
pub fn read() -> Result<Memory> {
    os::read()
}

#[cfg(target_os = "linux")]
mod os {
    use std::io::BufRead;
    use std::path::Path;

    use super::{Memory,PROC_MEMORY_NUMBER_OF_FIELDS,SYS_MEMORY_NUMBER_OF_FIELDS};
    use super::super::{ProbeError,Result,container};
    use super::super::{file_to_buf_reader,parse_u64};

    #[inline]
    pub fn read() -> Result<Memory> {
        if container::in_container() {
            read_and_parse_sys_memory(&Path::new("/sys/fs/cgroup/memory/"))
        } else {
            read_and_parse_proc_memory(&Path::new("/proc/memory.stat"))
        }
    }

    #[inline]
    pub fn read_and_parse_proc_memory(path: &Path) -> Result<Memory> {
        let mut memory = Memory {
            total: 0,
            free: 0,
            used: 0,
            buffers: 0,
            cached: 0,
            swap_total: 0,
            swap_free: 0,
            swap_used: 0,
        };
        let mut free = 0;

        let reader = try!(file_to_buf_reader(path));

        let mut fields_encountered = 0;
        for line in reader.lines() {
            let line = try!(line);
            let segments: Vec<&str> = line.split_whitespace().collect();
            let value: u64 = try!(parse_u64(segments[1]));

            // If this is a field we recognize set it's value and increment the
            // number of fields we encountered.
            fields_encountered += match segments[0] {
                "MemTotal:" => {
                    memory.total = value;
                    1
                },
                "MemFree:" => {
                    free = value;
                    1
                },
                "Buffers:" => {
                    memory.buffers = value;
                    1
                },
                "Cached:" => {
                    memory.cached = value;
                    1
                },
                "SwapTotal:" => {
                    memory.swap_total = value;
                    1
                },
                "SwapFree:" => {
                    memory.swap_free = value;
                    1
                },
                _ => 0
            };

            if fields_encountered == PROC_MEMORY_NUMBER_OF_FIELDS {
                break
            }
        }

        if fields_encountered != PROC_MEMORY_NUMBER_OF_FIELDS {
            return Err(ProbeError::UnexpectedContent("Did not encounter all expected fields".to_owned()))
        }

        // Total amount of free physical memory in Kb.
        // Includes buffers and caches, these will be freed
        // up by the OS when the memory is needed.
        memory.free = free + memory.buffers + memory.cached;
        memory.used = memory.total - memory.free;
        memory.swap_used = memory.swap_total - memory.swap_free;

        Ok(memory)
    }

    pub fn read_and_parse_sys_memory(path: &Path) -> Result<Memory> {
        let mut memory = Memory {
            total: 0,
            free: 0,
            used: 0,
            buffers: 0,
            cached: 0,
            swap_total: 0,
            swap_free: 0,
            swap_used: 0,
        };

        memory.total = bytes_to_kilo_bytes(try!(read_file_value_as_u64(&path.join("memory.limit_in_bytes"))));
        let used_memory = bytes_to_kilo_bytes(try!(read_file_value_as_u64(&path.join("memory.usage_in_bytes"))));

        let mut fields_encountered = 0;
        let reader = try!(file_to_buf_reader(&path.join("memory.stat")));
        for line in reader.lines() {
            let line = try!(line);
            let segments: Vec<&str> = line.split_whitespace().collect();
            let value = try!(parse_u64(&segments[1]));

            fields_encountered += match segments[0] {
                "cache" => {
                    memory.cached = bytes_to_kilo_bytes(value);
                    1
                },
                _ => 0
            };

            if fields_encountered == SYS_MEMORY_NUMBER_OF_FIELDS {
                break
            }
        }
        memory.used = used_memory - memory.cached;
        memory.free = memory.total - memory.used;

        // If swap is not configured for the container, read 0 as value
        memory.swap_total = match read_file_value_as_u64(&path.join("memory.memsw.limit_in_bytes")) {
            Ok(value) => bytes_to_kilo_bytes(value) - memory.total,
            Err(_) => 0
        };
        // If swap is not configured for the container, read 0 as value
        memory.swap_used = match read_file_value_as_u64(&path.join("memory.memsw.usage_in_bytes")) {
            Ok(value) => bytes_to_kilo_bytes(value) - used_memory,
            Err(_) => 0
        };
        memory.swap_free = memory.swap_total.checked_sub(memory.swap_used).unwrap_or(0);

        Ok(memory)
    }

    fn read_file_value_as_u64(path: &Path) -> Result<u64> {
        let mut reader = try!(file_to_buf_reader(path));
        let mut line = String::new();
        try!(reader.read_line(&mut line));
        parse_u64(&line.trim())
    }

    fn bytes_to_kilo_bytes(bytes: u64) -> u64 {
        bytes.checked_div(1024).unwrap_or(0)
    }
}

#[cfg(test)]
mod tests {
    use std::path::Path;
    use super::super::ProbeError;
    use super::Memory;

    #[test]
    fn test_read_memory() {
        assert!(super::read().is_ok());
    }

    #[test]
    fn test_read_and_parse_proc_memory() {
        let path = Path::new("fixtures/linux/memory/proc_meminfo");
        let memory = super::os::read_and_parse_proc_memory(&path).unwrap();

        let expected = Memory {
            total: 376072,
            free: 324248,
            used: 51824,
            buffers: 22820,
            cached: 176324,
            swap_total: 1101816,
            swap_free: 1100644,
            swap_used: 1172,
        };
        assert_eq!(expected, memory);
        assert_eq!(memory.total, memory.used + memory.free);
        assert_eq!(memory.swap_total, memory.swap_used + memory.swap_free);
    }

    #[test]
    fn test_read_and_parse_memory_wrong_path() {
        let path = Path::new("/nonsense");
        match super::os::read_and_parse_proc_memory(&path) {
            Err(ProbeError::IO(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_memory_incomplete() {
        let path = Path::new("fixtures/linux/memory/proc_meminfo_incomplete");
        match super::os::read_and_parse_proc_memory(&path) {
            Err(ProbeError::UnexpectedContent(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_memory_garbage() {
        let path = Path::new("fixtures/linux/memory/proc_meminfo_garbage");
        match super::os::read_and_parse_proc_memory(&path) {
            Err(ProbeError::UnexpectedContent(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_sys_memory() {
        let path = Path::new("fixtures/linux/sys/fs/cgroup/memory/");
        let memory = super::os::read_and_parse_sys_memory(&path).unwrap();

        let expected = Memory {
            total: 512000, // 500mb
            free: 503400, // total - used
            used: 8600,
            buffers: 0,
            cached: 58928,
            swap_total: 1_488_000, // reported swap total - reported memory total
            swap_free: 1_055_528,
            swap_used: 432_472, // reported swap used - (reported memory used, including cache)
        };
        assert_eq!(expected, memory);
        assert_eq!(memory.total, memory.used + memory.free);
        assert_eq!(memory.swap_total, memory.swap_used + memory.swap_free);
    }

    #[test]
    fn test_read_and_parse_sys_memory_wrong_path() {
        let path = Path::new("/nonsense");
        match super::os::read_and_parse_sys_memory(&path) {
            Err(ProbeError::IO(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_sys_memory_incomplete() {
        let path = Path::new("fixtures/linux/sys/fs/cgroup/memory_incomplete/");
        match super::os::read_and_parse_sys_memory(&path) {
            Err(ProbeError::UnexpectedContent(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_sys_memory_missing_files() {
        let path = Path::new("fixtures/linux/sys/fs/cgroup/memory_missing_files/");
        match super::os::read_and_parse_sys_memory(&path) {
            Err(ProbeError::IO(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_sys_memory_garbage() {
        let path = Path::new("fixtures/linux/sys/fs/cgroup/memory_garbage/");
        match super::os::read_and_parse_sys_memory(&path) {
            Err(ProbeError::UnexpectedContent(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_sys_memory_no_swap() {
        let path = Path::new("fixtures/linux/sys/fs/cgroup/memory_without_swap/");
        let memory = super::os::read_and_parse_sys_memory(&path).unwrap();

        let expected = Memory {
            total: 512000, // 500mb
            free: 503400, // total - used
            used: 8600,
            buffers: 0,
            cached: 58928,
            swap_total: 0, // Reads 0 swap
            swap_free: 0,  // Reads 0 swap
            swap_used: 0,
        };
        assert_eq!(expected, memory);
        assert_eq!(memory.total, memory.used + memory.free);
        assert_eq!(memory.swap_total, memory.swap_used + memory.swap_free);
    }
}
<commit_msg>Use checked_subs for swap calculations<commit_after>use super::Result;

#[derive(Debug,PartialEq)]
pub struct Memory {
    total: u64,
    free: u64,
    used: u64,
    buffers: u64,
    cached: u64,
    swap_total: u64,
    swap_free: u64,
    swap_used: u64,
}

const PROC_MEMORY_NUMBER_OF_FIELDS: usize = 6;
const SYS_MEMORY_NUMBER_OF_FIELDS: usize = 1;

impl Memory {
    /// Total amount of physical memory in Kb.
    pub fn total(&self) -> u64 {
        self.total
    }

    pub fn free(&self) -> u64 {
        self.free
    }

    /// Total amount of used physical memory in Kb.
    pub fn used(&self) -> u64 {
        self.used
    }

    /// Total amount of swap space in Kb.
    pub fn swap_total(&self) -> u64 {
        self.swap_total
    }

    /// Total amount of free swap space in Kb.
    pub fn swap_free(&self) -> u64 {
        self.swap_free
    }

    /// Total amount of used swap space in Kb.
    pub fn swap_used(&self) -> u64 {
        self.swap_used
    }
}

/// Read the current memory status of the system.
#[cfg(target_os = "linux")]
pub fn read() -> Result<Memory> {
    os::read()
}

#[cfg(target_os = "linux")]
mod os {
    use std::io::BufRead;
    use std::path::Path;

    use super::{Memory,PROC_MEMORY_NUMBER_OF_FIELDS,SYS_MEMORY_NUMBER_OF_FIELDS};
    use super::super::{ProbeError,Result,container};
    use super::super::{file_to_buf_reader,parse_u64};

    #[inline]
    pub fn read() -> Result<Memory> {
        if container::in_container() {
            read_and_parse_sys_memory(&Path::new("/sys/fs/cgroup/memory/"))
        } else {
            read_and_parse_proc_memory(&Path::new("/proc/memory.stat"))
        }
    }

    #[inline]
    pub fn read_and_parse_proc_memory(path: &Path) -> Result<Memory> {
        let mut memory = Memory {
            total: 0,
            free: 0,
            used: 0,
            buffers: 0,
            cached: 0,
            swap_total: 0,
            swap_free: 0,
            swap_used: 0,
        };
        let mut free = 0;

        let reader = try!(file_to_buf_reader(path));

        let mut fields_encountered = 0;
        for line in reader.lines() {
            let line = try!(line);
            let segments: Vec<&str> = line.split_whitespace().collect();
            let value: u64 = try!(parse_u64(segments[1]));

            // If this is a field we recognize set it's value and increment the
            // number of fields we encountered.
            fields_encountered += match segments[0] {
                "MemTotal:" => {
                    memory.total = value;
                    1
                },
                "MemFree:" => {
                    free = value;
                    1
                },
                "Buffers:" => {
                    memory.buffers = value;
                    1
                },
                "Cached:" => {
                    memory.cached = value;
                    1
                },
                "SwapTotal:" => {
                    memory.swap_total = value;
                    1
                },
                "SwapFree:" => {
                    memory.swap_free = value;
                    1
                },
                _ => 0
            };

            if fields_encountered == PROC_MEMORY_NUMBER_OF_FIELDS {
                break
            }
        }

        if fields_encountered != PROC_MEMORY_NUMBER_OF_FIELDS {
            return Err(ProbeError::UnexpectedContent("Did not encounter all expected fields".to_owned()))
        }

        // Total amount of free physical memory in Kb.
        // Includes buffers and caches, these will be freed
        // up by the OS when the memory is needed.
        memory.free = free + memory.buffers + memory.cached;
        memory.used = memory.total - memory.free;
        memory.swap_used = memory.swap_total - memory.swap_free;

        Ok(memory)
    }

    pub fn read_and_parse_sys_memory(path: &Path) -> Result<Memory> {
        let mut memory = Memory {
            total: 0,
            free: 0,
            used: 0,
            buffers: 0,
            cached: 0,
            swap_total: 0,
            swap_free: 0,
            swap_used: 0,
        };

        memory.total = bytes_to_kilo_bytes(try!(read_file_value_as_u64(&path.join("memory.limit_in_bytes"))));
        let used_memory = bytes_to_kilo_bytes(try!(read_file_value_as_u64(&path.join("memory.usage_in_bytes"))));

        let mut fields_encountered = 0;
        let reader = try!(file_to_buf_reader(&path.join("memory.stat")));
        for line in reader.lines() {
            let line = try!(line);
            let segments: Vec<&str> = line.split_whitespace().collect();
            let value = try!(parse_u64(&segments[1]));

            fields_encountered += match segments[0] {
                "cache" => {
                    memory.cached = bytes_to_kilo_bytes(value);
                    1
                },
                _ => 0
            };

            if fields_encountered == SYS_MEMORY_NUMBER_OF_FIELDS {
                break
            }
        }
        memory.used = used_memory - memory.cached;
        memory.free = memory.total - memory.used;

        // If swap is not configured for the container, read 0 as value
        memory.swap_total = match read_file_value_as_u64(&path.join("memory.memsw.limit_in_bytes")) {
            Ok(value) => bytes_to_kilo_bytes(value).checked_sub(memory.total).unwrap_or(0),
            Err(_) => 0
        };
        // If swap is not configured for the container, read 0 as value
        memory.swap_used = match read_file_value_as_u64(&path.join("memory.memsw.usage_in_bytes")) {
            Ok(value) => bytes_to_kilo_bytes(value).checked_sub(used_memory).unwrap_or(0),
            Err(_) => 0
        };
        memory.swap_free = memory.swap_total.checked_sub(memory.swap_used).unwrap_or(0);

        Ok(memory)
    }

    fn read_file_value_as_u64(path: &Path) -> Result<u64> {
        let mut reader = try!(file_to_buf_reader(path));
        let mut line = String::new();
        try!(reader.read_line(&mut line));
        parse_u64(&line.trim())
    }

    fn bytes_to_kilo_bytes(bytes: u64) -> u64 {
        bytes.checked_div(1024).unwrap_or(0)
    }
}

#[cfg(test)]
mod tests {
    use std::path::Path;
    use super::super::ProbeError;
    use super::Memory;

    #[test]
    fn test_read_memory() {
        assert!(super::read().is_ok());
    }

    #[test]
    fn test_read_and_parse_proc_memory() {
        let path = Path::new("fixtures/linux/memory/proc_meminfo");
        let memory = super::os::read_and_parse_proc_memory(&path).unwrap();

        let expected = Memory {
            total: 376072,
            free: 324248,
            used: 51824,
            buffers: 22820,
            cached: 176324,
            swap_total: 1101816,
            swap_free: 1100644,
            swap_used: 1172,
        };
        assert_eq!(expected, memory);
        assert_eq!(memory.total, memory.used + memory.free);
        assert_eq!(memory.swap_total, memory.swap_used + memory.swap_free);
    }

    #[test]
    fn test_read_and_parse_memory_wrong_path() {
        let path = Path::new("/nonsense");
        match super::os::read_and_parse_proc_memory(&path) {
            Err(ProbeError::IO(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_memory_incomplete() {
        let path = Path::new("fixtures/linux/memory/proc_meminfo_incomplete");
        match super::os::read_and_parse_proc_memory(&path) {
            Err(ProbeError::UnexpectedContent(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_memory_garbage() {
        let path = Path::new("fixtures/linux/memory/proc_meminfo_garbage");
        match super::os::read_and_parse_proc_memory(&path) {
            Err(ProbeError::UnexpectedContent(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_sys_memory() {
        let path = Path::new("fixtures/linux/sys/fs/cgroup/memory/");
        let memory = super::os::read_and_parse_sys_memory(&path).unwrap();

        let expected = Memory {
            total: 512000, // 500mb
            free: 503400, // total - used
            used: 8600,
            buffers: 0,
            cached: 58928,
            swap_total: 1_488_000, // reported swap total - reported memory total
            swap_free: 1_055_528,
            swap_used: 432_472, // reported swap used - (reported memory used, including cache)
        };
        assert_eq!(expected, memory);
        assert_eq!(memory.total, memory.used + memory.free);
        assert_eq!(memory.swap_total, memory.swap_used + memory.swap_free);
    }

    #[test]
    fn test_read_and_parse_sys_memory_wrong_path() {
        let path = Path::new("/nonsense");
        match super::os::read_and_parse_sys_memory(&path) {
            Err(ProbeError::IO(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_sys_memory_incomplete() {
        let path = Path::new("fixtures/linux/sys/fs/cgroup/memory_incomplete/");
        match super::os::read_and_parse_sys_memory(&path) {
            Err(ProbeError::UnexpectedContent(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_sys_memory_missing_files() {
        let path = Path::new("fixtures/linux/sys/fs/cgroup/memory_missing_files/");
        match super::os::read_and_parse_sys_memory(&path) {
            Err(ProbeError::IO(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_sys_memory_garbage() {
        let path = Path::new("fixtures/linux/sys/fs/cgroup/memory_garbage/");
        match super::os::read_and_parse_sys_memory(&path) {
            Err(ProbeError::UnexpectedContent(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_parse_sys_memory_no_swap() {
        let path = Path::new("fixtures/linux/sys/fs/cgroup/memory_without_swap/");
        let memory = super::os::read_and_parse_sys_memory(&path).unwrap();

        let expected = Memory {
            total: 512000, // 500mb
            free: 503400, // total - used
            used: 8600,
            buffers: 0,
            cached: 58928,
            swap_total: 0, // Reads 0 swap
            swap_free: 0,  // Reads 0 swap
            swap_used: 0,
        };
        assert_eq!(expected, memory);
        assert_eq!(memory.total, memory.used + memory.free);
        assert_eq!(memory.swap_total, memory.swap_used + memory.swap_free);
    }
}
<|endoftext|>


<commit_before>use gpu;
use rom;
use timer;

pub struct Memory {
    pub rom: rom::Rom,
    pub gpu: gpu::Gpu,
    pub timer: timer::Timer,
    ram: [u8, ..0x4000],
    hram: [u8, ..0x79],
    ie: u8
}

impl Memory {
    pub fn new(rom: rom::Rom, gpu: gpu::Gpu, timer: timer::Timer) -> Memory {
        Memory {
            rom: rom,
            gpu: gpu,
            timer: timer,
            ram: [0x00, ..0x4000],
            hram: [0x00, ..0x79],
            ie: 0x00
        }
    }

    pub fn read_byte(&self, address: u16) -> u8 {
        match address {
            0x0000..0x7FFF => self.rom.bytes[address as uint],
            0x8000..0x9FFF => self.gpu.read_byte(address),
            0xA000..0xBFFF => {
                print!("External RAM not implemented!");
                0x00
            },
            0xC000..0xDFFF => self.ram[address as uint - 0xC000],
            0xE000..0xFDFF => self.ram[address as uint - 0xE000],
            0xFE00..0xFE9F => self.gpu.read_byte(address),
            0xFEA0..0xFEFF => fail!("0xFEA0..0xFEFF segment is no usable"),
            0xFF00..0xFF7F => {
                match address {
                    0xFF04..0xFF07 => {
                        self.timer.read_byte(address)
                    },
                    0xFF40..0xFF4B => {
                        self.gpu.read_byte(address)
                    },
                    _ => {
                        print!("Implement I/O ports stuff... ");
                        0x00
                    }
                }
            },
            0xFF80..0xFFFE => self.hram[address as uint - 0xFF80],
            0xFFFF => self.ie,
            _ => unreachable!()
        }
    }

    pub fn write_byte(&mut self, address: u16, value: u8) {
        match address {
            0x0000..0x7FFF => {
                fail!("Address not writable in MBC0: {:#06x}", address);
            },
            0x8000..0x9FFF => {
                self.gpu.write_byte(address, value);
            },
            0xA000..0xBFFF => {
                print!("External RAM not implemented!");
            }
            0xC000..0xDFFF => {
                self.ram[address as uint - 0xC000] = value;
            },
            0xE000..0xFDFF => {
                self.ram[address as uint - 0xE000] = value;
            },
            0xFE00..0xFE9F => {
                self.gpu.write_byte(address, value);
            },
            0xFEA0..0xFEFF => fail!("0xFEA0..0xFEFF segment is no usable"),
            0xFF00..0xFF7F => {
                match address {
                    0xFF04..0xFF07 => {
                        self.timer.write_byte(address, value);
                    },
                    0xFF40..0xFF4B => {
                        self.gpu.write_byte(address, value);
                    },
                    _ => print!("Implement I/O ports stuff... ")
                }
            },
            0xFF80..0xFFFE => {
                self.hram[address as uint - 0xFF80] = value;
            },
            0xFFFF => {
                self.ie = value;
            },
            _ => unreachable!()
        }
    }

    pub fn read_word(&self, address: u16) -> u16 {
        self.read_byte(address) as u16 |
            self.read_byte(address + 1) as u16 << 8
    }

    pub fn write_word(&mut self, address: u16, value: u16) {
        self.write_byte(address, (value & 0xFF) as u8);
        self.write_byte(address + 1, (value >> 8) as u8);
    }
}
<commit_msg>Add IF to memory<commit_after>use gpu;
use rom;
use timer;

pub struct Memory {
    pub rom: rom::Rom,
    pub gpu: gpu::Gpu,
    pub timer: timer::Timer,
    ram: [u8, ..0x4000],
    hram: [u8, ..0x79],
    if_: u8,
    ie: u8
}

impl Memory {
    pub fn new(rom: rom::Rom, gpu: gpu::Gpu, timer: timer::Timer) -> Memory {
        Memory {
            rom: rom,
            gpu: gpu,
            timer: timer,
            ram: [0x00, ..0x4000],
            hram: [0x00, ..0x79],
            if_: 0x00,
            ie: 0x00
        }
    }

    pub fn read_byte(&self, address: u16) -> u8 {
        match address {
            0x0000..0x7FFF => self.rom.bytes[address as uint],
            0x8000..0x9FFF => self.gpu.read_byte(address),
            0xA000..0xBFFF => {
                print!("External RAM not implemented!");
                0x00
            },
            0xC000..0xDFFF => self.ram[address as uint - 0xC000],
            0xE000..0xFDFF => self.ram[address as uint - 0xE000],
            0xFE00..0xFE9F => self.gpu.read_byte(address),
            0xFEA0..0xFEFF => fail!("0xFEA0..0xFEFF segment is no usable"),
            0xFF00..0xFF7F => {
                match address {
                    0xFF04..0xFF07 => {
                        self.timer.read_byte(address)
                    },
                    0xFF0F => {
                        self.if_
                    },
                    0xFF40..0xFF4B => {
                        self.gpu.read_byte(address)
                    },
                    _ => {
                        print!("Implement I/O ports stuff... ");
                        0x00
                    }
                }
            },
            0xFF80..0xFFFE => self.hram[address as uint - 0xFF80],
            0xFFFF => self.ie,
            _ => unreachable!()
        }
    }

    pub fn write_byte(&mut self, address: u16, value: u8) {
        match address {
            0x0000..0x7FFF => {
                fail!("Address not writable in MBC0: {:#06x}", address);
            },
            0x8000..0x9FFF => {
                self.gpu.write_byte(address, value);
            },
            0xA000..0xBFFF => {
                print!("External RAM not implemented!");
            }
            0xC000..0xDFFF => {
                self.ram[address as uint - 0xC000] = value;
            },
            0xE000..0xFDFF => {
                self.ram[address as uint - 0xE000] = value;
            },
            0xFE00..0xFE9F => {
                self.gpu.write_byte(address, value);
            },
            0xFEA0..0xFEFF => fail!("0xFEA0..0xFEFF segment is no usable"),
            0xFF00..0xFF7F => {
                match address {
                    0xFF04..0xFF07 => {
                        self.timer.write_byte(address, value);
                    },
                    0xFF0F => {
                        self.if_ = value;
                    },
                    0xFF40..0xFF4B => {
                        self.gpu.write_byte(address, value);
                    },
                    _ => print!("Implement I/O ports stuff... ")
                }
            },
            0xFF80..0xFFFE => {
                self.hram[address as uint - 0xFF80] = value;
            },
            0xFFFF => {
                self.ie = value;
            },
            _ => unreachable!()
        }
    }

    pub fn read_word(&self, address: u16) -> u16 {
        self.read_byte(address) as u16 |
            self.read_byte(address + 1) as u16 << 8
    }

    pub fn write_word(&mut self, address: u16, value: u16) {
        self.write_byte(address, (value & 0xFF) as u8);
        self.write_byte(address + 1, (value >> 8) as u8);
    }
}
<|endoftext|>


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

//!
//! Merkle Tree implementation that uses the Roughtime leaf and node tweak values.
//!

use ring::digest;

use super::{TREE_LEAF_TWEAK, TREE_NODE_TWEAK};

type Data = Vec<u8>;
type Hash = Data;

///
/// Merkle Tree implementation that uses the Roughtime leaf and node tweak values.
///
pub struct MerkleTree {
    levels: Vec<Vec<Data>>,
    algorithm: &'static digest::Algorithm,
}

impl MerkleTree {
    ///
    /// Create a new empty Merkle Tree based on SHA-512/256
    /// **TODO** lobby to have "SHA512[0:32] replaced with SHA-512/256"
    ///
    pub fn new_sha512_256() -> MerkleTree {
        MerkleTree {
            levels: vec![vec![]],
            algorithm: &digest::SHA512_256,
        }
    }

    ///
    /// Create a new empty Merkle Tree based on SHA-512
    ///
    pub fn new_sha512() -> MerkleTree {
        MerkleTree {
            levels: vec![vec![]],
            algorithm: &digest::SHA512,
        }
    }

    pub fn push_leaf(&mut self, data: &[u8]) {
        assert_eq!(
            data.len(),
            self.algorithm.output_len,
            "nonce vs hash length inconsistency"
        );

        let hash = self.hash_leaf(data);
        self.levels[0].push(hash);
    }

    pub fn get_paths(&self, mut index: usize) -> Vec<u8> {
        let mut paths = Vec::with_capacity(self.levels.len() * self.algorithm.output_len);
        let mut level = 0;

        while !self.levels[level].is_empty() {
            let sibling = if index % 2 == 0 { index + 1 } else { index - 1 };

            paths.extend(self.levels[level][sibling].clone());
            level += 1;
            index /= 2;
        }
        paths
    }

    pub fn compute_root(&mut self) -> Hash {
        assert!(
            !self.levels[0].is_empty(),
            "Must have at least one leaf to hash!"
        );

        let mut level = 0;
        let mut node_count = self.levels[0].len();

        while node_count > 1 {
            level += 1;

            if self.levels.len() < level + 1 {
                self.levels.push(vec![]);
            }

            if node_count % 2 != 0 {
                self.levels[level - 1].push(vec![0; self.algorithm.output_len]);
                node_count += 1;
            }

            node_count /= 2;

            for i in 0..node_count {
                let hash = self.hash_nodes(
                    &self.levels[level - 1][i * 2],
                    &self.levels[level - 1][(i * 2) + 1],
                );
                self.levels[level].push(hash);
            }
        }

        assert_eq!(self.levels[level].len(), 1);
        self.levels[level].pop().unwrap()
    }

    pub fn reset(&mut self) {
        for level in &mut self.levels {
            level.clear();
        }
    }

    fn hash_leaf(&self, leaf: &[u8]) -> Data {
        self.hash(&[TREE_LEAF_TWEAK, leaf])
    }

    fn hash_nodes(&self, first: &[u8], second: &[u8]) -> Data {
        self.hash(&[TREE_NODE_TWEAK, first, second])
    }

    fn hash(&self, to_hash: &[&[u8]]) -> Data {
        let mut ctx = digest::Context::new(self.algorithm);
        for data in to_hash {
            ctx.update(data);
        }
        Data::from(ctx.finish().as_ref())
    }

    pub fn root_from_paths(&self, mut index: usize, data: &[u8], paths: &[u8]) -> Hash {
        let mut hash = {
            let mut ctx = digest::Context::new(self.algorithm);
            ctx.update(TREE_LEAF_TWEAK);
            ctx.update(data);
            Hash::from(ctx.finish().as_ref())
        };

        assert_eq!(paths.len() % self.algorithm.output_len, 0);

        for path in paths.chunks(self.algorithm.output_len) {
            let mut ctx = digest::Context::new(self.algorithm);
            ctx.update(TREE_NODE_TWEAK);

            if index & 1 == 0 {
                // Left
                ctx.update(&hash);
                ctx.update(path);
            } else {
                // Right
                ctx.update(path);
                ctx.update(&hash);
            }

            hash = Hash::from(ctx.finish().as_ref());
            index >>= 1;
        }

        hash
    }
}

#[cfg(test)]
mod test {
    use crate::merkle::*;

    fn test_paths_with_num(num: usize) {
        for mut merkle_impl in [MerkleTree::new_sha512_256(), MerkleTree::new_sha512()] {
            for i in 0..num {
                merkle_impl.push_leaf(&[i as u8]);
            }

            let root = merkle_impl.compute_root();

            for i in 0..num {
                println!("Testing {:?} {:?} {:?}", merkle_impl.algorithm, num, i);
                let paths: Vec<u8> = merkle_impl.get_paths(i);
                let computed_root = merkle_impl.root_from_paths(i, &[i as u8], &paths);

                assert_eq!(root, computed_root);
            }
        }
    }

    #[test]
    fn power_of_two() {
        test_paths_with_num(2);
        test_paths_with_num(4);
        test_paths_with_num(8);
        test_paths_with_num(16);
    }

    #[test]
    fn not_power_of_two() {
        test_paths_with_num(1);
        test_paths_with_num(20);
    }
}
<commit_msg>Back out over-eager assert<commit_after>// Copyright 2017-2022 int08h LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//!
//! Merkle Tree implementation that uses the Roughtime leaf and node tweak values.
//!

use ring::digest;

use super::{TREE_LEAF_TWEAK, TREE_NODE_TWEAK};

type Data = Vec<u8>;
type Hash = Data;

///
/// Merkle Tree implementation that uses the Roughtime leaf and node tweak values.
///
pub struct MerkleTree {
    levels: Vec<Vec<Data>>,
    algorithm: &'static digest::Algorithm,
}

impl MerkleTree {
    ///
    /// Create a new empty Merkle Tree based on SHA-512/256
    /// **TODO** lobby to have "SHA512[0:32] replaced with SHA-512/256"
    ///
    pub fn new_sha512_256() -> MerkleTree {
        MerkleTree {
            levels: vec![vec![]],
            algorithm: &digest::SHA512_256,
        }
    }

    ///
    /// Create a new empty Merkle Tree based on SHA-512
    ///
    pub fn new_sha512() -> MerkleTree {
        MerkleTree {
            levels: vec![vec![]],
            algorithm: &digest::SHA512,
        }
    }

    pub fn push_leaf(&mut self, data: &[u8]) {
        let hash = self.hash_leaf(data);
        self.levels[0].push(hash);
    }

    pub fn get_paths(&self, mut index: usize) -> Vec<u8> {
        let mut paths = Vec::with_capacity(self.levels.len() * self.algorithm.output_len);
        let mut level = 0;

        while !self.levels[level].is_empty() {
            let sibling = if index % 2 == 0 { index + 1 } else { index - 1 };

            paths.extend(self.levels[level][sibling].clone());
            level += 1;
            index /= 2;
        }
        paths
    }

    pub fn compute_root(&mut self) -> Hash {
        assert!(
            !self.levels[0].is_empty(),
            "Must have at least one leaf to hash!"
        );

        let mut level = 0;
        let mut node_count = self.levels[0].len();

        while node_count > 1 {
            level += 1;

            if self.levels.len() < level + 1 {
                self.levels.push(vec![]);
            }

            if node_count % 2 != 0 {
                self.levels[level - 1].push(vec![0; self.algorithm.output_len]);
                node_count += 1;
            }

            node_count /= 2;

            for i in 0..node_count {
                let hash = self.hash_nodes(
                    &self.levels[level - 1][i * 2],
                    &self.levels[level - 1][(i * 2) + 1],
                );
                self.levels[level].push(hash);
            }
        }

        assert_eq!(self.levels[level].len(), 1);
        self.levels[level].pop().unwrap()
    }

    pub fn reset(&mut self) {
        for level in &mut self.levels {
            level.clear();
        }
    }

    fn hash_leaf(&self, leaf: &[u8]) -> Data {
        self.hash(&[TREE_LEAF_TWEAK, leaf])
    }

    fn hash_nodes(&self, first: &[u8], second: &[u8]) -> Data {
        self.hash(&[TREE_NODE_TWEAK, first, second])
    }

    fn hash(&self, to_hash: &[&[u8]]) -> Data {
        let mut ctx = digest::Context::new(self.algorithm);
        for data in to_hash {
            ctx.update(data);
        }
        Data::from(ctx.finish().as_ref())
    }

    pub fn root_from_paths(&self, mut index: usize, data: &[u8], paths: &[u8]) -> Hash {
        let mut hash = {
            let mut ctx = digest::Context::new(self.algorithm);
            ctx.update(TREE_LEAF_TWEAK);
            ctx.update(data);
            Hash::from(ctx.finish().as_ref())
        };

        assert_eq!(paths.len() % self.algorithm.output_len, 0);

        for path in paths.chunks(self.algorithm.output_len) {
            let mut ctx = digest::Context::new(self.algorithm);
            ctx.update(TREE_NODE_TWEAK);

            if index & 1 == 0 {
                // Left
                ctx.update(&hash);
                ctx.update(path);
            } else {
                // Right
                ctx.update(path);
                ctx.update(&hash);
            }

            hash = Hash::from(ctx.finish().as_ref());
            index >>= 1;
        }

        hash
    }
}

#[cfg(test)]
mod test {
    use crate::merkle::*;

    fn test_paths_with_num(num: usize) {
        for mut merkle_impl in [MerkleTree::new_sha512_256(), MerkleTree::new_sha512()] {
            for i in 0..num {
                merkle_impl.push_leaf(&[i as u8]);
            }

            let root = merkle_impl.compute_root();

            for i in 0..num {
                println!("Testing {:?} {:?} {:?}", merkle_impl.algorithm, num, i);
                let paths: Vec<u8> = merkle_impl.get_paths(i);
                let computed_root = merkle_impl.root_from_paths(i, &[i as u8], &paths);

                assert_eq!(root, computed_root);
            }
        }
    }

    #[test]
    fn power_of_two() {
        test_paths_with_num(2);
        test_paths_with_num(4);
        test_paths_with_num(8);
        test_paths_with_num(16);
    }

    #[test]
    fn not_power_of_two() {
        test_paths_with_num(1);
        test_paths_with_num(20);
    }
}
<|endoftext|>


<commit_before>use serde::{Serialize, Serializer};
use std::collections::HashMap;

#[derive(Serialize, PartialEq, Debug)]
pub struct Replay {
    pub header_size: u32,
    pub header_crc: u32,
    pub major_version: u32,
    pub minor_version: u32,
    pub game_type: String,

    #[serde(serialize_with = "pair_vec")]
    pub properties: Vec<(String, HeaderProp)>,
    pub content_size: u32,
    pub content_crc: u32,
    pub levels: Vec<String>,
    pub keyframes: Vec<KeyFrame>,
    pub debug_info: Vec<DebugInfo>,
    pub tick_marks: Vec<TickMark>,
    pub packages: Vec<String>,
    pub objects: Vec<String>,
    pub names: Vec<String>,
    pub class_indices: Vec<ClassIndex>,
    pub net_cache: Vec<ClassNetCache>,
}

#[derive(Serialize, PartialEq, Debug)]
pub struct TickMark {
    pub description: String,
    pub frame: u32,
}

#[derive(Serialize, PartialEq, Debug)]
pub struct KeyFrame {
    pub time: f32,
    pub frame: u32,
    pub position: u32,
}

#[derive(PartialEq, Debug)]
pub enum HeaderProp {
    Array(Vec<Vec<(String, HeaderProp)>>),
    Bool(bool),
    Byte,
    Float(f32),
    Int(u32),
    Name(String),
    QWord(u64),
    Str(String),
}

#[derive(Serialize, PartialEq, Debug)]
pub struct DebugInfo {
    pub frame: u32,
    pub user: String,
    pub text: String,
}

#[derive(Serialize, PartialEq, Debug)]
pub struct ClassIndex {
    pub class: String,
    pub index: u32,
}

#[derive(Serialize, PartialEq, Debug)]
pub struct CacheProp {
    pub index: u32,
    pub id: u32,
}

#[derive(Serialize, PartialEq, Debug)]
pub struct ClassNetCache {
    pub index: u32,
    pub parent_id: u32,
    pub id: u32,
    pub properties: Vec<CacheProp>,
}

/// Serialize a vector of key value tuples into a map. This is useful when the data we're ingesting
/// (rocket league replay data) doesn't have a defined spec, so it may be assuming too much to
/// store it into an associative array, so it's stored as a normal sequence. Here we serialize as a
/// map structure because most replay parser do this, so we should be compliant and the data format
/// doesn't dictate that the keys in a sequence of key value pairs must be distinct. It's true,
/// JSON doesn't need the keys to be unique: http://stackoverflow.com/q/21832701/433785
fn pair_vec<K, V, S>(inp: &[(K, V)], serializer: &mut S) -> Result<(), S::Error>
    where K: Serialize,
          V: Serialize,
          S: Serializer
{
    let mut state = try!(serializer.serialize_map(Some(inp.len())));
    for &(ref key, ref val) in inp.iter() {
        try!(serializer.serialize_map_key(&mut state, key));
        try!(serializer.serialize_map_value(&mut state, val));
    }
    return serializer.serialize_map_end(state);
}

/// By default serde will generate a serialization method that writes out the enum as well as the
/// enum value. Since header values are self describing in JSON, we do not need to serialize the
/// enum type. This is slightly lossy as in the serialized format it will be ambiguous if a value
/// is a `Name` or `Str`, as well as `Byte`, `Float`, `Int`, or `QWord`.
impl Serialize for HeaderProp {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error>
        where S: Serializer
    {
        match *self {
            HeaderProp::Array(ref x) => {
                let mut state = try!(serializer.serialize_seq(Some(x.len())));
                for inner in x {
                    // Look for a better way to do this instead of allocating the intermediate map
                    let mut els = HashMap::new();
                    for &(ref key, ref val) in inner.iter() {
                        els.insert(key, val);
                    }
                    try!(serializer.serialize_seq_elt(&mut state, els));
                }
                return serializer.serialize_seq_end(state);
            }
            HeaderProp::Bool(ref x) => serializer.serialize_bool(*x),
            HeaderProp::Byte => serializer.serialize_u8(0),
            HeaderProp::Float(ref x) => serializer.serialize_f32(*x),
            HeaderProp::Int(ref x) => serializer.serialize_u32(*x),
            HeaderProp::QWord(ref x) => serializer.serialize_u64(*x),
            HeaderProp::Name(ref x) | HeaderProp::Str(ref x) => serializer.serialize_str(x)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde;
    use std;
    use serde_json;

    fn to_json<T: serde::Serialize>(input: &T) -> std::string::String {
        return serde_json::to_string(input).unwrap();
    }

    #[test]
    fn serialize_header_array() {
        let data = vec![
            vec![
                ("frame".to_string(), HeaderProp::Int(441)),
                ("PlayerName".to_string(), HeaderProp::Str("rust is awesome".to_string()))
            ], vec![
                ("frame".to_string(), HeaderProp::Int(1738)),
                ("PlayerName".to_string(), HeaderProp::Str("rusty".to_string()))
            ]
          ];
        let actual = HeaderProp::Array(data);
        assert_eq!(to_json(&actual), "[{\"PlayerName\":\"rust is awesome\",\"frame\":441},{\"PlayerName\":\"rusty\",\"frame\":1738}]");
    }

    #[test]
    fn serialize_header_bool() {
        assert_eq!(to_json(&HeaderProp::Bool(false)), "false");
        assert_eq!(to_json(&HeaderProp::Bool(true)), "true");
    }

    #[test]
    fn serialize_header_numbers() {
        assert_eq!(to_json(&HeaderProp::Byte), "0");
        assert_eq!(to_json(&HeaderProp::QWord(10)), "10");
        assert_eq!(to_json(&HeaderProp::Float(10.2)), "10.2");
        assert_eq!(to_json(&HeaderProp::Int(11)), "11");
    }

    #[test]
    fn serialize_header_str() {
        let val = "hello world";
        assert_eq!(to_json(&HeaderProp::Str(val.to_string())), "\"hello world\"");
        assert_eq!(to_json(&HeaderProp::Name(val.to_string())), "\"hello world\"");
    }
}
<commit_msg>Add structure comments to models.rs<commit_after>//! # Models
//!
//! Here lies the data structures that a rocket league replay is decoded into. All of the models
//! are contained in this one file because of serde. On Rust nightly, which is what this library
//! needs to compile, no special actions are needed for serde to work. However, to achieve
//! compilation on Rust stable, an extra [build step can be
//! introduced](https://serde.rs/codegen-stable.html). Part of the intermediate step is to split
//! all the serde specific code into a separate file, which is what I've done here. To me, it is
//! not critical to have this library compile on stable, hence this work remains unfinished.
//!
//! For serde, we only care about serialization, JSON serialization. Deserialization is not
//! implemented from our JSON output because it is lossy (JSON isn't the best with different
//! numeric/string types). Asking "why JSON" would be next logical step, and that other rocket
//! league replay parsers (like Octane) use JSON; however, the output of this library is not
//! compatible with that of other rocket league replay parsers.

use serde::{Serialize, Serializer};
use std::collections::HashMap;

/// The structure that a rocket league replay is parsed into.
#[derive(Serialize, PartialEq, Debug)]
pub struct Replay {
    pub header_size: u32,
    pub header_crc: u32,
    pub major_version: u32,
    pub minor_version: u32,
    pub game_type: String,

    /// Could use a map to represent properties but I don't want to assume that duplicate keys
    /// can't exist, so to be safe, use a traditional vector.
    #[serde(serialize_with = "pair_vec")]
    pub properties: Vec<(String, HeaderProp)>,
    pub content_size: u32,
    pub content_crc: u32,
    pub levels: Vec<String>,
    pub keyframes: Vec<KeyFrame>,
    pub debug_info: Vec<DebugInfo>,
    pub tick_marks: Vec<TickMark>,
    pub packages: Vec<String>,
    pub objects: Vec<String>,
    pub names: Vec<String>,
    pub class_indices: Vec<ClassIndex>,
    pub net_cache: Vec<ClassNetCache>,
}

/// In Rocket league replays, there are tickmarks that typically represent a significant event in
/// the game (eg. a goal). The tick mark is placed before the event happens so there is a ramp-up
/// time. For instance, a tickmark could be at frame 396 for a goal 441. At 30 fps, this would be
/// 1.5 seconds of ramp up time.
#[derive(Serialize, PartialEq, Debug)]
pub struct TickMark {
    pub description: String,
    pub frame: u32,
}

/// Keyframes as defined by the video compression section in the [wikipedia][] article, are the
/// main frames that are derived from in the following frame data. Since we are not decoding the
/// network stream, this is more a nice-to-decode than a necessity
///
/// [wikipedia]: https://en.wikipedia.org/wiki/Key_frame#Video_compression
#[derive(Serialize, PartialEq, Debug)]
pub struct KeyFrame {
    pub time: f32,
    pub frame: u32,
    pub position: u32,
}

/// All the interesting data are stored as properties in the header, properties such as:
/// - When and who scored a goal
/// - Player stats (goals, assists, score, etc).
/// - Date and level played on
/// A property can be a number, string, or a more complex object such as an array containing
/// additional properties.
#[derive(PartialEq, Debug)]
pub enum HeaderProp {
    Array(Vec<Vec<(String, HeaderProp)>>),
    Bool(bool),
    Byte,
    Float(f32),
    Int(u32),
    Name(String),
    QWord(u64),
    Str(String),
}

/// Debugging info stored in the replay if debugging is enabled.
#[derive(Serialize, PartialEq, Debug)]
pub struct DebugInfo {
    pub frame: u32,
    pub user: String,
    pub text: String,
}

/// Contains useful information when decoding the network stream, which we aren't
#[derive(Serialize, PartialEq, Debug)]
pub struct ClassIndex {
    pub class: String,
    pub index: u32,
}

/// Contains useful information when decoding the network stream, which we aren't
#[derive(Serialize, PartialEq, Debug)]
pub struct CacheProp {
    pub index: u32,
    pub id: u32,
}

/// Contains useful information when decoding the network stream, which we aren't
#[derive(Serialize, PartialEq, Debug)]
pub struct ClassNetCache {
    pub index: u32,
    pub parent_id: u32,
    pub id: u32,
    pub properties: Vec<CacheProp>,
}

/// Serialize a vector of key value tuples into a map. This is useful when the data we're ingesting
/// (rocket league replay data) doesn't have a defined spec, so it may be assuming too much to
/// store it into an associative array, so it's stored as a normal sequence. Here we serialize as a
/// map structure because most replay parser do this, so we should be compliant and the data format
/// doesn't dictate that the keys in a sequence of key value pairs must be distinct. It's true,
/// JSON doesn't need the keys to be unique: http://stackoverflow.com/q/21832701/433785
fn pair_vec<K, V, S>(inp: &[(K, V)], serializer: &mut S) -> Result<(), S::Error>
    where K: Serialize,
          V: Serialize,
          S: Serializer
{
    let mut state = try!(serializer.serialize_map(Some(inp.len())));
    for &(ref key, ref val) in inp.iter() {
        try!(serializer.serialize_map_key(&mut state, key));
        try!(serializer.serialize_map_value(&mut state, val));
    }
    return serializer.serialize_map_end(state);
}

/// By default serde will generate a serialization method that writes out the enum as well as the
/// enum value. Since header values are self describing in JSON, we do not need to serialize the
/// enum type. This is slightly lossy as in the serialized format it will be ambiguous if a value
/// is a `Name` or `Str`, as well as `Byte`, `Float`, `Int`, or `QWord`.
impl Serialize for HeaderProp {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error>
        where S: Serializer
    {
        match *self {
            HeaderProp::Array(ref x) => {
                let mut state = try!(serializer.serialize_seq(Some(x.len())));
                for inner in x {
                    // Look for a better way to do this instead of allocating the intermediate map
                    let mut els = HashMap::new();
                    for &(ref key, ref val) in inner.iter() {
                        els.insert(key, val);
                    }
                    try!(serializer.serialize_seq_elt(&mut state, els));
                }
                return serializer.serialize_seq_end(state);
            }
            HeaderProp::Bool(ref x) => serializer.serialize_bool(*x),
            HeaderProp::Byte => serializer.serialize_u8(0),
            HeaderProp::Float(ref x) => serializer.serialize_f32(*x),
            HeaderProp::Int(ref x) => serializer.serialize_u32(*x),
            HeaderProp::QWord(ref x) => serializer.serialize_u64(*x),
            HeaderProp::Name(ref x) | HeaderProp::Str(ref x) => serializer.serialize_str(x)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde;
    use std;
    use serde_json;

    fn to_json<T: serde::Serialize>(input: &T) -> std::string::String {
        return serde_json::to_string(input).unwrap();
    }

    #[test]
    fn serialize_header_array() {
        let data = vec![
            vec![
                ("frame".to_string(), HeaderProp::Int(441)),
                ("PlayerName".to_string(), HeaderProp::Str("rust is awesome".to_string()))
            ], vec![
                ("frame".to_string(), HeaderProp::Int(1738)),
                ("PlayerName".to_string(), HeaderProp::Str("rusty".to_string()))
            ]
          ];
        let actual = HeaderProp::Array(data);
        assert_eq!(to_json(&actual), "[{\"PlayerName\":\"rust is awesome\",\"frame\":441},{\"PlayerName\":\"rusty\",\"frame\":1738}]");
    }

    #[test]
    fn serialize_header_bool() {
        assert_eq!(to_json(&HeaderProp::Bool(false)), "false");
        assert_eq!(to_json(&HeaderProp::Bool(true)), "true");
    }

    #[test]
    fn serialize_header_numbers() {
        assert_eq!(to_json(&HeaderProp::Byte), "0");
        assert_eq!(to_json(&HeaderProp::QWord(10)), "10");
        assert_eq!(to_json(&HeaderProp::Float(10.2)), "10.2");
        assert_eq!(to_json(&HeaderProp::Int(11)), "11");
    }

    #[test]
    fn serialize_header_str() {
        let val = "hello world";
        assert_eq!(to_json(&HeaderProp::Str(val.to_string())), "\"hello world\"");
        assert_eq!(to_json(&HeaderProp::Name(val.to_string())), "\"hello world\"");
    }
}
<|endoftext|>


<commit_before>use libc::{c_char, c_uint};
use ffi::prelude::{LLVMValueRef, LLVMModuleRef};
use ffi::analysis::LLVMVerifierFailureAction;
use ffi::{analysis, core, linker, LLVMModule};
use ffi::transforms::pass_manager_builder as builder;
use ffi::bit_writer as writer;
use ffi::bit_reader as reader;
use cbox::{CBox, CSemiBox};
use std::ffi::CString;
use std::iter::{Iterator, IntoIterator};
use std::io::{Error, ErrorKind};
use std::io::Result as IoResult;
use std::{env, fmt, mem};
use std::marker::PhantomData;
use std::path::Path;
use std::process::Command;
use buffer::MemoryBuffer;
use context::{Context, GetContext};
use value::{Function, Value};
use ty::Type;
use util;

/// Represents a single compilation unit of code.
///
/// This is attached to the lifetime of the context that constructs it, but is owned by the `CSemiBox`.
pub struct Module;
native_ref!(&Module = LLVMModuleRef);
impl Module {
    /// Create a new module in the context given with the name given.
    ///
    /// The lifetime of the module will match the lifetime of the context
    /// you instance it in because the context contains it.
    ///
    /// ```rust
    /// use llvm::*;
    /// let context = Context::new();
    /// let context = context.as_semi();
    /// let module = Module::new("name", &context);
    /// println!("{:?}", module)
    /// ```
    pub fn new<'a>(name: &str, context: &'a Context) -> CSemiBox<'a, Module> {
        let c_name = CString::new(name).unwrap();
        unsafe { CSemiBox::new(core::LLVMModuleCreateWithNameInContext(c_name.as_ptr(), context.into())) }
    }
    /// Add a global to the module with the given type and name.
    pub fn add_global<'a>(&'a self, name: &str, ty: &'a Type) -> &'a Value {
        util::with_cstr(name, |ptr| unsafe {
            core::LLVMAddGlobal(self.into(), ty.into(), ptr).into()
        })
    }
    /// Add a constant global to the module with the given type, name and value.
    pub fn add_global_constant<'a>(&'a self, name: &str, val: &'a Value) -> &'a Value {
        util::with_cstr(name, |ptr| unsafe {
            let global = core::LLVMAddGlobal(self.into(), val.get_type().into(), ptr);
            core::LLVMSetInitializer (global.into(), val.into());
            global.into()
        })
    }
    /// Get the global with the name given, or `None` if no global with that name exists.
    pub fn get_global<'a>(&'a self, name: &str) -> Option<&'a Value> {
        util::with_cstr(name, |ptr| unsafe {
            let ptr = core::LLVMGetNamedGlobal(self.into(), ptr);
            util::ptr_to_null(ptr)
        })
    }
    /// Parse this bitcode file into a module, or return an error string.
    pub fn parse_bitcode<'a>(context: &'a Context, path: &str) -> Result<CSemiBox<'a, Module>, CBox<str>> {
        unsafe {
            let mut out = mem::uninitialized();
            let mut err = mem::uninitialized();
            let buf = try!(MemoryBuffer::new_from_file(path));
            if reader::LLVMParseBitcodeInContext(context.into(), buf.as_ptr(), &mut out, &mut err) == 1 {
                Err(CBox::new(err))
            } else {
                Ok(CSemiBox::new(out))
            }
        }
    }
    /// Write this module's bitcode to the path given.
    pub fn write_bitcode(&self, path: &str) -> IoResult<()> {
        util::with_cstr(path, |cpath| unsafe {
            if writer::LLVMWriteBitcodeToFile(self.into(), cpath) != 0 {
                Err(Error::new(ErrorKind::Other, &format!("could not write to {}", path) as &str))
            } else {
                Ok(())
            }
        })
    }
    /// Add a function to the module with the name given.
    pub fn add_function<'a>(&'a self, name: &str, sig: &'a Type) -> &'a mut Function {
        let c_name = CString::new(name).unwrap();
        unsafe { core::LLVMAddFunction(self.into(), c_name.as_ptr(), sig.into()) }.into()
    }
    /// Returns the function with the name given, or `None` if no function with that name exists.
    pub fn get_function<'a>(&'a self, name: &str) -> Option<&'a Function> {
        let c_name = CString::new(name).unwrap();
        unsafe {
            let ty = core::LLVMGetNamedFunction(self.into(), c_name.as_ptr());
            util::ptr_to_null(ty)
        }
    }
    /// Returns the type with the name given, or `None`` if no type with that name exists.
    pub fn get_type<'a>(&'a self, name: &str) -> Option<&'a Type> {
        let c_name = CString::new(name).unwrap();
        unsafe {
            let ty = core::LLVMGetTypeByName(self.into(), c_name.as_ptr());
            util::ptr_to_null(ty)
        }
    }
    /// Clone this module.
    pub fn clone<'a>(&'a self) -> CSemiBox<'a, Module> {
        CSemiBox::new(unsafe { core::LLVMCloneModule(self.into()) })
    }

    /// Optimize this module with the given optimization level and size level.
    ///
    /// This runs passes depending on the levels given.
    pub fn optimize(&self, opt_level: usize, size_level: usize) {
        unsafe {
            let builder = builder::LLVMPassManagerBuilderCreate();
            builder::LLVMPassManagerBuilderSetOptLevel(builder, opt_level as c_uint);
            builder::LLVMPassManagerBuilderSetSizeLevel(builder, size_level as c_uint);
            let pass_manager = core::LLVMCreatePassManager();
            builder::LLVMPassManagerBuilderPopulateModulePassManager(builder, pass_manager);
            builder::LLVMPassManagerBuilderDispose(builder);
            core::LLVMRunPassManager(pass_manager, self.into());
        }
    }

    /// Returns the target data of this module represented as a string
    pub fn get_target(&self) -> &str {
        unsafe {
            let target = core::LLVMGetTarget(self.into());
            util::to_str(target as *mut c_char)
        }
    }

    /// Set the target data of this module to the target data string given.
    pub fn set_target(&self, target: &str) {
        let c_target = CString::new(target).unwrap();
        unsafe { core::LLVMSetTarget(self.into(), c_target.as_ptr()) }
    }

    /// Verify that the module is safe to run, returning a string detailing the error
    /// when an error occurs.
    pub fn verify(&self) -> Result<(), CBox<str>> {
        unsafe {
            let mut error = mem::uninitialized();
            let action = LLVMVerifierFailureAction::LLVMReturnStatusAction;
            if analysis::LLVMVerifyModule(self.into(), action, &mut error) == 1 {
                Err(CBox::new(error))
            } else {
                Ok(())
            }
        }
    }

    /// Compile the module into an object file at the given location.
    ///
    /// Note that this uses the LLVM tool `llc` to do this, which may or may not be
    /// installed on the user's machine.
    pub fn compile(&self, path: &Path, opt_level: usize) -> IoResult<()> {
        let dir = env::temp_dir();
        let path = path.to_str().unwrap();
        let mod_path = dir.join("module.bc");
        let mod_path = mod_path.to_str().unwrap();
        try!(self.write_bitcode(mod_path));
        Command::new("llc")
            .arg(&format!("-O={}", opt_level))
            .arg("-filetype=obj")
            .arg("-o").arg(path)
            .arg(mod_path)
            .spawn()
            .map(|_| ())
    }

    /// Link a module into this module, returning an error string if an error occurs.
    ///
    /// This *does not* destroy the source module.
    pub fn link(&self, src: &Module) -> Result<(), CBox<str>> {
        unsafe {
            let dest = self.into();
            let src = src.into();
            let mut message = mem::uninitialized();
            if linker::LLVMLinkModules(dest, src, linker::LLVMLinkerMode::LLVMLinkerPreserveSource, &mut message) == 1 {
                Err(CBox::new(message))
            } else {
                Ok(())
            }
        }
    }

    /// Link a module into this module, returning an error string if an error occurs.
    ///
    /// This *does* destroy the source module.
    pub fn link_destroy(&self, src: CSemiBox<Module>) -> Result<(), CBox<str>> {
        unsafe {
            let dest = self.into();
            let src = src.as_ptr();
            let mut message = mem::uninitialized();
            if linker::LLVMLinkModules(dest, src, linker::LLVMLinkerMode::LLVMLinkerDestroySource, &mut message) == 1 {
                Err(CBox::new(message))
            } else {
                Ok(())
            }
        }
    }
}
impl<'a> IntoIterator for &'a Module {
    type Item = &'a Function;
    type IntoIter = Functions<'a>;
    /// Iterate through the functions in the module
    fn into_iter(self) -> Functions<'a> {
        Functions {
            value: unsafe { core::LLVMGetFirstFunction(self.into()) },
            marker: PhantomData
        }
    }
}
get_context!(Module, LLVMGetModuleContext);
to_str!(Module, LLVMPrintModuleToString);
dispose!(Module, LLVMModule, core::LLVMDisposeModule);
#[derive(Copy, Clone)]
/// An iterator through the functions contained in a module.
pub struct Functions<'a> {
    value: LLVMValueRef,
    marker: PhantomData<&'a ()>
}
impl<'a> Iterator for Functions<'a> {
    type Item = &'a Function;
    fn next(&mut self) -> Option<&'a Function> {
        if self.value.is_null() {
            None
        } else {
            let c_next = unsafe { core::LLVMGetNextFunction(self.value) };
            self.value = c_next;
            Some(self.value.into())
        }
    }
}
<commit_msg>Add parse_ir to Module.<commit_after>use libc::{c_char, c_uint};
use ffi::prelude::{LLVMValueRef, LLVMModuleRef};
use ffi::analysis::LLVMVerifierFailureAction;
use ffi::{analysis, core, linker, LLVMModule};
use ffi::transforms::pass_manager_builder as builder;
use ffi::bit_writer as writer;
use ffi::bit_reader as reader;
use ffi::ir_reader as ir_reader;
use cbox::{CBox, CSemiBox};
use std::ffi::CString;
use std::iter::{Iterator, IntoIterator};
use std::io::{Error, ErrorKind};
use std::io::Result as IoResult;
use std::{env, fmt, mem};
use std::marker::PhantomData;
use std::path::Path;
use std::process::Command;
use buffer::MemoryBuffer;
use context::{Context, GetContext};
use value::{Function, Value};
use ty::Type;
use util;

/// Represents a single compilation unit of code.
///
/// This is attached to the lifetime of the context that constructs it, but is owned by the `CSemiBox`.
pub struct Module;
native_ref!(&Module = LLVMModuleRef);
impl Module {
    /// Create a new module in the context given with the name given.
    ///
    /// The lifetime of the module will match the lifetime of the context
    /// you instance it in because the context contains it.
    ///
    /// ```rust
    /// use llvm::*;
    /// let context = Context::new();
    /// let context = context.as_semi();
    /// let module = Module::new("name", &context);
    /// println!("{:?}", module)
    /// ```
    pub fn new<'a>(name: &str, context: &'a Context) -> CSemiBox<'a, Module> {
        let c_name = CString::new(name).unwrap();
        unsafe { CSemiBox::new(core::LLVMModuleCreateWithNameInContext(c_name.as_ptr(), context.into())) }
    }
    /// Add a global to the module with the given type and name.
    pub fn add_global<'a>(&'a self, name: &str, ty: &'a Type) -> &'a Value {
        util::with_cstr(name, |ptr| unsafe {
            core::LLVMAddGlobal(self.into(), ty.into(), ptr).into()
        })
    }
    /// Add a constant global to the module with the given type, name and value.
    pub fn add_global_constant<'a>(&'a self, name: &str, val: &'a Value) -> &'a Value {
        util::with_cstr(name, |ptr| unsafe {
            let global = core::LLVMAddGlobal(self.into(), val.get_type().into(), ptr);
            core::LLVMSetInitializer (global.into(), val.into());
            global.into()
        })
    }
    /// Get the global with the name given, or `None` if no global with that name exists.
    pub fn get_global<'a>(&'a self, name: &str) -> Option<&'a Value> {
        util::with_cstr(name, |ptr| unsafe {
            let ptr = core::LLVMGetNamedGlobal(self.into(), ptr);
            util::ptr_to_null(ptr)
        })
    }

    pub fn parse_ir<'a>(context: &'a Context, path: &str) -> Result<CSemiBox<'a, Module>, CBox<str>> {
    	unsafe {
            let mut out = mem::uninitialized();
            let mut err = mem::uninitialized();
            let buf = try!(MemoryBuffer::new_from_file(path));
            if ir_reader::LLVMParseIRInContext(context.into(), buf.as_ptr(), &mut out, &mut err) == 1 {
            	Err(CBox::new(err))
            } else {
              Ok(CSemiBox::new(out))
            }
      }
    }

    /// Parse this bitcode file into a module, or return an error string.
    pub fn parse_bitcode<'a>(context: &'a Context, path: &str) -> Result<CSemiBox<'a, Module>, CBox<str>> {
        unsafe {
            let mut out = mem::uninitialized();
            let mut err = mem::uninitialized();
            let buf = try!(MemoryBuffer::new_from_file(path));
            if reader::LLVMParseBitcodeInContext(context.into(), buf.as_ptr(), &mut out, &mut err) == 1 {
                Err(CBox::new(err))
            } else {
                Ok(CSemiBox::new(out))
            }
        }
    }
    /// Write this module's bitcode to the path given.
    pub fn write_bitcode(&self, path: &str) -> IoResult<()> {
        util::with_cstr(path, |cpath| unsafe {
            if writer::LLVMWriteBitcodeToFile(self.into(), cpath) != 0 {
                Err(Error::new(ErrorKind::Other, &format!("could not write to {}", path) as &str))
            } else {
                Ok(())
            }
        })
    }
    /// Add a function to the module with the name given.
    pub fn add_function<'a>(&'a self, name: &str, sig: &'a Type) -> &'a mut Function {
        let c_name = CString::new(name).unwrap();
        unsafe { core::LLVMAddFunction(self.into(), c_name.as_ptr(), sig.into()) }.into()
    }
    /// Returns the function with the name given, or `None` if no function with that name exists.
    pub fn get_function<'a>(&'a self, name: &str) -> Option<&'a Function> {
        let c_name = CString::new(name).unwrap();
        unsafe {
            let ty = core::LLVMGetNamedFunction(self.into(), c_name.as_ptr());
            util::ptr_to_null(ty)
        }
    }
    /// Returns the type with the name given, or `None`` if no type with that name exists.
    pub fn get_type<'a>(&'a self, name: &str) -> Option<&'a Type> {
        let c_name = CString::new(name).unwrap();
        unsafe {
            let ty = core::LLVMGetTypeByName(self.into(), c_name.as_ptr());
            util::ptr_to_null(ty)
        }
    }
    /// Clone this module.
    pub fn clone<'a>(&'a self) -> CSemiBox<'a, Module> {
        CSemiBox::new(unsafe { core::LLVMCloneModule(self.into()) })
    }

    /// Optimize this module with the given optimization level and size level.
    ///
    /// This runs passes depending on the levels given.
    pub fn optimize(&self, opt_level: usize, size_level: usize) {
        unsafe {
            let builder = builder::LLVMPassManagerBuilderCreate();
            builder::LLVMPassManagerBuilderSetOptLevel(builder, opt_level as c_uint);
            builder::LLVMPassManagerBuilderSetSizeLevel(builder, size_level as c_uint);
            let pass_manager = core::LLVMCreatePassManager();
            builder::LLVMPassManagerBuilderPopulateModulePassManager(builder, pass_manager);
            builder::LLVMPassManagerBuilderDispose(builder);
            core::LLVMRunPassManager(pass_manager, self.into());
        }
    }

    /// Returns the target data of this module represented as a string
    pub fn get_target(&self) -> &str {
        unsafe {
            let target = core::LLVMGetTarget(self.into());
            util::to_str(target as *mut c_char)
        }
    }

    pub fn get_data_layout(&self) -> &str {
    	unsafe {
    		let layout = core::LLVMGetDataLayout(self.into());
    		util::to_str(layout as *mut c_char)
    	}
    }

    /// Set the target data of this module to the target data string given.
    pub fn set_target(&self, target: &str) {
        let c_target = CString::new(target).unwrap();
        unsafe { core::LLVMSetTarget(self.into(), c_target.as_ptr()) }
    }

    /// Verify that the module is safe to run, returning a string detailing the error
    /// when an error occurs.
    pub fn verify(&self) -> Result<(), CBox<str>> {
        unsafe {
            let mut error = mem::uninitialized();
            let action = LLVMVerifierFailureAction::LLVMReturnStatusAction;
            if analysis::LLVMVerifyModule(self.into(), action, &mut error) == 1 {
                Err(CBox::new(error))
            } else {
                Ok(())
            }
        }
    }

    /// Compile the module into an object file at the given location.
    ///
    /// Note that this uses the LLVM tool `llc` to do this, which may or may not be
    /// installed on the user's machine.
    pub fn compile(&self, path: &Path, opt_level: usize) -> IoResult<()> {
        let dir = env::temp_dir();
        let path = path.to_str().unwrap();
        let mod_path = dir.join("module.bc");
        let mod_path = mod_path.to_str().unwrap();
        try!(self.write_bitcode(mod_path));
        Command::new("llc")
            .arg(&format!("-O={}", opt_level))
            .arg("-filetype=obj")
            .arg("-o").arg(path)
            .arg(mod_path)
            .spawn()
            .map(|_| ())
    }

    /// Link a module into this module, returning an error string if an error occurs.
    ///
    /// This *does not* destroy the source module.
    pub fn link(&self, src: &Module) -> Result<(), CBox<str>> {
        unsafe {
            let dest = self.into();
            let src = src.into();
            let mut message = mem::uninitialized();
            if linker::LLVMLinkModules(dest, src, linker::LLVMLinkerMode::LLVMLinkerPreserveSource, &mut message) == 1 {
                Err(CBox::new(message))
            } else {
                Ok(())
            }
        }
    }

    /// Link a module into this module, returning an error string if an error occurs.
    ///
    /// This *does* destroy the source module.
    pub fn link_destroy(&self, src: CSemiBox<Module>) -> Result<(), CBox<str>> {
        unsafe {
            let dest = self.into();
            let src = src.as_ptr();
            let mut message = mem::uninitialized();
            if linker::LLVMLinkModules(dest, src, linker::LLVMLinkerMode::LLVMLinkerDestroySource, &mut message) == 1 {
                Err(CBox::new(message))
            } else {
                Ok(())
            }
        }
    }
}
impl<'a> IntoIterator for &'a Module {
    type Item = &'a Function;
    type IntoIter = Functions<'a>;
    /// Iterate through the functions in the module
    fn into_iter(self) -> Functions<'a> {
        Functions {
            value: unsafe { core::LLVMGetFirstFunction(self.into()) },
            marker: PhantomData
        }
    }
}
get_context!(Module, LLVMGetModuleContext);
to_str!(Module, LLVMPrintModuleToString);
dispose!(Module, LLVMModule, core::LLVMDisposeModule);
#[derive(Copy, Clone)]
/// An iterator through the functions contained in a module.
pub struct Functions<'a> {
    value: LLVMValueRef,
    marker: PhantomData<&'a ()>
}
impl<'a> Iterator for Functions<'a> {
    type Item = &'a Function;
    fn next(&mut self) -> Option<&'a Function> {
        if self.value.is_null() {
            None
        } else {
            let c_next = unsafe { core::LLVMGetNextFunction(self.value) };
            self.value = c_next;
            Some(self.value.into())
        }
    }
}
<|endoftext|>