<commit_before>use std::{fmt, i16, i32};
use std::io;
use std::str;
use uuid::Uuid;

// ========================================================================= //

/// Indicates the format of a string-typed database column.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Category {
    /// An unrestricted text string.
    ///
    /// # Examples
    ///
    /// ```
    /// assert!(msi::Category::Text.validate("Hello, World!"));
    /// ```
    Text,
    /// A text string containing no lowercase letters.
    ///
    /// # Examples
    ///
    /// ```
    /// assert!(msi::Category::UpperCase.validate("HELLO, WORLD!"));
    /// assert!(!msi::Category::UpperCase.validate("Hello, World!"));
    /// ```
    UpperCase,
    /// A text string containing no uppercase letters.
    ///
    /// # Examples
    ///
    /// ```
    /// assert!(msi::Category::LowerCase.validate("hello, world!"));
    /// assert!(!msi::Category::LowerCase.validate("Hello, World!"));
    /// ```
    LowerCase,
    /// A signed 16-bit integer.
    ///
    /// # Examples
    ///
    /// ```
    /// assert!(msi::Category::Integer.validate("32767"));
    /// assert!(msi::Category::Integer.validate("-47"));
    /// assert!(!msi::Category::Integer.validate("40000"));
    /// ```
    Integer,
    /// A signed 32-bit integer.
    ///
    /// # Examples
    ///
    /// ```
    /// assert!(msi::Category::DoubleInteger.validate("2147483647"));
    /// assert!(msi::Category::DoubleInteger.validate("-99999"));
    /// assert!(!msi::Category::DoubleInteger.validate("3000000000"));
    /// ```
    DoubleInteger,
    /// A string identifier (such as a table or column name).  May only contain
    /// alphanumerics, underscores, and periods, and must start with a letter
    /// or underscore.
    ///
    /// # Examples
    ///
    /// ```
    /// assert!(msi::Category::Identifier.validate("HelloWorld"));
    /// assert!(msi::Category::Identifier.validate("_99.Bottles"));
    /// assert!(!msi::Category::Identifier.validate("3.14159"));
    /// ```
    Identifier,
    /// A string that is either an identifier (see above), or a reference to an
    /// environment variable (which consists of a `%` character followed by an
    /// identifier).
    ///
    /// # Examples
    ///
    /// ```
    /// assert!(msi::Category::Property.validate("HelloWorld"));
    /// assert!(msi::Category::Property.validate("%HelloWorld"));
    /// assert!(!msi::Category::Property.validate("Hello%World"));
    /// ```
    Property,
    /// The name of a file or directory.
    Filename,
    /// A filename that can contain shell glob wildcards.
    WildCardFilename,
    /// A string containing an absolute filepath.
    Path,
    /// A string containing a semicolon-separated list of absolute filepaths.
    Paths,
    /// A string containing an absolute or relative filepath.
    AnyPath,
    /// A string containing either a filename or an identifier.
    DefaultDir,
    /// A string containing a registry path.
    RegPath,
    /// A string containing special formatting escapes, such as environment
    /// variables.
    Formatted,
    /// Unknown.
    KeyFormatted,
    /// Like `Formatted`, but allows additional escapes.
    Template,
    /// A string represeting a boolean predicate.
    Condition,
    /// A hyphenated, uppercase GUID string, enclosed in curly braces.
    ///
    /// # Examples
    ///
    /// ```
    /// assert!(msi::Category::Guid.validate(
    ///     "{34AB5C53-9B30-4E14-AEF0-2C1C7BA826C0}"));
    /// assert!(!msi::Category::Guid.validate(
    ///     "{34AB5C539B304E14AEF02C1C7BA826C0}")); // Must be hyphenated
    /// assert!(!msi::Category::Guid.validate(
    ///     "{34ab5c53-9b30-4e14-aef0-2c1c7ba826c0}")); // Must be uppercase
    /// assert!(!msi::Category::Guid.validate(
    ///     "34AB5C53-9B30-4E14-AEF0-2C1C7BA826C0")); // Must have braces
    /// assert!(!msi::Category::Guid.validate(
    ///     "{HELLOWO-RLDH-ELLO-WORL-DHELLOWORLD0}"));
    /// ```
    Guid,
    /// A string containing a version number.  The string must consist of at
    /// most four period-separated numbers, with each number being at most
    /// 65535.
    ///
    /// # Examples
    ///
    /// ```
    /// assert!(msi::Category::Version.validate("1"));
    /// assert!(msi::Category::Version.validate("1.22"));
    /// assert!(msi::Category::Version.validate("1.22.3"));
    /// assert!(msi::Category::Version.validate("1.22.3.444"));
    /// assert!(!msi::Category::Version.validate("1.99999"));
    /// assert!(!msi::Category::Version.validate(".12"));
    /// assert!(!msi::Category::Version.validate("1.2.3.4.5"));
    /// ```
    Version,
    /// A string containing a comma-separated list of deciaml language ID
    /// numbers.
    Language,
    /// A string that refers to a binary data stream.
    Binary,
    /// A string that refers to a custom source.
    CustomSource,
    /// A string that refers to a cabinet.  If it starts with a `#` character,
    /// then the rest of the string is an identifier (see above) indicating a
    /// data stream in the package where the cabinet is stored.  Otherwise, the
    /// string is a short filename (at most eight characters, a period, and a
    /// three-character extension).
    ///
    /// # Examples
    ///
    /// ```
    /// assert!(msi::Category::Cabinet.validate("hello.txt"));
    /// assert!(msi::Category::Cabinet.validate("#HelloWorld"));
    /// assert!(!msi::Category::Cabinet.validate("longfilename.long"));
    /// assert!(!msi::Category::Cabinet.validate("#123.456"));
    /// ```
    Cabinet,
    /// A string that refers to a shortcut.
    Shortcut,
    /// A string containing a URL.
    Url,
}

impl Category {
    pub(crate) fn all() -> Vec<Category> {
        vec![
            Category::Text,
            Category::UpperCase,
            Category::LowerCase,
            Category::Integer,
            Category::DoubleInteger,
            Category::Identifier,
            Category::Property,
            Category::Filename,
            Category::WildCardFilename,
            Category::Path,
            Category::Paths,
            Category::AnyPath,
            Category::DefaultDir,
            Category::RegPath,
            Category::Formatted,
            Category::KeyFormatted,
            Category::Template,
            Category::Condition,
            Category::Guid,
            Category::Version,
            Category::Language,
            Category::Binary,
            Category::CustomSource,
            Category::Cabinet,
            Category::Shortcut,
            Category::Url,
        ]
    }

    pub(crate) fn as_str(&self) -> &'static str {
        match *self {
            Category::AnyPath => "AnyPath",
            Category::Binary => "Binary",
            Category::Cabinet => "Cabinet",
            Category::Condition => "Condition",
            Category::CustomSource => "CustomSource",
            Category::DefaultDir => "DefaultDir",
            Category::DoubleInteger => "DoubleInteger",
            Category::Filename => "Filename",
            Category::Formatted => "Formatted",
            Category::Guid => "GUID",
            Category::Identifier => "Identifier",
            Category::Integer => "Integer",
            Category::KeyFormatted => "KeyFormatted",
            Category::Language => "Language",
            Category::LowerCase => "LowerCase",
            Category::Path => "Path",
            Category::Paths => "Paths",
            Category::Property => "Property",
            Category::RegPath => "RegPath",
            Category::Shortcut => "Shortcut",
            Category::Template => "Template",
            Category::Text => "Text",
            Category::UpperCase => "UpperCase",
            Category::Url => "URL",
            Category::Version => "Version",
            Category::WildCardFilename => "WildCardFilename",
        }
    }

    /// Returns true if the given string is valid to store in a database column
    /// with this category.
    pub fn validate(&self, string: &str) -> bool {
        match *self {
            Category::Text => true,
            Category::UpperCase => {
                !string.chars().any(|chr| chr >= 'a' && chr <= 'z')
            }
            Category::LowerCase => {
                !string.chars().any(|chr| chr >= 'A' && chr <= 'Z')
            }
            Category::Integer => string.parse::<i16>().is_ok(),
            Category::DoubleInteger => string.parse::<i32>().is_ok(),
            Category::Identifier => {
                string.starts_with(|chr| {
                                       chr >= 'A' && chr <= 'Z' ||
                                           chr >= 'a' && chr <= 'z' ||
                                           chr == '_'
                                   }) &&
                    !string.contains(|chr| {
                                         !(chr >= 'A' && chr <= 'Z' ||
                                               chr >= 'a' && chr <= 'z' ||
                                               chr >= '0' && chr <= '9' ||
                                               chr == '_' ||
                                               chr == '.')
                                     })
            }
            Category::Property => {
                let substr = if string.starts_with('%') {
                    &string[1..]
                } else {
                    string
                };
                Category::Identifier.validate(substr)
            }
            Category::Guid => {
                string.len() == 38 && string.starts_with('{') &&
                    string.ends_with('}') &&
                    !string.chars().any(|chr| chr >= 'a' && chr <= 'z') &&
                    Uuid::parse_str(&string[1..37]).is_ok()
            }
            Category::Version => {
                let mut parts = string.split('.');
                parts.clone().count() <= 4 &&
                    parts.all(|part| part.parse::<u16>().is_ok())
            }
            Category::Cabinet => {
                if string.starts_with('#') {
                    Category::Identifier.validate(&string[1..])
                } else {
                    let mut parts: Vec<&str> =
                        string.rsplitn(2, '.').collect();
                    parts.reverse();
                    parts.len() > 0 && parts[0].len() > 0 &&
                        parts[0].len() <= 8 &&
                        (parts.len() < 2 || parts[1].len() <= 3)
                }
            }
            // TODO: Validate other categories.
            _ => true,
        }
    }
}

