<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2015 Johan Johansson
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

//! Infer types where explicit type signatures are not available

// TODO: Check for mutual recursion when infering types. Maybe list of function call ancestry.
// Like: foo calling bar calling baz calling foo => ERROR

// TODO: Consider redisigning this module. Maybe have an Inferer that takes expressions instead
// of implementing inference for each expression type.

use std::collections::HashMap;
use super::{ FnDef, Item, Type, TypedBinding, Expr };

impl super::FnDef {
	// TODO: Infer types for incomplete function sig. E.g. inc: <→ u32 _> => inc: <→ u32 u32>
	fn infer_types(&mut self, def_map: HashMap<&str, FnDef>) {
		let mut arg_bindings = self.arg_bindings.clone();

		if let Some((fn_arg_types, fn_body_type)) = self.extract_type_sig() {
			// Type signature already exists. Just infer types for body and make sure
			// there are no collisions

			self.set_arg_types(fn_arg_types);

			self.body.infer_types(Some(fn_body_type), def_map, &mut arg_bindings);
		} else {
			// No type signature for function binding. Pass arg bindings to body and infer types
			// for body, then get function signature from updated arg types and body type

			let arg_bindings_old_len = arg_bindings.len();

			// args: type_to_infer_to, def_map, var_stack
			self.body.infer_types(None, def_map, &mut arg_bindings);

			if arg_bindings.len() != arg_bindings_old_len {
				panic!("FnDef::infer_types: arg_bindings.len() != arg_bindings_old_len");
			}

			let arg_types = arg_bindings.into_iter()
				.map(|b| b.type_sig.expect("FnDef::infer_types: Arg has no type"))
				.collect();

			self.set_arg_types(&mut arg_types);

			self.construct_fn_sig(
				arg_types,
				self.body.type_sig.expect("FnDef::infer_types: Body has no type"));
		}
	}
}

impl super::SExpr {
	fn body_type(&self) -> Option<&Type> {
		extract_fn_sig(self.func.type_sig).map(|(_, fn_body_type)| fn_body_type)
	}

	fn infer_types(
		&mut self,
		parent_expected_type: Option<&Type>,
		def_map: &mut HashMap<&str, FnDef>,
		var_stack: &mut Vec<TypedBinding>)
	{
		self.infer_arg_types(None, def_map, var_stack);

		// TODO: Partial inference when not all bindings have type signatures
		let expected_fn_type = if self.args.iter().all(TypedBinding::has_type)
			&& parent_expected_type.is_some()
		{
			Some(Type::construct("→", self.args.iter().map(|tb| tb.type_sig).collect()))
		} else {
			None
		};

		self.func.infer_types(expected_fn_type, def_map, var_stack);

		let found_fn_type = self.func.type_sig;

		// TODO: This only works for function pointers, i.e. lambdas will need some different type.
		//       When traits are added, use a function trait like Rusts Fn/FnMut/FnOnce
		if found_fn_type.is_some() && expected_fn_type.is_none() {
			if let Type::Construct(_, ref fn_arg_types) = found_fn_type {
				self.infer_arg_types(Some(fn_arg_types), def_map, var_stack);
			} else {
				panic!("SExpr::infer_types: `{:?}` is not a function type", found_fn_type);
			}
		} else if found_fn_type.is_none() {
			panic!("SExpr::infer_types: Could not infer type for function");
		} else if found_fn_type != expected_fn_type {
			panic!(
				"SExpr::infer_types: Function type mismatch. Expected `{:?}`, found `{:?}`",
				expected_fn_type,
				found_fn_type);
		}
	}
}

impl super::Block {
	fn get_type(&self) -> Option<&Type> {
		self.exprs.last().type_sig.as_ref()
	}
	fn infer_types(
		&mut self,
		parent_expected_type: Option<&Type>,
		def_map: &mut HashMap<&str, FnDef>,
		var_stack: &mut Vec<TypedBinding>)
	{
		join_map(def_map, self.constant_defs);

		let old_vars_len = var_stack.len();

		// First pass. If possible, all vars defined in block should have types infered.
		for expr in self.exprs.init_mut() {
			if let Expr::VarDef(ref mut var_def) = *expr {
				var_def.infer_types(def_map, var_stack);

				var_stack.push(var_def.binding);
			} else {
				expr.infer_types(None, def_map, var_stack);
			}
		}

		if self.exprs.last_mut().expr().is_var_def() {
			panic!("Block::infer_types: Last expression in block is var definition")
		} else {
			self.exprs.last_mut().infer_types(parent_expected_type, def_map, var_stack);
		}

		let block_defined_vars = var_stack.split_off(old_vars_len).into_iter();

		// Second pass. Infer types for all expressions in block now that types for all bindings
		// are, if possible, known.
		for expr in self.exprs.init_mut() {
			if let Expr::VarDef(ref mut var_def) = *expr {
				var_def.infer_types(def_map, var_stack);

				var_stack.push(block_defined_vars.next().unwrap());
			} else {
				expr.infer_types(None, def_map, var_stack);
			}
		}

		if self.exprs.last_mut().expr().is_var_def() {
			panic!("Block::infer_types: Last expression in block is var definition")
		} else {
			self.exprs.last_mut().infer_types(parent_expected_type, def_map, var_stack);
		}

		subtract_map(def_map, self.constant_defs.keys());
	}
}

impl super::Cond {
	fn infer_types(
		&mut self,
		parent_expected_type: Option<&Type>,
		def_map: &mut HashMap<&str, FnDef>,
		var_stack: &mut Vec<TypedBinding>,
		caller_stack: &mut Vec<&str>)
	{
		for (predicate, consequence) in self.clauses.iter_mut()
			.map(|&mut (ref mut p, ref mut c)| (p, c))
			.chain(self.else_clause.iter_mut().map(|c| (&mut ExprMeta::new_true(), c)))
		{
			predicate.infer_types(Some(&Type::bool()), def_map, var_stack, caller_stack);
			consequence.infer_types(parent_expected_type, def_map, var_stack, caller_stack);
		}

		if parent_expected_type.is_none() {
			match self.clauses.iter()
				.map(|&(_, ref c)| c)
				.chain(self.else_clause.iter())
				.find(|clause| clause.type_sig.is_some())
				.map(|t| t.cloned())
			{
				Some(expected_type) => if !self.clauses.iter().map(|&(_, ref c)| c)
					.chain(self.else_clause.iter())
					.all(|clause| clause.type_sig.is_some())
				{
					self.infer_types(expected_type.as_ref(), def_map, var_stack, caller_stack)
				},
				None => panic!("Cond::infer_types: Could not infer type for any clause"),
			}
		}
	}
}

fn get_stack_binding<'a>(var_stack: &'a mut [TypedBinding], bnd: &str) -> &'a mut TypedBinding {
	var_stack.iter_mut()
		.rev()
		.find(|stack_bnd| stack_bnd.ident == bnd)
		.expect(format!("get_stack_binding: Binding not in stack, `{}`", bnd))
}

fn update_binding_stack(var_stack: &mut [TypedBinding], bnd: &str, expected_ty: Type) {
	let stack_bnd = get_stack_binding(var_stack, bnd);

	if stack_bnd.type_sig.is_none() {
		stack_bnd.type_sig = Some(expected_ty)
	}
}

// TODO: Shouldn't necessarily panic if types differ. Add some sort of coercion and polymorphism.
/// Checks binding stack for the type of binding `bnd` and resolves conflicts.
/// If expected type is None, return stack type. If stack type is None, update stack type
/// with to expected type. If both are Some but differ in value, panic.
fn resolve_var_type(var_stack: &mut [TypedBinding], bnd: &str, expected_ty: Option<Type>)
	-> Option<Type>
{
	match expected_ty {
		Some(ty) => update_binding_stack(var_stack, bnd, ty),
		None => get_stack_binding(var_stack, bnd).clone(),
	}
}

impl super::ExprMeta {
	fn infer_types(
		&mut self,
		parent_expected_type: Option<&Type>,
		def_map: &mut HashMap<&str, FnDef>,
		var_stack: &mut Vec<TypedBinding>)
	{
		let expected_type = self.coerce_type.or(parent_expected_type);

		let found_type = match *self.value {
			// Doesn't have children to infer types for
			Expr::Nil => Type::nil(),
			// TODO: This should be an internal, more general integer type
			Expr::NumLit(_) => Type::basic("u64"),
			// TODO: This should be a construct somehow
			Expr::StrLit(_) => Type::basic("&str"),
			Expr::Binding(ref bnd) => resolve_var_type(var_stack, bnd, expected_type)
				.or(resolve_def_type(def_map, bnd, expected_type)),
			Expr::SExpr(ref mut sexpr) => {
				sexpr.infer_types(expected_type, var_stack);
				sexpr.body_type()
			},
			Expr::Block(ref mut block) => {
				block.infer_types(expected_type, var_stack);
				block.type_sig
			},
			Expr::Cond(ref mut cond) => {
				cond.infer_types(expected_type, var_stack);
				cond.type_sig
			},
			Expr::Lambda(ref mut lambda) => {
				lambda.infer_types(expected_type, var_stack);
				lambda.type_sig
			},
		};

		self.set_type(found_type);
	}
}

impl super::AST {
	pub fn infer_types(&mut self) {
		if let Some(main_def) = self.items.iter_mut()
			.find(|meta| if let Item::FnDef(ref def) = *meta.item {
				def.binding.ident == "main"
			} else {
				false
			})
		{
			main_def.infer_types(&mut self.items).unwrap()
		} else {
			panic!("AST::infer_types: No main function found")
		}
	}
}
<commit_msg>Add join_map and subtract_map for managing hashmaps of constant defs<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2015 Johan Johansson
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

//! Infer types where explicit type signatures are not available

// TODO: Check for mutual recursion when infering types. Maybe list of function call ancestry.
// Like: foo calling bar calling baz calling foo => ERROR

// TODO: Consider redisigning this module. Maybe have an Inferer that takes expressions instead
// of implementing inference for each expression type.

use std::collections::{ HashMap, HashSet };
use std::hash::Hash;
use super::{ FnDef, Item, Type, TypedBinding, Expr, ExprMeta };

impl super::FnDef {
	// TODO: Infer types for incomplete function sig. E.g. inc: <→ u32 _> => inc: <→ u32 u32>
	fn infer_types(&mut self, def_map: HashMap<&str, FnDef>) {
		let mut arg_bindings = self.arg_bindings.clone();

		if let Some((fn_arg_types, fn_body_type)) = self.extract_type_sig() {
			// Type signature already exists. Just infer types for body and make sure
			// there are no collisions

			self.set_arg_types(fn_arg_types);

			self.body.infer_types(Some(fn_body_type), def_map, &mut arg_bindings);
		} else {
			// No type signature for function binding. Pass arg bindings to body and infer types
			// for body, then get function signature from updated arg types and body type

			let arg_bindings_old_len = arg_bindings.len();

			// args: type_to_infer_to, def_map, var_stack
			self.body.infer_types(None, def_map, &mut arg_bindings);

			if arg_bindings.len() != arg_bindings_old_len {
				panic!("FnDef::infer_types: arg_bindings.len() != arg_bindings_old_len");
			}

			let arg_types = arg_bindings.into_iter()
				.map(|b| b.type_sig.expect("FnDef::infer_types: Arg has no type"))
				.collect();

			self.set_arg_types(&mut arg_types);

			self.construct_fn_sig(
				arg_types,
				self.body.type_sig.expect("FnDef::infer_types: Body has no type"));
		}
	}
}

