<commit_before>use crate::content::Content;
use crate::{Document, Object, ObjectId};
use crate::{Error, Result};
use log::info;
use std::collections::BTreeMap;
use std::fs::File;
use std::io::Write;

impl Document {
	/// Change producer of document information dictionary.
	pub fn change_producer(&mut self, producer: &str) {
		if let Ok(info) = self.trailer.get_mut(b"Info") {
			if let Some(dict) = match *info {
				Object::Dictionary(ref mut dict) => Some(dict),
				Object::Reference(ref id) => self.objects.get_mut(id).and_then(|o| o.as_dict_mut().ok()),
				_ => None,
			} {
				dict.set("Producer", Object::string_literal(producer));
			}
		}
	}

	/// Compress PDF stream objects.
	pub fn compress(&mut self) {
		for object in self.objects.values_mut() {
			if let Object::Stream(ref mut stream) = *object {
				if stream.allows_compression {
					// Ignore any error and continue to compress other streams.
					let _ = stream.compress();
				}
			}
		}
	}

	/// Decompress PDF stream objects.
	pub fn decompress(&mut self) {
		for object in self.objects.values_mut() {
			if let Object::Stream(ref mut stream) = *object {
				stream.decompress()
			}
		}
	}

	/// Delete pages.
	pub fn delete_pages(&mut self, page_numbers: &[u32]) {
		let pages = self.get_pages();
		for page_number in page_numbers {
			if let Some(page) = pages.get(&page_number).and_then(|page_id| self.delete_object(*page_id)) {
				let mut page_tree_ref = page.as_dict().and_then(|dict| dict.get(b"Parent")).and_then(Object::as_reference);
				while let Ok(page_tree_id) = page_tree_ref {
					if let Some(page_tree) = self.objects.get_mut(&page_tree_id).and_then(|pt| pt.as_dict_mut().ok()) {
						if let Ok(count) = page_tree.get(b"Count").and_then(Object::as_i64) {
							page_tree.set("Count", count - 1);
						}
						page_tree_ref = page_tree.get(b"Parent").and_then(Object::as_reference);
					} else {
						break;
					}
				}
			}
		}
	}

	/// Prune all unused objects.
	pub fn prune_objects(&mut self) -> Vec<ObjectId> {
		let mut ids = vec![];
		let refs = self.traverse_objects(|_| {});
		for id in self.objects.keys().cloned().collect::<Vec<ObjectId>>() {
			if !refs.contains(&id) {
				self.objects.remove(&id);
				ids.push(id);
			}
		}
		ids
	}

	/// Delete object by object ID.
	pub fn delete_object(&mut self, id: ObjectId) -> Option<Object> {
		let action = |object: &mut Object| match *object {
			Object::Array(ref mut array) => {
				if let Some(index) = array.iter().position(|item: &Object| match *item {
					Object::Reference(ref_id) => ref_id == id,
					_ => false,
				}) {
					array.remove(index);
				}
			}
			Object::Dictionary(ref mut dict) => {
				let keys: Vec<Vec<u8>> = dict
					.iter()
					.filter(|&(_, item): &(&Vec<u8>, &Object)| match *item {
						Object::Reference(ref_id) => ref_id == id,
						_ => false,
					})
					.map(|(k, _)| k.clone())
					.collect();
				for key in keys {
					dict.remove(&key);
				}
			}
			_ => {}
		};
		self.traverse_objects(action);
		self.objects.remove(&id)
	}

	/// Delete zero length stream objects.
	pub fn delete_zero_length_streams(&mut self) -> Vec<ObjectId> {
		let mut ids = vec![];
		for id in self.objects.keys().cloned().collect::<Vec<ObjectId>>() {
			if self.objects.get(&id).and_then(|o| Object::as_stream(o).ok()).map(|stream| stream.content.is_empty()).unwrap_or(false) {
				self.delete_object(id);
				ids.push(id);
			}
		}
		ids
	}

	/// Renumber objects, normally called after delete_unused_objects.
	pub fn renumber_objects(&mut self) {
		let mut replace = BTreeMap::new();
		let mut new_id = 1;
		let mut ids = self.objects.keys().cloned().collect::<Vec<ObjectId>>();
		ids.sort();

		for id in ids {
			if id.0 != new_id {
				replace.insert(id, (new_id, id.1));
			}
			new_id += 1;
		}

		// replace order is from small to big
		for (old, new) in &replace {
			if let Some(object) = self.objects.remove(old) {
				self.objects.insert(new.clone(), object);
			}
		}

		let action = |object: &mut Object| if let Object::Reference(ref mut id) = *object {
			if replace.contains_key(&id) {
				*id = replace[id];
			}
		};

		self.traverse_objects(action);
		self.max_id = new_id - 1;
	}

	pub fn extract_text(&self, page_numbers: &[u32]) -> Result<String> {
		fn collect_text(text: &mut String, encoding: Option<&str>, operands: &[Object]) {
			for operand in operands.iter() {
				match *operand {
					Object::String(ref bytes, _) => {
						let decoded_text = Document::decode_text(encoding, bytes);
						text.push_str(&decoded_text);
					}
					Object::Array(ref arr) => {
						collect_text(text, encoding, arr);
					}
					_ => {}
				}
			}
		}
		let mut text = String::new();
		let pages = self.get_pages();
		for page_number in page_numbers {
			let page_id = pages[page_number];
			let fonts = self.get_page_fonts(page_id);
			let encodings = fonts.into_iter().map(|(name, font)| (name, font.get_font_encoding())).collect::<BTreeMap<Vec<u8>, &str>>();
			let content_data = self.get_page_content(page_id)?;
			let content = Content::decode(&content_data)?;
			let mut current_encoding = None;
			for operation in &content.operations {
				match operation.operator.as_ref() {
					"Tf" => {
						let current_font = operation.operands[0].as_name()?;
						current_encoding = encodings.get(current_font).cloned();
					}
					"Tj" | "TJ" => {
						collect_text(&mut text, current_encoding, &operation.operands);
					}
					"ET" => {
						if !text.ends_with('\n') {
							text.push('\n')
						}
					}
					_ => {}
				}
			}
		}
		Ok(text)
	}

	pub fn change_content_stream(&mut self, stream_id: ObjectId, content: Vec<u8>) {
		if let Some(content_stream) = self.objects.get_mut(&stream_id) {
			if let Object::Stream(ref mut stream) = *content_stream {
				stream.set_plain_content(content);
				// Ignore any compression error.
				let _ = stream.compress();
			}
		}
	}

	pub fn change_page_content(&mut self, page_id: ObjectId, content: Vec<u8>) -> Result<()> {
		let contents = self.get_dictionary(page_id).and_then(|page| page.get(b"Contents"))?;
		match *contents {
			Object::Reference(id) => self.change_content_stream(id, content),
			Object::Array(ref arr) => {
				if arr.len() == 1 {
					if let Ok(id) = arr[0].as_reference() { self.change_content_stream(id, content) }
				} else {
					let new_stream = self.add_object(super::Stream::new(dictionary!{}, content));
					if let Ok(page) = self.get_object_mut(page_id) {
						if let Object::Dictionary(ref mut dict) = *page {
							dict.set("Contents", new_stream);
						}
					}
				}
			}
			_ => {}
		}
		Ok(())
	}

	pub fn replace_text(&mut self, page_number: u32, text: &str, other_text: &str) -> Result<()>{
		let pages = self.get_pages();
		let page_id = *pages.get(&page_number).ok_or(Error::PageNumberNotFound(page_number))?;
		let encodings = self
			.get_page_fonts(page_id)
			.into_iter()
			.map(|(name, font)| (name, font.get_font_encoding().to_owned()))
			.collect::<BTreeMap<Vec<u8>, String>>();
		let content_data = self.get_page_content(page_id)?;
		let mut content = Content::decode(&content_data)?;
		let mut current_encoding = None;
		for operation in &mut content.operations {
			match operation.operator.as_ref() {
				"Tf" => {
					let current_font = operation.operands[0].as_name()?;
					current_encoding = encodings.get(current_font).map(std::string::String::as_str);
				}
				"Tj" => {
					for operand in &mut operation.operands {
						if let Object::String(ref mut bytes, _) = *operand {
							let decoded_text = Document::decode_text(current_encoding, bytes);
							info!("{}", decoded_text);
							if decoded_text == text {
								let encoded_bytes = Document::encode_text(current_encoding, other_text);
								*bytes = encoded_bytes;
							}
						}
					}
				}
				_ => {}
			}
		}
		let modified_content = content.encode()?;
		self.change_page_content(page_id, modified_content)
	}

	pub fn extract_stream(&self, stream_id: ObjectId, decompress: bool) -> Result<()> {
		let mut file = File::create(format!("{:?}.bin", stream_id))?;
		if let Ok(stream_obj) = self.get_object(stream_id) {
			if let Object::Stream(ref stream) = *stream_obj {
				if decompress {
					if let Ok(data) = stream.decompressed_content() {
						file.write_all(&data)?;
					} else {
						file.write_all(&stream.content)?;
					}
				} else {
					file.write_all(&stream.content)?;
				}
			}
		}
		Ok(())
	}
}
<commit_msg>Avoid allocating an intermediate collection for iteration.<commit_after>use crate::content::Content;
use crate::{Document, Object, ObjectId};
use crate::{Error, Result};
use log::info;
use std::collections::BTreeMap;
use std::fs::File;
use std::io::Write;

impl Document {
	/// Change producer of document information dictionary.
	pub fn change_producer(&mut self, producer: &str) {
		if let Ok(info) = self.trailer.get_mut(b"Info") {
			if let Some(dict) = match *info {
				Object::Dictionary(ref mut dict) => Some(dict),
				Object::Reference(ref id) => self.objects.get_mut(id).and_then(|o| o.as_dict_mut().ok()),
				_ => None,
			} {
				dict.set("Producer", Object::string_literal(producer));
			}
		}
	}

	/// Compress PDF stream objects.
	pub fn compress(&mut self) {
		for object in self.objects.values_mut() {
			if let Object::Stream(ref mut stream) = *object {
				if stream.allows_compression {
					// Ignore any error and continue to compress other streams.
					let _ = stream.compress();
				}
			}
		}
	}

	/// Decompress PDF stream objects.
	pub fn decompress(&mut self) {
		for object in self.objects.values_mut() {
			if let Object::Stream(ref mut stream) = *object {
				stream.decompress()
			}
		}
	}

	/// Delete pages.
	pub fn delete_pages(&mut self, page_numbers: &[u32]) {
		let pages = self.get_pages();
		for page_number in page_numbers {
			if let Some(page) = pages.get(&page_number).and_then(|page_id| self.delete_object(*page_id)) {
				let mut page_tree_ref = page.as_dict().and_then(|dict| dict.get(b"Parent")).and_then(Object::as_reference);
				while let Ok(page_tree_id) = page_tree_ref {
					if let Some(page_tree) = self.objects.get_mut(&page_tree_id).and_then(|pt| pt.as_dict_mut().ok()) {
						if let Ok(count) = page_tree.get(b"Count").and_then(Object::as_i64) {
							page_tree.set("Count", count - 1);
						}
						page_tree_ref = page_tree.get(b"Parent").and_then(Object::as_reference);
					} else {
						break;
					}
				}
			}
		}
	}

	/// Prune all unused objects.
	pub fn prune_objects(&mut self) -> Vec<ObjectId> {
		let mut ids = vec![];
		let refs = self.traverse_objects(|_| {});
		for id in self.objects.keys() {
			if !refs.contains(id) {
				ids.push(*id);
			}
		}

		for id in &ids {
			self.objects.remove(id);
		}

		ids
	}