impl fmt::Display for Category {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        self.as_str().fmt(formatter)
    }
}

impl str::FromStr for Category {
    type Err = io::Error;

    fn from_str(string: &str) -> io::Result<Category> {
        match string {
            "AnyPath" => Ok(Category::AnyPath),
            "Binary" => Ok(Category::Binary),
            "Cabinet" => Ok(Category::Cabinet),
            "Condition" => Ok(Category::Condition),
            "CustomSource" => Ok(Category::CustomSource),
            "DefaultDir" => Ok(Category::DefaultDir),
            "DoubleInteger" => Ok(Category::DoubleInteger),
            "Filename" => Ok(Category::Filename),
            "Formatted" => Ok(Category::Formatted),
            "GUID" => Ok(Category::Guid),
            "Guid" => Ok(Category::Guid),
            "Identifier" => Ok(Category::Identifier),
            "Integer" => Ok(Category::Integer),
            "KeyFormatted" => Ok(Category::KeyFormatted),
            "Language" => Ok(Category::Language),
            "LowerCase" => Ok(Category::LowerCase),
            "Path" => Ok(Category::Path),
            "Paths" => Ok(Category::Paths),
            "Property" => Ok(Category::Property),
            "RegPath" => Ok(Category::RegPath),
            "Shortcut" => Ok(Category::Shortcut),
            "Template" => Ok(Category::Template),
            "Text" => Ok(Category::Text),
            "UpperCase" => Ok(Category::UpperCase),
            "URL" => Ok(Category::Url),
            "Url" => Ok(Category::Url),
            "Version" => Ok(Category::Version),
            "WildCardFilename" => Ok(Category::WildCardFilename),
            _ => invalid_data!("Invalid category: {:?}", string),
        }
    }
}

// ========================================================================= //

#[cfg(test)]
mod tests {
    use super::Category;

    #[test]
    fn category_string_round_trip() {
        for category in Category::all() {
            assert_eq!(category.to_string().parse::<Category>().unwrap(),
                       category);
        }
    }
}

// ========================================================================= //
<commit_msg>Implement Language column category validation<commit_after>use std::{fmt, i16, i32};
use std::io;
use std::str;
use uuid::Uuid;

// ========================================================================= //

/// Indicates the format of a string-typed database column.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Category {
    /// An unrestricted text string.
    ///
    /// # Examples
    ///
    /// ```
    /// assert!(msi::Category::Text.validate("Hello, World!"));
    /// ```
    Text,
    /// A text string containing no lowercase letters.
    ///
    /// # Examples
    ///
    /// ```
    /// // Valid:
    /// assert!(msi::Category::UpperCase.validate("HELLO, WORLD!"));
    /// // Invalid:
    /// assert!(!msi::Category::UpperCase.validate("Hello, World!"));
    /// ```
    UpperCase,
    /// A text string containing no uppercase letters.
    ///
    /// # Examples
    ///
    /// ```
    /// // Valid:
    /// assert!(msi::Category::LowerCase.validate("hello, world!"));
    /// // Invalid:
    /// assert!(!msi::Category::LowerCase.validate("Hello, World!"));
    /// ```
    LowerCase,
    /// A signed 16-bit integer.
    ///
    /// # Examples
    ///
    /// ```
    /// // Valid:
    /// assert!(msi::Category::Integer.validate("32767"));
    /// assert!(msi::Category::Integer.validate("-47"));
    /// // Invalid:
    /// assert!(!msi::Category::Integer.validate("40000"));
    /// ```
    Integer,
    /// A signed 32-bit integer.
    ///
    /// # Examples
    ///
    /// ```
    /// // Valid:
    /// assert!(msi::Category::DoubleInteger.validate("2147483647"));
    /// assert!(msi::Category::DoubleInteger.validate("-99999"));
    /// // Invalid:
    /// assert!(!msi::Category::DoubleInteger.validate("3000000000"));
    /// ```
    DoubleInteger,
    /// A string identifier (such as a table or column name).  May only contain
    /// alphanumerics, underscores, and periods, and must start with a letter
    /// or underscore.
    ///
    /// # Examples
    ///
    /// ```
    /// // Valid:
    /// assert!(msi::Category::Identifier.validate("HelloWorld"));
    /// assert!(msi::Category::Identifier.validate("_99.Bottles"));
    /// // Invalid:
    /// assert!(!msi::Category::Identifier.validate("$HELLO"));
    /// assert!(!msi::Category::Identifier.validate("3.14159"));
    /// ```
    Identifier,
    /// A string that is either an identifier (see above), or a reference to an
    /// environment variable (which consists of a `%` character followed by an
    /// identifier).
    ///
    /// # Examples
    ///
    /// ```
    /// // Valid:
    /// assert!(msi::Category::Property.validate("HelloWorld"));
    /// assert!(msi::Category::Property.validate("%HelloWorld"));
    /// // Invalid:
    /// assert!(!msi::Category::Property.validate("%"));
    /// assert!(!msi::Category::Property.validate("Hello%World"));
    /// ```
    Property,
    /// The name of a file or directory.
    Filename,
    /// A filename that can contain shell glob wildcards.
    WildCardFilename,
    /// A string containing an absolute filepath.
    Path,
    /// A string containing a semicolon-separated list of absolute filepaths.
    Paths,
    /// A string containing an absolute or relative filepath.
    AnyPath,
    /// A string containing either a filename or an identifier.
    DefaultDir,
    /// A string containing a registry path.
    RegPath,
    /// A string containing special formatting escapes, such as environment
    /// variables.
    Formatted,
    /// Unknown.
    KeyFormatted,
    /// Like `Formatted`, but allows additional escapes.
    Template,
    /// A string represeting a boolean predicate.
    Condition,
    /// A hyphenated, uppercase GUID string, enclosed in curly braces.
    ///
    /// # Examples
    ///
    /// ```
    /// // Valid:
    /// assert!(msi::Category::Guid.validate(
    ///     "{34AB5C53-9B30-4E14-AEF0-2C1C7BA826C0}"));
    /// // Invalid:
    /// assert!(!msi::Category::Guid.validate(
    ///     "{34AB5C539B304E14AEF02C1C7BA826C0}")); // Must be hyphenated
    /// assert!(!msi::Category::Guid.validate(
    ///     "{34ab5c53-9b30-4e14-aef0-2c1c7ba826c0}")); // Must be uppercase
    /// assert!(!msi::Category::Guid.validate(
    ///     "34AB5C53-9B30-4E14-AEF0-2C1C7BA826C0")); // Must have braces
    /// assert!(!msi::Category::Guid.validate(
    ///     "{HELLOWO-RLDH-ELLO-WORL-DHELLOWORLD0}"));
    /// ```
    Guid,
    /// A string containing a version number.  The string must consist of at
    /// most four period-separated numbers, with the value of each number being
    /// at most 65535.
    ///
    /// # Examples
    ///
    /// ```
    /// // Valid:
    /// assert!(msi::Category::Version.validate("1"));
    /// assert!(msi::Category::Version.validate("1.22"));
    /// assert!(msi::Category::Version.validate("1.22.3"));
    /// assert!(msi::Category::Version.validate("1.22.3.444"));
    /// // Invalid:
    /// assert!(!msi::Category::Version.validate("1.99999"));
    /// assert!(!msi::Category::Version.validate(".12"));
    /// assert!(!msi::Category::Version.validate("1.2.3.4.5"));
    /// ```
    Version,
    /// A string containing a comma-separated list of decimal language ID
    /// numbers.
    ///
    /// # Examples
    ///
    /// ```
    /// // Valid:
    /// assert!(msi::Category::Language.validate("1033"));
    /// assert!(msi::Category::Language.validate("1083,2107,3131"));
    /// // Invalid:
    /// assert!(!msi::Category::Language.validate(""));
    /// assert!(!msi::Category::Language.validate("1083,2107,3131,"));
    /// assert!(!msi::Category::Language.validate("1083,,3131"));
    /// assert!(!msi::Category::Language.validate("en-US"));
    /// ```
    Language,
    /// A string that refers to a binary data stream.
    Binary,
    /// A string that refers to a custom source.
    CustomSource,
    /// A string that refers to a cabinet.  If it starts with a `#` character,
    /// then the rest of the string is an identifier (see above) indicating a
    /// data stream in the package where the cabinet is stored.  Otherwise, the
    /// string is a short filename (at most eight characters, a period, and a
    /// three-character extension).
    ///
    /// # Examples
    ///
    /// ```
    /// // Valid:
    /// assert!(msi::Category::Cabinet.validate("hello.txt"));
    /// assert!(msi::Category::Cabinet.validate("#HelloWorld"));
    /// // Invalid:
    /// assert!(!msi::Category::Cabinet.validate("longfilename.long"));
    /// assert!(!msi::Category::Cabinet.validate("#123.456"));
    /// ```
    Cabinet,
    /// A string that refers to a shortcut.
    Shortcut,
    /// A string containing a URL.
    Url,
}

impl Category {
    pub(crate) fn all() -> Vec<Category> {
        vec![
            Category::Text,
            Category::UpperCase,
            Category::LowerCase,
            Category::Integer,
            Category::DoubleInteger,
            Category::Identifier,
            Category::Property,
            Category::Filename,
            Category::WildCardFilename,
            Category::Path,
            Category::Paths,
            Category::AnyPath,
            Category::DefaultDir,
            Category::RegPath,
            Category::Formatted,
            Category::KeyFormatted,
            Category::Template,
            Category::Condition,
            Category::Guid,
            Category::Version,
            Category::Language,
            Category::Binary,
            Category::CustomSource,
            Category::Cabinet,
            Category::Shortcut,
            Category::Url,
        ]
    }

