<commit_before>
use syntax::ast;
use syntax::ast::{MetaItemKind,NestedMetaItemKind};
use rustc::hir::def_id::DefId;
use rustc::hir::def;
use syntax::codemap::Span;
use rustc::lint::{self, LintContext, LintPass, LateLintPass, LintArray};
use rustc::ty::{self, TyCtxt};
use rustc::hir;

declare_lint!(NOT_TAGGED_SAFE, Warn, "Warn about use of non-tagged methods within tagged function");

#[derive(Default)]
pub struct Pass
{
    visit_stack: Vec<ast::NodeId>,
}

impl LintPass for Pass {
    fn get_lints(&self) -> LintArray {
        lint_array!(NOT_TAGGED_SAFE)
    }
}

impl LateLintPass for Pass {
    fn check_fn(&mut self, cx: &lint::LateContext, _kind: hir::intravisit::FnKind, _decl: &hir::FnDecl, body: &hir::Expr, _: Span, id: ast::NodeId) {
        let attrs = cx.tcx.map.attrs(id);

        // If this function is tagged with a particular safety, store
        {
            let mut lh = ::database::CACHE.write().unwrap();
            for tag_name in get_tags(attrs, "is_safe")
            {
                let tag = lh.get_tag_or_add(&tag_name.as_str());
                lh.mark(id, tag,  true);
            }
            for tag_name in get_tags(attrs, "not_safe")
            {
                let tag = lh.get_tag_or_add(&tag_name.as_str());
                lh.mark(id, tag,  false);
            }
        }

        // For each required safety, check
        for tag_name in get_tags(attrs, "req_safe")
        {
            let ty_tag = {
                let mut lh = ::database::CACHE.write().unwrap();
                let tag = lh.get_tag_or_add(&tag_name.as_str());
                //let tag = if let Some(v) = lh.get_tag_opt(&tag_name.as_str()) { v } else { error!("Tag {} unknown", ty_name);  continue; };
                lh.mark(id, tag,  true);
                tag
                };

            // Search body for calls to non safe methods
            let mut v = Visitor {
                    pass: self, tcx: &cx.tcx, tag: ty_tag,
                    cb: |span| {
                            cx.span_lint(NOT_TAGGED_SAFE, *span,
                                &format!("Calling {0}-unsafe method from a #[req_safe({0})] method", tag_name)[..]
                                );
                        },
                    };
            debug!("Method {:?} is marked safe '{}'", id, tag_name);
            hir::intravisit::walk_expr(&mut v, body);
        }

        // TODO: For all known safeties (that aren't already set) populate.
        // - Requires knowing all safeties (which we don't, ... yet)
    }
}

impl Pass
{
    fn fill_cache_for(&mut self, tcx: &TyCtxt, node_id: ast::NodeId)
    {
        debug!("Filling cache for node {:?}", node_id);
        let attrs = tcx.map.attrs(node_id);
        let mut lh = ::database::CACHE.write().unwrap();
        for tag_name in Iterator::chain( get_tags(attrs, "is_safe"), get_tags(attrs, "req_safe") )
        {
            debug!("#[is_safe/req_safe] {} - {}", tag_name, node_id);
            let tag = lh.get_tag_or_add(&tag_name.as_str());
            lh.mark(node_id, tag,  true);
        }
        for tag_name in get_tags(attrs, "not_safe")
        {
            debug!("#[not_safe] {} - {}", tag_name, node_id);
            let tag = lh.get_tag_or_add(&tag_name.as_str());
            lh.mark(node_id, tag,  false);
        }
    }

    /// Recursively check that the provided function is either safe or unsafe.
    // Used to avoid excessive annotating
    fn recurse_fcn_body(&mut self, tcx: &TyCtxt, node_id: ast::NodeId, tag: ::database::Tag) -> bool
    {
        // and apply a visitor to all
        match tcx.map.get(node_id)
        {
        hir::map::NodeItem(i) =>
            match i.node {
            hir::ItemFn(_, _, _, _, _, ref body) => {
                // Enumerate this function's code, recursively checking for a call to an unsafe method
                let mut is_safe = true;
                {
                    let mut v = Visitor {
                        pass: self, tcx: tcx, tag: tag,
                        cb: |_| { is_safe = false; }
                        };
                    hir::intravisit::walk_expr(&mut v, tcx.map.expr(*body));
                }
                is_safe
                },
            ref v @ _ => {
                error!("Node ID {} points to a non-function item {:?}", node_id, v);
                true
                },
            },
        hir::map::NodeImplItem(i) =>
            match i.node {
            hir::ImplItemKind::Method(_, ref body) => {

                let mut is_safe = true;
                {
                    let mut v = Visitor {
                        pass: self, tcx: tcx, tag: tag,
                        cb: |_| { is_safe = false; }
                        };
                    hir::intravisit::walk_expr(&mut v, tcx.map.expr(*body));
                }
                is_safe
                },
            _ => true,
            },
        //hir::map::NodeForeignItem(i) =>
        //    if Self::check_for_marker(tcx, i.id, "tag_safe", name) {
        //        true
        //    }
        //    else if Self::check_for_marker(tcx, i.id, "tag_unsafe", name) {
        //        false
        //    }
        //    else {
        //        unknown_assume
        //    },
        ref v @ _ => {
            error!("Node ID {} points to non-item {:?}", node_id, v);
            true
            }
        }
    }

    /// Locate a #[tag_safe(<name>)] attribute on the passed item
    pub fn method_is_safe(&mut self, tcx: &TyCtxt, id: DefId, tag: ::database::Tag) -> bool
    {
        if ! id.is_local()
        {
            // TODO: Get the entry from the crate cache
            if let Some(v) = ::database::CACHE.read().unwrap().get_extern(tcx,id.krate, id.index, tag) {
                debug!("{:?} - {} (extern cached)", id, v);
                v
            }
            else {
                debug!("{:?} - {} (extern assumed)", id, true);
                true
            }
        }
        else
        {
            let node_id = tcx.map.as_local_node_id(id).unwrap();
            let mut local_opt = ::database::CACHE.read().unwrap().get_local(node_id, tag);
            // NOTE: This only fires once (ideally)
            if local_opt.is_none() {
                self.fill_cache_for(tcx, node_id);
                local_opt = ::database::CACHE.read().unwrap().get_local(node_id, tag);
            }
            if let Some(v) = local_opt {
                debug!("{} - {} (cached)", node_id, v);
                v
            }
            else {
                // If this node is currently being checked, assume it's valid.
                // TODO: This can lead to a false positive being stored.
                if self.visit_stack.iter().position(|x| *x == node_id).is_some() {
                    warn!("Recursion, assuming true");
                    true
                }
                else {
                    self.visit_stack.push(node_id);
                    let rv = self.recurse_fcn_body(tcx, node_id, tag);
                    self.visit_stack.pop();

                    debug!("{} - {} (recursed)", node_id, rv);
                    ::database::CACHE.write().unwrap().mark(node_id, tag,  rv);
                    rv
                }
            }
        }
    }
}

struct Visitor<'a, 'gcx: 'a + 'tcx, 'tcx: 'a, F: FnMut(&Span) + 'a>
{
    pass: &'a mut Pass,
    tcx: &'a TyCtxt<'a, 'gcx, 'tcx>,
    tag: ::database::Tag,
    cb: F,
}

impl<'a, 'gcx: 'tcx + 'a, 'tcx: 'a, F: FnMut(&Span)> hir::intravisit::Visitor<'a> for Visitor<'a,'gcx, 'tcx, F>
{
	fn nested_visit_map<'this>(&'this mut self) -> hir::intravisit::NestedVisitorMap<'this, 'a> {
		hir::intravisit::NestedVisitorMap::None
	}

    // Locate function/method calls in a code block
    fn visit_expr_post(&mut self, ex: &'a hir::Expr) {
        debug!("visit node - {:?}", ex);
        match ex.node
        {
        // Call expressions - check that it's a path call
        hir::ExprCall(ref fcn, _) =>
            match fcn.node
            {
            hir::ExprPath(ref qp, ..) => {
				match self.tcx.tables.borrow().qpath_def(qp, fcn.id)
				{
				def::Def::Fn(did) | def::Def::Method(did) => {
					// Check for a safety tag
					if !self.pass.method_is_safe(self.tcx, did, self.tag)
					{
						(self.cb)(&ex.span);
					}
					else {
						debug!("Safe call {:?}", ex);
					}
					},
				_ => {
					info!("Call ExprPath with an unknown Def type");
					}
				}
                },
            _ => {
                info!("Call without ExprPath");
                },
            },

        // Method call expressions - get the relevant method
        hir::ExprMethodCall(ref _id, ref _tys, ref _exprs) =>
            {
                let tables = self.tcx.tables.borrow();
                let mm = &tables.method_map;

                let callee = mm.get( &ty::MethodCall::expr(ex.id) ).unwrap();
                let id = callee.def_id;

				// Check for a safety tag
				if !self.pass.method_is_safe(self.tcx, id, self.tag) {
					(self.cb)(&ex.span);
				}
            },

        // Ignore any other type of node
        _ => {},
        }
    }
}

fn get_tags<'a>(meta_items: &'a [ast::Attribute], attr_name: &'a str) -> impl Iterator<Item=::syntax::symbol::Symbol>+'a {
    meta_items.iter()
        .filter(move |attr| attr.value.name == attr_name)
        .flat_map(|attr|
            if let MetaItemKind::List(ref v) = attr.value.node {
                v.iter()
            }
            else {
                panic!("");
            }
            )
        .filter_map(|tag_meta|
            if let NestedMetaItemKind::MetaItem(ref ptr) = tag_meta.node {
                if let MetaItemKind::Word = ptr.node {
                    Some(ptr.name)
                }
                else {
                    warn!("");
                    None
                }
            }
            else {
                warn!("");
                None
            }
            )
}
<commit_msg>Add lifetime param to LateLintPass<commit_after>
use syntax::ast;
use syntax::ast::{MetaItemKind,NestedMetaItemKind};
use rustc::hir::def_id::DefId;
use rustc::hir::def;
use syntax::codemap::Span;
use rustc::lint::{self, LintContext, LintPass, LateLintPass, LintArray};
use rustc::ty::{self, TyCtxt};
use rustc::hir;

declare_lint!(NOT_TAGGED_SAFE, Warn, "Warn about use of non-tagged methods within tagged function");

#[derive(Default)]
pub struct Pass
{
    visit_stack: Vec<ast::NodeId>,
}

impl LintPass for Pass {
    fn get_lints(&self) -> LintArray {
        lint_array!(NOT_TAGGED_SAFE)
    }
}

impl<'a,'b> LateLintPass<'a,'b> for Pass {
    fn check_fn(&mut self, cx: &lint::LateContext, _kind: hir::intravisit::FnKind, _decl: &hir::FnDecl, body: &hir::Expr, _: Span, id: ast::NodeId) {
        let attrs = cx.tcx.map.attrs(id);

        // If this function is tagged with a particular safety, store
        {
            let mut lh = ::database::CACHE.write().unwrap();
            for tag_name in get_tags(attrs, "is_safe")
            {
                let tag = lh.get_tag_or_add(&tag_name.as_str());
                lh.mark(id, tag,  true);
            }
            for tag_name in get_tags(attrs, "not_safe")
            {
                let tag = lh.get_tag_or_add(&tag_name.as_str());
                lh.mark(id, tag,  false);
            }
        }

        // For each required safety, check
        for tag_name in get_tags(attrs, "req_safe")
        {
            let ty_tag = {
                let mut lh = ::database::CACHE.write().unwrap();
                let tag = lh.get_tag_or_add(&tag_name.as_str());
                //let tag = if let Some(v) = lh.get_tag_opt(&tag_name.as_str()) { v } else { error!("Tag {} unknown", ty_name);  continue; };
                lh.mark(id, tag,  true);
                tag
                };

            // Search body for calls to non safe methods
            let mut v = Visitor {
                    pass: self, tcx: &cx.tcx, tag: ty_tag,
                    cb: |span| {
                            cx.span_lint(NOT_TAGGED_SAFE, *span,
                                &format!("Calling {0}-unsafe method from a #[req_safe({0})] method", tag_name)[..]
                                );
                        },
                    };
            debug!("Method {:?} is marked safe '{}'", id, tag_name);
            hir::intravisit::walk_expr(&mut v, body);
        }

        // TODO: For all known safeties (that aren't already set) populate.
        // - Requires knowing all safeties (which we don't, ... yet)
    }
}