	/// Delete object by object ID.
	pub fn delete_object(&mut self, id: ObjectId) -> Option<Object> {
		let action = |object: &mut Object| match *object {
			Object::Array(ref mut array) => {
				if let Some(index) = array.iter().position(|item: &Object| match *item {
					Object::Reference(ref_id) => ref_id == id,
					_ => false,
				}) {
					array.remove(index);
				}
			}
			Object::Dictionary(ref mut dict) => {
				let keys: Vec<Vec<u8>> = dict
					.iter()
					.filter(|&(_, item): &(&Vec<u8>, &Object)| match *item {
						Object::Reference(ref_id) => ref_id == id,
						_ => false,
					})
					.map(|(k, _)| k.clone())
					.collect();
				for key in keys {
					dict.remove(&key);
				}
			}
			_ => {}
		};
		self.traverse_objects(action);
		self.objects.remove(&id)
	}

	/// Delete zero length stream objects.
	pub fn delete_zero_length_streams(&mut self) -> Vec<ObjectId> {
		let mut ids = vec![];
		for id in self.objects.keys() {
			if self.objects.get(id).and_then(|o| Object::as_stream(o).ok()).map(|stream| stream.content.is_empty()).unwrap_or(false) {
				ids.push(*id);
			}
		}

		for id in &ids {
			self.delete_object(*id);
		}

		ids
	}

	/// Renumber objects, normally called after delete_unused_objects.
	pub fn renumber_objects(&mut self) {
		let mut replace = BTreeMap::new();
		let mut new_id = 1;
		let mut ids = self.objects.keys().cloned().collect::<Vec<ObjectId>>();
		ids.sort();

		for id in ids {
			if id.0 != new_id {
				replace.insert(id, (new_id, id.1));
			}
			new_id += 1;
		}

		// replace order is from small to big
		for (old, new) in &replace {
			if let Some(object) = self.objects.remove(old) {
				self.objects.insert(new.clone(), object);
			}
		}

		let action = |object: &mut Object| if let Object::Reference(ref mut id) = *object {
			if replace.contains_key(&id) {
				*id = replace[id];
			}
		};

		self.traverse_objects(action);
		self.max_id = new_id - 1;
	}

	pub fn extract_text(&self, page_numbers: &[u32]) -> Result<String> {
		fn collect_text(text: &mut String, encoding: Option<&str>, operands: &[Object]) {
			for operand in operands.iter() {
				match *operand {
					Object::String(ref bytes, _) => {
						let decoded_text = Document::decode_text(encoding, bytes);
						text.push_str(&decoded_text);
					}
					Object::Array(ref arr) => {
						collect_text(text, encoding, arr);
					}
					_ => {}
				}
			}
		}
		let mut text = String::new();
		let pages = self.get_pages();
		for page_number in page_numbers {
			let page_id = pages[page_number];
			let fonts = self.get_page_fonts(page_id);
			let encodings = fonts.into_iter().map(|(name, font)| (name, font.get_font_encoding())).collect::<BTreeMap<Vec<u8>, &str>>();
			let content_data = self.get_page_content(page_id)?;
			let content = Content::decode(&content_data)?;
			let mut current_encoding = None;
			for operation in &content.operations {
				match operation.operator.as_ref() {
					"Tf" => {
						let current_font = operation.operands[0].as_name()?;
						current_encoding = encodings.get(current_font).cloned();
					}
					"Tj" | "TJ" => {
						collect_text(&mut text, current_encoding, &operation.operands);
					}
					"ET" => {
						if !text.ends_with('\n') {
							text.push('\n')
						}
					}
					_ => {}
				}
			}
		}
		Ok(text)
	}

	pub fn change_content_stream(&mut self, stream_id: ObjectId, content: Vec<u8>) {
		if let Some(content_stream) = self.objects.get_mut(&stream_id) {
			if let Object::Stream(ref mut stream) = *content_stream {
				stream.set_plain_content(content);
				// Ignore any compression error.
				let _ = stream.compress();
			}
		}
	}

	pub fn change_page_content(&mut self, page_id: ObjectId, content: Vec<u8>) -> Result<()> {
		let contents = self.get_dictionary(page_id).and_then(|page| page.get(b"Contents"))?;
		match *contents {
			Object::Reference(id) => self.change_content_stream(id, content),
			Object::Array(ref arr) => {
				if arr.len() == 1 {
					if let Ok(id) = arr[0].as_reference() { self.change_content_stream(id, content) }
				} else {
					let new_stream = self.add_object(super::Stream::new(dictionary!{}, content));
					if let Ok(page) = self.get_object_mut(page_id) {
						if let Object::Dictionary(ref mut dict) = *page {
							dict.set("Contents", new_stream);
						}
					}
				}
			}
			_ => {}
		}
		Ok(())
	}

	pub fn replace_text(&mut self, page_number: u32, text: &str, other_text: &str) -> Result<()>{
		let pages = self.get_pages();
		let page_id = *pages.get(&page_number).ok_or(Error::PageNumberNotFound(page_number))?;
		let encodings = self
			.get_page_fonts(page_id)
			.into_iter()
			.map(|(name, font)| (name, font.get_font_encoding().to_owned()))
			.collect::<BTreeMap<Vec<u8>, String>>();
		let content_data = self.get_page_content(page_id)?;
		let mut content = Content::decode(&content_data)?;
		let mut current_encoding = None;
		for operation in &mut content.operations {
			match operation.operator.as_ref() {
				"Tf" => {
					let current_font = operation.operands[0].as_name()?;
					current_encoding = encodings.get(current_font).map(std::string::String::as_str);
				}
				"Tj" => {
					for operand in &mut operation.operands {
						if let Object::String(ref mut bytes, _) = *operand {
							let decoded_text = Document::decode_text(current_encoding, bytes);
							info!("{}", decoded_text);
							if decoded_text == text {
								let encoded_bytes = Document::encode_text(current_encoding, other_text);
								*bytes = encoded_bytes;
							}
						}
					}
				}
				_ => {}
			}
		}
		let modified_content = content.encode()?;
		self.change_page_content(page_id, modified_content)
	}

	pub fn extract_stream(&self, stream_id: ObjectId, decompress: bool) -> Result<()> {
		let mut file = File::create(format!("{:?}.bin", stream_id))?;
		if let Ok(stream_obj) = self.get_object(stream_id) {
			if let Object::Stream(ref stream) = *stream_obj {
				if decompress {
					if let Ok(data) = stream.decompressed_content() {
						file.write_all(&data)?;
					} else {
						file.write_all(&stream.content)?;
					}
				} else {
					file.write_all(&stream.content)?;
				}
			}
		}
		Ok(())
	}
}
<|endoftext|>


<commit_before>use std::cell::UnsafeCell;
use std::io;
#[cfg(target_os = "linux")]
use std::os::unix::io::AsRawFd;
#[cfg(target_os = "linux")]
use std::convert::From;
use std::sync::Arc;
use std::thread;
#[cfg(target_os = "linux")]
use std::mem;

use mio::{EventLoop, Evented, Handler, Token, EventSet, PollOpt};
use mio::util::Slab;
#[cfg(target_os = "linux")]
use mio::Io;

use mio::util::BoundedQueue;

use scheduler::{Scheduler, CoroutineRefMut};
use coroutine::{self, Coroutine, State, Handle};

thread_local!(static PROCESSOR: UnsafeCell<Processor> = UnsafeCell::new(Processor::new()));

pub struct Processor {
    event_loop: EventLoop<IoHandler>,
    work_queue: Arc<BoundedQueue<CoroutineRefMut>>,
    handler: IoHandler,
    main_coro: Handle,
    cur_running: Option<CoroutineRefMut>,
    last_result: Option<coroutine::Result<State>>,
}

impl Processor {
    pub fn new() -> Processor {
        let main_coro = unsafe {
            Coroutine::empty()
        };

        Processor {
            event_loop: EventLoop::new().unwrap(),
            work_queue: Scheduler::get().get_queue(),
            handler: IoHandler::new(),
            main_coro: main_coro,
            cur_running: None,
            last_result: None,
        }
    }

    pub fn running(&mut self) -> Option<CoroutineRefMut> {
        self.cur_running
    }

    pub fn current() -> &'static mut Processor {
        PROCESSOR.with(|p| unsafe { &mut *p.get() })
    }

    pub fn set_last_result(&mut self, r: coroutine::Result<State>) {
        self.last_result = Some(r);
    }

    pub fn schedule(&mut self) -> io::Result<()> {
        loop {
            match self.work_queue.pop() {
                Some(hdl) => {
                    match self.resume(hdl) {
                        Ok(State::Suspended) => {
                            Scheduler::ready(hdl);
                        },
                        Ok(State::Finished) | Ok(State::Panicked) => {
                            Scheduler::finished(hdl);
                        },
                        Ok(State::Blocked) => (),
                        Err(err) => {
                            error!("Coroutine resume failed, {:?}", err);
                            Scheduler::finished(hdl);
                        }
                    }
                },
                None => {
                    if self.handler.slabs.count() != 0 {
                        try!(self.event_loop.run_once(&mut self.handler));
                    } else if Scheduler::get().work_count() == 0 {
                        break;
                    } else {
                        thread::sleep_ms(100);
                    }
                }
            }
        }

        Ok(())
    }

    pub fn resume(&mut self, coro_ref: CoroutineRefMut) -> coroutine::Result<State> {
        self.cur_running = Some(coro_ref);
        unsafe {
            self.main_coro.yield_to(&mut *coro_ref.coro_ptr);
        }

        match self.last_result.take() {
            None => Ok(State::Suspended),
            Some(r) => r,
        }
    }

    pub fn sched(&mut self) {
        match self.cur_running.take() {
            None => {},
            Some(coro_ref) => unsafe {
                self.set_last_result(Ok(State::Suspended));
                (&mut *coro_ref.coro_ptr).yield_to(&*self.main_coro)
            }
        }
    }

    pub fn block(&mut self) {
        match self.cur_running.take() {
            None => {},
            Some(coro_ref) => unsafe {
                self.set_last_result(Ok(State::Blocked));
                (&mut *coro_ref.coro_ptr).yield_to(&*self.main_coro)
            }
        }
    }

    pub fn yield_with(&mut self, r: coroutine::Result<State>) {
        match self.cur_running.take() {
            None => {},
            Some(coro_ref) => unsafe {
                self.set_last_result(r);
                (&mut *coro_ref.coro_ptr).yield_to(&*self.main_coro)
            }
        }
    }
}

const MAX_TOKEN_NUM: usize = 102400;
impl IoHandler {
    fn new() -> IoHandler {
        IoHandler {
            slabs: Slab::new(MAX_TOKEN_NUM),
        }
    }
}

#[cfg(any(target_os = "linux",
          target_os = "android"))]
impl Processor {
    pub fn wait_event<E: Evented + AsRawFd>(&mut self, fd: &E, interest: EventSet) -> io::Result<()> {
        let token = self.handler.slabs.insert((Processor::current().running().clone(),
                                               From::from(fd.as_raw_fd()))).unwrap();
        try!(self.event_loop.register_opt(fd, token, interest,
                                          PollOpt::edge()|PollOpt::oneshot()));

        debug!("wait_event: Blocked current Coroutine ...; token={:?}", token);
        Scheduler::block();
        debug!("wait_event: Waked up; token={:?}", token);

        Ok(())
    }
}

#[cfg(any(target_os = "linux",
          target_os = "android"))]
struct IoHandler {
    slabs: Slab<(Handle, Io)>,
}

#[cfg(any(target_os = "linux",
          target_os = "android"))]
impl Handler for IoHandler {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Self>, token: Token, events: EventSet) {
        debug!("Got {:?} for {:?}", events, token);

        match self.slabs.remove(token) {
            Some((hdl, fd)) => {
                // Linux EPoll needs to explicit EPOLL_CTL_DEL the fd
                event_loop.deregister(&fd).unwrap();
                mem::forget(fd);
                Scheduler::ready(hdl);
            },
            None => {
                warn!("No coroutine is waiting on readable {:?}", token);
            }
        }
    }
}

