<commit_before>use super::content::Content;
use super::encodings::{self, bytes_to_string, string_to_bytes};
use super::{Dictionary, Object, ObjectId};
use crate::xref::Xref;
use crate::{Error, Result};
use encoding::all::UTF_16BE;
use encoding::types::{DecoderTrap, EncoderTrap, Encoding};
use log::info;
use std::collections::BTreeMap;
use std::io::Write;
use std::str;

/// PDF document.
#[derive(Debug, Clone)]
pub struct Document {
	/// The version of the PDF specification to which the file conforms.
	pub version: String,

	/// The trailer gives the location of the cross-reference table and of certain special objects.
	pub trailer: Dictionary,

	/// The cross-reference table contains locations of the indirect objects.
	pub reference_table: Xref,

	/// The objects that make up the document contained in the file.
	pub objects: BTreeMap<ObjectId, Object>,

	/// Current maximum object id within the document.
	pub max_id: u32,
}

impl Document {
	/// Create new PDF document.
	pub fn new() -> Document {
		Document {
			version: "1.4".to_string(),
			trailer: Dictionary::new(),
			reference_table: Xref::new(0),
			objects: BTreeMap::new(),
			max_id: 0,
		}
	}

	const DEREF_LIMIT : usize = 128;

	/// Follow references if the supplied object is a reference.
	///
	/// Returns a tuple of an optional object id and final object.
	/// The object id will be None if the object was not a
	/// reference. Otherwise, it will be the last object id in the
	/// reference chain.
	pub fn dereference<'a>(&'a self, mut object: &'a Object) -> Result<(Option<ObjectId>, &'a Object)> {
		let mut nb_deref = 0;
		let mut id = None;

		while let Ok(ref_id) = object.as_reference() {
			id = Some(ref_id);
			object = self.objects.get(&ref_id).ok_or(Error::ObjectNotFound)?;

			nb_deref += 1;
			if nb_deref > Self::DEREF_LIMIT {
				return Err(Error::ReferenceLimit);
			}
		}

		Ok((id, object))
	}

	/// Get object by object id, will iteratively dereference a referenced object.
	pub fn get_object(&self, id: ObjectId) -> Result<&Object> {
		let object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
		self.dereference(object).map(|(_, object)| object)
	}

	/// Get mutable reference to object by object id, will iteratively dereference a referenced object.
	pub fn get_object_mut(&mut self, id: ObjectId) -> Result<&mut Object> {
		let object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
		let (ref_id, _) = self.dereference(object)?;

		Ok(self.objects.get_mut(&ref_id.unwrap_or(id)).unwrap())
	}

	/// Get dictionary object by id.
	pub fn get_dictionary(&self, id: ObjectId) -> Result<&Dictionary> {
		self.get_object(id).and_then(Object::as_dict)
	}

	/// Traverse objects from trailer recursively, return all referenced object IDs.
	pub fn traverse_objects<A: Fn(&mut Object) -> ()>(&mut self, action: A) -> Vec<ObjectId> {
		fn traverse_array<A: Fn(&mut Object) -> ()>(array: &mut Vec<Object>, action: &A, refs: &mut Vec<ObjectId>) {
			for item in array.iter_mut() {
				traverse_object(item, action, refs);
			}
		}
		fn traverse_dictionary<A: Fn(&mut Object) -> ()>(dict: &mut Dictionary, action: &A, refs: &mut Vec<ObjectId>) {
			for (_, v) in dict.iter_mut() {
				traverse_object(v, action, refs);
			}
		}
		fn traverse_object<A: Fn(&mut Object) -> ()>(object: &mut Object, action: &A, refs: &mut Vec<ObjectId>) {
			action(object);
			match *object {
				Object::Array(ref mut array) => traverse_array(array, action, refs),
				Object::Dictionary(ref mut dict) => traverse_dictionary(dict, action, refs),
				Object::Stream(ref mut stream) => traverse_dictionary(&mut stream.dict, action, refs),
				Object::Reference(id) => {
					if !refs.contains(&id) {
						refs.push(id);
					}
				}
				_ => {}
			}
		}
		let mut refs = vec![];
		traverse_dictionary(&mut self.trailer, &action, &mut refs);
		let mut index = 0;
		while index < refs.len() {
			if let Some(object) = self.objects.get_mut(&refs[index]) {
				traverse_object(object, &action, &mut refs);
			}
			index += 1;
		}
		refs
	}

	/// Get catalog dictionary.
	pub fn catalog(&self) -> Result<&Dictionary> {
		self.trailer.get(b"Root").and_then(Object::as_reference).and_then(|id| self.get_dictionary(id))
	}

	/// Get page numbers and corresponding object ids.
	pub fn get_pages(&self) -> BTreeMap<u32, ObjectId> {
		self.page_iter().enumerate().map(|(i, p)| ((i+1) as u32, p)).collect()
	}

	pub fn page_iter<'a>(&'a self) -> impl Iterator<Item=ObjectId> + 'a {
		PageTreeIter::new(self)
	}

	/// Get content stream object ids of a page.
	pub fn get_page_contents(&self, page_id: ObjectId) -> Vec<ObjectId> {
		let mut streams = vec![];
		if let Ok(page) = self.get_dictionary(page_id) {
			if let Ok(contents) = page.get(b"Contents") {
				match *contents {
					Object::Reference(ref id) => {
						streams.push(*id);
					}
					Object::Array(ref arr) => {
						for content in arr {
							if let Ok(id) = content.as_reference() { streams.push(id) }
						}
					}
					_ => {}
				}
			}
		}
		streams
	}

	/// Get content of a page.
	pub fn get_page_content(&self, page_id: ObjectId) -> Result<Vec<u8>> {
		let mut content = Vec::new();
		let content_streams = self.get_page_contents(page_id);
		for object_id in content_streams {
			if let Ok(content_stream) = self.get_object(object_id).and_then(Object::as_stream) {
				match content_stream.decompressed_content() {
					Ok(data) => content.write_all(&data)?,
					Err(_) => content.write_all(&content_stream.content)?,
				};
			}
		}
		Ok(content)
	}

	/// Get decoded page content;
	pub fn get_and_decode_page_content(&self, page_id: ObjectId) -> Result<Content> {
		let content_data = self.get_page_content(page_id)?;
		Content::decode(&content_data)
	}

	/// Get resources used by a page.
	pub fn get_page_resources(&self, page_id: ObjectId) -> (Option<&Dictionary>, Vec<ObjectId>) {
		fn collect_resources(page_node: &Dictionary, resource_ids: &mut Vec<ObjectId>, doc: &Document) {
			if let Ok(resources_id) = page_node.get(b"Resources").and_then(Object::as_reference) {
				resource_ids.push(resources_id);
			}
			if let Ok(page_tree) = page_node.get(b"Parent").and_then(Object::as_reference).and_then(|id| doc.get_dictionary(id)) {
				collect_resources(page_tree, resource_ids, doc);
			}
		};

		let mut resource_dict = None;
		let mut resource_ids = Vec::new();
		if let Ok(page) = self.get_dictionary(page_id) {
			resource_dict = page.get(b"Resources").and_then(Object::as_dict).ok();
			collect_resources(page, &mut resource_ids, self);
		}
		(resource_dict, resource_ids)
	}

	/// Get fonts used by a page.
	pub fn get_page_fonts(&self, page_id: ObjectId) -> BTreeMap<Vec<u8>, &Dictionary> {
		fn collect_fonts_from_resources<'a>(resources: &'a Dictionary, fonts: &mut BTreeMap<Vec<u8>, &'a Dictionary>, doc: &'a Document) {
			if let Ok(font_dict) = resources.get(b"Font").and_then(Object::as_dict) {
				for (name, value) in font_dict.iter() {
					let font = match *value {
						Object::Reference(id) => doc.get_dictionary(id).ok(),
						Object::Dictionary(ref dict) => Some(dict),
						_ => None,
					};
					if !fonts.contains_key(name) {
						font.map(|font| fonts.insert(name.clone(), font));
					}
				}
			}
		};

		let mut fonts = BTreeMap::new();
		let (resource_dict, resource_ids) = self.get_page_resources(page_id);
		if let Some(resources) = resource_dict {
			collect_fonts_from_resources(resources, &mut fonts, self);
		}
		for resource_id in resource_ids {
			if let Ok(resources) = self.get_dictionary(resource_id) {
				collect_fonts_from_resources(resources, &mut fonts, self);
			}
		}
		fonts
	}

	pub fn decode_text(encoding: Option<&str>, bytes: &[u8]) -> String {
		if let Some(encoding) = encoding {
			info!("{}", encoding);
			match encoding {
				"StandardEncoding" => bytes_to_string(encodings::STANDARD_ENCODING, bytes),
				"MacRomanEncoding" => bytes_to_string(encodings::MAC_ROMAN_ENCODING, bytes),
				"MacExpertEncoding" => bytes_to_string(encodings::MAC_EXPERT_ENCODING, bytes),
				"WinAnsiEncoding" => bytes_to_string(encodings::WIN_ANSI_ENCODING, bytes),
				"UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.decode(bytes, DecoderTrap::Ignore).unwrap(),
				"Identity-H" => "?Identity-H Unimplemented?".to_string(), // Unimplemented
				_ => String::from_utf8_lossy(bytes).to_string(),
			}
		} else {
			bytes_to_string(encodings::STANDARD_ENCODING, bytes)
		}
	}

	pub fn encode_text(encoding: Option<&str>, text: &str) -> Vec<u8> {
		if let Some(encoding) = encoding {
			match encoding {
				"StandardEncoding" => string_to_bytes(encodings::STANDARD_ENCODING, text),
				"MacRomanEncoding" => string_to_bytes(encodings::MAC_ROMAN_ENCODING, text),
				"MacExpertEncoding" => string_to_bytes(encodings::MAC_EXPERT_ENCODING, text),
				"WinAnsiEncoding" => string_to_bytes(encodings::WIN_ANSI_ENCODING, text),
				"UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.encode(text, EncoderTrap::Ignore).unwrap(),
				"Identity-H" => vec![], // Unimplemented
				_ => text.as_bytes().to_vec(),
			}
		} else {
			string_to_bytes(encodings::STANDARD_ENCODING, text)
		}
	}
}