/// Extract a function type signature in the form of <→ arg1 arg2 body> to (&[arg1, arg2], body)
fn extract_fn_sig(sig: &Type) -> (&[Type], &Type) {
	match sig {
		&Type::Construct(ref c, ref ts) if c == "fn" || c == "→" => (ts.init(), ts.last()),
		t => panic!("extract_fn_sig: `{:?}` is not a function type", t),
	}
}

impl super::SExpr {
	fn body_type(&self) -> Option<&Type> {
		self.func.type_sig.as_ref().map(extract_fn_sig).map(|(_, body_t)| body_t)
	}

	fn infer_types(
		&mut self,
		parent_expected_type: Option<&Type>,
		def_map: &mut HashMap<&str, FnDef>,
		var_stack: &mut Vec<TypedBinding>)
	{
		self.infer_arg_types(None, def_map, var_stack);

		// TODO: Partial inference when not all bindings have type signatures
		let expected_fn_type = if self.args.iter().all(TypedBinding::has_type)
			&& parent_expected_type.is_some()
		{
			Some(Type::construct("→", self.args.iter().map(|tb| tb.type_sig).collect()))
		} else {
			None
		};

		self.func.infer_types(expected_fn_type, def_map, var_stack);

		let found_fn_type = self.func.type_sig;

		// TODO: This only works for function pointers, i.e. lambdas will need some different type.
		//       When traits are added, use a function trait like Rusts Fn/FnMut/FnOnce
		if found_fn_type.is_some() && expected_fn_type.is_none() {
			if let Type::Construct(_, ref fn_arg_types) = found_fn_type {
				self.infer_arg_types(Some(fn_arg_types), def_map, var_stack);
			} else {
				panic!("SExpr::infer_types: `{:?}` is not a function type", found_fn_type);
			}
		} else if found_fn_type.is_none() {
			panic!("SExpr::infer_types: Could not infer type for function");
		} else if found_fn_type != expected_fn_type {
			panic!(
				"SExpr::infer_types: Function type mismatch. Expected `{:?}`, found `{:?}`",
				expected_fn_type,
				found_fn_type);
		}
	}
}

/// Moves all values in `rhs` to `main`, creating a union of the two maps.
/// If a key from `rhs` already exists in `main`, return `rhs` as an error.
/// On success, return the HashSet of the keys of `rhs`
fn join_map<K: Hash+Eq, V, S>(main: &mut HashMap<K, V, S>, rhs: HashMap<K, V, S>)
	-> Result(HashSet<K>, HashMap<K, V, S>)
{
	if rhs.keys().all(|key| main.contains_key(key)) {
		let mut set = HashSet::with_capacity(rhs.len());

		for (key, val) in rhs {
			main.insert(key, val);
			set.insert(key);
		}

		Ok(set)
	} else {
		Err(rhs)
	}
}

/// Subtract all entries in `map` of keys in `keys` and return the difference.
fn subtract_map<K: Hash+Eq, V, S>(map: &mut HashMap<K, V, S>, keys: &HashSet<K>)
	-> Option<HashMap<K, V>>
{
	if keys.iter().all(|key| map.contains_key(key)) {
		Some(HashMap::from_iter(keys.iter().map(|key| (key, map.remove(key).unwrap()))))
	} else {
		None
	}
}

impl super::Block {
	fn get_type(&self) -> Option<&Type> {
		self.exprs.last().type_sig.as_ref()
	}
	fn infer_types(
		&mut self,
		parent_expected_type: Option<&Type>,
		def_map: &mut HashMap<&str, FnDef>,
		var_stack: &mut Vec<TypedBinding>)
	{
		let constant_keys = join_map(def_map, self.constant_defs.drain());

		let old_vars_len = var_stack.len();

		// First pass. If possible, all vars defined in block should have types infered.
		for expr in self.exprs.init_mut() {
			if let Expr::VarDef(ref mut var_def) = *expr {
				var_def.infer_types(def_map, var_stack);

				var_stack.push(var_def.binding);
			} else {
				expr.infer_types(None, def_map, var_stack);
			}
		}

		if self.exprs.last_mut().expr().is_var_def() {
			panic!("Block::infer_types: Last expression in block is var definition")
		} else {
			self.exprs.last_mut().infer_types(parent_expected_type, def_map, var_stack);
		}

		let block_defined_vars = var_stack.split_off(old_vars_len).into_iter();

		// Second pass. Infer types for all expressions in block now that types for all bindings
		// are, if possible, known.
		for expr in self.exprs.init_mut() {
			if let Expr::VarDef(ref mut var_def) = *expr {
				var_def.infer_types(def_map, var_stack);

				var_stack.push(block_defined_vars.next().unwrap());
			} else {
				expr.infer_types(None, def_map, var_stack);
			}
		}

		if self.exprs.last_mut().expr().is_var_def() {
			panic!("Block::infer_types: Last expression in block is var definition")
		} else {
			self.exprs.last_mut().infer_types(parent_expected_type, def_map, var_stack);
		}

		self.constant_defs = subtract_map(def_map, constant_keys);
	}
}

impl super::Cond {
	fn infer_types(
		&mut self,
		parent_expected_type: Option<&Type>,
		def_map: &mut HashMap<&str, FnDef>,
		var_stack: &mut Vec<TypedBinding>,
		caller_stack: &mut Vec<&str>)
	{
		for (predicate, consequence) in self.clauses.iter_mut()
			.map(|&mut (ref mut p, ref mut c)| (p, c))
			.chain(self.else_clause.iter_mut().map(|c| (&mut ExprMeta::new_true(), c)))
		{
			predicate.infer_types(Some(&Type::bool()), def_map, var_stack, caller_stack);
			consequence.infer_types(parent_expected_type, def_map, var_stack, caller_stack);
		}

		if parent_expected_type.is_none() {
			match self.clauses.iter()
				.map(|&(_, ref c)| c)
				.chain(self.else_clause.iter())
				.find(|clause| clause.type_sig.is_some())
				.map(|t| t.cloned())
			{
				Some(expected_type) => if !self.clauses.iter().map(|&(_, ref c)| c)
					.chain(self.else_clause.iter())
					.all(|clause| clause.type_sig.is_some())
				{
					self.infer_types(expected_type.as_ref(), def_map, var_stack, caller_stack)
				},
				None => panic!("Cond::infer_types: Could not infer type for any clause"),
			}
		}
	}
}

fn get_stack_binding<'a>(var_stack: &'a mut [TypedBinding], bnd: &str) -> &'a mut TypedBinding {
	var_stack.iter_mut()
		.rev()
		.find(|stack_bnd| stack_bnd.ident == bnd)
		.expect(format!("get_stack_binding: Binding not in stack, `{}`", bnd))
}

fn update_binding_stack(var_stack: &mut [TypedBinding], bnd: &str, expected_ty: Type) {
	let stack_bnd = get_stack_binding(var_stack, bnd);

	if stack_bnd.type_sig.is_none() {
		stack_bnd.type_sig = Some(expected_ty)
	}
}

// TODO: Shouldn't necessarily panic if types differ. Add some sort of coercion and polymorphism.
/// Checks binding stack for the type of binding `bnd` and resolves conflicts.
/// If expected type is None, return stack type. If stack type is None, update stack type
/// with to expected type. If both are Some but differ in value, panic.
fn resolve_var_type(var_stack: &mut [TypedBinding], bnd: &str, expected_ty: Option<Type>)
	-> Option<Type>
{
	match expected_ty {
		Some(ty) => update_binding_stack(var_stack, bnd, ty),
		None => get_stack_binding(var_stack, bnd).clone(),
	}
}

impl ExprMeta {
	fn infer_types(
		&mut self,
		parent_expected_type: Option<&Type>,
		def_map: &mut HashMap<&str, FnDef>,
		var_stack: &mut Vec<TypedBinding>)
	{
		let expected_type = self.coerce_type.or(parent_expected_type);

		let found_type = match *self.value {
			// Doesn't have children to infer types for
			Expr::Nil => Type::nil(),
			// TODO: This should be an internal, more general integer type
			Expr::NumLit(_) => Type::basic("u64"),
			// TODO: This should be a construct somehow
			Expr::StrLit(_) => Type::basic("&str"),
			Expr::Binding(ref bnd) => resolve_var_type(var_stack, bnd, expected_type)
				.or(resolve_def_type(def_map, bnd, expected_type)),
			Expr::SExpr(ref mut sexpr) => {
				sexpr.infer_types(expected_type, var_stack);
				sexpr.body_type()
			},
			Expr::Block(ref mut block) => {
				block.infer_types(expected_type, var_stack);
				block.type_sig
			},
			Expr::Cond(ref mut cond) => {
				cond.infer_types(expected_type, var_stack);
				cond.type_sig
			},
			Expr::Lambda(ref mut lambda) => {
				lambda.infer_types(expected_type, var_stack);
				lambda.type_sig
			},
		};

		self.set_type(found_type);
	}
}

impl super::AST {
	pub fn infer_types(&mut self) {
		if let Some(main_def) = self.items.iter_mut()
			.find(|meta| if let Item::FnDef(ref def) = *meta.item {
				def.binding.ident == "main"
			} else {
				false
			})
		{
			main_def.infer_types(&mut self.items).unwrap()
		} else {
			panic!("AST::infer_types: No main function found")
		}
	}
}
<|endoftext|>


<commit_before>use std::fmt;
use std::ops::{Index, IndexMut};

use helper;

// A n-dimensional array
pub struct Array<T> {
    shape: Vec<usize>,
    dim_steps: Vec<usize>, // the ‘volume’ of 1 unit in each dimension.
    buffer: Vec<T>,
}

impl<T> Array<T> {
    pub fn new(shape: Vec<usize>, initial: T) -> Array<T> where T: Clone {
        let buf_size = shape.iter().fold(1, |a, b| a*b);
        let dim_steps = helper::compute_dim_steps(&shape);
        Array {
            shape: shape,
            dim_steps: dim_steps,
            buffer: vec![initial; buf_size],
        }
    }

    pub fn from_vec(shape: Vec<usize>, vec: Vec<T>) -> Array<T> {
        let dim_steps = helper::compute_dim_steps(&shape);
        Array {
            shape: shape,
            dim_steps: dim_steps,
            buffer: vec,
        }
    }

    pub fn get<'a, 'b, I: IntoIterator<Item=&'b usize>>(&'a self, coords: I) -> &'a T {
        let index: usize = coords.into_iter().zip(self.dim_steps.iter())
                                 .map(|(c, s)| (*c)*(*s))
                                 .sum();
        &self.buffer[index]
    }

    pub fn get_mut<'a, 'b, I: IntoIterator<Item=&'b usize>>(&'a mut self, coords: I) -> &'a mut T {
        let index: usize = coords.into_iter().zip(self.dim_steps.iter())
                                 .map(|(c, s)| (*c)*(*s))
                                 .sum();
        &mut self.buffer[index]
    }

    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    pub fn dim_steps(&self) -> &[usize] {
        &self.shape
    }

    pub fn buffer(&self) -> &[T] {
        &self.buffer
    }

    pub fn buffer_mut(&mut self) -> &mut [T] {
        &mut self.buffer
    }
}