#[cfg(any(target_os = "macos",
          target_os = "freebsd",
          target_os = "dragonfly",
          target_os = "ios",
          target_os = "bitrig",
          target_os = "openbsd"))]
impl Processor {
    pub fn wait_event<E: Evented>(&mut self, fd: &E, interest: EventSet) -> io::Result<()> {
        let token = self.handler.slabs.insert(Processor::current().running().unwrap()).unwrap();
        try!(self.event_loop.register_opt(fd, token, interest,
                                          PollOpt::edge()|PollOpt::oneshot()));

        debug!("wait_event: Blocked current Coroutine ...; token={:?}", token);
        // Coroutine::block();
        Scheduler::block();
        debug!("wait_event: Waked up; token={:?}", token);

        Ok(())
    }
}

#[cfg(any(target_os = "macos",
          target_os = "freebsd",
          target_os = "dragonfly",
          target_os = "ios",
          target_os = "bitrig",
          target_os = "openbsd"))]
struct IoHandler {
    slabs: Slab<CoroutineRefMut>,
}

#[cfg(any(target_os = "macos",
          target_os = "freebsd",
          target_os = "dragonfly",
          target_os = "ios",
          target_os = "bitrig",
          target_os = "openbsd"))]
impl Handler for IoHandler {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self, _: &mut EventLoop<Self>, token: Token, events: EventSet) {
        debug!("Got {:?} for {:?}", events, token);

        match self.slabs.remove(token) {
            Some(hdl) => {
                Scheduler::ready(hdl);
            },
            None => {
                warn!("No coroutine is waiting on readable {:?}", token);
            }
        }
    }
}
<commit_msg>fixed linux build<commit_after>use std::cell::UnsafeCell;
use std::io;
#[cfg(target_os = "linux")]
use std::os::unix::io::AsRawFd;
#[cfg(target_os = "linux")]
use std::convert::From;
use std::sync::Arc;
use std::thread;
#[cfg(target_os = "linux")]
use std::mem;

use mio::{EventLoop, Evented, Handler, Token, EventSet, PollOpt};
use mio::util::Slab;
#[cfg(target_os = "linux")]
use mio::Io;

use mio::util::BoundedQueue;

use scheduler::{Scheduler, CoroutineRefMut};
use coroutine::{self, Coroutine, State, Handle};

thread_local!(static PROCESSOR: UnsafeCell<Processor> = UnsafeCell::new(Processor::new()));

pub struct Processor {
    event_loop: EventLoop<IoHandler>,
    work_queue: Arc<BoundedQueue<CoroutineRefMut>>,
    handler: IoHandler,
    main_coro: Handle,
    cur_running: Option<CoroutineRefMut>,
    last_result: Option<coroutine::Result<State>>,
}

impl Processor {
    pub fn new() -> Processor {
        let main_coro = unsafe {
            Coroutine::empty()
        };

        Processor {
            event_loop: EventLoop::new().unwrap(),
            work_queue: Scheduler::get().get_queue(),
            handler: IoHandler::new(),
            main_coro: main_coro,
            cur_running: None,
            last_result: None,
        }
    }

    pub fn running(&mut self) -> Option<CoroutineRefMut> {
        self.cur_running
    }

    pub fn current() -> &'static mut Processor {
        PROCESSOR.with(|p| unsafe { &mut *p.get() })
    }

    pub fn set_last_result(&mut self, r: coroutine::Result<State>) {
        self.last_result = Some(r);
    }

    pub fn schedule(&mut self) -> io::Result<()> {
        loop {
            match self.work_queue.pop() {
                Some(hdl) => {
                    match self.resume(hdl) {
                        Ok(State::Suspended) => {
                            Scheduler::ready(hdl);
                        },
                        Ok(State::Finished) | Ok(State::Panicked) => {
                            Scheduler::finished(hdl);
                        },
                        Ok(State::Blocked) => (),
                        Err(err) => {
                            error!("Coroutine resume failed, {:?}", err);
                            Scheduler::finished(hdl);
                        }
                    }
                },
                None => {
                    if self.handler.slabs.count() != 0 {
                        try!(self.event_loop.run_once(&mut self.handler));
                    } else if Scheduler::get().work_count() == 0 {
                        break;
                    } else {
                        thread::sleep_ms(100);
                    }
                }
            }
        }

        Ok(())
    }

    pub fn resume(&mut self, coro_ref: CoroutineRefMut) -> coroutine::Result<State> {
        self.cur_running = Some(coro_ref);
        unsafe {
            self.main_coro.yield_to(&mut *coro_ref.coro_ptr);
        }

        match self.last_result.take() {
            None => Ok(State::Suspended),
            Some(r) => r,
        }
    }

    pub fn sched(&mut self) {
        match self.cur_running.take() {
            None => {},
            Some(coro_ref) => unsafe {
                self.set_last_result(Ok(State::Suspended));
                (&mut *coro_ref.coro_ptr).yield_to(&*self.main_coro)
            }
        }
    }

    pub fn block(&mut self) {
        match self.cur_running.take() {
            None => {},
            Some(coro_ref) => unsafe {
                self.set_last_result(Ok(State::Blocked));
                (&mut *coro_ref.coro_ptr).yield_to(&*self.main_coro)
            }
        }
    }

    pub fn yield_with(&mut self, r: coroutine::Result<State>) {
        match self.cur_running.take() {
            None => {},
            Some(coro_ref) => unsafe {
                self.set_last_result(r);
                (&mut *coro_ref.coro_ptr).yield_to(&*self.main_coro)
            }
        }
    }
}

const MAX_TOKEN_NUM: usize = 102400;
impl IoHandler {
    fn new() -> IoHandler {
        IoHandler {
            slabs: Slab::new(MAX_TOKEN_NUM),
        }
    }
}

#[cfg(any(target_os = "linux",
          target_os = "android"))]
impl Processor {
    pub fn wait_event<E: Evented + AsRawFd>(&mut self, fd: &E, interest: EventSet) -> io::Result<()> {
        let token = self.handler.slabs.insert((Processor::current().running().unwrap(),
                                               From::from(fd.as_raw_fd()))).unwrap();
        try!(self.event_loop.register_opt(fd, token, interest,
                                          PollOpt::edge()|PollOpt::oneshot()));

        debug!("wait_event: Blocked current Coroutine ...; token={:?}", token);
        Scheduler::block();
        debug!("wait_event: Waked up; token={:?}", token);

        Ok(())
    }
}

#[cfg(any(target_os = "linux",
          target_os = "android"))]
struct IoHandler {
    slabs: Slab<(CoroutineRefMut, Io)>,
}

#[cfg(any(target_os = "linux",
          target_os = "android"))]
impl Handler for IoHandler {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Self>, token: Token, events: EventSet) {
        debug!("Got {:?} for {:?}", events, token);

        match self.slabs.remove(token) {
            Some((hdl, fd)) => {
                // Linux EPoll needs to explicit EPOLL_CTL_DEL the fd
                event_loop.deregister(&fd).unwrap();
                mem::forget(fd);
                Scheduler::ready(hdl);
            },
            None => {
                warn!("No coroutine is waiting on readable {:?}", token);
            }
        }
    }
}

#[cfg(any(target_os = "macos",
          target_os = "freebsd",
          target_os = "dragonfly",
          target_os = "ios",
          target_os = "bitrig",
          target_os = "openbsd"))]
impl Processor {
    pub fn wait_event<E: Evented>(&mut self, fd: &E, interest: EventSet) -> io::Result<()> {
        let token = self.handler.slabs.insert(Processor::current().running().unwrap()).unwrap();
        try!(self.event_loop.register_opt(fd, token, interest,
                                          PollOpt::edge()|PollOpt::oneshot()));

        debug!("wait_event: Blocked current Coroutine ...; token={:?}", token);
        // Coroutine::block();
        Scheduler::block();
        debug!("wait_event: Waked up; token={:?}", token);

        Ok(())
    }
}

#[cfg(any(target_os = "macos",
          target_os = "freebsd",
          target_os = "dragonfly",
          target_os = "ios",
          target_os = "bitrig",
          target_os = "openbsd"))]
struct IoHandler {
    slabs: Slab<CoroutineRefMut>,
}

#[cfg(any(target_os = "macos",
          target_os = "freebsd",
          target_os = "dragonfly",
          target_os = "ios",
          target_os = "bitrig",
          target_os = "openbsd"))]
impl Handler for IoHandler {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self, _: &mut EventLoop<Self>, token: Token, events: EventSet) {
        debug!("Got {:?} for {:?}", events, token);

        match self.slabs.remove(token) {
            Some(hdl) => {
                Scheduler::ready(hdl);
            },
            None => {
                warn!("No coroutine is waiting on readable {:?}", token);
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Library service providers implementation.

extern crate hyper;

use hyper::client::{Client, Response};

/// Used to specify which provider to use to generate a short URL.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Provider {
    /// https://bn.gy provider
    BnGy,
    /// https://is.gd provider
    IsGd,
    /// http://psbe.co provider
    PsbeCo,
    /// http://readbility.com provider
    Rdd,
    /// http://rlu.ru provider
    /// * Attention! If you send a lot of requests from one IP, it can be
    /// blocked. If you plan to add more then 100 URLs in one hour, please let
    /// the technical support know. Otherwise your IP can be blocked
    /// unexpectedly. Prior added URLs can be deleted.
    Rlu,
    /// http://tinyurl.com provider
    /// * Note: this service does not provide any API.
    /// The implementation result depends on the service result web page.
    TinyUrl,
    /// https://v.gd provider
    VGd,
}

impl Provider {
    /// Converts the Provider variant into its domain name equivilant
    pub fn to_name(&self) -> &str {
        match *self {
            Provider::BnGy => "bn.gy",
            Provider::IsGd => "is.gd",
            Provider::PsbeCo => "psbe.co",
            Provider::Rdd => "readability.com",
            Provider::Rlu => "rlu.ru",
            Provider::TinyUrl => "tinyurl.com",
            Provider::VGd => "v.gd",
        }
    }
}

/// Returns a vector of all `Provider` variants.
///
/// The providers which are discouraged from use due to limitations such as
/// rate limitations are at the end of the resultant vector.
pub fn providers() -> Vec<Provider> {
    vec![
        Provider::IsGd,
        Provider::BnGy,
        Provider::VGd,
        Provider::Rdd,

        // The following list are items that are discouraged from use.
        // Reason: does not provide an api
        Provider::TinyUrl,
        // Reason: unstable work
        Provider::PsbeCo,
        // Reason: rate limit
        Provider::Rlu,
    ]
}

fn bngy_parse(res: &str) -> Option<String> {
    if res.is_empty() {
        return None
    }
    let string = res.to_owned();
    let iter = string.split("<ShortenedUrl>").skip(1).next();
    if iter.is_none() {
        return None
    }
    if let Some(string) = iter.unwrap().split("</ShortenedUrl>").next() {
        Some(string.to_owned())
    } else {
        None
    }
}

fn bngy_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("https://bn.gy/API.asmx/CreateUrl?real_url={}", url))
        .send()
        .ok()
}

fn isgd_parse(res: &str) -> Option<String> {
    Some(res.to_owned())
}

fn isgd_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("https://is.gd/create.php?format=simple&url={}", url))
        .send()
        .ok()
}

fn psbeco_parse(res: &str) -> Option<String> {
    if res.is_empty() {
        return None
    }
    let string = res.to_owned();
    let iter = string.split("<ShortUrl>").skip(1).next();
    if iter.is_none() {
        return None
    }
    if let Some(string) = iter.unwrap().split("</ShortUrl>").next() {
        Some(string.to_owned())
    } else {
        None
    }
}

fn psbeco_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("http://psbe.co/API.asmx/CreateUrl?real_url={}", url))
        .send()
        .ok()
}