impl Default for Document {
	fn default() -> Self {
		Self::new()
	}
}

struct PageTreeIter<'a> {
	doc: &'a Document,
	stack: Vec<&'a [Object]>,
	kids: Option<&'a [Object]>,
}

impl <'a> PageTreeIter<'a> {
	const PAGE_TREE_DEPTH_LIMIT: usize = 256;

	fn new(doc: &'a Document) -> Self {
		if let Ok(page_tree_id) = doc.catalog().and_then(|cat| cat.get(b"Pages")).and_then(Object::as_reference) {
			Self { doc, kids: Self::kids(doc, page_tree_id), stack: Vec::with_capacity(32) }
		} else {
			Self { doc, kids: None, stack: Vec::new() }
		}
	}

	fn kids(doc: &Document, page_tree_id: ObjectId) -> Option<&[Object]> {
		doc.get_dictionary(page_tree_id).and_then(|page_tree| page_tree.get(b"Kids")).and_then(Object::as_array).map(|k| k.as_slice()).ok()
	}
}

impl Iterator for PageTreeIter<'_> {
	type Item = ObjectId;

	fn next(&mut self) -> Option<Self::Item> {
		loop {
			while let Some((kid, new_kids)) = self.kids.and_then(|k| k.split_first()) {
				self.kids = Some(new_kids);

				if let Ok(kid_id) = kid.as_reference() {
					if let Ok(type_name) = self.doc.get_dictionary(kid_id).and_then(Dictionary::type_name) {
						match type_name {
							"Page" => {
								return Some(kid_id);
							},
							"Pages" => {
								if self.stack.len() < Self::PAGE_TREE_DEPTH_LIMIT {
									let kids = self.kids.unwrap();
									if !kids.is_empty() {
										self.stack.push(kids);
									}
									self.kids = Self::kids(self.doc, kid_id);
								}
							},
							_ => {}
						}
					}
				}
			}

			// Current level exhausted, try to pop.
			if let kids @ Some(_) = self.stack.pop() {
				self.kids = kids;
			} else {
				return None;
			}
		}
	}

	fn size_hint(&self) -> (usize, Option<usize>) {
		let stack_len: usize = self.stack.iter().map(|s| s.len()).sum();

		(self.kids.map(|k| k.len()).unwrap_or(0) + stack_len, None)
	}
}

impl std::iter::FusedIterator for PageTreeIter<'_> {}
<commit_msg>Use lifetime ellision.<commit_after>use super::content::Content;
use super::encodings::{self, bytes_to_string, string_to_bytes};
use super::{Dictionary, Object, ObjectId};
use crate::xref::Xref;
use crate::{Error, Result};
use encoding::all::UTF_16BE;
use encoding::types::{DecoderTrap, EncoderTrap, Encoding};
use log::info;
use std::collections::BTreeMap;
use std::io::Write;
use std::str;

/// PDF document.
#[derive(Debug, Clone)]
pub struct Document {
	/// The version of the PDF specification to which the file conforms.
	pub version: String,

	/// The trailer gives the location of the cross-reference table and of certain special objects.
	pub trailer: Dictionary,

	/// The cross-reference table contains locations of the indirect objects.
	pub reference_table: Xref,

	/// The objects that make up the document contained in the file.
	pub objects: BTreeMap<ObjectId, Object>,

	/// Current maximum object id within the document.
	pub max_id: u32,
}

impl Document {
	/// Create new PDF document.
	pub fn new() -> Document {
		Document {
			version: "1.4".to_string(),
			trailer: Dictionary::new(),
			reference_table: Xref::new(0),
			objects: BTreeMap::new(),
			max_id: 0,
		}
	}

	const DEREF_LIMIT : usize = 128;

	/// Follow references if the supplied object is a reference.
	///
	/// Returns a tuple of an optional object id and final object.
	/// The object id will be None if the object was not a
	/// reference. Otherwise, it will be the last object id in the
	/// reference chain.
	pub fn dereference<'a>(&'a self, mut object: &'a Object) -> Result<(Option<ObjectId>, &'a Object)> {
		let mut nb_deref = 0;
		let mut id = None;

		while let Ok(ref_id) = object.as_reference() {
			id = Some(ref_id);
			object = self.objects.get(&ref_id).ok_or(Error::ObjectNotFound)?;

			nb_deref += 1;
			if nb_deref > Self::DEREF_LIMIT {
				return Err(Error::ReferenceLimit);
			}
		}

		Ok((id, object))
	}

	/// Get object by object id, will iteratively dereference a referenced object.
	pub fn get_object(&self, id: ObjectId) -> Result<&Object> {
		let object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
		self.dereference(object).map(|(_, object)| object)
	}

	/// Get mutable reference to object by object id, will iteratively dereference a referenced object.
	pub fn get_object_mut(&mut self, id: ObjectId) -> Result<&mut Object> {
		let object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
		let (ref_id, _) = self.dereference(object)?;

		Ok(self.objects.get_mut(&ref_id.unwrap_or(id)).unwrap())
	}

	/// Get dictionary object by id.
	pub fn get_dictionary(&self, id: ObjectId) -> Result<&Dictionary> {
		self.get_object(id).and_then(Object::as_dict)
	}

	/// Traverse objects from trailer recursively, return all referenced object IDs.
	pub fn traverse_objects<A: Fn(&mut Object) -> ()>(&mut self, action: A) -> Vec<ObjectId> {
		fn traverse_array<A: Fn(&mut Object) -> ()>(array: &mut Vec<Object>, action: &A, refs: &mut Vec<ObjectId>) {
			for item in array.iter_mut() {
				traverse_object(item, action, refs);
			}
		}
		fn traverse_dictionary<A: Fn(&mut Object) -> ()>(dict: &mut Dictionary, action: &A, refs: &mut Vec<ObjectId>) {
			for (_, v) in dict.iter_mut() {
				traverse_object(v, action, refs);
			}
		}
		fn traverse_object<A: Fn(&mut Object) -> ()>(object: &mut Object, action: &A, refs: &mut Vec<ObjectId>) {
			action(object);
			match *object {
				Object::Array(ref mut array) => traverse_array(array, action, refs),
				Object::Dictionary(ref mut dict) => traverse_dictionary(dict, action, refs),
				Object::Stream(ref mut stream) => traverse_dictionary(&mut stream.dict, action, refs),
				Object::Reference(id) => {
					if !refs.contains(&id) {
						refs.push(id);
					}
				}
				_ => {}
			}
		}
		let mut refs = vec![];
		traverse_dictionary(&mut self.trailer, &action, &mut refs);
		let mut index = 0;
		while index < refs.len() {
			if let Some(object) = self.objects.get_mut(&refs[index]) {
				traverse_object(object, &action, &mut refs);
			}
			index += 1;
		}
		refs
	}

	/// Get catalog dictionary.
	pub fn catalog(&self) -> Result<&Dictionary> {
		self.trailer.get(b"Root").and_then(Object::as_reference).and_then(|id| self.get_dictionary(id))
	}

	/// Get page numbers and corresponding object ids.
	pub fn get_pages(&self) -> BTreeMap<u32, ObjectId> {
		self.page_iter().enumerate().map(|(i, p)| ((i+1) as u32, p)).collect()
	}

	pub fn page_iter(&self) -> impl Iterator<Item=ObjectId> + '_ {
		PageTreeIter::new(self)
	}

	/// Get content stream object ids of a page.
	pub fn get_page_contents(&self, page_id: ObjectId) -> Vec<ObjectId> {
		let mut streams = vec![];
		if let Ok(page) = self.get_dictionary(page_id) {
			if let Ok(contents) = page.get(b"Contents") {
				match *contents {
					Object::Reference(ref id) => {
						streams.push(*id);
					}
					Object::Array(ref arr) => {
						for content in arr {
							if let Ok(id) = content.as_reference() { streams.push(id) }
						}
					}
					_ => {}
				}
			}
		}
		streams
	}

	/// Get content of a page.
	pub fn get_page_content(&self, page_id: ObjectId) -> Result<Vec<u8>> {
		let mut content = Vec::new();
		let content_streams = self.get_page_contents(page_id);
		for object_id in content_streams {
			if let Ok(content_stream) = self.get_object(object_id).and_then(Object::as_stream) {
				match content_stream.decompressed_content() {
					Ok(data) => content.write_all(&data)?,
					Err(_) => content.write_all(&content_stream.content)?,
				};
			}
		}
		Ok(content)
	}

	/// Get decoded page content;
	pub fn get_and_decode_page_content(&self, page_id: ObjectId) -> Result<Content> {
		let content_data = self.get_page_content(page_id)?;
		Content::decode(&content_data)
	}

	/// Get resources used by a page.
	pub fn get_page_resources(&self, page_id: ObjectId) -> (Option<&Dictionary>, Vec<ObjectId>) {
		fn collect_resources(page_node: &Dictionary, resource_ids: &mut Vec<ObjectId>, doc: &Document) {
			if let Ok(resources_id) = page_node.get(b"Resources").and_then(Object::as_reference) {
				resource_ids.push(resources_id);
			}
			if let Ok(page_tree) = page_node.get(b"Parent").and_then(Object::as_reference).and_then(|id| doc.get_dictionary(id)) {
				collect_resources(page_tree, resource_ids, doc);
			}
		};

		let mut resource_dict = None;
		let mut resource_ids = Vec::new();
		if let Ok(page) = self.get_dictionary(page_id) {
			resource_dict = page.get(b"Resources").and_then(Object::as_dict).ok();
			collect_resources(page, &mut resource_ids, self);
		}
		(resource_dict, resource_ids)
	}

	/// Get fonts used by a page.
	pub fn get_page_fonts(&self, page_id: ObjectId) -> BTreeMap<Vec<u8>, &Dictionary> {
		fn collect_fonts_from_resources<'a>(resources: &'a Dictionary, fonts: &mut BTreeMap<Vec<u8>, &'a Dictionary>, doc: &'a Document) {
			if let Ok(font_dict) = resources.get(b"Font").and_then(Object::as_dict) {
				for (name, value) in font_dict.iter() {
					let font = match *value {
						Object::Reference(id) => doc.get_dictionary(id).ok(),
						Object::Dictionary(ref dict) => Some(dict),
						_ => None,
					};
					if !fonts.contains_key(name) {
						font.map(|font| fonts.insert(name.clone(), font));
					}
				}
			}
		};

		let mut fonts = BTreeMap::new();
		let (resource_dict, resource_ids) = self.get_page_resources(page_id);
		if let Some(resources) = resource_dict {
			collect_fonts_from_resources(resources, &mut fonts, self);
		}
		for resource_id in resource_ids {
			if let Ok(resources) = self.get_dictionary(resource_id) {
				collect_fonts_from_resources(resources, &mut fonts, self);
			}
		}
		fonts
	}

	pub fn decode_text(encoding: Option<&str>, bytes: &[u8]) -> String {
		if let Some(encoding) = encoding {
			info!("{}", encoding);
			match encoding {
				"StandardEncoding" => bytes_to_string(encodings::STANDARD_ENCODING, bytes),
				"MacRomanEncoding" => bytes_to_string(encodings::MAC_ROMAN_ENCODING, bytes),
				"MacExpertEncoding" => bytes_to_string(encodings::MAC_EXPERT_ENCODING, bytes),
				"WinAnsiEncoding" => bytes_to_string(encodings::WIN_ANSI_ENCODING, bytes),
				"UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.decode(bytes, DecoderTrap::Ignore).unwrap(),
				"Identity-H" => "?Identity-H Unimplemented?".to_string(), // Unimplemented
				_ => String::from_utf8_lossy(bytes).to_string(),
			}
		} else {
			bytes_to_string(encodings::STANDARD_ENCODING, bytes)
		}
	}

	pub fn encode_text(encoding: Option<&str>, text: &str) -> Vec<u8> {
		if let Some(encoding) = encoding {
			match encoding {
				"StandardEncoding" => string_to_bytes(encodings::STANDARD_ENCODING, text),
				"MacRomanEncoding" => string_to_bytes(encodings::MAC_ROMAN_ENCODING, text),
				"MacExpertEncoding" => string_to_bytes(encodings::MAC_EXPERT_ENCODING, text),
				"WinAnsiEncoding" => string_to_bytes(encodings::WIN_ANSI_ENCODING, text),
				"UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.encode(text, EncoderTrap::Ignore).unwrap(),
				"Identity-H" => vec![], // Unimplemented
				_ => text.as_bytes().to_vec(),
			}
		} else {
			string_to_bytes(encodings::STANDARD_ENCODING, text)
		}
	}
}