impl<'a, 'b, T, I: IntoIterator<Item=&'b usize>> Index<I> for Array<T> {
    type Output = T;

    fn index<'r>(&'r self, index: I) -> &'r T {
        self.get(index)
    }
}

impl<'a, 'b, T, I: IntoIterator<Item=&'b usize>> IndexMut<I> for Array<T> {
    fn index_mut<'r>(&'r mut self, index: I) -> &'r mut T {
        self.get_mut(index)
    }
}

impl<T: Clone+fmt::Debug> fmt::Debug for Array<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        try!(write!(f, "[\n"));
        for row in 0..self.shape[0] {
            try!(write!(f, "[{:?}", self.get(&[row, 0])));
            for col in 1..self.shape[1] {
                try!(write!(f, "\t{:?}", self.get(&[row, col])));
            }
            try!(write!(f, "]\n"));
        }
        try!(write!(f, "]\n"));
        Ok(())
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

#[test]
fn test_array_indexing() {
    let a = Array::from_vec(vec![2, 3, 4], vec![1, 2, 3, 4,
                                                5, 6, 7, 8,
                                                9, 10, 11, 12,

                                                13, 14, 15, 16,
                                                17, 18, 19, 20,
                                                21, 22, 23, 24]);
    assert!(a[&[0, 0, 0]] == 1);
    assert!(a[&[0, 0, 1]] == 2);
    assert!(a[&[0, 0, 2]] == 3);
    assert!(a[&[0, 0, 3]] == 4);
    assert!(a[&[0, 1, 0]] == 5);
    assert!(a[&[0, 1, 1]] == 6);
    assert!(a[&[0, 1, 2]] == 7);
    assert!(a[&[0, 1, 3]] == 8);
    assert!(a[&[0, 2, 0]] == 9);
    assert!(a[&[0, 2, 1]] == 10);
    assert!(a[&[0, 2, 2]] == 11);
    assert!(a[&[0, 2, 3]] == 12);
    assert!(a[&[1, 0, 0]] == 13);
    assert!(a[&[1, 0, 1]] == 14);
    assert!(a[&[1, 0, 2]] == 15);
    assert!(a[&[1, 0, 3]] == 16);
    assert!(a[&[1, 1, 0]] == 17);
    assert!(a[&[1, 1, 1]] == 18);
    assert!(a[&[1, 1, 2]] == 19);
    assert!(a[&[1, 1, 3]] == 20);
    assert!(a[&[1, 2, 0]] == 21);
    assert!(a[&[1, 2, 1]] == 22);
    assert!(a[&[1, 2, 2]] == 23);
    assert!(a[&[1, 2, 3]] == 24);
}

#[test]
fn test_array_indexing_mut() {
    let mut a = Array::from_vec(vec![3], vec![1, 2, 3]);
    a[&[1]] = 42;
    assert!(a[&[1]] == 42);
}
<commit_msg>Array::reshape<commit_after>use std::fmt;
use std::ops::{Index, IndexMut};

use helper;

// A n-dimensional array
pub struct Array<T> {
    shape: Vec<usize>,
    dim_steps: Vec<usize>, // the ‘volume’ of 1 unit in each dimension.
    buffer: Vec<T>,
}

impl<T> Array<T> {
    pub fn new(shape: Vec<usize>, initial: T) -> Array<T> where T: Clone {
        let buf_size = shape.iter().fold(1, |a, b| a*b);
        let dim_steps = helper::compute_dim_steps(&shape);
        Array {
            shape: shape,
            dim_steps: dim_steps,
            buffer: vec![initial; buf_size],
        }
    }

    pub fn from_vec(shape: Vec<usize>, vec: Vec<T>) -> Array<T> {
        let dim_steps = helper::compute_dim_steps(&shape);
        Array {
            shape: shape,
            dim_steps: dim_steps,
            buffer: vec,
        }
    }

    pub fn reshape(&mut self, new_shape: Vec<usize>) {
        let buf_size = new_shape.iter().fold(1, |a, b| a*b);

        if buf_size != self.buffer.len() {
            panic!("Failed to reshape Array of shape {:?} to {:?}", self.shape, new_shape);
        }
        self.dim_steps = helper::compute_dim_steps(&new_shape);
        self.shape = new_shape;
    }

    pub fn get<'a, 'b, I: IntoIterator<Item=&'b usize>>(&'a self, coords: I) -> &'a T {
        let index: usize = coords.into_iter().zip(self.dim_steps.iter())
                                 .map(|(c, s)| (*c)*(*s))
                                 .sum();
        &self.buffer[index]
    }

    pub fn get_mut<'a, 'b, I: IntoIterator<Item=&'b usize>>(&'a mut self, coords: I) -> &'a mut T {
        let index: usize = coords.into_iter().zip(self.dim_steps.iter())
                                 .map(|(c, s)| (*c)*(*s))
                                 .sum();
        &mut self.buffer[index]
    }

    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    pub fn dim_steps(&self) -> &[usize] {
        &self.shape
    }

    pub fn buffer(&self) -> &[T] {
        &self.buffer
    }

    pub fn buffer_mut(&mut self) -> &mut [T] {
        &mut self.buffer
    }
}

impl<'a, 'b, T, I: IntoIterator<Item=&'b usize>> Index<I> for Array<T> {
    type Output = T;

    fn index<'r>(&'r self, index: I) -> &'r T {
        self.get(index)
    }
}

impl<'a, 'b, T, I: IntoIterator<Item=&'b usize>> IndexMut<I> for Array<T> {
    fn index_mut<'r>(&'r mut self, index: I) -> &'r mut T {
        self.get_mut(index)
    }
}

impl<T: Clone+fmt::Debug> fmt::Debug for Array<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        try!(write!(f, "[\n"));
        for row in 0..self.shape[0] {
            try!(write!(f, "[{:?}", self.get(&[row, 0])));
            for col in 1..self.shape[1] {
                try!(write!(f, "\t{:?}", self.get(&[row, col])));
            }
            try!(write!(f, "]\n"));
        }
        try!(write!(f, "]\n"));
        Ok(())
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

#[test]
fn test_reshape() {
    let mut a = Array::from_vec(vec![2, 3], vec![1, 2, 3,
                                                 4, 5, 6]);
    a.reshape(vec![6]);

    assert!(a[&[0]] == 1);
    assert!(a[&[1]] == 2);
    assert!(a[&[2]] == 3);
    assert!(a[&[3]] == 4);
    assert!(a[&[4]] == 5);
    assert!(a[&[5]] == 6);
}

#[test]
fn test_array_indexing() {
    let a = Array::from_vec(vec![2, 3, 4], vec![1, 2, 3, 4,
                                                5, 6, 7, 8,
                                                9, 10, 11, 12,

                                                13, 14, 15, 16,
                                                17, 18, 19, 20,
                                                21, 22, 23, 24]);
    assert!(a[&[0, 0, 0]] == 1);
    assert!(a[&[0, 0, 1]] == 2);
    assert!(a[&[0, 0, 2]] == 3);
    assert!(a[&[0, 0, 3]] == 4);
    assert!(a[&[0, 1, 0]] == 5);
    assert!(a[&[0, 1, 1]] == 6);
    assert!(a[&[0, 1, 2]] == 7);
    assert!(a[&[0, 1, 3]] == 8);
    assert!(a[&[0, 2, 0]] == 9);
    assert!(a[&[0, 2, 1]] == 10);
    assert!(a[&[0, 2, 2]] == 11);
    assert!(a[&[0, 2, 3]] == 12);
    assert!(a[&[1, 0, 0]] == 13);
    assert!(a[&[1, 0, 1]] == 14);
    assert!(a[&[1, 0, 2]] == 15);
    assert!(a[&[1, 0, 3]] == 16);
    assert!(a[&[1, 1, 0]] == 17);
    assert!(a[&[1, 1, 1]] == 18);
    assert!(a[&[1, 1, 2]] == 19);
    assert!(a[&[1, 1, 3]] == 20);
    assert!(a[&[1, 2, 0]] == 21);
    assert!(a[&[1, 2, 1]] == 22);
    assert!(a[&[1, 2, 2]] == 23);
    assert!(a[&[1, 2, 3]] == 24);
}

#[test]
fn test_array_indexing_mut() {
    let mut a = Array::from_vec(vec![3], vec![1, 2, 3]);
    a[&[1]] = 42;
    assert!(a[&[1]] == 42);
}
<|endoftext|>


<commit_before>//! Transports for communicating with the docker daemon

extern crate hyper;

use self::hyper::buffer::BufReader;
use self::hyper::header::ContentType;
use self::hyper::status::StatusCode;
use self::super::{Error, Result};
use hyper::Client;
use hyper::client::Body;
use hyper::client::response::Response;
use hyper::header;
use hyper::method::Method;
use hyper::mime;
use hyperlocal::DomainUrl;
use rustc_serialize::json;
use std::fmt;
use std::io::Read;

pub fn tar() -> ContentType {
    ContentType(mime::Mime(
        mime::TopLevel::Application,
        mime::SubLevel::Ext(String::from("tar")),
        vec![],
    ))
}

/// Transports are types which define the means of communication
/// with the docker daemon
pub enum Transport {
    /// A network tcp interface
    Tcp { client: Client, host: String },
    /// A Unix domain socket
    Unix { client: Client, path: String },
}

impl fmt::Debug for Transport {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Transport::Tcp { ref host, .. } => write!(f, "Tcp({})", host),
            Transport::Unix { ref path, .. } => write!(f, "Unix({})", path),
        }
    }
}

impl Transport {
    pub fn request<'a, B>(
        &'a self,
        method: Method,
        endpoint: &str,
        body: Option<(B, ContentType)>,
    ) -> Result<String>
    where
        B: Into<Body<'a>>,
    {
        let mut res = self.stream(method, endpoint, body)?;
        let mut body = String::new();
        res.read_to_string(&mut body)?;
        debug!("{} raw response: {}", endpoint, body);
        Ok(body)
    }

    pub fn stream<'c, B>(
        &'c self,
        method: Method,
        endpoint: &str,
        body: Option<(B, ContentType)>,
    ) -> Result<Box<Read>>
    where
        B: Into<Body<'c>>,
    {
        let headers = {
            let mut headers = header::Headers::new();
            headers.set(header::Host {
                hostname: "".to_owned(),
                port: None,
            });
            headers
        };
        let req = match *self {
            Transport::Tcp {
                ref client,
                ref host,
            } => client.request(method, &format!("{}{}", host, endpoint)[..]),
            Transport::Unix {
                ref client,
                ref path,
            } => client.request(method, DomainUrl::new(&path, endpoint)),
        }.headers(headers);

        let embodied = match body {
            Some((b, c)) => req.header(c).body(b),
            _ => req,
        };
        let mut res = embodied.send()?;
        match res.status {
            StatusCode::Ok |
            StatusCode::Created |
            StatusCode::SwitchingProtocols => Ok(Box::new(res)),
            StatusCode::NoContent => Ok(
                Box::new(BufReader::new("".as_bytes())),
            ),
            // todo: constantize these
            StatusCode::BadRequest => {
                Err(Error::Fault {
                    code: res.status,
                    message: get_error_message(&mut res).unwrap_or(
                        "bad parameter"
                            .to_owned(),
                    ),
                })
            }
            StatusCode::NotFound => {
                Err(Error::Fault {
                    code: res.status,
                    message: get_error_message(&mut res).unwrap_or(
                        "not found".to_owned(),
                    ),
                })
            }
            StatusCode::NotAcceptable => {
                Err(Error::Fault {
                    code: res.status,
                    message: get_error_message(&mut res).unwrap_or(
                        "not acceptable"
                            .to_owned(),
                    ),
                })
            }
            StatusCode::Conflict => {
                Err(Error::Fault {
                    code: res.status,
                    message: get_error_message(&mut res).unwrap_or(
                        "conflict found"
                            .to_owned(),
                    ),
                })
            }
            StatusCode::InternalServerError => {
                Err(Error::Fault {
                    code: res.status,
                    message: get_error_message(&mut res).unwrap_or(
                        "internal server error"
                            .to_owned(),
                    ),
                })
            }
            _ => unreachable!(),
        }
    }
}

