<commit_before>use wrapped2d::b2;
use wrapped2d::user_data::*;
use wrapped2d::dynamics::world::callbacks::ContactAccess;
use backend::obj;
use backend::obj::{Solid, Geometry, Transformable};
use backend::world;
use backend::world::Intent;
use std::collections::HashMap;
use std::collections::HashSet;
use std::f32::consts;
use std::rc::Rc;
use std::cell::RefCell;
use core::geometry::*;
use super::*;

struct AgentData;

impl UserDataTypes for AgentData {
	type BodyData = world::AgentRefs;
	type JointData = ();
	type FixtureData = world::AgentRefs;
}

pub struct PhysicsSystem {
	world: b2::World<AgentData>,
	handles: HashMap<world::AgentRefs, b2::BodyHandle>,
	touched: Rc<RefCell<HashSet<world::AgentRefs>>>,
}

impl Updateable for PhysicsSystem {
	fn update(&mut self, state: &world::WorldState, dt: f32) {
		let mut forces = Vec::new();
		let mut impulses = Vec::new();

		for (h, b) in self.world.bodies() {
			let body = b.borrow();
			let center = (*body).world_center().clone();
			let key = (*body).user_data();
			if let Some(segment) = state.minion(key.agent_id).and_then(|c| c.segment(key.segment_index)) {
				match segment.state.intent {
					Intent::Move(force) => forces.push((h, center, force)),
					Intent::RunAway(impulse) => impulses.push((h, center, impulse)),
					Intent::Idle => {}
				}
			}
		}
		for (h, center, force) in forces {
			let b = &mut self.world.body_mut(h);
			b.apply_force(&b2::Vec2 {
				              x: force.x,
				              y: force.y,
			              },
			              &center,
			              true);
		}

		for (h, center, impulse) in impulses {
			let b = &mut self.world.body_mut(h);
			b.apply_linear_impulse(&b2::Vec2 {
				                       x: impulse.x,
				                       y: impulse.y,
			                       },
			                       &center,
			                       true);
		}

		self.world.step(dt, 8, 3);
	}
}

struct JointRef<'a> {
	refs: world::AgentRefs,
	handle: b2::BodyHandle,
	mesh: &'a obj::Mesh,
	flags: world::SegmentFlags,
	attachment: Option<world::Attachment>,
}

impl System for PhysicsSystem {
	fn register(&mut self, agent: &world::Agent) {
		// build fixtures
		let joint_refs = PhysicsSystem::build_fixtures(&mut self.world, &agent);
		// and then assemble them with joints
		PhysicsSystem::build_joints(&mut self.world, &joint_refs);
		// record them
		for JointRef { refs, handle, .. } in joint_refs {
			self.handles.insert(refs, handle);
		}
	}

	fn init(&mut self, world: &world::World) {
		self.init_extent(&world.extent);
	}

	fn to_world(&self, world: &mut world::World) {
		for (_, b) in self.world.bodies() {
			let body = b.borrow();
			let position = (*body).position();
			let angle = (*body).angle();
			let key = (*body).user_data();

			if let Some(agent) = world.minions.get_mut(key.agent_id) {
				if let Some(segment) = agent.segment_mut(key.segment_index) {
					let t = segment.transform();
					segment.transform_to(Transform {
						position: PhysicsSystem::from_vec2(&position),
						angle: angle,
						..t
					});
					segment.state.collision_detected = self.touched.borrow().contains(key);
				}
			}
		}
		self.touched.borrow_mut().clear();
	}
}

impl PhysicsSystem {
	pub fn new() -> Self {
		let touched = Rc::new(RefCell::new(HashSet::new()));
		PhysicsSystem {
			world: Self::new_world(touched.clone()),
			handles: HashMap::new(),
			touched: touched,
		}
	}

	fn to_vec2(p: &Position) -> b2::Vec2 {
		b2::Vec2 { x: p.x, y: p.y }
	}

	fn from_vec2(p: &b2::Vec2) -> Position {
		Position::new(p.x, p.y)
	}

	fn init_extent(&mut self, extent: &Rect) {
		let mut f_def = b2::FixtureDef::new();
		// f_def.density = material.density;
		// f_def.restitution = material.restitution;
		// f_def.friction = material.friction;

		let mut b_def = b2::BodyDef::new();
		b_def.body_type = b2::BodyType::Static;
		let refs = world::AgentRefs::with_id(0xFFFFFFFFusize);
		let handle = self.world.create_body_with(&b_def, refs);

		let mut rect = b2::ChainShape::new();
		rect.create_chain(&[Self::to_vec2(&extent.bottom_left()),
		                    Self::to_vec2(&extent.bottom_right()),
		                    Self::to_vec2(&extent.top_right()),
		                    Self::to_vec2(&extent.top_left())]);
		self.world.body_mut(handle).create_fixture_with(&rect, &mut f_def, refs);
	}

	fn build_fixtures<'a>(world: &mut b2::World<AgentData>, agent: &'a world::Agent) -> Vec<JointRef<'a>> {
		let object_id = agent.id();
		let segments = agent.segments();
		segments.into_iter()
			.enumerate()
			.map(|(segment_index, segment)| {
				let material = segment.material();
				let mut f_def = b2::FixtureDef::new();
				f_def.density = material.density;
				f_def.restitution = material.restitution;
				f_def.friction = material.friction;

				let transform = segment.transform();
				let mut b_def = b2::BodyDef::new();
				b_def.body_type = b2::BodyType::Dynamic;
				b_def.linear_damping = 0.8;
				b_def.angular_damping = 0.9;
				b_def.angle = transform.angle;
				b_def.position = b2::Vec2 {
					x: transform.position.x,
					y: transform.position.y,
				};
				let refs = world::AgentRefs::with_segment(object_id, segment_index as u8);
				let handle = world.create_body_with(&b_def, refs);

				let mesh = segment.mesh();
				let flags = segment.flags;
				let attached_to = segment.attached_to;
				match mesh.shape {
					obj::Shape::Ball { radius } => {
						let mut circle_shape = b2::CircleShape::new();
						circle_shape.set_radius(radius);
						world.body_mut(handle).create_fixture_with(&circle_shape, &mut f_def, refs);
					}
					obj::Shape::Box { radius, ratio } => {
						let mut rect_shape = b2::PolygonShape::new();
						rect_shape.set_as_box(radius * ratio, radius);
						world.body_mut(handle).create_fixture_with(&rect_shape, &mut f_def, refs);
					}
					obj::Shape::Star { radius, n, .. } => {
						let p = &mesh.vertices;
						for i in 0..n {
							let mut quad = b2::PolygonShape::new();
							let i1 = (i * 2 + 1) as usize;
							let i2 = (i * 2) as usize;
							let i3 = ((i * 2 + (n * 2) - 1) % (n * 2)) as usize;
							let (p1, p2, p3) = match mesh.winding {
								obj::Winding::CW => (&p[i1], &p[i2], &p[i3]),
								obj::Winding::CCW => (&p[i1], &p[i3], &p[i2]),
							};
							quad.set(&[b2::Vec2 { x: 0., y: 0. },
							           b2::Vec2 {
								           x: p1.x * radius,
								           y: p1.y * radius,
							           },
							           b2::Vec2 {
								           x: p2.x * radius,
								           y: p2.y * radius,
							           },
							           b2::Vec2 {
								           x: p3.x * radius,
								           y: p3.y * radius,
							           }]);
							let refs = world::AgentRefs::with_bone(object_id, segment_index as u8, i as u8);
							world.body_mut(handle).create_fixture_with(&quad, &mut f_def, refs);
						}
					}
					obj::Shape::Triangle { radius, .. } => {
						let p = &mesh.vertices;
						let mut tri = b2::PolygonShape::new();
						let (p1, p2, p3) = match mesh.winding {
							obj::Winding::CW => (&p[0], &p[2], &p[1]),
							obj::Winding::CCW => (&p[0], &p[1], &p[2]),
						};
						tri.set(&[b2::Vec2 {
							          x: p1.x * radius,
							          y: p1.y * radius,
						          },
						          b2::Vec2 {
							          x: p2.x * radius,
							          y: p2.y * radius,
						          },
						          b2::Vec2 {
							          x: p3.x * radius,
							          y: p3.y * radius,
						          }]);
						world.body_mut(handle).create_fixture_with(&tri, &mut f_def, refs);
					}
				};
				JointRef {
					refs: refs,
					handle: handle,
					mesh: mesh,
					flags: flags,
					attachment: attached_to,
				}
			})
			.collect::<Vec<_>>()
	}

	fn build_joints(world: &mut b2::World<AgentData>, joint_refs: &Vec<JointRef>) {
		for &JointRef { handle: distal, mesh, attachment, flags, .. } in joint_refs {
			if let Some(attachment) = attachment {
				let upstream = &joint_refs[attachment.index as usize];
				let medial = upstream.handle;
				let angle_delta = world.body(distal).angle() - world.body(medial).angle();

				let v0 = upstream.mesh.vertices[attachment.attachment_point as usize] * upstream.mesh.shape.radius();
				let v1 = mesh.vertices[0] * mesh.shape.radius();
				let a = b2::Vec2 { x: v0.x, y: v0.y };
				let b = b2::Vec2 { x: v1.x, y: v1.y };
				macro_rules! common_joint (
					($joint:ident) => {
						$joint.collide_connected = false;
						$joint.reference_angle = angle_delta;
						$joint.local_anchor_a = a;
						$joint.local_anchor_b = b;
						world.create_joint_with(&$joint, ())
					}
				);
				if flags.contains(world::JOINT) {
					let mut joint = b2::RevoluteJointDef::new(medial, distal);
					joint.enable_limit = true;
					joint.upper_angle = consts::PI / 6.;
					joint.lower_angle = -consts::PI / 6.;
					common_joint!(joint);
				} else {
					// TODO: how do we reduce the clutter?
					let mut joint = b2::WeldJointDef::new(medial, distal);
					joint.frequency = 5.0;
					joint.damping_ratio = 0.9;
					common_joint!(joint);
				}
			}
		}
	}

	fn new_world(touched: Rc<RefCell<HashSet<world::AgentRefs>>>) -> b2::World<AgentData> {
		let mut world = b2::World::new(&b2::Vec2 { x: 0.0, y: 0.0 });
		world.set_contact_listener(Box::new(ContactListener { touched: touched }));
		world
	}
}

struct ContactListener {
	touched: Rc<RefCell<HashSet<world::AgentRefs>>>,
}

impl b2::ContactListener<AgentData> for ContactListener {
	fn post_solve(&mut self, ca: ContactAccess<AgentData>, _: &b2::ContactImpulse) {
		let body_a = ca.fixture_a.user_data();
		let body_b = ca.fixture_b.user_data();
		if body_a.agent_id != body_b.agent_id {
			self.touched.borrow_mut().insert(body_a.no_bone());
			self.touched.borrow_mut().insert(body_b.no_bone());
			// println!("{:?} touched {:?}", body_a, body_b);
		}
	}
}
<commit_msg>Closed extent fence<commit_after>use wrapped2d::b2;
use wrapped2d::user_data::*;
use wrapped2d::dynamics::world::callbacks::ContactAccess;
use backend::obj;
use backend::obj::{Solid, Geometry, Transformable};
use backend::world;
use backend::world::Intent;
use std::collections::HashMap;
use std::collections::HashSet;
use std::f32::consts;
use std::rc::Rc;
use std::cell::RefCell;
use core::geometry::*;
use super::*;

struct AgentData;

impl UserDataTypes for AgentData {
	type BodyData = world::AgentRefs;
	type JointData = ();
	type FixtureData = world::AgentRefs;
}