impl Default for Document {
	fn default() -> Self {
		Self::new()
	}
}

struct PageTreeIter<'a> {
	doc: &'a Document,
	stack: Vec<&'a [Object]>,
	kids: Option<&'a [Object]>,
}

impl <'a> PageTreeIter<'a> {
	const PAGE_TREE_DEPTH_LIMIT: usize = 256;

	fn new(doc: &'a Document) -> Self {
		if let Ok(page_tree_id) = doc.catalog().and_then(|cat| cat.get(b"Pages")).and_then(Object::as_reference) {
			Self { doc, kids: Self::kids(doc, page_tree_id), stack: Vec::with_capacity(32) }
		} else {
			Self { doc, kids: None, stack: Vec::new() }
		}
	}

	fn kids(doc: &Document, page_tree_id: ObjectId) -> Option<&[Object]> {
		doc.get_dictionary(page_tree_id).and_then(|page_tree| page_tree.get(b"Kids")).and_then(Object::as_array).map(|k| k.as_slice()).ok()
	}
}

impl Iterator for PageTreeIter<'_> {
	type Item = ObjectId;

	fn next(&mut self) -> Option<Self::Item> {
		loop {
			while let Some((kid, new_kids)) = self.kids.and_then(|k| k.split_first()) {
				self.kids = Some(new_kids);

				if let Ok(kid_id) = kid.as_reference() {
					if let Ok(type_name) = self.doc.get_dictionary(kid_id).and_then(Dictionary::type_name) {
						match type_name {
							"Page" => {
								return Some(kid_id);
							},
							"Pages" => {
								if self.stack.len() < Self::PAGE_TREE_DEPTH_LIMIT {
									let kids = self.kids.unwrap();
									if !kids.is_empty() {
										self.stack.push(kids);
									}
									self.kids = Self::kids(self.doc, kid_id);
								}
							},
							_ => {}
						}
					}
				}
			}

			// Current level exhausted, try to pop.
			if let kids @ Some(_) = self.stack.pop() {
				self.kids = kids;
			} else {
				return None;
			}
		}
	}

	fn size_hint(&self) -> (usize, Option<usize>) {
		let stack_len: usize = self.stack.iter().map(|s| s.len()).sum();

		(self.kids.map(|k| k.len()).unwrap_or(0) + stack_len, None)
	}
}

impl std::iter::FusedIterator for PageTreeIter<'_> {}
<|endoftext|>


<commit_before>//! Main interface for a stream of events the Supervisor can send out
//! in the course of its operations.
//!
//! Currently, the Supervisor is able to send events to a [NATS
//! Streaming][1] server. The `init_stream` function must be called
//! before sending events to initialize the publishing thread in the
//! background. Thereafter, you can pass "event" structs to the
//! `event` function, which will publish the event to the stream.
//!
//! All events are published under the "habitat" subject.
//!
//! [1]:https://github.com/nats-io/nats-streaming-server

mod error;

#[cfg(feature = "nitox_stream")]
#[path = "event/nitox.rs"]
mod stream_impl;

// This is the default, but setting it as a default in Cargo.toml
// makes it trickier to use nitox instead.
#[cfg(any(feature = "ratsio_stream", not(feature = "nitox_stream")))]
#[path = "event/ratsio.rs"]
mod stream_impl;
mod types;

pub(crate) use self::types::ServiceMetadata;
use self::types::{EventMessage,
                  EventMetadata,
                  HealthCheckEvent,
                  ServiceStartedEvent,
                  ServiceStoppedEvent};
use crate::{error::Result as SupResult,
            manager::{service::{HealthCheckResult,
                                Service},
                      sys::Sys}};
use clap::ArgMatches;
pub use error::{Error,
                Result};
use futures::sync::mpsc::UnboundedSender;
use habitat_common::types::{AutomateAuthToken,
                            EventStreamMetadata};
use habitat_core::env::Config as EnvConfig;
use state::Container;
use std::{net::SocketAddr,
          num::ParseIntError,
          str::FromStr,
          sync::Once,
          time::Duration};

static INIT: Once = Once::new();
lazy_static! {
    // TODO (CM): When const fn support lands in stable, we can ditch
    // this lazy_static call.

    /// Reference to the event stream.
    static ref EVENT_STREAM: Container = Container::new();
    /// Core information that is shared between all events.
    static ref EVENT_CORE: Container = Container::new();
}

/// Starts a new thread for sending events to a NATS Streaming
/// server. Stashes the handle to the stream, as well as the core
/// event information that will be a part of all events, in a global
/// static reference for access later.
pub fn init_stream(config: EventStreamConfig, event_core: EventCore) -> Result<()> {
    // call_once can't return a Result (or anything), so we'll fake it
    // by hanging onto any error we might receive.
    let mut return_value: Result<()> = Ok(());

    INIT.call_once(|| {
            let conn_info = EventConnectionInfo::new(config.token, config.url);
            match stream_impl::init_stream(conn_info) {
                Ok(event_stream) => {
                    EVENT_STREAM.set(event_stream);
                    EVENT_CORE.set(event_core);
                }
                Err(e) => return_value = Err(e),
            }
        });

    return_value
}

/// Captures all event stream-related configuration options that would
/// be passed in by a user
#[derive(Clone, Debug)]
pub struct EventStreamConfig {
    environment: String,
    application: String,
    meta:        EventStreamMetadata,
    token:       AutomateAuthToken,
    url:         String,
}

impl EventStreamConfig {
    /// Create an instance from Clap arguments.
    // TODO (CM): result type!
    pub fn from_matches(m: &ArgMatches) -> SupResult<EventStreamConfig> {
        Ok(EventStreamConfig { environment: m.value_of("EVENT_STREAM_ENVIRONMENT")
                                             .map(str::to_string)
                                             .expect("Required option for EventStream feature"),
                               application: m.value_of("EVENT_STREAM_APPLICATION")
                                             .map(str::to_string)
                                             .expect("Required option for EventStream feature"),
                               meta:        EventStreamMetadata::from(m),
                               token:       AutomateAuthToken::from(m),
                               url:         m.value_of("EVENT_STREAM_URL")
                                             .map(str::to_string)
                                             .expect("Required option for EventStream feature"), })
    }
}

/// All the information needed to establish a connection to a NATS
/// Streaming server.
// TODO: This will change as we firm up what the interaction between
// Habitat and A2 looks like.
pub struct EventConnectionInfo {
    pub name:        String,
    pub verbose:     bool,
    pub cluster_uri: String,
    pub cluster_id:  String,
    pub auth_token:  AutomateAuthToken,
}

impl EventConnectionInfo {
    pub fn new(auth_token: AutomateAuthToken, cluster_uri: String) -> Self {
        EventConnectionInfo { name: "hab_client".to_string(),
                              verbose: true,
                              cluster_uri,
                              cluster_id: "event-service".to_string(),
                              auth_token }
    }
}

/// A collection of data that will be present in all events. Rather
/// than baking this into the structure of each event, we represent it
/// once and merge the information into the final rendered form of the
/// event.
///
/// This prevents us from having to thread information throughout the
/// system, just to get it to the places where the events are
/// generated (e.g., not all code has direct access to the
/// Supervisor's ID).
#[derive(Clone, Debug)]
pub struct EventCore {
    /// The unique identifier of the Supervisor sending the event.
    supervisor_id: String,
    ip_address: SocketAddr,
    // TODO (CM): could add application and environment to the meta
    // map directly... hrmm
    application: String,
    environment: String,
    meta:        EventStreamMetadata,
}