impl Pass
{
    fn fill_cache_for(&mut self, tcx: &TyCtxt, node_id: ast::NodeId)
    {
        debug!("Filling cache for node {:?}", node_id);
        let attrs = tcx.map.attrs(node_id);
        let mut lh = ::database::CACHE.write().unwrap();
        for tag_name in Iterator::chain( get_tags(attrs, "is_safe"), get_tags(attrs, "req_safe") )
        {
            debug!("#[is_safe/req_safe] {} - {}", tag_name, node_id);
            let tag = lh.get_tag_or_add(&tag_name.as_str());
            lh.mark(node_id, tag,  true);
        }
        for tag_name in get_tags(attrs, "not_safe")
        {
            debug!("#[not_safe] {} - {}", tag_name, node_id);
            let tag = lh.get_tag_or_add(&tag_name.as_str());
            lh.mark(node_id, tag,  false);
        }
    }

    /// Recursively check that the provided function is either safe or unsafe.
    // Used to avoid excessive annotating
    fn recurse_fcn_body(&mut self, tcx: &TyCtxt, node_id: ast::NodeId, tag: ::database::Tag) -> bool
    {
        // and apply a visitor to all
        match tcx.map.get(node_id)
        {
        hir::map::NodeItem(i) =>
            match i.node {
            hir::ItemFn(_, _, _, _, _, ref body) => {
                // Enumerate this function's code, recursively checking for a call to an unsafe method
                let mut is_safe = true;
                {
                    let mut v = Visitor {
                        pass: self, tcx: tcx, tag: tag,
                        cb: |_| { is_safe = false; }
                        };
                    hir::intravisit::walk_expr(&mut v, tcx.map.expr(*body));
                }
                is_safe
                },
            ref v @ _ => {
                error!("Node ID {} points to a non-function item {:?}", node_id, v);
                true
                },
            },
        hir::map::NodeImplItem(i) =>
            match i.node {
            hir::ImplItemKind::Method(_, ref body) => {

                let mut is_safe = true;
                {
                    let mut v = Visitor {
                        pass: self, tcx: tcx, tag: tag,
                        cb: |_| { is_safe = false; }
                        };
                    hir::intravisit::walk_expr(&mut v, tcx.map.expr(*body));
                }
                is_safe
                },
            _ => true,
            },
        //hir::map::NodeForeignItem(i) =>
        //    if Self::check_for_marker(tcx, i.id, "tag_safe", name) {
        //        true
        //    }
        //    else if Self::check_for_marker(tcx, i.id, "tag_unsafe", name) {
        //        false
        //    }
        //    else {
        //        unknown_assume
        //    },
        ref v @ _ => {
            error!("Node ID {} points to non-item {:?}", node_id, v);
            true
            }
        }
    }

    /// Locate a #[tag_safe(<name>)] attribute on the passed item
    pub fn method_is_safe(&mut self, tcx: &TyCtxt, id: DefId, tag: ::database::Tag) -> bool
    {
        if ! id.is_local()
        {
            // TODO: Get the entry from the crate cache
            if let Some(v) = ::database::CACHE.read().unwrap().get_extern(tcx,id.krate, id.index, tag) {
                debug!("{:?} - {} (extern cached)", id, v);
                v
            }
            else {
                debug!("{:?} - {} (extern assumed)", id, true);
                true
            }
        }
        else
        {
            let node_id = tcx.map.as_local_node_id(id).unwrap();
            let mut local_opt = ::database::CACHE.read().unwrap().get_local(node_id, tag);
            // NOTE: This only fires once (ideally)
            if local_opt.is_none() {
                self.fill_cache_for(tcx, node_id);
                local_opt = ::database::CACHE.read().unwrap().get_local(node_id, tag);
            }
            if let Some(v) = local_opt {
                debug!("{} - {} (cached)", node_id, v);
                v
            }
            else {
                // If this node is currently being checked, assume it's valid.
                // TODO: This can lead to a false positive being stored.
                if self.visit_stack.iter().position(|x| *x == node_id).is_some() {
                    warn!("Recursion, assuming true");
                    true
                }
                else {
                    self.visit_stack.push(node_id);
                    let rv = self.recurse_fcn_body(tcx, node_id, tag);
                    self.visit_stack.pop();

                    debug!("{} - {} (recursed)", node_id, rv);
                    ::database::CACHE.write().unwrap().mark(node_id, tag,  rv);
                    rv
                }
            }
        }
    }
}

struct Visitor<'a, 'gcx: 'a + 'tcx, 'tcx: 'a, F: FnMut(&Span) + 'a>
{
    pass: &'a mut Pass,
    tcx: &'a TyCtxt<'a, 'gcx, 'tcx>,
    tag: ::database::Tag,
    cb: F,
}

impl<'a, 'gcx: 'tcx + 'a, 'tcx: 'a, F: FnMut(&Span)> hir::intravisit::Visitor<'a> for Visitor<'a,'gcx, 'tcx, F>
{
	fn nested_visit_map<'this>(&'this mut self) -> hir::intravisit::NestedVisitorMap<'this, 'a> {
		hir::intravisit::NestedVisitorMap::None
	}

    // Locate function/method calls in a code block
    fn visit_expr_post(&mut self, ex: &'a hir::Expr) {
        debug!("visit node - {:?}", ex);
        match ex.node
        {
        // Call expressions - check that it's a path call
        hir::ExprCall(ref fcn, _) =>
            match fcn.node
            {
            hir::ExprPath(ref qp, ..) => {
				match self.tcx.tables.borrow().qpath_def(qp, fcn.id)
				{
				def::Def::Fn(did) | def::Def::Method(did) => {
					// Check for a safety tag
					if !self.pass.method_is_safe(self.tcx, did, self.tag)
					{
						(self.cb)(&ex.span);
					}
					else {
						debug!("Safe call {:?}", ex);
					}
					},
				_ => {
					info!("Call ExprPath with an unknown Def type");
					}
				}
                },
            _ => {
                info!("Call without ExprPath");
                },
            },

        // Method call expressions - get the relevant method
        hir::ExprMethodCall(ref _id, ref _tys, ref _exprs) =>
            {
                let tables = self.tcx.tables.borrow();
                let mm = &tables.method_map;

                let callee = mm.get( &ty::MethodCall::expr(ex.id) ).unwrap();
                let id = callee.def_id;

				// Check for a safety tag
				if !self.pass.method_is_safe(self.tcx, id, self.tag) {
					(self.cb)(&ex.span);
				}
            },

        // Ignore any other type of node
        _ => {},
        }
    }
}