pub struct PhysicsSystem {
	world: b2::World<AgentData>,
	handles: HashMap<world::AgentRefs, b2::BodyHandle>,
	touched: Rc<RefCell<HashSet<world::AgentRefs>>>,
}

impl Updateable for PhysicsSystem {
	fn update(&mut self, state: &world::WorldState, dt: f32) {
		let mut forces = Vec::new();
		let mut impulses = Vec::new();

		for (h, b) in self.world.bodies() {
			let body = b.borrow();
			let center = (*body).world_center().clone();
			let key = (*body).user_data();
			if let Some(segment) = state.minion(key.agent_id).and_then(|c| c.segment(key.segment_index)) {
				match segment.state.intent {
					Intent::Move(force) => forces.push((h, center, force)),
					Intent::RunAway(impulse) => impulses.push((h, center, impulse)),
					Intent::Idle => {}
				}
			}
		}
		for (h, center, force) in forces {
			let b = &mut self.world.body_mut(h);
			b.apply_force(&b2::Vec2 {
				              x: force.x,
				              y: force.y,
			              },
			              &center,
			              true);
		}

		for (h, center, impulse) in impulses {
			let b = &mut self.world.body_mut(h);
			b.apply_linear_impulse(&b2::Vec2 {
				                       x: impulse.x,
				                       y: impulse.y,
			                       },
			                       &center,
			                       true);
		}

		self.world.step(dt, 8, 3);
	}
}

struct JointRef<'a> {
	refs: world::AgentRefs,
	handle: b2::BodyHandle,
	mesh: &'a obj::Mesh,
	flags: world::SegmentFlags,
	attachment: Option<world::Attachment>,
}

impl System for PhysicsSystem {
	fn register(&mut self, agent: &world::Agent) {
		// build fixtures
		let joint_refs = PhysicsSystem::build_fixtures(&mut self.world, &agent);
		// and then assemble them with joints
		PhysicsSystem::build_joints(&mut self.world, &joint_refs);
		// record them
		for JointRef { refs, handle, .. } in joint_refs {
			self.handles.insert(refs, handle);
		}
	}

	fn init(&mut self, world: &world::World) {
		self.init_extent(&world.extent);
	}

	fn to_world(&self, world: &mut world::World) {
		for (_, b) in self.world.bodies() {
			let body = b.borrow();
			let position = (*body).position();
			let angle = (*body).angle();
			let key = (*body).user_data();

			if let Some(agent) = world.minions.get_mut(key.agent_id) {
				if let Some(segment) = agent.segment_mut(key.segment_index) {
					let t = segment.transform();
					segment.transform_to(Transform {
						position: PhysicsSystem::from_vec2(&position),
						angle: angle,
						..t
					});
					segment.state.collision_detected = self.touched.borrow().contains(key);
				}
			}
		}
		self.touched.borrow_mut().clear();
	}
}

impl PhysicsSystem {
	pub fn new() -> Self {
		let touched = Rc::new(RefCell::new(HashSet::new()));
		PhysicsSystem {
			world: Self::new_world(touched.clone()),
			handles: HashMap::new(),
			touched: touched,
		}
	}

	fn to_vec2(p: &Position) -> b2::Vec2 {
		b2::Vec2 { x: p.x, y: p.y }
	}

	fn from_vec2(p: &b2::Vec2) -> Position {
		Position::new(p.x, p.y)
	}

	fn init_extent(&mut self, extent: &Rect) {
		let mut f_def = b2::FixtureDef::new();
		// f_def.density = material.density;
		// f_def.restitution = material.restitution;
		// f_def.friction = material.friction;

		let mut b_def = b2::BodyDef::new();
		b_def.body_type = b2::BodyType::Static;
		let refs = world::AgentRefs::with_id(0xFFFFFFFFusize);
		let handle = self.world.create_body_with(&b_def, refs);

		let mut rect = b2::ChainShape::new();
		rect.create_loop(&[Self::to_vec2(&extent.bottom_left()),
		                   Self::to_vec2(&extent.bottom_right()),
		                   Self::to_vec2(&extent.top_right()),
		                   Self::to_vec2(&extent.top_left())]);
		self.world.body_mut(handle).create_fixture_with(&rect, &mut f_def, refs);
	}

	fn build_fixtures<'a>(world: &mut b2::World<AgentData>, agent: &'a world::Agent) -> Vec<JointRef<'a>> {
		let object_id = agent.id();
		let segments = agent.segments();
		segments.into_iter()
			.enumerate()
			.map(|(segment_index, segment)| {
				let material = segment.material();
				let mut f_def = b2::FixtureDef::new();
				f_def.density = material.density;
				f_def.restitution = material.restitution;
				f_def.friction = material.friction;

				let transform = segment.transform();
				let mut b_def = b2::BodyDef::new();
				b_def.body_type = b2::BodyType::Dynamic;
				b_def.linear_damping = 0.8;
				b_def.angular_damping = 0.9;
				b_def.angle = transform.angle;
				b_def.position = b2::Vec2 {
					x: transform.position.x,
					y: transform.position.y,
				};
				let refs = world::AgentRefs::with_segment(object_id, segment_index as u8);
				let handle = world.create_body_with(&b_def, refs);

				let mesh = segment.mesh();
				let flags = segment.flags;
				let attached_to = segment.attached_to;
				match mesh.shape {
					obj::Shape::Ball { radius } => {
						let mut circle_shape = b2::CircleShape::new();
						circle_shape.set_radius(radius);
						world.body_mut(handle).create_fixture_with(&circle_shape, &mut f_def, refs);
					}
					obj::Shape::Box { radius, ratio } => {
						let mut rect_shape = b2::PolygonShape::new();
						rect_shape.set_as_box(radius * ratio, radius);
						world.body_mut(handle).create_fixture_with(&rect_shape, &mut f_def, refs);
					}
					obj::Shape::Star { radius, n, .. } => {
						let p = &mesh.vertices;
						for i in 0..n {
							let mut quad = b2::PolygonShape::new();
							let i1 = (i * 2 + 1) as usize;
							let i2 = (i * 2) as usize;
							let i3 = ((i * 2 + (n * 2) - 1) % (n * 2)) as usize;
							let (p1, p2, p3) = match mesh.winding {
								obj::Winding::CW => (&p[i1], &p[i2], &p[i3]),
								obj::Winding::CCW => (&p[i1], &p[i3], &p[i2]),
							};
							quad.set(&[b2::Vec2 { x: 0., y: 0. },
							           b2::Vec2 {
								           x: p1.x * radius,
								           y: p1.y * radius,
							           },
							           b2::Vec2 {
								           x: p2.x * radius,
								           y: p2.y * radius,
							           },
							           b2::Vec2 {
								           x: p3.x * radius,
								           y: p3.y * radius,
							           }]);
							let refs = world::AgentRefs::with_bone(object_id, segment_index as u8, i as u8);
							world.body_mut(handle).create_fixture_with(&quad, &mut f_def, refs);
						}
					}
					obj::Shape::Triangle { radius, .. } => {
						let p = &mesh.vertices;
						let mut tri = b2::PolygonShape::new();
						let (p1, p2, p3) = match mesh.winding {
							obj::Winding::CW => (&p[0], &p[2], &p[1]),
							obj::Winding::CCW => (&p[0], &p[1], &p[2]),
						};
						tri.set(&[b2::Vec2 {
							          x: p1.x * radius,
							          y: p1.y * radius,
						          },
						          b2::Vec2 {
							          x: p2.x * radius,
							          y: p2.y * radius,
						          },
						          b2::Vec2 {
							          x: p3.x * radius,
							          y: p3.y * radius,
						          }]);
						world.body_mut(handle).create_fixture_with(&tri, &mut f_def, refs);
					}
				};
				JointRef {
					refs: refs,
					handle: handle,
					mesh: mesh,
					flags: flags,
					attachment: attached_to,
				}
			})
			.collect::<Vec<_>>()
	}

	fn build_joints(world: &mut b2::World<AgentData>, joint_refs: &Vec<JointRef>) {
		for &JointRef { handle: distal, mesh, attachment, flags, .. } in joint_refs {
			if let Some(attachment) = attachment {
				let upstream = &joint_refs[attachment.index as usize];
				let medial = upstream.handle;
				let angle_delta = world.body(distal).angle() - world.body(medial).angle();

				let v0 = upstream.mesh.vertices[attachment.attachment_point as usize] * upstream.mesh.shape.radius();
				let v1 = mesh.vertices[0] * mesh.shape.radius();
				let a = b2::Vec2 { x: v0.x, y: v0.y };
				let b = b2::Vec2 { x: v1.x, y: v1.y };
				macro_rules! common_joint (
					($joint:ident) => {
						$joint.collide_connected = false;
						$joint.reference_angle = angle_delta;
						$joint.local_anchor_a = a;
						$joint.local_anchor_b = b;
						world.create_joint_with(&$joint, ())
					}
				);
				if flags.contains(world::JOINT) {
					let mut joint = b2::RevoluteJointDef::new(medial, distal);
					joint.enable_limit = true;
					joint.upper_angle = consts::PI / 6.;
					joint.lower_angle = -consts::PI / 6.;
					common_joint!(joint);
				} else {
					// TODO: how do we reduce the clutter?
					let mut joint = b2::WeldJointDef::new(medial, distal);
					joint.frequency = 5.0;
					joint.damping_ratio = 0.9;
					common_joint!(joint);
				}
			}
		}
	}

	fn new_world(touched: Rc<RefCell<HashSet<world::AgentRefs>>>) -> b2::World<AgentData> {
		let mut world = b2::World::new(&b2::Vec2 { x: 0.0, y: 0.0 });
		world.set_contact_listener(Box::new(ContactListener { touched: touched }));
		world
	}
}

struct ContactListener {
	touched: Rc<RefCell<HashSet<world::AgentRefs>>>,
}

impl b2::ContactListener<AgentData> for ContactListener {
	fn post_solve(&mut self, ca: ContactAccess<AgentData>, _: &b2::ContactImpulse) {
		let body_a = ca.fixture_a.user_data();
		let body_b = ca.fixture_b.user_data();
		if body_a.agent_id != body_b.agent_id {
			self.touched.borrow_mut().insert(body_a.no_bone());
			self.touched.borrow_mut().insert(body_b.no_bone());
			// println!("{:?} touched {:?}", body_a, body_b);
		}
	}
}
<|endoftext|>


<commit_before>//! Abstractions for asynchronous programming
//!
//! This crate provides a number of core abstractions for writing asynchronous code:
//!
//! - [Futures](::Future) (sometimes called promises), which represent a single
//! asychronous computation that may result in a final value or an error.
//!
//! - [Streams](::Stream), which represent a series of values or errors produced asynchronously.
//!
//! - [Sinks](::Sink), which support asynchronous writing of data.
//!
//! - [Executors](::executor), which are responsible for running asynchronous tasks.
//!
//! The crate also contains abstractions for [asynchronous I/O](::io) and
//! [cross-task communication](::channel).
//!
//! Underlying all of this is the *task system*, which is a form of lightweight
//! threading. Large asynchronous computations are built up using futures,
//! streams and sinks, and then spawned as independent tasks that are run to
//! completion, but *do not block* the thread running them.

#![no_std]
#![doc(html_root_url = "https://docs.rs/futures/0.2.0-alpha")]

extern crate futures_core;
extern crate futures_channel;
extern crate futures_executor;
extern crate futures_io;
extern crate futures_sink;
extern crate futures_util;