impl EventCore {
    pub fn new(config: &EventStreamConfig, sys: &Sys) -> Self {
        EventCore { supervisor_id: sys.member_id.clone(),
                    ip_address:    sys.gossip_listen(),
                    environment:   config.environment.clone(),
                    application:   config.application.clone(),
                    meta:          config.meta.clone(), }
    }
}

/// Send an event for the start of a Service.
pub fn service_started(service: &Service) {
    if stream_initialized() {
        publish(ServiceStartedEvent { service_metadata: Some(service.to_service_metadata()),
                                      event_metadata:   None, });
    }
}

/// Send an event for the stop of a Service.
pub fn service_stopped(service: &Service) {
    if stream_initialized() {
        publish(ServiceStoppedEvent { service_metadata: Some(service.to_service_metadata()),
                                      event_metadata:   None, });
    }
}

// Takes metadata directly, rather than a `&Service` like other event
// functions, because of how the asynchronous health checking
// currently works. Revisit when async/await + Pin is all stabilized.
pub fn health_check(metadata: ServiceMetadata,
                    check_result: HealthCheckResult,
                    duration: Option<Duration>) {
    if stream_initialized() {
        let check_result: types::HealthCheck = check_result.into();
        publish(HealthCheckEvent { service_metadata: Some(metadata),
                                   event_metadata:   None,
                                   result:           i32::from(check_result),
                                   duration:         duration.map(Duration::into), });
    }
}

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

/// Internal helper function to know whether or not to go to the trouble of
/// creating event structures. If the event stream hasn't been
/// initialized, then we shouldn't need to do anything.
fn stream_initialized() -> bool { EVENT_STREAM.try_get::<EventStream>().is_some() }

/// Publish an event. This is the main interface that client code will
/// use.
///
/// If `init_stream` has not been called already, this function will
/// be a no-op.
fn publish(mut event: impl EventMessage) {
    if let Some(e) = EVENT_STREAM.try_get::<EventStream>() {
        // TODO (CM): Yeah... this is looking pretty gross. The
        // intention is to be able to timestamp the events right as
        // they go out.
        //
        // We *could* set the time when we convert the EventCore to a
        // EventMetadata struct, but that seems odd.
        //
        // It probably means that this structure just isn't the right
        // one.
        //
        // The ugliness is at least contained, though.
        event.event_metadata(EventMetadata { timestamp:
                                                 Some(std::time::SystemTime::now().into()),
                                             ..EVENT_CORE.get::<EventCore>().to_event_metadata() });

        e.send(event.to_bytes());
    }
}

/// A lightweight handle for the event stream. All events get to the
/// event stream through this.
struct EventStream(UnboundedSender<Vec<u8>>);

impl EventStream {
    /// Queues an event to be sent out.
    fn send(&self, event: Vec<u8>) {
        trace!("About to queue an event: {:?}", event);
        if let Err(e) = self.0.unbounded_send(event) {
            error!("Failed to queue event: {:?}", e);
        }
    }
}

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

/// How long should we for the event thread to start up before
/// abandoning it and shutting down?
#[derive(Clone, Debug)]
struct EventThreadStartupWait(u64);

impl Default for EventThreadStartupWait {
    fn default() -> Self { Self(5) }
}

impl FromStr for EventThreadStartupWait {
    type Err = ParseIntError;

    fn from_str(s: &str) -> ::std::result::Result<Self, Self::Err> { Ok(Self(s.parse()?)) }
}

impl EnvConfig for EventThreadStartupWait {
    const ENVVAR: &'static str = "HAB_EVENT_THREAD_STARTUP_WAIT_SEC";
}

impl Into<Duration> for EventThreadStartupWait {
    fn into(self) -> Duration { Duration::from_secs(self.0) }
}
<commit_msg>[PR Feedback] Use named `secs` field for reasonability<commit_after>//! Main interface for a stream of events the Supervisor can send out
//! in the course of its operations.
//!
//! Currently, the Supervisor is able to send events to a [NATS
//! Streaming][1] server. The `init_stream` function must be called
//! before sending events to initialize the publishing thread in the
//! background. Thereafter, you can pass "event" structs to the
//! `event` function, which will publish the event to the stream.
//!
//! All events are published under the "habitat" subject.
//!
//! [1]:https://github.com/nats-io/nats-streaming-server

mod error;

#[cfg(feature = "nitox_stream")]
#[path = "event/nitox.rs"]
mod stream_impl;

// This is the default, but setting it as a default in Cargo.toml
// makes it trickier to use nitox instead.
#[cfg(any(feature = "ratsio_stream", not(feature = "nitox_stream")))]
#[path = "event/ratsio.rs"]
mod stream_impl;
mod types;

pub(crate) use self::types::ServiceMetadata;
use self::types::{EventMessage,
                  EventMetadata,
                  HealthCheckEvent,
                  ServiceStartedEvent,
                  ServiceStoppedEvent};
use crate::{error::Result as SupResult,
            manager::{service::{HealthCheckResult,
                                Service},
                      sys::Sys}};
use clap::ArgMatches;
pub use error::{Error,
                Result};
use futures::sync::mpsc::UnboundedSender;
use habitat_common::types::{AutomateAuthToken,
                            EventStreamMetadata};
use habitat_core::env::Config as EnvConfig;
use state::Container;
use std::{net::SocketAddr,
          num::ParseIntError,
          str::FromStr,
          sync::Once,
          time::Duration};

static INIT: Once = Once::new();
lazy_static! {
    // TODO (CM): When const fn support lands in stable, we can ditch
    // this lazy_static call.

    /// Reference to the event stream.
    static ref EVENT_STREAM: Container = Container::new();
    /// Core information that is shared between all events.
    static ref EVENT_CORE: Container = Container::new();
}

/// Starts a new thread for sending events to a NATS Streaming
/// server. Stashes the handle to the stream, as well as the core
/// event information that will be a part of all events, in a global
/// static reference for access later.
pub fn init_stream(config: EventStreamConfig, event_core: EventCore) -> Result<()> {
    // call_once can't return a Result (or anything), so we'll fake it
    // by hanging onto any error we might receive.
    let mut return_value: Result<()> = Ok(());

    INIT.call_once(|| {
            let conn_info = EventConnectionInfo::new(config.token, config.url);
            match stream_impl::init_stream(conn_info) {
                Ok(event_stream) => {
                    EVENT_STREAM.set(event_stream);
                    EVENT_CORE.set(event_core);
                }
                Err(e) => return_value = Err(e),
            }
        });

    return_value
}

/// Captures all event stream-related configuration options that would
/// be passed in by a user
#[derive(Clone, Debug)]
pub struct EventStreamConfig {
    environment: String,
    application: String,
    meta:        EventStreamMetadata,
    token:       AutomateAuthToken,
    url:         String,
}

impl EventStreamConfig {
    /// Create an instance from Clap arguments.
    // TODO (CM): result type!
    pub fn from_matches(m: &ArgMatches) -> SupResult<EventStreamConfig> {
        Ok(EventStreamConfig { environment: m.value_of("EVENT_STREAM_ENVIRONMENT")
                                             .map(str::to_string)
                                             .expect("Required option for EventStream feature"),
                               application: m.value_of("EVENT_STREAM_APPLICATION")
                                             .map(str::to_string)
                                             .expect("Required option for EventStream feature"),
                               meta:        EventStreamMetadata::from(m),
                               token:       AutomateAuthToken::from(m),
                               url:         m.value_of("EVENT_STREAM_URL")
                                             .map(str::to_string)
                                             .expect("Required option for EventStream feature"), })
    }
}

/// All the information needed to establish a connection to a NATS
/// Streaming server.
// TODO: This will change as we firm up what the interaction between
// Habitat and A2 looks like.
pub struct EventConnectionInfo {
    pub name:        String,
    pub verbose:     bool,
    pub cluster_uri: String,
    pub cluster_id:  String,
    pub auth_token:  AutomateAuthToken,
}

impl EventConnectionInfo {
    pub fn new(auth_token: AutomateAuthToken, cluster_uri: String) -> Self {
        EventConnectionInfo { name: "hab_client".to_string(),
                              verbose: true,
                              cluster_uri,
                              cluster_id: "event-service".to_string(),
                              auth_token }
    }
}

/// A collection of data that will be present in all events. Rather
/// than baking this into the structure of each event, we represent it
/// once and merge the information into the final rendered form of the
/// event.
///
/// This prevents us from having to thread information throughout the
/// system, just to get it to the places where the events are
/// generated (e.g., not all code has direct access to the
/// Supervisor's ID).
#[derive(Clone, Debug)]
pub struct EventCore {
    /// The unique identifier of the Supervisor sending the event.
    supervisor_id: String,
    ip_address: SocketAddr,
    // TODO (CM): could add application and environment to the meta
    // map directly... hrmm
    application: String,
    environment: String,
    meta:        EventStreamMetadata,
}

impl EventCore {
    pub fn new(config: &EventStreamConfig, sys: &Sys) -> Self {
        EventCore { supervisor_id: sys.member_id.clone(),
                    ip_address:    sys.gossip_listen(),
                    environment:   config.environment.clone(),
                    application:   config.application.clone(),
                    meta:          config.meta.clone(), }
    }
}

/// Send an event for the start of a Service.
pub fn service_started(service: &Service) {
    if stream_initialized() {
        publish(ServiceStartedEvent { service_metadata: Some(service.to_service_metadata()),
                                      event_metadata:   None, });
    }
}

/// Send an event for the stop of a Service.
pub fn service_stopped(service: &Service) {
    if stream_initialized() {
        publish(ServiceStoppedEvent { service_metadata: Some(service.to_service_metadata()),
                                      event_metadata:   None, });
    }
}