    pub(crate) fn as_str(&self) -> &'static str {
        match *self {
            Category::AnyPath => "AnyPath",
            Category::Binary => "Binary",
            Category::Cabinet => "Cabinet",
            Category::Condition => "Condition",
            Category::CustomSource => "CustomSource",
            Category::DefaultDir => "DefaultDir",
            Category::DoubleInteger => "DoubleInteger",
            Category::Filename => "Filename",
            Category::Formatted => "Formatted",
            Category::Guid => "GUID",
            Category::Identifier => "Identifier",
            Category::Integer => "Integer",
            Category::KeyFormatted => "KeyFormatted",
            Category::Language => "Language",
            Category::LowerCase => "LowerCase",
            Category::Path => "Path",
            Category::Paths => "Paths",
            Category::Property => "Property",
            Category::RegPath => "RegPath",
            Category::Shortcut => "Shortcut",
            Category::Template => "Template",
            Category::Text => "Text",
            Category::UpperCase => "UpperCase",
            Category::Url => "URL",
            Category::Version => "Version",
            Category::WildCardFilename => "WildCardFilename",
        }
    }

    /// Returns true if the given string is valid to store in a database column
    /// with this category.
    pub fn validate(&self, string: &str) -> bool {
        match *self {
            Category::Text => true,
            Category::UpperCase => {
                !string.chars().any(|chr| chr >= 'a' && chr <= 'z')
            }
            Category::LowerCase => {
                !string.chars().any(|chr| chr >= 'A' && chr <= 'Z')
            }
            Category::Integer => string.parse::<i16>().is_ok(),
            Category::DoubleInteger => string.parse::<i32>().is_ok(),
            Category::Identifier => {
                string.starts_with(|chr| {
                                       chr >= 'A' && chr <= 'Z' ||
                                           chr >= 'a' && chr <= 'z' ||
                                           chr == '_'
                                   }) &&
                    !string.contains(|chr| {
                                         !(chr >= 'A' && chr <= 'Z' ||
                                               chr >= 'a' && chr <= 'z' ||
                                               chr >= '0' && chr <= '9' ||
                                               chr == '_' ||
                                               chr == '.')
                                     })
            }
            Category::Property => {
                let substr = if string.starts_with('%') {
                    &string[1..]
                } else {
                    string
                };
                Category::Identifier.validate(substr)
            }
            Category::Guid => {
                string.len() == 38 && string.starts_with('{') &&
                    string.ends_with('}') &&
                    !string.chars().any(|chr| chr >= 'a' && chr <= 'z') &&
                    Uuid::parse_str(&string[1..37]).is_ok()
            }
            Category::Version => {
                let mut parts = string.split('.');
                parts.clone().count() <= 4 &&
                    parts.all(|part| part.parse::<u16>().is_ok())
            }
            Category::Language => {
                let mut parts = string.split(',');
                parts.all(|part| part.parse::<u16>().is_ok())
            }
            Category::Cabinet => {
                if string.starts_with('#') {
                    Category::Identifier.validate(&string[1..])
                } else {
                    let mut parts: Vec<&str> =
                        string.rsplitn(2, '.').collect();
                    parts.reverse();
                    parts.len() > 0 && parts[0].len() > 0 &&
                        parts[0].len() <= 8 &&
                        (parts.len() < 2 || parts[1].len() <= 3)
                }
            }
            // TODO: Validate other categories.
            _ => true,
        }
    }
}

impl fmt::Display for Category {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        self.as_str().fmt(formatter)
    }
}

impl str::FromStr for Category {
    type Err = io::Error;

    fn from_str(string: &str) -> io::Result<Category> {
        match string {
            "AnyPath" => Ok(Category::AnyPath),
            "Binary" => Ok(Category::Binary),
            "Cabinet" => Ok(Category::Cabinet),
            "Condition" => Ok(Category::Condition),
            "CustomSource" => Ok(Category::CustomSource),
            "DefaultDir" => Ok(Category::DefaultDir),
            "DoubleInteger" => Ok(Category::DoubleInteger),
            "Filename" => Ok(Category::Filename),
            "Formatted" => Ok(Category::Formatted),
            "GUID" => Ok(Category::Guid),
            "Guid" => Ok(Category::Guid),
            "Identifier" => Ok(Category::Identifier),
            "Integer" => Ok(Category::Integer),
            "KeyFormatted" => Ok(Category::KeyFormatted),
            "Language" => Ok(Category::Language),
            "LowerCase" => Ok(Category::LowerCase),
            "Path" => Ok(Category::Path),
            "Paths" => Ok(Category::Paths),
            "Property" => Ok(Category::Property),
            "RegPath" => Ok(Category::RegPath),
            "Shortcut" => Ok(Category::Shortcut),
            "Template" => Ok(Category::Template),
            "Text" => Ok(Category::Text),
            "UpperCase" => Ok(Category::UpperCase),
            "URL" => Ok(Category::Url),
            "Url" => Ok(Category::Url),
            "Version" => Ok(Category::Version),
            "WildCardFilename" => Ok(Category::WildCardFilename),
            _ => invalid_data!("Invalid category: {:?}", string),
        }
    }
}

// ========================================================================= //

#[cfg(test)]
mod tests {
    use super::Category;

    #[test]
    fn category_string_round_trip() {
        for category in Category::all() {
            assert_eq!(category.to_string().parse::<Category>().unwrap(),
                       category);
        }
    }
}

// ========================================================================= //
<|endoftext|>


<commit_before>//
// Copyright 2018 The Project Oak Authors
//
// 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.
//

// This model was inspired by
// https://github.com/AtheMathmo/rusty-machine/blob/master/examples/naive_bayes_dogs.rs .

extern crate oak;
extern crate oak_derive;
extern crate rand;
extern crate rusty_machine;

use oak_derive::OakNode;
use rand::distributions::Distribution;
use rand::distributions::Normal;
use rand::distributions::Standard;
use rand::prelude::*;
use rusty_machine::learning::naive_bayes::{self, NaiveBayes};
use rusty_machine::learning::SupModel;
use rusty_machine::linalg::{BaseMatrix, Matrix};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Type {
    Cat,
    Dog,
}

#[derive(Clone, Debug)]
struct Animal {
    type_: Type,
    friendliness: f64,
    furriness: f64,
    speed: f64,
}

impl Distribution<Animal> for Standard {
    /// Generate a random animal.
    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> Animal {
        // Friendliness, furriness, and speed are normally distributed and (given type) independent.
        let cat_friendliness = Normal::new(0., 1.);
        let cat_furriness = Normal::new(0., 1.);
        let cat_speed = Normal::new(0., 1.);

        let dog_friendliness = Normal::new(1., 1.);
        let dog_furriness = Normal::new(1., 1.);
        let dog_speed = Normal::new(-1., 1.);

        // Flip a coin to decide whether to generate a cat or a dog.
        let coin: f64 = rng.gen();
        let color = if coin < 0.5 { Type::Cat } else { Type::Dog };

        match color {
            Type::Cat => Animal {
                type_: Type::Cat,
                friendliness: cat_friendliness.sample(rng),
                furriness: cat_furriness.sample(rng),
                speed: cat_speed.sample(rng),
            },
            Type::Dog => Animal {
                type_: Type::Dog,
                friendliness: dog_friendliness.sample(rng),
                furriness: dog_furriness.sample(rng),
                speed: dog_speed.sample(rng),
            },
        }
    }
}

// FIXME: move data generation to client once we can write clients in Rust.
fn generate_animal_data(
    training_set_size: usize,
    test_set_size: usize,
) -> (Matrix<f64>, Matrix<f64>, Matrix<f64>, Vec<Animal>) {
    oak::print("rnd xxx\n");
    //let mut rng = rand::thread_rng();
    let mut rng = rand::rngs::StdRng::seed_from_u64(123);
    oak::print("rnd OK\n");

    // We'll train the model on these dogs
    let training_animals = (0..training_set_size)
        .map(|_| rng.gen())
        .collect::<Vec<Animal>>();

    // ... and then use the model to make predictions about these dogs' color
    // given only their trait measurements.
    let test_animals = (0..test_set_size)
        .map(|_| rng.gen())
        .collect::<Vec<Animal>>();

    // The model's `.train` method will take two matrices, each with a row for
    // each dog in the training set: the rows in the first matrix contain the
    // trait measurements; the rows in the second are either [1, 0] or [0, 1]
    // to indicate color.
    let training_data: Vec<f64> = training_animals
        .iter()
        .flat_map(|animal| vec![animal.friendliness, animal.furriness, animal.speed])
        .collect();
    let training_matrix: Matrix<f64> = training_data.chunks(3).collect();
    let target_data: Vec<f64> = training_animals
        .iter()
        .flat_map(|animal| match animal.type_ {
            Type::Cat => vec![1., 0.],
            Type::Dog => vec![0., 1.],
        })
        .collect();
    let target_matrix: Matrix<f64> = target_data.chunks(2).collect();

    // Build another matrix for the test set of dogs to make predictions about.
    let test_data: Vec<f64> = test_animals
        .iter()
        .flat_map(|animal| vec![animal.friendliness, animal.furriness, animal.speed])
        .collect();
    let test_matrix: Matrix<f64> = test_data.chunks(3).collect();

    (training_matrix, target_matrix, test_matrix, test_animals)
}