pub use futures_core::future::{Future, IntoFuture};
pub use futures_util::future::FutureExt;
pub use futures_core::stream::Stream;
pub use futures_util::stream::StreamExt;
pub use futures_sink::Sink;
pub use futures_util::sink::SinkExt;

// Macros redefined here because macro re-exports are unstable.

/// A macro for extracting the successful type of a `Poll<T, E>`.
///
/// This macro bakes in propagation of both errors and `Pending` signals by
/// returning early.
#[macro_export]
macro_rules! try_ready {
    ($e:expr) => (match $e {
        Ok($crate::prelude::Async::Ready(t)) => t,
        Ok($crate::prelude::Async::Pending) => return Ok($crate::prelude::Async::Pending),
        Err(e) => return Err(From::from(e)),
    })
}

/// A macro to create a `static` of type `LocalKey`
///
/// This macro is intentionally similar to the `thread_local!`, and creates a
/// `static` which has a `get_mut` method to access the data on a task.
///
/// The data associated with each task local is per-task, so different tasks
/// will contain different values.
#[macro_export]
macro_rules! task_local {
    (static $NAME:ident: $t:ty = $e:expr) => (
        static $NAME: $crate::task::LocalKey<$t> = {
            fn __init() -> $t { $e }
            fn __key() -> ::std::any::TypeId {
                struct __A;
                ::std::any::TypeId::of::<__A>()
            }
            $crate::task::LocalKey {
                __init: __init,
                __key: __key,
            }
        };
    )
}

pub use futures_core::{Async, Poll, Never};

#[cfg(feature = "std")]
pub mod channel {
    //! Cross-task communciation.
    //!
    //! Like threads, concurrent tasks sometimes need to communicate with each
    //! other. This module contains two basic abstractions for doing so:
    //!
    //! - [oneshot](::channel::oneshot), a way of sending a single value from
    //! one task to another.
    //!
    //! - [mpsc](::channel::mpsc), a multi-producer, single-consumer channel for
    //! sending values between tasks, analogous to the similarly-named structure
    //! in the standard library.

    pub use futures_channel::{oneshot, mpsc};
}

#[cfg(feature = "std")]
pub mod executor {
    //! Task execution.
    //!
    //! All asynchronous computation occurs within an executor, which is
    //! capable of spawning futures as tasks. This module provides several
    //! built-in executors, as well as tools for building your own.
    //!
    //! # Using a thread pool (M:N task scheduling)
    //!
    //! Most of the time tasks should be executed on a [thread
    //! pool](::executor::ThreadPool). A small set of worker threads can handle
    //! a very large set of spawned tasks (which are much lighter weight than
    //! threads).
    //!
    //! The simplest way to use a thread pool is to
    //! [`run`](::executor::ThreadPool::run) an initial task on it, which can
    //! then spawn further tasks back onto the pool to complete its work:
    //!
    //! ```
    //! use futures::executor::ThreadPool;
    //! # use futures::future::{Future, lazy};
    //! # let my_app: Box<Future<Item = (), Error = ()>> = Box::new(lazy(|| Ok(())));
    //!
    //! // assumping `my_app: Future`
    //! ThreadPool::new().run(my_app);
    //! ```
    //!
    //! The call to [`run`](::executor::ThreadPool::run) will block the current
    //! thread until the future defined by `my_app` completes, and will return
    //! the result of that future.
    //!
    //! # Spawning additional tasks
    //!
    //! There are two ways to spawn a task:
    //!
    //! - Spawn onto a "default" execuctor by calling the top-level
    //! [`spawn`](::executor::spawn) function or [pulling the executor from the
    //! task context](::task::Context::executor).
    //!
    //! - Spawn onto a specific executor by calling its
    //! [`spawn`](::executor::Executor::spawn) method directly.
    //!
    //! Every task always has an associated default executor, which is usually
    //! the executor on which the task is running.
    //!
    //! # Single-threaded execution
    //!
    //! In addition to thread pools, it's possible to run a task (and the tasks
    //! it spawns) entirely within a single thread via the
    //! [`LocalPool`](::executor::LocalPool) executor. Aside from cutting down
    //! on synchronization costs, this executor also makes it possible to
    //! execute non-`Send` tasks, via
    //! [`spawn_local`](::executor::LocalExecutor::spawn_local). The `LocalPool`
    //! is best suited for running I/O-bound tasks that do relatively little
    //! work between I/O operations.
    //!
    //! There is also a convenience function,
    //! [`block_on`](::executor::block_on), for simply running a future to
    //! completion on the current thread. However, any tasks spawned by the
    //! future will *also* be run on the current thread (via `LocalPool`), and
    //! will be dropped when the primary future completes.
    // TODO: add docs (or link to apr) for implementing an executor

    pub use futures_executor::{
        Enter, EnterError, LocalExecutor, LocalPool, Spawn, SpawnWithHandle, ThreadPool,
        ThreadPoolBuilder, block_on, enter, spawn, spawn_with_handle
    };
    pub use futures_core::executor::{SpawnError, Executor};
}

pub mod future {
    //! Asynchronous values.
    //!
    //! This module contains:
    //!
    //! - The [`Future` trait](::Future).
    //!
    //! - The [`FutureExt`](::future::FutureExt) trait, which provides adapters
    //! for chaining and composing futures.
    //!
    //! - Top-level future combinators like [`lazy`](::future::lazy) which
    //! creates a future from a closure that defines its return value, and
    //! [`result`](::future::result), which constructs a future with an
    //! immediate defined value.

    pub use futures_core::future::{
        FutureOption, FutureResult, Future, FutureFrom, IntoFuture, err, ok, result
    };
    pub use futures_util::future::{
        AndThen, Empty, Flatten, FlattenStream, ErrInto, Fuse,
        Inspect, IntoStream, Join, Join3, Join4, Join5, Lazy, LoopFn,
        Map, MapErr, OrElse, PollFn, Select, Then, Either, Loop, FutureExt, empty,
        lazy, loop_fn, poll_fn
    };

    #[cfg(feature = "std")]
    pub use futures_util::future::{
        CatchUnwind, JoinAll, SelectAll, SelectOk, Shared, SharedError, SharedItem,
        join_all, select_all, select_ok
    };
}

#[cfg(feature = "std")]
pub mod io {
    //! Asynchronous I/O.
    //!
    //! This module is the asynchronous version of `std::io`. It defines two
    //! traits, [`AsyncRead`](::io::AsyncRead) and
    //! [`AsyncWrite`](::io::AsyncWrite), which mirror the `Read` and `Write`
    //! traits of the standard library. However, these traits integrate with the
    //! asynchronous task system, so that if an I/O object isn't ready for
    //! reading (or writing), the thread is not blocked, and instead the current
    //! task is queued to be woken when I/O is ready.
    //!
    //! In addition, the [`AsyncReadExt`](::io::AsyncReadExt) and
    //! [`AsyncWriteExt`](::io::AsyncWriteExt) extension traits offer a variety
    //! of useful combinators for operating with asynchronous I/O objects,
    //! including ways to work with them using futures, streams and sinks.

    pub use futures_io::{
        Error, Initializer, IoVec, ErrorKind, AsyncRead, AsyncWrite, Result
    };
    pub use futures_util::io::{
        AsyncReadExt, AsyncWriteExt, AllowStdIo, Close, CopyInto, Flush,
        Read, ReadExact, ReadHalf, ReadToEnd, Window, WriteAll, WriteHalf,
    };
}

#[cfg(feature = "std")]
pub mod never {
    //! This module contains the `Never` type.
    //!
    //! Values of this type can never be created and will never exist.
    pub use futures_core::never::*;
}

pub mod prelude {
    //! A "prelude" for crates using the `futures` crate.
    //!
    //! This prelude is similar to the standard library's prelude in that you'll
    //! almost always want to import its entire contents, but unlike the standard
    //! library's prelude you'll have to do so manually:
    //!
    //! ```
    //! use futures::prelude::*;
    //! ```
    //!
    //! The prelude may grow over time as additional items see ubiquitous use.

    pub use futures_core::{
        Future,
        IntoFuture,
        Stream,
        Async,
        Poll,
        Never,
        task,
    };

    #[cfg(feature = "std")]
    pub use futures_core::executor::Executor;

    pub use futures_sink::Sink;

    #[cfg(feature = "std")]
    pub use futures_io::{
        AsyncRead,
        AsyncWrite,
    };

    pub use futures_util::{
        FutureExt,
        StreamExt,
        SinkExt,
    };

    #[cfg(feature = "std")]
    pub use futures_util::{
        AsyncReadExt,
        AsyncWriteExt,
    };
}

pub mod sink {
    //! Asynchronous sinks.
    //!
    //! This module contains:
    //!
    //! - The [`Sink` trait](::Sink), which allows you to asynchronously write data.
    //!
    //! - The [`SinkExt`](::sink::SinkExt) trait, which provides adapters
    //! for chaining and composing sinks.

    pub use futures_sink::Sink;

    pub use futures_util::sink::{
        Close, Fanout, Flush, Send, SendAll, SinkErrInto, SinkMapErr, With,
        WithFlatMap, SinkExt, close, flush,
    };

    #[cfg(feature = "std")]
    pub use futures_util::sink::Buffer;
}

pub mod stream {
    //! Asynchronous streams.
    //!
    //! This module contains:
    //!
    //! - The [`Stream` trait](::Stream), for objects that can asynchronously
    //! produce a sequence of values.
    //!
    //! - The [`StreamExt`](::future::StreamExt) trait, which provides adapters
    //! for chaining and composing streams.
    //!
    //! - Top-level stream contructors like [`iter_ok`](::stream::iter_ok) which
    //! creates a stream from an iterator, and
    //! [`futures_unordered`](::stream::futures_unordered()), which constructs a
    //! stream from a collection of futures.

    pub use futures_core::stream::Stream;

    pub use futures_util::stream::{
        AndThen, Chain, Concat, Empty, Filter, FilterMap, Flatten, Fold,
        ForEach, Forward, ErrInto, Fuse, Inspect, InspectErr, IterOk,
        IterResult, Map, MapErr, Once, OrElse, Peekable, PollFn, Repeat, Select,
        Skip, SkipWhile, StreamFuture, Take, TakeWhile, Then, Unfold, Zip,
        StreamExt, empty, iter_ok, iter_result, once, poll_fn, repeat, unfold,
    };

    #[cfg(feature = "std")]
    pub use futures_util::stream::{
        futures_unordered, BufferUnordered, Buffered, CatchUnwind, Chunks, Collect,
        FuturesUnordered, FuturesOrdered, ReuniteError, SelectAll, SplitSink, SplitStream,
        futures_ordered,
    };
}

pub mod task {
    //! Tools for working with tasks.
    //!
    //! This module contains:
    //!
    //! - [`Context`](::task::Context), which provides contextual data present
    //! for every task, including a handle for waking up the task.
    //!
    //! - [`Waker`](::task::Waker), a handle for waking up a task.
    //!
    //! - [`LocalKey`](::task::LocalKey), a key for task-local data; you should
    //! use the [`task_local` macro](../macro.task_local.html) to set up such keys.
    //!
    //! Tasks themselves are generally created by spawning a future onto [an
    //! executor](::executor). However, you can manually construct a task by
    //! creating your own `Context` instance, and polling a future with it.
    //!
    //! The remaining types and traits in the module are used for implementing
    //! executors or dealing with synchronization issues around task wakeup.

    pub use futures_core::task::{
        AtomicWaker, Context, LocalMap, Waker, UnsafeWake,
    };