// Takes metadata directly, rather than a `&Service` like other event
// functions, because of how the asynchronous health checking
// currently works. Revisit when async/await + Pin is all stabilized.
pub fn health_check(metadata: ServiceMetadata,
                    check_result: HealthCheckResult,
                    duration: Option<Duration>) {
    if stream_initialized() {
        let check_result: types::HealthCheck = check_result.into();
        publish(HealthCheckEvent { service_metadata: Some(metadata),
                                   event_metadata:   None,
                                   result:           i32::from(check_result),
                                   duration:         duration.map(Duration::into), });
    }
}

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

/// Internal helper function to know whether or not to go to the trouble of
/// creating event structures. If the event stream hasn't been
/// initialized, then we shouldn't need to do anything.
fn stream_initialized() -> bool { EVENT_STREAM.try_get::<EventStream>().is_some() }

/// Publish an event. This is the main interface that client code will
/// use.
///
/// If `init_stream` has not been called already, this function will
/// be a no-op.
fn publish(mut event: impl EventMessage) {
    if let Some(e) = EVENT_STREAM.try_get::<EventStream>() {
        // TODO (CM): Yeah... this is looking pretty gross. The
        // intention is to be able to timestamp the events right as
        // they go out.
        //
        // We *could* set the time when we convert the EventCore to a
        // EventMetadata struct, but that seems odd.
        //
        // It probably means that this structure just isn't the right
        // one.
        //
        // The ugliness is at least contained, though.
        event.event_metadata(EventMetadata { timestamp:
                                                 Some(std::time::SystemTime::now().into()),
                                             ..EVENT_CORE.get::<EventCore>().to_event_metadata() });

        e.send(event.to_bytes());
    }
}

/// A lightweight handle for the event stream. All events get to the
/// event stream through this.
struct EventStream(UnboundedSender<Vec<u8>>);

impl EventStream {
    /// Queues an event to be sent out.
    fn send(&self, event: Vec<u8>) {
        trace!("About to queue an event: {:?}", event);
        if let Err(e) = self.0.unbounded_send(event) {
            error!("Failed to queue event: {:?}", e);
        }
    }
}

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

/// How long should we for the event thread to start up before
/// abandoning it and shutting down?
#[derive(Clone, Debug)]
struct EventThreadStartupWait {
    secs: u64,
}

impl Default for EventThreadStartupWait {
    fn default() -> Self { Self { secs: 5 } }
}

impl FromStr for EventThreadStartupWait {
    type Err = ParseIntError;

    fn from_str(s: &str) -> ::std::result::Result<Self, Self::Err> { Ok(Self { secs: s.parse()? }) }
}

impl EnvConfig for EventThreadStartupWait {
    const ENVVAR: &'static str = "HAB_EVENT_THREAD_STARTUP_WAIT_SEC";
}

impl Into<Duration> for EventThreadStartupWait {
    fn into(self) -> Duration { Duration::from_secs(self.secs) }
}
<|endoftext|>


<commit_before>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.3

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::SocketAddr;
use std::io::println;
use std::cell::Cell;
use std::{os, str, io};
use extra::arc;
use std::comm::*;
use extra::priority_queue;
use extra::md4;
use extra::sort;
use std::path;

static PORT:    int = 4414;
static IP: &'static str = "127.0.0.1";

static MAX_CACHE_SIZE_BYTES: u64 = 50000000;
static CACHE_MANAGER_A_RATE: u64 = 2000;
static CACHE_MANAGER_B_RATE: u64 = 6000;

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath,
    topPriority: int,
    fileSize: u64,
    httpHeader: ~str
}

impl Ord for sched_msg {
    fn lt(&self, other: &sched_msg) -> bool {

        let mut retVal:bool = false;

        //If self has a lower geographic priority, then it is less than
        if (self.topPriority < other.topPriority) {
            retVal = true;
        }
        else if (self.topPriority == other.topPriority) {

            //If the two have equal priorities, then we will sort by file size
            //Smaller files get a higher priority
            if (self.fileSize > other.fileSize) {
                retVal = true;
            }
        }

        retVal
    }
}

struct cache_item {
    name: ~str,
    in_use_flag: bool,
    ssi_flag: bool,
    hash: ~str,
    data: ~[u8],
    count: uint,
    size: u64
}

fn le(this: &cache_item, other: &cache_item) -> bool {
        let mut retVal:bool = false;

        if(this.count <= other.count) {
            retVal = true;
        }

        retVal
    }

fn main() {
    let req_vec: ~priority_queue::PriorityQueue<sched_msg> = ~priority_queue::PriorityQueue::new();
    let shared_req_vec = arc::RWArc::new(req_vec);
    let add_vec = shared_req_vec.clone();
    let take_vec = shared_req_vec.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    //Variables pertaining to safe visitor counting
    let safe_count: ~[uint] = ~[0];
    let shared_visit = arc::RWArc::new(safe_count);

    //Variables pertaining to cache and cache manager
    let cache_list: ~[cache_item] = ~[];
    let shared_cache_list = arc::RWArc::new(cache_list);
    let cache_manager_a = shared_cache_list.clone();
    let cache_manager_b = shared_cache_list.clone();

    //CACHE UPDATE MANAGER (Manager B)
    //This will handle making sure that items in the cache are up-to-date in case they are changed
    do spawn {
        loop {
            do cache_manager_b.write |vec| {
                for i in range(0, (*vec).len()) {
                    //If it is in the cache and in use, we will check to see if the file has been updated
                    if((*vec)[i].in_use_flag) {
                        let curr_path = ~path::Path((*vec)[i].name);

                        if os::path_exists(curr_path) {
                            match io::read_whole_file(curr_path) {
                                    Ok(file_data) => {
                                        let temp_md4 = md4::md4_str(file_data.to_owned());

                                        if(temp_md4 != (*vec)[i].hash)
                                        {
                                            let fileInfo = match std::rt::io::file::stat(curr_path) {
                                                Some(s) => s,
                                                None => fail!("Could not access file stats for cache")
                                            };

                                            println(fmt!("===== UPDATING FILE: %?", (*vec)[i].name));

                                            (*vec)[i].data = file_data;
                                            (*vec)[i].size = fileInfo.size;
                                            (*vec)[i].hash = temp_md4;
                                        }
                                    }
                                    Err(err) => {
                                        println("ERROR IN UPDATE CACHE");
                                        println(err);
                                    }
                                }
                        }
                    }
                }
            }

            timer::sleep(CACHE_MANAGER_B_RATE);
        }
    }

    //MAIN CACHE MANAGER (Manager A)
    do spawn {
        loop {
            do cache_manager_a.write |vec| {
                //Quick sort sorts in-place, so we don't need to worry about memory overhead
                //Just time overhead
                sort::quick_sort((*vec), le);

                let mut cache_remaining = MAX_CACHE_SIZE_BYTES;

                for i in range(0, (*vec).len()) {
                    if((*vec)[i].size <= cache_remaining && !(*vec)[i].in_use_flag) {
                        let curr_path = ~path::Path((*vec)[i].name);

                        if os::path_exists(curr_path) {
                            let fileInfo = match std::rt::io::file::stat(curr_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats for cache")
                            };

                            if(fileInfo.size <= cache_remaining) {
                                match io::read_whole_file(curr_path) {
                                    Ok(file_data) => {

                                        //NEED TO ADD CHECK IF IT HAS SSI
                                        //PERHAPS EXCLUDE ALL SSI FROM CACHE?
                                        //ADD TO IF: size < remaining, !in use, !ssi?

                                        (*vec)[i].data = file_data.to_owned();
                                        (*vec)[i].size = fileInfo.size;
                                        (*vec)[i].hash = md4::md4_str(file_data);
                                        (*vec)[i].in_use_flag = true;

                                        cache_remaining = cache_remaining - fileInfo.size;
                                    }
                                    Err(err) => {
                                        println("ERROR IN UPDATE CACHE");
                                        println(err);
                                    }
                                }
                            }
                        }
                    }
                    else if((*vec)[i].size <= cache_remaining && (*vec)[i].in_use_flag) {
                        cache_remaining = cache_remaining - (*vec)[i].size;
                    }
                    else if((*vec)[i].size > cache_remaining) {
                        (*vec)[i].in_use_flag = false;
                        (*vec)[i].data = ~[];
                    }

                }
            }

            timer::sleep(CACHE_MANAGER_A_RATE);
        }
    }

    // dequeue file requests, and send responses.
    // FIFO
    do spawn {
        let (sm_port, sm_chan) = stream();

        // a task for sending responses.
        do spawn {
            loop {
                let mut tf: sched_msg = sm_port.recv(); // wait for the dequeued request to handle

                //Check if file is in cache
                //Will do so using RWArc
                let mut serve_from_cache: bool = false;

                do shared_cache_list.write |vec| {
                    let mut found: bool = false;

                    for i in range(0, (*vec).len()) {
                        if( (*vec)[i].name == tf.filepath.to_str() && (*vec)[i].in_use_flag) {
                            serve_from_cache = true;
                            found = true;

                            println(fmt!("===== SERVING FROM CACHE: %?", tf.filepath.to_str()));

                            tf.stream.write(tf.httpHeader.as_bytes());
                            tf.stream.write((*vec)[i].data);

                            (*vec)[i].count += 1;
                        }
                        else if( (*vec)[i].name == tf.filepath.to_str() && !(*vec)[i].in_use_flag) {
                            (*vec)[i].count += 1;
                            found = true;
                        }
                    }

                    //If it isn't found, we will create a blank entry and enter basic info, and add the
                    //data and md4 later
                    if(!found) {
                        println(fmt!("===== ADDING ITEM %?", tf.filepath.to_str()));

                        let new_cache_item: cache_item = cache_item{name: tf.filepath.to_str(), in_use_flag: false,
                            ssi_flag: false, hash: ~"", data: ~[], count: 1, size: tf.fileSize};

                        (*vec).push(new_cache_item);
                    }
                }

                if(!serve_from_cache) {
                    match io::read_whole_file(tf.filepath) { // killed if file size is larger than memory size.
                        Ok(file_data) => {
                            tf.stream.write(tf.httpHeader.as_bytes());
                            tf.stream.write(file_data);
                            println(fmt!("===== SERVING FROM DISK: %?", tf.filepath.to_str()));

                        }
                        Err(err) => {
                            println("ERROR IN SEND");
                            println(err);
                        }
                    }
                }
            }
        }

        loop {
            port.recv(); // wait for arrving notification
            do take_vec.write |vec| {
                if ((*vec).len() > 0) {

                    //Since we are using a priority queue, we will use pop()
                    let tf = (*vec).pop();
                    //println(fmt!("shift from queue, size: %ud", (*vec).len()));
                    sm_chan.send(tf); // send the request to send-response-task to serve.
                }
            }
        }
    }

    let ip = match FromStr::from_str(IP) { Some(ip) => ip,
                                           None => { println(fmt!("Error: Invalid IP address <%s>", IP));
                                                     return;},
                                         };

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: ip, port: PORT as u16});

    println(fmt!("Listening on %s:%d ...", ip.to_str(), PORT));
    let mut acceptor = socket.listen().unwrap();

    loop {
        let stream = acceptor.accept();
        let stream = Cell::new(stream);

        // Start a new task to handle the each connection
        let child_chan = chan.clone();
        let child_add_vec = add_vec.clone();
        let child_arc = shared_visit.clone();

        do spawn {

            let mut actual_count:uint = 0;

            //We will use the cloned RWARC
            do child_arc.write |count_vec| {
                let prev_count = (*count_vec).pop();
                let new_count: uint = prev_count + 1;

                actual_count = new_count;
                (*count_vec).push(new_count);
            }

            let stream = stream.take();

            match stream {
                Some(s) => {

                    let mut stream = s;

                    let mut buf = [0, ..500];
                    stream.read(buf);
                    let request_str = str::from_utf8(buf);

                    let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
                    if req_group.len() > 2 {
                        let path = req_group[1];

                        let file_path = ~os::getcwd().push(path.replace("/../", ""));

                        if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                            let response: ~str = fmt!(
                                "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                                 <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                                 \r\n", actual_count);

                            stream.write(response.as_bytes());
                        }
                        else {
                            // Requests scheduling

                            let mut streamPriority: int = 0;

                            //Retrieving the requesting IP address
                            let ipStr: ~str = match (stream).peer_name() {
                                Some(pr) => pr.ip.to_str(),
                                None => ~"0.0.0.0"
                            };

                            //Split the IP address by '.' so that we can compare
                            let ipSplit: ~[~str] = ipStr.split_iter('.').filter(|&x| x != "")
                                 .map(|x| x.to_owned()).collect();

                            //Assign priority based on geography or if localhost
                            if ( (ipSplit[0] == ~"127" && ipSplit[1] == ~"0") || (ipSplit[0] == ~"128" && ipSplit[1] == ~"143")
                                || (ipSplit[0] == ~"137" && ipSplit[1] == ~"54") ) {

                                streamPriority = 1;
                            }

                            let fileName: ~str = file_path.filename().unwrap().to_owned();
                            let fileNameSplit: ~[~str] = fileName.split_iter('.').filter(|&x| x != "").map(|x| x.to_owned()).collect();

                            //In order to optimize for the benchmark, we will send the HTTP header quickly before adding to the queue
                            let httpHeader: ~str = match fileNameSplit[fileNameSplit.len()-1] {
                                ~"html" | ~"htm" | ~"php" => ~"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n",
                                _ => ~"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream; charset=UTF-8\r\n\r\n"
                            };

                            //stream.write(httpHeader.as_bytes());
                            //stream.flush();

                            //Retrieve file info for additional latency fixes
                            let fileInfo = match std::rt::io::file::stat(file_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats")
                            };

                            let msg: sched_msg = sched_msg{stream: Some(stream), filepath: file_path.clone(), topPriority: streamPriority, fileSize: fileInfo.size, httpHeader: httpHeader};
                            let (sm_port, sm_chan) = std::comm::stream();
                            sm_chan.send(msg);

                            do child_add_vec.write |vec| {
                                let msg = sm_port.recv();
                                (*vec).push(msg); // enqueue new request.
                            }
                            child_chan.send(""); //notify the new arriving request.
                        }
                    }
                    println!("connection terminates")

                },
                None => ()
            }
        }
    }
}<commit_msg>Modified cache so that files with server-side includes are never cached<commit_after>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.3

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::SocketAddr;
use std::io::println;
use std::cell::Cell;
use std::{os, str, io};
use extra::arc;
use std::comm::*;
use extra::priority_queue;
use extra::md4;
use extra::sort;
use std::path;

