<commit_before>use token::Token;
use logger::Logger;
use error::{Error, Result, Source};

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::borrow::Cow;
use std::fs;
use std::io::Read;

use walkdir::WalkDir;
use rustc_serialize::base64::{self, ToBase64};
use mime_guess;

/// Resource Handler.
///
/// Its task is to make sure that some resource (image, link) is available
/// for the book and to list images used in Markdown files so they can be used for the book
#[derive(Debug)]
pub struct ResourceHandler<'r> {
    /// Maps an original url (e.g.) "foo/Readme.md" to a valid link
    /// (e.g.) chapter3.html
    links: HashMap<String, String>,
    map_images: bool,
    logger: &'r Logger,
    base64: bool,

    /// Maps an original (local) file name to a new file name. Allows to
    /// make sure all image files will be included in e.g. the Epub document.
    #[doc(hidden)]
    pub images: HashMap<String, String>,
}

impl<'r> ResourceHandler<'r> {
    /// Creates a new, empty Resource Handler
    pub fn new(logger: &'r Logger) -> ResourceHandler {
        ResourceHandler {
            links: HashMap::new(),
            images: HashMap::new(),
            map_images: false,
            base64: false,
            logger: logger,
        }
    }

    /// Turns on mapping for image files
    ///
    /// # Argument: an offset (should be book.root)
    pub fn set_images_mapping(&mut self, b: bool) -> &mut Self {
        self.map_images = b;
        self
    }

    /// Sets base64 mode for image mapping
    ///
    /// If set to true, instead of returning a destination file path,
    /// `map_image` will include the image as base64
    pub fn set_base64(&mut self, b: bool) {
        self.base64 = b;
    }

    /// Add a local image file and get the resulting transformed
    /// file name
    pub fn map_image<'a, S: Into<Cow<'a, str>>>(&'a mut self,
                                                source: &Source,
                                                file: S)
                                                -> Result<Cow<'a, str>> {
        // If image is not local, do nothing much
        let file = file.into();
        if !Self::is_local(file.as_ref()) {
            self.logger
                .warning(lformat!("Resources: book includes non-local image {file}, which might \
                                   cause problem for proper inclusion.",
                                  file = file));
            return Ok(file);
        }

        // Check exisence of the file
        if fs::metadata(file.as_ref()).is_err() {
            return Err(Error::file_not_found(source, lformat!("image"), format!("{}", file)));
        }

        // if image mapping is not activated do nothing else
        if !self.map_images {
            return Ok(file);
        }

        // If this image has already been registered, returns it
        if self.images.contains_key(file.as_ref()) {
            return Ok(Cow::Borrowed(&self.images[file.as_ref()]));
        }

        // Else, create a new file name that has same extension
        // (or a base64 version of the file)
        let dest_file = if !(self.base64) {
            if let Some(extension) = Path::new(file.as_ref()).extension() {
                format!("images/image_{}.{}",
                        self.images.len(),
                        extension.to_string_lossy())
            } else {
                self.logger
                    .warning(lformat!("Resources: book includes image {file} which doesn't have \
                                       an extension",
                                      file = file));
                format!("images/image_{}", self.images.len())
            }
        } else {
            let mut f = match fs::File::open(file.as_ref()) {
                Ok(f) => f,
                Err(_) => {
                    return Err(Error::file_not_found(source,
                                                     lformat!("image"),
                                                     format!("{}", file)));
                }
            };
            let mut content: Vec<u8> = vec![];
            if f.read_to_end(&mut content).is_err() {
                self.logger.error(lformat!("Resources: could not read file {file}", file = file));
                return Ok(file);
            }
            let base64 = content.to_base64(base64::STANDARD);
            match mime_guess::guess_mime_type_opt(file.as_ref()) {
                None => {
                    self.logger
                        .error(lformat!("Resources: could not guess mime type of file {file}",
                                        file = file));
                    return Ok(file);
                }
                Some(s) => format!("data:{};base64,{}", s.to_string(), base64),
            }
        };

        self.images.insert(file.into_owned(), dest_file.clone());
        Ok(Cow::Owned(dest_file))
    }

    /// Returns an iterator the the images files mapping
    #[doc(hidden)]
    pub fn images_mapping(&self) -> &HashMap<String, String> {
        &self.images
    }

    /// Add a match between an original file and a dest file
    pub fn add_link<S1: Into<String>, S2:Into<String>>(&mut self, from: S1, to: S2) {
        self.links.insert(from.into(), to.into());
    }

    /// Get a destination link from an original link
    pub fn get_link<'a>(&'a self, from: &'a str) -> &'a str {
        if let Some(link) = self.links.get(from) {
            link
        } else {
            // Try to get a link by changing the extension
            let new_from = format!("{}", Path::new(from)
                .with_extension("md")
                .display());
            println!("trying {}", new_from);
            if let Some(link) = self.links.get(&new_from) {
                link
            } else {
                self.logger.warning(lformat!("Resources: could not find an in-book match for link \
                                        {file}",
                                             file = from));
                from
            }
        }
    }

    /// Tell whether a file name is a local resource or net
    pub fn contains_link(&self, from: &str) -> bool {
        if self.links.contains_key(from) {
            true
        } else {
            // Try to get a link by changing the extension
            let new_from = format!("{}", Path::new(from)
                .with_extension("md")
                .display());
            if self.links.contains_key(&new_from) {
                true
            } else {
                false
            }
        }
    }

    pub fn is_local(path: &str) -> bool {
        !path.contains("://") // todo: use better algorithm
    }

    /// Add a path offset to all linked urls and images src
    pub fn add_offset(link_offset: &Path, image_offset: &Path, ast: &mut [Token]) {
        if link_offset == Path::new("") && image_offset == Path::new("") {
            // nothing do to
            return;
        }
        for mut token in ast {
            match *token {
                Token::Link(ref mut url, _, ref mut v) => {
                    if ResourceHandler::is_local(url) {
                        let new_url = format!("{}", link_offset.join(&url).display());
                        *url = new_url;
                    }
                    Self::add_offset(link_offset, image_offset, v);
                }
                Token::Image(ref mut url, _, ref mut v) |
                Token::StandaloneImage(ref mut url, _, ref mut v) => {
                    if ResourceHandler::is_local(url) {
                        let new_url = format!("{}", image_offset.join(&url).display());
                        *url = new_url;
                    }
                    Self::add_offset(link_offset, image_offset, v);
                }
                _ => {
                    if let Some(ref mut inner) = token.inner_mut() {
                        Self::add_offset(link_offset, image_offset, inner);
                    }
                }
            }
        }
    }
}

/// Get the list of all files, walking recursively in directories
///
/// # Arguments
/// - list: a list of files
/// - base: the path where to get them
///
/// # Returns
/// A list of files (relative to `base`), or an error.
pub fn get_files(list: Vec<String>, base: &str) -> Result<Vec<String>> {
    let mut out: Vec<String> = vec![];
    let base = Path::new(base);
    for path in list {
        let abs_path = base.join(&path);
        let res = fs::metadata(&abs_path);
        match res {
            Err(err) => {
                return Err(Error::render(Source::empty(),
                                         lformat!("error reading file {file}: {error}",
                                                  file = abs_path.display(),
                                                  error = err)))
            }
            Ok(metadata) => {
                if metadata.is_file() {
                    out.push(path);
                } else if metadata.is_dir() {
                    let files = WalkDir::new(&abs_path)
                        .follow_links(true)
                        .into_iter()
                        .filter_map(|e| e.ok())
                        .filter(|e| e.file_type().is_file())
                        .map(|e| {
                            PathBuf::from(e.path()
                                .strip_prefix(base)
                                .unwrap())
                        });
                    for file in files {
                        out.push(file.to_string_lossy().into_owned());
                    }
                } else {
                    return Err(Error::render(Source::empty(),
                                             lformat!("error in epub rendering: {path} is \
                                                       neither a file nor a directory",
                                                      path = &path)));
                }
            }
        }
    }
    Ok(out)
}
<commit_msg>Remove debug println<commit_after>use token::Token;
use logger::Logger;
use error::{Error, Result, Source};

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::borrow::Cow;
use std::fs;
use std::io::Read;

use walkdir::WalkDir;
use rustc_serialize::base64::{self, ToBase64};
use mime_guess;

/// Resource Handler.
///
/// Its task is to make sure that some resource (image, link) is available
/// for the book and to list images used in Markdown files so they can be used for the book
#[derive(Debug)]
pub struct ResourceHandler<'r> {
    /// Maps an original url (e.g.) "foo/Readme.md" to a valid link
    /// (e.g.) chapter3.html
    links: HashMap<String, String>,
    map_images: bool,
    logger: &'r Logger,
    base64: bool,

    /// Maps an original (local) file name to a new file name. Allows to
    /// make sure all image files will be included in e.g. the Epub document.
    #[doc(hidden)]
    pub images: HashMap<String, String>,
}

impl<'r> ResourceHandler<'r> {
    /// Creates a new, empty Resource Handler
    pub fn new(logger: &'r Logger) -> ResourceHandler {
        ResourceHandler {
            links: HashMap::new(),
            images: HashMap::new(),
            map_images: false,
            base64: false,
            logger: logger,
        }
    }

    /// Turns on mapping for image files
    ///
    /// # Argument: an offset (should be book.root)
    pub fn set_images_mapping(&mut self, b: bool) -> &mut Self {
        self.map_images = b;
        self
    }

    /// Sets base64 mode for image mapping
    ///
    /// If set to true, instead of returning a destination file path,
    /// `map_image` will include the image as base64
    pub fn set_base64(&mut self, b: bool) {
        self.base64 = b;
    }

    /// Add a local image file and get the resulting transformed
    /// file name
    pub fn map_image<'a, S: Into<Cow<'a, str>>>(&'a mut self,
                                                source: &Source,
                                                file: S)
                                                -> Result<Cow<'a, str>> {
        // If image is not local, do nothing much
        let file = file.into();
        if !Self::is_local(file.as_ref()) {
            self.logger
                .warning(lformat!("Resources: book includes non-local image {file}, which might \
                                   cause problem for proper inclusion.",
                                  file = file));
            return Ok(file);
        }

        // Check exisence of the file
        if fs::metadata(file.as_ref()).is_err() {
            return Err(Error::file_not_found(source, lformat!("image"), format!("{}", file)));
        }

        // if image mapping is not activated do nothing else
        if !self.map_images {
            return Ok(file);
        }

        // If this image has already been registered, returns it
        if self.images.contains_key(file.as_ref()) {
            return Ok(Cow::Borrowed(&self.images[file.as_ref()]));
        }

        // Else, create a new file name that has same extension
        // (or a base64 version of the file)
        let dest_file = if !(self.base64) {
            if let Some(extension) = Path::new(file.as_ref()).extension() {
                format!("images/image_{}.{}",
                        self.images.len(),
                        extension.to_string_lossy())
            } else {
                self.logger
                    .warning(lformat!("Resources: book includes image {file} which doesn't have \
                                       an extension",
                                      file = file));
                format!("images/image_{}", self.images.len())
            }
        } else {
            let mut f = match fs::File::open(file.as_ref()) {
                Ok(f) => f,
                Err(_) => {
                    return Err(Error::file_not_found(source,
                                                     lformat!("image"),
                                                     format!("{}", file)));
                }
            };
            let mut content: Vec<u8> = vec![];
            if f.read_to_end(&mut content).is_err() {
                self.logger.error(lformat!("Resources: could not read file {file}", file = file));
                return Ok(file);
            }
            let base64 = content.to_base64(base64::STANDARD);
            match mime_guess::guess_mime_type_opt(file.as_ref()) {
                None => {
                    self.logger
                        .error(lformat!("Resources: could not guess mime type of file {file}",
                                        file = file));
                    return Ok(file);
                }
                Some(s) => format!("data:{};base64,{}", s.to_string(), base64),
            }
        };