    #[cfg(feature = "std")]
    pub use futures_core::task::{LocalKey, Wake};
}
<commit_msg>rustdoc: minor typos fixing<commit_after>//! Abstractions for asynchronous programming.
//!
//! This crate provides a number of core abstractions for writing asynchronous code:
//!
//! - [Futures](::Future) (sometimes called promises), which represent a single
//! asychronous computation that may result in a final value or an error.
//!
//! - [Streams](::Stream), which represent a series of values or errors produced asynchronously.
//!
//! - [Sinks](::Sink), which support asynchronous writing of data.
//!
//! - [Executors](::executor), which are responsible for running asynchronous tasks.
//!
//! The crate also contains abstractions for [asynchronous I/O](::io) and
//! [cross-task communication](::channel).
//!
//! Underlying all of this is the *task system*, which is a form of lightweight
//! threading. Large asynchronous computations are built up using futures,
//! streams and sinks, and then spawned as independent tasks that are run to
//! completion, but *do not block* the thread running them.

#![no_std]
#![doc(html_root_url = "https://docs.rs/futures/0.2.0-alpha")]

extern crate futures_core;
extern crate futures_channel;
extern crate futures_executor;
extern crate futures_io;
extern crate futures_sink;
extern crate futures_util;

pub use futures_core::future::{Future, IntoFuture};
pub use futures_util::future::FutureExt;
pub use futures_core::stream::Stream;
pub use futures_util::stream::StreamExt;
pub use futures_sink::Sink;
pub use futures_util::sink::SinkExt;

// Macros redefined here because macro re-exports are unstable.

/// A macro for extracting the successful type of a `Poll<T, E>`.
///
/// This macro bakes in propagation of both errors and `Pending` signals by
/// returning early.
#[macro_export]
macro_rules! try_ready {
    ($e:expr) => (match $e {
        Ok($crate::prelude::Async::Ready(t)) => t,
        Ok($crate::prelude::Async::Pending) => return Ok($crate::prelude::Async::Pending),
        Err(e) => return Err(From::from(e)),
    })
}

/// A macro to create a `static` of type `LocalKey`.
///
/// This macro is intentionally similar to the `thread_local!`, and creates a
/// `static` which has a `get_mut` method to access the data on a task.
///
/// The data associated with each task local is per-task, so different tasks
/// will contain different values.
#[macro_export]
macro_rules! task_local {
    (static $NAME:ident: $t:ty = $e:expr) => (
        static $NAME: $crate::task::LocalKey<$t> = {
            fn __init() -> $t { $e }
            fn __key() -> ::std::any::TypeId {
                struct __A;
                ::std::any::TypeId::of::<__A>()
            }
            $crate::task::LocalKey {
                __init: __init,
                __key: __key,
            }
        };
    )
}

pub use futures_core::{Async, Poll, Never};

#[cfg(feature = "std")]
pub mod channel {
    //! Cross-task communication.
    //!
    //! Like threads, concurrent tasks sometimes need to communicate with each
    //! other. This module contains two basic abstractions for doing so:
    //!
    //! - [oneshot](::channel::oneshot), a way of sending a single value from
    //! one task to another.
    //!
    //! - [mpsc](::channel::mpsc), a multi-producer, single-consumer channel for
    //! sending values between tasks, analogous to the similarly-named structure
    //! in the standard library.

    pub use futures_channel::{oneshot, mpsc};
}

#[cfg(feature = "std")]
pub mod executor {
    //! Task execution.
    //!
    //! All asynchronous computation occurs within an executor, which is
    //! capable of spawning futures as tasks. This module provides several
    //! built-in executors, as well as tools for building your own.
    //!
    //! # Using a thread pool (M:N task scheduling)
    //!
    //! Most of the time tasks should be executed on a [thread
    //! pool](::executor::ThreadPool). A small set of worker threads can handle
    //! a very large set of spawned tasks (which are much lighter weight than
    //! threads).
    //!
    //! The simplest way to use a thread pool is to
    //! [`run`](::executor::ThreadPool::run) an initial task on it, which can
    //! then spawn further tasks back onto the pool to complete its work:
    //!
    //! ```
    //! use futures::executor::ThreadPool;
    //! # use futures::future::{Future, lazy};
    //! # let my_app: Box<Future<Item = (), Error = ()>> = Box::new(lazy(|| Ok(())));
    //!
    //! // assumping `my_app: Future`
    //! ThreadPool::new().run(my_app);
    //! ```
    //!
    //! The call to [`run`](::executor::ThreadPool::run) will block the current
    //! thread until the future defined by `my_app` completes, and will return
    //! the result of that future.
    //!
    //! # Spawning additional tasks
    //!
    //! There are two ways to spawn a task:
    //!
    //! - Spawn onto a "default" execuctor by calling the top-level
    //! [`spawn`](::executor::spawn) function or [pulling the executor from the
    //! task context](::task::Context::executor).
    //!
    //! - Spawn onto a specific executor by calling its
    //! [`spawn`](::executor::Executor::spawn) method directly.
    //!
    //! Every task always has an associated default executor, which is usually
    //! the executor on which the task is running.
    //!
    //! # Single-threaded execution
    //!
    //! In addition to thread pools, it's possible to run a task (and the tasks
    //! it spawns) entirely within a single thread via the
    //! [`LocalPool`](::executor::LocalPool) executor. Aside from cutting down
    //! on synchronization costs, this executor also makes it possible to
    //! execute non-`Send` tasks, via
    //! [`spawn_local`](::executor::LocalExecutor::spawn_local). The `LocalPool`
    //! is best suited for running I/O-bound tasks that do relatively little
    //! work between I/O operations.
    //!
    //! There is also a convenience function,
    //! [`block_on`](::executor::block_on), for simply running a future to
    //! completion on the current thread. However, any tasks spawned by the
    //! future will *also* be run on the current thread (via `LocalPool`), and
    //! will be dropped when the primary future completes.
    // TODO: add docs (or link to apr) for implementing an executor

    pub use futures_executor::{
        Enter, EnterError, LocalExecutor, LocalPool, Spawn, SpawnWithHandle, ThreadPool,
        ThreadPoolBuilder, block_on, enter, spawn, spawn_with_handle
    };
    pub use futures_core::executor::{SpawnError, Executor};
}

pub mod future {
    //! Asynchronous values.
    //!
    //! This module contains:
    //!
    //! - The [`Future` trait](::Future).
    //!
    //! - The [`FutureExt`](::future::FutureExt) trait, which provides adapters
    //! for chaining and composing futures.
    //!
    //! - Top-level future combinators like [`lazy`](::future::lazy) which
    //! creates a future from a closure that defines its return value, and
    //! [`result`](::future::result), which constructs a future with an
    //! immediate defined value.

    pub use futures_core::future::{
        FutureOption, FutureResult, Future, FutureFrom, IntoFuture, err, ok, result
    };
    pub use futures_util::future::{
        AndThen, Empty, Flatten, FlattenStream, ErrInto, Fuse,
        Inspect, IntoStream, Join, Join3, Join4, Join5, Lazy, LoopFn,
        Map, MapErr, OrElse, PollFn, Select, Then, Either, Loop, FutureExt, empty,
        lazy, loop_fn, poll_fn
    };

    #[cfg(feature = "std")]
    pub use futures_util::future::{
        CatchUnwind, JoinAll, SelectAll, SelectOk, Shared, SharedError, SharedItem,
        join_all, select_all, select_ok
    };
}

#[cfg(feature = "std")]
pub mod io {
    //! Asynchronous I/O.
    //!
    //! This module is the asynchronous version of `std::io`. It defines two
    //! traits, [`AsyncRead`](::io::AsyncRead) and
    //! [`AsyncWrite`](::io::AsyncWrite), which mirror the `Read` and `Write`
    //! traits of the standard library. However, these traits integrate with the
    //! asynchronous task system, so that if an I/O object isn't ready for
    //! reading (or writing), the thread is not blocked, and instead the current
    //! task is queued to be woken when I/O is ready.
    //!
    //! In addition, the [`AsyncReadExt`](::io::AsyncReadExt) and
    //! [`AsyncWriteExt`](::io::AsyncWriteExt) extension traits offer a variety
    //! of useful combinators for operating with asynchronous I/O objects,
    //! including ways to work with them using futures, streams and sinks.

    pub use futures_io::{
        Error, Initializer, IoVec, ErrorKind, AsyncRead, AsyncWrite, Result
    };
    pub use futures_util::io::{
        AsyncReadExt, AsyncWriteExt, AllowStdIo, Close, CopyInto, Flush,
        Read, ReadExact, ReadHalf, ReadToEnd, Window, WriteAll, WriteHalf,
    };
}

#[cfg(feature = "std")]
pub mod never {
    //! This module contains the `Never` type.
    //!
    //! Values of this type can never be created and will never exist.
    pub use futures_core::never::*;
}

pub mod prelude {
    //! A "prelude" for crates using the `futures` crate.
    //!
    //! This prelude is similar to the standard library's prelude in that you'll
    //! almost always want to import its entire contents, but unlike the standard
    //! library's prelude you'll have to do so manually:
    //!
    //! ```
    //! use futures::prelude::*;
    //! ```
    //!
    //! The prelude may grow over time as additional items see ubiquitous use.

    pub use futures_core::{
        Future,
        IntoFuture,
        Stream,
        Async,
        Poll,
        Never,
        task,
    };

    #[cfg(feature = "std")]
    pub use futures_core::executor::Executor;

    pub use futures_sink::Sink;

    #[cfg(feature = "std")]
    pub use futures_io::{
        AsyncRead,
        AsyncWrite,
    };

    pub use futures_util::{
        FutureExt,
        StreamExt,
        SinkExt,
    };

    #[cfg(feature = "std")]
    pub use futures_util::{
        AsyncReadExt,
        AsyncWriteExt,
    };
}

pub mod sink {
    //! Asynchronous sinks.
    //!
    //! This module contains:
    //!
    //! - The [`Sink` trait](::Sink), which allows you to asynchronously write data.
    //!
    //! - The [`SinkExt`](::sink::SinkExt) trait, which provides adapters
    //! for chaining and composing sinks.

    pub use futures_sink::Sink;

    pub use futures_util::sink::{
        Close, Fanout, Flush, Send, SendAll, SinkErrInto, SinkMapErr, With,
        WithFlatMap, SinkExt, close, flush,
    };

    #[cfg(feature = "std")]
    pub use futures_util::sink::Buffer;
}

pub mod stream {
    //! Asynchronous streams.
    //!
    //! This module contains:
    //!
    //! - The [`Stream` trait](::Stream), for objects that can asynchronously
    //! produce a sequence of values.
    //!
    //! - The [`StreamExt`](::future::StreamExt) trait, which provides adapters
    //! for chaining and composing streams.
    //!
    //! - Top-level stream contructors like [`iter_ok`](::stream::iter_ok) which
    //! creates a stream from an iterator, and
    //! [`futures_unordered`](::stream::futures_unordered()), which constructs a
    //! stream from a collection of futures.

    pub use futures_core::stream::Stream;

    pub use futures_util::stream::{
        AndThen, Chain, Concat, Empty, Filter, FilterMap, Flatten, Fold,
        ForEach, Forward, ErrInto, Fuse, Inspect, InspectErr, IterOk,
        IterResult, Map, MapErr, Once, OrElse, Peekable, PollFn, Repeat, Select,
        Skip, SkipWhile, StreamFuture, Take, TakeWhile, Then, Unfold, Zip,
        StreamExt, empty, iter_ok, iter_result, once, poll_fn, repeat, unfold,
    };

    #[cfg(feature = "std")]
    pub use futures_util::stream::{
        futures_unordered, BufferUnordered, Buffered, CatchUnwind, Chunks, Collect,
        FuturesUnordered, FuturesOrdered, ReuniteError, SelectAll, SplitSink, SplitStream,
        futures_ordered,
    };
}