fn evaluate_prediction(hits: &mut u32, animal: &Animal, prediction: &[f64]) -> (Type, bool) {
    let predicted_type = animal.type_;
    let actual_type = if prediction[0] == 1. {
        Type::Cat
    } else {
        Type::Dog
    };
    let accurate = predicted_type == actual_type;
    if accurate {
        *hits += 1;
    }
    (actual_type, accurate)
}

struct Config {
    training_matrix: Matrix<f64>,
    target_matrix: Matrix<f64>,
    test_matrix: Matrix<f64>,
    test_animals: Vec<Animal>,
}

#[derive(OakNode)]
struct Node {
    training_set_size: usize,
    test_set_size: usize,
    config: Option<Config>,
    model: NaiveBayes<naive_bayes::Gaussian>,
}

impl oak::Node for Node {
    fn new() -> Self {
        Node {
            training_set_size: 1000,
            test_set_size: 1000,
            config: None,
            model: NaiveBayes::new(),
        }
    }
    fn invoke(
        &mut self,
        method_name: &str,
        _request: &mut oak::Reader,
        _response: &mut oak::Writer,
    ) {
        oak::print(method_name);
        match method_name {
            "/oak.examples.machine_learning.MachineLearning/Data" => {
                oak::print("Data");
                //(self.training_set_size, self.test_set_size) = (1000, 1000);
                // Generate all of our train and test data
                let (training_matrix, target_matrix, test_matrix, test_animals) =
                    generate_animal_data(self.training_set_size, self.test_set_size);
                self.config = Some(Config {
                    training_matrix: training_matrix,
                    target_matrix: target_matrix,
                    test_matrix: test_matrix,
                    test_animals: test_animals,
                });
            }
            "/oak.examples.machine_learning.MachineLearning/Learn" => {
                oak::print("Training model\n");
                //self.model = NaiveBayes::<naive_bayes::Gaussian>::new();
                match self.config {
                    Some(ref c) => self
                        .model
                        .train(&c.training_matrix, &c.target_matrix)
                        .expect("failed to train model of dogs"),
                    None => oak::print("config not set"),
                }
            }
            "/oak.examples.machine_learning.MachineLearning/Predict" => {
                oak::print("Predicting\n");
                let mut predictions = None;
                match self.config {
                    Some(ref c) => {
                        predictions = Some(
                            self.model
                                .predict(&c.test_matrix)
                                .expect("failed to predict dogs!?"),
                        )
                    }
                    None => oak::print("config not set"),
                }
                // Score how well we did.
                let mut hits = 0;
                let unprinted_total = self.test_set_size.saturating_sub(10) as usize;
                match self.config {
                    Some(ref c) => {
                        if let Some(ref p) = predictions {
                            for (animal, prediction) in c
                                .test_animals
                                .iter()
                                .zip(p.iter_rows())
                                .take(unprinted_total)
                            {
                                evaluate_prediction(&mut hits, animal, prediction);
                            }
                        }
                    }
                    None => {
                        oak::print("test_animals not set");
                        return;
                    }
                }

                if unprinted_total > 0 {
                    println!("...");
                }

                if let Some(ref c) = self.config {
                    if let Some(ref p) = predictions {
                        for (animal, prediction) in c
                            .test_animals
                            .iter()
                            .zip(p.iter_rows())
                            .skip(unprinted_total)
                        {
                            let (actual_type, accurate) =
                                evaluate_prediction(&mut hits, animal, prediction);
                            println!(
                                "Predicted: {:?}; Actual: {:?}; Accurate? {:?}",
                                animal.type_, actual_type, accurate
                            );
                        }
                    }
                }
                oak::print(&format!(
                    "Accuracy: {}/{} = {:.1}%",
                    hits,
                    self.test_set_size,
                    (f64::from(hits)) / (f64::from(self.test_set_size as u32)) * 100.
                ));
            }
            _ => {
                panic!("unknown method name");
            }
        }
    }
}