        self.images.insert(file.into_owned(), dest_file.clone());
        Ok(Cow::Owned(dest_file))
    }

    /// Returns an iterator the the images files mapping
    #[doc(hidden)]
    pub fn images_mapping(&self) -> &HashMap<String, String> {
        &self.images
    }

    /// Add a match between an original file and a dest file
    pub fn add_link<S1: Into<String>, S2:Into<String>>(&mut self, from: S1, to: S2) {
        self.links.insert(from.into(), to.into());
    }

    /// Get a destination link from an original link
    pub fn get_link<'a>(&'a self, from: &'a str) -> &'a str {
        if let Some(link) = self.links.get(from) {
            link
        } else {
            // Try to get a link by changing the extension
            let new_from = format!("{}", Path::new(from)
                .with_extension("md")
                .display());
            if let Some(link) = self.links.get(&new_from) {
                link
            } else {
                self.logger.warning(lformat!("Resources: could not find an in-book match for link \
                                        {file}",
                                             file = from));
                from
            }
        }
    }

    /// Tell whether a file name is a local resource or net
    pub fn contains_link(&self, from: &str) -> bool {
        if self.links.contains_key(from) {
            true
        } else {
            // Try to get a link by changing the extension
            let new_from = format!("{}", Path::new(from)
                .with_extension("md")
                .display());
            if self.links.contains_key(&new_from) {
                true
            } else {
                false
            }
        }
    }

    pub fn is_local(path: &str) -> bool {
        !path.contains("://") // todo: use better algorithm
    }

    /// Add a path offset to all linked urls and images src
    pub fn add_offset(link_offset: &Path, image_offset: &Path, ast: &mut [Token]) {
        if link_offset == Path::new("") && image_offset == Path::new("") {
            // nothing do to
            return;
        }
        for mut token in ast {
            match *token {
                Token::Link(ref mut url, _, ref mut v) => {
                    if ResourceHandler::is_local(url) {
                        let new_url = format!("{}", link_offset.join(&url).display());
                        *url = new_url;
                    }
                    Self::add_offset(link_offset, image_offset, v);
                }
                Token::Image(ref mut url, _, ref mut v) |
                Token::StandaloneImage(ref mut url, _, ref mut v) => {
                    if ResourceHandler::is_local(url) {
                        let new_url = format!("{}", image_offset.join(&url).display());
                        *url = new_url;
                    }
                    Self::add_offset(link_offset, image_offset, v);
                }
                _ => {
                    if let Some(ref mut inner) = token.inner_mut() {
                        Self::add_offset(link_offset, image_offset, inner);
                    }
                }
            }
        }
    }
}

/// Get the list of all files, walking recursively in directories
///
/// # Arguments
/// - list: a list of files
/// - base: the path where to get them
///
/// # Returns
/// A list of files (relative to `base`), or an error.
pub fn get_files(list: Vec<String>, base: &str) -> Result<Vec<String>> {
    let mut out: Vec<String> = vec![];
    let base = Path::new(base);
    for path in list {
        let abs_path = base.join(&path);
        let res = fs::metadata(&abs_path);
        match res {
            Err(err) => {
                return Err(Error::render(Source::empty(),
                                         lformat!("error reading file {file}: {error}",
                                                  file = abs_path.display(),
                                                  error = err)))
            }
            Ok(metadata) => {
                if metadata.is_file() {
                    out.push(path);
                } else if metadata.is_dir() {
                    let files = WalkDir::new(&abs_path)
                        .follow_links(true)
                        .into_iter()
                        .filter_map(|e| e.ok())
                        .filter(|e| e.file_type().is_file())
                        .map(|e| {
                            PathBuf::from(e.path()
                                .strip_prefix(base)
                                .unwrap())
                        });
                    for file in files {
                        out.push(file.to_string_lossy().into_owned());
                    }
                } else {
                    return Err(Error::render(Source::empty(),
                                             lformat!("error in epub rendering: {path} is \
                                                       neither a file nor a directory",
                                                      path = &path)));
                }
            }
        }
    }
    Ok(out)
}
<|endoftext|>


<commit_before>// Copyright 2015 Pierre Talbot (IRCAM)

// 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.

use solver::fd::event::*;
use solver::event::*;
use solver::variable::Variable;
use interval::ncollections::ops::*;
use interval::ops::*;

use std::fmt::{Formatter, Display, Error};

#[derive(Copy, PartialEq, Eq, Debug, Clone)]
pub struct FDVar<Domain> {
  idx: usize,
  dom: Domain
}

impl<Domain: Display> Display for FDVar<Domain> {
  fn fmt(&self, formatter: &mut Formatter) -> Result<(), Error> {
    formatter.write_fmt(format_args!("({}, {})", self.idx, self.dom))
  }
}

impl<Domain> Variable for FDVar<Domain> where
  Domain: Cardinality
{
  type Domain = Domain;
  type Event = FDEvent;

  fn new(idx: usize, dom: Domain) -> FDVar<Domain> {
    assert!(!dom.is_empty());
    FDVar {
      idx: idx,
      dom: dom
    }
  }
}

pub trait VarIndex
{
  fn index(&self) -> usize;
}

impl<Domain> VarIndex for FDVar<Domain>
{
  fn index(&self) -> usize {
    self.idx
  }
}

pub trait Failure
{
  fn is_failed(&self) -> bool;
}

impl<Domain> Failure for FDVar<Domain> where
  Domain: Cardinality
{
  fn is_failed(&self) -> bool {
    self.dom.is_empty()
  }
}

trait VarDomain: Bounded + Cardinality + Subset
{}

impl<R> VarDomain for R where
  R: Bounded + Cardinality + Subset
{}

pub trait EventUpdate<Item>
{
  fn event_update(&mut self, value: Item,
    events: &mut Vec<(usize, FDEvent)>) -> bool;
}

impl<Domain> EventUpdate<Domain> for FDVar<Domain> where
  Domain: VarDomain
{
  fn event_update(&mut self, dom: Domain,
    events: &mut Vec<(usize, FDEvent)>) -> bool
  {
    assert!(dom.is_subset(&self.dom), "Domain update must be monotonic.");
    if dom.is_empty() { false } // Failure
    else {
      if let Some(event) = FDEvent::new(&dom, &self.dom) {
        events.push((self.idx, event));
        self.dom = dom;
      }
      true
    }
  }
}

impl<Domain> Bounded for FDVar<Domain> where
  Domain: Bounded
{
  type Bound = Domain::Bound;

  fn lower(&self) -> Domain::Bound {
    self.dom.lower()
  }

  fn upper(&self) -> Domain::Bound {
    self.dom.upper()
  }
}

impl<Domain> Disjoint for FDVar<Domain> where
  Domain: Disjoint
{
  fn is_disjoint(&self, other: &FDVar<Domain>) -> bool {
    self.dom.is_disjoint(&other.dom)
  }
}

impl<Domain> Contains<Domain::Bound> for FDVar<Domain> where
  Domain: Bounded + Contains<<Domain as Bounded>::Bound>
{
  fn contains(&self, value: &Domain::Bound) -> bool {
    self.dom.contains(value)
  }
}

pub trait EventShrinkLeft<Bound>
{
  fn event_shrink_left(&mut self, lb: Bound,
    events: &mut Vec<(usize, FDEvent)>) -> bool;
}

impl<Domain> EventShrinkLeft<Domain::Bound> for FDVar<Domain> where
  Domain: VarDomain + ShrinkLeft<<Domain as Bounded>::Bound>
{
  fn event_shrink_left(&mut self, lb: Domain::Bound,
    events: &mut Vec<(usize, FDEvent)>) -> bool
  {
    let new = self.dom.shrink_left(lb);
    self.event_update(new, events)
  }
}

pub trait EventShrinkRight<Bound>
{
  fn event_shrink_right(&mut self, ub: Bound,
    events: &mut Vec<(usize, FDEvent)>) -> bool;
}

impl<Domain> EventShrinkRight<Domain::Bound> for FDVar<Domain> where
  Domain: VarDomain + ShrinkRight<<Domain as Bounded>::Bound>
{
  fn event_shrink_right(&mut self, ub: Domain::Bound,
    events: &mut Vec<(usize, FDEvent)>) -> bool
  {
    let new = self.dom.shrink_right(ub);
    self.event_update(new, events)
  }
}

pub trait EventRemove<Item>
{
  fn event_remove(&mut self, value: Item,
    events: &mut Vec<(usize, FDEvent)>) -> bool;
}

impl<Domain> EventRemove<Domain::Bound> for FDVar<Domain> where
  Domain: VarDomain + Difference<<Domain as Bounded>::Bound, Output=Domain>
{
  fn event_remove(&mut self, value: Domain::Bound,
    events: &mut Vec<(usize, FDEvent)>) -> bool
  {
    let new = self.dom.difference(&value);
    self.event_update(new, events)
  }
}

pub trait EventIntersection<RHS = Self>
{
  fn event_intersection(&mut self, other: &mut RHS,
    events: &mut Vec<(usize, FDEvent)>) -> bool;
}

impl<Domain> EventIntersection for FDVar<Domain> where
  Domain: VarDomain + Intersection<Output=Domain> + Clone
{
  fn event_intersection(&mut self, other: &mut FDVar<Domain>,
    events: &mut Vec<(usize, FDEvent)>) -> bool
  {
    let new = self.dom.intersection(&other.dom);
    self.event_update(new.clone(), events) &&
    other.event_update(new, events)
  }
}

#[cfg(test)]
mod test {
  use super::*;
  use solver::fd::event::*;
  use solver::fd::event::FDEvent::*;
  use solver::variable::Variable;
  use interval::interval::*;
  use interval::ncollections::ops::*;

  #[test]
  fn var_update_test() {
    let dom0_10 = (0,10).to_interval();
    let dom0_9 = (0,5).to_interval();
    let dom1_10 = (5,10).to_interval();
    let dom1_9 = (1,9).to_interval();
    let dom0_0 = (0,0).to_interval();
    let empty = Interval::empty();
    let var0_10 = Variable::new(0, dom0_10);

    var_update_test_one(var0_10, dom0_10, vec![], true);
    var_update_test_one(var0_10, empty, vec![], false);
    var_update_test_one(var0_10, dom0_0, vec![Assignment], true);
    var_update_test_one(var0_10, dom1_10, vec![Bound], true);
    var_update_test_one(var0_10, dom0_9, vec![Bound], true);
    var_update_test_one(var0_10, dom1_9, vec![Bound], true);
  }

  fn var_update_test_one(mut var: FDVar<Interval<i32>>, dom: Interval<i32>, expect: Vec<FDEvent>, expect_success: bool) {
    let mut events = vec![];
    assert_eq!(var.event_update(dom, &mut events), expect_success);
    if expect_success {
      assert_eq_events(events, expect);
      assert_eq!(var.dom, dom);
    }
  }

  fn assert_eq_events(events: Vec<(usize, FDEvent)>, expect: Vec<FDEvent>) {
    for ((_,ev), expect) in events.into_iter().zip(expect.into_iter()) {
      assert_eq!(ev, expect);
    }
  }