pub mod task {
    //! Tools for working with tasks.
    //!
    //! This module contains:
    //!
    //! - [`Context`](::task::Context), which provides contextual data present
    //! for every task, including a handle for waking up the task.
    //!
    //! - [`Waker`](::task::Waker), a handle for waking up a task.
    //!
    //! - [`LocalKey`](::task::LocalKey), a key for task-local data; you should
    //! use the [`task_local` macro](../macro.task_local.html) to set up such keys.
    //!
    //! Tasks themselves are generally created by spawning a future onto [an
    //! executor](::executor). However, you can manually construct a task by
    //! creating your own `Context` instance, and polling a future with it.
    //!
    //! The remaining types and traits in the module are used for implementing
    //! executors or dealing with synchronization issues around task wakeup.

    pub use futures_core::task::{
        AtomicWaker, Context, LocalMap, Waker, UnsafeWake,
    };

    #[cfg(feature = "std")]
    pub use futures_core::task::{LocalKey, Wake};
}
<|endoftext|>


<commit_before>//! Infers how each function uses every register

use frontend::imports::ImportInfo;
use frontend::radeco_containers::{RadecoFunction, RadecoModule};
use middle::dce;
use middle::ir;
use middle::regfile::{RegisterId, RegisterUsage, SubRegisterFile};
use middle::ssa::cfg_traits::*;
use middle::ssa::ssa_traits::*;
use middle::ssa::ssastorage::SSAStorage;

use r2api::structs::LCCInfo;

use petgraph::prelude::*;
use petgraph::visit::{DfsPostOrder, VisitMap, Walker};

use std::collections::{hash_map, BTreeMap, HashMap, HashSet};
use std::iter::{self, Extend, FromIterator};

/// For every function, patch all of its call sites to ignore registers that the
/// callee doesn't read and to preserve register values that the callee
/// preserves. Then, record which registers it reads and which registers it
/// preserves.
///
/// After this, all functions should have mutable [`regusage`][RadecoFunction::regusage]s.
///
/// This analysis is super conservative; for example, if a function preserves a
/// register by pushing it onto the stack and popping it back right before
/// returning, it is considered to be read and not preserved because we can't
/// guarantee that that stack location is never subsequently read or modified.
pub fn run(rmod: &mut RadecoModule, reginfo: &SubRegisterFile) -> () {
    let mut a = Inferer::new();
    a.run(rmod, reginfo);
    for (_, rfn) in &rmod.functions {
        eprintln!("{:?} (@ {:#X}):", rfn.name, rfn.offset);
        for (i, regname) in reginfo.whole_names.iter().enumerate() {
            let rid = RegisterId::from_usize(i);
            let is_ignored = rfn.regusage.is_ignored(rid);
            let is_preserved = rfn.regusage.is_preserved(rid);
            eprintln!(
                "  {:>8}: {} {}",
                regname,
                if is_ignored { 'I' } else { 'R' },
                if is_preserved { 'P' } else { 'C' }
            );
        }
    }
}

struct Inferer {
    /// Addresses of the functions we've already analyzed
    analyzed: HashSet<u64>,
}

impl Inferer {
    fn new() -> Inferer {
        Inferer {
            analyzed: HashSet::new(),
        }
    }

    /// Calls `patch_fn`, `dce::collect`, and `analyze_fn` on every function,
    /// callees first
    fn run(&mut self, rmod: &mut RadecoModule, reginfo: &SubRegisterFile) -> () {
        // for imports, *ASSUME* that the callconv that r2 says is correct
        {
            let imp_ru_iter = rmod.imports.iter().filter_map(|(&imp_addr, imp_info)| {
                let imp_rfn = imp_info.rfn.borrow();
                let regusage = reginfo.r2callconv_to_register_usage(
                    imp_rfn.callconv.as_ref()?, // ignore imports without callconvs
                    &*imp_rfn.callconv_name,
                )?;
                Some((imp_addr, regusage))
            });
            for (imp_addr, imp_ru) in imp_ru_iter {
                rmod.functions.get_mut(&imp_addr).unwrap().regusage = imp_ru;
                self.analyzed.insert(imp_addr);
            }
        }

        let mut dfs_wi = DfsPostOrder::empty(&rmod.callgraph).iter(&rmod.callgraph);
        // pick a function ...
        for fn_ni in rmod.callgraph.node_indices() {
            // ... and start a dfs on it
            dfs_wi.inner_mut().move_to(fn_ni);
            while let Some(fn_to_anal) = dfs_wi.next() {
                let fn_addr = rmod.callgraph[fn_to_anal];

                // ignore functions already in `call_convs` (probably because its an import)
                if !self.analyzed.contains(&fn_addr) {
                    self.patch_fn(fn_addr, &mut rmod.functions);

                    let rfn = &mut rmod.functions.get_mut(&fn_addr).unwrap();
                    dce::collect(rfn.ssa_mut());
                    // TODO: inst_combine

                    let ru = self.analyze_fn(rfn, reginfo).unwrap_or_else(|| {
                        radeco_err!("Failed to analyze fn: {:?} (@ {:#X})", rfn.name, fn_addr);
                        // if analysis failed, default to "reads and clobbers everything"
                        reginfo.new_register_usage()
                    });

                    rfn.regusage = ru;
                    self.analyzed.insert(fn_addr);
                }
            }
        }
    }

    /// Using the callconv info we've gathered so far, patch-up call sites to
    /// to remove arguments that the callee doesn't read and make values in
    /// callee-saved registers be preserved across the call.
    fn patch_fn(&self, fn_addr: u64, fn_map: &mut BTreeMap<u64, RadecoFunction>) -> () {
        radeco_trace!("patching calls in fn: {}", rfn.name);
        for node in fn_map[&fn_addr].ssa().inorder_walk() {
            if let Ok(NodeType::Op(ir::MOpcode::OpCall)) =
                fn_map[&fn_addr].ssa().node_data(node).map(|nd| nd.nt)
            {
                self.patch_call_node(node, fn_addr, fn_map)
                    .unwrap_or_else(|| {
                        radeco_warn!(
                            "failed to remove unused args for call at {:#X}",
                            ssa.address(node).unwrap()
                        );
                    });
            }
        }
    }

    fn patch_call_node(
        &self,
        call_node: <SSAStorage as SSA>::ValueRef,
        fn_addr: u64,
        fn_map: &mut BTreeMap<u64, RadecoFunction>,
    ) -> Option<()> {
        // bail on indirect or weird call
        let (call_tgt_addr, call_reg_map) = direct_call_info(fn_map[&fn_addr].ssa(), call_node)?;

        // remove unread args
        for (reg_idx, &op_node) in (0..).zip(&call_reg_map) {
            if fn_map[&call_tgt_addr]
                .regusage
                .is_ignored(RegisterId::from_usize(reg_idx))
            {
                fn_map
                    .get_mut(&fn_addr)
                    .unwrap()
                    .ssa_mut()
                    .op_unuse(call_node, op_node);
            }
        }

        // bridge preserved registers
        for use_node in fn_map[&fn_addr].ssa().uses_of(call_node) {
            let reg_idx = match fn_map[&fn_addr]
                .ssa()
                .sparse_operands_of(use_node)
                .as_slice()
            {
                &[(reg_idx, _)] => reg_idx,
                _ => panic!("invalid use of call as operand"),
            };
            if fn_map[&call_tgt_addr]
                .regusage
                .is_preserved(RegisterId::from_u8(reg_idx))
            {
                fn_map
                    .get_mut(&fn_addr)
                    .unwrap()
                    .ssa_mut()
                    .replace_value(use_node, call_reg_map[reg_idx as usize]);
            }
        }

        Some(())
    }

    fn analyze_fn(&self, rfn: &RadecoFunction, reginfo: &SubRegisterFile) -> Option<RegisterUsage> {
        eprintln!("analyzing fn: {}", rfn.name);
        let ssa = rfn.ssa();
        let entry_regstate_node = ssa.registers_in(ssa.entry_node()?)?;
        let exit_regstate_node = ssa.registers_in(ssa.exit_node()?)?;
        // some registers may not be present in the entry node;
        // this means that the function neither reads nor preserves that register
        let entry_regstate: Vec<Option<_>> = register_state_info(ssa, entry_regstate_node);
        // every register must be present in the exit node though
        let exit_regstate: Vec<_> = register_state_info(ssa, exit_regstate_node)
            .into_iter()
            .collect::<Option<_>>()?;

        let mut ret = reginfo.new_register_usage();
        ret.set_all_ignored();

        // ignore registers not in entry regstate
        let mut regstate_iter = (0..)
            .zip(exit_regstate)
            .zip(entry_regstate)
            .filter_map(|((i, x), on)| on.map(|n| (i, n, x)));

        for (i, reg_val_entry, reg_val_exit) in regstate_iter {
            if reg_val_exit == reg_val_entry {
                eprintln!(
                    "  set_preserved({})",
                    ssa.regnames.get(i as usize).map(|x| &**x).unwrap_or("mem")
                );
                ret.set_preserved(RegisterId::from_u8(i));
            }

            // find all uses, ignoring entry/exit register state
            let mut uses_iter = ssa.uses_of(reg_val_entry)
                .into_iter()
                .filter(|&n| n != entry_regstate_node && n != exit_regstate_node);
            if uses_iter.next().is_some() {
                ret.set_read(RegisterId::from_u8(i));
            }
        }

        Some(ret)
    }
}

/// Extracts the value of all registers at a `RegisterState` SSA node.
/// The length of the returned `Vec` is exactly `ssa.regnames.len() + 1`.
fn register_state_info(
    ssa: &SSAStorage,
    regstate_node: <SSAStorage as SSA>::ValueRef,
) -> Vec<Option<<SSAStorage as SSA>::ValueRef>> {
    let mut reg_opt_map: Vec<Option<_>> = vec![None; ssa.regnames.len() + 1];
    for (op_idx, op_node) in ssa.sparse_operands_of(regstate_node) {
        if let Some(p) = reg_opt_map.get_mut(op_idx as usize) {
            *p = Some(op_node);
        }
    }
    reg_opt_map
}

/// Extracts the call target address and the value of all registers
/// The length of the returned `Vec` is exactly `ssa.regnames.len() + 1`.
/// Returns `None` if the call is indirect or if not all registers have a value.
fn direct_call_info(
    ssa: &SSAStorage,
    call_node: <SSAStorage as SSA>::ValueRef,
) -> Option<(u64, Vec<<SSAStorage as SSA>::ValueRef>)> {
    let mut tgt_opt: Option<u64> = None;
    let mut reg_opt_map: Vec<Option<_>> = vec![None; ssa.regnames.len() + 1];

    for (op_idx, op_node) in ssa.sparse_operands_of(call_node) {
        if op_idx == 0 {
            tgt_opt = tgt_opt.or_else(|| ssa.constant(op_node));
        } else {
            if let Some(p) = reg_opt_map.get_mut((op_idx - 1) as usize) {
                *p = Some(op_node);
            }
        }
    }

    let reg_map_opt: Option<Vec<_>> = reg_opt_map.into_iter().collect();
    Some((tgt_opt?, reg_map_opt?))
}
<commit_msg>Fix logging<commit_after>//! Infers how each function uses every register

use frontend::imports::ImportInfo;
use frontend::radeco_containers::{RadecoFunction, RadecoModule};
use middle::dce;
use middle::ir;
use middle::regfile::{RegisterId, RegisterUsage, SubRegisterFile};
use middle::ssa::cfg_traits::*;
use middle::ssa::ssa_traits::*;
use middle::ssa::ssastorage::SSAStorage;