fn rdd_parse(res: &str) -> Option<String> {
    if res.is_empty() {
        return None
    }
    let string = res.to_owned();
    let value = string.split("\"rdd_url\"")
                      .nth(1).unwrap_or("")
                      .split(",").next().unwrap_or("")
                      .split("\"").nth(1);
    if let Some(string) = value {
        let mut short_url = string.to_owned();
        let _ = short_url.pop();
        Some(short_url)
    } else {
        None
    }
}

fn rdd_request(url: &str, client: &Client) -> Option<Response> {
    client.post("https://readability.com/api/shortener/v1/urls")
        .body(&format!("url={}", url))
        .send()
        .ok()
}

fn rlu_parse(res: &str) -> Option<String> {
    Some(res.to_owned())
}

fn rlu_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("http://rlu.ru/index.sema?a=api&link={}", url))
        .send()
        .ok()
}

fn tinyurl_parse(res: &str) -> Option<String> {
    if res.is_empty() {
        return None
    }
    let string = res.to_owned();
    let value = string.split("data-clipboard-text=\"")
                      .nth(1).unwrap_or("")
                      .split("\">").next();
    if let Some(string) = value {
        Some(string.to_owned())
    } else {
        None
    }
}

fn tinyurl_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("http://tinyurl.com/create.php?url={}", url))
        .send()
        .ok()
}

fn vgd_parse(res: &str) -> Option<String> {
    Some(res.to_owned())
}

fn vgd_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("http://v.gd/create.php?format=simple&url={}", url))
        .send()
        .ok()
}

/// Parses the response from a successful request to a provider into the
/// URL-shortened string.
pub fn parse(res: &str, provider: Provider) -> Option<String> {
    match provider {
        Provider::BnGy => bngy_parse(res),
        Provider::IsGd => isgd_parse(res),
        Provider::PsbeCo => psbeco_parse(res),
        Provider::Rdd => rdd_parse(res),
        Provider::Rlu => rlu_parse(res),
        Provider::TinyUrl => tinyurl_parse(res),
        Provider::VGd => vgd_parse(res),
    }
}

/// Performs a request to the short link provider.
/// Response to be parsed or `None` on a error.
pub fn request(url: &str, client: &Client, provider: Provider) -> Option<Response> {
    match provider {
        Provider::BnGy => bngy_request(url, client),
        Provider::IsGd => isgd_request(url, client),
        Provider::PsbeCo => psbeco_request(url, client),
        Provider::Rdd => rdd_request(url, client),
        Provider::Rlu => rlu_request(url, client),
        Provider::TinyUrl => tinyurl_request(url, client),
        Provider::VGd => vgd_request(url, client),
    }
}
<commit_msg>HecSu provider implemented<commit_after>//! Library service providers implementation.

extern crate hyper;

use hyper::client::{Client, Response};

/// Used to specify which provider to use to generate a short URL.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Provider {
    /// https://bn.gy provider
    BnGy,
    /// https://hec.su provider
    ///
    /// * Limited to 3000 API requests per day
    HecSu,
    /// https://is.gd provider
    IsGd,
    /// http://psbe.co provider
    PsbeCo,
    /// http://readbility.com provider
    Rdd,
    /// http://rlu.ru provider
    ///
    /// * Attention! If you send a lot of requests from one IP, it can be
    /// blocked. If you plan to add more then 100 URLs in one hour, please let
    /// the technical support know. Otherwise your IP can be blocked
    /// unexpectedly. Prior added URLs can be deleted.
    Rlu,
    /// http://tinyurl.com provider
    /// * Note: this service does not provide any API.
    /// The implementation result depends on the service result web page.
    TinyUrl,
    /// https://v.gd provider
    VGd,
}

impl Provider {
    /// Converts the Provider variant into its domain name equivilant
    pub fn to_name(&self) -> &str {
        match *self {
            Provider::BnGy => "bn.gy",
            Provider::HecSu => "hec.su",
            Provider::IsGd => "is.gd",
            Provider::PsbeCo => "psbe.co",
            Provider::Rdd => "readability.com",
            Provider::Rlu => "rlu.ru",
            Provider::TinyUrl => "tinyurl.com",
            Provider::VGd => "v.gd",
        }
    }
}

/// Returns a vector of all `Provider` variants.
///
/// The providers which are discouraged from use due to limitations such as
/// rate limitations are at the end of the resultant vector.
pub fn providers() -> Vec<Provider> {
    vec![
        Provider::IsGd,
        Provider::BnGy,
        Provider::VGd,
        Provider::Rdd,

        // The following list are items that are discouraged from use.
        // Reason: rate limit (3000 requests per day)
        Provider::HecSu,
        // Reason: does not provide an api
        Provider::TinyUrl,
        // Reason: unstable work
        Provider::PsbeCo,
        // Reason: rate limit (100 requests per hour)
        Provider::Rlu,
    ]
}

fn bngy_parse(res: &str) -> Option<String> {
    if res.is_empty() {
        return None
    }
    let string = res.to_owned();
    let iter = string.split("<ShortenedUrl>").skip(1).next();
    if iter.is_none() {
        return None
    }
    if let Some(string) = iter.unwrap().split("</ShortenedUrl>").next() {
        Some(string.to_owned())
    } else {
        None
    }
}

fn bngy_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("https://bn.gy/API.asmx/CreateUrl?real_url={}", url))
        .send()
        .ok()
}

fn hecsu_parse(res: &str) -> Option<String> {
    if res.is_empty() {
        return None
    }
    let string = res.to_owned();
    let iter = string.split("<short>").skip(1).next();
    if iter.is_none() {
        return None
    }
    if let Some(string) = iter.unwrap().split("</short>").next() {
        Some(string.to_owned())
    } else {
        None
    }
}

fn hecsu_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("https://hec.su/api?url={}&method=xml", url))
        .send()
        .ok()
}

fn isgd_parse(res: &str) -> Option<String> {
    Some(res.to_owned())
}

fn isgd_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("https://is.gd/create.php?format=simple&url={}", url))
        .send()
        .ok()
}

fn psbeco_parse(res: &str) -> Option<String> {
    if res.is_empty() {
        return None
    }
    let string = res.to_owned();
    let iter = string.split("<ShortUrl>").skip(1).next();
    if iter.is_none() {
        return None
    }
    if let Some(string) = iter.unwrap().split("</ShortUrl>").next() {
        Some(string.to_owned())
    } else {
        None
    }
}

fn psbeco_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("http://psbe.co/API.asmx/CreateUrl?real_url={}", url))
        .send()
        .ok()
}

fn rdd_parse(res: &str) -> Option<String> {
    if res.is_empty() {
        return None
    }
    let string = res.to_owned();
    let value = string.split("\"rdd_url\"")
                      .nth(1).unwrap_or("")
                      .split(",").next().unwrap_or("")
                      .split("\"").nth(1);
    if let Some(string) = value {
        let mut short_url = string.to_owned();
        let _ = short_url.pop();
        Some(short_url)
    } else {
        None
    }
}

fn rdd_request(url: &str, client: &Client) -> Option<Response> {
    client.post("https://readability.com/api/shortener/v1/urls")
        .body(&format!("url={}", url))
        .send()
        .ok()
}

fn rlu_parse(res: &str) -> Option<String> {
    Some(res.to_owned())
}

fn rlu_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("http://rlu.ru/index.sema?a=api&link={}", url))
        .send()
        .ok()
}

fn tinyurl_parse(res: &str) -> Option<String> {
    if res.is_empty() {
        return None
    }
    let string = res.to_owned();
    let value = string.split("data-clipboard-text=\"")
                      .nth(1).unwrap_or("")
                      .split("\">").next();
    if let Some(string) = value {
        Some(string.to_owned())
    } else {
        None
    }
}

fn tinyurl_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("http://tinyurl.com/create.php?url={}", url))
        .send()
        .ok()
}

fn vgd_parse(res: &str) -> Option<String> {
    Some(res.to_owned())
}

fn vgd_request(url: &str, client: &Client) -> Option<Response> {
    client.get(&format!("http://v.gd/create.php?format=simple&url={}", url))
        .send()
        .ok()
}

/// Parses the response from a successful request to a provider into the
/// URL-shortened string.
pub fn parse(res: &str, provider: Provider) -> Option<String> {
    match provider {
        Provider::BnGy => bngy_parse(res),
        Provider::HecSu => hecsu_parse(res),
        Provider::IsGd => isgd_parse(res),
        Provider::PsbeCo => psbeco_parse(res),
        Provider::Rdd => rdd_parse(res),
        Provider::Rlu => rlu_parse(res),
        Provider::TinyUrl => tinyurl_parse(res),
        Provider::VGd => vgd_parse(res),
    }
}

/// Performs a request to the short link provider.
/// Response to be parsed or `None` on a error.
pub fn request(url: &str, client: &Client, provider: Provider) -> Option<Response> {
    match provider {
        Provider::BnGy => bngy_request(url, client),
        Provider::HecSu => hecsu_request(url, client),
        Provider::IsGd => isgd_request(url, client),
        Provider::PsbeCo => psbeco_request(url, client),
        Provider::Rdd => rdd_request(url, client),
        Provider::Rlu => rlu_request(url, client),
        Provider::TinyUrl => tinyurl_request(url, client),
        Provider::VGd => vgd_request(url, client),
    }
}
<|endoftext|>


<commit_before>extern crate csv;

use config::Config;
use chrono::{Utc, Timelike};
use feed::Feed;
use std::collections::HashMap;
use std::path::PathBuf;
use ::math;

error_chain! {
    foreign_links {
        Csv(csv::Error);
        ParseInt(::std::num::ParseIntError);
        ParseFloat(::std::num::ParseFloatError);
        Utf8Error(::std::string::FromUtf8Error);
        Io(::std::io::Error);
    }

    errors {
        TooFewRows {
            display("csv file contains record with too few rows")
        }
    }
}

type FeedID = u32;

/// An interface to save and load ListenerStats data.
pub struct AverageData {
    /// The path to the file to save and load data from.
    pub path: PathBuf,
    /// The data to save and load.
    pub data: HashMap<FeedID, ListenerStats>,
}

impl AverageData {
    pub fn new(path: PathBuf) -> AverageData {
        AverageData {
            path,
            data: HashMap::new(),
        }
    }

    pub fn load(&mut self) -> Result<()> {
        let hour = Utc::now().hour() as usize;
        let mut rdr = csv::Reader::from_path(&self.path)?;

        for result in rdr.records() {
            let record = result?;

            if record.len() < 1 + ListenerStats::HOURLY_SIZE {
                bail!(ErrorKind::TooFewRows);
            }

            let id = record[0].parse()?;
            let mut averages = [0.0; ListenerStats::HOURLY_SIZE];

            for i in 0..ListenerStats::HOURLY_SIZE {
                // Use an offset of 1 to avoid capturing the feed ID field
                averages[i] = record[1 + i].parse::<f32>()?;
            }

            let listeners = averages[hour] as i32;
            self.data.insert(id, ListenerStats::with_data(listeners, averages));
        }

        Ok(())
    }

    pub fn save(&self) -> Result<()> {
        let mut wtr = csv::Writer::from_path(&self.path)?;
        let mut fields = Vec::with_capacity(1 + ListenerStats::HOURLY_SIZE);

        for (id, stats) in &self.data {
            fields.push(id.to_string());

            for average in stats.average_hourly.iter() {
                fields.push(average.round().to_string());
            }

            wtr.write_record(&fields)?;
            fields.clear();
        }

        wtr.flush()?;
        Ok(())
    }
}

/// Represents an average set of data that wraps around its specified sample size.
#[derive(Debug, Clone)]
pub struct Average {
    /// The current average. It is updated after calling to self.add_sample().
    pub current: f32,
    /// The current average before the last call to self.add_sample().
    pub last: f32,
    /// The raw data that is used to calculate the current and last average.
    pub data: Vec<i32>,
    /// The current data index.
    index: usize,
    /// This keeps track of how many samples have been added since the struct
    /// was created, up until it reaches the specified sample size.
    /// It prevents data values that haven't been added from being averaged.
    populated: usize,
}