  #[test]
  fn var_update_bound() {
    let dom0_10 = (0,10).to_interval();
    let var0_10 = Variable::new(0, dom0_10);

    var_update_lb_test_one(var0_10, 0, vec![], true);
    var_update_lb_test_one(var0_10, 10, vec![Assignment], true);
    var_update_lb_test_one(var0_10, 1, vec![Bound], true);
    var_update_lb_test_one(var0_10, 11, vec![], false);

    var_update_ub_test_one(var0_10, 10, vec![], true);
    var_update_ub_test_one(var0_10, 0, vec![Assignment], true);
    var_update_ub_test_one(var0_10, 1, vec![Bound], true);
    var_update_ub_test_one(var0_10, -1, vec![], false);
  }

  fn var_update_lb_test_one(mut var: FDVar<Interval<i32>>, lb: i32, expect: Vec<FDEvent>, expect_success: bool) {
    let ub = var.upper();
    let mut events = vec![];
    assert_eq!(var.event_shrink_left(lb, &mut events), expect_success);
    if expect_success {
      assert_eq_events(events, expect);
      assert_eq!(var.dom, (lb,ub).to_interval());
    }
  }

  fn var_update_ub_test_one(mut var: FDVar<Interval<i32>>, ub: i32, expect: Vec<FDEvent>, expect_success: bool) {
    let lb = var.lower();
    let mut events = vec![];
    assert_eq!(var.event_shrink_right(ub, &mut events), expect_success);
    if expect_success {
      assert_eq_events(events, expect);
      assert_eq!(var.dom, (lb,ub).to_interval());
    }
  }

  #[test]
  fn var_intersection_test() {
    let var0_10 = Variable::new(1, (0,10).to_interval());
    let var10_20 = Variable::new(2, (10,20).to_interval());
    let var11_20 = Variable::new(3, (11,20).to_interval());
    let var1_9 = Variable::new(4, (1,9).to_interval());

    var_intersection_test_one(var0_10, var10_20, Some(vec![(1, Assignment), (2, Assignment)]));
    var_intersection_test_one(var0_10, var1_9, Some(vec![(1, Bound)]));
    var_intersection_test_one(var1_9, var0_10, Some(vec![(1, Bound)]));
    var_intersection_test_one(var0_10, var11_20, None);
  }

  fn var_intersection_test_one(mut v1: FDVar<Interval<i32>>, mut v2: FDVar<Interval<i32>>, events: Option<Vec<(usize, FDEvent)>>) {
    let mut ev = vec![];
    let res =
      if v1.event_intersection(&mut v2, &mut ev) {
        Some(ev)
      }
      else {
        None
      };
    assert_eq!(res, events);
  }

  #[test]
  #[should_panic]
  fn var_non_monotonic_update_singleton() {
    let dom0_10 = (0,10).to_interval();
    let mut var0_10: FDVar<Interval<i32>> = Variable::new(0, dom0_10);
    let dom11_11 = 11.to_interval();

    var0_10.event_update(dom11_11, &mut vec![]);
  }

  #[test]
  #[should_panic]
  fn var_non_monotonic_update_widen() {
    let dom0_10 = (0,10).to_interval();
    let mut var0_10: FDVar<Interval<i32>> = Variable::new(0, dom0_10);
    let domm5_15 = (-5, 15).to_interval();

    var0_10.event_update(domm5_15, &mut vec![]);
  }
}<commit_msg>Event traits function takes the Event as a type parameter.<commit_after>// Copyright 2015 Pierre Talbot (IRCAM)

// 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.

use solver::fd::event::*;
use solver::event::*;
use solver::variable::Variable;
use interval::ncollections::ops::*;
use interval::ops::*;

use std::fmt::{Formatter, Display, Error};

#[derive(Copy, PartialEq, Eq, Debug, Clone)]
pub struct FDVar<Domain> {
  idx: usize,
  dom: Domain
}

impl<Domain: Display> Display for FDVar<Domain> {
  fn fmt(&self, formatter: &mut Formatter) -> Result<(), Error> {
    formatter.write_fmt(format_args!("({}, {})", self.idx, self.dom))
  }
}

impl<Domain> Variable for FDVar<Domain> where
  Domain: Cardinality
{
  type Domain = Domain;
  type Event = FDEvent;

  fn new(idx: usize, dom: Domain) -> FDVar<Domain> {
    assert!(!dom.is_empty());
    FDVar {
      idx: idx,
      dom: dom
    }
  }
}

pub trait VarIndex
{
  fn index(&self) -> usize;
}

impl<Domain> VarIndex for FDVar<Domain>
{
  fn index(&self) -> usize {
    self.idx
  }
}

pub trait Failure
{
  fn is_failed(&self) -> bool;
}

impl<Domain> Failure for FDVar<Domain> where
  Domain: Cardinality
{
  fn is_failed(&self) -> bool {
    self.dom.is_empty()
  }
}

pub trait VarDomain: Bounded + Cardinality + Subset
{}

impl<R> VarDomain for R where
  R: Bounded + Cardinality + Subset
{}

pub trait EventUpdate<Domain: VarDomain>
{
  fn event_update<Event>(&mut self, value: Domain,
    events: &mut Vec<(usize, Event)>) -> bool
   where
    Event: MonotonicEvent<Domain>;
}

impl<Domain> EventUpdate<Domain> for FDVar<Domain> where
  Domain: VarDomain
{
  fn event_update<Event>(&mut self, dom: Domain,
    events: &mut Vec<(usize, Event)>) -> bool
   where
    Event: MonotonicEvent<Domain>
  {
    assert!(dom.is_subset(&self.dom), "Domain update must be monotonic.");
    if dom.is_empty() { false } // Failure
    else {
      if let Some(event) = Event::new(&dom, &self.dom) {
        events.push((self.idx, event));
        self.dom = dom;
      }
      true
    }
  }
}

impl<Domain> Bounded for FDVar<Domain> where
  Domain: Bounded
{
  type Bound = Domain::Bound;

  fn lower(&self) -> Domain::Bound {
    self.dom.lower()
  }

  fn upper(&self) -> Domain::Bound {
    self.dom.upper()
  }
}

impl<Domain> Disjoint for FDVar<Domain> where
  Domain: Disjoint
{
  fn is_disjoint(&self, other: &FDVar<Domain>) -> bool {
    self.dom.is_disjoint(&other.dom)
  }
}

impl<Domain> Contains<Domain::Bound> for FDVar<Domain> where
  Domain: Bounded + Contains<<Domain as Bounded>::Bound>
{
  fn contains(&self, value: &Domain::Bound) -> bool {
    self.dom.contains(value)
  }
}

pub trait EventShrinkLeft<Domain: VarDomain>
{
  fn event_shrink_left<Event>(&mut self, lb: Domain::Bound,
    events: &mut Vec<(usize, Event)>) -> bool
   where
    Event: MonotonicEvent<Domain>;
}

impl<Domain> EventShrinkLeft<Domain> for FDVar<Domain> where
  Domain: VarDomain + ShrinkLeft<<Domain as Bounded>::Bound>
{
  fn event_shrink_left<Event>(&mut self, lb: Domain::Bound,
    events: &mut Vec<(usize, Event)>) -> bool
   where
    Event: MonotonicEvent<Domain>
  {
    let new = self.dom.shrink_left(lb);
    self.event_update(new, events)
  }
}

pub trait EventShrinkRight<Domain: VarDomain>
{
  fn event_shrink_right<Event>(&mut self, ub: Domain::Bound,
    events: &mut Vec<(usize, Event)>) -> bool
   where
    Event: MonotonicEvent<Domain>;
}

impl<Domain> EventShrinkRight<Domain> for FDVar<Domain> where
  Domain: VarDomain + ShrinkRight<<Domain as Bounded>::Bound>
{
  fn event_shrink_right<Event>(&mut self, ub: Domain::Bound,
    events: &mut Vec<(usize, Event)>) -> bool
   where
    Event: MonotonicEvent<Domain>
  {
    let new = self.dom.shrink_right(ub);
    self.event_update(new, events)
  }
}

pub trait EventRemove<Domain: VarDomain>
{
  fn event_remove<Event>(&mut self, value: Domain::Bound,
    events: &mut Vec<(usize, Event)>) -> bool
  where
    Event: MonotonicEvent<Domain>;
}

impl<Domain> EventRemove<Domain> for FDVar<Domain> where
  Domain: VarDomain + Difference<<Domain as Bounded>::Bound, Output=Domain>
{
  fn event_remove<Event>(&mut self, value: Domain::Bound,
    events: &mut Vec<(usize, Event)>) -> bool
   where
    Event: MonotonicEvent<Domain>
  {
    let new = self.dom.difference(&value);
    self.event_update(new, events)
  }
}

pub trait EventIntersection<Domain, RHS = Self>
{
  fn event_intersection<Event>(&mut self, other: &mut RHS,
    events: &mut Vec<(usize, Event)>) -> bool
   where
    Event: MonotonicEvent<Domain>;
}

impl<Domain> EventIntersection<Domain> for FDVar<Domain> where
  Domain: VarDomain + Intersection<Output=Domain> + Clone
{
  fn event_intersection<Event>(&mut self, other: &mut FDVar<Domain>,
    events: &mut Vec<(usize, Event)>) -> bool
   where
    Event: MonotonicEvent<Domain>
  {
    let new = self.dom.intersection(&other.dom);
    self.event_update(new.clone(), events) &&
    other.event_update(new, events)
  }
}

#[cfg(test)]
mod test {
  use super::*;
  use solver::fd::event::*;
  use solver::fd::event::FDEvent::*;
  use solver::variable::Variable;
  use interval::interval::*;
  use interval::ncollections::ops::*;

  #[test]
  fn var_update_test() {
    let dom0_10 = (0,10).to_interval();
    let dom0_9 = (0,5).to_interval();
    let dom1_10 = (5,10).to_interval();
    let dom1_9 = (1,9).to_interval();
    let dom0_0 = (0,0).to_interval();
    let empty = Interval::empty();
    let var0_10 = Variable::new(0, dom0_10);

    var_update_test_one(var0_10, dom0_10, vec![], true);
    var_update_test_one(var0_10, empty, vec![], false);
    var_update_test_one(var0_10, dom0_0, vec![Assignment], true);
    var_update_test_one(var0_10, dom1_10, vec![Bound], true);
    var_update_test_one(var0_10, dom0_9, vec![Bound], true);
    var_update_test_one(var0_10, dom1_9, vec![Bound], true);
  }

  fn var_update_test_one(mut var: FDVar<Interval<i32>>, dom: Interval<i32>, expect: Vec<FDEvent>, expect_success: bool) {
    let mut events = vec![];
    assert_eq!(var.event_update(dom, &mut events), expect_success);
    if expect_success {
      assert_eq_events(events, expect);
      assert_eq!(var.dom, dom);
    }
  }

  fn assert_eq_events(events: Vec<(usize, FDEvent)>, expect: Vec<FDEvent>) {
    for ((_,ev), expect) in events.into_iter().zip(expect.into_iter()) {
      assert_eq!(ev, expect);
    }
  }

  #[test]
  fn var_update_bound() {
    let dom0_10 = (0,10).to_interval();
    let var0_10 = Variable::new(0, dom0_10);

    var_update_lb_test_one(var0_10, 0, vec![], true);
    var_update_lb_test_one(var0_10, 10, vec![Assignment], true);
    var_update_lb_test_one(var0_10, 1, vec![Bound], true);
    var_update_lb_test_one(var0_10, 11, vec![], false);

    var_update_ub_test_one(var0_10, 10, vec![], true);
    var_update_ub_test_one(var0_10, 0, vec![Assignment], true);
    var_update_ub_test_one(var0_10, 1, vec![Bound], true);
    var_update_ub_test_one(var0_10, -1, vec![], false);
  }