fn get_tags<'a>(meta_items: &'a [ast::Attribute], attr_name: &'a str) -> impl Iterator<Item=::syntax::symbol::Symbol>+'a {
    meta_items.iter()
        .filter(move |attr| attr.value.name == attr_name)
        .flat_map(|attr|
            if let MetaItemKind::List(ref v) = attr.value.node {
                v.iter()
            }
            else {
                panic!("");
            }
            )
        .filter_map(|tag_meta|
            if let NestedMetaItemKind::MetaItem(ref ptr) = tag_meta.node {
                if let MetaItemKind::Word = ptr.node {
                    Some(ptr.name)
                }
                else {
                    warn!("");
                    None
                }
            }
            else {
                warn!("");
                None
            }
            )
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::fmt;
use syntax::{Expr, ExprF, Literal};
use typed_arena::Arena;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Error<'str, 'ty> {
  Unify(&'ty Ty<'ty>, &'ty Ty<'ty>),
  Var(&'str str),
}

#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct ID(usize);

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Ty<'ty> {
  Deferred(ID),
  Func(&'ty Ty<'ty>, &'ty Ty<'ty>),
  Bool,
  Int,
  Str,
  Bytes,
  Tuple(Vec<&'ty Ty<'ty>>),
}

pub static TY_BOOL:  Ty<'static> = Ty::Bool;
pub static TY_INT:   Ty<'static> = Ty::Int;
pub static TY_STR:   Ty<'static> = Ty::Str;
pub static TY_BYTES: Ty<'static> = Ty::Bytes;

impl<'ty> Ty<'ty> {
  pub fn pretty<Purge, W>(&'ty self, purge: &Purge, into: &mut W) -> fmt::Result
    where
      Purge: Fn(&'ty Ty<'ty>) -> &'ty Ty<'ty>,
      W: fmt::Write {
    match *purge(self) {
      Ty::Deferred(ID(id)) =>
        write!(into, "t{}", id),
      Ty::Func(a, b) => {
        try!(write!(into, "("));
        try!(a.pretty(purge, into));
        try!(write!(into, " -> "));
        try!(b.pretty(purge, into));
        try!(write!(into, ")"));
        Ok(())
      },
      Ty::Bool =>
        write!(into, "bool"),
      Ty::Int =>
        write!(into, "int"),
      Ty::Str =>
        write!(into, "str"),
      Ty::Bytes =>
        write!(into, "bytes"),
      Ty::Tuple(ref elem_tys) => {
        try!(write!(into, "{{"));
        for elem_ty in elem_tys {
          try!(elem_ty.pretty(purge, into));
          try!(write!(into, ","));
        }
        try!(write!(into, "}}"));
        Ok(())
      },
    }
  }

  pub fn pretty_string<Purge>(&'ty self, purge: &Purge) -> String
    where Purge: Fn(&'ty Ty<'ty>) -> &'ty Ty<'ty> {
    let mut pretty = String::new();
    self.pretty(purge, &mut pretty).unwrap();
    pretty
  }
}

pub struct Check<'ty> {
  arena: &'ty Arena<Ty<'ty>>,
  next_id: usize,
  solved: HashMap<ID, &'ty Ty<'ty>>,
}

impl<'ty> Check<'ty> {
  pub fn new(arena: &'ty Arena<Ty<'ty>>) -> Self {
    Check{
      arena: arena,
      next_id: 0,
      solved: HashMap::new(),
    }
  }

  pub fn fresh(&mut self) -> &'ty Ty<'ty> {
    let ty = self.arena.alloc(Ty::Deferred(ID(self.next_id)));
    self.next_id += 1;
    ty
  }

  fn solve(&mut self, id: ID, ty: &'ty Ty<'ty>) {
    self.solved.insert(id, ty);
  }

  pub fn purge(&self, ty: &'ty Ty<'ty>) -> &'ty Ty<'ty> {
    match ty {
      &Ty::Deferred(id) =>
        match self.solved.get(&id) {
          Some(uy) => self.purge(uy),
          None => ty,
        },
      _ => ty,
    }
  }

  pub fn unify<'str>(&mut self, ty: &'ty Ty<'ty>, uy: &'ty Ty<'ty>)
    -> Result<(), Error<'str, 'ty>> {
    match (self.purge(ty), self.purge(uy)) {
      (&Ty::Deferred(ty_id), &Ty::Deferred(uy_id))
        if ty_id == uy_id =>
        Ok(()),
      (&Ty::Deferred(ty_id), uy_p) => {
        self.solve(ty_id, uy_p);
        Ok(())
      },
      (ty_p, &Ty::Deferred(uy_id)) => {
        self.solve(uy_id, ty_p);
        Ok(())
      },
      (&Ty::Func(ty_a, ty_b), &Ty::Func(uy_a, uy_b)) => {
        try!(self.unify(ty_a, uy_a));
        try!(self.unify(ty_b, uy_b));
        Ok(())
      },
      (&Ty::Bool, &Ty::Bool) =>
        Ok(()),
      (&Ty::Int, &Ty::Int) =>
        Ok(()),
      (&Ty::Str, &Ty::Str) =>
        Ok(()),
      (&Ty::Bytes, &Ty::Bytes) =>
        Ok(()),
      (a, b) => {
        Err(Error::Unify(a, b))
      },
    }
  }

  pub fn infer<'str, 'expr, T>(
    &mut self,
    env: &HashMap<&str, &'ty Ty<'ty>>,
    expr: &Expr<'str, 'expr, T>,
  ) -> Result<&'ty Ty<'ty>, Error<'str, 'ty>> {
    match expr.1 {
      ExprF::Lit(Literal::Bool(_)) =>
        Ok(&TY_BOOL),
      ExprF::Lit(Literal::Int(_)) =>
        Ok(&TY_INT),
      ExprF::Lit(Literal::Str(_)) =>
        Ok(&TY_STR),
      ExprF::Var(name) =>
        match env.get(&name) {
          Some(ty) => Ok(ty),
          None => Err(Error::Var(name)),
        },
      ExprF::Abs(param, body) => {
        let param_ty = self.fresh();

        let mut body_env = env.clone();
        body_env.insert(param, param_ty);
        let result_ty = try!(self.infer(&body_env, body));

        let func_ty = self.arena.alloc(Ty::Func(param_ty, result_ty));
        Ok(func_ty)
      },
      ExprF::App(callee, argument) => {
        let callee_ty = try!(self.infer(env, callee));
        let argument_ty = try!(self.infer(env, argument));

        let result_ty = self.fresh();
        let func_ty = self.arena.alloc(Ty::Func(argument_ty, result_ty));
        try!(self.unify(callee_ty, func_ty));

        Ok(result_ty)
      },
      ExprF::Tup(ref elems) => {
        let mut elem_tys = Vec::with_capacity(elems.len());
        for elem in elems {
          let elem_ty = try!(self.infer(env, elem));
          elem_tys.push(elem_ty);
        }
        let tuple_ty = self.arena.alloc(Ty::Tuple(elem_tys));
        Ok(tuple_ty)
      },
    }
  }
}

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

  #[test]
  fn test_unify() {
    let arena = Arena::new();
    let mut check = Check::new(&arena);

    let ty1 = &TY_BOOL;
    let ty2 = arena.alloc(Ty::Func(ty1, ty1));
    let ty3 = check.fresh();
    let ty4 = check.fresh();
    let ty5 = check.fresh();
    let ty6 = check.fresh();

    assert_eq!(check.unify(ty1, ty1), Ok(()));

    assert_eq!(check.unify(ty2, ty2), Ok(()));

    assert_eq!(check.unify(ty1, ty2), Err(Error::Unify(ty1, ty2)));

    assert_eq!(check.unify(ty3, ty3), Ok(()));
    assert_eq!(check.unify(ty3, ty4), Ok(()));
    assert_eq!(check.unify(ty4, ty1), Ok(()));
    assert_eq!(check.purge(ty3),      ty1);
    assert_eq!(check.unify(ty5, ty1), Ok(()));
    assert_eq!(check.unify(ty5, ty3), Ok(()));
    assert_eq!(check.unify(ty6, ty2), Ok(()));
    assert_eq!(check.unify(ty5, ty6), Err(Error::Unify(ty1, ty2)));
  }

  #[test]
  fn test_infer() {
    let expr_arena = Arena::new();
    let ty_arena = Arena::new();

    let mut check = Check::new(&ty_arena);

    let mut env = HashMap::new();
    env.insert("vari", &TY_INT);
    env.insert("even", ty_arena.alloc(Ty::Func(&TY_INT, &TY_BOOL)));

    let expr = Expr(
      (),
      ExprF::App(
        expr_arena.alloc(Expr((), ExprF::Var("even"))),
        expr_arena.alloc(Expr((), ExprF::Var("vari"))),
      ),
    );
    let ty = check.infer(&env, &expr);
    assert_eq!(ty.map(|t| check.purge(t)), Ok(&TY_BOOL));
  }
}
<commit_msg>Add unification for tuple types<commit_after>use std::collections::HashMap;
use std::fmt;
use syntax::{Expr, ExprF, Literal};
use typed_arena::Arena;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Error<'str, 'ty> {
  Unify(&'ty Ty<'ty>, &'ty Ty<'ty>),
  Var(&'str str),
}

#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct ID(usize);

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Ty<'ty> {
  Deferred(ID),
  Func(&'ty Ty<'ty>, &'ty Ty<'ty>),
  Bool,
  Int,
  Str,
  Bytes,
  Tuple(Vec<&'ty Ty<'ty>>),
}

pub static TY_BOOL:  Ty<'static> = Ty::Bool;
pub static TY_INT:   Ty<'static> = Ty::Int;
pub static TY_STR:   Ty<'static> = Ty::Str;
pub static TY_BYTES: Ty<'static> = Ty::Bytes;

impl<'ty> Ty<'ty> {
  pub fn pretty<Purge, W>(&'ty self, purge: &Purge, into: &mut W) -> fmt::Result
    where
      Purge: Fn(&'ty Ty<'ty>) -> &'ty Ty<'ty>,
      W: fmt::Write {
    match *purge(self) {
      Ty::Deferred(ID(id)) =>
        write!(into, "t{}", id),
      Ty::Func(a, b) => {
        try!(write!(into, "("));
        try!(a.pretty(purge, into));
        try!(write!(into, " -> "));
        try!(b.pretty(purge, into));
        try!(write!(into, ")"));
        Ok(())
      },
      Ty::Bool =>
        write!(into, "bool"),
      Ty::Int =>
        write!(into, "int"),
      Ty::Str =>
        write!(into, "str"),
      Ty::Bytes =>
        write!(into, "bytes"),
      Ty::Tuple(ref elem_tys) => {
        try!(write!(into, "{{"));
        for elem_ty in elem_tys {
          try!(elem_ty.pretty(purge, into));
          try!(write!(into, ","));
        }
        try!(write!(into, "}}"));
        Ok(())
      },
    }
  }

  pub fn pretty_string<Purge>(&'ty self, purge: &Purge) -> String
    where Purge: Fn(&'ty Ty<'ty>) -> &'ty Ty<'ty> {
    let mut pretty = String::new();
    self.pretty(purge, &mut pretty).unwrap();
    pretty
  }
}

pub struct Check<'ty> {
  arena: &'ty Arena<Ty<'ty>>,
  next_id: usize,
  solved: HashMap<ID, &'ty Ty<'ty>>,
}

impl<'ty> Check<'ty> {
  pub fn new(arena: &'ty Arena<Ty<'ty>>) -> Self {
    Check{
      arena: arena,
      next_id: 0,
      solved: HashMap::new(),
    }
  }

  pub fn fresh(&mut self) -> &'ty Ty<'ty> {
    let ty = self.arena.alloc(Ty::Deferred(ID(self.next_id)));
    self.next_id += 1;
    ty
  }

  fn solve(&mut self, id: ID, ty: &'ty Ty<'ty>) {
    self.solved.insert(id, ty);
  }

  pub fn purge(&self, ty: &'ty Ty<'ty>) -> &'ty Ty<'ty> {
    match ty {
      &Ty::Deferred(id) =>
        match self.solved.get(&id) {
          Some(uy) => self.purge(uy),
          None => ty,
        },
      _ => ty,
    }
  }

  pub fn unify<'str>(&mut self, ty: &'ty Ty<'ty>, uy: &'ty Ty<'ty>)
    -> Result<(), Error<'str, 'ty>> {
    match (self.purge(ty), self.purge(uy)) {
      (&Ty::Deferred(ty_id), &Ty::Deferred(uy_id))
        if ty_id == uy_id =>
        Ok(()),
      (&Ty::Deferred(ty_id), uy_p) => {
        self.solve(ty_id, uy_p);
        Ok(())
      },
      (ty_p, &Ty::Deferred(uy_id)) => {
        self.solve(uy_id, ty_p);
        Ok(())
      },
      (&Ty::Func(ty_a, ty_b), &Ty::Func(uy_a, uy_b)) => {
        try!(self.unify(ty_a, uy_a));
        try!(self.unify(ty_b, uy_b));
        Ok(())
      },
      (&Ty::Bool, &Ty::Bool) =>
        Ok(()),
      (&Ty::Int, &Ty::Int) =>
        Ok(()),
      (&Ty::Str, &Ty::Str) =>
        Ok(()),
      (&Ty::Bytes, &Ty::Bytes) =>
        Ok(()),
      (&Ty::Tuple(ref elem_tys), &Ty::Tuple(ref elem_uys)) => {
        if elem_tys.len() != elem_uys.len() {
          Err(Error::Unify(&TY_BOOL, &TY_INT)) // FIXME: Tuple types.
        } else {
          for (ty, uy) in elem_tys.iter().zip(elem_uys.iter()) {
            try!(self.unify(ty, uy));
          }
          Ok(())
        }
      },
      (a, b) => {
        Err(Error::Unify(a, b))
      },
    }
  }

  pub fn infer<'str, 'expr, T>(
    &mut self,
    env: &HashMap<&str, &'ty Ty<'ty>>,
    expr: &Expr<'str, 'expr, T>,
  ) -> Result<&'ty Ty<'ty>, Error<'str, 'ty>> {
    match expr.1 {
      ExprF::Lit(Literal::Bool(_)) =>
        Ok(&TY_BOOL),
      ExprF::Lit(Literal::Int(_)) =>
        Ok(&TY_INT),
      ExprF::Lit(Literal::Str(_)) =>
        Ok(&TY_STR),
      ExprF::Var(name) =>
        match env.get(&name) {
          Some(ty) => Ok(ty),
          None => Err(Error::Var(name)),
        },
      ExprF::Abs(param, body) => {
        let param_ty = self.fresh();

        let mut body_env = env.clone();
        body_env.insert(param, param_ty);
        let result_ty = try!(self.infer(&body_env, body));

        let func_ty = self.arena.alloc(Ty::Func(param_ty, result_ty));
        Ok(func_ty)
      },
      ExprF::App(callee, argument) => {
        let callee_ty = try!(self.infer(env, callee));
        let argument_ty = try!(self.infer(env, argument));

        let result_ty = self.fresh();
        let func_ty = self.arena.alloc(Ty::Func(argument_ty, result_ty));
        try!(self.unify(callee_ty, func_ty));

        Ok(result_ty)
      },
      ExprF::Tup(ref elems) => {
        let mut elem_tys = Vec::with_capacity(elems.len());
        for elem in elems {
          let elem_ty = try!(self.infer(env, elem));
          elem_tys.push(elem_ty);
        }
        let tuple_ty = self.arena.alloc(Ty::Tuple(elem_tys));
        Ok(tuple_ty)
      },
    }
  }
}

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

  #[test]
  fn test_unify() {
    let arena = Arena::new();
    let mut check = Check::new(&arena);

    let ty1 = &TY_BOOL;
    let ty2 = arena.alloc(Ty::Func(ty1, ty1));
    let ty3 = check.fresh();
    let ty4 = check.fresh();
    let ty5 = check.fresh();
    let ty6 = check.fresh();

    assert_eq!(check.unify(ty1, ty1), Ok(()));

    assert_eq!(check.unify(ty2, ty2), Ok(()));

    assert_eq!(check.unify(ty1, ty2), Err(Error::Unify(ty1, ty2)));

    assert_eq!(check.unify(ty3, ty3), Ok(()));
    assert_eq!(check.unify(ty3, ty4), Ok(()));
    assert_eq!(check.unify(ty4, ty1), Ok(()));
    assert_eq!(check.purge(ty3),      ty1);
    assert_eq!(check.unify(ty5, ty1), Ok(()));
    assert_eq!(check.unify(ty5, ty3), Ok(()));
    assert_eq!(check.unify(ty6, ty2), Ok(()));
    assert_eq!(check.unify(ty5, ty6), Err(Error::Unify(ty1, ty2)));
  }

  #[test]
  fn test_infer() {
    let expr_arena = Arena::new();
    let ty_arena = Arena::new();

    let mut check = Check::new(&ty_arena);

    let mut env = HashMap::new();
    env.insert("vari", &TY_INT);
    env.insert("even", ty_arena.alloc(Ty::Func(&TY_INT, &TY_BOOL)));

    let expr = Expr(
      (),
      ExprF::App(
        expr_arena.alloc(Expr((), ExprF::Var("even"))),
        expr_arena.alloc(Expr((), ExprF::Var("vari"))),
      ),
    );
    let ty = check.infer(&env, &expr);
    assert_eq!(ty.map(|t| check.purge(t)), Ok(&TY_BOOL));
  }
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with
// this file, You can obtain one at https://mozilla.org/MPL/2.0/.

use common::is_whitespace;

use regex::Regex;
use std::borrow::Cow;

/// Removes unnecessary whitespaces from a String.
///
/// # Example
///
/// ```
/// use crowbook_text_processing::clean::remove_whitespaces;
/// let s = remove_whitespaces("  A  string   with   more   whitespaces  than  needed   ");
/// assert_eq!(&s, " A string with more whitespaces than needed ");
/// ```
pub fn remove_whitespaces<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new(r"[  \x{202F}\x{2002}]{2,}?").unwrap();
    }
    let input = input.into();
    let first = REGEX.find(&input);
    if let Some((first, _)) = first {
        let mut new_s = String::with_capacity(input.len());
        new_s.push_str(&input[0..first]);
        let mut previous_space = false;
        for c in input[first..].chars() {
            if is_whitespace(c) {
                if previous_space {
                    // previous char already a space, don't copy it
                } else {
                    new_s.push(c);
                    previous_space = true;
                }
            } else {
                previous_space = false;
                new_s.push(c);
            }
        }
        Cow::Owned(new_s)
    } else {
        input
    }
}