oak::oak_node!(Node);
<commit_msg>Fix machine learning example (#67)<commit_after>//
// Copyright 2018 The Project Oak Authors
//
// 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.
//

// This model was inspired by
// https://github.com/AtheMathmo/rusty-machine/blob/master/examples/naive_bayes_dogs.rs .

extern crate oak;
extern crate oak_derive;
extern crate rand;
extern crate rusty_machine;

use oak_derive::OakNode;
use rand::distributions::Distribution;
use rand::distributions::Normal;
use rand::distributions::Standard;
use rand::prelude::*;
use rusty_machine::learning::naive_bayes::{self, NaiveBayes};
use rusty_machine::learning::SupModel;
use rusty_machine::linalg::{BaseMatrix, Matrix};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Type {
    Cat,
    Dog,
}

#[derive(Clone, Debug)]
struct Animal {
    type_: Type,
    friendliness: f64,
    furriness: f64,
    speed: f64,
}

impl Distribution<Animal> for Standard {
    /// Generate a random animal.
    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> Animal {
        // Friendliness, furriness, and speed are normally distributed and (given type) independent.
        let cat_friendliness = Normal::new(0., 1.);
        let cat_furriness = Normal::new(0., 1.);
        let cat_speed = Normal::new(0., 1.);

        let dog_friendliness = Normal::new(1., 1.);
        let dog_furriness = Normal::new(1., 1.);
        let dog_speed = Normal::new(-1., 1.);

        // Flip a coin to decide whether to generate a cat or a dog.
        let coin: f64 = rng.gen();
        let color = if coin < 0.5 { Type::Cat } else { Type::Dog };

        match color {
            Type::Cat => Animal {
                type_: Type::Cat,
                friendliness: cat_friendliness.sample(rng),
                furriness: cat_furriness.sample(rng),
                speed: cat_speed.sample(rng),
            },
            Type::Dog => Animal {
                type_: Type::Dog,
                friendliness: dog_friendliness.sample(rng),
                furriness: dog_furriness.sample(rng),
                speed: dog_speed.sample(rng),
            },
        }
    }
}

// FIXME: move data generation to client once we can write clients in Rust.
fn generate_animal_data(
    training_set_size: usize,
    test_set_size: usize,
) -> (Matrix<f64>, Matrix<f64>, Matrix<f64>, Vec<Animal>) {
    oak::print("rnd xxx\n");
    //let mut rng = rand::thread_rng();
    let mut rng = rand::rngs::StdRng::seed_from_u64(123);
    oak::print("rnd OK\n");

    // We'll train the model on these dogs
    let training_animals = (0..training_set_size)
        .map(|_| rng.gen())
        .collect::<Vec<Animal>>();

    // ... and then use the model to make predictions about these dogs' color
    // given only their trait measurements.
    let test_animals = (0..test_set_size)
        .map(|_| rng.gen())
        .collect::<Vec<Animal>>();

    // The model's `.train` method will take two matrices, each with a row for
    // each dog in the training set: the rows in the first matrix contain the
    // trait measurements; the rows in the second are either [1, 0] or [0, 1]
    // to indicate color.
    let training_data: Vec<f64> = training_animals
        .iter()
        .flat_map(|animal| vec![animal.friendliness, animal.furriness, animal.speed])
        .collect();
    let training_matrix: Matrix<f64> = training_data.chunks(3).collect();
    let target_data: Vec<f64> = training_animals
        .iter()
        .flat_map(|animal| match animal.type_ {
            Type::Cat => vec![1., 0.],
            Type::Dog => vec![0., 1.],
        })
        .collect();
    let target_matrix: Matrix<f64> = target_data.chunks(2).collect();

    // Build another matrix for the test set of dogs to make predictions about.
    let test_data: Vec<f64> = test_animals
        .iter()
        .flat_map(|animal| vec![animal.friendliness, animal.furriness, animal.speed])
        .collect();
    let test_matrix: Matrix<f64> = test_data.chunks(3).collect();

    (training_matrix, target_matrix, test_matrix, test_animals)
}

fn evaluate_prediction(hits: &mut u32, animal: &Animal, prediction: &[f64]) -> (Type, bool) {
    let predicted_type = animal.type_;
    let actual_type = if prediction[0] == 1. {
        Type::Cat
    } else {
        Type::Dog
    };
    let accurate = predicted_type == actual_type;
    if accurate {
        *hits += 1;
    }
    (actual_type, accurate)
}

struct Config {
    training_matrix: Matrix<f64>,
    target_matrix: Matrix<f64>,
    test_matrix: Matrix<f64>,
    test_animals: Vec<Animal>,
}

#[derive(OakNode)]
struct Node {
    training_set_size: usize,
    test_set_size: usize,
    config: Option<Config>,
    model: NaiveBayes<naive_bayes::Gaussian>,
}

impl oak::Node for Node {
    fn new() -> Self {
        Node {
            training_set_size: 1000,
            test_set_size: 1000,
            config: None,
            model: NaiveBayes::new(),
        }
    }
    fn invoke(
        &mut self,
        method_name: &str,
        _request: &mut oak::Reader,
        _response: &mut oak::Writer,
    ) {
        oak::print(method_name);
        match method_name {
            "/oak.examples.machine_learning.MachineLearning/Data" => {
                oak::print("Data");
                //(self.training_set_size, self.test_set_size) = (1000, 1000);
                // Generate all of our train and test data
                let (training_matrix, target_matrix, test_matrix, test_animals) =
                    generate_animal_data(self.training_set_size, self.test_set_size);
                self.config = Some(Config {
                    training_matrix: training_matrix,
                    target_matrix: target_matrix,
                    test_matrix: test_matrix,
                    test_animals: test_animals,
                });
            }
            "/oak.examples.machine_learning.MachineLearning/Learn" => {
                oak::print("Training model\n");
                //self.model = NaiveBayes::<naive_bayes::Gaussian>::new();
                match self.config {
                    Some(ref c) => self
                        .model
                        .train(&c.training_matrix, &c.target_matrix)
                        .expect("failed to train model of dogs"),
                    None => oak::print("config not set"),
                }
            }
            "/oak.examples.machine_learning.MachineLearning/Predict" => {
                oak::print("Predicting\n");
                let mut predictions = None;
                match self.config {
                    Some(ref c) => {
                        predictions = Some(
                            self.model
                                .predict(&c.test_matrix)
                                .expect("failed to predict dogs!?"),
                        )
                    }
                    None => oak::print("config not set"),
                }
                // Score how well we did.
                let mut hits = 0;
                let unprinted_total = self.test_set_size.saturating_sub(10) as usize;
                match self.config {
                    Some(ref c) => {
                        if let Some(ref p) = predictions {
                            for (animal, prediction) in c
                                .test_animals
                                .iter()
                                .zip(p.iter_rows())
                                .take(unprinted_total)
                            {
                                evaluate_prediction(&mut hits, animal, prediction);
                            }
                        }
                    }
                    None => {
                        oak::print("test_animals not set");
                        return;
                    }
                }

                if unprinted_total > 0 {
                    println!("...");
                }

                if let Some(ref c) = self.config {
                    if let Some(ref p) = predictions {
                        for (animal, prediction) in c
                            .test_animals
                            .iter()
                            .zip(p.iter_rows())
                            .skip(unprinted_total)
                        {
                            let (actual_type, accurate) =
                                evaluate_prediction(&mut hits, animal, prediction);
                            println!(
                                "Predicted: {:?}; Actual: {:?}; Accurate? {:?}",
                                animal.type_, actual_type, accurate
                            );
                        }
                    }
                }
                oak::print(&format!(
                    "Accuracy: {}/{} = {:.1}%",
                    hits,
                    self.test_set_size,
                    (f64::from(hits)) / (f64::from(self.test_set_size as u32)) * 100.
                ));
            }
            _ => {
                panic!("unknown method name");
            }
        }
    }
}
<|endoftext|>


<commit_before>//! A `ControlFlowGraph` is a directed `Graph` of `Block` and `Edge`.

use il::*;
use std::collections::{BTreeMap};
use std::fmt;

/// A directed graph of types `Block` and `Edge`.
///
/// # Entry and Exit
/// A `ControlFlowGraph` has an optional, "Entry," and an optional, "Exit." When these are
/// provided, certain convenience functions become available.
///
/// For example, when translating a native instruction to Falcon IL, it can be useful to consider
/// an instruction as its own `ControlFlowGraph`. `rep scasb` is a great example of when this
/// pattern is helpful. Instructions in a `Block` will have one entry, and one exit. Explicitly
/// declaring these makes merging `ControlFlowGraph`s easier.
#[derive(Clone, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
pub struct ControlFlowGraph {
    // The internal graph used to store our blocks.
    graph: graph::Graph<Block, Edge>,
    // The next index to use when creating a basic block.
    next_index: u64,
    // The index for the next temp variable to create.
    next_temp_index: u64,
    // An optional entry index for the graph.
    entry: Option<u64>,
    // An optional exit index for the graph.
    exit: Option<u64>,
    // True if SSA has been applied to the ControlFlowGraph
    ssa_form: bool,
}

impl ControlFlowGraph {
    pub fn new() -> ControlFlowGraph {
        ControlFlowGraph {
            graph: graph::Graph::new(),
            next_index: 0,
            next_temp_index: 0,
            entry: None,
            exit: None,
            ssa_form: false,
        }
    }

    /// Returns the underlying graph
    pub fn graph(&self) -> &graph::Graph<Block, Edge> {
        &self.graph
    }

    /// Sets the entry point for this `ControlFlowGraph` to the given `Block` index.
    pub fn set_entry(&mut self, entry: u64) -> Result<()> {
        if self.graph.has_vertex(entry) {
            self.entry = Some(entry);
            return Ok(());
        }
        Err("Index does not exist for set_entry".into())
    }

    /// Sets the exit point for this `ControlFlowGraph` to the given `Block` index.
    pub fn set_exit(&mut self, exit: u64) -> Result<()> {
        if self.graph.has_vertex(exit) {
            self.exit = Some(exit);
            return Ok(());
        }
        Err("Index does not exist for set_exit".into())
    }

    /// Get the entry `Block` index for this `ControlFlowGraph`.
    pub fn entry(&self) -> Option<u64> {
        self.entry
    }

    /// Get the exit `Block` index for this `ControlFlowGraph`.
    pub fn exit(&self) -> Option<u64> {
        self.exit
    }

    /// Get a `Block` by index.
    pub fn block(&self, index: u64) -> Option<&Block> {
        self.graph.vertex(index)
    }

    /// Get a mutable reference to a `Block` by index.
    pub fn block_mut(&mut self, index: u64) -> Option<&mut Block> {
        self.graph.vertex_mut(index)
    }

    /// Get every `Block` in this `ControlFlowGraph`.
    pub fn blocks(&self) -> Vec<&Block> {
        self.graph.vertices()
    }

    /// Get a mutable reference to every `Block` in this `ControlFlowGraph`.
    pub fn blocks_mut(&mut self) -> Vec<&mut Block> {
        self.graph.vertices_mut()
    }

    /// Get an `Edge` by its head and tail `Block` indices.
    pub fn edge(&self, head: u64, tail: u64) -> Option<&Edge> {
        self.graph.edge(head, tail)
    }

    /// Get a mutable reference to an `Edge` by its head and tail `Block` indices.
    pub fn edge_mut(&mut self, head: u64, tail: u64) -> Option<&mut Edge> {
        self.graph.edge_mut(head, tail)
    }

    /// Get every `Edge` in thie `ControlFlowGraph`.
    pub fn edges(&self) -> Vec<&Edge> {
        self.graph.edges()
    }

    /// Get a mutable reference to every `Edge` in this `ControlFlowGraph`.
    pub fn edges_mut(&mut self) -> Vec<&mut Edge> {
        self.graph.edges_mut()
    }

    /// Get every incoming edge to a block
    pub fn edges_in(&self, index: u64) -> Option<&Vec<Edge>> {
        self.graph.edges_in(index)
    }

    /// Get every outgoing edge from a block
    pub fn edges_out(&self, index: u64) -> Option<&Vec<Edge>> {
        self.graph.edges_out(index)
    }

    /// Sets the address for all instructions in this `ControlFlowGraph`.
    ///
    /// Useful for translators to set address information.
    pub fn set_address(&mut self, address: Option<u64>) {
        for block in self.blocks_mut() {
            for instruction in block.instructions_mut() {
                instruction.set_address(address);
            }
        }
    }

    /// Returns the entry block for this ControlFlowGraph
    pub fn entry_block(&self) -> Option<&Block> {
        if self.entry.is_none() {
            None
        }
        else {
            self.block(self.entry.unwrap())
        }
    }

    /// Generates a temporary scalar unique to this control flow graph.
    pub fn temp(&mut self, bits: usize) -> Scalar {
        let next_index = self.next_temp_index;
        self.next_temp_index = next_index + 1;
        Scalar::new(format!("temp_{}", next_index), bits)
    }

    /// Creates a new basic block, adds it to the graph, and returns it
    pub fn new_block(&mut self) -> Result<&mut Block> {
        let next_index = self.next_index;
        self.next_index += 1;
        let block = Block::new(next_index);
        self.graph.insert_vertex(block)?;
        Ok(self.graph.vertex_mut(next_index).unwrap())
    }

    /// Creates an unconditional edge from one block to another block
    pub fn unconditional_edge(&mut self, head: u64, tail: u64) -> Result<()> {
        let edge = Edge::new(head, tail, None);
        self.graph.insert_edge(edge)
    }

    /// Creates a conditional edge from one block to another block
    pub fn conditional_edge(&mut self, head: u64, tail: u64, condition: Expression) -> Result<()> {
        let edge = Edge::new(head, tail, Some(condition));
        self.graph.insert_edge(edge)
    }

    /// Merge `Block`s.
    ///
    /// When a `Block` as only one successor, and that successor has only one predecessor, we
    /// merge both into one `Block`.
    pub fn merge(&mut self) -> Result<()> {
        loop {
            let mut merge_index = None;
            let mut successor_index = None;
            for block in self.blocks() {
                // check to see how many successors we have
                let successors = self.graph.edges_out(block.index()).unwrap();

                // if we do not have just one successor, we will not merge this block
                if successors.len() != 1 {
                    continue;
                }

                // If this successor has a condition, we will not merge this block
                if successors.first().unwrap().condition().is_some() {
                    continue;
                }

                // get the vertex for this successor
                let successor: u64 = match successors.first() {
                    Some(successor) => successor.tail(),
                    None => bail!("successor not found")
                };

                // get all predecessors for this successor
                let predecessors = self.graph.edges_in(successor).unwrap();

                // if this successor does not have exactly one predecessor, we
                // will not merge this block
                if predecessors.len() != 1 {
                    continue;
                }

                successor_index = Some(successor);
                merge_index = Some(block.index());
                break;
            }

            // we have a block to merge
            if let Some(merge_index) = merge_index {
                let successor_index = match successor_index {
                    Some(successor_index) => successor_index,
                    None => bail!("merge_index set, but not successor_index")
                };

                // merge the blocks
                let successor_block = self.graph
                                          .vertex(successor_index)
                                          .ok_or("Could not find block")?
                                          .clone();
                self.graph
                    .vertex_mut(merge_index)
                    .ok_or("Could not find block")?
                    .append(&successor_block);

                // all of successor's successors become merge_block's successors
                let mut new_edges = Vec::new();
                for edge in self.graph.edges_out(successor_index).unwrap() {
                    let head = merge_index;
                    let tail = edge.tail();
                    let condition = edge.condition().clone();
                    let edge = Edge::new(head, tail, condition.cloned());
                    new_edges.push(edge);
                }
                for edge in new_edges {
                    self.graph.insert_edge(edge)?;
                }

                // remove the block we just merged
                self.graph.remove_vertex(successor_index)?;
            } else {
                break;
            }
        }
        Ok(())
    }

    /// Appends a control flow graph to this control flow graph.
    ///
    /// In order for this to work, the entry and exit of boths graphs must be
    /// set, which should be the case for all conformant translators. You can
    /// also append to an empty ControlFlowGraph.
    pub fn append(&mut self, other: &ControlFlowGraph) -> Result<()> {
        let is_empty = match self.graph.num_vertices() {
            0 => true,
            _ => false
        };

        if !is_empty && (self.entry().is_none() || self.exit().is_none()) {
            return Err("entry/exit not set for dest ControlFlowGraph::append".into());
        }

        if other.entry().is_none() || other.exit().is_none() {
            return Err("entry/exit not set for src ControlFlowGraph::append".into());
        }

        // Bring in new blocks
        let mut block_map: BTreeMap<u64, u64> = BTreeMap::new();
        for block in other.graph().vertices() {
            // we need to clone the underlying block
            let new_block = block.clone_new_index(self.next_index);
            block_map.insert(block.index(), self.next_index);
            self.next_index += 1;
            self.graph.insert_vertex(new_block)?;
        }

        // Now set all new edges
        for edge in other.graph().edges() {
            let new_head: u64 = block_map[&edge.head()];
            let new_tail: u64 = block_map[&edge.tail()];
            let new_edge = Edge::new(new_head, new_tail, edge.condition().cloned());
            self.graph.insert_edge(new_edge)?;
        }

        if is_empty {
            self.entry = Some(block_map[&other.entry().unwrap()]);
        }
        else {
            // Create an edge from the exit of this graph to the head of the other
            // graph
            let transition_edge = Edge::new(
                self.exit.unwrap(),
                block_map[&(other.entry().unwrap())],
                None
            );
            self.graph.insert_edge(transition_edge)?;
        }

        self.exit = Some(block_map[&other.exit().unwrap()]);

        Ok(())
    }

    /// Inserts a control flow graph into this control flow graph, and returns
    /// the entry and exit indices for inserted graph.
    ///
    /// Requires the graph being inserted to have entry set.
    ///
    /// This function causes the `ControlFlowGraph` to become disconnected.
    ///
    /// This function is useful for inserting multiple `ControlFlowGraph`s into
    /// one before adding all `Edge`s in a subsequent pass.
    ///
    /// # Warnings
    /// This invalidates the entry and exit of the control flow graph.
    pub fn insert(&mut self, other: &ControlFlowGraph) -> Result<(u64, u64)> {
        if other.entry().is_none() || other.exit().is_none() {
            bail!("entry/exit not set on control flow graph");
        }

        // our entry and exit our no longer valid
        self.entry = None;
        self.exit = None;

        // Options to store the other graph entry/exit indices
        let mut entry_index = None;
        let mut exit_index = None;

        // keep track of mapping between old indices and new indices
        let mut block_map: BTreeMap<u64, u64> = BTreeMap::new();

        // insert all the blocks
        for block in other.graph().vertices() {
            let new_block = block.clone_new_index(self.next_index);
            block_map.insert(block.index(), self.next_index);
            if block.index() == other.entry().unwrap() {
                entry_index = Some(self.next_index);
            }
            if block.index() == other.exit().unwrap() {
                exit_index = Some(self.next_index);
            }
            self.next_index += 1;
            self.graph.insert_vertex(new_block)?;
        }

        // insert edges
        for edge in other.graph().edges() {
            let new_head: u64 = block_map[&edge.head()];
            let new_tail: u64 = block_map[&edge.tail()];
            let new_edge = Edge::new(new_head, new_tail, edge.condition().cloned());
            self.graph.insert_edge(new_edge)?;
        }

        if entry_index.is_none() || exit_index.is_none() {
            bail!("failed to get entry/exit indices");
        }

        Ok((entry_index.unwrap(), exit_index.unwrap()))
    }
}

impl fmt::Display for ControlFlowGraph {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for block in self.blocks() {
            match writeln!(f, "{}", block) {
                Err(e) => return Err(e),
                Ok(_) => {}
            }
        }
        for edge in self.edges() {
            match writeln!(f, "edge {}", edge) {
                Err(e) => return Err(e),
                Ok(_) => {}
            }
        }
        Ok(())
    }
}<commit_msg>improvements to il::ControlFlowGraph::merge, closes #14<commit_after>//! A `ControlFlowGraph` is a directed `Graph` of `Block` and `Edge`.