use r2api::structs::LCCInfo;

use petgraph::prelude::*;
use petgraph::visit::{DfsPostOrder, VisitMap, Walker};

use std::collections::{hash_map, BTreeMap, HashMap, HashSet};
use std::iter::{self, Extend, FromIterator};

/// For every function, patch all of its call sites to ignore registers that the
/// callee doesn't read and to preserve register values that the callee
/// preserves. Then, record which registers it reads and which registers it
/// preserves.
///
/// After this, all functions should have mutable [`regusage`][RadecoFunction::regusage]s.
///
/// This analysis is super conservative; for example, if a function preserves a
/// register by pushing it onto the stack and popping it back right before
/// returning, it is considered to be read and not preserved because we can't
/// guarantee that that stack location is never subsequently read or modified.
pub fn run(rmod: &mut RadecoModule, reginfo: &SubRegisterFile) -> () {
    let mut a = Inferer::new();
    a.run(rmod, reginfo);
    for (_, rfn) in &rmod.functions {
        eprintln!("{:?} (@ {:#X}):", rfn.name, rfn.offset);
        for (i, regname) in reginfo.whole_names.iter().enumerate() {
            let rid = RegisterId::from_usize(i);
            let is_ignored = rfn.regusage.is_ignored(rid);
            let is_preserved = rfn.regusage.is_preserved(rid);
            eprintln!(
                "  {:>8}: {} {}",
                regname,
                if is_ignored { 'I' } else { 'R' },
                if is_preserved { 'P' } else { 'C' }
            );
        }
    }
}

struct Inferer {
    /// Addresses of the functions we've already analyzed
    analyzed: HashSet<u64>,
}

impl Inferer {
    fn new() -> Inferer {
        Inferer {
            analyzed: HashSet::new(),
        }
    }

    /// Calls `patch_fn`, `dce::collect`, and `analyze_fn` on every function,
    /// callees first
    fn run(&mut self, rmod: &mut RadecoModule, reginfo: &SubRegisterFile) -> () {
        // for imports, *ASSUME* that the callconv that r2 says is correct
        {
            let imp_ru_iter = rmod.imports.iter().filter_map(|(&imp_addr, imp_info)| {
                let imp_rfn = imp_info.rfn.borrow();
                let regusage = reginfo.r2callconv_to_register_usage(
                    imp_rfn.callconv.as_ref()?, // ignore imports without callconvs
                    &*imp_rfn.callconv_name,
                )?;
                Some((imp_addr, regusage))
            });
            for (imp_addr, imp_ru) in imp_ru_iter {
                rmod.functions.get_mut(&imp_addr).unwrap().regusage = imp_ru;
                self.analyzed.insert(imp_addr);
            }
        }

        let mut dfs_wi = DfsPostOrder::empty(&rmod.callgraph).iter(&rmod.callgraph);
        // pick a function ...
        for fn_ni in rmod.callgraph.node_indices() {
            // ... and start a dfs on it
            dfs_wi.inner_mut().move_to(fn_ni);
            while let Some(fn_to_anal) = dfs_wi.next() {
                let fn_addr = rmod.callgraph[fn_to_anal];

                // ignore functions already in `call_convs` (probably because its an import)
                if !self.analyzed.contains(&fn_addr) {
                    self.patch_fn(fn_addr, &mut rmod.functions);

                    let rfn = &mut rmod.functions.get_mut(&fn_addr).unwrap();
                    dce::collect(rfn.ssa_mut());
                    // TODO: inst_combine

                    let ru = self.analyze_fn(rfn, reginfo).unwrap_or_else(|| {
                        radeco_err!("Failed to analyze fn: {:?} (@ {:#X})", rfn.name, fn_addr);
                        // if analysis failed, default to "reads and clobbers everything"
                        reginfo.new_register_usage()
                    });

                    rfn.regusage = ru;
                    self.analyzed.insert(fn_addr);
                }
            }
        }
    }

    /// Using the callconv info we've gathered so far, patch-up call sites to
    /// to remove arguments that the callee doesn't read and make values in
    /// callee-saved registers be preserved across the call.
    fn patch_fn(&self, fn_addr: u64, fn_map: &mut BTreeMap<u64, RadecoFunction>) -> () {
        radeco_trace!("patching calls in fn: {}", fn_map[&fn_addr].name);
        for node in fn_map[&fn_addr].ssa().inorder_walk() {
            if let Ok(NodeType::Op(ir::MOpcode::OpCall)) =
                fn_map[&fn_addr].ssa().node_data(node).map(|nd| nd.nt)
            {
                self.patch_call_node(node, fn_addr, fn_map)
                    .unwrap_or_else(|| {
                        radeco_warn!(
                            "failed to remove unused args for call at {:#X}",
                            fn_map[&fn_addr].ssa().address(node).unwrap()
                        );
                    });
            }
        }
    }

    fn patch_call_node(
        &self,
        call_node: <SSAStorage as SSA>::ValueRef,
        fn_addr: u64,
        fn_map: &mut BTreeMap<u64, RadecoFunction>,
    ) -> Option<()> {
        // bail on indirect or weird call
        let (call_tgt_addr, call_reg_map) = direct_call_info(fn_map[&fn_addr].ssa(), call_node)?;

        // remove unread args
        for (reg_idx, &op_node) in (0..).zip(&call_reg_map) {
            if fn_map[&call_tgt_addr]
                .regusage
                .is_ignored(RegisterId::from_usize(reg_idx))
            {
                fn_map
                    .get_mut(&fn_addr)
                    .unwrap()
                    .ssa_mut()
                    .op_unuse(call_node, op_node);
            }
        }

        // bridge preserved registers
        for use_node in fn_map[&fn_addr].ssa().uses_of(call_node) {
            let reg_idx = match fn_map[&fn_addr]
                .ssa()
                .sparse_operands_of(use_node)
                .as_slice()
            {
                &[(reg_idx, _)] => reg_idx,
                _ => panic!("invalid use of call as operand"),
            };
            if fn_map[&call_tgt_addr]
                .regusage
                .is_preserved(RegisterId::from_u8(reg_idx))
            {
                fn_map
                    .get_mut(&fn_addr)
                    .unwrap()
                    .ssa_mut()
                    .replace_value(use_node, call_reg_map[reg_idx as usize]);
            }
        }

        Some(())
    }

    fn analyze_fn(&self, rfn: &RadecoFunction, reginfo: &SubRegisterFile) -> Option<RegisterUsage> {
        eprintln!("analyzing fn: {}", rfn.name);
        let ssa = rfn.ssa();
        let entry_regstate_node = ssa.registers_in(ssa.entry_node()?)?;
        let exit_regstate_node = ssa.registers_in(ssa.exit_node()?)?;
        // some registers may not be present in the entry node;
        // this means that the function neither reads nor preserves that register
        let entry_regstate: Vec<Option<_>> = register_state_info(ssa, entry_regstate_node);
        // every register must be present in the exit node though
        let exit_regstate: Vec<_> = register_state_info(ssa, exit_regstate_node)
            .into_iter()
            .collect::<Option<_>>()?;

        let mut ret = reginfo.new_register_usage();
        ret.set_all_ignored();

        // ignore registers not in entry regstate
        let mut regstate_iter = (0..)
            .zip(exit_regstate)
            .zip(entry_regstate)
            .filter_map(|((i, x), on)| on.map(|n| (i, n, x)));

        for (i, reg_val_entry, reg_val_exit) in regstate_iter {
            if reg_val_exit == reg_val_entry {
                eprintln!(
                    "  set_preserved({})",
                    ssa.regnames.get(i as usize).map(|x| &**x).unwrap_or("mem")
                );
                ret.set_preserved(RegisterId::from_u8(i));
            }

            // find all uses, ignoring entry/exit register state
            let mut uses_iter = ssa.uses_of(reg_val_entry)
                .into_iter()
                .filter(|&n| n != entry_regstate_node && n != exit_regstate_node);
            if uses_iter.next().is_some() {
                ret.set_read(RegisterId::from_u8(i));
            }
        }

        Some(ret)
    }
}

/// Extracts the value of all registers at a `RegisterState` SSA node.
/// The length of the returned `Vec` is exactly `ssa.regnames.len() + 1`.
fn register_state_info(
    ssa: &SSAStorage,
    regstate_node: <SSAStorage as SSA>::ValueRef,
) -> Vec<Option<<SSAStorage as SSA>::ValueRef>> {
    let mut reg_opt_map: Vec<Option<_>> = vec![None; ssa.regnames.len() + 1];
    for (op_idx, op_node) in ssa.sparse_operands_of(regstate_node) {
        if let Some(p) = reg_opt_map.get_mut(op_idx as usize) {
            *p = Some(op_node);
        }
    }
    reg_opt_map
}

/// Extracts the call target address and the value of all registers
/// The length of the returned `Vec` is exactly `ssa.regnames.len() + 1`.
/// Returns `None` if the call is indirect or if not all registers have a value.
fn direct_call_info(
    ssa: &SSAStorage,
    call_node: <SSAStorage as SSA>::ValueRef,
) -> Option<(u64, Vec<<SSAStorage as SSA>::ValueRef>)> {
    let mut tgt_opt: Option<u64> = None;
    let mut reg_opt_map: Vec<Option<_>> = vec![None; ssa.regnames.len() + 1];

    for (op_idx, op_node) in ssa.sparse_operands_of(call_node) {
        if op_idx == 0 {
            tgt_opt = tgt_opt.or_else(|| ssa.constant(op_node));
        } else {
            if let Some(p) = reg_opt_map.get_mut((op_idx - 1) as usize) {
                *p = Some(op_node);
            }
        }
    }

    let reg_map_opt: Option<Vec<_>> = reg_opt_map.into_iter().collect();
    Some((tgt_opt?, reg_map_opt?))
}
<|endoftext|>


<commit_before>/// Filters operating on numbers
use std::collections::HashMap;

use serde_json::value::Value;

use errors::Result;

/// Returns a suffix if the value is greater or equal than 2. Suffix defaults to `s`
pub fn get(value: Value, args: HashMap<String, Value>) -> Result<Value> {
    let key = match args.get("key") {
        Some(val) => try_get_value!("get", "key", String, val),
        None => "s".to_string(),
    };

    match value.as_object() {
        Some(o) => o.get(&key).cloned().ok_or_else(|| {
            format!("Filter `get` tried to get key `{}` but it wasn't found", &key).into()
        }),
        None => bail!("Filter `get` was used on a value that isn't an object"),
    }
}

#[cfg(test)]
mod tests {
    use super::get;
    use serde_json::value::to_value;
    use std::collections::HashMap;

    #[test]
    fn test_get_filter_exists() {
        let mut obj = HashMap::new();
        obj.insert("1".to_string(), "first".to_string());
        obj.insert("2".to_string(), "second".to_string());

        let mut args = HashMap::new();
        args.insert("key".to_string(), to_value("1").unwrap());
        let result = get(to_value(&obj).unwrap(), args);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), to_value("first").unwrap());
    }

    #[test]
    fn test_get_filter_doesnt_exist() {
        let mut obj = HashMap::new();
        obj.insert("1".to_string(), "first".to_string());
        obj.insert("2".to_string(), "second".to_string());

        let mut args = HashMap::new();
        args.insert("key".to_string(), to_value("3").unwrap());
        let result = get(to_value(&obj).unwrap(), args);
        assert!(result.is_err());
    }
}
<commit_msg>fix object get filter<commit_after>/// Filters operating on numbers
use std::collections::HashMap;

use serde_json::value::Value;

use errors::Result;