/// Extract the error message content from an HTTP response that
/// contains a Docker JSON error structure.
fn get_error_message(res: &mut Response) -> Option<String> {
    let mut output = String::new();
    if res.read_to_string(&mut output).is_ok() {
        let json_response = json::Json::from_str(output.as_str()).ok();
        let message = json_response
            .as_ref()
            .and_then(|x| x.as_object())
            .and_then(|x| x.get("message"))
            .and_then(|x| x.as_string())
            .map(|x| x.to_owned());

        message
    } else {
        None
    }
}
<commit_msg>Add support for the 304 Not Modified HTTP status code<commit_after>//! Transports for communicating with the docker daemon

extern crate hyper;

use self::hyper::buffer::BufReader;
use self::hyper::header::ContentType;
use self::hyper::status::StatusCode;
use self::super::{Error, Result};
use hyper::Client;
use hyper::client::Body;
use hyper::client::response::Response;
use hyper::header;
use hyper::method::Method;
use hyper::mime;
use hyperlocal::DomainUrl;
use rustc_serialize::json;
use std::fmt;
use std::io::Read;

pub fn tar() -> ContentType {
    ContentType(mime::Mime(
        mime::TopLevel::Application,
        mime::SubLevel::Ext(String::from("tar")),
        vec![],
    ))
}

/// Transports are types which define the means of communication
/// with the docker daemon
pub enum Transport {
    /// A network tcp interface
    Tcp { client: Client, host: String },
    /// A Unix domain socket
    Unix { client: Client, path: String },
}

impl fmt::Debug for Transport {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Transport::Tcp { ref host, .. } => write!(f, "Tcp({})", host),
            Transport::Unix { ref path, .. } => write!(f, "Unix({})", path),
        }
    }
}

impl Transport {
    pub fn request<'a, B>(
        &'a self,
        method: Method,
        endpoint: &str,
        body: Option<(B, ContentType)>,
    ) -> Result<String>
    where
        B: Into<Body<'a>>,
    {
        let mut res = self.stream(method, endpoint, body)?;
        let mut body = String::new();
        res.read_to_string(&mut body)?;
        debug!("{} raw response: {}", endpoint, body);
        Ok(body)
    }

    pub fn stream<'c, B>(
        &'c self,
        method: Method,
        endpoint: &str,
        body: Option<(B, ContentType)>,
    ) -> Result<Box<Read>>
    where
        B: Into<Body<'c>>,
    {
        let headers = {
            let mut headers = header::Headers::new();
            headers.set(header::Host {
                hostname: "".to_owned(),
                port: None,
            });
            headers
        };
        let req = match *self {
            Transport::Tcp {
                ref client,
                ref host,
            } => client.request(method, &format!("{}{}", host, endpoint)[..]),
            Transport::Unix {
                ref client,
                ref path,
            } => client.request(method, DomainUrl::new(&path, endpoint)),
        }.headers(headers);

        let embodied = match body {
            Some((b, c)) => req.header(c).body(b),
            _ => req,
        };
        let mut res = embodied.send()?;
        match res.status {
            StatusCode::Ok |
            StatusCode::Created |
            StatusCode::SwitchingProtocols => Ok(Box::new(res)),
            StatusCode::NoContent => Ok(
                Box::new(BufReader::new("".as_bytes())),
            ),
            // todo: constantize these
            StatusCode::BadRequest => {
                Err(Error::Fault {
                    code: res.status,
                    message: get_error_message(&mut res).unwrap_or(
                        "bad parameter"
                            .to_owned(),
                    ),
                })
            }
            StatusCode::NotFound => {
                Err(Error::Fault {
                    code: res.status,
                    message: get_error_message(&mut res).unwrap_or(
                        "not found".to_owned(),
                    ),
                })
            }
            StatusCode::NotModified => {
                Err(Error::Fault {
                    code: res.status,
                    message: get_error_message(&mut res).unwrap_or(
                        "not modified"
                            .to_owned(),
                    ),
                })
            },
            StatusCode::NotAcceptable => {
                Err(Error::Fault {
                    code: res.status,
                    message: get_error_message(&mut res).unwrap_or(
                        "not acceptable"
                            .to_owned(),
                    ),
                })
            }
            StatusCode::Conflict => {
                Err(Error::Fault {
                    code: res.status,
                    message: get_error_message(&mut res).unwrap_or(
                        "conflict found"
                            .to_owned(),
                    ),
                })
            }
            StatusCode::InternalServerError => {
                Err(Error::Fault {
                    code: res.status,
                    message: get_error_message(&mut res).unwrap_or(
                        "internal server error"
                            .to_owned(),
                    ),
                })
            }
            _ => unreachable!(),
        }
    }
}

/// Extract the error message content from an HTTP response that
/// contains a Docker JSON error structure.
fn get_error_message(res: &mut Response) -> Option<String> {
    let mut output = String::new();
    if res.read_to_string(&mut output).is_ok() {
        let json_response = json::Json::from_str(output.as_str()).ok();
        let message = json_response
            .as_ref()
            .and_then(|x| x.as_object())
            .and_then(|x| x.get("message"))
            .and_then(|x| x.as_string())
            .map(|x| x.to_owned());

        message
    } else {
        None
    }
}
<|endoftext|>


<commit_before>use std::{
    env,
    ffi::OsStr,
    fs::{self, File},
    io::{BufReader, Cursor, Read, Seek},
    path::{Path, PathBuf},
};

use app_dirs::{get_app_root, AppDataType};
use log::debug;
use reqwest::{blocking::Client, Proxy};
use std::time::{Duration, SystemTime};
use walkdir::{DirEntry, WalkDir};
use zip::ZipArchive;

use crate::{
    error::TealdeerError::{self, CacheError, UpdateError},
    types::{PathSource, PlatformType},
};

static CACHE_DIR_ENV_VAR: &str = "TEALDEER_CACHE_DIR";

pub static TLDR_PAGES_DIR: &str = "tldr-pages";
static TLDR_OLD_PAGES_DIR: &str = "tldr-master";

#[derive(Debug)]
pub struct Cache {
    url: String,
    platform: PlatformType,
}

#[derive(Debug)]
pub struct PageLookupResult {
    pub page_path: PathBuf,
    pub patch_path: Option<PathBuf>,
}

impl PageLookupResult {
    pub fn with_page(page_path: PathBuf) -> Self {
        Self {
            page_path,
            patch_path: None,
        }
    }

    pub fn with_optional_patch(mut self, patch_path: Option<PathBuf>) -> Self {
        self.patch_path = patch_path;
        self
    }

    /// Create a buffered reader that sequentially reads from the page and the
    /// patch, as if they were concatenated.
    ///
    /// This will return an error if either the page file or the patch file
    /// cannot be opened.
    pub fn reader(&self) -> Result<BufReader<Box<dyn Read>>, String> {
        // Open page file
        let page_file = File::open(&self.page_path)
            .map_err(|msg| format!("Could not open page file at {:?}: {}", self.page_path, msg))?;

        // Open patch file
        let patch_file_opt = match &self.patch_path {
            Some(path) => Some(
                File::open(path)
                    .map_err(|msg| format!("Could not open patch file at {:?}: {}", path, msg))?,
            ),
            None => None,
        };

        // Create chained reader from file(s)
        //
        // Note: It might be worthwhile to create our own struct that accepts
        // the page and patch files and that will read them sequentially,
        // because it avoids the boxing below. However, the performance impact
        // would first need to be shown to be significant using a benchmark.
        Ok(BufReader::new(if let Some(patch_file) = patch_file_opt {
            Box::new(page_file.chain(&b"\n"[..]).chain(patch_file)) as Box<dyn Read>
        } else {
            Box::new(page_file) as Box<dyn Read>
        }))
    }
}

pub enum CacheFreshness {
    /// The cache is still fresh (less than MAX_CACHE_AGE old)
    Fresh,
    /// The cache is stale and should be updated
    Stale(Duration),
    /// The cache is missing
    Missing,
}

impl Cache {
    pub fn new<S>(url: S, platform: PlatformType) -> Self
    where
        S: Into<String>,
    {
        Self {
            url: url.into(),
            platform,
        }
    }

    /// Return the path to the cache directory.
    pub fn get_cache_dir() -> Result<(PathBuf, PathSource), TealdeerError> {
        // Allow overriding the cache directory by setting the env variable.
        if let Ok(value) = env::var(CACHE_DIR_ENV_VAR) {
            let path = PathBuf::from(value);
            let (path_exists, path_is_dir) = path
                .metadata()
                .map_or((false, false), |md| (true, md.is_dir()));
            if path_exists && !path_is_dir {
                return Err(CacheError(format!(
                    "Path specified by ${} is not a directory.",
                    CACHE_DIR_ENV_VAR
                )));
            }
            if !path_exists {
                // Try to create the complete directory path.
                fs::create_dir_all(&path).map_err(|_| {
                    CacheError(format!(
                        "Directory path specified by ${} cannot be created.",
                        CACHE_DIR_ENV_VAR
                    ))
                })?;
                eprintln!(
                    "Successfully created cache directory path `{}`.",
                    path.to_str().unwrap()
                );
            }
            return Ok((path, PathSource::EnvVar));
        };

        // Otherwise, fall back to user cache directory.
        match get_app_root(AppDataType::UserCache, &crate::APP_INFO) {
            Ok(dirs) => Ok((dirs, PathSource::OsConvention)),
            Err(_) => Err(CacheError(
                "Could not determine user cache directory.".into(),
            )),
        }
    }

    /// Download the archive
    fn download(&self) -> Result<Vec<u8>, TealdeerError> {
        let mut builder = Client::builder();
        if let Ok(ref host) = env::var("HTTP_PROXY") {
            if let Ok(proxy) = Proxy::http(host) {
                builder = builder.proxy(proxy);
            }
        }
        if let Ok(ref host) = env::var("HTTPS_PROXY") {
            if let Ok(proxy) = Proxy::https(host) {
                builder = builder.proxy(proxy);
            }
        }
        let client = builder.build().unwrap_or_else(|_| Client::new());
        let mut resp = client.get(&self.url).send()?;
        let mut buf: Vec<u8> = vec![];
        let bytes_downloaded = resp.copy_to(&mut buf)?;
        debug!("{} bytes downloaded", bytes_downloaded);
        Ok(buf)
    }