impl Average {
    pub fn new(sample_size: usize) -> Average {
        Average {
            current:   0.0,
            last:      0.0,
            data:      vec![0; sample_size],
            index:     0,
            populated: 0,
        }
    }

    /// Creates a new Average with one initial sample.
    pub fn with_value(sample_size: usize, value: i32) -> Average {
        let mut average = Average::new(sample_size);
        average.add_sample(value);

        average
    }

    /// Adds a new sample to the data and calculates the new average.
    pub fn add_sample(&mut self, value: i32) {
        if self.index >= self.data.len() {
            self.index = 0;
        }

        self.data[self.index] = value;
        self.index += 1;

        if self.populated < self.index {
            self.populated += 1;
        }

        self.last = self.current;

        self.current = self.data
            .iter()
            .take(self.populated)
            .sum::<i32>() as f32 / self.populated as f32;
    }
}

/// Represents general statistical data for feeds.
#[derive(Debug, Clone)]
pub struct ListenerStats {
    /// Represents the average number of listeners.
    pub average: Average,
    /// Represents the average number of listeners before a consistent spike occured.
    pub unskewed_average: Option<f32>,
    /// Contains the average number of listeners for any given hour.
    pub average_hourly: [f32; ListenerStats::HOURLY_SIZE],
    /// Indicates whether or not the listner count has spiked since the last update.
    pub has_spiked: bool,
    /// Represents the number of times the feed has spiked consecutively.
    pub spike_count: u32,
}

impl ListenerStats {
    const AVERAGE_SIZE: usize = 5;
    const HOURLY_SIZE:  usize = 24;

    pub fn new() -> ListenerStats {
        ListenerStats {
            average:          Average::new(ListenerStats::AVERAGE_SIZE),
            unskewed_average: None,
            average_hourly:   [0.0; ListenerStats::HOURLY_SIZE],
            has_spiked:       false,
            spike_count:      0,
        }
    }

    /// Creates a new ListenerStats struct with existing listener and hourly listener data.
    pub fn with_data(listeners: i32, hourly: [f32; ListenerStats::HOURLY_SIZE]) -> ListenerStats {
        ListenerStats {
            average:          Average::with_value(ListenerStats::AVERAGE_SIZE, listeners),
            unskewed_average: None,
            average_hourly:   hourly,
            has_spiked:       false,
            spike_count:      0,
        }
    }

    /// Updates the listener data and determines if the feed has spiked
    pub fn update(&mut self, hour: usize, feed: &Feed, config: &Config) {
        self.has_spiked = self.is_spiking(feed, config);

        self.spike_count = if self.has_spiked {
            self.spike_count + 1
        } else {
            0
        };

        self.average.add_sample(feed.listeners as i32);
        self.update_unskewed_average(feed.listeners as f32, config);
        self.average_hourly[hour] = self.get_unskewed_avg();
    }

    /// Returns true if the specified feed is currently spiking in listeners
    /// based off of previous data collected by self.update().
    fn is_spiking(&self, feed: &Feed, config: &Config) -> bool {
        if self.average.current == 0.0 {
            return false;
        }

        let spike = config.get_feed_spike(&feed);
        let listeners = feed.listeners as f32;

        // If a feed has a low number of listeners, use a higher threshold to
        // make the calculation less sensitive to very small listener jumps
        let threshold = if listeners < 50.0 {
            spike.jump + (50.0 - listeners) * spike.low_listener_increase
        } else {
            // Otherwise, use a lower threshold based off of how fast the feed's
            // listeners are rising to encourage more updates during large incidents
            let delta = self.get_jump(feed.listeners);
            let rise_amount = delta / spike.high_listener_dec_every * spike.high_listener_dec;

            spike.jump - rise_amount.min(spike.jump - 0.01)
        };

        listeners - self.average.current >= listeners * threshold
    }

    fn update_unskewed_average(&mut self, listeners: f32, config: &Config) {
        if let Some(unskewed) = self.unskewed_average {
            // Remove the unskewed average if the current average is close to it
            if self.average.current - unskewed < unskewed * config.unskewed_avg.reset_pcnt {
                self.unskewed_average = None;
                return
            }

            // Otherwise, if there isn't a huge jump in listeners, slowly increase
            // the unskewed average to adjust to natural listener increases
            if listeners - unskewed < unskewed * config.unskewed_avg.jump_required {
                self.unskewed_average = Some(math::lerp(
                    unskewed,
                    self.average.current,
                    config.unskewed_avg.adjust_pcnt
                ));
            }
        } else if self.has_spiked && self.average.last > 0.0 {
            // This is used to set the unskewed average if the listener count is
            // much higher than the average to avoid polluting the average listener
            // count with a very high value
            let has_large_jump = listeners > self.average.last * config.unskewed_avg.jump_required;

            let has_spiked_enough = self.spike_count > config.unskewed_avg.spikes_required;

            if has_spiked_enough || has_large_jump {
                self.unskewed_average = Some(self.average.last);
            }
        }
    }

    /// Returns the unskewed average if it is set, or the current average otherwise.
    pub fn get_unskewed_avg(&self) -> f32 {
        self.unskewed_average.unwrap_or(self.average.current)
    }

    /// Returns the difference in listeners from the unskewed average.
    pub fn get_jump(&self, listeners: u32) -> f32 {
        listeners as f32 - self.get_unskewed_avg()
    }
}<commit_msg>Fix nonexistent data being used in the first average sample (which was causing a LOT of feed updates on the first update)<commit_after>extern crate csv;

use config::Config;
use chrono::{Utc, Timelike};
use feed::Feed;
use std::collections::HashMap;
use std::path::PathBuf;
use ::math;

error_chain! {
    foreign_links {
        Csv(csv::Error);
        ParseInt(::std::num::ParseIntError);
        ParseFloat(::std::num::ParseFloatError);
        Utf8Error(::std::string::FromUtf8Error);
        Io(::std::io::Error);
    }

    errors {
        TooFewRows {
            display("csv file contains record with too few rows")
        }
    }
}

type FeedID = u32;

/// An interface to save and load ListenerStats data.
pub struct AverageData {
    /// The path to the file to save and load data from.
    pub path: PathBuf,
    /// The data to save and load.
    pub data: HashMap<FeedID, ListenerStats>,
}

impl AverageData {
    pub fn new(path: PathBuf) -> AverageData {
        AverageData {
            path,
            data: HashMap::new(),
        }
    }

    pub fn load(&mut self) -> Result<()> {
        let hour = Utc::now().hour() as usize;
        let mut rdr = csv::Reader::from_path(&self.path)?;

        for result in rdr.records() {
            let record = result?;

            if record.len() < 1 + ListenerStats::HOURLY_SIZE {
                bail!(ErrorKind::TooFewRows);
            }

            let id = record[0].parse()?;
            let mut averages = [0.0; ListenerStats::HOURLY_SIZE];

            for i in 0..ListenerStats::HOURLY_SIZE {
                // Use an offset of 1 to avoid capturing the feed ID field
                averages[i] = record[1 + i].parse::<f32>()?;
            }

            let listeners = averages[hour] as i32;

            // Zero listeners means that data for the current hour doesn't exist yet
            let stats = if listeners == 0 {
                ListenerStats::with_hourly(averages)
            } else {
                ListenerStats::with_data(listeners, averages)
            };

            self.data.insert(id, stats);
        }

        Ok(())
    }

    pub fn save(&self) -> Result<()> {
        let mut wtr = csv::Writer::from_path(&self.path)?;
        let mut fields = Vec::with_capacity(1 + ListenerStats::HOURLY_SIZE);

        for (id, stats) in &self.data {
            fields.push(id.to_string());

            for average in stats.average_hourly.iter() {
                fields.push(average.round().to_string());
            }

            wtr.write_record(&fields)?;
            fields.clear();
        }

        wtr.flush()?;
        Ok(())
    }
}

/// Represents an average set of data that wraps around its specified sample size.
#[derive(Debug, Clone)]
pub struct Average {
    /// The current average. It is updated after calling to self.add_sample().
    pub current: f32,
    /// The current average before the last call to self.add_sample().
    pub last: f32,
    /// The raw data that is used to calculate the current and last average.
    pub data: Vec<i32>,
    /// The current data index.
    index: usize,
    /// This keeps track of how many samples have been added since the struct
    /// was created, up until it reaches the specified sample size.
    /// It prevents data values that haven't been added from being averaged.
    populated: usize,
}

impl Average {
    pub fn new(sample_size: usize) -> Average {
        Average {
            current:   0.0,
            last:      0.0,
            data:      vec![0; sample_size],
            index:     0,
            populated: 0,
        }
    }

    /// Creates a new Average with one initial sample.
    pub fn with_value(sample_size: usize, value: i32) -> Average {
        let mut average = Average::new(sample_size);
        average.add_sample(value);

        average
    }

    /// Adds a new sample to the data and calculates the new average.
    pub fn add_sample(&mut self, value: i32) {
        if self.index >= self.data.len() {
            self.index = 0;
        }

        self.data[self.index] = value;
        self.index += 1;

        if self.populated < self.index {
            self.populated += 1;
        }

        self.last = self.current;

        self.current = self.data
            .iter()
            .take(self.populated)
            .sum::<i32>() as f32 / self.populated as f32;
    }
}

/// Represents general statistical data for feeds.
#[derive(Debug, Clone)]
pub struct ListenerStats {
    /// Represents the average number of listeners.
    pub average: Average,
    /// Represents the average number of listeners before a consistent spike occured.
    pub unskewed_average: Option<f32>,
    /// Contains the average number of listeners for any given hour.
    pub average_hourly: [f32; ListenerStats::HOURLY_SIZE],
    /// Indicates whether or not the listner count has spiked since the last update.
    pub has_spiked: bool,
    /// Represents the number of times the feed has spiked consecutively.
    pub spike_count: u32,
}

impl ListenerStats {
    const AVERAGE_SIZE: usize = 5;
    const HOURLY_SIZE:  usize = 24;

    pub fn new() -> ListenerStats {
        ListenerStats::with_hourly([0.0; ListenerStats::HOURLY_SIZE])
    }

    /// Creates a new ListenerStats struct with existing hourly data.
    pub fn with_hourly(hourly: [f32; ListenerStats::HOURLY_SIZE]) -> ListenerStats {
        ListenerStats {
            average:          Average::new(ListenerStats::AVERAGE_SIZE),
            unskewed_average: None,
            average_hourly:   hourly,
            has_spiked:       false,
            spike_count:      0,
        }
    }

    /// Creates a new ListenerStats struct with existing listener and hourly listener data.
    pub fn with_data(listeners: i32, hourly: [f32; ListenerStats::HOURLY_SIZE]) -> ListenerStats {
        ListenerStats {
            average:          Average::with_value(ListenerStats::AVERAGE_SIZE, listeners),
            unskewed_average: None,
            average_hourly:   hourly,
            has_spiked:       false,
            spike_count:      0,
        }
    }

    /// Updates the listener data and determines if the feed has spiked
    pub fn update(&mut self, hour: usize, feed: &Feed, config: &Config) {
        self.has_spiked = self.is_spiking(feed, config);

        self.spike_count = if self.has_spiked {
            self.spike_count + 1
        } else {
            0
        };

        self.average.add_sample(feed.listeners as i32);
        self.update_unskewed_average(feed.listeners as f32, config);
        self.average_hourly[hour] = self.get_unskewed_avg();
    }

    /// Returns true if the specified feed is currently spiking in listeners
    /// based off of previous data collected by self.update().
    fn is_spiking(&self, feed: &Feed, config: &Config) -> bool {
        if self.average.current == 0.0 {
            return false;
        }

        let spike = config.get_feed_spike(&feed);
        let listeners = feed.listeners as f32;

        // If a feed has a low number of listeners, use a higher threshold to
        // make the calculation less sensitive to very small listener jumps
        let threshold = if listeners < 50.0 {
            spike.jump + (50.0 - listeners) * spike.low_listener_increase
        } else {
            // Otherwise, use a lower threshold based off of how fast the feed's
            // listeners are rising to encourage more updates during large incidents
            let delta = self.get_jump(feed.listeners);
            let rise_amount = delta / spike.high_listener_dec_every * spike.high_listener_dec;

            spike.jump - rise_amount.min(spike.jump - 0.01)
        };

        listeners - self.average.current >= listeners * threshold
    }