  fn var_update_lb_test_one(mut var: FDVar<Interval<i32>>, lb: i32, expect: Vec<FDEvent>, expect_success: bool) {
    let ub = var.upper();
    let mut events = vec![];
    assert_eq!(var.event_shrink_left(lb, &mut events), expect_success);
    if expect_success {
      assert_eq_events(events, expect);
      assert_eq!(var.dom, (lb,ub).to_interval());
    }
  }

  fn var_update_ub_test_one(mut var: FDVar<Interval<i32>>, ub: i32, expect: Vec<FDEvent>, expect_success: bool) {
    let lb = var.lower();
    let mut events = vec![];
    assert_eq!(var.event_shrink_right(ub, &mut events), expect_success);
    if expect_success {
      assert_eq_events(events, expect);
      assert_eq!(var.dom, (lb,ub).to_interval());
    }
  }

  #[test]
  fn var_intersection_test() {
    let var0_10 = Variable::new(1, (0,10).to_interval());
    let var10_20 = Variable::new(2, (10,20).to_interval());
    let var11_20 = Variable::new(3, (11,20).to_interval());
    let var1_9 = Variable::new(4, (1,9).to_interval());

    var_intersection_test_one(var0_10, var10_20, Some(vec![(1, Assignment), (2, Assignment)]));
    var_intersection_test_one(var0_10, var1_9, Some(vec![(1, Bound)]));
    var_intersection_test_one(var1_9, var0_10, Some(vec![(1, Bound)]));
    var_intersection_test_one(var0_10, var11_20, None);
  }

  fn var_intersection_test_one(mut v1: FDVar<Interval<i32>>, mut v2: FDVar<Interval<i32>>, events: Option<Vec<(usize, FDEvent)>>) {
    let mut ev = vec![];
    let res =
      if v1.event_intersection(&mut v2, &mut ev) {
        Some(ev)
      }
      else {
        None
      };
    assert_eq!(res, events);
  }

  #[test]
  #[should_panic]
  fn var_non_monotonic_update_singleton() {
    let dom0_10 = (0,10).to_interval();
    let mut var0_10: FDVar<Interval<i32>> = Variable::new(0, dom0_10);
    let dom11_11 = 11.to_interval();

    var0_10.event_update::<FDEvent>(dom11_11, &mut vec![]);
  }

  #[test]
  #[should_panic]
  fn var_non_monotonic_update_widen() {
    let dom0_10 = (0,10).to_interval();
    let mut var0_10: FDVar<Interval<i32>> = Variable::new(0, dom0_10);
    let domm5_15 = (-5, 15).to_interval();

    var0_10.event_update::<FDEvent>(domm5_15, &mut vec![]);
  }
}<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The local, garbage collected heap

use cast;
use iter::Iterator;
use mem;
use ops::Drop;
use option::{Option, None, Some};
use ptr;
use ptr::RawPtr;
use rt::env;
use rt::global_heap;
use rt::local::Local;
use rt::task::Task;
use unstable::raw;
use vec::ImmutableVector;

// This has no meaning with out rtdebug also turned on.
#[cfg(rtdebug)]
static TRACK_ALLOCATIONS: int = 0;
#[cfg(rtdebug)]
static MAGIC: u32 = 0xbadc0ffe;

pub type Box = raw::Box<()>;

pub struct MemoryRegion {
    priv allocations: ~[*AllocHeader],
    priv live_allocations: uint,
}

pub struct LocalHeap {
    priv memory_region: MemoryRegion,

    priv poison_on_free: bool,
    priv live_allocs: *mut raw::Box<()>,
}

impl LocalHeap {
    #[inline]
    pub fn new() -> LocalHeap {
        let region = MemoryRegion {
            allocations: ~[],
            live_allocations: 0,
        };
        LocalHeap {
            memory_region: region,
            poison_on_free: env::poison_on_free(),
            live_allocs: ptr::mut_null(),
        }
    }

    #[inline]
    pub fn alloc(&mut self, drop_glue: fn(*mut u8), size: uint, align: uint) -> *mut Box {
        let total_size = global_heap::get_box_size(size, align);
        let alloc = self.memory_region.malloc(total_size);
        {
            // Make sure that we can't use `mybox` outside of this scope
            let mybox: &mut Box = unsafe { cast::transmute(alloc) };
            // Clear out this box, and move it to the front of the live
            // allocations list
            mybox.drop_glue = drop_glue;
            mybox.ref_count = 1;
            mybox.prev = ptr::mut_null();
            mybox.next = self.live_allocs;
            if !self.live_allocs.is_null() {
                unsafe { (*self.live_allocs).prev = alloc; }
            }
            self.live_allocs = alloc;
        }
        return alloc;
    }

    #[inline]
    pub fn realloc(&mut self, ptr: *mut Box, size: uint) -> *mut Box {
        // Make sure that we can't use `mybox` outside of this scope
        let total_size = size + mem::size_of::<Box>();
        let new_box = self.memory_region.realloc(ptr, total_size);
        {
            // Fix links because we could have moved around
            let mybox: &mut Box = unsafe { cast::transmute(new_box) };
            if !mybox.prev.is_null() {
                unsafe { (*mybox.prev).next = new_box; }
            }
            if !mybox.next.is_null() {
                unsafe { (*mybox.next).prev = new_box; }
            }
        }
        if self.live_allocs == ptr {
            self.live_allocs = new_box;
        }
        return new_box;
    }

    #[inline]
    pub fn free(&mut self, alloc: *mut Box) {
        {
            // Make sure that we can't use `mybox` outside of this scope
            let mybox: &mut Box = unsafe { cast::transmute(alloc) };

            // Unlink it from the linked list
            if !mybox.prev.is_null() {
                unsafe { (*mybox.prev).next = mybox.next; }
            }
            if !mybox.next.is_null() {
                unsafe { (*mybox.next).prev = mybox.prev; }
            }
            if self.live_allocs == alloc {
                self.live_allocs = mybox.next;
            }
        }

        self.memory_region.free(alloc);
    }
}

impl Drop for LocalHeap {
    fn drop(&mut self) {
        assert!(self.live_allocs.is_null());
    }
}

#[cfg(rtdebug)]
struct AllocHeader {
    magic: u32,
    index: i32,
    size: u32,
}
#[cfg(not(rtdebug))]
struct AllocHeader;

impl AllocHeader {
    #[cfg(rtdebug)]
    fn init(&mut self, size: u32) {
        if TRACK_ALLOCATIONS > 0 {
            self.magic = MAGIC;
            self.index = -1;
            self.size = size;
        }
    }
    #[cfg(not(rtdebug))]
    fn init(&mut self, _size: u32) {}

    #[cfg(rtdebug)]
    fn assert_sane(&self) {
        if TRACK_ALLOCATIONS > 0 {
            rtassert!(self.magic == MAGIC);
        }
    }
    #[cfg(not(rtdebug))]
    fn assert_sane(&self) {}

    #[cfg(rtdebug)]
    fn update_size(&mut self, size: u32) {
        if TRACK_ALLOCATIONS > 0 {
            self.size = size;
        }
    }
    #[cfg(not(rtdebug))]
    fn update_size(&mut self, _size: u32) {}

    fn as_box(&mut self) -> *mut Box {
        let myaddr: uint = unsafe { cast::transmute(self) };
        (myaddr + AllocHeader::size()) as *mut Box
    }

    fn size() -> uint {
        // For some platforms, 16 byte alignment is required.
        let ptr_size = 16;
        let header_size = mem::size_of::<AllocHeader>();
        return (header_size + ptr_size - 1) / ptr_size * ptr_size;
    }

    fn from(a_box: *mut Box) -> *mut AllocHeader {
        (a_box as uint - AllocHeader::size()) as *mut AllocHeader
    }
}

impl MemoryRegion {
    #[inline]
    fn malloc(&mut self, size: uint) -> *mut Box {
        let total_size = size + AllocHeader::size();
        let alloc: *AllocHeader = unsafe {
            global_heap::malloc_raw(total_size) as *AllocHeader
        };

        let alloc: &mut AllocHeader = unsafe { cast::transmute(alloc) };
        alloc.init(size as u32);
        self.claim(alloc);
        self.live_allocations += 1;

        return alloc.as_box();
    }

    #[inline]
    fn realloc(&mut self, alloc: *mut Box, size: uint) -> *mut Box {
        rtassert!(!alloc.is_null());
        let orig_alloc = AllocHeader::from(alloc);
        unsafe { (*orig_alloc).assert_sane(); }

        let total_size = size + AllocHeader::size();
        let alloc: *AllocHeader = unsafe {
            global_heap::realloc_raw(orig_alloc as *mut u8,
                                     total_size) as *AllocHeader
        };

        let alloc: &mut AllocHeader = unsafe { cast::transmute(alloc) };
        alloc.assert_sane();
        alloc.update_size(size as u32);
        self.update(alloc, orig_alloc as *AllocHeader);
        return alloc.as_box();
    }

    #[inline]
    fn free(&mut self, alloc: *mut Box) {
        rtassert!(!alloc.is_null());
        let alloc = AllocHeader::from(alloc);
        unsafe {
            (*alloc).assert_sane();
            self.release(cast::transmute(alloc));
            rtassert!(self.live_allocations > 0);
            self.live_allocations -= 1;
            global_heap::exchange_free(alloc as *u8)
        }
    }

    #[cfg(rtdebug)]
    fn claim(&mut self, alloc: &mut AllocHeader) {
        alloc.assert_sane();
        if TRACK_ALLOCATIONS > 1 {
            alloc.index = self.allocations.len() as i32;
            self.allocations.push(&*alloc as *AllocHeader);
        }
    }
    #[cfg(not(rtdebug))]
    #[inline]
    fn claim(&mut self, _alloc: &mut AllocHeader) {}

    #[cfg(rtdebug)]
    fn release(&mut self, alloc: &AllocHeader) {
        alloc.assert_sane();
        if TRACK_ALLOCATIONS > 1 {
            rtassert!(self.allocations[alloc.index] == alloc as *AllocHeader);
            self.allocations[alloc.index] = ptr::null();
        }
    }
    #[cfg(not(rtdebug))]
    #[inline]
    fn release(&mut self, _alloc: &AllocHeader) {}

    #[cfg(rtdebug)]
    fn update(&mut self, alloc: &mut AllocHeader, orig: *AllocHeader) {
        alloc.assert_sane();
        if TRACK_ALLOCATIONS > 1 {
            rtassert!(self.allocations[alloc.index] == orig);
            self.allocations[alloc.index] = &*alloc as *AllocHeader;
        }
    }
    #[cfg(not(rtdebug))]
    #[inline]
    fn update(&mut self, _alloc: &mut AllocHeader, _orig: *AllocHeader) {}
}

impl Drop for MemoryRegion {
    fn drop(&mut self) {
        if self.live_allocations != 0 {
            rtabort!("leaked managed memory ({} objects)", self.live_allocations);
        }
        rtassert!(self.allocations.iter().all(|s| s.is_null()));
    }
}

#[inline]
pub unsafe fn local_malloc(drop_glue: fn(*mut u8), size: uint, align: uint) -> *u8 {
    // FIXME: Unsafe borrow for speed. Lame.
    let task: Option<*mut Task> = Local::try_unsafe_borrow();
    match task {
        Some(task) => {
            (*task).heap.alloc(drop_glue, size, align) as *u8
        }
        None => rtabort!("local malloc outside of task")
    }
}