    /// Decompress and open the archive
    fn decompress<R: Read + Seek>(reader: R) -> ZipArchive<R> {
        ZipArchive::new(reader).unwrap()
    }

    /// Update the pages cache.
    pub fn update(&self) -> Result<(), TealdeerError> {
        // First, download the compressed data
        let bytes: Vec<u8> = self.download()?;

        // Decompress the response body into an `Archive`
        let mut archive = Self::decompress(Cursor::new(bytes));

        // Determine paths
        let (cache_dir, _) = Self::get_cache_dir()?;
        let pages_dir = cache_dir.join(TLDR_PAGES_DIR);

        // Make sure that cache directory exists
        debug!("Ensure cache directory {:?} exists", &cache_dir);
        fs::create_dir_all(&cache_dir)
            .map_err(|e| UpdateError(format!("Could not create cache directory: {}", e)))?;

        // Clear cache directory
        // Note: This is not the best solution. Ideally we would download the
        // archive to a temporary directory and then swap the two directories.
        // But renaming a directory doesn't work across filesystems and Rust
        // does not yet offer a recursive directory copying function. So for
        // now, we'll use this approach.
        Self::clear()?;

        // Extract archive
        archive
            .extract(&pages_dir)
            .map_err(|e| UpdateError(format!("Could not unpack compressed data: {}", e)))?;

        Ok(())
    }

    /// Return the duration since the cache directory was last modified.
    pub fn last_update() -> Option<Duration> {
        if let Ok((cache_dir, _)) = Self::get_cache_dir() {
            if let Ok(metadata) = fs::metadata(cache_dir.join(TLDR_PAGES_DIR)) {
                if let Ok(mtime) = metadata.modified() {
                    let now = SystemTime::now();
                    return now.duration_since(mtime).ok();
                };
            };
        };
        None
    }

    /// Return the freshness of the cache (fresh, stale or missing).
    pub fn freshness() -> CacheFreshness {
        match Cache::last_update() {
            Some(ago) if ago > crate::config::MAX_CACHE_AGE => CacheFreshness::Stale(ago),
            Some(_) => CacheFreshness::Fresh,
            None => CacheFreshness::Missing,
        }
    }

    /// Return the platform directory.
    fn get_platform_dir(&self) -> &'static str {
        match self.platform {
            PlatformType::Linux => "linux",
            PlatformType::OsX => "osx",
            PlatformType::SunOs => "sunos",
            PlatformType::Windows => "windows",
        }
    }

    /// Check for pages for a given platform in one of the given languages.
    fn find_page_for_platform(
        page_name: &str,
        cache_dir: &Path,
        platform: &str,
        language_dirs: &[String],
    ) -> Option<PathBuf> {
        language_dirs
            .iter()
            .map(|lang_dir| cache_dir.join(lang_dir).join(platform).join(page_name))
            .find(|path| path.exists() && path.is_file())
    }

    /// Look up custom patch (<name>.patch). If it exists, store it in a variable.
    fn find_patch(patch_name: &str, custom_pages_dir: Option<&Path>) -> Option<PathBuf> {
        custom_pages_dir
            .map(|custom_dir| custom_dir.join(patch_name))
            .filter(|path| path.exists() && path.is_file())
    }

    /// Search for a page and return the path to it.
    pub fn find_page(
        &self,
        name: &str,
        languages: &[String],
        custom_pages_dir: Option<&Path>,
    ) -> Option<PageLookupResult> {
        let page_filename = format!("{}.md", name);
        let patch_filename = format!("{}.patch", name);
        let custom_filename = format!("{}.page", name);

        // Get cache dir
        let cache_dir = match Self::get_cache_dir() {
            Ok((cache_dir, _)) => cache_dir.join(TLDR_PAGES_DIR),
            Err(e) => {
                log::error!("Could not get cache directory: {}", e);
                return None;
            }
        };

        let lang_dirs: Vec<String> = languages
            .iter()
            .map(|lang| {
                if lang == "en" {
                    String::from("pages")
                } else {
                    format!("pages.{}", lang)
                }
            })
            .collect();

        // Look up custom page (<name>.page). If it exists, return it directly
        if let Some(config_dir) = custom_pages_dir {
            let custom_page = config_dir.join(custom_filename);
            if custom_page.exists() && custom_page.is_file() {
                return Some(PageLookupResult::with_page(custom_page));
            }
        }

        let patch_path = Self::find_patch(&patch_filename, custom_pages_dir);

        // Try to find a platform specific path next, append custom patch to it.
        let platform_dir = self.get_platform_dir();
        if let Some(page) =
            Self::find_page_for_platform(&page_filename, &cache_dir, platform_dir, &lang_dirs)
        {
            return Some(PageLookupResult::with_page(page).with_optional_patch(patch_path));
        }

        // Did not find platform specific results, fall back to "common"
        Self::find_page_for_platform(&page_filename, &cache_dir, "common", &lang_dirs)
            .map(|page| PageLookupResult::with_page(page).with_optional_patch(patch_path))
    }

    /// Return the available pages.
    pub fn list_pages(&self) -> Result<Vec<String>, TealdeerError> {
        // Determine platforms directory and platform
        let (cache_dir, _) = Self::get_cache_dir()?;
        let platforms_dir = cache_dir.join(TLDR_PAGES_DIR).join("pages");
        let platform_dir = self.get_platform_dir();

        // Closure that allows the WalkDir instance to traverse platform
        // specific and common page directories, but not others.
        let should_walk = |entry: &DirEntry| -> bool {
            let file_type = entry.file_type();
            let file_name = match entry.file_name().to_str() {
                Some(name) => name,
                None => return false,
            };
            if file_type.is_dir() {
                return file_name == "common" || file_name == platform_dir;
            } else if file_type.is_file() {
                return true;
            }
            false
        };

        // Recursively walk through common and (if applicable) platform specific directory
        let mut pages = WalkDir::new(platforms_dir)
            .min_depth(1) // Skip root directory
            .into_iter()
            .filter_entry(|e| should_walk(e)) // Filter out pages for other architectures
            .filter_map(Result::ok) // Convert results to options, filter out errors
            .filter_map(|e| {
                let path = e.path();
                let extension = &path.extension().and_then(OsStr::to_str).unwrap_or("");
                if e.file_type().is_file() && extension == &"md" {
                    path.file_stem()
                        .and_then(|stem| stem.to_str().map(Into::into))
                } else {
                    None
                }
            })
            .collect::<Vec<String>>();
        pages.sort();
        pages.dedup();
        Ok(pages)
    }

    /// Delete the cache directory.
    pub fn clear() -> Result<(), TealdeerError> {
        let (path, _) = Self::get_cache_dir()?;
        if path.exists() && path.is_dir() {
            // Delete old tldr-pages cache location as well if present
            // TODO: To be removed in the future
            for pages_dir_name in [TLDR_PAGES_DIR, TLDR_OLD_PAGES_DIR] {
                let pages_dir = path.join(pages_dir_name);

                if pages_dir.exists() {
                    fs::remove_dir_all(&pages_dir).map_err(|e| {
                        CacheError(format!(
                            "Could not remove cache directory ({}): {}",
                            pages_dir.display(),
                            e
                        ))
                    })?;
                }
            }
        } else if path.exists() {
            return Err(CacheError(format!(
                "Cache path ({}) is not a directory.",
                path.display()
            )));
        } else {
            return Err(CacheError(format!(
                "Cache path ({}) does not exist.",
                path.display()
            )));
        };
        Ok(())
    }
}

/// Unit Tests for cache module
#[cfg(test)]
mod tests {
    use super::*;

    use std::{
        fs::File,
        io::{Read, Write},
    };

    #[test]
    fn test_reader_with_patch() {
        // Write test files
        let dir = tempfile::tempdir().unwrap();
        let page_path = dir.path().join("test.page");
        let patch_path = dir.path().join("test.patch");
        {
            let mut f1 = File::create(&page_path).unwrap();
            f1.write_all(b"Hello\n").unwrap();
            let mut f2 = File::create(&patch_path).unwrap();
            f2.write_all(b"World").unwrap();
        }

        // Create chained reader from lookup result
        let lr = PageLookupResult::with_page(page_path).with_optional_patch(Some(patch_path));
        let mut reader = lr.reader().unwrap();

        // Read into a Vec
        let mut buf = Vec::new();
        reader.read_to_end(&mut buf).unwrap();

        assert_eq!(&buf, b"Hello\n\nWorld");
    }

    #[test]
    fn test_reader_without_patch() {
        // Write test file
        let dir = tempfile::tempdir().unwrap();
        let page_path = dir.path().join("test.page");
        {
            let mut f = File::create(&page_path).unwrap();
            f.write_all(b"Hello\n").unwrap();
        }

        // Create chained reader from lookup result
        let lr = PageLookupResult::with_page(page_path);
        let mut reader = lr.reader().unwrap();

        // Read into a Vec
        let mut buf = Vec::new();
        reader.read_to_end(&mut buf).unwrap();

        assert_eq!(&buf, b"Hello\n");
    }
}
<commit_msg>Cache: Return error if HTTP client cannot be created (#247)<commit_after>use std::{
    env,
    ffi::OsStr,
    fs::{self, File},
    io::{BufReader, Cursor, Read, Seek},
    path::{Path, PathBuf},
};

use app_dirs::{get_app_root, AppDataType};
use log::debug;
use reqwest::{blocking::Client, Proxy};
use std::time::{Duration, SystemTime};
use walkdir::{DirEntry, WalkDir};
use zip::ZipArchive;

use crate::{
    error::TealdeerError::{self, CacheError, UpdateError},
    types::{PathSource, PlatformType},
};

static CACHE_DIR_ENV_VAR: &str = "TEALDEER_CACHE_DIR";

pub static TLDR_PAGES_DIR: &str = "tldr-pages";
static TLDR_OLD_PAGES_DIR: &str = "tldr-master";

#[derive(Debug)]
pub struct Cache {
    url: String,
    platform: PlatformType,
}

#[derive(Debug)]
pub struct PageLookupResult {
    pub page_path: PathBuf,
    pub patch_path: Option<PathBuf>,
}

impl PageLookupResult {
    pub fn with_page(page_path: PathBuf) -> Self {
        Self {
            page_path,
            patch_path: None,
        }
    }

    pub fn with_optional_patch(mut self, patch_path: Option<PathBuf>) -> Self {
        self.patch_path = patch_path;
        self
    }

    /// Create a buffered reader that sequentially reads from the page and the
    /// patch, as if they were concatenated.
    ///
    /// This will return an error if either the page file or the patch file
    /// cannot be opened.
    pub fn reader(&self) -> Result<BufReader<Box<dyn Read>>, String> {
        // Open page file
        let page_file = File::open(&self.page_path)
            .map_err(|msg| format!("Could not open page file at {:?}: {}", self.page_path, msg))?;

        // Open patch file
        let patch_file_opt = match &self.patch_path {
            Some(path) => Some(
                File::open(path)
                    .map_err(|msg| format!("Could not open patch file at {:?}: {}", path, msg))?,
            ),
            None => None,
        };

        // Create chained reader from file(s)
        //
        // Note: It might be worthwhile to create our own struct that accepts
        // the page and patch files and that will read them sequentially,
        // because it avoids the boxing below. However, the performance impact
        // would first need to be shown to be significant using a benchmark.
        Ok(BufReader::new(if let Some(patch_file) = patch_file_opt {
            Box::new(page_file.chain(&b"\n"[..]).chain(patch_file)) as Box<dyn Read>
        } else {
            Box::new(page_file) as Box<dyn Read>
        }))
    }
}