/// Class of a character
#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Copy)]
enum CharClass {
    Whitespace = 0,
    Punctuation,
    Alphanumeric
}

/// Get class of a character
fn char_class(c: char) -> CharClass {
    if c.is_alphanumeric() {
        CharClass::Alphanumeric
    } else if c.is_whitespace() {
        CharClass::Whitespace
    } else {
        CharClass::Punctuation
    }
}

/// Replace quotes with more typographic variants
///
/// While it should work pretty well for double quotes (`"`), the rules for single
/// quote (`'`) are more ambiguous, as it can be a quote, or an apostrophe and it's not
/// that easy to get right.
///
/// # Example
///
/// ```
/// use crowbook_text_processing::typographic_quotes;
/// let s = typographic_quotes("\"foo\"");
/// assert_eq!(&s, "“foo”");
/// let s = typographic_quotes("'foo'");
/// assert_eq!(&s, "‘foo’");
/// ```
pub fn typographic_quotes<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    /// Custom whitespace-detection function, including `,`, `;`, `.`, `!`, `?`, and `:`
    fn is_whitespace(c: char) -> bool {
        match c {
            ',' | '.' | ';' | '!' | '?' | ':' | '"'
                | '(' | ')' => true,
            _ => c.is_whitespace()
        }
    }

    lazy_static! {
        static ref REGEX: Regex = Regex::new("[\"\']").unwrap();
    }
    let input = input.into();
    let first = REGEX.find(&input);
    if let Some((mut first, _)) = first {
        let mut new_s = String::with_capacity(input.len());
        if first > 0 {
            // Move one step backward since we might need to know if previous char was
            // a letter or not
            first -= 1;
        }
        new_s.push_str(&input[0..first]);
        let chars = input[first..].chars().collect::<Vec<_>>();
        let mut closing_quote = None;
        let mut opened_doubles = 0;
        for i in 0..chars.len() {
            let c = chars[i];
            let has_opened_quote = if let Some(n) = closing_quote {
                i <= n
            } else {
                false
            };
            match c {
                '"' => {
                    let prev = if i > 0 {
                        char_class(chars[i-1])
                    } else {
                        CharClass::Whitespace
                    };
                    let next = if i < chars.len() - 1 {
                        char_class(chars[i+1])
                    } else {
                        CharClass::Whitespace
                    };

                    if prev < next {
                        opened_doubles += 1;
                        new_s.push('“');
                    } else {
                        if opened_doubles > 0 {
                            opened_doubles -= 1;
                            new_s.push('”');
                        } else {
                            new_s.push('"');
                        }
                    }
                },
                '\'' => {
                    let prev = if i > 0 {
                        Some(!chars[i - 1].is_whitespace())
                    } else {
                        None
                    };
                    let next = if i < chars.len() - 1 {
                        Some(!chars[i+1].is_whitespace())
                    } else {
                        None
                    };
                    let replacement = match (prev, next) {
                        // Elision, it's closing
                        (Some(true), Some(true)) => '’',

                        // Beginning of word, it's opening (not always though)
                        (Some(false), Some(true))
                            | (None, Some(true))
                            => {
                                let mut is_next_closing = false;
                                for j in (i + 1)..chars.len() {
                                    if chars[j] == '\'' {
                                        if chars[j-1].is_whitespace() {
                                            continue;
                                        } else {
                                            if j >= chars.len() - 1
                                                || is_whitespace(chars[j+1]) {
                                                    is_next_closing = true;
                                                    closing_quote = Some(j);
                                                    break;
                                                }
                                        }
                                    }
                                }
                                if is_next_closing && !has_opened_quote {
                                    '‘'
                                } else {
                                    '’'
                                }
                            }

                        // Apostrophe at end of word, it's closing
                        (Some(true), Some(false))
                            | (Some(true), None)
                            => {
                                '’'
                            },

                        _ => '\'',
                    };
                    new_s.push(replacement);
                    // if i > 0 && !chars[i - 1].is_whitespace() {
                    //     new_s.push('’');
                    // } else if i < chars.len() - 1 && !chars[i + 1].is_whitespace() {
                    //     new_s.push('‘');
                    // } else {
                    //     new_s.push('\'');
                    // }
                },
                _ => new_s.push(c)
            }
        }
        Cow::Owned(new_s)
    } else {
        input
    }
}

#[test]
fn remove_whitespaces_1() {
    let s = "   Remove    supplementary   spaces    but    don't    trim     either   ";
    let res = remove_whitespaces(s);
    assert_eq!(&res, " Remove supplementary spaces but don't trim either ");
}

#[test]
fn typographic_quotes_1() {
    let s = "Some string without ' typographic ' quotes";
    let res = typographic_quotes(s);
    assert_eq!(&res, s);
}

#[test]
fn typographic_quotes_2() {
    let s = typographic_quotes("\"foo\"");
    assert_eq!(&s, "“foo”");
    let s = typographic_quotes("'foo'");
    assert_eq!(&s, "‘foo’");
}

#[test]
fn typographic_quotes_3() {
    let s = typographic_quotes("\'mam, how are you?");
    assert_eq!(&s, "’mam, how are you?");
}

#[test]
fn typographic_quotes_4() {
    let s = typographic_quotes("some char: 'c', '4', '&'");
    assert_eq!(&s, "some char: ‘c’, ‘4’, ‘&’");
}

#[test]
fn typographic_quotes_5() {
    let s = typographic_quotes("It's a good day to say 'hi'");
    assert_eq!(&s, "It’s a good day to say ‘hi’");
}

#[test]
fn typographic_quotes_6() {
    let s = typographic_quotes("The '60s were nice, weren't they?");
    assert_eq!(&s, "The ’60s were nice, weren’t they?");
}

#[test]
fn typographic_quotes_7() {
    let s = typographic_quotes("Plurals' possessive");
    assert_eq!(&s, "Plurals’ possessive");
}

#[test]
fn typographic_quotes_8() {
    let s = typographic_quotes("\"I like 'That '70s show'\", she said");
    assert_eq!(&s, "“I like ‘That ’70s show’”, she said");
}

#[test]
fn typographic_quotes_9() {
    let s = typographic_quotes("some char: '!', '?', ','");
    assert_eq!(&s, "some char: ‘!’, ‘?’, ‘,’");
}
<commit_msg>Enhance typographic_quotes function<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with
// this file, You can obtain one at https://mozilla.org/MPL/2.0/.

use common::is_whitespace;

use regex::Regex;
use std::borrow::Cow;

/// Removes unnecessary whitespaces from a String.
///
/// # Example
///
/// ```
/// use crowbook_text_processing::clean::remove_whitespaces;
/// let s = remove_whitespaces("  A  string   with   more   whitespaces  than  needed   ");
/// assert_eq!(&s, " A string with more whitespaces than needed ");
/// ```
pub fn remove_whitespaces<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new(r"[  \x{202F}\x{2002}]{2,}?").unwrap();
    }
    let input = input.into();
    let first = REGEX.find(&input);
    if let Some((first, _)) = first {
        let mut new_s = String::with_capacity(input.len());
        new_s.push_str(&input[0..first]);
        let mut previous_space = false;
        for c in input[first..].chars() {
            if is_whitespace(c) {
                if previous_space {
                    // previous char already a space, don't copy it
                } else {
                    new_s.push(c);
                    previous_space = true;
                }
            } else {
                previous_space = false;
                new_s.push(c);
            }
        }
        Cow::Owned(new_s)
    } else {
        input
    }
}

/// Class of a character
#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Copy)]
enum CharClass {
    Whitespace = 0,
    Punctuation,
    Alphanumeric
}

/// Get class of a character
fn char_class(c: char) -> CharClass {
    if c.is_alphanumeric() {
        CharClass::Alphanumeric
    } else if c.is_whitespace() {
        CharClass::Whitespace
    } else {
        CharClass::Punctuation
    }
}