// A little compatibility function
#[inline]
pub unsafe fn local_free(ptr: *u8) {
    // FIXME: Unsafe borrow for speed. Lame.
    let task_ptr: Option<*mut Task> = Local::try_unsafe_borrow();
    match task_ptr {
        Some(task) => {
            (*task).heap.free(ptr as *mut Box)
        }
        None => rtabort!("local free outside of task")
    }
}

pub fn live_allocs() -> *mut Box {
    let mut task = Local::borrow(None::<Task>);
    task.get().heap.live_allocs
}

#[cfg(test)]
mod bench {
    use extra::test::BenchHarness;

    #[bench]
    fn alloc_managed_small(bh: &mut BenchHarness) {
        bh.iter(|| { @10; });
    }

    #[bench]
    fn alloc_managed_big(bh: &mut BenchHarness) {
        bh.iter(|| { @([10, ..1000]); });
    }
}
<commit_msg>Don't allocate in LocalHeap::new()<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The local, garbage collected heap

use cast;
use iter::Iterator;
use mem;
use ops::Drop;
use option::{Option, None, Some};
use ptr;
use ptr::RawPtr;
use rt::env;
use rt::global_heap;
use rt::local::Local;
use rt::task::Task;
use unstable::raw;
use vec::ImmutableVector;
use vec_ng::Vec;

// This has no meaning with out rtdebug also turned on.
#[cfg(rtdebug)]
static TRACK_ALLOCATIONS: int = 0;
#[cfg(rtdebug)]
static MAGIC: u32 = 0xbadc0ffe;

pub type Box = raw::Box<()>;

pub struct MemoryRegion {
    priv allocations: Vec<*AllocHeader>,
    priv live_allocations: uint,
}

pub struct LocalHeap {
    priv memory_region: MemoryRegion,

    priv poison_on_free: bool,
    priv live_allocs: *mut raw::Box<()>,
}

impl LocalHeap {
    #[inline]
    pub fn new() -> LocalHeap {
        let region = MemoryRegion {
            allocations: Vec::new(),
            live_allocations: 0,
        };
        LocalHeap {
            memory_region: region,
            poison_on_free: env::poison_on_free(),
            live_allocs: ptr::mut_null(),
        }
    }

    #[inline]
    pub fn alloc(&mut self, drop_glue: fn(*mut u8), size: uint, align: uint) -> *mut Box {
        let total_size = global_heap::get_box_size(size, align);
        let alloc = self.memory_region.malloc(total_size);
        {
            // Make sure that we can't use `mybox` outside of this scope
            let mybox: &mut Box = unsafe { cast::transmute(alloc) };
            // Clear out this box, and move it to the front of the live
            // allocations list
            mybox.drop_glue = drop_glue;
            mybox.ref_count = 1;
            mybox.prev = ptr::mut_null();
            mybox.next = self.live_allocs;
            if !self.live_allocs.is_null() {
                unsafe { (*self.live_allocs).prev = alloc; }
            }
            self.live_allocs = alloc;
        }
        return alloc;
    }

    #[inline]
    pub fn realloc(&mut self, ptr: *mut Box, size: uint) -> *mut Box {
        // Make sure that we can't use `mybox` outside of this scope
        let total_size = size + mem::size_of::<Box>();
        let new_box = self.memory_region.realloc(ptr, total_size);
        {
            // Fix links because we could have moved around
            let mybox: &mut Box = unsafe { cast::transmute(new_box) };
            if !mybox.prev.is_null() {
                unsafe { (*mybox.prev).next = new_box; }
            }
            if !mybox.next.is_null() {
                unsafe { (*mybox.next).prev = new_box; }
            }
        }
        if self.live_allocs == ptr {
            self.live_allocs = new_box;
        }
        return new_box;
    }

    #[inline]
    pub fn free(&mut self, alloc: *mut Box) {
        {
            // Make sure that we can't use `mybox` outside of this scope
            let mybox: &mut Box = unsafe { cast::transmute(alloc) };

            // Unlink it from the linked list
            if !mybox.prev.is_null() {
                unsafe { (*mybox.prev).next = mybox.next; }
            }
            if !mybox.next.is_null() {
                unsafe { (*mybox.next).prev = mybox.prev; }
            }
            if self.live_allocs == alloc {
                self.live_allocs = mybox.next;
            }
        }

        self.memory_region.free(alloc);
    }
}

impl Drop for LocalHeap {
    fn drop(&mut self) {
        assert!(self.live_allocs.is_null());
    }
}

#[cfg(rtdebug)]
struct AllocHeader {
    magic: u32,
    index: i32,
    size: u32,
}
#[cfg(not(rtdebug))]
struct AllocHeader;

impl AllocHeader {
    #[cfg(rtdebug)]
    fn init(&mut self, size: u32) {
        if TRACK_ALLOCATIONS > 0 {
            self.magic = MAGIC;
            self.index = -1;
            self.size = size;
        }
    }
    #[cfg(not(rtdebug))]
    fn init(&mut self, _size: u32) {}

    #[cfg(rtdebug)]
    fn assert_sane(&self) {
        if TRACK_ALLOCATIONS > 0 {
            rtassert!(self.magic == MAGIC);
        }
    }
    #[cfg(not(rtdebug))]
    fn assert_sane(&self) {}

    #[cfg(rtdebug)]
    fn update_size(&mut self, size: u32) {
        if TRACK_ALLOCATIONS > 0 {
            self.size = size;
        }
    }
    #[cfg(not(rtdebug))]
    fn update_size(&mut self, _size: u32) {}

    fn as_box(&mut self) -> *mut Box {
        let myaddr: uint = unsafe { cast::transmute(self) };
        (myaddr + AllocHeader::size()) as *mut Box
    }

    fn size() -> uint {
        // For some platforms, 16 byte alignment is required.
        let ptr_size = 16;
        let header_size = mem::size_of::<AllocHeader>();
        return (header_size + ptr_size - 1) / ptr_size * ptr_size;
    }

    fn from(a_box: *mut Box) -> *mut AllocHeader {
        (a_box as uint - AllocHeader::size()) as *mut AllocHeader
    }
}

impl MemoryRegion {
    #[inline]
    fn malloc(&mut self, size: uint) -> *mut Box {
        let total_size = size + AllocHeader::size();
        let alloc: *AllocHeader = unsafe {
            global_heap::malloc_raw(total_size) as *AllocHeader
        };

        let alloc: &mut AllocHeader = unsafe { cast::transmute(alloc) };
        alloc.init(size as u32);
        self.claim(alloc);
        self.live_allocations += 1;

        return alloc.as_box();
    }

    #[inline]
    fn realloc(&mut self, alloc: *mut Box, size: uint) -> *mut Box {
        rtassert!(!alloc.is_null());
        let orig_alloc = AllocHeader::from(alloc);
        unsafe { (*orig_alloc).assert_sane(); }

        let total_size = size + AllocHeader::size();
        let alloc: *AllocHeader = unsafe {
            global_heap::realloc_raw(orig_alloc as *mut u8,
                                     total_size) as *AllocHeader
        };

        let alloc: &mut AllocHeader = unsafe { cast::transmute(alloc) };
        alloc.assert_sane();
        alloc.update_size(size as u32);
        self.update(alloc, orig_alloc as *AllocHeader);
        return alloc.as_box();
    }

    #[inline]
    fn free(&mut self, alloc: *mut Box) {
        rtassert!(!alloc.is_null());
        let alloc = AllocHeader::from(alloc);
        unsafe {
            (*alloc).assert_sane();
            self.release(cast::transmute(alloc));
            rtassert!(self.live_allocations > 0);
            self.live_allocations -= 1;
            global_heap::exchange_free(alloc as *u8)
        }
    }

    #[cfg(rtdebug)]
    fn claim(&mut self, alloc: &mut AllocHeader) {
        alloc.assert_sane();
        if TRACK_ALLOCATIONS > 1 {
            alloc.index = self.allocations.len() as i32;
            self.allocations.push(&*alloc as *AllocHeader);
        }
    }
    #[cfg(not(rtdebug))]
    #[inline]
    fn claim(&mut self, _alloc: &mut AllocHeader) {}

    #[cfg(rtdebug)]
    fn release(&mut self, alloc: &AllocHeader) {
        alloc.assert_sane();
        if TRACK_ALLOCATIONS > 1 {
            rtassert!(self.allocations.as_slice()[alloc.index] == alloc as *AllocHeader);
            self.allocations.as_mut_slice()[alloc.index] = ptr::null();
        }
    }
    #[cfg(not(rtdebug))]
    #[inline]
    fn release(&mut self, _alloc: &AllocHeader) {}

    #[cfg(rtdebug)]
    fn update(&mut self, alloc: &mut AllocHeader, orig: *AllocHeader) {
        alloc.assert_sane();
        if TRACK_ALLOCATIONS > 1 {
            rtassert!(self.allocations.as_slice()[alloc.index] == orig);
            self.allocations.as_mut_slice()[alloc.index] = &*alloc as *AllocHeader;
        }
    }
    #[cfg(not(rtdebug))]
    #[inline]
    fn update(&mut self, _alloc: &mut AllocHeader, _orig: *AllocHeader) {}
}

impl Drop for MemoryRegion {
    fn drop(&mut self) {
        if self.live_allocations != 0 {
            rtabort!("leaked managed memory ({} objects)", self.live_allocations);
        }
        rtassert!(self.allocations.as_slice().iter().all(|s| s.is_null()));
    }
}

#[inline]
pub unsafe fn local_malloc(drop_glue: fn(*mut u8), size: uint, align: uint) -> *u8 {
    // FIXME: Unsafe borrow for speed. Lame.
    let task: Option<*mut Task> = Local::try_unsafe_borrow();
    match task {
        Some(task) => {
            (*task).heap.alloc(drop_glue, size, align) as *u8
        }
        None => rtabort!("local malloc outside of task")
    }
}

// A little compatibility function
#[inline]
pub unsafe fn local_free(ptr: *u8) {
    // FIXME: Unsafe borrow for speed. Lame.
    let task_ptr: Option<*mut Task> = Local::try_unsafe_borrow();
    match task_ptr {
        Some(task) => {
            (*task).heap.free(ptr as *mut Box)
        }
        None => rtabort!("local free outside of task")
    }
}

pub fn live_allocs() -> *mut Box {
    let mut task = Local::borrow(None::<Task>);
    task.get().heap.live_allocs
}

#[cfg(test)]
mod bench {
    use extra::test::BenchHarness;

    #[bench]
    fn alloc_managed_small(bh: &mut BenchHarness) {
        bh.iter(|| { @10; });
    }

    #[bench]
    fn alloc_managed_big(bh: &mut BenchHarness) {
        bh.iter(|| { @([10, ..1000]); });
    }
}
<|endoftext|>


<commit_before>/// this analyzer parses the Control Flow Guard function table.
/// it populates functions in the workspace.
/// the purpose of the table is to enumerate the functions that may be called indirectly.
///
/// references:
///   - https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory
///   - https://lucasg.github.io/2017/02/05/Control-Flow-Guard/
use log::{debug};
use goblin::{Object};
use failure::{Error};

use super::super::super::arch::{RVA, VA};
use super::super::super::loader::{Permissions};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

pub struct CFGuardTableAnalyzer {}

impl CFGuardTableAnalyzer {
    pub fn new() -> CFGuardTableAnalyzer {
        CFGuardTableAnalyzer {}
    }
}

const IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_MASK: u32 = 0xF0000000;
const IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT: u32 = 28;