    fn update_unskewed_average(&mut self, listeners: f32, config: &Config) {
        if let Some(unskewed) = self.unskewed_average {
            // Remove the unskewed average if the current average is close to it
            if self.average.current - unskewed < unskewed * config.unskewed_avg.reset_pcnt {
                self.unskewed_average = None;
                return
            }

            // Otherwise, if there isn't a huge jump in listeners, slowly increase
            // the unskewed average to adjust to natural listener increases
            if listeners - unskewed < unskewed * config.unskewed_avg.jump_required {
                self.unskewed_average = Some(math::lerp(
                    unskewed,
                    self.average.current,
                    config.unskewed_avg.adjust_pcnt
                ));
            }
        } else if self.has_spiked && self.average.last > 0.0 {
            // This is used to set the unskewed average if the listener count is
            // much higher than the average to avoid polluting the average listener
            // count with a very high value
            let has_large_jump = listeners > self.average.last * config.unskewed_avg.jump_required;

            let has_spiked_enough = self.spike_count > config.unskewed_avg.spikes_required;

            if has_spiked_enough || has_large_jump {
                self.unskewed_average = Some(self.average.last);
            }
        }
    }

    /// Returns the unskewed average if it is set, or the current average otherwise.
    pub fn get_unskewed_avg(&self) -> f32 {
        self.unskewed_average.unwrap_or(self.average.current)
    }

    /// Returns the difference in listeners from the unskewed average.
    pub fn get_jump(&self, listeners: u32) -> f32 {
        listeners as f32 - self.get_unskewed_avg()
    }
}<|endoftext|>


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

// Functions for dealing with devices.

use std::{collections::HashMap, fs::OpenOptions, path::Path};

use devicemapper::{devnode_to_devno, Device};

use crate::{
    engine::{
        strat_engine::backstore::{metadata::device_identifiers, util::get_udev_block_device},
        types::{DevUuid, PoolUuid},
    },
    stratis::{ErrorEnum, StratisError, StratisResult},
};

/// Resolve a list of Paths of some sort to a set of unique Devices.
/// Return an IOError if there was a problem resolving any particular device.
/// The set of devices maps each device to one of the paths passed.
/// Returns an error if any path does not correspond to a block device.
pub fn resolve_devices<'a>(paths: &'a [&Path]) -> StratisResult<HashMap<Device, &'a Path>> {
    let mut map = HashMap::new();
    for path in paths {
        match devnode_to_devno(path)? {
            Some(devno) => {
                let _ = map.insert(Device::from(devno), *path);
            }
            None => {
                let err_msg = format!("path {} does not refer to a block device", path.display());
                return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
            }
        }
    }
    Ok(map)
}

#[derive(Debug, PartialEq, Eq)]
pub enum DevOwnership {
    Ours(PoolUuid, DevUuid),
    Unowned,
    Theirs(String), // String is something useful to give back to end user about what's on device
}

impl DevOwnership {
    pub fn stratis_identifiers(&self) -> Option<(PoolUuid, DevUuid)> {
        match self {
            DevOwnership::Ours(pool_uuid, dev_uuid) => Some((*pool_uuid, *dev_uuid)),
            _ => None,
        }
    }
}

/// Returns true if a device has no signature and is not one of the paths of a multipath device,
/// yes this is a bit convoluted.  Logic gleaned from blivet library.
fn empty(device: &HashMap<String, String>) -> bool {
    device
        .get("DM_MULTIPATH_DEVICE_PATH")
        .map_or(true, |v| v != "1")
        && !((device.contains_key("ID_PART_TABLE_TYPE")
            && !device.contains_key("ID_PART_ENTRY_DISK"))
            || device.contains_key("ID_FS_USAGE"))
}

/// Generate some kind of human readable text about what's on a device.
fn signature(device: &HashMap<String, String>) -> String {
    if empty(device) {
        String::from("empty")
    } else {
        device
            .iter()
            .filter(|&(k, _)| k.contains("ID_FS_") || k.contains("ID_PART_TABLE_"))
            .map(|(k, v)| format!("{}={}", k, v))
            .collect::<Vec<String>>()
            .join(" ")
    }
}

/// Determine what a block device is used for.
pub fn identify(devnode: &Path) -> StratisResult<DevOwnership> {
    if let Some(device) = get_udev_block_device(devnode)? {
        if empty(&device) {
            // The device is either really empty or we are running on a distribution that hasn't
            // picked up the latest libblkid, lets read down to the device and find out for sure.
            // TODO: At some point in the future we can remove this and just return Unowned.
            if let Some((pool_uuid, device_uuid)) =
                device_identifiers(&mut OpenOptions::new().read(true).open(&devnode)?)?
            {
                Ok(DevOwnership::Ours(pool_uuid, device_uuid))
            } else {
                Ok(DevOwnership::Unowned)
            }
        } else if device
            .get("DM_MULTIPATH_DEVICE_PATH")
            .map_or(false, |v| v == "1")
        {
            Ok(DevOwnership::Theirs(String::from("multipath path")))
        } else if device.contains_key("ID_FS_TYPE") && device["ID_FS_TYPE"] == "stratis" {
            // Device is ours, but we don't get everything we need from udev db, lets go to disk.
            if let Some((pool_uuid, device_uuid)) =
                device_identifiers(&mut OpenOptions::new().read(true).open(&devnode)?)?
            {
                Ok(DevOwnership::Ours(pool_uuid, device_uuid))
            } else {
                // In this case the udev db says it's ours, but our check says otherwise.  We should
                // trust ourselves.  Should we raise an error here?
                Ok(DevOwnership::Theirs(String::from(
                    "Udev db says stratis, disk meta says no",
                )))
            }
        } else {
            Ok(DevOwnership::Theirs(signature(&device)))
        }
    } else {
        Err(StratisError::Engine(
            ErrorEnum::NotFound,
            format!(
                "We expected to find the block device {:?} \
                 in the udev db",
                devnode
            ),
        ))
    }
}

#[cfg(test)]
mod test {
    use std::path::Path;

    use crate::engine::strat_engine::{
        cmd,
        tests::{loopbacked, real},
    };

    use super::*;

    /// Verify that the device is not stratis by creating a device with XFS fs.
    fn test_other_ownership(paths: &[&Path]) {
        cmd::create_fs(paths[0], None).unwrap();

        cmd::udev_settle().unwrap();

        match identify(paths[0]).unwrap() {
            DevOwnership::Theirs(identity) => {
                assert!(identity.contains("ID_FS_USAGE=filesystem"));
                assert!(identity.contains("ID_FS_TYPE=xfs"));
                assert!(identity.contains("ID_FS_UUID"));
            }
            // This must fail, and will give a helpful error message
            id => assert_matches!(id, DevOwnership::Theirs(_)),
        }
    }

    /// Test a blank device and ensure it comes up as device::Usage::Unowned
    fn test_empty(paths: &[&Path]) {
        cmd::udev_settle().unwrap();
        assert_matches!(identify(paths[0]).unwrap(), DevOwnership::Unowned);
    }

    #[test]
    pub fn loop_test_device_other_ownership() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Range(1, 3, None),
            test_other_ownership,
        );
    }

    #[test]
    pub fn real_test_device_other_ownership() {
        real::test_with_spec(
            &real::DeviceLimits::AtLeast(1, None, None),
            test_other_ownership,
        );
    }

    #[test]
    pub fn loop_test_device_empty() {
        loopbacked::test_with_spec(&loopbacked::DeviceLimits::Range(1, 3, None), test_empty);
    }

    #[test]
    pub fn real_test_device_empty() {
        real::test_with_spec(&real::DeviceLimits::AtLeast(1, None, None), test_empty);
    }
}
<commit_msg>Fix a comment<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Functions for dealing with devices.

use std::{collections::HashMap, fs::OpenOptions, path::Path};

use devicemapper::{devnode_to_devno, Device};

use crate::{
    engine::{
        strat_engine::backstore::{metadata::device_identifiers, util::get_udev_block_device},
        types::{DevUuid, PoolUuid},
    },
    stratis::{ErrorEnum, StratisError, StratisResult},
};

/// Resolve a list of Paths of some sort to a set of unique Devices.
/// Return an IOError if there was a problem resolving any particular device.
/// The set of devices maps each device to one of the paths passed.
/// Returns an error if any path does not correspond to a block device.
pub fn resolve_devices<'a>(paths: &'a [&Path]) -> StratisResult<HashMap<Device, &'a Path>> {
    let mut map = HashMap::new();
    for path in paths {
        match devnode_to_devno(path)? {
            Some(devno) => {
                let _ = map.insert(Device::from(devno), *path);
            }
            None => {
                let err_msg = format!("path {} does not refer to a block device", path.display());
                return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
            }
        }
    }
    Ok(map)
}

#[derive(Debug, PartialEq, Eq)]
pub enum DevOwnership {
    Ours(PoolUuid, DevUuid),
    Unowned,
    Theirs(String), // String is something useful to give back to end user about what's on device
}

impl DevOwnership {
    pub fn stratis_identifiers(&self) -> Option<(PoolUuid, DevUuid)> {
        match self {
            DevOwnership::Ours(pool_uuid, dev_uuid) => Some((*pool_uuid, *dev_uuid)),
            _ => None,
        }
    }
}

/// Returns true if a device has no signature and is not one of the paths of a multipath device,
/// yes this is a bit convoluted.  Logic gleaned from blivet library.
fn empty(device: &HashMap<String, String>) -> bool {
    device
        .get("DM_MULTIPATH_DEVICE_PATH")
        .map_or(true, |v| v != "1")
        && !((device.contains_key("ID_PART_TABLE_TYPE")
            && !device.contains_key("ID_PART_ENTRY_DISK"))
            || device.contains_key("ID_FS_USAGE"))
}

/// Generate some kind of human readable text about what's on a device.
fn signature(device: &HashMap<String, String>) -> String {
    if empty(device) {
        String::from("empty")
    } else {
        device
            .iter()
            .filter(|&(k, _)| k.contains("ID_FS_") || k.contains("ID_PART_TABLE_"))
            .map(|(k, v)| format!("{}={}", k, v))
            .collect::<Vec<String>>()
            .join(" ")
    }
}

/// Determine what a block device is used for.
pub fn identify(devnode: &Path) -> StratisResult<DevOwnership> {
    if let Some(device) = get_udev_block_device(devnode)? {
        if empty(&device) {
            // The device is either really empty or we are running on a distribution that hasn't
            // picked up the latest libblkid, lets read down to the device and find out for sure.
            // TODO: At some point in the future we can remove this and just return Unowned.
            if let Some((pool_uuid, device_uuid)) =
                device_identifiers(&mut OpenOptions::new().read(true).open(&devnode)?)?
            {
                Ok(DevOwnership::Ours(pool_uuid, device_uuid))
            } else {
                Ok(DevOwnership::Unowned)
            }
        } else if device
            .get("DM_MULTIPATH_DEVICE_PATH")
            .map_or(false, |v| v == "1")
        {
            Ok(DevOwnership::Theirs(String::from("multipath path")))
        } else if device.contains_key("ID_FS_TYPE") && device["ID_FS_TYPE"] == "stratis" {
            // Device is ours, but we don't get everything we need from udev db, lets go to disk.
            if let Some((pool_uuid, device_uuid)) =
                device_identifiers(&mut OpenOptions::new().read(true).open(&devnode)?)?
            {
                Ok(DevOwnership::Ours(pool_uuid, device_uuid))
            } else {
                // In this case the udev db says it's ours, but our check says otherwise.  We should
                // trust ourselves.  Should we raise an error here?
                Ok(DevOwnership::Theirs(String::from(
                    "Udev db says stratis, disk meta says no",
                )))
            }
        } else {
            Ok(DevOwnership::Theirs(signature(&device)))
        }
    } else {
        Err(StratisError::Engine(
            ErrorEnum::NotFound,
            format!(
                "We expected to find the block device {:?} \
                 in the udev db",
                devnode
            ),
        ))
    }
}