/// Replace quotes with more typographic variants
///
/// While it should work pretty well for double quotes (`"`), the rules for single
/// quote (`'`) are more ambiguous, as it can be a quote, or an apostrophe and it's not
/// that easy to get right.
///
/// # Example
///
/// ```
/// use crowbook_text_processing::typographic_quotes;
/// let s = typographic_quotes("\"foo\"");
/// assert_eq!(&s, "“foo”");
/// let s = typographic_quotes("'foo'");
/// assert_eq!(&s, "‘foo’");
/// ```
pub fn typographic_quotes<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    /// Custom whitespace-detection function, including `,`, `;`, `.`, `!`, `?`, and `:`
    fn is_whitespace(c: char) -> bool {
        match c {
            ',' | '.' | ';' | '!' | '?' | ':' | '"'
                | '(' | ')' => true,
            _ => c.is_whitespace()
        }
    }

    lazy_static! {
        static ref REGEX: Regex = Regex::new("[\"\']").unwrap();
    }
    let input = input.into();
    let first = REGEX.find(&input);
    if let Some((mut first, _)) = first {
        let mut new_s = String::with_capacity(input.len());
        if first > 0 {
            // Move one step backward since we might need to know if previous char was
            // a letter or not
            first -= 1;
        }
        new_s.push_str(&input[0..first]);
        let mut chars = input[first..].chars().collect::<Vec<_>>();
        let mut closing_quote = None;
        let mut opened_doubles = 0;
        for i in 0..chars.len() {
            let c = chars[i];
            let has_opened_quote = if let Some(n) = closing_quote {
                i <= n
            } else {
                false
            };
            match c {
                '"' => {
                    let prev = if i > 0 {
                        char_class(chars[i-1])
                    } else {
                        CharClass::Whitespace
                    };
                    let next = if i < chars.len() - 1 {
                        char_class(chars[i+1])
                    } else {
                        CharClass::Whitespace
                    };

                    if prev < next {
                        opened_doubles += 1;
                        new_s.push('“');
                    } else {
                        if opened_doubles > 0 {
                            opened_doubles -= 1;
                            new_s.push('”');
                        } else {
                            new_s.push('"');
                        }
                    }
                },
                '\'' => {
                    let prev = if i > 0 {
                        char_class(chars[i - 1])
                    } else {
                        CharClass::Whitespace
                    };
                    let next = if i < chars.len() - 1 {
                        char_class(chars[i + 1])
                    } else {
                        CharClass::Whitespace
                    };

                    let replacement = match (prev, next) {
                        // Elision or possessive
                        (CharClass::Alphanumeric, CharClass::Alphanumeric)
//                            | (CharClass::Punctuation, CharClass::Alphanumeric)
                            => '’',

                        // Beginning of word, it's opening (not always though)
                        (x, y) if x < y
                            => {
                                let mut is_next_closing = false;
                                for j in (i + 1)..chars.len() {
                                    if chars[j] == '\'' {
                                        if chars[j-1].is_whitespace() {
                                            continue;
                                        } else {
                                            if j >= chars.len() - 1
                                                || char_class(chars[j+1]) != CharClass::Alphanumeric {
                                                    is_next_closing = true;
                                                    closing_quote = Some(j);
                                                    chars[j] = '’';
                                                    break;
                                                }
                                        }
                                    }
                                }
                                if is_next_closing && !has_opened_quote {
                                    '‘'
                                } else {
                                    '’'
                                }
                            }

                        // Apostrophe at end of word, it's closing
                        (x, y) if x > y
                            => {
                                '’'
                            },
                        _ => '\'',
                    };
                    new_s.push(replacement);
                    // if i > 0 && !chars[i - 1].is_whitespace() {
                    //     new_s.push('’');
                    // } else if i < chars.len() - 1 && !chars[i + 1].is_whitespace() {
                    //     new_s.push('‘');
                    // } else {
                    //     new_s.push('\'');
                    // }
                },
                _ => new_s.push(c)
            }
        }
        Cow::Owned(new_s)
    } else {
        input
    }
}

#[test]
fn remove_whitespaces_1() {
    let s = "   Remove    supplementary   spaces    but    don't    trim     either   ";
    let res = remove_whitespaces(s);
    assert_eq!(&res, " Remove supplementary spaces but don't trim either ");
}

#[test]
fn typographic_quotes_1() {
    let s = "Some string without ' typographic ' quotes";
    let res = typographic_quotes(s);
    assert_eq!(&res, s);
}

#[test]
fn typographic_quotes_2() {
    let s = typographic_quotes("\"foo\"");
    assert_eq!(&s, "“foo”");
    let s = typographic_quotes("'foo'");
    assert_eq!(&s, "‘foo’");
}

#[test]
fn typographic_quotes_3() {
    let s = typographic_quotes("\'mam, how are you?");
    assert_eq!(&s, "’mam, how are you?");
}

#[test]
fn typographic_quotes_4() {
    let s = typographic_quotes("some char: 'c', '4', '&'");
    assert_eq!(&s, "some char: ‘c’, ‘4’, ‘&’");
}

#[test]
fn typographic_quotes_5() {
    let s = typographic_quotes("It's a good day to say 'hi'");
    assert_eq!(&s, "It’s a good day to say ‘hi’");
}

#[test]
fn typographic_quotes_6() {
    let s = typographic_quotes("The '60s were nice, weren't they?");
    assert_eq!(&s, "The ’60s were nice, weren’t they?");
}

#[test]
fn typographic_quotes_7() {
    let s = typographic_quotes("Plurals' possessive");
    assert_eq!(&s, "Plurals’ possessive");
}

#[test]
fn typographic_quotes_8() {
    let s = typographic_quotes("\"I like 'That '70s show'\", she said");
    assert_eq!(&s, "“I like ‘That ’70s show’”, she said");
}

#[test]
fn typographic_quotes_9() {
    let s = typographic_quotes("some char: '!', '?', ','");
    assert_eq!(&s, "some char: ‘!’, ‘?’, ‘,’");
}

#[test]
fn typographic_quotes_10() {
    let s = typographic_quotes("\"'Let's try \"nested\" quotes,' he said.\"");
    assert_eq!(&s, "“‘Let’s try “nested” quotes,’ he said.”");
}
<|endoftext|>


<commit_before>//! Colors used for drawing to the screen
//! TODO This needs to be in a helper library

use std::convert::From;

/// Color to draw to the screen, including the alpha channel.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct Color {
    red: u8,
    green: u8,
    blue: u8,
    alpha: u8
}

impl Color {
    /// Makes a new solid color, with no transparency.
    pub fn solid_color(red: u8, green: u8, blue: u8) -> Self {
        Color {
            red,
            green,
            blue,
            alpha: 255
        }
    }

    /// Gets the values of the colors, in this order:
    /// (Red, Green, Blue, Alpha)
    #[allow(dead_code)]
    pub fn to_u8s(self) -> (u8, u8, u8, u8) {
        (self.red, self.green, self.blue, self.alpha)
    }

    pub fn from_u8s(red: u8, green: u8, blue: u8) -> Self {
        Color::solid_color(red, green, blue)
    }

    /// To a u32, represented as in RGBA format
    pub fn to_u32(self) -> u32 {
        let values = [self.red, self.green, self.blue, self.alpha];
        unsafe {
            ::std::mem::transmute(values)
        }
    }
}

impl From<u32> for Color {
    fn from(val: u32) -> Self {
        let blue = ((val & 0xff0000) >> 16) as u8;
        let green = ((val & 0x00ff00) >> 8) as u8;
        let red = (val & 0x0000ff) as u8;
        Color::solid_color(red, green, blue)
    }
}
<commit_msg>Fixed warning<commit_after>//! Colors used for drawing to the screen
//! TODO This needs to be in a helper library

use std::convert::From;

/// Color to draw to the screen, including the alpha channel.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct Color {
    red: u8,
    green: u8,
    blue: u8,
    alpha: u8
}

impl Color {
    /// Makes a new solid color, with no transparency.
    pub fn solid_color(red: u8, green: u8, blue: u8) -> Self {
        Color {
            red,
            green,
            blue,
            alpha: 255
        }
    }

    /// Gets the values of the colors, in this order:
    /// (Red, Green, Blue, Alpha)
    #[allow(dead_code)]
    pub fn to_u8s(self) -> (u8, u8, u8, u8) {
        (self.red, self.green, self.blue, self.alpha)
    }

    /// To a u32, represented as in RGBA format
    pub fn to_u32(self) -> u32 {
        let values = [self.red, self.green, self.blue, self.alpha];
        unsafe {
            ::std::mem::transmute(values)
        }
    }
}

impl From<u32> for Color {
    fn from(val: u32) -> Self {
        let blue = ((val & 0xff0000) >> 16) as u8;
        let green = ((val & 0x00ff00) >> 8) as u8;
        let red = (val & 0x0000ff) as u8;
        Color::solid_color(red, green, blue)
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::rc::Rc;
use std::hash;
use digit_slice::{DigitSlice, FromDigits};

/// A wrapper for referencing Noun-like patterns.
pub enum Shape<A, N> {
    Atom(A),
    Cell(N, N),
}

/// A Nock noun, the basic unit of representation.
///
/// A noun is an atom or a cell. An atom is any natural number. A cell is any
/// ordered pair of nouns.
///
/// Atoms are represented by a little-endian byte array of 8-bit digits.
#[derive(Clone, PartialEq, Eq)]
pub struct Noun(Inner);

#[derive(Clone, PartialEq, Eq)]
enum Inner {
    Atom(Rc<Vec<u8>>),
    Cell(Rc<Noun>, Rc<Noun>),
}

impl Noun {
    fn get<'a>(&'a self) -> Shape<&'a [u8], &'a Noun> {
        match self.0 {
            Inner::Atom(ref v) => Shape::Atom(&v),
            Inner::Cell(ref a, ref b) => Shape::Cell(&*a, &*b),
        }
    }

    /// Memory address or other unique identifier for the noun.
    fn addr(&self) -> usize {
        &*self as *const _ as usize
    }

    /// Build a new atom noun from a little-endian 8-bit digit sequence.
    pub fn atom(digits: &[u8]) -> Noun {
        Noun(Inner::Atom(Rc::new(digits.to_vec())))
    }

    /// Build a new cell noun from two existing nouns.
    pub fn cell(a: Noun, b: Noun) -> Noun {
        Noun(Inner::Cell(Rc::new(a), Rc::new(b)))
    }

    /// Build a noun from a convertible value.
    pub fn from<T: ToNoun>(item: T) -> Noun {
        item.to_noun()
    }

    /// Run a memoizing fold over the noun
    fn fold<'a, F, T>(&'a self, mut f: F) -> T
        where F: FnMut(Shape<&'a [u8], T>) -> T,
              T: Clone
    {
        fn h<'a, F, T>(noun: &'a Noun,
                       memo: &mut HashMap<usize, T>,
                       f: &mut F)
                       -> T
            where F: FnMut(Shape<&'a [u8], T>) -> T,
                  T: Clone
        {
            let key = noun.addr();

            if memo.contains_key(&key) {
                memo.get(&key).unwrap().clone()
            } else {
                let ret = match noun.get() {
                    Shape::Atom(x) => f(Shape::Atom(x)),
                    Shape::Cell(ref a, ref b) => {
                        let a = h(*a, memo, f);
                        let b = h(*b, memo, f);
                        let ret = f(Shape::Cell(a, b));
                        ret
                    }
                };
                memo.insert(key, ret.clone());
                ret
            }
        }

        h(self, &mut HashMap::new(), &mut f)
    }

}

impl hash::Hash for Noun {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        fn f<H: hash::Hasher>(state: &mut H, shape: Shape<&[u8], u64>) -> u64 {
            match shape {
                Shape::Atom(x) => x.hash(state),
                Shape::Cell(a, b) => {
                    a.hash(state);
                    b.hash(state);
                }
            }
            state.finish()
        }
        self.fold(|x| f(state, x))
            .hash(state);
    }
}

/// Trait for types that can convert themselves to a noun.
pub trait ToNoun {
    fn to_noun(&self) -> Noun;
}

impl<T> ToNoun for T where T: DigitSlice
{
    fn to_noun(&self) -> Noun {
        Noun::atom(self.as_digits())
    }
}

/// A trait for types that can be instantiated from a Nock noun.
pub trait FromNoun: Sized {
    /// The associated error.
    type Err;

    /// Try to convert a noun to an instance of the type.
    fn from_noun(n: &Noun) -> Result<Self, Self::Err>;
}

impl<T> FromNoun for T where T: FromDigits {
    type Err = ();

    fn from_noun(n: &Noun) -> Result<Self, Self::Err> {
        match n.get() {
            Shape::Atom(x) => T::from_digits(x).map_err(|_| ()),
            _ => Err(())
        }
    }
}

impl<T, U> FromNoun for (T, U) where T: FromNoun, U: FromNoun {
    type Err = ();

    fn from_noun(n: &Noun) -> Result<Self, Self::Err> {
        match n.get() {
            Shape::Cell(a, b) => {
                let t = try!(T::from_noun(a).map_err(|_| ()));
                let u = try!(U::from_noun(b).map_err(|_| ()));
                Ok((t, u))
            }
            _ => Err(())
        }
    }
}