/// Returns a value by a `key` argument from a given object
pub fn get(value: Value, args: HashMap<String, Value>) -> Result<Value> {
    let key = match args.get("key") {
        Some(val) => try_get_value!("get", "key", String, val),
        None => bail!("The `get` filter has to have an `key` argument"),
    };

    match value.as_object() {
        Some(o) => o.get(&key).cloned().ok_or_else(|| {
            format!("Filter `get` tried to get key `{}` but it wasn't found", &key).into()
        }),
        None => bail!("Filter `get` was used on a value that isn't an object"),
    }
}

#[cfg(test)]
mod tests {
    use super::get;
    use serde_json::value::to_value;
    use std::collections::HashMap;

    #[test]
    fn test_get_filter_exists() {
        let mut obj = HashMap::new();
        obj.insert("1".to_string(), "first".to_string());
        obj.insert("2".to_string(), "second".to_string());

        let mut args = HashMap::new();
        args.insert("key".to_string(), to_value("1").unwrap());
        let result = get(to_value(&obj).unwrap(), args);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), to_value("first").unwrap());
    }

    #[test]
    fn test_get_filter_doesnt_exist() {
        let mut obj = HashMap::new();
        obj.insert("1".to_string(), "first".to_string());
        obj.insert("2".to_string(), "second".to_string());

        let mut args = HashMap::new();
        args.insert("key".to_string(), to_value("3").unwrap());
        let result = get(to_value(&obj).unwrap(), args);
        assert!(result.is_err());
    }
}
<|endoftext|>


<commit_before>use super::{Open, Sink};
use libpulse_binding as pulse;
use libpulse_simple_binding::Simple;
use std::io;

const APP_NAME: &str = "librespot";
const STREAM_NAME: &str = "Spotify endpoint";

pub struct PulseAudioSink {
    s: Option<Simple>,
    ss: pulse::sample::Spec,
    device: Option<String>,
}

impl Open for PulseAudioSink {
    fn open(device: Option<String>) -> PulseAudioSink {
        debug!("Using PulseAudio sink");

        let ss = pulse::sample::Spec {
            format: pulse::sample::Format::S16le,
            channels: 2, // stereo
            rate: 44100,
        };
        debug_assert!(ss.is_valid());

        PulseAudioSink {
            s: None,
            ss: ss,
            device: device,
        }
    }
}

impl Sink for PulseAudioSink {
    fn start(&mut self) -> io::Result<()> {
        if self.s.is_some() {
            return Ok(());
        }

        let device = self.device.as_ref().map(|s| (*s).as_str());
        let result = Simple::new(
            None,               // Use the default server.
            APP_NAME,           // Our application's name.
            pulse::stream::Direction::Playback,
            device,
            STREAM_NAME,        // desc of our stream.
            &self.ss,           // Our sample format.
            None,               // Use default channel map
            None,               // Use default buffering attributes.
        );
        match result {
            Ok(s) => {
                self.s = Some(s);
                Ok(())
            }
            Err(e) => Err(io::Error::new(
                io::ErrorKind::ConnectionRefused,
                e.to_string().unwrap(),
            )),
        }
    }

    fn stop(&mut self) -> io::Result<()> {
        self.s = None;
        Ok(())
    }

    fn write(&mut self, data: &[i16]) -> io::Result<()> {
        if let Some(s) = &self.s {
            let d: &[u8] = unsafe { std::mem::transmute(data) };

            match s.write(d) {
                Ok(_) => Ok(()),
                Err(e) => Err(io::Error::new(
                    io::ErrorKind::BrokenPipe,
                    e.to_string().unwrap(),
                )),
            }
        } else {
            Err(io::Error::new(
                io::ErrorKind::NotConnected,
                "Not connected to pulseaudio",
            ))
        }
    }
}
<commit_msg>fix alignment causing format check failure<commit_after>use super::{Open, Sink};
use libpulse_binding::{self as pulse, stream::Direction};
use libpulse_simple_binding::Simple;
use std::io;

const APP_NAME: &str = "librespot";
const STREAM_NAME: &str = "Spotify endpoint";

pub struct PulseAudioSink {
    s: Option<Simple>,
    ss: pulse::sample::Spec,
    device: Option<String>,
}

impl Open for PulseAudioSink {
    fn open(device: Option<String>) -> PulseAudioSink {
        debug!("Using PulseAudio sink");

        let ss = pulse::sample::Spec {
            format: pulse::sample::Format::S16le,
            channels: 2, // stereo
            rate: 44100,
        };
        debug_assert!(ss.is_valid());

        PulseAudioSink {
            s: None,
            ss: ss,
            device: device,
        }
    }
}

impl Sink for PulseAudioSink {
    fn start(&mut self) -> io::Result<()> {
        if self.s.is_some() {
            return Ok(());
        }

        let device = self.device.as_ref().map(|s| (*s).as_str());
        let result = Simple::new(
            None,                // Use the default server.
            APP_NAME,            // Our application's name.
            Direction::Playback, // Direction.
            device,              // Our device (sink) name.
            STREAM_NAME,         // Description of our stream.
            &self.ss,            // Our sample format.
            None,                // Use default channel map.
            None,                // Use default buffering attributes.
        );
        match result {
            Ok(s) => {
                self.s = Some(s);
                Ok(())
            }
            Err(e) => Err(io::Error::new(
                io::ErrorKind::ConnectionRefused,
                e.to_string().unwrap(),
            )),
        }
    }

    fn stop(&mut self) -> io::Result<()> {
        self.s = None;
        Ok(())
    }