impl Analyzer for CFGuardTableAnalyzer {
    fn get_name(&self) -> String {
        "CF Guard Table analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::CFGuardTableAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// CFGuardTableAnalyzer::new().analyze(&mut ws);
    ///
    /// // export: RtlVirtualUnwind
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x1010)).is_some());
    ///
    /// // __guard_check_icall
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x21960)).is_some());
    ///
    /// // __guard_dispatch_icall
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x21B40)).is_some());
    /// ```
    fn analyze(&self, ws: &mut Workspace) -> Result<(), Error> {
        let (load_config_directory, is_64) = {
            let pe = match Object::parse(&ws.buf) {
                Ok(Object::PE(pe)) => pe,
                _ => panic!("can't analyze unexpected format"),
            };

            let opt_header = match pe.header.optional_header {
                Some(opt_header) => opt_header,
                _ => return Ok(()),
            };

            let load_config_directory = match opt_header.data_directories.get_load_config_table() {
                Some(load_config_directory) => load_config_directory,
                _ => return Ok(()),
            };

            (RVA::from(load_config_directory.virtual_address as i64), pe.is_64)
        };

        debug!("load config directory: {}", load_config_directory);

        // offsets defined here:
        // https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory

        // according to IDA, the first DWORD is `Size` not `Characteristics` (unused).
        let size = ws.read_u32(load_config_directory)?;

        // max offset into the config directory that we'll read.
        let max_config_directory_offset = match is_64 {
            true => 0x94,   // CFG flags
            false => 0x58,  // CFG flags
        };

        // in `d3d11sdklayers.dll` for example, the config size is 0x70,
        // which is much too small to read the CFG table options.
        if max_config_directory_offset > size {
            debug!("no CF Guard table: load config directory too small");
            return Ok(())
        }

        let cfg_flags = match is_64 {
            true => ws.read_u32(load_config_directory + 0x90)?,
            false => ws.read_u32(load_config_directory + 0x58)?,
        };

        let stride = (cfg_flags & IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_MASK) >> IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT;

        let cfg_table_va = match is_64 {
            true => ws.read_va(load_config_directory + 0x80)?,
            false => ws.read_va(load_config_directory + 0x50)?,
        };
        debug!("CF Guard table: {}", cfg_table_va);

        let cfg_table_count = match is_64 {
            true => ws.read_u32(load_config_directory + 0x88)?,
            false => ws.read_u32(load_config_directory + 0x54)?,
        };
        debug!("CF Guard table count: {}", cfg_table_count);

        if cfg_table_va == VA(0x0) {
            debug!("CF guard table empty");
            return Ok(())
        };

        let cfg_table_rva = ws.rva(cfg_table_va).unwrap();
        debug!("CF guard table: {}", cfg_table_rva);

        let mut offset = cfg_table_rva;
        for _ in 0..cfg_table_count {
            let function = RVA::from(ws.read_i32(offset)?);

            debug!("CF guard function: {}", function);
            ws.make_function(function)?;
            ws.analyze()?;

            // 4 == sizeof(32-bit RVA)
            offset = offset + (4 + stride as usize);
        }

        // add function pointed to by GuardCFCheckFunctionPointer
        let guard_check_icall_fptr = match is_64 {
            true => ws.rva(ws.read_va(load_config_directory + 0x70)?).unwrap(),
            false => ws.rva(ws.read_va(load_config_directory + 0x48)?).unwrap(),
        };
        if ws.probe(guard_check_icall_fptr, 8, Permissions::R) {
            let guard_check_icall = ws.rva(ws.read_va(guard_check_icall_fptr)?).unwrap();
            if ws.probe(guard_check_icall, 1, Permissions::X) {
                debug!("CF guard check function: {:#x}", guard_check_icall);
                ws.make_function(guard_check_icall)?;
                ws.analyze()?;
            }
        };

        // add function pointed to by GuardCFDispatchFunctionPointer
        let guard_dispatch_icall_fptr = match is_64 {
            true => ws.rva(ws.read_va(load_config_directory + 120)?).unwrap(),
            false => ws.rva(ws.read_va(load_config_directory + 76)?).unwrap(),
        };
        if ws.probe(guard_dispatch_icall_fptr, 8, Permissions::R) {
            let guard_dispatch_icall = ws.rva(ws.read_va(guard_dispatch_icall_fptr)?).unwrap();
            if ws.probe(guard_dispatch_icall, 1, Permissions::X) {
                debug!("CF guard dispatch function: {:#x}", guard_dispatch_icall);
                ws.make_function(guard_dispatch_icall)?;
                ws.analyze()?;
            }
        };

        Ok(())
    }
}
<commit_msg>cfguard: ensure the GuardCFDispatchFunctionPointer is valid<commit_after>/// this analyzer parses the Control Flow Guard function table.
/// it populates functions in the workspace.
/// the purpose of the table is to enumerate the functions that may be called indirectly.
///
/// references:
///   - https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory
///   - https://lucasg.github.io/2017/02/05/Control-Flow-Guard/
use log::{debug};
use goblin::{Object};
use failure::{Error};

use super::super::super::arch::{RVA, VA};
use super::super::super::loader::{Permissions};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

pub struct CFGuardTableAnalyzer {}

impl CFGuardTableAnalyzer {
    pub fn new() -> CFGuardTableAnalyzer {
        CFGuardTableAnalyzer {}
    }
}

const IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_MASK: u32 = 0xF0000000;
const IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT: u32 = 28;