static PORT:    int = 4414;
static IP: &'static str = "127.0.0.1";

static MAX_CACHE_SIZE_BYTES: u64 = 50000000;
static CACHE_MANAGER_A_RATE: u64 = 2000;
static CACHE_MANAGER_B_RATE: u64 = 6000;

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath,
    topPriority: int,
    fileSize: u64,
    httpHeader: ~str
}

impl Ord for sched_msg {
    fn lt(&self, other: &sched_msg) -> bool {

        let mut retVal:bool = false;

        //If self has a lower geographic priority, then it is less than
        if (self.topPriority < other.topPriority) {
            retVal = true;
        }
        else if (self.topPriority == other.topPriority) {

            //If the two have equal priorities, then we will sort by file size
            //Smaller files get a higher priority
            if (self.fileSize > other.fileSize) {
                retVal = true;
            }
        }

        retVal
    }
}

struct cache_item {
    name: ~str,
    in_use_flag: bool,
    ssi_flag: bool,
    hash: ~str,
    data: ~[u8],
    count: uint,
    size: u64
}

fn le(this: &cache_item, other: &cache_item) -> bool {
        let mut retVal:bool = false;

        if(this.count <= other.count) {
            retVal = true;
        }

        retVal
    }

fn main() {
    let req_vec: ~priority_queue::PriorityQueue<sched_msg> = ~priority_queue::PriorityQueue::new();
    let shared_req_vec = arc::RWArc::new(req_vec);
    let add_vec = shared_req_vec.clone();
    let take_vec = shared_req_vec.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    //Variables pertaining to safe visitor counting
    let safe_count: ~[uint] = ~[0];
    let shared_visit = arc::RWArc::new(safe_count);

    //Variables pertaining to cache and cache manager
    let cache_list: ~[cache_item] = ~[];
    let shared_cache_list = arc::RWArc::new(cache_list);
    let cache_manager_a = shared_cache_list.clone();
    let cache_manager_b = shared_cache_list.clone();

    //CACHE UPDATE MANAGER (Manager B)
    //This will handle making sure that items in the cache are up-to-date in case they are changed
    do spawn {
        loop {
            do cache_manager_b.write |vec| {
                for i in range(0, (*vec).len()) {
                    //If it is in the cache and in use, we will check to see if the file has been updated
                    if((*vec)[i].in_use_flag) {
                        let curr_path = ~path::Path((*vec)[i].name);

                        if os::path_exists(curr_path) {
                            match io::read_whole_file(curr_path) {
                                    Ok(file_data) => {
                                        let temp_md4 = md4::md4_str(file_data.to_owned());

                                        if(temp_md4 != (*vec)[i].hash)
                                        {
                                            let fileInfo = match std::rt::io::file::stat(curr_path) {
                                                Some(s) => s,
                                                None => fail!("Could not access file stats for cache")
                                            };

                                            println(fmt!("===== UPDATING FILE: %?", (*vec)[i].name));

                                            (*vec)[i].data = file_data;
                                            (*vec)[i].size = fileInfo.size;
                                            (*vec)[i].hash = temp_md4;
                                        }
                                    }
                                    Err(err) => {
                                        println("ERROR IN UPDATE CACHE");
                                        println(err);
                                    }
                                }
                        }
                    }
                }
            }

            timer::sleep(CACHE_MANAGER_B_RATE);
        }
    }

    //MAIN CACHE MANAGER (Manager A)
    do spawn {
        loop {
            do cache_manager_a.write |vec| {
                //Quick sort sorts in-place, so we don't need to worry about memory overhead
                //Just time overhead
                sort::quick_sort((*vec), le);

                let mut cache_remaining = MAX_CACHE_SIZE_BYTES;

                for i in range(0, (*vec).len()) {
                    if((*vec)[i].size <= cache_remaining && !(*vec)[i].in_use_flag) {
                        let curr_path = ~path::Path((*vec)[i].name);

                        if os::path_exists(curr_path) {
                            let fileInfo = match std::rt::io::file::stat(curr_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats for cache")
                            };

                            if(fileInfo.size <= cache_remaining) {
                                match io::read_whole_file(curr_path) {
                                    Ok(file_data) => {

                                        let fileName: ~str = curr_path.filename().unwrap().to_owned();
                                        let fileNameSplit: ~[~str] = fileName.split_iter('.').filter(|&x| x != "").map(|x| x.to_owned()).collect();

                                        let is_html: bool = match fileNameSplit[fileNameSplit.len()-1] {
                                            ~"html" | ~"htm" => true,
                                            _ => false
                                        };

                                        let mut has_ssi: bool = false;

                                        //We will only check html files for server-side includes
                                        if(is_html) {
                                            has_ssi = std::str::from_utf8_slice(file_data).contains("


<commit_before>pub trait Parse : Sized {

    fn parse(s: &str) -> Option<Self>;

}

<commit_msg>Add time parse regex<commit_after>pub trait Parse : Sized {

    fn parse(s: &str) -> Option<Self>;

}

pub fn time_parse_regex() -> &'static str {
    r#"(?P<Y>\d{4})-(?P<M>\d{2})-(?P<D>\d{2})(T(?P<h>\d{2})(:(?P<m>\d{2})(:(?P<s>\d{2}))?)?)?$"#
}

<|endoftext|>


<commit_before>//! **GraphMap\<N, E\>** is an undirected graph where node values are mapping keys.

use std::hash::{Hash};
use std::collections::HashMap;
use std::iter::Map;
use std::collections::hash_map::{
    Keys,
};
use std::collections::hash_map::Entry::{
    Occupied,
    Vacant,
};
use std::slice::{
    Iter,
};
use std::fmt;
use std::ops::{Index, IndexMut};

/// **GraphMap\<N, E\>** is an undirected graph, with generic node values **N** and edge weights **E**.
///
/// It uses an combined adjacency list and sparse adjacency matrix representation, using **O(|V|
/// + |E|)** space, and allows testing for edge existance in constant time.
///
/// The node type **N** must implement **Copy** and will be used as node identifier, duplicated
/// into several places in the data structure.
/// It must be suitable as a hash table key (implementing **Eq + Hash**).
/// The node type must also implement **Ord** so that the implementation can
/// order the pair (**a**, **b**) for an edge connecting any two nodes **a** and **b**.
///
/// **GraphMap** does not allow parallel edges, but self loops are allowed.
#[derive(Clone)]
pub struct GraphMap<N: Eq + Hash, E> {
    nodes: HashMap<N, Vec<N>>,
    edges: HashMap<(N, N), E>,
}

impl<N: Eq + Hash + fmt::Debug, E: fmt::Debug> fmt::Debug for GraphMap<N, E>
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.nodes.fmt(f)
    }
}

#[inline]
fn edge_key<N: Copy + Ord>(a: N, b: N) -> (N, N)
{
    if a <= b { (a, b) } else { (b, a) }
}

#[inline]
fn copy<N: Copy>(n: &N) -> N { *n }

/// A trait group for **GraphMap**'s node identifier.
pub trait NodeTrait : Copy + Ord + Eq + Hash {}
impl<N> NodeTrait for N where N: Copy + Ord + Eq + Hash {}

impl<N, E> GraphMap<N, E> where N: NodeTrait
{
    /// Create a new **GraphMap**.
    pub fn new() -> Self
    {
        GraphMap {
            nodes: HashMap::new(),
            edges: HashMap::new(),
        }
    }

    /// Create a new **GraphMap** with estimated capacity.
    pub fn with_capacity(nodes: usize, edges: usize) -> Self
    {
        GraphMap {
            nodes: HashMap::with_capacity(nodes),
            edges: HashMap::with_capacity(edges),
        }
    }

    /// Return the number of nodes in the graph.
    pub fn node_count(&self) -> usize
    {
        self.nodes.len()
    }

    /// Return the number of edges in the graph.
    pub fn edge_count(&self) -> usize
    {
        self.edges.len()
    }

    /// Remove all nodes and edges
    pub fn clear(&mut self)
    {
        self.nodes.clear();
        self.edges.clear();
    }

    /// Add node **n** to the graph.
    pub fn add_node(&mut self, n: N) -> N {
        match self.nodes.entry(n) {
            Occupied(_) => {}
            Vacant(ent) => { ent.insert(Vec::new()); }
        }
        n
    }

    /// Return **true** if node **n** was removed.
    pub fn remove_node(&mut self, n: N) -> bool {
        let successors = match self.nodes.remove(&n) {
            None => return false,
            Some(sus) => sus,
        };
        for succ in successors.into_iter() {
            // remove all successor links
            self.remove_single_edge(&succ, &n);
            // Remove all edge values
            self.edges.remove(&edge_key(n, succ));
        }
        true
    }

    /// Return **true** if the node is contained in the graph.
    pub fn contains_node(&self, n: N) -> bool {
        self.nodes.contains_key(&n)
    }

    /// Add an edge connecting **a** and **b** to the graph.
    ///
    /// Inserts nodes **a** and/or **b** if they aren't already part of the graph.
    ///
    /// Return **true** if edge did not previously exist.
    ///
    /// ## Example
    /// ```
    /// use petgraph::GraphMap;
    ///
    /// let mut g = GraphMap::new();
    /// g.add_edge(1, 2, -1);
    /// assert_eq!(g.node_count(), 2);
    /// assert_eq!(g.edge_count(), 1);
    /// ```
    pub fn add_edge(&mut self, a: N, b: N, edge: E) -> bool
    {
        // Use Ord to order the edges
        match self.nodes.entry(a) {
            Occupied(ent) => { ent.into_mut().push(b); }
            Vacant(ent) => { ent.insert(vec![b]); }
        }
        match self.nodes.entry(b) {
            Occupied(ent) => { ent.into_mut().push(a); }
            Vacant(ent) => { ent.insert(vec![a]); }
        }
        self.edges.insert(edge_key(a, b), edge).is_none()
    }

    /// Remove successor relation from a to b
    ///
    /// Return **true** if it did exist.
    fn remove_single_edge(&mut self, a: &N, b: &N) -> bool
    {
        match self.nodes.get_mut(a) {
            None => false,
            Some(sus) => {
                match sus.iter().position(|elt| elt == b) {
                    Some(index) => { sus.swap_remove(index); true }
                    None => false,
                }
            }
        }
    }

    /// Remove edge from **a** to **b** from the graph and return the edge weight.
    ///
    /// Return **None** if the edge didn't exist.
    ///
    /// ## Example
    ///
    /// ```
    /// use petgraph::GraphMap;
    ///
    /// let mut g = GraphMap::new();
    /// g.add_node(1);
    /// g.add_node(2);
    /// g.add_edge(1, 2, -1);
    ///
    /// let edge = g.remove_edge(2, 1);
    /// assert_eq!(edge, Some(-1));
    /// assert_eq!(g.edge_count(), 0);
    /// ```
    pub fn remove_edge(&mut self, a: N, b: N) -> Option<E>
    {
        let exist1 = self.remove_single_edge(&a, &b);
        let exist2 = self.remove_single_edge(&b, &a);
        let weight = self.edges.remove(&edge_key(a, b));
        debug_assert!(exist1 == exist2 && exist1 == weight.is_some());
        weight
    }

    /// Return **true** if the edge connecting **a** with **b** is contained in the graph.
    pub fn contains_edge(&self, a: N, b: N) -> bool {
        self.edges.contains_key(&edge_key(a, b))
    }

    /// Return an iterator over the nodes of the graph.
    ///
    /// Iterator element type is **N**.
    pub fn nodes<'a>(&'a self) -> Nodes<'a, N>
    {
        Nodes{iter: self.nodes.keys().map(copy)}
    }

    /// Return an iterator over the nodes that are connected with **from** by edges.
    ///
    /// If the node **from** does not exist in the graph, return an empty iterator.
    ///
    /// Iterator element type is **N**.
    pub fn neighbors<'a>(&'a self, from: N) -> Neighbors<'a, N>
    {
        Neighbors{iter:
            match self.nodes.get(&from) {
                Some(neigh) => neigh.iter(),
                None => [].iter(),
            }.map(copy)
        }
    }

    /// Return an iterator over the nodes that are connected with **from** by edges,
    /// paired with the edge weight.
    ///
    /// If the node **from** does not exist in the graph, return an empty iterator.
    ///
    /// Iterator element type is **(N, &'a E)**.
    pub fn edges<'a>(&'a self, from: N) -> Edges<'a, N, E>
    {
        Edges {
            from: from,
            iter: self.neighbors(from),
            edges: &self.edges,
        }
    }

    /// Return a reference to the edge weight connecting **a** with **b**, or
    /// **None** if the edge does not exist in the graph.
    pub fn edge_weight<'a>(&'a self, a: N, b: N) -> Option<&'a E>
    {
        self.edges.get(&edge_key(a, b))
    }

    /// Return a mutable reference to the edge weight connecting **a** with **b**, or
    /// **None** if the edge does not exist in the graph.
    pub fn edge_weight_mut<'a>(&'a mut self, a: N, b: N) -> Option<&'a mut E>
    {
        self.edges.get_mut(&edge_key(a, b))
    }
}

macro_rules! iterator_methods {
    () => (
        #[inline]
        fn next(&mut self) -> Option<Self::Item>
        {
            self.iter.next()
        }

        #[inline]
        fn size_hint(&self) -> (usize, Option<usize>)
        {
            self.iter.size_hint()
        }
    )
}

pub struct Nodes<'a, N: 'a> {
    iter: Map<Keys<'a, N, Vec<N>>, fn(&N) -> N>,
}

impl<'a, N> Iterator for Nodes<'a, N>
{
    type Item = N;
    iterator_methods!();
}

pub struct Neighbors<'a, N: 'a> {
    iter: Map<Iter<'a, N>, fn(&N) -> N>,
}

impl<'a, N> Iterator for Neighbors<'a, N>
{
    type Item = N;
    iterator_methods!();
}

pub struct Edges<'a, N, E: 'a> where N: 'a + NodeTrait
{
    pub from: N,
    pub edges: &'a HashMap<(N, N), E>,
    pub iter: Neighbors<'a, N>,
}

impl<'a, N, E> Iterator for Edges<'a, N, E>
    where N: 'a + NodeTrait, E: 'a
{
    type Item = (N, &'a E);
    fn next(&mut self) -> Option<(N, &'a E)>
    {
        match self.iter.next() {
            None => None,
            Some(b) => {
                let a = self.from;
                match self.edges.get(&edge_key(a, b)) {
                    None => unreachable!(),
                    Some(edge) => {
                        Some((b, edge))
                    }
                }
            }
        }
    }
}

/// Index **GraphMap** by node pairs to access edge weights.
impl<N, E> Index<(N, N)> for GraphMap<N, E> where N: NodeTrait
{
    type Output = E;
    fn index(&self, index: (N, N)) -> &E
    {
        self.edge_weight(index.0, index.1).expect("GraphMap::index: no such edge")
    }
}

/// Index **GraphMap** by node pairs to access edge weights.
impl<N, E> IndexMut<(N, N)> for GraphMap<N, E> where N: NodeTrait
{
    fn index_mut(&mut self, index: (N, N)) -> &mut E
    {
        self.edge_weight_mut(index.0, index.1).expect("GraphMap::index: no such edge")
    }
}
<commit_msg>Improve the iterator wrapper macro<commit_after>//! **GraphMap\<N, E\>** is an undirected graph where node values are mapping keys.

use std::hash::{Hash};
use std::collections::HashMap;
use std::iter::Map;
use std::collections::hash_map::{
    Keys,
};
use std::collections::hash_map::Entry::{
    Occupied,
    Vacant,
};
use std::slice::{
    Iter,
};
use std::fmt;
use std::ops::{Index, IndexMut};

/// **GraphMap\<N, E\>** is an undirected graph, with generic node values **N** and edge weights **E**.
///
/// It uses an combined adjacency list and sparse adjacency matrix representation, using **O(|V|
/// + |E|)** space, and allows testing for edge existance in constant time.
///
/// The node type **N** must implement **Copy** and will be used as node identifier, duplicated
/// into several places in the data structure.
/// It must be suitable as a hash table key (implementing **Eq + Hash**).
/// The node type must also implement **Ord** so that the implementation can
/// order the pair (**a**, **b**) for an edge connecting any two nodes **a** and **b**.
///
/// **GraphMap** does not allow parallel edges, but self loops are allowed.
#[derive(Clone)]
pub struct GraphMap<N: Eq + Hash, E> {
    nodes: HashMap<N, Vec<N>>,
    edges: HashMap<(N, N), E>,
}

impl<N: Eq + Hash + fmt::Debug, E: fmt::Debug> fmt::Debug for GraphMap<N, E>
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.nodes.fmt(f)
    }
}

#[inline]
fn edge_key<N: Copy + Ord>(a: N, b: N) -> (N, N)
{
    if a <= b { (a, b) } else { (b, a) }
}

#[inline]
fn copy<N: Copy>(n: &N) -> N { *n }

/// A trait group for **GraphMap**'s node identifier.
pub trait NodeTrait : Copy + Ord + Eq + Hash {}
impl<N> NodeTrait for N where N: Copy + Ord + Eq + Hash {}

impl<N, E> GraphMap<N, E> where N: NodeTrait
{
    /// Create a new **GraphMap**.
    pub fn new() -> Self
    {
        GraphMap {
            nodes: HashMap::new(),
            edges: HashMap::new(),
        }
    }

    /// Create a new **GraphMap** with estimated capacity.
    pub fn with_capacity(nodes: usize, edges: usize) -> Self
    {
        GraphMap {
            nodes: HashMap::with_capacity(nodes),
            edges: HashMap::with_capacity(edges),
        }
    }

    /// Return the number of nodes in the graph.
    pub fn node_count(&self) -> usize
    {
        self.nodes.len()
    }

    /// Return the number of edges in the graph.
    pub fn edge_count(&self) -> usize
    {
        self.edges.len()
    }

    /// Remove all nodes and edges
    pub fn clear(&mut self)
    {
        self.nodes.clear();
        self.edges.clear();
    }

    /// Add node **n** to the graph.
    pub fn add_node(&mut self, n: N) -> N {
        match self.nodes.entry(n) {
            Occupied(_) => {}
            Vacant(ent) => { ent.insert(Vec::new()); }
        }
        n
    }

    /// Return **true** if node **n** was removed.
    pub fn remove_node(&mut self, n: N) -> bool {
        let successors = match self.nodes.remove(&n) {
            None => return false,
            Some(sus) => sus,
        };
        for succ in successors.into_iter() {
            // remove all successor links
            self.remove_single_edge(&succ, &n);
            // Remove all edge values
            self.edges.remove(&edge_key(n, succ));
        }
        true
    }

    /// Return **true** if the node is contained in the graph.
    pub fn contains_node(&self, n: N) -> bool {
        self.nodes.contains_key(&n)
    }

    /// Add an edge connecting **a** and **b** to the graph.
    ///
    /// Inserts nodes **a** and/or **b** if they aren't already part of the graph.
    ///
    /// Return **true** if edge did not previously exist.
    ///
    /// ## Example
    /// ```
    /// use petgraph::GraphMap;
    ///
    /// let mut g = GraphMap::new();
    /// g.add_edge(1, 2, -1);
    /// assert_eq!(g.node_count(), 2);
    /// assert_eq!(g.edge_count(), 1);
    /// ```
    pub fn add_edge(&mut self, a: N, b: N, edge: E) -> bool
    {
        // Use Ord to order the edges
        match self.nodes.entry(a) {
            Occupied(ent) => { ent.into_mut().push(b); }
            Vacant(ent) => { ent.insert(vec![b]); }
        }
        match self.nodes.entry(b) {
            Occupied(ent) => { ent.into_mut().push(a); }
            Vacant(ent) => { ent.insert(vec![a]); }
        }
        self.edges.insert(edge_key(a, b), edge).is_none()
    }

    /// Remove successor relation from a to b
    ///
    /// Return **true** if it did exist.
    fn remove_single_edge(&mut self, a: &N, b: &N) -> bool
    {
        match self.nodes.get_mut(a) {
            None => false,
            Some(sus) => {
                match sus.iter().position(|elt| elt == b) {
                    Some(index) => { sus.swap_remove(index); true }
                    None => false,
                }
            }
        }
    }

    /// Remove edge from **a** to **b** from the graph and return the edge weight.
    ///
    /// Return **None** if the edge didn't exist.
    ///
    /// ## Example
    ///
    /// ```
    /// use petgraph::GraphMap;
    ///
    /// let mut g = GraphMap::new();
    /// g.add_node(1);
    /// g.add_node(2);
    /// g.add_edge(1, 2, -1);
    ///
    /// let edge = g.remove_edge(2, 1);
    /// assert_eq!(edge, Some(-1));
    /// assert_eq!(g.edge_count(), 0);
    /// ```
    pub fn remove_edge(&mut self, a: N, b: N) -> Option<E>
    {
        let exist1 = self.remove_single_edge(&a, &b);
        let exist2 = self.remove_single_edge(&b, &a);
        let weight = self.edges.remove(&edge_key(a, b));
        debug_assert!(exist1 == exist2 && exist1 == weight.is_some());
        weight
    }

    /// Return **true** if the edge connecting **a** with **b** is contained in the graph.
    pub fn contains_edge(&self, a: N, b: N) -> bool {
        self.edges.contains_key(&edge_key(a, b))
    }

    /// Return an iterator over the nodes of the graph.
    ///
    /// Iterator element type is **N**.
    pub fn nodes<'a>(&'a self) -> Nodes<'a, N>
    {
        Nodes{iter: self.nodes.keys().map(copy)}
    }

    /// Return an iterator over the nodes that are connected with **from** by edges.
    ///
    /// If the node **from** does not exist in the graph, return an empty iterator.
    ///
    /// Iterator element type is **N**.
    pub fn neighbors<'a>(&'a self, from: N) -> Neighbors<'a, N>
    {
        Neighbors{iter:
            match self.nodes.get(&from) {
                Some(neigh) => neigh.iter(),
                None => [].iter(),
            }.map(copy)
        }
    }

    /// Return an iterator over the nodes that are connected with **from** by edges,
    /// paired with the edge weight.
    ///
    /// If the node **from** does not exist in the graph, return an empty iterator.
    ///
    /// Iterator element type is **(N, &'a E)**.
    pub fn edges<'a>(&'a self, from: N) -> Edges<'a, N, E>
    {
        Edges {
            from: from,
            iter: self.neighbors(from),
            edges: &self.edges,
        }
    }

    /// Return a reference to the edge weight connecting **a** with **b**, or
    /// **None** if the edge does not exist in the graph.
    pub fn edge_weight<'a>(&'a self, a: N, b: N) -> Option<&'a E>
    {
        self.edges.get(&edge_key(a, b))
    }

    /// Return a mutable reference to the edge weight connecting **a** with **b**, or
    /// **None** if the edge does not exist in the graph.
    pub fn edge_weight_mut<'a>(&'a mut self, a: N, b: N) -> Option<&'a mut E>
    {
        self.edges.get_mut(&edge_key(a, b))
    }
}