    fn write(&mut self, data: &[i16]) -> io::Result<()> {
        if let Some(s) = &self.s {
            let d: &[u8] = unsafe { std::mem::transmute(data) };

            match s.write(d) {
                Ok(_) => Ok(()),
                Err(e) => Err(io::Error::new(
                    io::ErrorKind::BrokenPipe,
                    e.to_string().unwrap(),
                )),
            }
        } else {
            Err(io::Error::new(
                io::ErrorKind::NotConnected,
                "Not connected to pulseaudio",
            ))
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! Structs for creating and using a Resolver
use std::io;
use std::net::IpAddr;
use std::sync::Mutex;

use proto::rr::RecordType;
use tokio::runtime::{self, Runtime};

use config::{ResolverConfig, ResolverOpts};
use error::*;
use lookup;
use lookup::Lookup;
use lookup_ip::LookupIp;
use AsyncResolver;

/// The Resolver is used for performing DNS queries.
///
/// For forward (A) lookups, hostname -> IP address, see: `Resolver::lookup_ip`
///
/// Special note about resource consumption. The Resolver and all Trust-DNS software is built around the Tokio async-io library. This synchronous Resolver is intended to be a simpler wrapper for of the [`AsyncResolver`]. To allow the `Resolver` to be [`Send`] + [`Sync`], the construction of the `AsyncResolver` is lazy, this means some of the features of the `AsyncResolver`, like performance based resolution via the most efficient `NameServer` will be lost (the lookup cache is shared across invocations of the `Resolver`). If these other features of the Trust-DNS Resolver are desired, please use the tokio based [`AsyncResolver`].
///
/// *Note: Threaded/Sync usage*: In multithreaded scenarios, the internal Tokio Runtime will block on an internal Mutex for the tokio Runtime in use. For higher performance, it's recommended to use the [`AsyncResolver`].
pub struct Resolver {
    // TODO: Mutex allows this to be Sync, another option would be to instantiate a thread_local, but that has other
    //   drawbacks. One major issues, is if this Resolver is shared across threads, it will cause all to block on any
    //   query. A TLS on the other hand would not, at the cost of only allowing a Resolver to be configured once per Thread
    runtime: Mutex<Runtime>,
    async_resolver: AsyncResolver,
}

macro_rules! lookup_fn {
    ($p:ident, $l:ty) => {
/// Performs a lookup for the associated type.
///
/// *hint* queries that end with a '.' are fully qualified names and are cheaper lookups
///
/// # Arguments
///
/// * `query` - a `&str` which parses to a domain name, failure to parse will return an error
pub fn $p(&self, query: &str) -> ResolveResult<$l> {
    let lookup = self.async_resolver.$p(query);
    self.runtime.lock()?.block_on(lookup)
}
    };
    ($p:ident, $l:ty, $t:ty) => {
/// Performs a lookup for the associated type.
///
/// # Arguments
///
/// * `query` - a type which can be converted to `Name` via `From`.
pub fn $p(&self, query: $t) -> ResolveResult<$l> {
    let lookup = self.async_resolver.$p(query);
    self.runtime.lock()?.block_on(lookup)
}
    };
}

impl Resolver {
    /// Constructs a new Resolver with the specified configuration.
    ///
    /// # Arguments
    /// * `config` - configuration for the resolver
    /// * `options` - resolver options for performing lookups
    /// * `client_connection` - ClientConnection for establishing the connection to the DNS server
    ///
    /// # Returns
    ///
    /// A new `Resolver` or an error if there was an error with the configuration.
    pub fn new(config: ResolverConfig, options: ResolverOpts) -> io::Result<Self> {
        let mut builder = runtime::Builder::new();
        builder.core_threads(1);

        let mut runtime = builder.build()?;
        let (async_resolver, bg) = AsyncResolver::new(config, options);

        runtime.spawn(bg);

        Ok(Resolver {
            runtime: Mutex::new(runtime),
            async_resolver,
        })
    }

    /// Constructs a new Resolver with default config and default options.
    ///
    /// See [`ResolverConfig::default`] and [`ResolverOpts::default`] for more information.
    ///
    /// # Returns
    ///
    /// A new `Resolver` or an error if there was an error with the configuration.
    pub fn default() -> io::Result<Self> {
        Self::new(ResolverConfig::default(), ResolverOpts::default())
    }

    /// Constructs a new Resolver with the system configuration.
    ///
    /// This will use `/etc/resolv.conf` on Unix OSes and the registry on Windows.
    #[cfg(any(unix, target_os = "windows"))]
    pub fn from_system_conf() -> io::Result<Self> {
        let (config, options) = super::system_conf::read_system_conf()?;
        Self::new(config, options)
    }

    /// Generic lookup for any RecordType
    ///
    /// *WARNING* This interface may change in the future, please use [`Self::lookup_ip`] or another variant for more stable interfaces.
    ///
    /// # Arguments
    ///
    /// * `name` - name of the record to lookup, if name is not a valid domain name, an error will be returned
    /// * `record_type` - type of record to lookup
    pub fn lookup(&self, name: &str, record_type: RecordType) -> ResolveResult<Lookup> {
        let lookup = self.async_resolver.lookup(name, record_type);
        self.runtime.lock()?.block_on(lookup)
    }

    /// Performs a dual-stack DNS lookup for the IP for the given hostname.
    ///
    /// See the configuration and options parameters for controlling the way in which A(Ipv4) and AAAA(Ipv6) lookups will be performed. For the least expensive query a fully-qualified-domain-name, FQDN, which ends in a final `.`, e.g. `www.example.com.`, will only issue one query. Anything else will always incur the cost of querying the `ResolverConfig::domain` and `ResolverConfig::search`.
    ///
    /// # Arguments
    ///
    /// * `host` - string hostname, if this is an invalid hostname, an error will be returned.
    pub fn lookup_ip(&self, host: &str) -> ResolveResult<LookupIp> {
        let lookup = self.async_resolver.lookup_ip(host);
        self.runtime.lock()?.block_on(lookup)
    }

    /// Performs a DNS lookup for an SRV record for the specified service type and protocol at the given name.
    ///
    /// This is a convenience method over `lookup_srv`, it combines the service, protocol and name into a single name: `_service._protocol.name`.
    ///
    /// # Arguments
    ///
    /// * `service` - service to lookup, e.g. ldap or http
    /// * `protocol` - wire protocol, e.g. udp or tcp
    /// * `name` - zone or other name at which the service is located.
    #[deprecated(note = "use lookup_srv instead, this interface is not ideal")]
    pub fn lookup_service(
        &self,
        service: &str,
        protocol: &str,
        name: &str,
    ) -> ResolveResult<lookup::SrvLookup> {
        #[allow(deprecated)]
        let lookup = self.async_resolver.lookup_service(service, protocol, name);
        self.runtime.lock()?.block_on(lookup)
    }

    /// Lookup an SRV record.
    pub fn lookup_srv(&self, name: &str) -> ResolveResult<lookup::SrvLookup> {
        let lookup = self.async_resolver.lookup_srv(name);
        self.runtime.lock()?.block_on(lookup)
    }

    lookup_fn!(reverse_lookup, lookup::ReverseLookup, IpAddr);
    lookup_fn!(ipv4_lookup, lookup::Ipv4Lookup);
    lookup_fn!(ipv6_lookup, lookup::Ipv6Lookup);
    lookup_fn!(mx_lookup, lookup::MxLookup);
    #[deprecated(note = "use lookup_srv instead, this interface is not ideal")]
    lookup_fn!(srv_lookup, lookup::SrvLookup);
    lookup_fn!(txt_lookup, lookup::TxtLookup);
}

#[cfg(test)]
mod tests {
    use std::net::*;

    use super::*;

    fn require_send_sync<S: Send + Sync>() {}

    #[test]
    fn test_resolver_sendable() {
        require_send_sync::<Resolver>();
    }

    #[test]
    fn test_lookup() {
        let resolver = Resolver::new(ResolverConfig::default(), ResolverOpts::default()).unwrap();

        let response = resolver.lookup_ip("www.example.com.").unwrap();
        println!("response records: {:?}", response);

        assert_eq!(response.iter().count(), 1);
        for address in response.iter() {
            if address.is_ipv4() {
                assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
            } else {
                assert_eq!(
                    address,
                    IpAddr::V6(Ipv6Addr::new(
                        0x2606, 0x2800, 0x220, 0x1, 0x248, 0x1893, 0x25c8, 0x1946,
                    ))
                );
            }
        }
    }

    #[test]
    #[ignore]
    #[cfg(any(unix, target_os = "windows"))]
    fn test_system_lookup() {
        let resolver = Resolver::from_system_conf().unwrap();

        let response = resolver.lookup_ip("www.example.com.").unwrap();
        println!("response records: {:?}", response);

        assert_eq!(response.iter().count(), 1);
        for address in response.iter() {
            if address.is_ipv4() {
                assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
            } else {
                assert_eq!(
                    address,
                    IpAddr::V6(Ipv6Addr::new(
                        0x2606, 0x2800, 0x220, 0x1, 0x248, 0x1893, 0x25c8, 0x1946,
                    ))
                );
            }
        }
    }
}
<commit_msg>Remove arg from doc comments that isn't present any longer.<commit_after>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! Structs for creating and using a Resolver
use std::io;
use std::net::IpAddr;
use std::sync::Mutex;

use proto::rr::RecordType;
use tokio::runtime::{self, Runtime};

use config::{ResolverConfig, ResolverOpts};
use error::*;
use lookup;
use lookup::Lookup;
use lookup_ip::LookupIp;
use AsyncResolver;

/// The Resolver is used for performing DNS queries.
///
/// For forward (A) lookups, hostname -> IP address, see: `Resolver::lookup_ip`
///
/// Special note about resource consumption. The Resolver and all Trust-DNS software is built around the Tokio async-io library. This synchronous Resolver is intended to be a simpler wrapper for of the [`AsyncResolver`]. To allow the `Resolver` to be [`Send`] + [`Sync`], the construction of the `AsyncResolver` is lazy, this means some of the features of the `AsyncResolver`, like performance based resolution via the most efficient `NameServer` will be lost (the lookup cache is shared across invocations of the `Resolver`). If these other features of the Trust-DNS Resolver are desired, please use the tokio based [`AsyncResolver`].
///
/// *Note: Threaded/Sync usage*: In multithreaded scenarios, the internal Tokio Runtime will block on an internal Mutex for the tokio Runtime in use. For higher performance, it's recommended to use the [`AsyncResolver`].
pub struct Resolver {
    // TODO: Mutex allows this to be Sync, another option would be to instantiate a thread_local, but that has other
    //   drawbacks. One major issues, is if this Resolver is shared across threads, it will cause all to block on any
    //   query. A TLS on the other hand would not, at the cost of only allowing a Resolver to be configured once per Thread
    runtime: Mutex<Runtime>,
    async_resolver: AsyncResolver,
}

macro_rules! lookup_fn {
    ($p:ident, $l:ty) => {
/// Performs a lookup for the associated type.
///
/// *hint* queries that end with a '.' are fully qualified names and are cheaper lookups
///
/// # Arguments
///
/// * `query` - a `&str` which parses to a domain name, failure to parse will return an error
pub fn $p(&self, query: &str) -> ResolveResult<$l> {
    let lookup = self.async_resolver.$p(query);
    self.runtime.lock()?.block_on(lookup)
}
    };
    ($p:ident, $l:ty, $t:ty) => {
/// Performs a lookup for the associated type.
///
/// # Arguments
///
/// * `query` - a type which can be converted to `Name` via `From`.
pub fn $p(&self, query: $t) -> ResolveResult<$l> {
    let lookup = self.async_resolver.$p(query);
    self.runtime.lock()?.block_on(lookup)
}
    };
}

impl Resolver {
    /// Constructs a new Resolver with the specified configuration.
    ///
    /// # Arguments
    /// * `config` - configuration for the resolver
    /// * `options` - resolver options for performing lookups
    ///
    /// # Returns
    ///
    /// A new `Resolver` or an error if there was an error with the configuration.
    pub fn new(config: ResolverConfig, options: ResolverOpts) -> io::Result<Self> {
        let mut builder = runtime::Builder::new();
        builder.core_threads(1);

        let mut runtime = builder.build()?;
        let (async_resolver, bg) = AsyncResolver::new(config, options);

        runtime.spawn(bg);

        Ok(Resolver {
            runtime: Mutex::new(runtime),
            async_resolver,
        })
    }

    /// Constructs a new Resolver with default config and default options.
    ///
    /// See [`ResolverConfig::default`] and [`ResolverOpts::default`] for more information.
    ///
    /// # Returns
    ///
    /// A new `Resolver` or an error if there was an error with the configuration.
    pub fn default() -> io::Result<Self> {
        Self::new(ResolverConfig::default(), ResolverOpts::default())
    }

    /// Constructs a new Resolver with the system configuration.
    ///
    /// This will use `/etc/resolv.conf` on Unix OSes and the registry on Windows.
    #[cfg(any(unix, target_os = "windows"))]
    pub fn from_system_conf() -> io::Result<Self> {
        let (config, options) = super::system_conf::read_system_conf()?;
        Self::new(config, options)
    }

    /// Generic lookup for any RecordType
    ///
    /// *WARNING* This interface may change in the future, please use [`Self::lookup_ip`] or another variant for more stable interfaces.
    ///
    /// # Arguments
    ///
    /// * `name` - name of the record to lookup, if name is not a valid domain name, an error will be returned
    /// * `record_type` - type of record to lookup
    pub fn lookup(&self, name: &str, record_type: RecordType) -> ResolveResult<Lookup> {
        let lookup = self.async_resolver.lookup(name, record_type);
        self.runtime.lock()?.block_on(lookup)
    }

    /// Performs a dual-stack DNS lookup for the IP for the given hostname.
    ///
    /// See the configuration and options parameters for controlling the way in which A(Ipv4) and AAAA(Ipv6) lookups will be performed. For the least expensive query a fully-qualified-domain-name, FQDN, which ends in a final `.`, e.g. `www.example.com.`, will only issue one query. Anything else will always incur the cost of querying the `ResolverConfig::domain` and `ResolverConfig::search`.
    ///
    /// # Arguments
    ///
    /// * `host` - string hostname, if this is an invalid hostname, an error will be returned.
    pub fn lookup_ip(&self, host: &str) -> ResolveResult<LookupIp> {
        let lookup = self.async_resolver.lookup_ip(host);
        self.runtime.lock()?.block_on(lookup)
    }

    /// Performs a DNS lookup for an SRV record for the specified service type and protocol at the given name.
    ///
    /// This is a convenience method over `lookup_srv`, it combines the service, protocol and name into a single name: `_service._protocol.name`.
    ///
    /// # Arguments
    ///
    /// * `service` - service to lookup, e.g. ldap or http
    /// * `protocol` - wire protocol, e.g. udp or tcp
    /// * `name` - zone or other name at which the service is located.
    #[deprecated(note = "use lookup_srv instead, this interface is not ideal")]
    pub fn lookup_service(
        &self,
        service: &str,
        protocol: &str,
        name: &str,
    ) -> ResolveResult<lookup::SrvLookup> {
        #[allow(deprecated)]
        let lookup = self.async_resolver.lookup_service(service, protocol, name);
        self.runtime.lock()?.block_on(lookup)
    }

    /// Lookup an SRV record.
    pub fn lookup_srv(&self, name: &str) -> ResolveResult<lookup::SrvLookup> {
        let lookup = self.async_resolver.lookup_srv(name);
        self.runtime.lock()?.block_on(lookup)
    }

    lookup_fn!(reverse_lookup, lookup::ReverseLookup, IpAddr);
    lookup_fn!(ipv4_lookup, lookup::Ipv4Lookup);
    lookup_fn!(ipv6_lookup, lookup::Ipv6Lookup);
    lookup_fn!(mx_lookup, lookup::MxLookup);
    #[deprecated(note = "use lookup_srv instead, this interface is not ideal")]
    lookup_fn!(srv_lookup, lookup::SrvLookup);
    lookup_fn!(txt_lookup, lookup::TxtLookup);
}

#[cfg(test)]
mod tests {
    use std::net::*;

    use super::*;

    fn require_send_sync<S: Send + Sync>() {}

    #[test]
    fn test_resolver_sendable() {
        require_send_sync::<Resolver>();
    }

    #[test]
    fn test_lookup() {
        let resolver = Resolver::new(ResolverConfig::default(), ResolverOpts::default()).unwrap();

        let response = resolver.lookup_ip("www.example.com.").unwrap();
        println!("response records: {:?}", response);

        assert_eq!(response.iter().count(), 1);
        for address in response.iter() {
            if address.is_ipv4() {
                assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
            } else {
                assert_eq!(
                    address,
                    IpAddr::V6(Ipv6Addr::new(
                        0x2606, 0x2800, 0x220, 0x1, 0x248, 0x1893, 0x25c8, 0x1946,
                    ))
                );
            }
        }
    }

    #[test]
    #[ignore]
    #[cfg(any(unix, target_os = "windows"))]
    fn test_system_lookup() {
        let resolver = Resolver::from_system_conf().unwrap();

        let response = resolver.lookup_ip("www.example.com.").unwrap();
        println!("response records: {:?}", response);

        assert_eq!(response.iter().count(), 1);
        for address in response.iter() {
            if address.is_ipv4() {
                assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
            } else {
                assert_eq!(
                    address,
                    IpAddr::V6(Ipv6Addr::new(
                        0x2606, 0x2800, 0x220, 0x1, 0x248, 0x1893, 0x25c8, 0x1946,
                    ))
                );
            }
        }
    }
}
<|endoftext|>