pub enum CacheFreshness {
    /// The cache is still fresh (less than MAX_CACHE_AGE old)
    Fresh,
    /// The cache is stale and should be updated
    Stale(Duration),
    /// The cache is missing
    Missing,
}

impl Cache {
    pub fn new<S>(url: S, platform: PlatformType) -> Self
    where
        S: Into<String>,
    {
        Self {
            url: url.into(),
            platform,
        }
    }

    /// Return the path to the cache directory.
    pub fn get_cache_dir() -> Result<(PathBuf, PathSource), TealdeerError> {
        // Allow overriding the cache directory by setting the env variable.
        if let Ok(value) = env::var(CACHE_DIR_ENV_VAR) {
            let path = PathBuf::from(value);
            let (path_exists, path_is_dir) = path
                .metadata()
                .map_or((false, false), |md| (true, md.is_dir()));
            if path_exists && !path_is_dir {
                return Err(CacheError(format!(
                    "Path specified by ${} is not a directory.",
                    CACHE_DIR_ENV_VAR
                )));
            }
            if !path_exists {
                // Try to create the complete directory path.
                fs::create_dir_all(&path).map_err(|_| {
                    CacheError(format!(
                        "Directory path specified by ${} cannot be created.",
                        CACHE_DIR_ENV_VAR
                    ))
                })?;
                eprintln!(
                    "Successfully created cache directory path `{}`.",
                    path.to_str().unwrap()
                );
            }
            return Ok((path, PathSource::EnvVar));
        };

        // Otherwise, fall back to user cache directory.
        match get_app_root(AppDataType::UserCache, &crate::APP_INFO) {
            Ok(dirs) => Ok((dirs, PathSource::OsConvention)),
            Err(_) => Err(CacheError(
                "Could not determine user cache directory.".into(),
            )),
        }
    }

    /// Download the archive
    fn download(&self) -> Result<Vec<u8>, TealdeerError> {
        let mut builder = Client::builder();
        if let Ok(ref host) = env::var("HTTP_PROXY") {
            if let Ok(proxy) = Proxy::http(host) {
                builder = builder.proxy(proxy);
            }
        }
        if let Ok(ref host) = env::var("HTTPS_PROXY") {
            if let Ok(proxy) = Proxy::https(host) {
                builder = builder.proxy(proxy);
            }
        }
        let client = builder.build().map_err(|e| {
            TealdeerError::UpdateError(format!("Could not instantiate HTTP client: {}", e))
        })?;
        let mut resp = client.get(&self.url).send()?;
        let mut buf: Vec<u8> = vec![];
        let bytes_downloaded = resp.copy_to(&mut buf)?;
        debug!("{} bytes downloaded", bytes_downloaded);
        Ok(buf)
    }

    /// Decompress and open the archive
    fn decompress<R: Read + Seek>(reader: R) -> ZipArchive<R> {
        ZipArchive::new(reader).unwrap()
    }

    /// Update the pages cache.
    pub fn update(&self) -> Result<(), TealdeerError> {
        // First, download the compressed data
        let bytes: Vec<u8> = self.download()?;

        // Decompress the response body into an `Archive`
        let mut archive = Self::decompress(Cursor::new(bytes));

        // Determine paths
        let (cache_dir, _) = Self::get_cache_dir()?;
        let pages_dir = cache_dir.join(TLDR_PAGES_DIR);

        // Make sure that cache directory exists
        debug!("Ensure cache directory {:?} exists", &cache_dir);
        fs::create_dir_all(&cache_dir)
            .map_err(|e| UpdateError(format!("Could not create cache directory: {}", e)))?;

        // Clear cache directory
        // Note: This is not the best solution. Ideally we would download the
        // archive to a temporary directory and then swap the two directories.
        // But renaming a directory doesn't work across filesystems and Rust
        // does not yet offer a recursive directory copying function. So for
        // now, we'll use this approach.
        Self::clear()?;

        // Extract archive
        archive
            .extract(&pages_dir)
            .map_err(|e| UpdateError(format!("Could not unpack compressed data: {}", e)))?;

        Ok(())
    }

    /// Return the duration since the cache directory was last modified.
    pub fn last_update() -> Option<Duration> {
        if let Ok((cache_dir, _)) = Self::get_cache_dir() {
            if let Ok(metadata) = fs::metadata(cache_dir.join(TLDR_PAGES_DIR)) {
                if let Ok(mtime) = metadata.modified() {
                    let now = SystemTime::now();
                    return now.duration_since(mtime).ok();
                };
            };
        };
        None
    }

    /// Return the freshness of the cache (fresh, stale or missing).
    pub fn freshness() -> CacheFreshness {
        match Cache::last_update() {
            Some(ago) if ago > crate::config::MAX_CACHE_AGE => CacheFreshness::Stale(ago),
            Some(_) => CacheFreshness::Fresh,
            None => CacheFreshness::Missing,
        }
    }

    /// Return the platform directory.
    fn get_platform_dir(&self) -> &'static str {
        match self.platform {
            PlatformType::Linux => "linux",
            PlatformType::OsX => "osx",
            PlatformType::SunOs => "sunos",
            PlatformType::Windows => "windows",
        }
    }

    /// Check for pages for a given platform in one of the given languages.
    fn find_page_for_platform(
        page_name: &str,
        cache_dir: &Path,
        platform: &str,
        language_dirs: &[String],
    ) -> Option<PathBuf> {
        language_dirs
            .iter()
            .map(|lang_dir| cache_dir.join(lang_dir).join(platform).join(page_name))
            .find(|path| path.exists() && path.is_file())
    }

    /// Look up custom patch (<name>.patch). If it exists, store it in a variable.
    fn find_patch(patch_name: &str, custom_pages_dir: Option<&Path>) -> Option<PathBuf> {
        custom_pages_dir
            .map(|custom_dir| custom_dir.join(patch_name))
            .filter(|path| path.exists() && path.is_file())
    }

    /// Search for a page and return the path to it.
    pub fn find_page(
        &self,
        name: &str,
        languages: &[String],
        custom_pages_dir: Option<&Path>,
    ) -> Option<PageLookupResult> {
        let page_filename = format!("{}.md", name);
        let patch_filename = format!("{}.patch", name);
        let custom_filename = format!("{}.page", name);

        // Get cache dir
        let cache_dir = match Self::get_cache_dir() {
            Ok((cache_dir, _)) => cache_dir.join(TLDR_PAGES_DIR),
            Err(e) => {
                log::error!("Could not get cache directory: {}", e);
                return None;
            }
        };

        let lang_dirs: Vec<String> = languages
            .iter()
            .map(|lang| {
                if lang == "en" {
                    String::from("pages")
                } else {
                    format!("pages.{}", lang)
                }
            })
            .collect();

        // Look up custom page (<name>.page). If it exists, return it directly
        if let Some(config_dir) = custom_pages_dir {
            let custom_page = config_dir.join(custom_filename);
            if custom_page.exists() && custom_page.is_file() {
                return Some(PageLookupResult::with_page(custom_page));
            }
        }

        let patch_path = Self::find_patch(&patch_filename, custom_pages_dir);

        // Try to find a platform specific path next, append custom patch to it.
        let platform_dir = self.get_platform_dir();
        if let Some(page) =
            Self::find_page_for_platform(&page_filename, &cache_dir, platform_dir, &lang_dirs)
        {
            return Some(PageLookupResult::with_page(page).with_optional_patch(patch_path));
        }

        // Did not find platform specific results, fall back to "common"
        Self::find_page_for_platform(&page_filename, &cache_dir, "common", &lang_dirs)
            .map(|page| PageLookupResult::with_page(page).with_optional_patch(patch_path))
    }

    /// Return the available pages.
    pub fn list_pages(&self) -> Result<Vec<String>, TealdeerError> {
        // Determine platforms directory and platform
        let (cache_dir, _) = Self::get_cache_dir()?;
        let platforms_dir = cache_dir.join(TLDR_PAGES_DIR).join("pages");
        let platform_dir = self.get_platform_dir();

        // Closure that allows the WalkDir instance to traverse platform
        // specific and common page directories, but not others.
        let should_walk = |entry: &DirEntry| -> bool {
            let file_type = entry.file_type();
            let file_name = match entry.file_name().to_str() {
                Some(name) => name,
                None => return false,
            };
            if file_type.is_dir() {
                return file_name == "common" || file_name == platform_dir;
            } else if file_type.is_file() {
                return true;
            }
            false
        };

        // Recursively walk through common and (if applicable) platform specific directory
        let mut pages = WalkDir::new(platforms_dir)
            .min_depth(1) // Skip root directory
            .into_iter()
            .filter_entry(|e| should_walk(e)) // Filter out pages for other architectures
            .filter_map(Result::ok) // Convert results to options, filter out errors
            .filter_map(|e| {
                let path = e.path();
                let extension = &path.extension().and_then(OsStr::to_str).unwrap_or("");
                if e.file_type().is_file() && extension == &"md" {
                    path.file_stem()
                        .and_then(|stem| stem.to_str().map(Into::into))
                } else {
                    None
                }
            })
            .collect::<Vec<String>>();
        pages.sort();
        pages.dedup();
        Ok(pages)
    }

    /// Delete the cache directory.
    pub fn clear() -> Result<(), TealdeerError> {
        let (path, _) = Self::get_cache_dir()?;
        if path.exists() && path.is_dir() {
            // Delete old tldr-pages cache location as well if present
            // TODO: To be removed in the future
            for pages_dir_name in [TLDR_PAGES_DIR, TLDR_OLD_PAGES_DIR] {
                let pages_dir = path.join(pages_dir_name);

                if pages_dir.exists() {
                    fs::remove_dir_all(&pages_dir).map_err(|e| {
                        CacheError(format!(
                            "Could not remove cache directory ({}): {}",
                            pages_dir.display(),
                            e
                        ))
                    })?;
                }
            }
        } else if path.exists() {
            return Err(CacheError(format!(
                "Cache path ({}) is not a directory.",
                path.display()
            )));
        } else {
            return Err(CacheError(format!(
                "Cache path ({}) does not exist.",
                path.display()
            )));
        };
        Ok(())
    }
}

/// Unit Tests for cache module
#[cfg(test)]
mod tests {
    use super::*;

    use std::{
        fs::File,
        io::{Read, Write},
    };

    #[test]
    fn test_reader_with_patch() {
        // Write test files
        let dir = tempfile::tempdir().unwrap();
        let page_path = dir.path().join("test.page");
        let patch_path = dir.path().join("test.patch");
        {
            let mut f1 = File::create(&page_path).unwrap();
            f1.write_all(b"Hello\n").unwrap();
            let mut f2 = File::create(&patch_path).unwrap();
            f2.write_all(b"World").unwrap();
        }

        // Create chained reader from lookup result
        let lr = PageLookupResult::with_page(page_path).with_optional_patch(Some(patch_path));
        let mut reader = lr.reader().unwrap();

        // Read into a Vec
        let mut buf = Vec::new();
        reader.read_to_end(&mut buf).unwrap();

        assert_eq!(&buf, b"Hello\n\nWorld");
    }