use il::*;
use std::collections::{BTreeMap};
use std::fmt;

/// A directed graph of types `Block` and `Edge`.
///
/// # Entry and Exit
/// A `ControlFlowGraph` has an optional, "Entry," and an optional, "Exit." When these are
/// provided, certain convenience functions become available.
///
/// For example, when translating a native instruction to Falcon IL, it can be useful to consider
/// an instruction as its own `ControlFlowGraph`. `rep scasb` is a great example of when this
/// pattern is helpful. Instructions in a `Block` will have one entry, and one exit. Explicitly
/// declaring these makes merging `ControlFlowGraph`s easier.
#[derive(Clone, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
pub struct ControlFlowGraph {
    // The internal graph used to store our blocks.
    graph: graph::Graph<Block, Edge>,
    // The next index to use when creating a basic block.
    next_index: u64,
    // The index for the next temp variable to create.
    next_temp_index: u64,
    // An optional entry index for the graph.
    entry: Option<u64>,
    // An optional exit index for the graph.
    exit: Option<u64>,
    // True if SSA has been applied to the ControlFlowGraph
    ssa_form: bool,
}

impl ControlFlowGraph {
    pub fn new() -> ControlFlowGraph {
        ControlFlowGraph {
            graph: graph::Graph::new(),
            next_index: 0,
            next_temp_index: 0,
            entry: None,
            exit: None,
            ssa_form: false,
        }
    }

    /// Returns the underlying graph
    pub fn graph(&self) -> &graph::Graph<Block, Edge> {
        &self.graph
    }

    /// Sets the entry point for this `ControlFlowGraph` to the given `Block` index.
    pub fn set_entry(&mut self, entry: u64) -> Result<()> {
        if self.graph.has_vertex(entry) {
            self.entry = Some(entry);
            return Ok(());
        }
        Err("Index does not exist for set_entry".into())
    }

    /// Sets the exit point for this `ControlFlowGraph` to the given `Block` index.
    pub fn set_exit(&mut self, exit: u64) -> Result<()> {
        if self.graph.has_vertex(exit) {
            self.exit = Some(exit);
            return Ok(());
        }
        Err("Index does not exist for set_exit".into())
    }

    /// Get the entry `Block` index for this `ControlFlowGraph`.
    pub fn entry(&self) -> Option<u64> {
        self.entry
    }

    /// Get the exit `Block` index for this `ControlFlowGraph`.
    pub fn exit(&self) -> Option<u64> {
        self.exit
    }

    /// Get a `Block` by index.
    pub fn block(&self, index: u64) -> Option<&Block> {
        self.graph.vertex(index)
    }

    /// Get a mutable reference to a `Block` by index.
    pub fn block_mut(&mut self, index: u64) -> Option<&mut Block> {
        self.graph.vertex_mut(index)
    }

    /// Get every `Block` in this `ControlFlowGraph`.
    pub fn blocks(&self) -> Vec<&Block> {
        self.graph.vertices()
    }

    /// Get a mutable reference to every `Block` in this `ControlFlowGraph`.
    pub fn blocks_mut(&mut self) -> Vec<&mut Block> {
        self.graph.vertices_mut()
    }

    /// Get an `Edge` by its head and tail `Block` indices.
    pub fn edge(&self, head: u64, tail: u64) -> Option<&Edge> {
        self.graph.edge(head, tail)
    }

    /// Get a mutable reference to an `Edge` by its head and tail `Block` indices.
    pub fn edge_mut(&mut self, head: u64, tail: u64) -> Option<&mut Edge> {
        self.graph.edge_mut(head, tail)
    }

    /// Get every `Edge` in thie `ControlFlowGraph`.
    pub fn edges(&self) -> Vec<&Edge> {
        self.graph.edges()
    }

    /// Get a mutable reference to every `Edge` in this `ControlFlowGraph`.
    pub fn edges_mut(&mut self) -> Vec<&mut Edge> {
        self.graph.edges_mut()
    }

    /// Get every incoming edge to a block
    pub fn edges_in(&self, index: u64) -> Option<&Vec<Edge>> {
        self.graph.edges_in(index)
    }

    /// Get every outgoing edge from a block
    pub fn edges_out(&self, index: u64) -> Option<&Vec<Edge>> {
        self.graph.edges_out(index)
    }

    /// Sets the address for all instructions in this `ControlFlowGraph`.
    ///
    /// Useful for translators to set address information.
    pub fn set_address(&mut self, address: Option<u64>) {
        for block in self.blocks_mut() {
            for instruction in block.instructions_mut() {
                instruction.set_address(address);
            }
        }
    }