// TODO: FromNoun for T: FromIterator<U: FromNoun>. Pair impl should give us
// a HashMap derivation then. Use ~-terminated cell sequence as backend.

// TODO: Turn a ~-terminated noun into a vec or an iter. Can fail if the last
// element isn't a ~, and we'll only know when we hit the last element...
// Return type is Option<Vec<&'a Noun>>?

// TODO: ToNoun for T: IntoIterator<U: ToNoun>.

// TODO: FromNoun/ToNoun for String, compatible with cord datatype.

// TODO: FromNoun/ToNoun for signed numbers

// For the n! macro to work.
impl Into<Noun> for u64 {
    fn into(self) -> Noun {
        Noun::from(self)
    }
}

/// Macro for noun literals.
///
/// Rust n![1, 2, 3] corresponds to Nock [1 2 3]
#[macro_export]
macro_rules! n {
    [$x:expr, $y:expr] => { $crate::draft::Noun::cell($x.into(), $y.into()) };
    [$x:expr, $y:expr, $($ys:expr),+] => { $crate::draft::Noun::cell($x.into(), n![$y, $($ys),+]) };
}

#[cfg(test)]
mod tests {
    use std::hash;
    use super::Noun;

    #[test]
    fn scratch() {
        let x = Noun::from(123u32);
        assert!(x == Noun::from(123u8));
    }

    fn hash<T: hash::Hash>(t: &T) -> u64 {
        use std::hash::Hasher;
        let mut s = hash::SipHasher::new();
        t.hash(&mut s);
        s.finish()
    }

    #[test]
    fn test_fold() {
        assert_eq!(hash(&n![1, 2, 3]), hash(&n![1, 2, 3]));
        assert!(hash(&n![n![1, 2], 3]) != hash(&n![1, 2, 3]));
        assert!(hash(&n![1, 2, 3]) != hash(&n![1, 2]));
    }
}
<commit_msg>Simple pattern-matching for nouns<commit_after>use std::collections::HashMap;
use std::rc::Rc;
use std::hash;
use digit_slice::{DigitSlice, FromDigits};

/// A wrapper for referencing Noun-like patterns.
#[derive(Copy, Clone)]
pub enum Shape<A, N> {
    Atom(A),
    Cell(N, N),
}

/// A Nock noun, the basic unit of representation.
///
/// A noun is an atom or a cell. An atom is any natural number. A cell is any
/// ordered pair of nouns.
///
/// Atoms are represented by a little-endian byte array of 8-bit digits.
#[derive(Clone, PartialEq, Eq)]
pub struct Noun(Inner);

#[derive(Clone, PartialEq, Eq)]
enum Inner {
    Atom(Rc<Vec<u8>>),
    Cell(Rc<Noun>, Rc<Noun>),
}

pub type NounShape<'a> = Shape<&'a [u8], &'a Noun>;

impl Noun {
    fn get<'a>(&'a self) -> NounShape<'a> {
        match self.0 {
            Inner::Atom(ref v) => Shape::Atom(&v),
            Inner::Cell(ref a, ref b) => Shape::Cell(&*a, &*b),
        }
    }

    /// Pattern-match a noun with shape [p q r].
    ///
    /// The digit sequence shows the branch length of each leaf node in the
    /// expression being matched. 122 has the leftmost leaf 1 step away from
    /// the root and the two leaves on the right both 2 steps away from the
    /// root.
    pub fn get_122<'a>(&'a self) -> Option<(NounShape<'a>, NounShape<'a>, NounShape<'a>)> {
        if let Shape::Cell(ref a, ref b) = self.get() {
            if let Shape::Cell(ref b, ref c) = b.get() {
                return Some((a.get(), b.get(), c.get()));
            }
        }
        None
    }

    /// Pattern-match a noun with shape [[p q] r].
    pub fn get_221<'a>(&'a self) -> Option<(NounShape<'a>, NounShape<'a>, NounShape<'a>)> {
        if let Shape::Cell(ref a, ref c) = self.get() {
            if let Shape::Cell(ref a, ref b) = a.get() {
                return Some((a.get(), b.get(), c.get()));
            }
        }
        None
    }

    /// Memory address or other unique identifier for the noun.
    fn addr(&self) -> usize {
        &*self as *const _ as usize
    }

    /// Build a new atom noun from a little-endian 8-bit digit sequence.
    pub fn atom(digits: &[u8]) -> Noun {
        Noun(Inner::Atom(Rc::new(digits.to_vec())))
    }

    /// Build a new cell noun from two existing nouns.
    pub fn cell(a: Noun, b: Noun) -> Noun {
        Noun(Inner::Cell(Rc::new(a), Rc::new(b)))
    }

    /// Build a noun from a convertible value.
    pub fn from<T: ToNoun>(item: T) -> Noun {
        item.to_noun()
    }

    /// Run a memoizing fold over the noun
    fn fold<'a, F, T>(&'a self, mut f: F) -> T
        where F: FnMut(Shape<&'a [u8], T>) -> T,
              T: Clone
    {
        fn h<'a, F, T>(noun: &'a Noun,
                       memo: &mut HashMap<usize, T>,
                       f: &mut F)
                       -> T
            where F: FnMut(Shape<&'a [u8], T>) -> T,
                  T: Clone
        {
            let key = noun.addr();

            if memo.contains_key(&key) {
                memo.get(&key).unwrap().clone()
            } else {
                let ret = match noun.get() {
                    Shape::Atom(x) => f(Shape::Atom(x)),
                    Shape::Cell(ref a, ref b) => {
                        let a = h(*a, memo, f);
                        let b = h(*b, memo, f);
                        let ret = f(Shape::Cell(a, b));
                        ret
                    }
                };
                memo.insert(key, ret.clone());
                ret
            }
        }

        h(self, &mut HashMap::new(), &mut f)
    }

}

impl hash::Hash for Noun {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        fn f<H: hash::Hasher>(state: &mut H, shape: Shape<&[u8], u64>) -> u64 {
            match shape {
                Shape::Atom(x) => x.hash(state),
                Shape::Cell(a, b) => {
                    a.hash(state);
                    b.hash(state);
                }
            }
            state.finish()
        }
        self.fold(|x| f(state, x))
            .hash(state);
    }
}

/// Trait for types that can convert themselves to a noun.
pub trait ToNoun {
    fn to_noun(&self) -> Noun;
}

impl<T> ToNoun for T where T: DigitSlice
{
    fn to_noun(&self) -> Noun {
        Noun::atom(self.as_digits())
    }
}

/// A trait for types that can be instantiated from a Nock noun.
pub trait FromNoun: Sized {
    /// The associated error.
    type Err;

    /// Try to convert a noun to an instance of the type.
    fn from_noun(n: &Noun) -> Result<Self, Self::Err>;
}

impl<T> FromNoun for T where T: FromDigits {
    type Err = ();

    fn from_noun(n: &Noun) -> Result<Self, Self::Err> {
        match n.get() {
            Shape::Atom(x) => T::from_digits(x).map_err(|_| ()),
            _ => Err(())
        }
    }
}

impl<T, U> FromNoun for (T, U) where T: FromNoun, U: FromNoun {
    type Err = ();

    fn from_noun(n: &Noun) -> Result<Self, Self::Err> {
        match n.get() {
            Shape::Cell(a, b) => {
                let t = try!(T::from_noun(a).map_err(|_| ()));
                let u = try!(U::from_noun(b).map_err(|_| ()));
                Ok((t, u))
            }
            _ => Err(())
        }
    }
}

// TODO: FromNoun for T: FromIterator<U: FromNoun>. Pair impl should give us
// a HashMap derivation then. Use ~-terminated cell sequence as backend.

// TODO: Turn a ~-terminated noun into a vec or an iter. Can fail if the last
// element isn't a ~, and we'll only know when we hit the last element...
// Return type is Option<Vec<&'a Noun>>?

// TODO: ToNoun for T: IntoIterator<U: ToNoun>.

// TODO: FromNoun/ToNoun for String, compatible with cord datatype.

// TODO: FromNoun/ToNoun for signed numbers using the Urbit representation
// convention.

// Into-conversion is only used so that we can put untyped numeric literals in
// the noun-constructing macro and have them typed as unsigned. If the noun
// constructor uses ToNoun, literals are assumed to be i32, which does not map
// to atoms in quite the way we want.
impl Into<Noun> for u64 {
    fn into(self) -> Noun {
        Noun::from(self)
    }
}

/// Macro for noun literals.
///
/// Rust n![1, 2, 3] corresponds to Nock [1 2 3]
#[macro_export]
macro_rules! n {
    [$x:expr, $y:expr] => { $crate::draft::Noun::cell($x.into(), $y.into()) };
    [$x:expr, $y:expr, $($ys:expr),+] => { $crate::draft::Noun::cell($x.into(), n![$y, $($ys),+]) };
}

#[cfg(test)]
mod tests {
    use std::hash;
    use super::Noun;

    #[test]
    fn scratch() {
        let x = Noun::from(123u32);
        assert!(x == Noun::from(123u8));
    }

    fn hash<T: hash::Hash>(t: &T) -> u64 {
        use std::hash::Hasher;
        let mut s = hash::SipHasher::new();
        t.hash(&mut s);
        s.finish()
    }

    #[test]
    fn test_fold() {
        assert_eq!(hash(&n![1, 2, 3]), hash(&n![1, 2, 3]));
        assert!(hash(&n![n![1, 2], 3]) != hash(&n![1, 2, 3]));
        assert!(hash(&n![1, 2, 3]) != hash(&n![1, 2]));
    }
}
<|endoftext|>


<commit_before>#![crate_name = "uu_du"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Derek Chiang <derekchiang93@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate time;

#[macro_use]
extern crate uucore;

use std::fs;
use std::iter;
use std::io::{stderr, Result, Write};
use std::os::unix::fs::MetadataExt;
use std::path::PathBuf;
use time::Timespec;

const NAME: &'static str = "du";
const SUMMARY: &'static str = "estimate file space usage";
const LONG_HELP: &'static str = "
 Display  values  are  in  units  of  the  first  available  SIZE from
 --block-size,  and the DU_BLOCK_SIZE, BLOCK_SIZE and BLOCKSIZE environ‐
 ment variables.  Otherwise, units default to  1024  bytes  (or  512  if
 POSIXLY_CORRECT is set).

 SIZE  is  an  integer and optional unit (example: 10M is 10*1024*1024).
 Units are K, M, G, T, P, E, Z, Y (powers of 1024) or KB, MB, ...  (pow‐
 ers of 1000).
";

struct Options {
    all: bool,
    program_name: String,
    max_depth: Option<usize>,
    total: bool,
    separate_dirs: bool,
}

struct Stat {
    path: PathBuf,
    is_dir: bool,
    size: u64,
    blocks: u64,
    nlink: u64,
    created: u64,
    accessed: u64,
    modified: u64,
}

impl Stat {
    fn new(path: PathBuf) -> Result<Stat> {
        match fs::symlink_metadata(&path) {
            Ok(metadata) => {
                return Ok(Stat {
                    path: path,
                    is_dir: metadata.is_dir(),
                    size: metadata.len(),
                    blocks: metadata.blocks() as u64,
                    nlink: metadata.nlink() as u64,
                    created: metadata.mtime() as u64,
                    accessed: metadata.atime() as u64,
                    modified: metadata.mtime() as u64,
                })
            }
            Err(e) => Err(e),
        }
    }
}