    #[test]
    fn test_reader_without_patch() {
        // Write test file
        let dir = tempfile::tempdir().unwrap();
        let page_path = dir.path().join("test.page");
        {
            let mut f = File::create(&page_path).unwrap();
            f.write_all(b"Hello\n").unwrap();
        }

        // Create chained reader from lookup result
        let lr = PageLookupResult::with_page(page_path);
        let mut reader = lr.reader().unwrap();

        // Read into a Vec
        let mut buf = Vec::new();
        reader.read_to_end(&mut buf).unwrap();

        assert_eq!(&buf, b"Hello\n");
    }
}
<|endoftext|>


<commit_before>use num::{PrimInt, Signed};
use num::traits::CheckedShr;
use std::mem;
use std::cmp::{min};

/// Right-shift an integer, but round to nearest, with ties rounding to even.
///
/// This minimizes error, and avoids a directional bias.
pub fn right_shift_nicely_rounded <T: PrimInt> (input: T, shift: usize)->T {
  let divisor = T::one() << shift;
  let mask = divisor - T::one();
  if (input & mask) << 1 == divisor {
    let shifted = input >> shift;
    shifted + (shifted & T::one())
  } else {
    (input + (divisor >> 1)) >> shift
  }
}

/// Approximately evaluate an integer polynomial at an input in the range [-0.5, 0.5].
///
/// The input is represented as an integer combined with a right-shift size.
/// Each coefficient, except the constant one, must obey coefficient.abs() <= T::max_value() >> shift.
/// The constant coefficient must obey coefficient.abs() <= T::max_value() - (T::max_value() >> shift),
/// i.e. be small enough that you can add one of the other coefficients to it without overflowing.
/// Given these conditions, the result is guaranteed to be strictly within 1 of the ideal result.
/// For instance, if the ideal result was 2.125, this function could return 2 or 3,
/// but if it was 2, this function can only return exactly 2.
pub fn evaluate_polynomial_at_small_input <Coefficient: Copy, T: PrimInt + Signed + From <Coefficient>> (coefficients: & [Coefficient], input: T, shift: usize)->T {
  let half = (T::one() << shift) >> 1;
  assert!(-half <= input && input <= half, "inputs to evaluate_polynomial_at_small_input must be in the range [-0.5, 0.5]");

  if coefficients.len () == 0 {return T::zero()}

  let mut result = T::zero();
  for coefficient in coefficients.iter().skip(1).rev() {
    result = right_shift_nicely_rounded ((result + (*coefficient).into())*input, shift);
  }
  result + coefficients [0].into()
}

/// Translate a polynomial horizontally, letting it take inputs relative to a given input instead of relative to 0.
///
/// This is equivalent to replacing its coefficients with its Taylor coefficients at the given input.
/// This is equivalent to translating the graph *left* by the given input amount.
///
/// For constant polynomials, this is a no-op.
///
/// Otherwise, to avoid overflow, each term, when evaluating at the given input,
/// must obey term.abs() <= T::max_value() >> (coefficients.len() - 1).
/// If this condition is broken, translate_polynomial() returns Err and makes no changes.
pub fn translate_polynomial <T: PrimInt + Signed + CheckedShr> (coefficients: &mut [T], input: T)->Result <(),()> {
  if coefficients.len() <= 1 {return Ok(());}
  if !safe_to_translate_polynomial_to (coefficients, input, |_| T::max_value()) {return Err (())}
  translate_polynomial_unchecked (coefficients, input) ;
  Ok (())
}

/// Same as translate_polynomial(), but assume no overflow will occur.
///
/// Useful in performance-critical code when you already know there won't be overflow,
/// such as in the range returned by conservative_safe_polynomial_translation_range.
pub fn translate_polynomial_unchecked <T: PrimInt + Signed> (coefficients: &mut [T], input: T) {
  coefficients.reverse();
  for index in 0..coefficients.len() {
    let coefficient = mem::replace (&mut coefficients [index], T::zero());
    for derivative in (1..(index + 1)).rev() {
      coefficients [derivative] = coefficients [derivative]*input + coefficients [derivative - 1]
    }
    coefficients [0] = coefficients [0]*input + coefficient
  }
}

pub fn safe_to_translate_polynomial_to <T: PrimInt + Signed + CheckedShr, MaximumFn: Fn (usize)->T> (coefficients: & [T], input: T, maximum: MaximumFn)->bool {
  if coefficients.len() <= 1 {return true;}
  let mut factor = T::one();
  let input = input.abs();
  let sum_safety_shift = coefficients.len() as u32 - 1;
  let mut running_maximum = T::max_value();
  for (power, coefficient) in coefficients.iter().enumerate() {
    running_maximum = min (running_maximum, maximum (power).checked_shr (sum_safety_shift).unwrap_or (T::zero()));
    if coefficient == &T::min_value() {return false;}
    match coefficient.abs().checked_mul (&factor) {
      None => return false,
      Some (term) => if term > running_maximum {return false;},
    }
    if power + 1 < coefficients.len() { match factor.checked_mul (&input) {
      None => return false,
      Some (next_factor) => factor = next_factor,
    }}
  }
  true
}

/// Find the range of inputs to which the polynomial can safely be translated.
///
/// Magnitudes <= result are safe to translate to; magnitudes > result aren't.
/// If even the initial value is out of bounds, this function returns -1.
///
/// The third argument can impose a stricter maximum on the resulting coefficients.
pub fn exact_safe_polynomial_translation_range <T: PrimInt + Signed + CheckedShr, MaximumFn: Fn (usize)->T> (coefficients: & [T], maximum: MaximumFn)->T {
  if coefficients.len() <= 1 || coefficients [1..].iter().all (| coefficient | coefficient == &T::zero()) {return T::max_value()}
  // if we know that there's at least one nonzero non-constant coefficient, T::max_value is never legal, so we can pick a min that definitely legal and a max that's definitely not
  let mut min = -T::one();
  let mut max = T::max_value();
  while min + T::one() < max {
    let mid = (min >> 1u32) + (max >> 1u32) + (min & max & T::one());
    if safe_to_translate_polynomial_to (coefficients, mid, & maximum) {
      min = mid;
    } else {
      max = mid;
    }
  }
  min
}

/// Find a range of inputs to which the polynomial can safely be translated.
///
/// Magnitudes <= result are safe to translate to.
/// Magnitudes >= result*2 aren't.
/// If even the initial value is out of bounds, this function returns -1.
///
/// The third argument can impose a stricter maximum on the resulting coefficients.
///
/// This function is much faster than exact_safe_polynomial_translation_range, at the cost of being less precise.
pub fn conservative_safe_polynomial_translation_range <T: PrimInt + Signed + CheckedShr, MaximumFn: Fn (usize)->T> (coefficients: &[T], maximum: MaximumFn)->T {
  if coefficients.len() <= 1 {return T::max_value()}
  let sum_safety_shift = coefficients.len() as u32 - 1;
  let mut running_maximum = T::max_value();
  let mut result_shift = mem::size_of::<T>()*8 - 2;
  for (power, coefficient) in coefficients.iter().enumerate() {
    let magnitude = coefficient.abs();
    running_maximum = min (running_maximum, maximum (power).checked_shr (sum_safety_shift).unwrap_or (T::zero()));
    while magnitude > running_maximum.checked_shr((result_shift*power) as u32).unwrap_or (T::zero()) {
      if result_shift == 0 {return -T::one();}
      result_shift -= 1;
    }
  }
  T::one() << result_shift
}

#[cfg (test)]
mod tests {
  use super::*;
  use num::{Zero, One, FromPrimitive};
  use num::bigint::BigInt;
  use num::rational::{Ratio, BigRational};
  use rand::distributions::Distribution;
  use rand::distributions::range::{Range};
  use rand::{Rng, SeedableRng};

  fn evaluate_polynomial_exactly <Coefficient: Copy, T: PrimInt + Signed + From <Coefficient>> (coefficients: & [Coefficient], input: T, shift: usize)->BigRational
    where BigInt: From <Coefficient> + From <T> {
    let mut result: BigRational = Ratio::zero();
    let input = Ratio::new(BigInt::from (input), BigInt::one() << shift);
    for coefficient in coefficients.iter().rev() {
      result = result*&input + BigInt::from(*coefficient);
    }
    result
  }

  #[test]
  fn test_right_shift_nicely_rounded() {
    let inputs: Vec<(i64, usize, i64)> = vec![
      (0, 0, 0), (0, 5, 0), (1, 3, 0), (4, 3, 0), (5, 3, 1),
      (999, 1, 500), (998, 1, 499), (997, 1, 498)
    ];
    for (input, shift, result) in inputs {
      assert_eq!(right_shift_nicely_rounded (input, shift), result);
      assert_eq!(right_shift_nicely_rounded (-input, shift), -result);
    }
  }

  #[test]
  fn test_evaluate_polynomial_at_small_input() {
    let mut generator =::rand::chacha::ChaChaRng::from_seed ([33; 32]) ;
    for shift in 0..5 {
      let shift = 1 << shift;
      let half = 1i64 << (shift - 1);
      let maximum = i64::max_value() >> shift;
      let constant_maximum = i64::max_value() - maximum;
      let range = Range::new_inclusive (- maximum, maximum) ;
      let constant_range = Range::new_inclusive (- constant_maximum, constant_maximum);
      let mut test_inputs = vec![- half, - half + 1, - 1, 0, 1, half - 1, half];
      if half > 2 {for _ in 0..6 {test_inputs.push (generator.gen_range (- half + 2, half - 1));}}
      for _ in 0..40 {
        let mut coefficients = vec![constant_range.sample (&mut generator)];
        while generator.gen() {coefficients.push (range.sample (&mut generator));}
        for input in test_inputs.iter() {
          let result = evaluate_polynomial_at_small_input (& coefficients,*input, shift);

          let perfect_result = evaluate_polynomial_exactly (& coefficients,*input, shift);
          let difference = Ratio::from_integer (FromPrimitive::from_i64 (result).unwrap()) - perfect_result;
          assert!(difference < Ratio::from_integer (FromPrimitive::from_i64 (1).unwrap()));
        }
      }
    }
  }

  quickcheck! {
    fn quickcheck_exact_safe_translation_range_is_safe (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1;
      let range = exact_safe_polynomial_translation_range (& coefficients, maxima);
      println!( "{:?}", range);
      range < 0 || safe_to_translate_polynomial_to (& coefficients, range, maxima)
    }

    fn quickcheck_exact_safe_translation_range_is_maximal (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1;
      let range = exact_safe_polynomial_translation_range (& coefficients, maxima);
      println!( "{:?}", range);
      range == i64::max_value() || !safe_to_translate_polynomial_to (& coefficients, range + 1, maxima)
    }

    fn quickcheck_conservative_safe_translation_range_is_safe (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1;
      let exact_range = exact_safe_polynomial_translation_range (& coefficients, maxima);
      let conservative_range = conservative_safe_polynomial_translation_range (& coefficients, maxima);
      println!( "{:?}", (exact_range, conservative_range));
      conservative_range <= exact_range
    }

    fn quickcheck_conservative_safe_translation_range_is_within_half_maximal (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1;
      let exact_range = exact_safe_polynomial_translation_range (& coefficients, maxima);
      let conservative_range = conservative_safe_polynomial_translation_range (& coefficients, maxima);
      println!( "{:?}", (exact_range, conservative_range));
      conservative_range == exact_range || conservative_range > exact_range >> 1
    }
  }
}
<commit_msg>Fixed some special cases that were caught by quickcheck<commit_after>use num::{PrimInt, Signed};
use num::traits::CheckedShr;
use std::mem;
use std::cmp::{min};