impl Analyzer for CFGuardTableAnalyzer {
    fn get_name(&self) -> String {
        "CF Guard Table analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::CFGuardTableAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// CFGuardTableAnalyzer::new().analyze(&mut ws);
    ///
    /// // export: RtlVirtualUnwind
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x1010)).is_some());
    ///
    /// // __guard_check_icall
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x21960)).is_some());
    ///
    /// // __guard_dispatch_icall
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x21B40)).is_some());
    /// ```
    fn analyze(&self, ws: &mut Workspace) -> Result<(), Error> {
        let (load_config_directory, is_64) = {
            let pe = match Object::parse(&ws.buf) {
                Ok(Object::PE(pe)) => pe,
                _ => panic!("can't analyze unexpected format"),
            };

            let opt_header = match pe.header.optional_header {
                Some(opt_header) => opt_header,
                _ => return Ok(()),
            };

            let load_config_directory = match opt_header.data_directories.get_load_config_table() {
                Some(load_config_directory) => load_config_directory,
                _ => return Ok(()),
            };

            (RVA::from(load_config_directory.virtual_address as i64), pe.is_64)
        };

        debug!("load config directory: {}", load_config_directory);

        // offsets defined here:
        // https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory

        // according to IDA, the first DWORD is `Size` not `Characteristics` (unused).
        let size = ws.read_u32(load_config_directory)?;

        // max offset into the config directory that we'll read.
        let max_config_directory_offset = match is_64 {
            true => 0x94,   // CFG flags
            false => 0x58,  // CFG flags
        };

        // in `d3d11sdklayers.dll` for example, the config size is 0x70,
        // which is much too small to read the CFG table options.
        if max_config_directory_offset > size {
            debug!("no CF Guard table: load config directory too small");
            return Ok(())
        }

        let cfg_flags = match is_64 {
            true => ws.read_u32(load_config_directory + 0x90)?,
            false => ws.read_u32(load_config_directory + 0x58)?,
        };

        let stride = (cfg_flags & IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_MASK) >> IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT;

        let cfg_table_va = match is_64 {
            true => ws.read_va(load_config_directory + 0x80)?,
            false => ws.read_va(load_config_directory + 0x50)?,
        };
        debug!("CF Guard table: {}", cfg_table_va);

        let cfg_table_count = match is_64 {
            true => ws.read_u32(load_config_directory + 0x88)?,
            false => ws.read_u32(load_config_directory + 0x54)?,
        };
        debug!("CF Guard table count: {}", cfg_table_count);

        if cfg_table_va == VA(0x0) {
            debug!("CF guard table empty");
            return Ok(())
        };

        let cfg_table_rva = ws.rva(cfg_table_va).unwrap();
        debug!("CF guard table: {}", cfg_table_rva);

        let mut offset = cfg_table_rva;
        for _ in 0..cfg_table_count {
            let function = RVA::from(ws.read_i32(offset)?);

            debug!("CF guard function: {}", function);
            ws.make_function(function)?;
            ws.analyze()?;

            // 4 == sizeof(32-bit RVA)
            offset = offset + (4 + stride as usize);
        }

        // add function pointed to by GuardCFCheckFunctionPointer
        let guard_check_icall_fptr = match is_64 {
            true => ws.rva(ws.read_va(load_config_directory + 0x70)?).unwrap(),
            false => ws.rva(ws.read_va(load_config_directory + 0x48)?).unwrap(),
        };
        if ws.probe(guard_check_icall_fptr, 8, Permissions::R) {
            let guard_check_icall = ws.rva(ws.read_va(guard_check_icall_fptr)?).unwrap();
            if ws.probe(guard_check_icall, 1, Permissions::X) {
                debug!("CF guard check function: {:#x}", guard_check_icall);
                ws.make_function(guard_check_icall)?;
                ws.analyze()?;
            }
        };

        // add function pointed to by GuardCFDispatchFunctionPointer
        //
        // set to 0x0 when not used.
        let guard_dispatch_icall_fptr = match is_64 {
            true => ws.read_va(load_config_directory + 0x78)?,
            false => ws.read_va(load_config_directory + 0x4c)?,
        };
        // first: is the field initialized?
        if guard_dispatch_icall_fptr != VA(0x0) {

            // and is the pointer valid?
            if let Some(guard_dispatch_icall_fptr) = ws.rva(guard_dispatch_icall_fptr) {

                // good.
                // now, does the pointer point to something readable in the image?
                if ws.probe(guard_dispatch_icall_fptr, 8, Permissions::R) {

                    // if so, dereference it, and this is the GuardCFDispatchFunction
                    // TODO: dangerous unwrap.
                    let guard_dispatch_icall = ws.rva(ws.read_va(guard_dispatch_icall_fptr)?).unwrap();
                    if ws.probe(guard_dispatch_icall, 1, Permissions::X) {
                        debug!("CF guard dispatch function: {:#x}", guard_dispatch_icall);
                        ws.make_function(guard_dispatch_icall)?;
                        ws.analyze()?;
                    }
                };
            }
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>use crate::prelude::*;
use std::fs;
use std::path::{Path, PathBuf};

#[cfg(unix)]
pub fn pipe() -> Result<(fs::File, fs::File), Error> {
    use fs::File;
    use std::os::unix::prelude::*;

    unsafe {
        let mut fds = [0; 2];
        if libc::pipe(&mut fds[0] as *mut _) == 0 {
            let r = File::from_raw_fd(fds[0]);
            let w = File::from_raw_fd(fds[1]);
            Ok((r, w))
        } else {
            bail!("Failed to create a pipe.");
        }
    }
}

#[cfg(windows)]
pub fn pipe() -> Result<(fs::File, fs::File), Error> {
    use fs::File;
    use std::os::windows::prelude::*;
    use std::os::windows::raw::HANDLE;
    use std::ptr;
    use winapi::um::namedpipeapi::CreatePipe;

    unsafe {
        let mut r: HANDLE = ptr::null_mut();
        let mut w: HANDLE = ptr::null_mut();
        if CreatePipe(&mut r, &mut w, ptr::null_mut(), 4096) != 0 {
            let r = File::from_raw_handle(r);
            let w = File::from_raw_handle(w);
            Ok((r, w))
        } else {
            bail!("Failed to create a pipe.");
        }
    }
}

#[cfg(unix)]
pub fn sink() -> Result<fs::File, Error> {
    Ok(fs::File::create("/dev/null")?)
}

#[cfg(windows)]
pub fn sink() -> Result<fs::File, Error> {
    Ok(fs::File::create(r#"\\.\NUL"#)?)
}

/// Returns file path with the actual casing, as stored on disk.
#[cfg(windows)]
pub fn get_fs_path_case(path: &Path) -> Result<PathBuf, std::io::Error> {
    use std::ffi::OsString;
    use std::os::windows::ffi::{OsStrExt, OsStringExt};
    use winapi::um::fileapi::GetLongPathNameW;
    let mut wpath: Vec<u16> = path.as_os_str().encode_wide().collect();
    wpath.push(0);
    let mut buffer: Vec<u16> = Vec::with_capacity(256);
    unsafe {
        let mut size = GetLongPathNameW(wpath.as_ptr(), buffer.as_mut_ptr(), buffer.capacity() as u32) as usize;
        if size == 0 {
            return Err(std::io::Error::last_os_error());
        }
        if size > buffer.capacity() {
            buffer.reserve(size - buffer.capacity());
            size = GetLongPathNameW(wpath.as_ptr(), buffer.as_mut_ptr(), buffer.capacity() as u32) as usize;
            if size == 0 {
                return Err(std::io::Error::last_os_error());
            }
            assert!(size <= buffer.capacity());
        }
        buffer.set_len(size as usize);
    }
    Ok(PathBuf::from(OsString::from_wide(&buffer)))
}

#[cfg(unix)]
pub fn get_fs_path_case(path: &Path) -> Result<PathBuf, std::io::Error> {
    Ok(path.into())
}

#[cfg(all(unix, not(test)))]
pub fn make_case_folder() -> impl FnMut(&str) -> String {
    |k| k.into()
}

#[cfg(any(windows, test))]
pub fn make_case_folder() -> impl FnMut(&str) -> String {
    use std::collections::hash_map::Entry;
    use std::collections::HashMap;

    let mut case_map: HashMap<String, String> = HashMap::new();
    move |k| {
        let uk = k.to_uppercase();
        match case_map.entry(uk) {
            Entry::Occupied(e) => e.get().into(),
            Entry::Vacant(e) => {
                e.insert(k.into());
                k.into()
            }
        }
    }
}

#[test]
fn case_folder() {
    let mut folder = make_case_folder();
    assert_eq!(folder("Path"), "Path");
    assert_eq!(folder("PATH"), "Path");

    assert_eq!(folder("Foo"), "Foo");
    assert_eq!(folder("foo"), "Foo");
}

// #[cfg(unix)]
// pub fn waitpid(pid: u32) -> Result<(), io::Error> {
//     use std::ptr;

//     unsafe {
//         if libc::waitpid(pid as libc::pid_t, ptr::null_mut(), 0) < 0 {
//             return Err(io::Error::last_os_error()).into();
//         }
//     }
//     Ok(())
// }

// #[cfg(windows)]
// pub fn waitpid(pid: u32) -> Result<(), io::Error> {
//     use winapi::um::{
//         errhandlingapi::GetLastError, handleapi::CloseHandle, processthreadsapi::OpenProcess,
//         synchapi::WaitForSingleObject, winbase::INFINITE, winnt::PROCESS_QUERY_INFORMATION,
//     };

//     unsafe {
//         let handle = OpenProcess(PROCESS_QUERY_INFORMATION, false, pid);
//         if handle == ptr::null_mut() {
//             return Err(io::Error::last_os_error()).into();
//         }
//         WaitForSingleObject(handle, INFINITE);
//         CloseHandle(handle);
//     }
//     Ok(())
// }

// #[cfg(windows)]
// fn put_env(key: &CStr, value: &CStr) {
//     use std::os::raw::{c_char, c_int};
//     extern "C" {
//         fn _putenv_s(key: *const c_char, value: *const c_char) -> c_int;
//     }
//     unsafe {
//         _putenv_s(key.as_ptr(), value.as_ptr());
//     }
// }
<commit_msg>Disable unused sink()<commit_after>use crate::prelude::*;
use std::fs;
use std::path::{Path, PathBuf};

#[cfg(unix)]
pub fn pipe() -> Result<(fs::File, fs::File), Error> {
    use fs::File;
    use std::os::unix::prelude::*;

    unsafe {
        let mut fds = [0; 2];
        if libc::pipe(&mut fds[0] as *mut _) == 0 {
            let r = File::from_raw_fd(fds[0]);
            let w = File::from_raw_fd(fds[1]);
            Ok((r, w))
        } else {
            bail!("Failed to create a pipe.");
        }
    }
}

#[cfg(windows)]
pub fn pipe() -> Result<(fs::File, fs::File), Error> {
    use fs::File;
    use std::os::windows::prelude::*;
    use std::os::windows::raw::HANDLE;
    use std::ptr;
    use winapi::um::namedpipeapi::CreatePipe;

    unsafe {
        let mut r: HANDLE = ptr::null_mut();
        let mut w: HANDLE = ptr::null_mut();
        if CreatePipe(&mut r, &mut w, ptr::null_mut(), 4096) != 0 {
            let r = File::from_raw_handle(r);
            let w = File::from_raw_handle(w);
            Ok((r, w))
        } else {
            bail!("Failed to create a pipe.");
        }
    }
}

// #[cfg(unix)]
// pub fn sink() -> Result<fs::File, Error> {
//     Ok(fs::File::create("/dev/null")?)
// }

// #[cfg(windows)]
// pub fn sink() -> Result<fs::File, Error> {
//     Ok(fs::File::create(r#"\\.\NUL"#)?)
// }

/// Returns the actual file path casing.
#[cfg(windows)]
pub fn get_fs_path_case(path: &Path) -> Result<PathBuf, std::io::Error> {
    use std::ffi::OsString;
    use std::os::windows::ffi::{OsStrExt, OsStringExt};
    use winapi::um::fileapi::GetLongPathNameW;
    let mut wpath: Vec<u16> = path.as_os_str().encode_wide().collect();
    wpath.push(0);
    let mut buffer: Vec<u16> = Vec::with_capacity(256);
    unsafe {
        let mut size = GetLongPathNameW(wpath.as_ptr(), buffer.as_mut_ptr(), buffer.capacity() as u32) as usize;
        if size == 0 {
            return Err(std::io::Error::last_os_error());
        }
        if size > buffer.capacity() {
            buffer.reserve(size - buffer.capacity());
            size = GetLongPathNameW(wpath.as_ptr(), buffer.as_mut_ptr(), buffer.capacity() as u32) as usize;
            if size == 0 {
                return Err(std::io::Error::last_os_error());
            }
            assert!(size <= buffer.capacity());
        }
        buffer.set_len(size as usize);
    }
    Ok(PathBuf::from(OsString::from_wide(&buffer)))
}

#[cfg(unix)]
pub fn get_fs_path_case(path: &Path) -> Result<PathBuf, std::io::Error> {
    Ok(path.into())
}

#[cfg(all(unix, not(test)))]
pub fn make_case_folder() -> impl FnMut(&str) -> String {
    |k| k.into()
}

#[cfg(any(windows, test))]
pub fn make_case_folder() -> impl FnMut(&str) -> String {
    use std::collections::hash_map::Entry;
    use std::collections::HashMap;

    let mut case_map: HashMap<String, String> = HashMap::new();
    move |k| {
        let uk = k.to_uppercase();
        match case_map.entry(uk) {
            Entry::Occupied(e) => e.get().into(),
            Entry::Vacant(e) => {
                e.insert(k.into());
                k.into()
            }
        }
    }
}

#[test]
fn case_folder() {
    let mut folder = make_case_folder();
    assert_eq!(folder("Path"), "Path");
    assert_eq!(folder("PATH"), "Path");

    assert_eq!(folder("Foo"), "Foo");
    assert_eq!(folder("foo"), "Foo");
}

// #[cfg(unix)]
// pub fn waitpid(pid: u32) -> Result<(), io::Error> {
//     use std::ptr;

//     unsafe {
//         if libc::waitpid(pid as libc::pid_t, ptr::null_mut(), 0) < 0 {
//             return Err(io::Error::last_os_error()).into();
//         }
//     }
//     Ok(())
// }

// #[cfg(windows)]
// pub fn waitpid(pid: u32) -> Result<(), io::Error> {
//     use winapi::um::{
//         errhandlingapi::GetLastError, handleapi::CloseHandle, processthreadsapi::OpenProcess,
//         synchapi::WaitForSingleObject, winbase::INFINITE, winnt::PROCESS_QUERY_INFORMATION,
//     };

//     unsafe {
//         let handle = OpenProcess(PROCESS_QUERY_INFORMATION, false, pid);
//         if handle == ptr::null_mut() {
//             return Err(io::Error::last_os_error()).into();
//         }
//         WaitForSingleObject(handle, INFINITE);
//         CloseHandle(handle);
//     }
//     Ok(())
// }

// #[cfg(windows)]
// fn put_env(key: &CStr, value: &CStr) {
//     use std::os::raw::{c_char, c_int};
//     extern "C" {
//         fn _putenv_s(key: *const c_char, value: *const c_char) -> c_int;
//     }
//     unsafe {
//         _putenv_s(key.as_ptr(), value.as_ptr());
//     }
// }
<|endoftext|>


<commit_before>extern mod std;
priv use std::time::{ precise_time_ns };

priv use common::problem::{ Problem };

#[path="./common/mod.rs"]
mod common;

#[path="./prob0001_0050/mod.rs"]
mod prob0001_0050;
#[path="./prob0051_0100/mod.rs"]
mod prob0051_0100;

priv static problem_sets: &'static [&'static [&'static Problem<'static>]] = &[
    prob0001_0050::problems,
    prob0051_0100::problems
];

priv fn each_problems(f: &fn(&Problem) -> bool) {
    for problem_sets.each_val |ps| {
        for ps.each_val |p| {
            if !f(p) { return; }
        }
    }
}

priv static NSEC_PER_SEC: u64 = 1000000000;
priv fn nanosec_to_str(nsec: u64) -> ~str {
    return fmt!("%3u.%09u",
         (nsec / NSEC_PER_SEC) as uint,
         (nsec % NSEC_PER_SEC) as uint);
}

priv fn solve(p: &Problem) -> u64 {
    let start_time = precise_time_ns();
    let comp_answer = (p.solver)();
    let calc_time   = precise_time_ns() - start_time;

    assert_eq!(comp_answer, p.answer.to_str());

    io::println(fmt!("%-5u %-12s %s",
                     p.id, p.answer, nanosec_to_str(calc_time)));

    return calc_time;
}

fn main() {
    let nums = os::args().filter_mapped(|&s| uint::from_str(s));

    let mut total_time = 0;
    let mut solve_cnt = 0;
    if nums.is_empty() {
        for each_problems |p| {
            total_time += solve(p);
            solve_cnt += 1;
        }
    } else {
        for nums.each_val |n| {
            for each_problems |p| {
                if p.id == n {
                    total_time += solve(p);
                    solve_cnt += 1;
                }
            }
        }
    }

    if solve_cnt > 1 {
        io::println(fmt!("TOTAL %-12s %s", "", nanosec_to_str(total_time)));
    }
}
<commit_msg>Modify output style<commit_after>extern mod std;
priv use std::time::{ precise_time_ns };

priv use common::problem::{ Problem };

#[path="./common/mod.rs"]
mod common;

#[path="./prob0001_0050/mod.rs"]
mod prob0001_0050;
#[path="./prob0051_0100/mod.rs"]
mod prob0051_0100;

priv static problem_sets: &'static [&'static [&'static Problem<'static>]] = &[
    prob0001_0050::problems,
    prob0051_0100::problems
];

priv fn each_problems(f: &fn(&Problem) -> bool) {
    for problem_sets.each_val |ps| {
        for ps.each_val |p| {
            if !f(p) { return; }
        }
    }
}

priv static NSEC_PER_SEC: u64 = 1000000000;
priv fn nanosec_to_str(nsec: u64) -> ~str {
    return fmt!("%3u.%09u",
         (nsec / NSEC_PER_SEC) as uint,
         (nsec % NSEC_PER_SEC) as uint);
}

priv fn solve(p: &Problem) -> u64 {
    let start_time = precise_time_ns();
    let comp_answer = (p.solver)();
    let calc_time   = precise_time_ns() - start_time;

    io::println(fmt!("%-5u %-12s %s",
                     p.id, p.answer, nanosec_to_str(calc_time)));

    assert_eq!(comp_answer, p.answer.to_str());

    return calc_time;
}

fn main() {
    let nums = os::args().filter_mapped(|&s| uint::from_str(s));

    let mut total_time = 0;
    let mut solve_cnt = 0;
    if nums.is_empty() {
        for each_problems |p| {
            total_time += solve(p);
            solve_cnt += 1;
        }
    } else {
        for nums.each_val |n| {
            for each_problems |p| {
                if p.id == n {
                    total_time += solve(p);
                    solve_cnt += 1;
                }
            }
        }
    }

    if solve_cnt > 1 {
        io::println(fmt!("TOTAL %-12s %s", "", nanosec_to_str(total_time)));
    }
}
<|endoftext|>


<commit_before>use crate::err::Result;
use crate::file::{SaveDir, TomlFile};
use serde::de::{self, Deserializer, Visitor};
use serde::ser::Serializer;
use serde_derive::{Deserialize, Serialize};
use std::borrow::Cow;
use std::ops::Mul;
use std::path::{Path, PathBuf};
use std::result;

#[derive(Debug, Deserialize, Serialize)]
pub struct Config {
    pub series_dir: PathBuf,
    pub reset_dates_on_rewatch: bool,
    pub episode: EpisodeConfig,
    pub tui: TuiConfig,
}

impl Config {
    pub fn new<P>(series_dir: P) -> Self
    where
        P: Into<PathBuf>,
    {
        Self {
            series_dir: series_dir.into(),
            ..Default::default()
        }
    }

    pub fn load_or_create() -> Result<Self> {
        match Self::load() {
            Ok(config) => Ok(config),
            Err(ref err) if err.is_file_nonexistant() => {
                // Fallback path is ~/anime/
                let mut dir = dirs_next::home_dir().unwrap_or_else(|| PathBuf::from("~/"));
                dir.push("anime");

                let config = Config::new(dir);
                config.save()?;

                Ok(config)
            }
            Err(err) => Err(err),
        }
    }

    pub fn stripped_path<'a, P>(&self, path: P) -> PathBuf
    where
        P: Into<Cow<'a, Path>>,
    {
        let path = path.into();

        match path.strip_prefix(&self.series_dir) {
            Ok(stripped) => stripped.into(),
            Err(_) => path.into_owned(),
        }
    }
}

impl Default for Config {
    fn default() -> Self {
        // Default series dir is ~/anime/
        let mut series_dir = dirs_next::home_dir().unwrap_or_else(|| PathBuf::from("~/"));
        series_dir.push("anime");

        Self {
            series_dir,
            reset_dates_on_rewatch: false,
            episode: EpisodeConfig::default(),
            tui: TuiConfig::default(),
        }
    }
}

impl TomlFile for Config {
    fn filename() -> &'static str {
        "config"
    }

    fn save_dir() -> SaveDir {
        SaveDir::Config
    }
}

#[derive(Debug, Deserialize, Serialize)]
pub struct EpisodeConfig {
    #[serde(rename = "percent_watched_to_progress")]
    pub pcnt_must_watch: Percentage,
    pub player: String,
    pub player_args: Vec<String>,
}

impl Default for EpisodeConfig {
    fn default() -> EpisodeConfig {
        EpisodeConfig {
            pcnt_must_watch: Percentage::new(50.0),
            player: String::from("mpv"),
            player_args: Vec::new(),
        }
    }
}

#[derive(Copy, Clone, Debug, Deserialize, Serialize)]
pub struct Percentage(#[serde(with = "Percentage")] f32);

impl Percentage {
    pub fn new(value: f32) -> Percentage {
        Percentage(value / 100.0)
    }

    fn deserialize<'de, D>(de: D) -> result::Result<f32, D::Error>
    where
        D: Deserializer<'de>,
    {
        use std::fmt;

        struct PercentageVisitor;

        impl<'de> Visitor<'de> for PercentageVisitor {
            type Value = f32;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("a positive percentage number")
            }

            fn visit_f32<E>(self, value: f32) -> result::Result<Self::Value, E>
            where
                E: de::Error,
            {
                if value.is_sign_negative() {
                    return Err(E::custom(format!(
                        "percentage must be greater than 0: {}",
                        value
                    )));
                }

                Ok(value / 100.0)
            }

            fn visit_f64<E>(self, value: f64) -> result::Result<Self::Value, E>
            where
                E: de::Error,
            {
                if value.is_sign_negative() {
                    return Err(E::custom(format!(
                        "percentage must be greater than 0: {}",
                        value
                    )));
                }

                Ok(value as f32 / 100.0)
            }
        }

        de.deserialize_f32(PercentageVisitor)
    }

    #[allow(clippy::trivially_copy_pass_by_ref)]
    fn serialize<S>(value: &f32, ser: S) -> result::Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        ser.serialize_f32(value * 100.0)
    }

    #[inline(always)]
    pub fn as_multiplier(self) -> f32 {
        self.0
    }
}