/// Utitily macro -- reinterpret passed in macro arguments as items
macro_rules! items {
    ($($item:item)*) => ($($item)*);
}

macro_rules! iterator_wrap {
    ($name: ident <$($typarm:tt),*> where { $($bounds: tt)* }
     item: $item: ty,
     iter: $iter: ty,
     ) => (
        items! {
            pub struct $name <$($typarm),*> where $($bounds)* {
                iter: $iter,
            }
            impl<$($typarm),*> Iterator for $name <$($typarm),*>
            {
                type Item = $item;
                #[inline]
                fn next(&mut self) -> Option<Self::Item> {
                    self.iter.next()
                }

                #[inline]
                fn size_hint(&self) -> (usize, Option<usize>) {
                    self.iter.size_hint()
                }
            }
        }
    );
}

iterator_wrap! {
    Nodes <'a, N> where { N: 'a }
    item: N,
    iter: Map<Keys<'a, N, Vec<N>>, fn(&N) -> N>,
}

iterator_wrap! {
    Neighbors <'a, N> where { N: 'a }
    item: N,
    iter: Map<Iter<'a, N>, fn(&N) -> N>,
}

pub struct Edges<'a, N, E: 'a> where N: 'a + NodeTrait
{
    pub from: N,
    pub edges: &'a HashMap<(N, N), E>,
    pub iter: Neighbors<'a, N>,
}

impl<'a, N, E> Iterator for Edges<'a, N, E>
    where N: 'a + NodeTrait, E: 'a
{
    type Item = (N, &'a E);
    fn next(&mut self) -> Option<(N, &'a E)>
    {
        match self.iter.next() {
            None => None,
            Some(b) => {
                let a = self.from;
                match self.edges.get(&edge_key(a, b)) {
                    None => unreachable!(),
                    Some(edge) => {
                        Some((b, edge))
                    }
                }
            }
        }
    }
}

/// Index **GraphMap** by node pairs to access edge weights.
impl<N, E> Index<(N, N)> for GraphMap<N, E> where N: NodeTrait
{
    type Output = E;
    fn index(&self, index: (N, N)) -> &E
    {
        self.edge_weight(index.0, index.1).expect("GraphMap::index: no such edge")
    }
}

/// Index **GraphMap** by node pairs to access edge weights.
impl<N, E> IndexMut<(N, N)> for GraphMap<N, E> where N: NodeTrait
{
    fn index_mut(&mut self, index: (N, N)) -> &mut E
    {
        self.edge_weight_mut(index.0, index.1).expect("GraphMap::index: no such edge")
    }
}
<|endoftext|>