/// Right-shift an integer, but round to nearest, with ties rounding to even.
///
/// This minimizes error, and avoids a directional bias.
pub fn right_shift_nicely_rounded <T: PrimInt> (input: T, shift: usize)->T {
  let divisor = T::one() << shift;
  let mask = divisor - T::one();
  if (input & mask) << 1 == divisor {
    let shifted = input >> shift;
    shifted + (shifted & T::one())
  } else {
    (input + (divisor >> 1)) >> shift
  }
}

/// Approximately evaluate an integer polynomial at an input in the range [-0.5, 0.5].
///
/// The input is represented as an integer combined with a right-shift size.
/// Each coefficient, except the constant one, must obey coefficient.abs() <= T::max_value() >> shift.
/// The constant coefficient must obey coefficient.abs() <= T::max_value() - (T::max_value() >> shift),
/// i.e. be small enough that you can add one of the other coefficients to it without overflowing.
/// Given these conditions, the result is guaranteed to be strictly within 1 of the ideal result.
/// For instance, if the ideal result was 2.125, this function could return 2 or 3,
/// but if it was 2, this function can only return exactly 2.
pub fn evaluate_polynomial_at_small_input <Coefficient: Copy, T: PrimInt + Signed + From <Coefficient>> (coefficients: & [Coefficient], input: T, shift: usize)->T {
  let half = (T::one() << shift) >> 1;
  assert!(-half <= input && input <= half, "inputs to evaluate_polynomial_at_small_input must be in the range [-0.5, 0.5]");

  if coefficients.len () == 0 {return T::zero()}

  let mut result = T::zero();
  for coefficient in coefficients.iter().skip(1).rev() {
    result = right_shift_nicely_rounded ((result + (*coefficient).into())*input, shift);
  }
  result + coefficients [0].into()
}

/// Translate a polynomial horizontally, letting it take inputs relative to a given input instead of relative to 0.
///
/// This is equivalent to replacing its coefficients with its Taylor coefficients at the given input.
/// This is equivalent to translating the graph *left* by the given input amount.
///
/// For constant polynomials, this is a no-op.
///
/// To avoid overflow, each term, when evaluating at the given input,
/// must obey term.abs() <= T::max_value() >> (coefficients.len() - 1).
/// If this condition is broken, translate_polynomial() returns Err and makes no changes.
pub fn translate_polynomial <T: PrimInt + Signed + CheckedShr> (coefficients: &mut [T], input: T)->Result <(),()> {
  if !safe_to_translate_polynomial_to (coefficients, input, |_| T::max_value()) {return Err (())}
  translate_polynomial_unchecked (coefficients, input) ;
  Ok (())
}

/// Same as translate_polynomial(), but assume no overflow will occur.
///
/// Useful in performance-critical code when you already know there won't be overflow,
/// such as in the range returned by conservative_safe_polynomial_translation_range.
pub fn translate_polynomial_unchecked <T: PrimInt + Signed> (coefficients: &mut [T], input: T) {
  coefficients.reverse();
  for index in 0..coefficients.len() {
    let coefficient = mem::replace (&mut coefficients [index], T::zero());
    for derivative in (1..(index + 1)).rev() {
      coefficients [derivative] = coefficients [derivative]*input + coefficients [derivative - 1]
    }
    coefficients [0] = coefficients [0]*input + coefficient
  }
}

pub fn safe_to_translate_polynomial_to <T: PrimInt + Signed + CheckedShr, MaximumFn: Fn (usize)->T> (coefficients: & [T], input: T, maximum: MaximumFn)->bool {
  if coefficients.len() == 0 {return true;}
  let mut factor = T::one();
  let input = input.abs();
  let sum_safety_shift = coefficients.len() as u32 - 1;
  let mut running_maximum = T::max_value();
  for (power, coefficient) in coefficients.iter().enumerate() {
    running_maximum = min (running_maximum, maximum (power).checked_shr (sum_safety_shift).unwrap_or (T::zero()));
    if coefficient == &T::min_value() {return false;}
    let magnitude = coefficient.abs();
    if magnitude > running_maximum {return false;}
    match magnitude.checked_mul (&factor) {
      None => return false,
      Some (term) => if term > running_maximum {return false;},
    }
    if power + 1 < coefficients.len() { match factor.checked_mul (&input) {
      None => return false,
      Some (next_factor) => factor = next_factor,
    }}
  }
  true
}

/// Find the range of inputs to which the polynomial can safely be translated.
///
/// Magnitudes <= result are safe to translate to; magnitudes > result aren't.
/// If even the initial value is out of bounds, this function returns -1.
///
/// The third argument can impose a stricter maximum on the resulting coefficients.
pub fn exact_safe_polynomial_translation_range <T: PrimInt + Signed + CheckedShr, MaximumFn: Fn (usize)->T> (coefficients: & [T], maximum: MaximumFn)->T {
  if !safe_to_translate_polynomial_to (coefficients, T::zero(), & maximum) {return -T::one();}
  // pick a min that's definitely legal and a max where max + 1 is definitely illegal
  let mut min = T::zero();
  let mut max = T::max_value();
  while min < max {
    let mid = (min >> 1u32) + (max >> 1u32) + ((min | max) & T::one());
    if safe_to_translate_polynomial_to (coefficients, mid, & maximum) {
      min = mid;
    } else {
      max = mid - T::one();
    }
  }
  min
}

/// Find a range of inputs to which the polynomial can safely be translated.
///
/// Magnitudes <= result are safe to translate to.
/// Magnitudes >= result*2 aren't.
/// If even the initial value is out of bounds, this function returns -1.
///
/// The third argument can impose a stricter maximum on the resulting coefficients.
///
/// This function is much faster than exact_safe_polynomial_translation_range, at the cost of being less precise.
pub fn conservative_safe_polynomial_translation_range <T: PrimInt + Signed + CheckedShr, MaximumFn: Fn (usize)->T> (coefficients: &[T], maximum: MaximumFn)->T {
  if coefficients.len() == 0 {return T::max_value()}
  let sum_safety_shift = coefficients.len() as u32 - 1;
  let mut running_maximum = T::max_value();
  let mut result_shift = mem::size_of::<T>()*8 - 2;
  for (power, coefficient) in coefficients.iter().enumerate() {
    let magnitude = coefficient.abs();
    running_maximum = min (running_maximum, maximum (power).checked_shr (sum_safety_shift).unwrap_or (T::zero()));
    while magnitude > running_maximum.checked_shr((result_shift*power) as u32).unwrap_or (T::zero()) {
      if result_shift == 0 {return -T::one();}
      result_shift -= 1;
    }
  }
  T::one() << result_shift
}

#[cfg (test)]
mod tests {
  use super::*;
  use num::{Zero, One, FromPrimitive};
  use num::bigint::BigInt;
  use num::rational::{Ratio, BigRational};
  use rand::distributions::Distribution;
  use rand::distributions::range::{Range};
  use rand::{Rng, SeedableRng};

  fn evaluate_polynomial_exactly <Coefficient: Copy, T: PrimInt + Signed + From <Coefficient>> (coefficients: & [Coefficient], input: T, shift: usize)->BigRational
    where BigInt: From <Coefficient> + From <T> {
    let mut result: BigRational = Ratio::zero();
    let input = Ratio::new(BigInt::from (input), BigInt::one() << shift);
    for coefficient in coefficients.iter().rev() {
      result = result*&input + BigInt::from(*coefficient);
    }
    result
  }

  #[test]
  fn test_right_shift_nicely_rounded() {
    let inputs: Vec<(i64, usize, i64)> = vec![
      (0, 0, 0), (0, 5, 0), (1, 3, 0), (4, 3, 0), (5, 3, 1),
      (999, 1, 500), (998, 1, 499), (997, 1, 498)
    ];
    for (input, shift, result) in inputs {
      assert_eq!(right_shift_nicely_rounded (input, shift), result);
      assert_eq!(right_shift_nicely_rounded (-input, shift), -result);
    }
  }

  #[test]
  fn test_evaluate_polynomial_at_small_input() {
    let mut generator =::rand::chacha::ChaChaRng::from_seed ([33; 32]) ;
    for shift in 0..5 {
      let shift = 1 << shift;
      let half = 1i64 << (shift - 1);
      let maximum = i64::max_value() >> shift;
      let constant_maximum = i64::max_value() - maximum;
      let range = Range::new_inclusive (- maximum, maximum) ;
      let constant_range = Range::new_inclusive (- constant_maximum, constant_maximum);
      let mut test_inputs = vec![- half, - half + 1, - 1, 0, 1, half - 1, half];
      if half > 2 {for _ in 0..6 {test_inputs.push (generator.gen_range (- half + 2, half - 1));}}
      for _ in 0..40 {
        let mut coefficients = vec![constant_range.sample (&mut generator)];
        while generator.gen() {coefficients.push (range.sample (&mut generator));}
        for input in test_inputs.iter() {
          let result = evaluate_polynomial_at_small_input (& coefficients,*input, shift);

          let perfect_result = evaluate_polynomial_exactly (& coefficients,*input, shift);
          let difference = Ratio::from_integer (FromPrimitive::from_i64 (result).unwrap()) - perfect_result;
          assert!(difference < Ratio::from_integer (FromPrimitive::from_i64 (1).unwrap()));
        }
      }
    }
  }

  quickcheck! {
    fn quickcheck_exact_safe_translation_range_is_safe (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1;
      let range = exact_safe_polynomial_translation_range (& coefficients, maxima);
      println!( "{:?}", range);
      range < 0 || safe_to_translate_polynomial_to (& coefficients, range, maxima)
    }

    fn quickcheck_exact_safe_translation_range_is_maximal (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1;
      let range = exact_safe_polynomial_translation_range (& coefficients, maxima);
      println!( "{:?}", range);
      range == i64::max_value() || !safe_to_translate_polynomial_to (& coefficients, range + 1, maxima)
    }

    fn quickcheck_conservative_safe_translation_range_is_safe (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1;
      let exact_range = exact_safe_polynomial_translation_range (& coefficients, maxima);
      let conservative_range = conservative_safe_polynomial_translation_range (& coefficients, maxima);
      println!( "{:?}", (exact_range, conservative_range));
      conservative_range <= exact_range
    }

    fn quickcheck_conservative_safe_translation_range_is_within_half_maximal (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1;
      let exact_range = exact_safe_polynomial_translation_range (& coefficients, maxima);
      let conservative_range = conservative_safe_polynomial_translation_range (& coefficients, maxima);
      println!( "{:?}", (exact_range, conservative_range));
      conservative_range == exact_range || conservative_range > exact_range >> 1
    }
  }
}
<|endoftext|>