#[cfg(test)]
mod test {
    use std::path::Path;

    use crate::engine::strat_engine::{
        cmd,
        tests::{loopbacked, real},
    };

    use super::*;

    /// Verify that the device is not stratis by creating a device with XFS fs.
    fn test_other_ownership(paths: &[&Path]) {
        cmd::create_fs(paths[0], None).unwrap();

        cmd::udev_settle().unwrap();

        match identify(paths[0]).unwrap() {
            DevOwnership::Theirs(identity) => {
                assert!(identity.contains("ID_FS_USAGE=filesystem"));
                assert!(identity.contains("ID_FS_TYPE=xfs"));
                assert!(identity.contains("ID_FS_UUID"));
            }
            // This must fail, and will give a helpful error message
            id => assert_matches!(id, DevOwnership::Theirs(_)),
        }
    }

    /// Verify that identify() recognizes a blank device as unowned
    fn test_empty(paths: &[&Path]) {
        cmd::udev_settle().unwrap();
        assert_matches!(identify(paths[0]).unwrap(), DevOwnership::Unowned);
    }

    #[test]
    pub fn loop_test_device_other_ownership() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Range(1, 3, None),
            test_other_ownership,
        );
    }

    #[test]
    pub fn real_test_device_other_ownership() {
        real::test_with_spec(
            &real::DeviceLimits::AtLeast(1, None, None),
            test_other_ownership,
        );
    }

    #[test]
    pub fn loop_test_device_empty() {
        loopbacked::test_with_spec(&loopbacked::DeviceLimits::Range(1, 3, None), test_empty);
    }

    #[test]
    pub fn real_test_device_empty() {
        real::test_with_spec(&real::DeviceLimits::AtLeast(1, None, None), test_empty);
    }
}
<|endoftext|>


<commit_before>use crate::deriving::generic::ty::*;
use crate::deriving::generic::*;

use rustc_ast::ptr::P;
use rustc_ast::EnumDef;
use rustc_ast::VariantData;
use rustc_ast::{Expr, MetaItem};
use rustc_errors::Applicability;
use rustc_expand::base::{Annotatable, DummyResult, ExtCtxt};
use rustc_span::symbol::Ident;
use rustc_span::symbol::{kw, sym};
use rustc_span::Span;
use smallvec::SmallVec;

pub fn expand_deriving_default(
    cx: &mut ExtCtxt<'_>,
    span: Span,
    mitem: &MetaItem,
    item: &Annotatable,
    push: &mut dyn FnMut(Annotatable),
) {
    let inline = cx.meta_word(span, sym::inline);
    let attrs = vec![cx.attribute(inline)];
    let trait_def = TraitDef {
        span,
        attributes: Vec::new(),
        path: Path::new(vec![kw::Default, sym::Default]),
        additional_bounds: Vec::new(),
        generics: Bounds::empty(),
        is_unsafe: false,
        supports_unions: false,
        methods: vec![MethodDef {
            name: kw::Default,
            generics: Bounds::empty(),
            explicit_self: None,
            args: Vec::new(),
            ret_ty: Self_,
            attributes: attrs,
            is_unsafe: false,
            unify_fieldless_variants: false,
            combine_substructure: combine_substructure(Box::new(|cx, trait_span, substr| {
                match substr.fields {
                    StaticStruct(_, fields) => {
                        default_struct_substructure(cx, trait_span, substr, fields)
                    }
                    StaticEnum(enum_def, _) => {
                        if !cx.sess.features_untracked().derive_default_enum {
                            rustc_session::parse::feature_err(
                                cx.parse_sess(),
                                sym::derive_default_enum,
                                span,
                                "deriving `Default` on enums is experimental",
                            )
                            .emit();
                        }
                        default_enum_substructure(cx, trait_span, enum_def)
                    }
                    _ => cx.span_bug(trait_span, "method in `derive(Default)`"),
                }
            })),
        }],
        associated_types: Vec::new(),
    };
    trait_def.expand(cx, mitem, item, push)
}

fn default_struct_substructure(
    cx: &mut ExtCtxt<'_>,
    trait_span: Span,
    substr: &Substructure<'_>,
    summary: &StaticFields,
) -> P<Expr> {
    // Note that `kw::Default` is "default" and `sym::Default` is "Default"!
    let default_ident = cx.std_path(&[kw::Default, sym::Default, kw::Default]);
    let default_call = |span| cx.expr_call_global(span, default_ident.clone(), Vec::new());

    match summary {
        Unnamed(ref fields, is_tuple) => {
            if !is_tuple {
                cx.expr_ident(trait_span, substr.type_ident)
            } else {
                let exprs = fields.iter().map(|sp| default_call(*sp)).collect();
                cx.expr_call_ident(trait_span, substr.type_ident, exprs)
            }
        }
        Named(ref fields) => {
            let default_fields = fields
                .iter()
                .map(|&(ident, span)| cx.field_imm(span, ident, default_call(span)))
                .collect();
            cx.expr_struct_ident(trait_span, substr.type_ident, default_fields)
        }
    }
}

fn default_enum_substructure(
    cx: &mut ExtCtxt<'_>,
    trait_span: Span,
    enum_def: &EnumDef,
) -> P<Expr> {
    let default_variant = match extract_default_variant(cx, enum_def, trait_span) {
        Ok(value) => value,
        Err(()) => return DummyResult::raw_expr(trait_span, true),
    };

    // At this point, we know that there is exactly one variant with a `#[default]` attribute. The
    // attribute hasn't yet been validated.

    if let Err(()) = validate_default_attribute(cx, default_variant) {
        return DummyResult::raw_expr(trait_span, true);
    }

    // We now know there is exactly one unit variant with exactly one `#[default]` attribute.

    cx.expr_path(cx.path(
        default_variant.span,
        vec![Ident::new(kw::SelfUpper, default_variant.span), default_variant.ident],
    ))
}

fn extract_default_variant<'a>(
    cx: &mut ExtCtxt<'_>,
    enum_def: &'a EnumDef,
    trait_span: Span,
) -> Result<&'a rustc_ast::Variant, ()> {
    let default_variants: SmallVec<[_; 1]> = enum_def
        .variants
        .iter()
        .filter(|variant| cx.sess.contains_name(&variant.attrs, kw::Default))
        .collect();

    let variant = match default_variants.as_slice() {
        [variant] => variant,
        [] => {
            cx.struct_span_err(trait_span, "no default declared")
                .help("make a unit variant default by placing `#[default]` above it")
                .emit();

            return Err(());
        }
        [first, rest @ ..] => {
            cx.struct_span_err(trait_span, "multiple declared defaults")
                .span_label(first.span, "first default")
                .span_labels(rest.iter().map(|variant| variant.span), "additional default")
                .note("only one variant can be default")
                .emit();

            return Err(());
        }
    };

    if !matches!(variant.data, VariantData::Unit(..)) {
        cx.struct_span_err(variant.ident.span, "`#[default]` may only be used on unit variants")
            .help("consider a manual implementation of `Default`")
            .emit();

        return Err(());
    }

    if let Some(non_exhaustive_attr) = cx.sess.find_by_name(&variant.attrs, sym::non_exhaustive) {
        cx.struct_span_err(variant.ident.span, "default variant must be exhaustive")
            .span_label(non_exhaustive_attr.span, "declared `#[non_exhaustive]` here")
            .help("consider a manual implementation of `Default`")
            .emit();

        return Err(());
    }

    Ok(variant)
}

fn validate_default_attribute(
    cx: &mut ExtCtxt<'_>,
    default_variant: &rustc_ast::Variant,
) -> Result<(), ()> {
    let attrs: SmallVec<[_; 1]> =
        cx.sess.filter_by_name(&default_variant.attrs, kw::Default).collect();

    let attr = match attrs.as_slice() {
        [attr] => attr,
        [] => cx.bug(
            "this method must only be called with a variant that has a `#[default]` attribute",
        ),
        [first, rest @ ..] => {
            // FIXME(jhpratt) Do we want to perform this check? It doesn't exist
            // for `#[inline]`, `#[non_exhaustive]`, and presumably others.

            let suggestion_text =
                if rest.len() == 1 { "try removing this" } else { "try removing these" };

            cx.struct_span_err(default_variant.ident.span, "multiple `#[default]` attributes")
                .note("only one `#[default]` attribute is needed")
                .span_label(first.span, "`#[default]` used here")
                .span_label(rest[0].span, "`#[default]` used again here")
                .span_help(rest.iter().map(|attr| attr.span).collect::<Vec<_>>(), suggestion_text)
                // This would otherwise display the empty replacement, hence the otherwise
                // repetitive `.span_help` call above.
                .tool_only_multipart_suggestion(
                    suggestion_text,
                    rest.iter().map(|attr| (attr.span, String::new())).collect(),
                    Applicability::MachineApplicable,
                )
                .emit();

            return Err(());
        }
    };
    if !attr.is_word() {
        cx.struct_span_err(attr.span, "`#[default]` attribute does not accept a value")
            .span_suggestion_hidden(
                attr.span,
                "try using `#[default]`",
                "#[default]".into(),
                Applicability::MaybeIncorrect,
            )
            .emit();

        return Err(());
    }
    Ok(())
}
<commit_msg>Add machine-applicable suggestions<commit_after>use crate::deriving::generic::ty::*;
use crate::deriving::generic::*;

use rustc_ast::ptr::P;
use rustc_ast::EnumDef;
use rustc_ast::VariantData;
use rustc_ast::{Expr, MetaItem};
use rustc_errors::Applicability;
use rustc_expand::base::{Annotatable, DummyResult, ExtCtxt};
use rustc_span::symbol::Ident;
use rustc_span::symbol::{kw, sym};
use rustc_span::Span;
use smallvec::SmallVec;

pub fn expand_deriving_default(
    cx: &mut ExtCtxt<'_>,
    span: Span,
    mitem: &MetaItem,
    item: &Annotatable,
    push: &mut dyn FnMut(Annotatable),
) {
    let inline = cx.meta_word(span, sym::inline);
    let attrs = vec![cx.attribute(inline)];
    let trait_def = TraitDef {
        span,
        attributes: Vec::new(),
        path: Path::new(vec![kw::Default, sym::Default]),
        additional_bounds: Vec::new(),
        generics: Bounds::empty(),
        is_unsafe: false,
        supports_unions: false,
        methods: vec![MethodDef {
            name: kw::Default,
            generics: Bounds::empty(),
            explicit_self: None,
            args: Vec::new(),
            ret_ty: Self_,
            attributes: attrs,
            is_unsafe: false,
            unify_fieldless_variants: false,
            combine_substructure: combine_substructure(Box::new(|cx, trait_span, substr| {
                match substr.fields {
                    StaticStruct(_, fields) => {
                        default_struct_substructure(cx, trait_span, substr, fields)
                    }
                    StaticEnum(enum_def, _) => {
                        if !cx.sess.features_untracked().derive_default_enum {
                            rustc_session::parse::feature_err(
                                cx.parse_sess(),
                                sym::derive_default_enum,
                                span,
                                "deriving `Default` on enums is experimental",
                            )
                            .emit();
                        }
                        default_enum_substructure(cx, trait_span, enum_def)
                    }
                    _ => cx.span_bug(trait_span, "method in `derive(Default)`"),
                }
            })),
        }],
        associated_types: Vec::new(),
    };
    trait_def.expand(cx, mitem, item, push)
}