impl Mul<Percentage> for f32 {
    type Output = f32;

    fn mul(self, other: Percentage) -> Self::Output {
        self * other.as_multiplier()
    }
}

#[derive(Debug, Deserialize, Serialize)]
pub struct TuiConfig {
    pub keys: TuiKeys,
}

impl Default for TuiConfig {
    fn default() -> TuiConfig {
        TuiConfig {
            keys: TuiKeys::default(),
        }
    }
}

#[derive(Debug, Deserialize, Serialize)]
pub struct TuiKeys {
    pub play_next_episode: char,
    pub run_last_command: char,
}

impl Default for TuiKeys {
    fn default() -> TuiKeys {
        TuiKeys {
            play_next_episode: '\n',
            run_last_command: ';',
        }
    }
}
<commit_msg>anup: Use Config::default() when creating a new config file<commit_after>use crate::err::Result;
use crate::file::{SaveDir, TomlFile};
use serde::de::{self, Deserializer, Visitor};
use serde::ser::Serializer;
use serde_derive::{Deserialize, Serialize};
use std::borrow::Cow;
use std::ops::Mul;
use std::path::{Path, PathBuf};
use std::result;

#[derive(Debug, Deserialize, Serialize)]
pub struct Config {
    pub series_dir: PathBuf,
    pub reset_dates_on_rewatch: bool,
    pub episode: EpisodeConfig,
    pub tui: TuiConfig,
}

impl Config {
    pub fn new<P>(series_dir: P) -> Self
    where
        P: Into<PathBuf>,
    {
        Self {
            series_dir: series_dir.into(),
            ..Default::default()
        }
    }

    pub fn load_or_create() -> Result<Self> {
        match Self::load() {
            Ok(config) => Ok(config),
            Err(ref err) if err.is_file_nonexistant() => {
                let config = Config::default();
                config.save()?;
                Ok(config)
            }
            Err(err) => Err(err),
        }
    }

    pub fn stripped_path<'a, P>(&self, path: P) -> PathBuf
    where
        P: Into<Cow<'a, Path>>,
    {
        let path = path.into();

        match path.strip_prefix(&self.series_dir) {
            Ok(stripped) => stripped.into(),
            Err(_) => path.into_owned(),
        }
    }
}

impl Default for Config {
    fn default() -> Self {
        // Default series dir is ~/anime/
        let mut series_dir = dirs_next::home_dir().unwrap_or_else(|| PathBuf::from("~/"));
        series_dir.push("anime");

        Self {
            series_dir,
            reset_dates_on_rewatch: false,
            episode: EpisodeConfig::default(),
            tui: TuiConfig::default(),
        }
    }
}

impl TomlFile for Config {
    fn filename() -> &'static str {
        "config"
    }

    fn save_dir() -> SaveDir {
        SaveDir::Config
    }
}

#[derive(Debug, Deserialize, Serialize)]
pub struct EpisodeConfig {
    #[serde(rename = "percent_watched_to_progress")]
    pub pcnt_must_watch: Percentage,
    pub player: String,
    pub player_args: Vec<String>,
}

impl Default for EpisodeConfig {
    fn default() -> EpisodeConfig {
        EpisodeConfig {
            pcnt_must_watch: Percentage::new(50.0),
            player: String::from("mpv"),
            player_args: Vec::new(),
        }
    }
}

#[derive(Copy, Clone, Debug, Deserialize, Serialize)]
pub struct Percentage(#[serde(with = "Percentage")] f32);

impl Percentage {
    pub fn new(value: f32) -> Percentage {
        Percentage(value / 100.0)
    }

    fn deserialize<'de, D>(de: D) -> result::Result<f32, D::Error>
    where
        D: Deserializer<'de>,
    {
        use std::fmt;

        struct PercentageVisitor;

        impl<'de> Visitor<'de> for PercentageVisitor {
            type Value = f32;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("a positive percentage number")
            }

            fn visit_f32<E>(self, value: f32) -> result::Result<Self::Value, E>
            where
                E: de::Error,
            {
                if value.is_sign_negative() {
                    return Err(E::custom(format!(
                        "percentage must be greater than 0: {}",
                        value
                    )));
                }

                Ok(value / 100.0)
            }

            fn visit_f64<E>(self, value: f64) -> result::Result<Self::Value, E>
            where
                E: de::Error,
            {
                if value.is_sign_negative() {
                    return Err(E::custom(format!(
                        "percentage must be greater than 0: {}",
                        value
                    )));
                }

                Ok(value as f32 / 100.0)
            }
        }

        de.deserialize_f32(PercentageVisitor)
    }

    #[allow(clippy::trivially_copy_pass_by_ref)]
    fn serialize<S>(value: &f32, ser: S) -> result::Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        ser.serialize_f32(value * 100.0)
    }

    #[inline(always)]
    pub fn as_multiplier(self) -> f32 {
        self.0
    }
}

impl Mul<Percentage> for f32 {
    type Output = f32;

    fn mul(self, other: Percentage) -> Self::Output {
        self * other.as_multiplier()
    }
}

#[derive(Debug, Deserialize, Serialize)]
pub struct TuiConfig {
    pub keys: TuiKeys,
}

impl Default for TuiConfig {
    fn default() -> TuiConfig {
        TuiConfig {
            keys: TuiKeys::default(),
        }
    }
}

#[derive(Debug, Deserialize, Serialize)]
pub struct TuiKeys {
    pub play_next_episode: char,
    pub run_last_command: char,
}

impl Default for TuiKeys {
    fn default() -> TuiKeys {
        TuiKeys {
            play_next_episode: '\n',
            run_last_command: ';',
        }
    }
}
<|endoftext|>