    /// Returns the entry block for this ControlFlowGraph
    pub fn entry_block(&self) -> Option<&Block> {
        if self.entry.is_none() {
            None
        }
        else {
            self.block(self.entry.unwrap())
        }
    }

    /// Generates a temporary scalar unique to this control flow graph.
    pub fn temp(&mut self, bits: usize) -> Scalar {
        let next_index = self.next_temp_index;
        self.next_temp_index = next_index + 1;
        Scalar::new(format!("temp_{}", next_index), bits)
    }

    /// Creates a new basic block, adds it to the graph, and returns it
    pub fn new_block(&mut self) -> Result<&mut Block> {
        let next_index = self.next_index;
        self.next_index += 1;
        let block = Block::new(next_index);
        self.graph.insert_vertex(block)?;
        Ok(self.graph.vertex_mut(next_index).unwrap())
    }

    /// Creates an unconditional edge from one block to another block
    pub fn unconditional_edge(&mut self, head: u64, tail: u64) -> Result<()> {
        let edge = Edge::new(head, tail, None);
        self.graph.insert_edge(edge)
    }

    /// Creates a conditional edge from one block to another block
    pub fn conditional_edge(&mut self, head: u64, tail: u64, condition: Expression) -> Result<()> {
        let edge = Edge::new(head, tail, Some(condition));
        self.graph.insert_edge(edge)
    }

    /// Merge `Block`s.
    ///
    /// When a `Block` as only one successor, and that successor has only one predecessor, we
    /// merge both into one `Block`.
    pub fn merge(&mut self) -> Result<()> {
        use std::collections::HashSet;

        loop {

            let mut blocks_being_merged: HashSet<u64> = HashSet::new();
            let mut merges: Vec<(u64, u64)> = Vec::new();

            for block in self.blocks() {
                // If we are already merging this block this iteration, skip it
                if blocks_being_merged.contains(&block.index()) {
                    continue;
                }

                // check to see how many successors we have
                let successors = self.graph.edges_out(block.index()).unwrap();

                // if we do not have just one successor, we will not merge this block
                if successors.len() != 1 {
                    continue;
                }

                // If this successor has a condition, we will not merge this block
                if successors.first().unwrap().condition().is_some() {
                    continue;
                }

                // get the vertex for this successor
                let successor: u64 = match successors.first() {
                    Some(successor) => successor.tail(),
                    None => bail!("successor not found")
                };

                // If this successor is already being merged, skip it
                if blocks_being_merged.contains(&successor) {
                    continue;
                }

                // get all predecessors for this successor
                let predecessors = self.graph.edges_in(successor).unwrap();

                // if this successor does not have exactly one predecessor, we
                // will not merge this block
                if predecessors.len() != 1 {
                    continue;
                }

                blocks_being_merged.insert(block.index());
                blocks_being_merged.insert(successor);

                merges.push((block.index(), successor));
            }

            if merges.is_empty() {
                break;
            }

            for (merge_index, successor_index) in merges {
                // merge the blocks
                let successor_block = self.graph
                                          .vertex(successor_index)
                                          .ok_or("Could not find block")?
                                          .clone();
                self.graph
                    .vertex_mut(merge_index)
                    .ok_or("Could not find block")?
                    .append(&successor_block);

                // all of successor's successors become merge_block's successors
                let mut new_edges = Vec::new();
                for edge in self.graph.edges_out(successor_index).unwrap() {
                    let head = merge_index;
                    let tail = edge.tail();
                    let condition = edge.condition().clone();
                    let edge = Edge::new(head, tail, condition.cloned());
                    new_edges.push(edge);
                }
                for edge in new_edges {
                    self.graph.insert_edge(edge)?;
                }

                // remove the block we just merged
                self.graph.remove_vertex(successor_index)?;
            }
        }
        Ok(())
    }

    /// Appends a control flow graph to this control flow graph.
    ///
    /// In order for this to work, the entry and exit of boths graphs must be
    /// set, which should be the case for all conformant translators. You can
    /// also append to an empty ControlFlowGraph.
    pub fn append(&mut self, other: &ControlFlowGraph) -> Result<()> {
        let is_empty = match self.graph.num_vertices() {
            0 => true,
            _ => false
        };

        if !is_empty && (self.entry().is_none() || self.exit().is_none()) {
            return Err("entry/exit not set for dest ControlFlowGraph::append".into());
        }

        if other.entry().is_none() || other.exit().is_none() {
            return Err("entry/exit not set for src ControlFlowGraph::append".into());
        }

        // Bring in new blocks
        let mut block_map: BTreeMap<u64, u64> = BTreeMap::new();
        for block in other.graph().vertices() {
            // we need to clone the underlying block
            let new_block = block.clone_new_index(self.next_index);
            block_map.insert(block.index(), self.next_index);
            self.next_index += 1;
            self.graph.insert_vertex(new_block)?;
        }

        // Now set all new edges
        for edge in other.graph().edges() {
            let new_head: u64 = block_map[&edge.head()];
            let new_tail: u64 = block_map[&edge.tail()];
            let new_edge = Edge::new(new_head, new_tail, edge.condition().cloned());
            self.graph.insert_edge(new_edge)?;
        }

        if is_empty {
            self.entry = Some(block_map[&other.entry().unwrap()]);
        }
        else {
            // Create an edge from the exit of this graph to the head of the other
            // graph
            let transition_edge = Edge::new(
                self.exit.unwrap(),
                block_map[&(other.entry().unwrap())],
                None
            );
            self.graph.insert_edge(transition_edge)?;
        }

        self.exit = Some(block_map[&other.exit().unwrap()]);

        Ok(())
    }

    /// Inserts a control flow graph into this control flow graph, and returns
    /// the entry and exit indices for inserted graph.
    ///
    /// Requires the graph being inserted to have entry set.
    ///
    /// This function causes the `ControlFlowGraph` to become disconnected.
    ///
    /// This function is useful for inserting multiple `ControlFlowGraph`s into
    /// one before adding all `Edge`s in a subsequent pass.
    ///
    /// # Warnings
    /// This invalidates the entry and exit of the control flow graph.
    pub fn insert(&mut self, other: &ControlFlowGraph) -> Result<(u64, u64)> {
        if other.entry().is_none() || other.exit().is_none() {
            bail!("entry/exit not set on control flow graph");
        }

        // our entry and exit our no longer valid
        self.entry = None;
        self.exit = None;

        // Options to store the other graph entry/exit indices
        let mut entry_index = None;
        let mut exit_index = None;

        // keep track of mapping between old indices and new indices
        let mut block_map: BTreeMap<u64, u64> = BTreeMap::new();

        // insert all the blocks
        for block in other.graph().vertices() {
            let new_block = block.clone_new_index(self.next_index);
            block_map.insert(block.index(), self.next_index);
            if block.index() == other.entry().unwrap() {
                entry_index = Some(self.next_index);
            }
            if block.index() == other.exit().unwrap() {
                exit_index = Some(self.next_index);
            }
            self.next_index += 1;
            self.graph.insert_vertex(new_block)?;
        }

        // insert edges
        for edge in other.graph().edges() {
            let new_head: u64 = block_map[&edge.head()];
            let new_tail: u64 = block_map[&edge.tail()];
            let new_edge = Edge::new(new_head, new_tail, edge.condition().cloned());
            self.graph.insert_edge(new_edge)?;
        }

        if entry_index.is_none() || exit_index.is_none() {
            bail!("failed to get entry/exit indices");
        }

        Ok((entry_index.unwrap(), exit_index.unwrap()))
    }
}

impl fmt::Display for ControlFlowGraph {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for block in self.blocks() {
            match writeln!(f, "{}", block) {
                Err(e) => return Err(e),
                Ok(_) => {}
            }
        }
        for edge in self.edges() {
            match writeln!(f, "edge {}", edge) {
                Err(e) => return Err(e),
                Ok(_) => {}
            }
        }
        Ok(())
    }
}<|endoftext|>


<commit_before>generate_error_imports!();

generate_error_types!(TagError, TagErrorKind,
    TagTypeError     => "Entry Header Tag Type wrong",
    HeaderReadError  => "Error while reading entry header",
    HeaderWriteError => "Error while writing entry header",
    NotATag          => "String is not a tag"
);

<commit_msg>libimagentrytag: use generate_error_module!() macro and reexport generated types<commit_after>generate_error_module!(
    generate_error_types!(TagError, TagErrorKind,
        TagTypeError     => "Entry Header Tag Type wrong",
        HeaderReadError  => "Error while reading entry header",
        HeaderWriteError => "Error while writing entry header",
        NotATag          => "String is not a tag"
    );
);

pub use self::error::TagError;
pub use self::error::TagErrorKind;

<|endoftext|>


<commit_before>pub use self::fs::LazyFile;

#[cfg(test)]
mod fs {
    use error::StoreError as SE;
    use std::io::Cursor;
    use std::path::PathBuf;

    use std::collections::HashMap;
    use std::sync::Mutex;

    lazy_static! {
        static ref MAP: Mutex<HashMap<PathBuf, Cursor<Vec<u8>>>> = {
            Mutex::new(HashMap::new())
        };
    }

    /// `LazyFile` type, this is the Test version!
    ///
    /// A lazy file is either absent, but a path to it is available, or it is present.
    #[derive(Debug)]
    pub enum LazyFile {
        Absent(PathBuf),
    }

    impl LazyFile {