fn default_struct_substructure(
    cx: &mut ExtCtxt<'_>,
    trait_span: Span,
    substr: &Substructure<'_>,
    summary: &StaticFields,
) -> P<Expr> {
    // Note that `kw::Default` is "default" and `sym::Default` is "Default"!
    let default_ident = cx.std_path(&[kw::Default, sym::Default, kw::Default]);
    let default_call = |span| cx.expr_call_global(span, default_ident.clone(), Vec::new());

    match summary {
        Unnamed(ref fields, is_tuple) => {
            if !is_tuple {
                cx.expr_ident(trait_span, substr.type_ident)
            } else {
                let exprs = fields.iter().map(|sp| default_call(*sp)).collect();
                cx.expr_call_ident(trait_span, substr.type_ident, exprs)
            }
        }
        Named(ref fields) => {
            let default_fields = fields
                .iter()
                .map(|&(ident, span)| cx.field_imm(span, ident, default_call(span)))
                .collect();
            cx.expr_struct_ident(trait_span, substr.type_ident, default_fields)
        }
    }
}

fn default_enum_substructure(
    cx: &mut ExtCtxt<'_>,
    trait_span: Span,
    enum_def: &EnumDef,
) -> P<Expr> {
    let default_variant = match extract_default_variant(cx, enum_def, trait_span) {
        Ok(value) => value,
        Err(()) => return DummyResult::raw_expr(trait_span, true),
    };

    // At this point, we know that there is exactly one variant with a `#[default]` attribute. The
    // attribute hasn't yet been validated.

    if let Err(()) = validate_default_attribute(cx, default_variant) {
        return DummyResult::raw_expr(trait_span, true);
    }

    // We now know there is exactly one unit variant with exactly one `#[default]` attribute.

    cx.expr_path(cx.path(
        default_variant.span,
        vec![Ident::new(kw::SelfUpper, default_variant.span), default_variant.ident],
    ))
}

fn extract_default_variant<'a>(
    cx: &mut ExtCtxt<'_>,
    enum_def: &'a EnumDef,
    trait_span: Span,
) -> Result<&'a rustc_ast::Variant, ()> {
    let default_variants: SmallVec<[_; 1]> = enum_def
        .variants
        .iter()
        .filter(|variant| cx.sess.contains_name(&variant.attrs, kw::Default))
        .collect();

    let variant = match default_variants.as_slice() {
        [variant] => variant,
        [] => {
            let possible_defaults = enum_def
                .variants
                .iter()
                .filter(|variant| matches!(variant.data, VariantData::Unit(..)))
                .filter(|variant| !cx.sess.contains_name(&variant.attrs, sym::non_exhaustive));

            let mut diag = cx.struct_span_err(trait_span, "no default declared");
            diag.help("make a unit variant default by placing `#[default]` above it");
            for variant in possible_defaults {
                // Suggest making each unit variant default.
                diag.tool_only_span_suggestion(
                    variant.span,
                    &format!("make `{}` default", variant.ident),
                    format!("#[default] {}", variant.ident),
                    Applicability::MaybeIncorrect,
                );
            }
            diag.emit();

            return Err(());
        }
        [first, rest @ ..] => {
            let mut diag = cx.struct_span_err(trait_span, "multiple declared defaults");
            diag.span_label(first.span, "first default");
            diag.span_labels(rest.iter().map(|variant| variant.span), "additional default");
            diag.note("only one variant can be default");
            for variant in &default_variants {
                // Suggest making each variant already tagged default.
                let suggestion = default_variants
                    .iter()
                    .filter_map(|v| {
                        if v.ident == variant.ident {
                            None
                        } else {
                            Some((cx.sess.find_by_name(&v.attrs, kw::Default)?.span, String::new()))
                        }
                    })
                    .collect();

                diag.tool_only_multipart_suggestion(
                    &format!("make `{}` default", variant.ident),
                    suggestion,
                    Applicability::MaybeIncorrect,
                );
            }
            diag.emit();

            return Err(());
        }
    };

    if !matches!(variant.data, VariantData::Unit(..)) {
        cx.struct_span_err(variant.ident.span, "`#[default]` may only be used on unit variants")
            .help("consider a manual implementation of `Default`")
            .emit();

        return Err(());
    }

    if let Some(non_exhaustive_attr) = cx.sess.find_by_name(&variant.attrs, sym::non_exhaustive) {
        cx.struct_span_err(variant.ident.span, "default variant must be exhaustive")
            .span_label(non_exhaustive_attr.span, "declared `#[non_exhaustive]` here")
            .help("consider a manual implementation of `Default`")
            .emit();

        return Err(());
    }

    Ok(variant)
}

fn validate_default_attribute(
    cx: &mut ExtCtxt<'_>,
    default_variant: &rustc_ast::Variant,
) -> Result<(), ()> {
    let attrs: SmallVec<[_; 1]> =
        cx.sess.filter_by_name(&default_variant.attrs, kw::Default).collect();

    let attr = match attrs.as_slice() {
        [attr] => attr,
        [] => cx.bug(
            "this method must only be called with a variant that has a `#[default]` attribute",
        ),
        [first, rest @ ..] => {
            // FIXME(jhpratt) Do we want to perform this check? It doesn't exist
            // for `#[inline]`, `#[non_exhaustive]`, and presumably others.

            let suggestion_text =
                if rest.len() == 1 { "try removing this" } else { "try removing these" };

            cx.struct_span_err(default_variant.ident.span, "multiple `#[default]` attributes")
                .note("only one `#[default]` attribute is needed")
                .span_label(first.span, "`#[default]` used here")
                .span_label(rest[0].span, "`#[default]` used again here")
                .span_help(rest.iter().map(|attr| attr.span).collect::<Vec<_>>(), suggestion_text)
                // This would otherwise display the empty replacement, hence the otherwise
                // repetitive `.span_help` call above.
                .tool_only_multipart_suggestion(
                    suggestion_text,
                    rest.iter().map(|attr| (attr.span, String::new())).collect(),
                    Applicability::MachineApplicable,
                )
                .emit();

            return Err(());
        }
    };
    if !attr.is_word() {
        cx.struct_span_err(attr.span, "`#[default]` attribute does not accept a value")
            .span_suggestion_hidden(
                attr.span,
                "try using `#[default]`",
                "#[default]".into(),
                Applicability::MaybeIncorrect,
            )
            .emit();

        return Err(());
    }
    Ok(())
}
<|endoftext|>


<commit_before>use std;
import int;

tag tree { nil; node(~tree, ~tree, int); }

fn item_check(t: ~tree) -> int {
    alt *t {
      nil. { ret 0; }
      node(left, right, item) {
        ret item + item_check(left) - item_check(right);
      }
    }
}

fn bottom_up_tree(item: int, depth: int) -> ~tree {
    if depth > 0 {
        ret ~node(bottom_up_tree(2 * item - 1, depth - 1),
                  bottom_up_tree(2 * item, depth - 1), item);
    } else { ret ~nil; }
}

fn main() {
    let n = 8;
    let min_depth = 4;
    let max_depth;
    if min_depth + 2 > n {
        max_depth = min_depth + 2;
    } else { max_depth = n; }
    let stretch_depth = max_depth + 1;
    let stretch_tree = bottom_up_tree(0, stretch_depth);
    #debug("stretch tree of depth %d\t check: %d", stretch_depth,
           item_check(stretch_tree));
    let long_lived_tree = bottom_up_tree(0, max_depth);
    let depth = min_depth;
    while depth <= max_depth {
        let iterations = int::pow(2, max_depth - depth + min_depth as uint);
        let chk = 0;
        let i = 1;
        while i <= iterations {
            let temp_tree = bottom_up_tree(i, depth);
            chk += item_check(temp_tree);
            temp_tree = bottom_up_tree(-i, depth);
            chk += item_check(temp_tree);
            i += 1;
        }
        #debug("%d\t trees of depth %d\t check: %d", iterations * 2, depth,
               chk);
        depth += 2;
    }
    #debug("long lived trees of depth %d\t check: %d", max_depth,
           item_check(long_lived_tree));
}
<commit_msg>bench: Update shootout-binarytrees for style<commit_after>use std;
import int;

tag tree { nil; node(~tree, ~tree, int); }

fn item_check(t: ~tree) -> int {
    alt *t {
      nil. { ret 0; }
      node(left, right, item) {
        ret item + item_check(left) - item_check(right);
      }
    }
}

fn bottom_up_tree(item: int, depth: int) -> ~tree {
    if depth > 0 {
        ret ~node(bottom_up_tree(2 * item - 1, depth - 1),
                  bottom_up_tree(2 * item, depth - 1), item);
    } else { ret ~nil; }
}

fn main(args: [str]) {
    let n = if vec::len(args) == 2u {
        int::from_str(args[1])
    } else {
        15
    };
    let min_depth = 4;
    let max_depth;
    if min_depth + 2 > n {
        max_depth = min_depth + 2;
    } else { max_depth = n; }
    let stretch_depth = max_depth + 1;
    let stretch_tree = bottom_up_tree(0, stretch_depth);
    std::io::println(#fmt("stretch tree of depth %d\t check: %d",
                          stretch_depth,
                          item_check(stretch_tree)));
    let long_lived_tree = bottom_up_tree(0, max_depth);
    let depth = min_depth;
    while depth <= max_depth {
        let iterations = int::pow(2, max_depth - depth + min_depth as uint);
        let chk = 0;
        let i = 1;
        while i <= iterations {
            let temp_tree = bottom_up_tree(i, depth);
            chk += item_check(temp_tree);
            temp_tree = bottom_up_tree(-i, depth);
            chk += item_check(temp_tree);
            i += 1;
        }
        std::io::println(#fmt("%d\t trees of depth %d\t check: %d",
                               iterations * 2, depth,
                               chk));
        depth += 2;
    }
    std::io::println(#fmt("long lived trees of depth %d\t check: %d",
                          max_depth,
                          item_check(long_lived_tree)));
}
<|endoftext|>


<commit_before>#![feature(cfg_target_thread_local, thread_local_internals)]

// On platforms *without* `#[thread_local]`, use
// a custom non-`Sync` type to fake the same error.
#[cfg(not(target_thread_local))]
struct Key<T> {
    _data: std::cell::UnsafeCell<Option<T>>,
    _flag: std::cell::Cell<bool>,
}

#[cfg(not(target_thread_local))]
impl<T> Key<T> {
    const fn new() -> Self {
        Key {
            _data: std::cell::UnsafeCell::new(None),
            _flag: std::cell::Cell::new(false),
        }
    }
}

#[cfg(target_thread_local)]
use std::thread::__FastLocalKeyInner as Key;

static __KEY: Key<()> = Key::new();
//~^ ERROR `std::cell::UnsafeCell<std::option::Option<()>>` cannot be shared between threads
//~| ERROR `std::cell::Cell<std::thread::local::fast::DtorState>` cannot be shared between threads safely [E0277]

fn main() {}
<commit_msg>clang tidy fixes<commit_after>#![feature(cfg_target_thread_local, thread_local_internals)]

// On platforms *without* `#[thread_local]`, use
// a custom non-`Sync` type to fake the same error.
#[cfg(not(target_thread_local))]
struct Key<T> {
    _data: std::cell::UnsafeCell<Option<T>>,
    _flag: std::cell::Cell<bool>,
}

#[cfg(not(target_thread_local))]
impl<T> Key<T> {
    const fn new() -> Self {
        Key {
            _data: std::cell::UnsafeCell::new(None),
            _flag: std::cell::Cell::new(false),
        }
    }
}

#[cfg(target_thread_local)]
use std::thread::__FastLocalKeyInner as Key;

static __KEY: Key<()> = Key::new();
//~^ ERROR `std::cell::UnsafeCell<std::option::Option<()>>` cannot be shared between threads
//~| ERROR `std::cell::Cell<std::thread::local::fast::DtorState>` cannot be shared between threads
// safely [E0277]

fn main() {}
<|endoftext|>