// this takes `my_stat` to avoid having to stat files multiple times.
// XXX: this should use the impl Trait return type when it is stabilized
fn du(mut my_stat: Stat, options: &Options, depth: usize) -> Box<DoubleEndedIterator<Item = Stat>> {
    let mut stats = vec![];
    let mut futures = vec![];

    if my_stat.is_dir {
        let read = match fs::read_dir(&my_stat.path) {
            Ok(read) => read,
            Err(e) => {
                safe_writeln!(
                    stderr(),
                    "{}: cannot read directory ‘{}‘: {}",
                    options.program_name,
                    my_stat.path.display(),
                    e
                );
                return Box::new(iter::once(my_stat));
            }
        };

        for f in read.into_iter() {
            match f {
                Ok(entry) => match Stat::new(entry.path()) {
                    Ok(this_stat) => {
                        if this_stat.is_dir {
                            futures.push(du(this_stat, options, depth + 1));
                        } else {
                            my_stat.size += this_stat.size;
                            my_stat.blocks += this_stat.blocks;
                            if options.all {
                                stats.push(this_stat);
                            }
                        }
                    }
                    Err(error) => show_error!("{}", error),
                },
                Err(error) => show_error!("{}", error),
            }
        }
    }

    stats.extend(
        futures
            .into_iter()
            .flat_map(|val| val)
            .rev()
            .filter_map(|stat| {
                if !options.separate_dirs && stat.path.parent().unwrap() == my_stat.path {
                    my_stat.size += stat.size;
                    my_stat.blocks += stat.blocks;
                }
                if options.max_depth == None || depth < options.max_depth.unwrap() {
                    Some(stat)
                } else {
                    None
                }
            }),
    );
    stats.push(my_stat);
    Box::new(stats.into_iter())
}