        /**
         * Get the mutable file behind a LazyFile object
         */
        pub fn get_file_content(&mut self) -> Result<Cursor<Vec<u8>>, SE> {
            debug!("Getting lazy file: {:?}", self);
            match *self {
                LazyFile::Absent(ref f) => {
                    let map = MAP.lock().unwrap();
                    return Ok(map.get(f).unwrap().clone());
                },
            };
        }

        pub fn write_file_content(&mut self, buf: &[u8]) -> Result<(), SE> {
            match *self {
                LazyFile::Absent(ref f) => {
                    let mut map = MAP.lock().unwrap();
                    if let Some(ref mut cur) = map.get_mut(f) {
                        let mut vec = cur.get_mut();
                        vec.clear();
                        vec.extend_from_slice(buf);
                        return Ok(());
                    }
                    let vec = Vec::from(buf);
                    map.insert(f.clone(), Cursor::new(vec));
                    return Ok(());
                },
            };
        }
    }
}

#[cfg(not(test))]
mod fs {
    use error::{MapErrInto, StoreError as SE, StoreErrorKind as SEK};
    use std::io::{Seek, SeekFrom, Read};
    use std::path::{Path, PathBuf};
    use std::fs::{File, OpenOptions, create_dir_all};

    /// `LazyFile` type
    ///
    /// A lazy file is either absent, but a path to it is available, or it is present.
    #[derive(Debug)]
    pub enum LazyFile {
        Absent(PathBuf),
        File(File)
    }

    fn open_file<A: AsRef<Path>>(p: A) -> ::std::io::Result<File> {
        OpenOptions::new().write(true).read(true).open(p)
    }

    fn create_file<A: AsRef<Path>>(p: A) -> ::std::io::Result<File> {
        if let Some(parent) = p.as_ref().parent() {
            debug!("Implicitely creating directory: {:?}", parent);
            if let Err(e) = create_dir_all(parent) {
                return Err(e);
            }
        }
        OpenOptions::new().write(true).read(true).create(true).open(p)
    }

    impl LazyFile {

        /**
         * Get the content behind this file
         */
        pub fn get_file_content(&mut self) -> Result<&mut Read, SE> {
            debug!("Getting lazy file: {:?}", self);
            let file = match *self {
                LazyFile::File(ref mut f) => return {
                    // We seek to the beginning of the file since we expect each
                    // access to the file to be in a different context
                    try!(f.seek(SeekFrom::Start(0))
                        .map_err_into(SEK::FileNotSeeked));
                    Ok(f)
                },
                LazyFile::Absent(ref p) => try!(open_file(p).map_err_into(SEK::FileNotFound)),
            };
            *self = LazyFile::File(file);
            if let LazyFile::File(ref mut f) = *self {
                return Ok(f);
            }
            unreachable!()
        }

        /**
         * Write the content of this file
         */
        pub fn write_file_content(&mut self, buf: &[u8]) -> Result<(), SE> {
            use std::io::Write;
            let file = match *self {
                LazyFile::File(ref mut f) => return {
                    // We seek to the beginning of the file since we expect each
                    // access to the file to be in a different context
                    try!(f.seek(SeekFrom::Start(0))
                        .map_err_into(SEK::FileNotCreated));
                    f.write_all(buf).map_err_into(SEK::FileNotWritten)
                },
                LazyFile::Absent(ref p) => try!(create_file(p).map_err_into(SEK::FileNotCreated)),
            };
            *self = LazyFile::File(file);
            if let LazyFile::File(ref mut f) = *self {
                return f.write_all(buf).map_err_into(SEK::FileNotWritten);
            }
            unreachable!();
        }
    }
}

#[cfg(test)]
mod test {
    // use super::LazyFile;
    // use std::io::{Read, Write};
    // use std::path::PathBuf;
    // use tempdir::TempDir;

    // fn get_dir() -> TempDir {
    //     TempDir::new("test-image").unwrap()
    // }

    // #[test]
    // fn lazy_file() {
    //     let dir = get_dir();
    //     let mut path = PathBuf::from(dir.path());
    //     path.set_file_name("test1");
    //     let mut lf = LazyFile::Absent(path);

    //     write!(lf.create_file().unwrap(), "Hello World").unwrap();
    //     dir.close().unwrap();
    // }

    // #[test]
    // fn lazy_file_with_file() {
    //     let dir = get_dir();
    //     let mut path = PathBuf::from(dir.path());
    //     path.set_file_name("test2");
    //     let mut lf = LazyFile::Absent(path.clone());

    //     {
    //         let mut file = lf.create_file().unwrap();

    //         file.write(b"Hello World").unwrap();
    //         file.sync_all().unwrap();
    //     }

    //     {
    //         let mut file = lf.get_file_mut().unwrap();
    //         let mut s = Vec::new();
    //         file.read_to_end(&mut s).unwrap();
    //         assert_eq!(s, "Hello World".to_string().into_bytes());
    //     }

    //     dir.close().unwrap();
    // }
}
<commit_msg>Add simple test<commit_after>pub use self::fs::LazyFile;

#[cfg(test)]
mod fs {
    use error::StoreError as SE;
    use std::io::Cursor;
    use std::path::PathBuf;

    use std::collections::HashMap;
    use std::sync::Mutex;

    lazy_static! {
        static ref MAP: Mutex<HashMap<PathBuf, Cursor<Vec<u8>>>> = {
            Mutex::new(HashMap::new())
        };
    }

    /// `LazyFile` type, this is the Test version!
    ///
    /// A lazy file is either absent, but a path to it is available, or it is present.
    #[derive(Debug)]
    pub enum LazyFile {
        Absent(PathBuf),
    }

    impl LazyFile {

        /**
         * Get the mutable file behind a LazyFile object
         */
        pub fn get_file_content(&mut self) -> Result<Cursor<Vec<u8>>, SE> {
            debug!("Getting lazy file: {:?}", self);
            match *self {
                LazyFile::Absent(ref f) => {
                    let map = MAP.lock().unwrap();
                    return Ok(map.get(f).unwrap().clone());
                },
            };
        }

        pub fn write_file_content(&mut self, buf: &[u8]) -> Result<(), SE> {
            match *self {
                LazyFile::Absent(ref f) => {
                    let mut map = MAP.lock().unwrap();
                    if let Some(ref mut cur) = map.get_mut(f) {
                        let mut vec = cur.get_mut();
                        vec.clear();
                        vec.extend_from_slice(buf);
                        return Ok(());
                    }
                    let vec = Vec::from(buf);
                    map.insert(f.clone(), Cursor::new(vec));
                    return Ok(());
                },
            };
        }
    }
}

#[cfg(not(test))]
mod fs {
    use error::{MapErrInto, StoreError as SE, StoreErrorKind as SEK};
    use std::io::{Seek, SeekFrom, Read};
    use std::path::{Path, PathBuf};
    use std::fs::{File, OpenOptions, create_dir_all};

    /// `LazyFile` type
    ///
    /// A lazy file is either absent, but a path to it is available, or it is present.
    #[derive(Debug)]
    pub enum LazyFile {
        Absent(PathBuf),
        File(File)
    }

    fn open_file<A: AsRef<Path>>(p: A) -> ::std::io::Result<File> {
        OpenOptions::new().write(true).read(true).open(p)
    }

    fn create_file<A: AsRef<Path>>(p: A) -> ::std::io::Result<File> {
        if let Some(parent) = p.as_ref().parent() {
            debug!("Implicitely creating directory: {:?}", parent);
            if let Err(e) = create_dir_all(parent) {
                return Err(e);
            }
        }
        OpenOptions::new().write(true).read(true).create(true).open(p)
    }

    impl LazyFile {

        /**
         * Get the content behind this file
         */
        pub fn get_file_content(&mut self) -> Result<&mut Read, SE> {
            debug!("Getting lazy file: {:?}", self);
            let file = match *self {
                LazyFile::File(ref mut f) => return {
                    // We seek to the beginning of the file since we expect each
                    // access to the file to be in a different context
                    try!(f.seek(SeekFrom::Start(0))
                        .map_err_into(SEK::FileNotSeeked));
                    Ok(f)
                },
                LazyFile::Absent(ref p) => try!(open_file(p).map_err_into(SEK::FileNotFound)),
            };
            *self = LazyFile::File(file);
            if let LazyFile::File(ref mut f) = *self {
                return Ok(f);
            }
            unreachable!()
        }

        /**
         * Write the content of this file
         */
        pub fn write_file_content(&mut self, buf: &[u8]) -> Result<(), SE> {
            use std::io::Write;
            let file = match *self {
                LazyFile::File(ref mut f) => return {
                    // We seek to the beginning of the file since we expect each
                    // access to the file to be in a different context
                    try!(f.seek(SeekFrom::Start(0))
                        .map_err_into(SEK::FileNotCreated));
                    f.write_all(buf).map_err_into(SEK::FileNotWritten)
                },
                LazyFile::Absent(ref p) => try!(create_file(p).map_err_into(SEK::FileNotCreated)),
            };
            *self = LazyFile::File(file);
            if let LazyFile::File(ref mut f) = *self {
                return f.write_all(buf).map_err_into(SEK::FileNotWritten);
            }
            unreachable!();
        }
    }
}

#[cfg(test)]
mod test {
    use super::LazyFile;
    use std::io::Read;
    use std::path::PathBuf;

    #[test]
    fn lazy_file() {
        let mut path = PathBuf::from("/tests");
        path.set_file_name("test1");
        let mut lf = LazyFile::Absent(path);
        lf.write_file_content(b"Hello World").unwrap();
        let mut bah = Vec::new();
        lf.get_file_content().unwrap().read_to_end(&mut bah).unwrap();
        assert_eq!(bah, b"Hello World");
    }

}
<|endoftext|>