pub fn uumain(args: Vec<String>) -> i32 {
    let syntax = format!(
        "[OPTION]... [FILE]...
 {0} [OPTION]... --files0-from=F",
        NAME
    );
    let matches = new_coreopts!(&syntax, SUMMARY, LONG_HELP)
    // In task
        .optflag("a", "all", " write counts for all files, not just directories")
    // In main
        .optflag("", "apparent-size", "print apparent sizes,  rather  than  disk  usage
            although  the apparent  size is usually smaller, it may be larger due to holes
            in ('sparse') files, internal  fragmentation,  indirect  blocks, and the like")
    // In main
        .optopt("B", "block-size", "scale sizes  by  SIZE before printing them.
            E.g., '-BM' prints sizes in units of 1,048,576 bytes.  See SIZE format below.",
            "SIZE")
    // In main
        .optflag("b", "bytes", "equivalent to '--apparent-size --block-size=1'")
    // In main
        .optflag("c", "total", "produce a grand total")
    // In task
    // opts.optflag("D", "dereference-args", "dereference only symlinks that are listed
    //     on the command line"),
    // In main
    // opts.optopt("", "files0-from", "summarize disk usage of the NUL-terminated file
    //                   names specified in file F;
    //                   If F is - then read names from standard input", "F"),
    // // In task
    // opts.optflag("H", "", "equivalent to --dereference-args (-D)"),
    // In main
        .optflag("h", "human-readable", "print sizes in human readable format (e.g., 1K 234M 2G)")
    // In main
        .optflag("", "si", "like -h, but use powers of 1000 not 1024")
    // In main
        .optflag("k", "", "like --block-size=1K")
    // In task
        .optflag("l", "count-links", "count sizes many times if hard linked")
    // // In main
        .optflag("m", "", "like --block-size=1M")
    // // In task
    // opts.optflag("L", "dereference", "dereference all symbolic links"),
    // // In task
    // opts.optflag("P", "no-dereference", "don't follow any symbolic links (this is the default)"),
    // // In main
        .optflag("0", "null", "end each output line with 0 byte rather than newline")
    // In main
        .optflag("S", "separate-dirs", "do not include size of subdirectories")
    // In main
        .optflag("s", "summarize", "display only a total for each argument")
    // // In task
    // opts.optflag("x", "one-file-system", "skip directories on different file systems"),
    // // In task
    // opts.optopt("X", "exclude-from", "exclude files that match any pattern in FILE", "FILE"),
    // // In task
    // opts.optopt("", "exclude", "exclude files that match PATTERN", "PATTERN"),
    // In main
        .optopt("d", "max-depth", "print the total for a directory (or file, with --all)
            only if it is N or fewer levels below the command
            line argument;  --max-depth=0 is the same as --summarize", "N")
    // In main
        .optflagopt("", "time", "show time of the last modification of any file in the
            directory, or any of its subdirectories.  If WORD is given, show time as WORD instead
            of modification time: atime, access, use, ctime or status", "WORD")
    // In main
        .optopt("", "time-style", "show times using style STYLE:
            full-iso, long-iso, iso, +FORMAT FORMAT is interpreted like 'date'", "STYLE")
        .parse(args);

    let summarize = matches.opt_present("summarize");

    let max_depth_str = matches.opt_str("max-depth");
    let max_depth = max_depth_str.as_ref().and_then(|s| s.parse::<usize>().ok());
    match (max_depth_str, max_depth) {
        (Some(ref s), _) if summarize => {
            show_error!("summarizing conflicts with --max-depth={}", *s);
            return 1;
        }
        (Some(ref s), None) => {
            show_error!("invalid maximum depth '{}'", *s);
            return 1;
        }
        (Some(_), Some(_)) | (None, _) => { /* valid */ }
    }

    let options = Options {
        all: matches.opt_present("all"),
        program_name: NAME.to_owned(),
        max_depth: max_depth,
        total: matches.opt_present("total"),
        separate_dirs: matches.opt_present("S"),
    };

    let strs = if matches.free.is_empty() {
        vec!["./".to_owned()]
    } else {
        matches.free.clone()
    };

    let mb = if matches.opt_present("si") {
        1000 * 1000
    } else {
        1024 * 1024
    };
    let kb = if matches.opt_present("si") {
        1000
    } else {
        1024
    };

    let block_size = match matches.opt_str("block-size") {
        Some(s) => {
            let mut found_number = false;
            let mut found_letter = false;
            let mut numbers = String::new();
            let mut letters = String::new();
            for c in s.chars() {
                if found_letter && c.is_digit(10) || !found_number && !c.is_digit(10) {
                    show_error!("invalid --block-size argument '{}'", s);
                    return 1;
                } else if c.is_digit(10) {
                    found_number = true;
                    numbers.push(c);
                } else if c.is_alphabetic() {
                    found_letter = true;
                    letters.push(c);
                }
            }
            let number = numbers.parse::<u64>().unwrap();
            let multiple = match &letters[..] {
                "K" => 1024u64.pow(1),
                "M" => 1024u64.pow(2),
                "G" => 1024u64.pow(3),
                "T" => 1024u64.pow(4),
                "P" => 1024u64.pow(5),
                "E" => 1024u64.pow(6),
                "Z" => 1024u64.pow(7),
                "Y" => 1024u64.pow(8),
                "KB" => 1000u64.pow(1),
                "MB" => 1000u64.pow(2),
                "GB" => 1000u64.pow(3),
                "TB" => 1000u64.pow(4),
                "PB" => 1000u64.pow(5),
                "EB" => 1000u64.pow(6),
                "ZB" => 1000u64.pow(7),
                "YB" => 1000u64.pow(8),
                _ => {
                    show_error!("invalid --block-size argument '{}'", s);
                    return 1;
                }
            };
            number * multiple
        }
        None => 1024,
    };

    let convert_size = |size: u64| -> String {
        if matches.opt_present("human-readable") || matches.opt_present("si") {
            if size >= mb {
                format!("{:.1}M", (size as f64) / (mb as f64))
            } else if size >= kb {
                format!("{:.1}K", (size as f64) / (kb as f64))
            } else {
                format!("{}B", size)
            }
        } else if matches.opt_present("k") {
            format!("{}", ((size as f64) / (kb as f64)).ceil())
        } else if matches.opt_present("m") {
            format!("{}", ((size as f64) / (mb as f64)).ceil())
        } else {
            format!("{}", ((size as f64) / (block_size as f64)).ceil())
        }
    };

    let time_format_str = match matches.opt_str("time-style") {
        Some(s) => match &s[..] {
            "full-iso" => "%Y-%m-%d %H:%M:%S.%f %z",
            "long-iso" => "%Y-%m-%d %H:%M",
            "iso" => "%Y-%m-%d",
            _ => {
                show_error!(
                    "invalid argument '{}' for 'time style'
Valid arguments are:
- 'full-iso'
- 'long-iso'
- 'iso'
Try '{} --help' for more information.",
                    s,
                    NAME
                );
                return 1;
            }
        },
        None => "%Y-%m-%d %H:%M",
    };

    let line_separator = if matches.opt_present("0") { "\0" } else { "\n" };

    let mut grand_total = 0;
    for path_str in strs.into_iter() {
        let path = PathBuf::from(path_str);
        match Stat::new(path) {
            Ok(stat) => {
                let iter = du(stat, &options, 0).into_iter();
                let (_, len) = iter.size_hint();
                let len = len.unwrap();
                for (index, stat) in iter.enumerate() {
                    let size = if matches.opt_present("apparent-size") {
                        stat.nlink * stat.size
                    } else {
                        // C's stat is such that each block is assume to be 512 bytes
                        // See: http://linux.die.net/man/2/stat
                        stat.blocks * 512
                    };
                    if matches.opt_present("time") {
                        let tm = {
                            let (secs, nsecs) = {
                                let time = match matches.opt_str("time") {
                                    Some(s) => match &s[..] {
                                        "accessed" => stat.accessed,
                                        "created" => stat.created,
                                        "modified" => stat.modified,
                                        _ => {
                                            show_error!(
                                                "invalid argument 'modified' for '--time'
    Valid arguments are:
      - 'accessed', 'created', 'modified'
    Try '{} --help' for more information.",
                                                NAME
                                            );
                                            return 1;
                                        }
                                    },
                                    None => stat.modified,
                                };
                                ((time / 1000) as i64, (time % 1000 * 1000000) as i32)
                            };
                            time::at(Timespec::new(secs, nsecs))
                        };
                        if !summarize || (summarize && index == len - 1) {
                            let time_str = tm.strftime(time_format_str).unwrap();
                            print!(
                                "{}\t{}\t{}{}",
                                convert_size(size),
                                time_str,
                                stat.path.display(),
                                line_separator
                            );
                        }
                    } else {
                        if !summarize || (summarize && index == len - 1) {
                            print!(
                                "{}\t{}{}",
                                convert_size(size),
                                stat.path.display(),
                                line_separator
                            );
                        }
                    }
                    if options.total && index == (len - 1) {
                        // The last element will be the total size of the the path under
                        // path_str.  We add it to the grand total.
                        grand_total += size;
                    }
                }
            }
            Err(error) => show_error!("{}", error),
        }
    }

    if options.total {
        print!("{}\ttotal", convert_size(grand_total));
        print!("{}", line_separator);
    }

    0
}
<commit_msg>du: refactor to use ? shortcut<commit_after>#![crate_name = "uu_du"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Derek Chiang <derekchiang93@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate time;

#[macro_use]
extern crate uucore;

use std::fs;
use std::iter;
use std::io::{stderr, Result, Write};
use std::os::unix::fs::MetadataExt;
use std::path::PathBuf;
use time::Timespec;

const NAME: &'static str = "du";
const SUMMARY: &'static str = "estimate file space usage";
const LONG_HELP: &'static str = "
 Display  values  are  in  units  of  the  first  available  SIZE from
 --block-size,  and the DU_BLOCK_SIZE, BLOCK_SIZE and BLOCKSIZE environ‐
 ment variables.  Otherwise, units default to  1024  bytes  (or  512  if
 POSIXLY_CORRECT is set).

 SIZE  is  an  integer and optional unit (example: 10M is 10*1024*1024).
 Units are K, M, G, T, P, E, Z, Y (powers of 1024) or KB, MB, ...  (pow‐
 ers of 1000).
";

struct Options {
    all: bool,
    program_name: String,
    max_depth: Option<usize>,
    total: bool,
    separate_dirs: bool,
}

struct Stat {
    path: PathBuf,
    is_dir: bool,
    size: u64,
    blocks: u64,
    nlink: u64,
    created: u64,
    accessed: u64,
    modified: u64,
}

impl Stat {
    fn new(path: PathBuf) -> Result<Stat> {
        let metadata = fs::symlink_metadata(&path)?;
        Ok(Stat {
            path: path,
            is_dir: metadata.is_dir(),
            size: metadata.len(),
            blocks: metadata.blocks() as u64,
            nlink: metadata.nlink() as u64,
            created: metadata.mtime() as u64,
            accessed: metadata.atime() as u64,
            modified: metadata.mtime() as u64,
        })
    }
}

// this takes `my_stat` to avoid having to stat files multiple times.
// XXX: this should use the impl Trait return type when it is stabilized
fn du(mut my_stat: Stat, options: &Options, depth: usize) -> Box<DoubleEndedIterator<Item = Stat>> {
    let mut stats = vec![];
    let mut futures = vec![];

    if my_stat.is_dir {
        let read = match fs::read_dir(&my_stat.path) {
            Ok(read) => read,
            Err(e) => {
                safe_writeln!(
                    stderr(),
                    "{}: cannot read directory ‘{}‘: {}",
                    options.program_name,
                    my_stat.path.display(),
                    e
                );
                return Box::new(iter::once(my_stat));
            }
        };

        for f in read.into_iter() {
            match f {
                Ok(entry) => match Stat::new(entry.path()) {
                    Ok(this_stat) => {
                        if this_stat.is_dir {
                            futures.push(du(this_stat, options, depth + 1));
                        } else {
                            my_stat.size += this_stat.size;
                            my_stat.blocks += this_stat.blocks;
                            if options.all {
                                stats.push(this_stat);
                            }
                        }
                    }
                    Err(error) => show_error!("{}", error),
                },
                Err(error) => show_error!("{}", error),
            }
        }
    }

    stats.extend(
        futures
            .into_iter()
            .flat_map(|val| val)
            .rev()
            .filter_map(|stat| {
                if !options.separate_dirs && stat.path.parent().unwrap() == my_stat.path {
                    my_stat.size += stat.size;
                    my_stat.blocks += stat.blocks;
                }
                if options.max_depth == None || depth < options.max_depth.unwrap() {
                    Some(stat)
                } else {
                    None
                }
            }),
    );
    stats.push(my_stat);
    Box::new(stats.into_iter())
}

pub fn uumain(args: Vec<String>) -> i32 {
    let syntax = format!(
        "[OPTION]... [FILE]...
 {0} [OPTION]... --files0-from=F",
        NAME
    );
    let matches = new_coreopts!(&syntax, SUMMARY, LONG_HELP)
    // In task
        .optflag("a", "all", " write counts for all files, not just directories")
    // In main
        .optflag("", "apparent-size", "print apparent sizes,  rather  than  disk  usage
            although  the apparent  size is usually smaller, it may be larger due to holes
            in ('sparse') files, internal  fragmentation,  indirect  blocks, and the like")
    // In main
        .optopt("B", "block-size", "scale sizes  by  SIZE before printing them.
            E.g., '-BM' prints sizes in units of 1,048,576 bytes.  See SIZE format below.",
            "SIZE")
    // In main
        .optflag("b", "bytes", "equivalent to '--apparent-size --block-size=1'")
    // In main
        .optflag("c", "total", "produce a grand total")
    // In task
    // opts.optflag("D", "dereference-args", "dereference only symlinks that are listed
    //     on the command line"),
    // In main
    // opts.optopt("", "files0-from", "summarize disk usage of the NUL-terminated file
    //                   names specified in file F;
    //                   If F is - then read names from standard input", "F"),
    // // In task
    // opts.optflag("H", "", "equivalent to --dereference-args (-D)"),
    // In main
        .optflag("h", "human-readable", "print sizes in human readable format (e.g., 1K 234M 2G)")
    // In main
        .optflag("", "si", "like -h, but use powers of 1000 not 1024")
    // In main
        .optflag("k", "", "like --block-size=1K")
    // In task
        .optflag("l", "count-links", "count sizes many times if hard linked")
    // // In main
        .optflag("m", "", "like --block-size=1M")
    // // In task
    // opts.optflag("L", "dereference", "dereference all symbolic links"),
    // // In task
    // opts.optflag("P", "no-dereference", "don't follow any symbolic links (this is the default)"),
    // // In main
        .optflag("0", "null", "end each output line with 0 byte rather than newline")
    // In main
        .optflag("S", "separate-dirs", "do not include size of subdirectories")
    // In main
        .optflag("s", "summarize", "display only a total for each argument")
    // // In task
    // opts.optflag("x", "one-file-system", "skip directories on different file systems"),
    // // In task
    // opts.optopt("X", "exclude-from", "exclude files that match any pattern in FILE", "FILE"),
    // // In task
    // opts.optopt("", "exclude", "exclude files that match PATTERN", "PATTERN"),
    // In main
        .optopt("d", "max-depth", "print the total for a directory (or file, with --all)
            only if it is N or fewer levels below the command
            line argument;  --max-depth=0 is the same as --summarize", "N")
    // In main
        .optflagopt("", "time", "show time of the last modification of any file in the
            directory, or any of its subdirectories.  If WORD is given, show time as WORD instead
            of modification time: atime, access, use, ctime or status", "WORD")
    // In main
        .optopt("", "time-style", "show times using style STYLE:
            full-iso, long-iso, iso, +FORMAT FORMAT is interpreted like 'date'", "STYLE")
        .parse(args);

    let summarize = matches.opt_present("summarize");

    let max_depth_str = matches.opt_str("max-depth");
    let max_depth = max_depth_str.as_ref().and_then(|s| s.parse::<usize>().ok());
    match (max_depth_str, max_depth) {
        (Some(ref s), _) if summarize => {
            show_error!("summarizing conflicts with --max-depth={}", *s);
            return 1;
        }
        (Some(ref s), None) => {
            show_error!("invalid maximum depth '{}'", *s);
            return 1;
        }
        (Some(_), Some(_)) | (None, _) => { /* valid */ }
    }

    let options = Options {
        all: matches.opt_present("all"),
        program_name: NAME.to_owned(),
        max_depth: max_depth,
        total: matches.opt_present("total"),
        separate_dirs: matches.opt_present("S"),
    };

    let strs = if matches.free.is_empty() {
        vec!["./".to_owned()]
    } else {
        matches.free.clone()
    };

    let mb = if matches.opt_present("si") {
        1000 * 1000
    } else {
        1024 * 1024
    };
    let kb = if matches.opt_present("si") {
        1000
    } else {
        1024
    };

    let block_size = match matches.opt_str("block-size") {
        Some(s) => {
            let mut found_number = false;
            let mut found_letter = false;
            let mut numbers = String::new();
            let mut letters = String::new();
            for c in s.chars() {
                if found_letter && c.is_digit(10) || !found_number && !c.is_digit(10) {
                    show_error!("invalid --block-size argument '{}'", s);
                    return 1;
                } else if c.is_digit(10) {
                    found_number = true;
                    numbers.push(c);
                } else if c.is_alphabetic() {
                    found_letter = true;
                    letters.push(c);
                }
            }
            let number = numbers.parse::<u64>().unwrap();
            let multiple = match &letters[..] {
                "K" => 1024u64.pow(1),
                "M" => 1024u64.pow(2),
                "G" => 1024u64.pow(3),
                "T" => 1024u64.pow(4),
                "P" => 1024u64.pow(5),
                "E" => 1024u64.pow(6),
                "Z" => 1024u64.pow(7),
                "Y" => 1024u64.pow(8),
                "KB" => 1000u64.pow(1),
                "MB" => 1000u64.pow(2),
                "GB" => 1000u64.pow(3),
                "TB" => 1000u64.pow(4),
                "PB" => 1000u64.pow(5),
                "EB" => 1000u64.pow(6),
                "ZB" => 1000u64.pow(7),
                "YB" => 1000u64.pow(8),
                _ => {
                    show_error!("invalid --block-size argument '{}'", s);
                    return 1;
                }
            };
            number * multiple
        }
        None => 1024,
    };

    let convert_size = |size: u64| -> String {
        if matches.opt_present("human-readable") || matches.opt_present("si") {
            if size >= mb {
                format!("{:.1}M", (size as f64) / (mb as f64))
            } else if size >= kb {
                format!("{:.1}K", (size as f64) / (kb as f64))
            } else {
                format!("{}B", size)
            }
        } else if matches.opt_present("k") {
            format!("{}", ((size as f64) / (kb as f64)).ceil())
        } else if matches.opt_present("m") {
            format!("{}", ((size as f64) / (mb as f64)).ceil())
        } else {
            format!("{}", ((size as f64) / (block_size as f64)).ceil())
        }
    };

    let time_format_str = match matches.opt_str("time-style") {
        Some(s) => match &s[..] {
            "full-iso" => "%Y-%m-%d %H:%M:%S.%f %z",
            "long-iso" => "%Y-%m-%d %H:%M",
            "iso" => "%Y-%m-%d",
            _ => {
                show_error!(
                    "invalid argument '{}' for 'time style'
Valid arguments are:
- 'full-iso'
- 'long-iso'
- 'iso'
Try '{} --help' for more information.",
                    s,
                    NAME
                );
                return 1;
            }
        },
        None => "%Y-%m-%d %H:%M",
    };

    let line_separator = if matches.opt_present("0") { "\0" } else { "\n" };

    let mut grand_total = 0;
    for path_str in strs.into_iter() {
        let path = PathBuf::from(path_str);
        match Stat::new(path) {
            Ok(stat) => {
                let iter = du(stat, &options, 0).into_iter();
                let (_, len) = iter.size_hint();
                let len = len.unwrap();
                for (index, stat) in iter.enumerate() {
                    let size = if matches.opt_present("apparent-size") {
                        stat.nlink * stat.size
                    } else {
                        // C's stat is such that each block is assume to be 512 bytes
                        // See: http://linux.die.net/man/2/stat
                        stat.blocks * 512
                    };
                    if matches.opt_present("time") {
                        let tm = {
                            let (secs, nsecs) = {
                                let time = match matches.opt_str("time") {
                                    Some(s) => match &s[..] {
                                        "accessed" => stat.accessed,
                                        "created" => stat.created,
                                        "modified" => stat.modified,
                                        _ => {
                                            show_error!(
                                                "invalid argument 'modified' for '--time'
    Valid arguments are:
      - 'accessed', 'created', 'modified'
    Try '{} --help' for more information.",
                                                NAME
                                            );
                                            return 1;
                                        }
                                    },
                                    None => stat.modified,
                                };
                                ((time / 1000) as i64, (time % 1000 * 1000000) as i32)
                            };
                            time::at(Timespec::new(secs, nsecs))
                        };
                        if !summarize || (summarize && index == len - 1) {
                            let time_str = tm.strftime(time_format_str).unwrap();
                            print!(
                                "{}\t{}\t{}{}",
                                convert_size(size),
                                time_str,
                                stat.path.display(),
                                line_separator
                            );
                        }
                    } else {
                        if !summarize || (summarize && index == len - 1) {
                            print!(
                                "{}\t{}{}",
                                convert_size(size),
                                stat.path.display(),
                                line_separator
                            );
                        }
                    }
                    if options.total && index == (len - 1) {
                        // The last element will be the total size of the the path under
                        // path_str.  We add it to the grand total.
                        grand_total += size;
                    }
                }
            }
            Err(error) => show_error!("{}", error),
        }
    }

    if options.total {
        print!("{}\ttotal", convert_size(grand_total));
        print!("{}", line_separator);
    }

    0
}
<|endoftext|>