<commit_before>// Copyright 2020 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Provides a connection wrapper that handles the lower level tasks in sending
//! or receiving data from the TCP socket, as well as dealing with timeouts.
//!
//! Because of a few idiosyncracies in the Rust `TcpStream`, this has to use
//! async I/O to be able to both read *and* write on the connection. Which
//! forces us to go through some additional gymnastic to loop over the async
//! stream and make sure we get the right number of bytes out.

use crate::core::ser;
use crate::core::ser::ProtocolVersion;
use crate::msg::{
	read_body, read_discard, read_header, read_item, write_message, Msg, MsgHeader,
	MsgHeaderWrapper,
};
use crate::types::Error;
use crate::util::{RateCounter, RwLock};
use std::io::{self, Read, Write};
use std::net::{Shutdown, TcpStream};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{mpsc, Arc};
use std::time::Duration;
use std::{
	cmp,
	thread::{self, JoinHandle},
};

pub const SEND_CHANNEL_CAP: usize = 100;

const HEADER_IO_TIMEOUT: Duration = Duration::from_millis(2000);
const CHANNEL_TIMEOUT: Duration = Duration::from_millis(1000);
const BODY_IO_TIMEOUT: Duration = Duration::from_millis(60000);

/// A trait to be implemented in order to receive messages from the
/// connection. Allows providing an optional response.
pub trait MessageHandler: Send + 'static {
	fn consume<'a>(
		&self,
		msg: Message<'a>,
		stopped: Arc<AtomicBool>,
		tracker: Arc<Tracker>,
	) -> Result<Option<Msg>, Error>;
}

// Macro to simplify the boilerplate around I/O and Grin error handling
macro_rules! try_break {
	($inner:expr) => {
		match $inner {
			Ok(v) => Some(v),
			Err(Error::Connection(ref e)) if e.kind() == io::ErrorKind::TimedOut => None,
			Err(Error::Connection(ref e)) if e.kind() == io::ErrorKind::WouldBlock => {
				// to avoid the heavy polling which will consume CPU 100%
				thread::sleep(Duration::from_millis(10));
				None
				}
			Err(Error::Store(_))
			| Err(Error::Chain(_))
			| Err(Error::Internal)
			| Err(Error::NoDandelionRelay) => None,
			Err(ref e) => {
				debug!("try_break: exit the loop: {:?}", e);
				break;
				}
			}
	};
}

macro_rules! try_header {
	($res:expr, $conn: expr) => {{
		$conn
			.set_read_timeout(Some(HEADER_IO_TIMEOUT))
			.expect("set timeout");
		try_break!($res)
		}};
}

/// A message as received by the connection. Provides access to the message
/// header lazily consumes the message body, handling its deserialization.
pub struct Message<'a> {
	pub header: MsgHeader,
	stream: &'a mut dyn Read,
	version: ProtocolVersion,
}

impl<'a> Message<'a> {
	fn from_header(
		header: MsgHeader,
		stream: &'a mut dyn Read,
		version: ProtocolVersion,
	) -> Message<'a> {
		Message {
			header,
			stream,
			version,
		}
	}

	/// Read the message body from the underlying connection
	pub fn body<T: ser::Readable>(&mut self) -> Result<T, Error> {
		read_body(&self.header, self.stream, self.version)
	}

	/// Read a single "thing" from the underlying connection.
	/// Return the thing and the total bytes read.
	pub fn streaming_read<T: ser::Readable>(&mut self) -> Result<(T, u64), Error> {
		read_item(self.stream, self.version)
	}

	pub fn copy_attachment(&mut self, len: usize, writer: &mut dyn Write) -> Result<usize, Error> {
		let mut written = 0;
		while written < len {
			let read_len = cmp::min(8000, len - written);
			let mut buf = vec![0u8; read_len];
			self.stream.read_exact(&mut buf[..])?;
			writer.write_all(&buf)?;
			written += read_len;
		}
		Ok(written)
	}
}

pub struct StopHandle {
	/// Channel to close the connection
	stopped: Arc<AtomicBool>,
	// we need Option to take ownhership of the handle in stop()
	reader_thread: Option<JoinHandle<()>>,
	writer_thread: Option<JoinHandle<()>>,
}

impl StopHandle {
	/// Schedule this connection to safely close via the async close_channel.
	pub fn stop(&self) {
		self.stopped.store(true, Ordering::Relaxed);
	}

	pub fn wait(&mut self) {
		if let Some(reader_thread) = self.reader_thread.take() {
			self.join_thread(reader_thread);
		}
		if let Some(writer_thread) = self.writer_thread.take() {
			self.join_thread(writer_thread);
		}
	}

	fn join_thread(&self, peer_thread: JoinHandle<()>) {
		// wait only if other thread is calling us, eg shutdown
		if thread::current().id() != peer_thread.thread().id() {
			debug!("waiting for thread {:?} exit", peer_thread.thread().id());
			if let Err(e) = peer_thread.join() {
				error!("failed to stop peer thread: {:?}", e);
			}
		} else {
			debug!(
				"attempt to stop thread {:?} from itself",
				peer_thread.thread().id()
			);
		}
	}
}

#[derive(Clone)]
pub struct ConnHandle {
	/// Channel to allow sending data through the connection
	pub send_channel: mpsc::SyncSender<Msg>,
}

impl ConnHandle {
	/// Send msg via the synchronous, bounded channel (sync_sender).
	/// Two possible failure cases -
	/// * Disconnected: Propagate this up to the caller so the peer connection can be closed.
	/// * Full: Our internal msg buffer is full. This is not a problem with the peer connection
	/// and we do not want to close the connection. We drop the msg rather than blocking here.
	/// If the buffer is full because there is an underlying issue with the peer
	/// and potentially the peer connection. We assume this will be handled at the peer level.
	pub fn send(&self, msg: Msg) -> Result<(), Error> {
		match self.send_channel.try_send(msg) {
			Ok(()) => Ok(()),
			Err(mpsc::TrySendError::Disconnected(_)) => {
				Err(Error::Send("try_send disconnected".to_owned()))
			}
			Err(mpsc::TrySendError::Full(_)) => {
				debug!("conn_handle: try_send but buffer is full, dropping msg");
				Ok(())
			}
		}
	}
}

pub struct Tracker {
	/// Bytes we've sent.
	pub sent_bytes: Arc<RwLock<RateCounter>>,
	/// Bytes we've received.
	pub received_bytes: Arc<RwLock<RateCounter>>,
}

impl Tracker {
	pub fn new() -> Tracker {
		let received_bytes = Arc::new(RwLock::new(RateCounter::new()));
		let sent_bytes = Arc::new(RwLock::new(RateCounter::new()));
		Tracker {
			received_bytes,
			sent_bytes,
		}
	}

	pub fn inc_received(&self, size: u64) {
		self.received_bytes.write().inc(size);
	}

	pub fn inc_sent(&self, size: u64) {
		self.sent_bytes.write().inc(size);
	}

	pub fn inc_quiet_received(&self, size: u64) {
		self.received_bytes.write().inc_quiet(size);
	}

	pub fn inc_quiet_sent(&self, size: u64) {
		self.sent_bytes.write().inc_quiet(size);
	}
}

/// Start listening on the provided connection and wraps it. Does not hang
/// the current thread, instead just returns a future and the Connection
/// itself.
pub fn listen<H>(
	stream: TcpStream,
	version: ProtocolVersion,
	tracker: Arc<Tracker>,
	handler: H,
) -> io::Result<(ConnHandle, StopHandle)>
where
	H: MessageHandler,
{
	let (send_tx, send_rx) = mpsc::sync_channel(SEND_CHANNEL_CAP);

	let stopped = Arc::new(AtomicBool::new(false));

	let conn_handle = ConnHandle {
		send_channel: send_tx,
	};

	let (reader_thread, writer_thread) = poll(
		stream,
		conn_handle.clone(),
		version,
		handler,
		send_rx,
		stopped.clone(),
		tracker,
	)?;

	Ok((
		conn_handle,
		StopHandle {
			stopped,
			reader_thread: Some(reader_thread),
			writer_thread: Some(writer_thread),
		},
	))
}

fn poll<H>(
	conn: TcpStream,
	conn_handle: ConnHandle,
	version: ProtocolVersion,
	handler: H,
	send_rx: mpsc::Receiver<Msg>,
	stopped: Arc<AtomicBool>,
	tracker: Arc<Tracker>,
) -> io::Result<(JoinHandle<()>, JoinHandle<()>)>
where
	H: MessageHandler,
{
	// Split out tcp stream out into separate reader/writer halves.
	let mut reader = conn.try_clone().expect("clone conn for reader failed");
	let mut writer = conn.try_clone().expect("clone conn for writer failed");
	let reader_stopped = stopped.clone();

	let reader_tracker = tracker.clone();
	let writer_tracker = tracker;

	let reader_thread = thread::Builder::new()
		.name("peer_read".to_string())
		.spawn(move || {
			loop {
				// check the read end
				match try_header!(read_header(&mut reader, version), &reader) {
					Some(MsgHeaderWrapper::Known(header)) => {
						reader
							.set_read_timeout(Some(BODY_IO_TIMEOUT))
							.expect("set timeout");
						let msg = Message::from_header(header, &mut reader, version);

						trace!(
							"Received message header, type {:?}, len {}.",
							msg.header.msg_type,
							msg.header.msg_len
						);

						// Increase received bytes counter
						reader_tracker.inc_received(MsgHeader::LEN as u64 + msg.header.msg_len);

						let resp_msg = try_break!(handler.consume(
							msg,
							reader_stopped.clone(),
							reader_tracker.clone()
						));
						if let Some(Some(resp_msg)) = resp_msg {
							try_break!(conn_handle.send(resp_msg));
						}
					}
					Some(MsgHeaderWrapper::Unknown(msg_len, type_byte)) => {
						debug!(
							"Received unknown message header, type {:?}, len {}.",
							type_byte, msg_len
						);
						// Increase received bytes counter
						reader_tracker.inc_received(MsgHeader::LEN as u64 + msg_len);

						try_break!(read_discard(msg_len, &mut reader));
					}
					None => {}
				}

				// check the close channel
				if reader_stopped.load(Ordering::Relaxed) {
					break;
				}
			}

			debug!(
				"Shutting down reader connection with {}",
				reader
					.peer_addr()
					.map(|a| a.to_string())
					.unwrap_or_else(|_| "?".to_owned())
			);
			let _ = reader.shutdown(Shutdown::Both);
		})?;

	let writer_thread = thread::Builder::new()
		.name("peer_write".to_string())
		.spawn(move || {
			let mut retry_send = Err(());
			writer
				.set_write_timeout(Some(BODY_IO_TIMEOUT))
				.expect("set timeout");
			loop {
				let maybe_data = retry_send.or_else(|_| send_rx.recv_timeout(CHANNEL_TIMEOUT));
				retry_send = Err(());
				if let Ok(data) = maybe_data {
					let written =
						try_break!(write_message(&mut writer, &data, writer_tracker.clone()));
					if written.is_none() {
						retry_send = Ok(data);
					}
				}
				// check the close channel
				if stopped.load(Ordering::Relaxed) {
					break;
				}
			}

			debug!(
				"Shutting down writer connection with {}",
				writer
					.peer_addr()
					.map(|a| a.to_string())
					.unwrap_or_else(|_| "?".to_owned())
			);
		})?;
	Ok((reader_thread, writer_thread))
}
<commit_msg>handle mpsc channel disconnect from peer_write thread (#3241)<commit_after>// Copyright 2020 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Provides a connection wrapper that handles the lower level tasks in sending
//! or receiving data from the TCP socket, as well as dealing with timeouts.
//!
//! Because of a few idiosyncracies in the Rust `TcpStream`, this has to use
//! async I/O to be able to both read *and* write on the connection. Which
//! forces us to go through some additional gymnastic to loop over the async
//! stream and make sure we get the right number of bytes out.

use crate::core::ser;
use crate::core::ser::ProtocolVersion;
use crate::msg::{
	read_body, read_discard, read_header, read_item, write_message, Msg, MsgHeader,
	MsgHeaderWrapper,
};
use crate::types::Error;
use crate::util::{RateCounter, RwLock};
use std::io::{self, Read, Write};
use std::net::{Shutdown, TcpStream};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::RecvTimeoutError;
use std::sync::{mpsc, Arc};
use std::time::Duration;
use std::{
	cmp,
	thread::{self, JoinHandle},
};

pub const SEND_CHANNEL_CAP: usize = 100;

const HEADER_IO_TIMEOUT: Duration = Duration::from_millis(2000);
const CHANNEL_TIMEOUT: Duration = Duration::from_millis(1000);
const BODY_IO_TIMEOUT: Duration = Duration::from_millis(60000);

/// A trait to be implemented in order to receive messages from the
/// connection. Allows providing an optional response.
pub trait MessageHandler: Send + 'static {
	fn consume<'a>(
		&self,
		msg: Message<'a>,
		stopped: Arc<AtomicBool>,
		tracker: Arc<Tracker>,
	) -> Result<Option<Msg>, Error>;
}

// Macro to simplify the boilerplate around I/O and Grin error handling
macro_rules! try_break {
	($inner:expr) => {
		match $inner {
			Ok(v) => Some(v),
			Err(Error::Connection(ref e)) if e.kind() == io::ErrorKind::TimedOut => None,
			Err(Error::Connection(ref e)) if e.kind() == io::ErrorKind::WouldBlock => {
				// to avoid the heavy polling which will consume CPU 100%
				thread::sleep(Duration::from_millis(10));
				None
				}
			Err(Error::Store(_))
			| Err(Error::Chain(_))
			| Err(Error::Internal)
			| Err(Error::NoDandelionRelay) => None,
			Err(ref e) => {
				debug!("try_break: exit the loop: {:?}", e);
				break;
				}
			}
	};
}

macro_rules! try_header {
	($res:expr, $conn: expr) => {{
		$conn
			.set_read_timeout(Some(HEADER_IO_TIMEOUT))
			.expect("set timeout");
		try_break!($res)
		}};
}

/// A message as received by the connection. Provides access to the message
/// header lazily consumes the message body, handling its deserialization.
pub struct Message<'a> {
	pub header: MsgHeader,
	stream: &'a mut dyn Read,
	version: ProtocolVersion,
}

impl<'a> Message<'a> {
	fn from_header(
		header: MsgHeader,
		stream: &'a mut dyn Read,
		version: ProtocolVersion,
	) -> Message<'a> {
		Message {
			header,
			stream,
			version,
		}
	}

	/// Read the message body from the underlying connection
	pub fn body<T: ser::Readable>(&mut self) -> Result<T, Error> {
		read_body(&self.header, self.stream, self.version)
	}

	/// Read a single "thing" from the underlying connection.
	/// Return the thing and the total bytes read.
	pub fn streaming_read<T: ser::Readable>(&mut self) -> Result<(T, u64), Error> {
		read_item(self.stream, self.version)
	}

	pub fn copy_attachment(&mut self, len: usize, writer: &mut dyn Write) -> Result<usize, Error> {
		let mut written = 0;
		while written < len {
			let read_len = cmp::min(8000, len - written);
			let mut buf = vec![0u8; read_len];
			self.stream.read_exact(&mut buf[..])?;
			writer.write_all(&buf)?;
			written += read_len;
		}
		Ok(written)
	}
}

pub struct StopHandle {
	/// Channel to close the connection
	stopped: Arc<AtomicBool>,
	// we need Option to take ownhership of the handle in stop()
	reader_thread: Option<JoinHandle<()>>,
	writer_thread: Option<JoinHandle<()>>,
}

impl StopHandle {
	/// Schedule this connection to safely close via the async close_channel.
	pub fn stop(&self) {
		self.stopped.store(true, Ordering::Relaxed);
	}

	pub fn wait(&mut self) {
		if let Some(reader_thread) = self.reader_thread.take() {
			self.join_thread(reader_thread);
		}
		if let Some(writer_thread) = self.writer_thread.take() {
			self.join_thread(writer_thread);
		}
	}

	fn join_thread(&self, peer_thread: JoinHandle<()>) {
		// wait only if other thread is calling us, eg shutdown
		if thread::current().id() != peer_thread.thread().id() {
			debug!("waiting for thread {:?} exit", peer_thread.thread().id());
			if let Err(e) = peer_thread.join() {
				error!("failed to stop peer thread: {:?}", e);
			}
		} else {
			debug!(
				"attempt to stop thread {:?} from itself",
				peer_thread.thread().id()
			);
		}
	}
}

#[derive(Clone)]
pub struct ConnHandle {
	/// Channel to allow sending data through the connection
	pub send_channel: mpsc::SyncSender<Msg>,
}

impl ConnHandle {
	/// Send msg via the synchronous, bounded channel (sync_sender).
	/// Two possible failure cases -
	/// * Disconnected: Propagate this up to the caller so the peer connection can be closed.
	/// * Full: Our internal msg buffer is full. This is not a problem with the peer connection
	/// and we do not want to close the connection. We drop the msg rather than blocking here.
	/// If the buffer is full because there is an underlying issue with the peer
	/// and potentially the peer connection. We assume this will be handled at the peer level.
	pub fn send(&self, msg: Msg) -> Result<(), Error> {
		match self.send_channel.try_send(msg) {
			Ok(()) => Ok(()),
			Err(mpsc::TrySendError::Disconnected(_)) => {
				Err(Error::Send("try_send disconnected".to_owned()))
			}
			Err(mpsc::TrySendError::Full(_)) => {
				debug!("conn_handle: try_send but buffer is full, dropping msg");
				Ok(())
			}
		}
	}
}

pub struct Tracker {
	/// Bytes we've sent.
	pub sent_bytes: Arc<RwLock<RateCounter>>,
	/// Bytes we've received.
	pub received_bytes: Arc<RwLock<RateCounter>>,
}

impl Tracker {
	pub fn new() -> Tracker {
		let received_bytes = Arc::new(RwLock::new(RateCounter::new()));
		let sent_bytes = Arc::new(RwLock::new(RateCounter::new()));
		Tracker {
			received_bytes,
			sent_bytes,
		}
	}

	pub fn inc_received(&self, size: u64) {
		self.received_bytes.write().inc(size);
	}

	pub fn inc_sent(&self, size: u64) {
		self.sent_bytes.write().inc(size);
	}

	pub fn inc_quiet_received(&self, size: u64) {
		self.received_bytes.write().inc_quiet(size);
	}

	pub fn inc_quiet_sent(&self, size: u64) {
		self.sent_bytes.write().inc_quiet(size);
	}
}

/// Start listening on the provided connection and wraps it. Does not hang
/// the current thread, instead just returns a future and the Connection
/// itself.
pub fn listen<H>(
	stream: TcpStream,
	version: ProtocolVersion,
	tracker: Arc<Tracker>,
	handler: H,
) -> io::Result<(ConnHandle, StopHandle)>
where
	H: MessageHandler,
{
	let (send_tx, send_rx) = mpsc::sync_channel(SEND_CHANNEL_CAP);

	let stopped = Arc::new(AtomicBool::new(false));

	let conn_handle = ConnHandle {
		send_channel: send_tx,
	};

	let (reader_thread, writer_thread) = poll(
		stream,
		conn_handle.clone(),
		version,
		handler,
		send_rx,
		stopped.clone(),
		tracker,
	)?;

	Ok((
		conn_handle,
		StopHandle {
			stopped,
			reader_thread: Some(reader_thread),
			writer_thread: Some(writer_thread),
		},
	))
}

fn poll<H>(
	conn: TcpStream,
	conn_handle: ConnHandle,
	version: ProtocolVersion,
	handler: H,
	send_rx: mpsc::Receiver<Msg>,
	stopped: Arc<AtomicBool>,
	tracker: Arc<Tracker>,
) -> io::Result<(JoinHandle<()>, JoinHandle<()>)>
where
	H: MessageHandler,
{
	// Split out tcp stream out into separate reader/writer halves.
	let mut reader = conn.try_clone().expect("clone conn for reader failed");
	let mut writer = conn.try_clone().expect("clone conn for writer failed");
	let reader_stopped = stopped.clone();

	let reader_tracker = tracker.clone();
	let writer_tracker = tracker;

	let reader_thread = thread::Builder::new()
		.name("peer_read".to_string())
		.spawn(move || {
			loop {
				// check the read end
				match try_header!(read_header(&mut reader, version), &reader) {
					Some(MsgHeaderWrapper::Known(header)) => {
						reader
							.set_read_timeout(Some(BODY_IO_TIMEOUT))
							.expect("set timeout");
						let msg = Message::from_header(header, &mut reader, version);

						trace!(
							"Received message header, type {:?}, len {}.",
							msg.header.msg_type,
							msg.header.msg_len
						);

						// Increase received bytes counter
						reader_tracker.inc_received(MsgHeader::LEN as u64 + msg.header.msg_len);

						let resp_msg = try_break!(handler.consume(
							msg,
							reader_stopped.clone(),
							reader_tracker.clone()
						));
						if let Some(Some(resp_msg)) = resp_msg {
							try_break!(conn_handle.send(resp_msg));
						}
					}
					Some(MsgHeaderWrapper::Unknown(msg_len, type_byte)) => {
						debug!(
							"Received unknown message header, type {:?}, len {}.",
							type_byte, msg_len
						);
						// Increase received bytes counter
						reader_tracker.inc_received(MsgHeader::LEN as u64 + msg_len);

						try_break!(read_discard(msg_len, &mut reader));
					}
					None => {}
				}

				// check the close channel
				if reader_stopped.load(Ordering::Relaxed) {
					break;
				}
			}

			debug!(
				"Shutting down reader connection with {}",
				reader
					.peer_addr()
					.map(|a| a.to_string())
					.unwrap_or_else(|_| "?".to_owned())
			);
			let _ = reader.shutdown(Shutdown::Both);
		})?;

	let writer_thread = thread::Builder::new()
		.name("peer_write".to_string())
		.spawn(move || {
			let mut retry_send = Err(());
			writer
				.set_write_timeout(Some(BODY_IO_TIMEOUT))
				.expect("set timeout");
			loop {
				let maybe_data = retry_send.or_else(|_| send_rx.recv_timeout(CHANNEL_TIMEOUT));
				retry_send = Err(());
				match maybe_data {
					Ok(data) => {
						let written =
							try_break!(write_message(&mut writer, &data, writer_tracker.clone()));
						if written.is_none() {
							retry_send = Ok(data);
						}
					}
					Err(RecvTimeoutError::Disconnected) => {
						debug!("peer_write: mpsc channel disconnected during recv_timeout");
						break;
					}
					Err(RecvTimeoutError::Timeout) => {}
				}

				// check the close channel
				if stopped.load(Ordering::Relaxed) {
					break;
				}
			}

			debug!(
				"Shutting down writer connection with {}",
				writer
					.peer_addr()
					.map(|a| a.to_string())
					.unwrap_or_else(|_| "?".to_owned())
			);
			let _ = writer.shutdown(Shutdown::Both);
		})?;
	Ok((reader_thread, writer_thread))
}
<|endoftext|>


<commit_before>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Provides a connection wrapper that handles the lower level tasks in sending
//! or receiving data from the TCP socket, as well as dealing with timeouts.
//!
//! Because of a few idiosyncracies in the Rust `TcpStream`, this has to use
//! async I/O to be able to both read *and* write on the connection. Which
//! forces us to go through some additional gymnastic to loop over the async
//! stream and make sure we get the right number of bytes out.

use std::fs::File;
use std::io::{self, Read, Write};
use std::net::TcpStream;
use std::sync::{mpsc, Arc};
use std::{cmp, thread, time};

use crate::core::ser;
use crate::core::ser::FixedLength;
use crate::msg::{read_body, read_header, read_item, write_to_buf, MsgHeader, Type};
use crate::types::Error;
use crate::util::read_write::{read_exact, write_all};
use crate::util::{RateCounter, RwLock};

/// A trait to be implemented in order to receive messages from the
/// connection. Allows providing an optional response.
pub trait MessageHandler: Send + 'static {
	fn consume<'a>(
		&self,
		msg: Message<'a>,
		writer: &'a mut dyn Write,
		received_bytes: Arc<RwLock<RateCounter>>,
	) -> Result<Option<Response<'a>>, Error>;
}

// Macro to simplify the boilerplate around async I/O error handling,
// especially with WouldBlock kind of errors.
macro_rules! try_break {
	($chan:ident, $inner:expr) => {
		match $inner {
			Ok(v) => Some(v),
			Err(Error::Connection(ref e)) if e.kind() == io::ErrorKind::WouldBlock => None,
			Err(e) => {
				let _ = $chan.send(e);
				break;
				}
			}
	};
}

/// A message as received by the connection. Provides access to the message
/// header lazily consumes the message body, handling its deserialization.
pub struct Message<'a> {
	pub header: MsgHeader,
	stream: &'a mut dyn Read,
}

impl<'a> Message<'a> {
	fn from_header(header: MsgHeader, stream: &'a mut dyn Read) -> Message<'a> {
		Message { header, stream }
	}

	/// Read the message body from the underlying connection
	pub fn body<T: ser::Readable>(&mut self) -> Result<T, Error> {
		read_body(&self.header, self.stream)
	}

	/// Read a single "thing" from the underlying connection.
	/// Return the thing and the total bytes read.
	pub fn streaming_read<T: ser::Readable>(&mut self) -> Result<(T, u64), Error> {
		read_item(self.stream)
	}

	pub fn copy_attachment(&mut self, len: usize, writer: &mut dyn Write) -> Result<usize, Error> {
		let mut written = 0;
		while written < len {
			let read_len = cmp::min(8000, len - written);
			let mut buf = vec![0u8; read_len];
			read_exact(
				&mut self.stream,
				&mut buf[..],
				time::Duration::from_secs(10),
				true,
			)?;
			writer.write_all(&mut buf)?;
			written += read_len;
		}
		Ok(written)
	}
}

/// Response to a `Message`.
pub struct Response<'a> {
	resp_type: Type,
	body: Vec<u8>,
	stream: &'a mut dyn Write,
	attachment: Option<File>,
}

impl<'a> Response<'a> {
	pub fn new<T: ser::Writeable>(
		resp_type: Type,
		body: T,
		stream: &'a mut dyn Write,
	) -> Response<'a> {
		let body = ser::ser_vec(&body).unwrap();
		Response {
			resp_type,
			body,
			stream,
			attachment: None,
		}
	}

	fn write(mut self, sent_bytes: Arc<RwLock<RateCounter>>) -> Result<(), Error> {
		let mut msg =
			ser::ser_vec(&MsgHeader::new(self.resp_type, self.body.len() as u64)).unwrap();
		msg.append(&mut self.body);
		write_all(&mut self.stream, &msg[..], time::Duration::from_secs(10))?;
		// Increase sent bytes counter
		{
			let mut sent_bytes = sent_bytes.write();
			sent_bytes.inc(msg.len() as u64);
		}
		if let Some(mut file) = self.attachment {
			let mut buf = [0u8; 8000];
			loop {
				match file.read(&mut buf[..]) {
					Ok(0) => break,
					Ok(n) => {
						write_all(&mut self.stream, &buf[..n], time::Duration::from_secs(10))?;
						// Increase sent bytes "quietly" without incrementing the counter.
						// (In a loop here for the single attachment).
						let mut sent_bytes = sent_bytes.write();
						sent_bytes.inc_quiet(n as u64);
					}
					Err(e) => return Err(From::from(e)),
				}
			}
		}
		Ok(())
	}

	pub fn add_attachment(&mut self, file: File) {
		self.attachment = Some(file);
	}
}

pub const SEND_CHANNEL_CAP: usize = 10;

pub struct Tracker {
	/// Bytes we've sent.
	pub sent_bytes: Arc<RwLock<RateCounter>>,
	/// Bytes we've received.
	pub received_bytes: Arc<RwLock<RateCounter>>,
	/// Channel to allow sending data through the connection
	pub send_channel: mpsc::SyncSender<Vec<u8>>,
	/// Channel to close the connection
	pub close_channel: mpsc::Sender<()>,
	/// Channel to check for errors on the connection
	pub error_channel: mpsc::Receiver<Error>,
}

impl Tracker {
	pub fn send<T>(&self, body: T, msg_type: Type) -> Result<(), Error>
	where
		T: ser::Writeable,
	{
		let buf = write_to_buf(body, msg_type);
		let buf_len = buf.len();
		self.send_channel.try_send(buf)?;

		// Increase sent bytes counter
		let mut sent_bytes = self.sent_bytes.write();
		sent_bytes.inc(buf_len as u64);

		Ok(())
	}
}

/// Start listening on the provided connection and wraps it. Does not hang
/// the current thread, instead just returns a future and the Connection
/// itself.
pub fn listen<H>(stream: TcpStream, handler: H) -> Tracker
where
	H: MessageHandler,
{
	let (send_tx, send_rx) = mpsc::sync_channel(SEND_CHANNEL_CAP);
	let (close_tx, close_rx) = mpsc::channel();
	let (error_tx, error_rx) = mpsc::channel();

	// Counter of number of bytes received
	let received_bytes = Arc::new(RwLock::new(RateCounter::new()));
	// Counter of number of bytes sent
	let sent_bytes = Arc::new(RwLock::new(RateCounter::new()));

	stream
		.set_nonblocking(true)
		.expect("Non-blocking IO not available.");
	poll(
		stream,
		handler,
		send_rx,
		error_tx,
		close_rx,
		received_bytes.clone(),
		sent_bytes.clone(),
	);

	Tracker {
		sent_bytes: sent_bytes.clone(),
		received_bytes: received_bytes.clone(),
		send_channel: send_tx,
		close_channel: close_tx,
		error_channel: error_rx,
	}
}

fn poll<H>(
	conn: TcpStream,
	handler: H,
	send_rx: mpsc::Receiver<Vec<u8>>,
	error_tx: mpsc::Sender<Error>,
	close_rx: mpsc::Receiver<()>,
	received_bytes: Arc<RwLock<RateCounter>>,
	sent_bytes: Arc<RwLock<RateCounter>>,
) where
	H: MessageHandler,
{
	// Split out tcp stream out into separate reader/writer halves.
	let mut reader = conn.try_clone().expect("clone conn for reader failed");
	let mut writer = conn.try_clone().expect("clone conn for writer failed");

	let _ = thread::Builder::new()
		.name("peer".to_string())
		.spawn(move || {
			let sleep_time = time::Duration::from_millis(1);
			let mut retry_send = Err(());
			loop {
				// check the read end
				if let Some(h) = try_break!(error_tx, read_header(&mut reader, None)) {
					let msg = Message::from_header(h, &mut reader);

					trace!(
						"Received message header, type {:?}, len {}.",
						msg.header.msg_type,
						msg.header.msg_len
					);

					// Increase received bytes counter
					let received = received_bytes.clone();
					{
						let mut received_bytes = received_bytes.write();
						received_bytes.inc(MsgHeader::LEN as u64 + msg.header.msg_len);
					}

					if let Some(Some(resp)) =
						try_break!(error_tx, handler.consume(msg, &mut writer, received))
					{
						try_break!(error_tx, resp.write(sent_bytes.clone()));
					}
				}

				// check the write end, use or_else so try_recv is lazily eval'd
				let maybe_data = retry_send.or_else(|_| send_rx.try_recv());
				retry_send = Err(());
				if let Ok(data) = maybe_data {
					let written =
						try_break!(error_tx, writer.write_all(&data[..]).map_err(&From::from));
					if written.is_none() {
						retry_send = Ok(data);
					}
				}

				// check the close channel
				if let Ok(_) = close_rx.try_recv() {
					debug!(
						"Connection close with {} initiated by us",
						conn.peer_addr()
							.map(|a| a.to_string())
							.unwrap_or("?".to_owned())
					);
					break;
				}

				thread::sleep(sleep_time);
			}
		});
}
<commit_msg>fix: shutdown the tcpstream when we initiate to close a peer connection (#2288)<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Provides a connection wrapper that handles the lower level tasks in sending
//! or receiving data from the TCP socket, as well as dealing with timeouts.
//!
//! Because of a few idiosyncracies in the Rust `TcpStream`, this has to use
//! async I/O to be able to both read *and* write on the connection. Which
//! forces us to go through some additional gymnastic to loop over the async
//! stream and make sure we get the right number of bytes out.

use std::fs::File;
use std::io::{self, Read, Write};
use std::net::{Shutdown, TcpStream};
use std::sync::{mpsc, Arc};
use std::{cmp, thread, time};

use crate::core::ser;
use crate::core::ser::FixedLength;
use crate::msg::{read_body, read_header, read_item, write_to_buf, MsgHeader, Type};
use crate::types::Error;
use crate::util::read_write::{read_exact, write_all};
use crate::util::{RateCounter, RwLock};

/// A trait to be implemented in order to receive messages from the
/// connection. Allows providing an optional response.
pub trait MessageHandler: Send + 'static {
	fn consume<'a>(
		&self,
		msg: Message<'a>,
		writer: &'a mut dyn Write,
		received_bytes: Arc<RwLock<RateCounter>>,
	) -> Result<Option<Response<'a>>, Error>;
}

// Macro to simplify the boilerplate around async I/O error handling,
// especially with WouldBlock kind of errors.
macro_rules! try_break {
	($chan:ident, $inner:expr) => {
		match $inner {
			Ok(v) => Some(v),
			Err(Error::Connection(ref e)) if e.kind() == io::ErrorKind::WouldBlock => None,
			Err(e) => {
				let _ = $chan.send(e);
				break;
				}
			}
	};
}

/// A message as received by the connection. Provides access to the message
/// header lazily consumes the message body, handling its deserialization.
pub struct Message<'a> {
	pub header: MsgHeader,
	stream: &'a mut dyn Read,
}

impl<'a> Message<'a> {
	fn from_header(header: MsgHeader, stream: &'a mut dyn Read) -> Message<'a> {
		Message { header, stream }
	}

	/// Read the message body from the underlying connection
	pub fn body<T: ser::Readable>(&mut self) -> Result<T, Error> {
		read_body(&self.header, self.stream)
	}

	/// Read a single "thing" from the underlying connection.
	/// Return the thing and the total bytes read.
	pub fn streaming_read<T: ser::Readable>(&mut self) -> Result<(T, u64), Error> {
		read_item(self.stream)
	}

	pub fn copy_attachment(&mut self, len: usize, writer: &mut dyn Write) -> Result<usize, Error> {
		let mut written = 0;
		while written < len {
			let read_len = cmp::min(8000, len - written);
			let mut buf = vec![0u8; read_len];
			read_exact(
				&mut self.stream,
				&mut buf[..],
				time::Duration::from_secs(10),
				true,
			)?;
			writer.write_all(&mut buf)?;
			written += read_len;
		}
		Ok(written)
	}
}

/// Response to a `Message`.
pub struct Response<'a> {
	resp_type: Type,
	body: Vec<u8>,
	stream: &'a mut dyn Write,
	attachment: Option<File>,
}

impl<'a> Response<'a> {
	pub fn new<T: ser::Writeable>(
		resp_type: Type,
		body: T,
		stream: &'a mut dyn Write,
	) -> Response<'a> {
		let body = ser::ser_vec(&body).unwrap();
		Response {
			resp_type,
			body,
			stream,
			attachment: None,
		}
	}

	fn write(mut self, sent_bytes: Arc<RwLock<RateCounter>>) -> Result<(), Error> {
		let mut msg =
			ser::ser_vec(&MsgHeader::new(self.resp_type, self.body.len() as u64)).unwrap();
		msg.append(&mut self.body);
		write_all(&mut self.stream, &msg[..], time::Duration::from_secs(10))?;
		// Increase sent bytes counter
		{
			let mut sent_bytes = sent_bytes.write();
			sent_bytes.inc(msg.len() as u64);
		}
		if let Some(mut file) = self.attachment {
			let mut buf = [0u8; 8000];
			loop {
				match file.read(&mut buf[..]) {
					Ok(0) => break,
					Ok(n) => {
						write_all(&mut self.stream, &buf[..n], time::Duration::from_secs(10))?;
						// Increase sent bytes "quietly" without incrementing the counter.
						// (In a loop here for the single attachment).
						let mut sent_bytes = sent_bytes.write();
						sent_bytes.inc_quiet(n as u64);
					}
					Err(e) => return Err(From::from(e)),
				}
			}
		}
		Ok(())
	}

	pub fn add_attachment(&mut self, file: File) {
		self.attachment = Some(file);
	}
}

pub const SEND_CHANNEL_CAP: usize = 10;

pub struct Tracker {
	/// Bytes we've sent.
	pub sent_bytes: Arc<RwLock<RateCounter>>,
	/// Bytes we've received.
	pub received_bytes: Arc<RwLock<RateCounter>>,
	/// Channel to allow sending data through the connection
	pub send_channel: mpsc::SyncSender<Vec<u8>>,
	/// Channel to close the connection
	pub close_channel: mpsc::Sender<()>,
	/// Channel to check for errors on the connection
	pub error_channel: mpsc::Receiver<Error>,
}

impl Tracker {
	pub fn send<T>(&self, body: T, msg_type: Type) -> Result<(), Error>
	where
		T: ser::Writeable,
	{
		let buf = write_to_buf(body, msg_type);
		let buf_len = buf.len();
		self.send_channel.try_send(buf)?;

		// Increase sent bytes counter
		let mut sent_bytes = self.sent_bytes.write();
		sent_bytes.inc(buf_len as u64);

		Ok(())
	}
}

/// Start listening on the provided connection and wraps it. Does not hang
/// the current thread, instead just returns a future and the Connection
/// itself.
pub fn listen<H>(stream: TcpStream, handler: H) -> Tracker
where
	H: MessageHandler,
{
	let (send_tx, send_rx) = mpsc::sync_channel(SEND_CHANNEL_CAP);
	let (close_tx, close_rx) = mpsc::channel();
	let (error_tx, error_rx) = mpsc::channel();

	// Counter of number of bytes received
	let received_bytes = Arc::new(RwLock::new(RateCounter::new()));
	// Counter of number of bytes sent
	let sent_bytes = Arc::new(RwLock::new(RateCounter::new()));

	stream
		.set_nonblocking(true)
		.expect("Non-blocking IO not available.");
	poll(
		stream,
		handler,
		send_rx,
		error_tx,
		close_rx,
		received_bytes.clone(),
		sent_bytes.clone(),
	);

	Tracker {
		sent_bytes: sent_bytes.clone(),
		received_bytes: received_bytes.clone(),
		send_channel: send_tx,
		close_channel: close_tx,
		error_channel: error_rx,
	}
}

fn poll<H>(
	conn: TcpStream,
	handler: H,
	send_rx: mpsc::Receiver<Vec<u8>>,
	error_tx: mpsc::Sender<Error>,
	close_rx: mpsc::Receiver<()>,
	received_bytes: Arc<RwLock<RateCounter>>,
	sent_bytes: Arc<RwLock<RateCounter>>,
) where
	H: MessageHandler,
{
	// Split out tcp stream out into separate reader/writer halves.
	let mut reader = conn.try_clone().expect("clone conn for reader failed");
	let mut writer = conn.try_clone().expect("clone conn for writer failed");

	let _ = thread::Builder::new()
		.name("peer".to_string())
		.spawn(move || {
			let sleep_time = time::Duration::from_millis(1);
			let mut retry_send = Err(());
			loop {
				// check the read end
				if let Some(h) = try_break!(error_tx, read_header(&mut reader, None)) {
					let msg = Message::from_header(h, &mut reader);

					trace!(
						"Received message header, type {:?}, len {}.",
						msg.header.msg_type,
						msg.header.msg_len
					);

					// Increase received bytes counter
					let received = received_bytes.clone();
					{
						let mut received_bytes = received_bytes.write();
						received_bytes.inc(MsgHeader::LEN as u64 + msg.header.msg_len);
					}

					if let Some(Some(resp)) =
						try_break!(error_tx, handler.consume(msg, &mut writer, received))
					{
						try_break!(error_tx, resp.write(sent_bytes.clone()));
					}
				}

				// check the write end, use or_else so try_recv is lazily eval'd
				let maybe_data = retry_send.or_else(|_| send_rx.try_recv());
				retry_send = Err(());
				if let Ok(data) = maybe_data {
					let written =
						try_break!(error_tx, writer.write_all(&data[..]).map_err(&From::from));
					if written.is_none() {
						retry_send = Ok(data);
					}
				}

				// check the close channel
				if let Ok(_) = close_rx.try_recv() {
					debug!(
						"Connection close with {} initiated by us",
						conn.peer_addr()
							.map(|a| a.to_string())
							.unwrap_or("?".to_owned())
					);
					let _ = conn.shutdown(Shutdown::Both);
					break;
				}

				thread::sleep(sleep_time);
			}
		});
}
<|endoftext|>


<commit_before>use std::option::Option;

/// Represents an inline keyboard that appears right next to the message it belongs to.
#[derive(Serialize, Debug)]
pub struct InlineKeyboardMarkup {
    pub inline_keyboard: Vec<Vec<InlineKeyboardButton>>,
}

/// Represents one button of an `InlineKeyboardMarkup` You must use exactly one of the optional fields.
#[derive(Serialize, Debug)]
pub struct InlineKeyboardButton {
    text: String,
    url: Option<String>,
    callback_data: Option<String>,
    switch_inline_query: Option<String>,
    switch_inline_query_current_chat: Option<String>,
}

impl InlineKeyboardMarkup {
    /// Constructs a new `InlineKeyboardMarkup`.
    pub fn new(inline_keyboard: Vec<Vec<InlineKeyboardButton>>) -> InlineKeyboardMarkup {
        InlineKeyboardMarkup { inline_keyboard: inline_keyboard }
    }
}

impl InlineKeyboardButton {
    /// Constructs a new `InlineKeyboardButton`.
    pub fn new(text: String,
               url: Option<String>,
               callback_data: Option<String>,
               switch_inline_query: Option<String>,
               switch_inline_query_current_chat: Option<String>)
               -> InlineKeyboardButton {
        InlineKeyboardButton {
            text: text,
            url: Some(url.unwrap_or("".to_string())),
            callback_data: Some(callback_data.unwrap_or("".to_string())),
            switch_inline_query: Some(switch_inline_query.unwrap_or("".to_string())),
            switch_inline_query_current_chat: Some(switch_inline_query_current_chat
                                                       .unwrap_or("".to_string())),
        }
    }
}
<commit_msg>added missing ','<commit_after>use std::option::Option;

/// Represents an inline keyboard that appears right next to the message it belongs to.
#[derive(Serialize, Debug)]
pub struct InlineKeyboardMarkup {
    pub inline_keyboard: Vec<Vec<InlineKeyboardButton>>,
}

/// Represents one button of an `InlineKeyboardMarkup`, you must use exactly one of the optional fields.
#[derive(Serialize, Debug)]
pub struct InlineKeyboardButton {
    text: String,
    url: Option<String>,
    callback_data: Option<String>,
    switch_inline_query: Option<String>,
    switch_inline_query_current_chat: Option<String>,
}

impl InlineKeyboardMarkup {
    /// Constructs a new `InlineKeyboardMarkup`.
    pub fn new(inline_keyboard: Vec<Vec<InlineKeyboardButton>>) -> InlineKeyboardMarkup {
        InlineKeyboardMarkup { inline_keyboard: inline_keyboard }
    }
}

impl InlineKeyboardButton {
    /// Constructs a new `InlineKeyboardButton`.
    pub fn new(text: String,
               url: Option<String>,
               callback_data: Option<String>,
               switch_inline_query: Option<String>,
               switch_inline_query_current_chat: Option<String>)
               -> InlineKeyboardButton {
        InlineKeyboardButton {
            text: text,
            url: Some(url.unwrap_or("".to_string())),
            callback_data: Some(callback_data.unwrap_or("".to_string())),
            switch_inline_query: Some(switch_inline_query.unwrap_or("".to_string())),
            switch_inline_query_current_chat: Some(switch_inline_query_current_chat
                                                       .unwrap_or("".to_string())),
        }
    }
}
<|endoftext|>


<commit_before>use std::mem;
use std::marker;

pub trait MULX {
    fn mulx(self, Self) -> (Self, Self) where Self: marker::Sized;
}

impl MULX for u8 {
    fn mulx(self, y: u8) -> (u8, u8) {
        let result: u16 = (self as u16) * (y as u16);
        let hi = (result >> 8) as u8;
        (result as u8, hi)
    }
}

impl MULX for u16 {
    fn mulx(self, y: u16) -> (u16, u16) {
        let result: u32 = (self as u32) * (y as u32);
        let hi = (result >> 16) as u16;
        (result as u16, hi)
    }
}

impl MULX for u32 {
    fn mulx(self, y: u32) -> (u32, u32) {
        let result: u64 = (self as u64) * (y as u64);
        let hi = (result >> 32) as u32;
        (result as u32, hi)
    }
}

impl MULX for u64 {
    fn mulx(self, y: u64) -> (u64, u64) {
        let result: u128 = (self as u128) * (y as u128);
        let hi = (result >> 64) as u64;
        (result as u64, hi)
    }
}

impl MULX for usize {
    fn mulx(self, y: usize) -> (usize, usize) {
        match mem::size_of::<usize>() * 8 {
            8 => { let tmp = (self as u8).mulx(y as u8); (tmp.0 as usize, tmp.1 as usize) },
            16 => { let tmp = (self as u16).mulx(y as u16); (tmp.0 as usize, tmp.1 as usize) },
            32 => { let tmp = (self as u32).mulx(y as u32); (tmp.0 as usize, tmp.1 as usize) },
            64 => { let tmp = (self as u64).mulx(y as u64); (tmp.0 as usize, tmp.1 as usize) },
            _ => unreachable!(),
        }
    }
}

impl MULX for i8 {
    fn mulx(self, y: i8) -> (i8, i8) {
        unsafe { mem::transmute((self as u8).mulx(y as u8)) }
    }
}

impl MULX for i16 {
    fn mulx(self, y: i16) -> (i16, i16) {
        unsafe { mem::transmute((self as u16).mulx(y as u16)) }
    }
}

impl MULX for i32 {
    fn mulx(self, y: i32) -> (i32, i32) {
        unsafe { mem::transmute((self as u32).mulx(y as u32)) }
    }
}

impl MULX for i64 {
    fn mulx(self, y: i64) -> (i64, i64) {
        unsafe { mem::transmute((self as u64).mulx(y as u64)) }
    }
}

impl MULX for isize {
    fn mulx(self, y: isize) -> (isize, isize) {
        unsafe { mem::transmute((self as usize).mulx(y as usize)) }
    }
}

/// Unsigned multiply without affecting flags.
///
/// Unsigned multiplication of `x` with `y` returning a pair `(lo, hi)` with
/// the low half and the high half of the result.
///
/// # Intrinsics (when available BMI2)
///
/// [`MULX`](http://www.felixcloutier.com/x86/MULX.html): Unsigned Multiply
/// Without Affecting Flags (supports 32/64 bit registers).
///
/// # Examples
///
/// ```
/// use bitintr::bmi2::mulx;
///
/// { // 8-bit
///   let a: u8 = 128;
///   let b: u8 = 128;
///   let (lo, hi): (u8, u8) = mulx(a, b);
///   // result = 16384 = 0b0100_0000_0000_0000u16
///   //                    ^~hi~~~~~ ^~lo~~~~~
///   assert_eq!(lo, 0b0000_0000);
///   assert_eq!(hi, 0b0100_0000);
/// }
/// { // 16-bit
///   let a: u16 = 65_500;
///   let b: u16 = 65_500;
///   let (lo, hi): (u16, u16)  = mulx(a, b);
///   // result = 4290250000 = 0b1111_1111_1011_1000_0000_0101_0001_0000u32
///   //                         ^~hi~~~~~~~~~~~~~~~ ^~lo~~~~~~~~~~~~~~~
///   assert_eq!(lo, 0b0000_0101_0001_0000);
///   assert_eq!(hi, 0b1111_1111_1011_1000);
/// }
/// { // 32-bit
///   let a: u32 = 4_294_967_200;
///   let b: u32 = 2;
///   let (lo, hi): (u32, u32)  = mulx(a, b);
///   // result = 8589934400
///   //        = 0b0001_1111_1111_1111_1111_1111_1111_0100_0000u64
///   //            ^~hi ^~lo~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
///   assert_eq!(lo, 0b1111_1111_1111_1111_1111_1111_0100_0000u32);
///   assert_eq!(hi, 0b0001u32);
/// }
/// { // 64-bit
///   let a: u64 = 9_223_372_036_854_775_800;
///   let b: u64 = 100;
///   let (lo, hi): (u64, u64)  = mulx(a, b);
///   // result = 922337203685477580000
///   //        = 0b00110001_11111111_11111111_11111111_11111111_11111111_11111111_11111100_11100000u128
///   //            ^~hi~~~~ ^~lo~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
///   assert_eq!(lo, 0b11111111_11111111_11111111_11111111_11111111_11111111_11111100_11100000u64);
///   assert_eq!(hi, 0b00110001u64);
/// }

/// ```
pub fn mulx<T: MULX>(x: T, y: T) -> (T, T) {
    x.mulx(y)
}
<commit_msg>add some signed integer examples to mulx<commit_after>use std::mem;
use std::marker;

pub trait MULX {
    fn mulx(self, Self) -> (Self, Self) where Self: marker::Sized;
}

impl MULX for u8 {
    fn mulx(self, y: u8) -> (u8, u8) {
        let result: u16 = (self as u16) * (y as u16);
        let hi = (result >> 8) as u8;
        (result as u8, hi)
    }
}

impl MULX for u16 {
    fn mulx(self, y: u16) -> (u16, u16) {
        let result: u32 = (self as u32) * (y as u32);
        let hi = (result >> 16) as u16;
        (result as u16, hi)
    }
}

impl MULX for u32 {
    fn mulx(self, y: u32) -> (u32, u32) {
        let result: u64 = (self as u64) * (y as u64);
        let hi = (result >> 32) as u32;
        (result as u32, hi)
    }
}

impl MULX for u64 {
    fn mulx(self, y: u64) -> (u64, u64) {
        let result: u128 = (self as u128) * (y as u128);
        let hi = (result >> 64) as u64;
        (result as u64, hi)
    }
}

impl MULX for usize {
    fn mulx(self, y: usize) -> (usize, usize) {
        match mem::size_of::<usize>() * 8 {
            8 => { let tmp = (self as u8).mulx(y as u8); (tmp.0 as usize, tmp.1 as usize) },
            16 => { let tmp = (self as u16).mulx(y as u16); (tmp.0 as usize, tmp.1 as usize) },
            32 => { let tmp = (self as u32).mulx(y as u32); (tmp.0 as usize, tmp.1 as usize) },
            64 => { let tmp = (self as u64).mulx(y as u64); (tmp.0 as usize, tmp.1 as usize) },
            _ => unreachable!(),
        }
    }
}

impl MULX for i8 {
    fn mulx(self, y: i8) -> (i8, i8) {
        unsafe { mem::transmute((self as u8).mulx(y as u8)) }
    }
}

impl MULX for i16 {
    fn mulx(self, y: i16) -> (i16, i16) {
        unsafe { mem::transmute((self as u16).mulx(y as u16)) }
    }
}

impl MULX for i32 {
    fn mulx(self, y: i32) -> (i32, i32) {
        unsafe { mem::transmute((self as u32).mulx(y as u32)) }
    }
}

impl MULX for i64 {
    fn mulx(self, y: i64) -> (i64, i64) {
        unsafe { mem::transmute((self as u64).mulx(y as u64)) }
    }
}

impl MULX for isize {
    fn mulx(self, y: isize) -> (isize, isize) {
        unsafe { mem::transmute((self as usize).mulx(y as usize)) }
    }
}

/// Unsigned multiply without affecting flags.
///
/// Unsigned multiplication of `x` with `y` returning a pair `(lo, hi)` with
/// the low half and the high half of the result.
///
/// # Intrinsics (when available BMI2)
///
/// [`MULX`](http://www.felixcloutier.com/x86/MULX.html): Unsigned Multiply
/// Without Affecting Flags (supports 32/64 bit registers).
///
/// # Examples
///
/// ```
/// use bitintr::bmi2::mulx;
///
/// { // 8-bit
///   let a: u8 = 128;
///   let b: u8 = 128;
///   let (lo, hi): (u8, u8) = mulx(a, b);
///   // result = 16384 = 0b0100_0000_0000_0000u16
///   //                    ^~hi~~~~~ ^~lo~~~~~
///   assert_eq!(lo, 0b0000_0000);
///   assert_eq!(hi, 0b0100_0000);
/// }
/// { // 16-bit
///   let a: u16 = 65_500;
///   let b: u16 = 65_500;
///   let (lo, hi): (u16, u16)  = mulx(a, b);
///   // result = 4290250000 = 0b1111_1111_1011_1000_0000_0101_0001_0000u32
///   //                         ^~hi~~~~~~~~~~~~~~~ ^~lo~~~~~~~~~~~~~~~
///   assert_eq!(lo, 0b0000_0101_0001_0000);
///   assert_eq!(hi, 0b1111_1111_1011_1000);
/// }
/// { // 32-bit
///   let a: u32 = 4_294_967_200;
///   let b: u32 = 2;
///   let (lo, hi): (u32, u32)  = mulx(a, b);
///   // result = 8589934400
///   //        = 0b0001_1111_1111_1111_1111_1111_1111_0100_0000u64
///   //            ^~hi ^~lo~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
///   assert_eq!(lo, 0b1111_1111_1111_1111_1111_1111_0100_0000u32);
///   assert_eq!(hi, 0b0001u32);
/// }
/// { // 64-bit
///   let a: u64 = 9_223_372_036_854_775_800;
///   let b: u64 = 100;
///   let (lo, hi): (u64, u64)  = mulx(a, b);
///   // result = 922337203685477580000
///   //        = 0b00110001_11111111_11111111_11111111_11111111_11111111_11111111_11111100_11100000u128
///   //            ^~hi~~~~ ^~lo~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
///   assert_eq!(lo, 0b11111111_11111111_11111111_11111111_11111111_11111111_11111100_11100000u64);
///   assert_eq!(hi, 0b00110001u64);
/// }
/// { // 8-bit
///   let a: i8 = 128;
///   let b: i8 = 128;
///   let (lo, hi): (i8, i8) = mulx(a, b);
///   // result = _____ = 0b0100_0000_0000_0000u16
///   //                    ^~hi~~~~~ ^~lo~~~~~
///   assert_eq!(lo, 0b0000_0000);
///   assert_eq!(hi, 0b0100_0000);
/// }
/// { // 16-bit
///   let a: i16 = 65_500;
///   let b: i16 = 65_500;
///   let (lo, hi): (i16, i16)  = mulx(a, b);
///   // result = 4290250000 = 0b1111_1111_1011_1000_0000_0101_0001_0000u32
///   //                         ^~hi~~~~~~~~~~~~~~~ ^~lo~~~~~~~~~~~~~~~
///   assert_eq!(lo, 0b0000_0101_0001_0000);
///   assert_eq!(hi, 0b1111_1111_1011_1000);
/// }
/// ```
pub fn mulx<T: MULX>(x: T, y: T) -> (T, T) {
    x.mulx(y)
}
<|endoftext|>


<commit_before>use crate::api::internal::InterConfig;
use crate::config::EncoderConfig;
use crate::context::{
  BlockOffset, FrameBlocks, SuperBlockOffset, TileBlockOffset,
  TileSuperBlockOffset,
};
use crate::cpu_features::CpuFeatureLevel;
use crate::dist::get_satd;
use crate::encoder::{
  build_coarse_pmvs, build_full_res_pmvs, build_half_res_pmvs,
  FrameInvariants, FrameState, Sequence, IMPORTANCE_BLOCK_SIZE,
};
use crate::frame::{AsRegion, PlaneOffset};
use crate::hawktracer::*;
use crate::partition::{get_intra_edges, BlockSize};
use crate::predict::{IntraParam, PredictionMode};
use crate::rayon::iter::*;
use crate::tiling::{Area, TileRect};
use crate::transform::TxSize;
use crate::{Frame, Pixel};
use std::sync::Arc;

pub(crate) const IMP_BLOCK_MV_UNITS_PER_PIXEL: i64 = 8;
pub(crate) const IMP_BLOCK_SIZE_IN_MV_UNITS: i64 =
  IMPORTANCE_BLOCK_SIZE as i64 * IMP_BLOCK_MV_UNITS_PER_PIXEL;
pub(crate) const IMP_BLOCK_AREA_IN_MV_UNITS: i64 =
  IMP_BLOCK_SIZE_IN_MV_UNITS * IMP_BLOCK_SIZE_IN_MV_UNITS;

pub(crate) fn estimate_intra_costs<T: Pixel>(
  frame: &Frame<T>, bit_depth: usize, cpu_feature_level: CpuFeatureLevel,
) -> Box<[u32]> {
  let plane = &frame.planes[0];
  let mut plane_after_prediction = frame.planes[0].clone();

  let bsize = BlockSize::from_width_and_height(
    IMPORTANCE_BLOCK_SIZE,
    IMPORTANCE_BLOCK_SIZE,
  );
  let tx_size = bsize.tx_size();

  let h_in_imp_b = plane.cfg.height / IMPORTANCE_BLOCK_SIZE;
  let w_in_imp_b = plane.cfg.width / IMPORTANCE_BLOCK_SIZE;
  let mut intra_costs = Vec::with_capacity(h_in_imp_b * w_in_imp_b);

  for y in 0..h_in_imp_b {
    for x in 0..w_in_imp_b {
      let plane_org = plane.region(Area::Rect {
        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      // TODO: other intra prediction modes.
      let edge_buf = get_intra_edges(
        &plane.as_region(),
        TileBlockOffset(BlockOffset { x, y }),
        0,
        0,
        bsize,
        PlaneOffset {
          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        },
        TxSize::TX_8X8,
        bit_depth,
        Some(PredictionMode::DC_PRED),
        false,
        IntraParam::None,
      );

      let mut plane_after_prediction_region = plane_after_prediction
        .region_mut(Area::Rect {
          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
          width: IMPORTANCE_BLOCK_SIZE,
          height: IMPORTANCE_BLOCK_SIZE,
        });

      PredictionMode::DC_PRED.predict_intra(
        TileRect {
          x: x * IMPORTANCE_BLOCK_SIZE,
          y: y * IMPORTANCE_BLOCK_SIZE,
          width: IMPORTANCE_BLOCK_SIZE,
          height: IMPORTANCE_BLOCK_SIZE,
        },
        &mut plane_after_prediction_region,
        tx_size,
        bit_depth,
        &[], // Not used by DC_PRED
        IntraParam::None,
        None, // Not used by DC_PRED
        &edge_buf,
        cpu_feature_level,
      );

      let plane_after_prediction_region =
        plane_after_prediction.region(Area::Rect {
          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
          width: IMPORTANCE_BLOCK_SIZE,
          height: IMPORTANCE_BLOCK_SIZE,
        });

      let intra_cost = get_satd(
        &plane_org,
        &plane_after_prediction_region,
        bsize,
        bit_depth,
        cpu_feature_level,
      );

      intra_costs.push(intra_cost);
    }
  }

  intra_costs.into_boxed_slice()
}

pub(crate) fn estimate_inter_costs<T: Pixel>(
  frame: Arc<Frame<T>>, ref_frame: Arc<Frame<T>>, bit_depth: usize,
  mut config: EncoderConfig, sequence: Sequence,
) -> Box<[u32]> {
  config.low_latency = true;
  config.speed_settings.multiref = false;
  let inter_cfg = InterConfig::new(&config);
  let last_fi = FrameInvariants::new_key_frame(config, sequence, 0);
  let mut fi =
    FrameInvariants::new_inter_frame(&last_fi, &inter_cfg, 0, 1, 2, false);

  // Compute the motion vectors.
  let mut fs = FrameState::new_with_frame(&fi, frame.clone());
  compute_motion_vectors(&mut fi, &mut fs, &inter_cfg);

  // Estimate inter costs
  let plane_org = &frame.planes[0];
  let plane_ref = &ref_frame.planes[0];
  let h_in_imp_b = plane_org.cfg.height / IMPORTANCE_BLOCK_SIZE;
  let w_in_imp_b = plane_org.cfg.width / IMPORTANCE_BLOCK_SIZE;
  let mut inter_costs = Vec::with_capacity(h_in_imp_b * w_in_imp_b);
  let mvs = &fs.frame_mvs[0];
  let bsize = BlockSize::from_width_and_height(
    IMPORTANCE_BLOCK_SIZE,
    IMPORTANCE_BLOCK_SIZE,
  );
  (0..h_in_imp_b).for_each(|y| {
    (0..w_in_imp_b).for_each(|x| {
      let mv = mvs[y * 2][x * 2];

      // Coordinates of the top-left corner of the reference block, in MV
      // units.
      let reference_x = x as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.col as i64;
      let reference_y = y as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.row as i64;

      let region_org = plane_org.region(Area::Rect {
        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      let region_ref = plane_ref.region(Area::Rect {
        x: reference_x as isize / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
        y: reference_y as isize / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      inter_costs.push(get_satd(
        &region_org,
        &region_ref,
        bsize,
        bit_depth,
        fi.cpu_feature_level,
      ));
    });
  });
  inter_costs.into_boxed_slice()
}

#[hawktracer(compute_motion_vectors)]
pub(crate) fn compute_motion_vectors<T: Pixel>(
  fi: &mut FrameInvariants<T>, fs: &mut FrameState<T>, inter_cfg: &InterConfig,
) {
  let mut blocks = FrameBlocks::new(fi.w_in_b, fi.h_in_b);
  fs.half_res_pmvs = fi
    .tiling
    .tile_iter_mut(fs, &mut blocks)
    .collect::<Vec<_>>()
    .into_par_iter()
    .map(|mut ctx| {
      let ts = &mut ctx.ts;

      // Compute the quarter-resolution motion vectors.
      let tile_pmvs = build_coarse_pmvs(fi, ts, inter_cfg);

      // Compute the half-resolution motion vectors.
      let mut half_res_pmvs = Vec::with_capacity(ts.sb_height * ts.sb_width);
      for sby in 0..ts.sb_height {
        for sbx in 0..ts.sb_width {
          let tile_sbo =
            TileSuperBlockOffset(SuperBlockOffset { x: sbx, y: sby });
          half_res_pmvs
            .push(build_half_res_pmvs(fi, ts, tile_sbo, &tile_pmvs));
        }
      }

      // Compute the full-resolution motion vectors.
      for sby in 0..ts.sb_height {
        for sbx in 0..ts.sb_width {
          let tile_sbo =
            TileSuperBlockOffset(SuperBlockOffset { x: sbx, y: sby });
          build_full_res_pmvs(fi, ts, tile_sbo, &half_res_pmvs);
        }
      }

      (ts.sbo, half_res_pmvs)
    })
    .collect::<Vec<_>>();
}
<commit_msg>Refactor compute_motion_vectors_per_tile<commit_after>use crate::api::internal::InterConfig;
use crate::config::EncoderConfig;
use crate::context::{
  BlockOffset, FrameBlocks, PlaneSuperBlockOffset, SuperBlockOffset,
  TileBlockOffset, TileSuperBlockOffset,
};
use crate::cpu_features::CpuFeatureLevel;
use crate::dist::get_satd;
use crate::encoder::BlockPmv;
use crate::encoder::{
  build_coarse_pmvs, build_full_res_pmvs, build_half_res_pmvs,
  FrameInvariants, FrameState, Sequence, IMPORTANCE_BLOCK_SIZE,
};
use crate::frame::{AsRegion, PlaneOffset};
use crate::hawktracer::*;
use crate::partition::{get_intra_edges, BlockSize};
use crate::predict::{IntraParam, PredictionMode};
use crate::rayon::iter::*;
use crate::tiling::{Area, TileRect, TileStateMut};
use crate::transform::TxSize;
use crate::{Frame, Pixel};
use std::sync::Arc;

pub(crate) const IMP_BLOCK_MV_UNITS_PER_PIXEL: i64 = 8;
pub(crate) const IMP_BLOCK_SIZE_IN_MV_UNITS: i64 =
  IMPORTANCE_BLOCK_SIZE as i64 * IMP_BLOCK_MV_UNITS_PER_PIXEL;
pub(crate) const IMP_BLOCK_AREA_IN_MV_UNITS: i64 =
  IMP_BLOCK_SIZE_IN_MV_UNITS * IMP_BLOCK_SIZE_IN_MV_UNITS;

pub(crate) fn estimate_intra_costs<T: Pixel>(
  frame: &Frame<T>, bit_depth: usize, cpu_feature_level: CpuFeatureLevel,
) -> Box<[u32]> {
  let plane = &frame.planes[0];
  let mut plane_after_prediction = frame.planes[0].clone();

  let bsize = BlockSize::from_width_and_height(
    IMPORTANCE_BLOCK_SIZE,
    IMPORTANCE_BLOCK_SIZE,
  );
  let tx_size = bsize.tx_size();

  let h_in_imp_b = plane.cfg.height / IMPORTANCE_BLOCK_SIZE;
  let w_in_imp_b = plane.cfg.width / IMPORTANCE_BLOCK_SIZE;
  let mut intra_costs = Vec::with_capacity(h_in_imp_b * w_in_imp_b);

  for y in 0..h_in_imp_b {
    for x in 0..w_in_imp_b {
      let plane_org = plane.region(Area::Rect {
        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      // TODO: other intra prediction modes.
      let edge_buf = get_intra_edges(
        &plane.as_region(),
        TileBlockOffset(BlockOffset { x, y }),
        0,
        0,
        bsize,
        PlaneOffset {
          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        },
        TxSize::TX_8X8,
        bit_depth,
        Some(PredictionMode::DC_PRED),
        false,
        IntraParam::None,
      );

      let mut plane_after_prediction_region = plane_after_prediction
        .region_mut(Area::Rect {
          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
          width: IMPORTANCE_BLOCK_SIZE,
          height: IMPORTANCE_BLOCK_SIZE,
        });

      PredictionMode::DC_PRED.predict_intra(
        TileRect {
          x: x * IMPORTANCE_BLOCK_SIZE,
          y: y * IMPORTANCE_BLOCK_SIZE,
          width: IMPORTANCE_BLOCK_SIZE,
          height: IMPORTANCE_BLOCK_SIZE,
        },
        &mut plane_after_prediction_region,
        tx_size,
        bit_depth,
        &[], // Not used by DC_PRED
        IntraParam::None,
        None, // Not used by DC_PRED
        &edge_buf,
        cpu_feature_level,
      );

      let plane_after_prediction_region =
        plane_after_prediction.region(Area::Rect {
          x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
          y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
          width: IMPORTANCE_BLOCK_SIZE,
          height: IMPORTANCE_BLOCK_SIZE,
        });

      let intra_cost = get_satd(
        &plane_org,
        &plane_after_prediction_region,
        bsize,
        bit_depth,
        cpu_feature_level,
      );

      intra_costs.push(intra_cost);
    }
  }

  intra_costs.into_boxed_slice()
}

pub(crate) fn estimate_inter_costs<T: Pixel>(
  frame: Arc<Frame<T>>, ref_frame: Arc<Frame<T>>, bit_depth: usize,
  mut config: EncoderConfig, sequence: Sequence,
) -> Box<[u32]> {
  config.low_latency = true;
  config.speed_settings.multiref = false;
  let inter_cfg = InterConfig::new(&config);
  let last_fi = FrameInvariants::new_key_frame(config, sequence, 0);
  let mut fi =
    FrameInvariants::new_inter_frame(&last_fi, &inter_cfg, 0, 1, 2, false);

  // Compute the motion vectors.
  let mut fs = FrameState::new_with_frame(&fi, frame.clone());
  compute_motion_vectors(&mut fi, &mut fs, &inter_cfg);

  // Estimate inter costs
  let plane_org = &frame.planes[0];
  let plane_ref = &ref_frame.planes[0];
  let h_in_imp_b = plane_org.cfg.height / IMPORTANCE_BLOCK_SIZE;
  let w_in_imp_b = plane_org.cfg.width / IMPORTANCE_BLOCK_SIZE;
  let mut inter_costs = Vec::with_capacity(h_in_imp_b * w_in_imp_b);
  let mvs = &fs.frame_mvs[0];
  let bsize = BlockSize::from_width_and_height(
    IMPORTANCE_BLOCK_SIZE,
    IMPORTANCE_BLOCK_SIZE,
  );
  (0..h_in_imp_b).for_each(|y| {
    (0..w_in_imp_b).for_each(|x| {
      let mv = mvs[y * 2][x * 2];

      // Coordinates of the top-left corner of the reference block, in MV
      // units.
      let reference_x = x as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.col as i64;
      let reference_y = y as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.row as i64;

      let region_org = plane_org.region(Area::Rect {
        x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
        y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      let region_ref = plane_ref.region(Area::Rect {
        x: reference_x as isize / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
        y: reference_y as isize / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
        width: IMPORTANCE_BLOCK_SIZE,
        height: IMPORTANCE_BLOCK_SIZE,
      });

      inter_costs.push(get_satd(
        &region_org,
        &region_ref,
        bsize,
        bit_depth,
        fi.cpu_feature_level,
      ));
    });
  });
  inter_costs.into_boxed_slice()
}

#[hawktracer(compute_motion_vectors_per_tile)]
fn compute_motion_vectors_per_tile<T: Pixel>(
  ts: &mut TileStateMut<T>, fi: &FrameInvariants<T>, inter_cfg: &InterConfig,
) -> (PlaneSuperBlockOffset, Vec<BlockPmv>) {
  // Compute the quarter-resolution motion vectors.
  let tile_pmvs = build_coarse_pmvs(fi, ts, inter_cfg);

  // Compute the half-resolution motion vectors.
  let mut half_res_pmvs = Vec::with_capacity(ts.sb_height * ts.sb_width);
  for sby in 0..ts.sb_height {
    for sbx in 0..ts.sb_width {
      let tile_sbo = TileSuperBlockOffset(SuperBlockOffset { x: sbx, y: sby });
      half_res_pmvs.push(build_half_res_pmvs(fi, ts, tile_sbo, &tile_pmvs));
    }
  }

  // Compute the full-resolution motion vectors.
  for sby in 0..ts.sb_height {
    for sbx in 0..ts.sb_width {
      let tile_sbo = TileSuperBlockOffset(SuperBlockOffset { x: sbx, y: sby });
      build_full_res_pmvs(fi, ts, tile_sbo, &half_res_pmvs);
    }
  }

  (ts.sbo, half_res_pmvs)
}

#[hawktracer(compute_motion_vectors)]
pub(crate) fn compute_motion_vectors<T: Pixel>(
  fi: &mut FrameInvariants<T>, fs: &mut FrameState<T>, inter_cfg: &InterConfig,
) {
  let mut blocks = FrameBlocks::new(fi.w_in_b, fi.h_in_b);
  fs.half_res_pmvs = fi
    .tiling
    .tile_iter_mut(fs, &mut blocks)
    .collect::<Vec<_>>()
    .into_par_iter()
    .map(|mut ctx| {
      let ts = &mut ctx.ts;
      compute_motion_vectors_per_tile(ts, fi, inter_cfg)
    })
    .collect::<Vec<_>>();
}
<|endoftext|>


<commit_before>extern crate rusqlite;
use super::sql_queries::sqlite3 as queries;
pub use super::base::SKStore;
use rusqlite::SqliteConnection;
// mod sql_queries;
// mod base;

pub struct SKSqliteStore<'a, 'b> {
    path: &'a str,
    conn: Option<&'b SqliteConnection>
}

impl<'a, 'b> SKSqliteStore<'a, 'b> {
    pub fn new(file: &'a str) -> SKSqliteStore {
        let store = SKSqliteStore { path: file, conn: None };
        store
    }
}

impl<'a, 'b> SKStore for SKSqliteStore<'a, 'b> {
    fn init(&self) {
        println!("Initializing database.");
    }
    fn connect(&self) {
        println!("Connecting to {}.", self.path);
    }
    fn disconnect(&self) {
        println!("Disconnecting from {}.", self.path);
    }
}

mod tests {
    use super::*;

    #[test]
    fn test_setup_new_db() {
        let store = SKSqliteStore::new("test.db");
        assert_eq!(store.path, "test.db");
    }
}
<commit_msg>Started implementing Sqlite connect and disconnect, but getting lifetime errors ;-(<commit_after>extern crate rusqlite;
use super::sql_queries::sqlite3 as queries;
pub use super::base::SKStore;
use rusqlite::SqliteConnection;
// mod sql_queries;
// mod base;

pub struct SKSqliteStore<'a, 'b> {
    path: &'a str,
    conn: Option<&'b SqliteConnection>
}

impl<'a, 'b> SKSqliteStore<'a, 'b> {
    pub fn new(file: &'a str) -> SKSqliteStore {
        let store = SKSqliteStore { path: file, conn: None };
        store
    }
}

impl<'a, 'b> SKStore for SKSqliteStore<'a, 'b> {
    fn init(&self) {
        println!("Initializing database.");
    }
    fn connect(&self) {
        let conn =
            match self.conn {
                Some(&c) => c,
                None => {
                    let c = SqliteConnection::open(self.path).unwrap();
                    self.conn = Some(&c);
                    c
                }
            };
        println!("Connected to {}.", self.path);
    }
    fn disconnect(&self) {
        match self.conn {
            Some(&c) => { c.close(); self.conn = None; true }
            None => false
        };
        println!("Disconnected from {}.", self.path);
    }
}

mod tests {
    use super::*;

    #[test]
    fn test_setup_new_db() {
        let store = SKSqliteStore::new("test.db");
        assert_eq!(store.path, "test.db");
    }
}
<|endoftext|>


<commit_before>use error::AttributeError;
use parsers::{Parse, ParseError};

/// Defines the units to be used for things that can consider a
/// coordinate system in terms of the current transformation, or in
/// terms of the current object's bounding box.
///
/// Keep in sync with rsvg-private.h:RsvgCoordUnits
#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum CoordUnits {
    UserSpaceOnUse,
    ObjectBoundingBox,
}

impl Parse for CoordUnits {
    type Data = ();
    type Err = AttributeError;

    fn parse(s: &str, _: ()) -> Result<CoordUnits, AttributeError> {
        match s {
            "userSpaceOnUse" => Ok(CoordUnits::UserSpaceOnUse),
            "objectBoundingBox" => Ok(CoordUnits::ObjectBoundingBox),
            _ => Err(AttributeError::Parse(ParseError::new(
                "expected 'userSpaceOnUse' or 'objectBoundingBox'",
            ))),
        }
    }
}

/// Creates a newtype around `CoordUnits`, with a default value.
///
/// SVG attributes that can take `userSpaceOnUse` or
/// `objectBoundingBox` values often have different default values
/// depending on the type of SVG element.  We use this macro to create
/// a newtype for each SVG element and attribute that requires values
/// of this type.  The newtype provides an `impl Default` with the
/// specified `$default` value.
#[macro_export]
macro_rules! coord_units {
    ($name:ident, $default:expr) => {
        #[derive(Debug, Copy, Clone, PartialEq, Eq)]
        pub struct $name(pub CoordUnits);

        impl Default for $name {
            fn default() -> Self {
                $name($default)
            }
        }

        impl From<$name> for CoordUnits {
            fn from(u: $name) -> Self {
                u.0
            }
        }

        impl $crate::parsers::Parse for $name {
            type Data = ();
            type Err = $crate::error::AttributeError;

            fn parse(s: &str, _: ()) -> Result<Self, $crate::error::AttributeError> {
                Ok($name($crate::coord_units::CoordUnits::parse(s, ())?))
            }
        }
    };
}

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

    coord_units!(MyUnits, CoordUnits::ObjectBoundingBox);

    #[test]
    fn parsing_invalid_strings_yields_error() {
        assert!(MyUnits::parse("", ()).is_err());
        assert!(MyUnits::parse("foo", ()).is_err());
    }

    #[test]
    fn parses_paint_server_units() {
        assert_eq!(
            MyUnits::parse("userSpaceOnUse", ()),
            Ok(MyUnits(CoordUnits::UserSpaceOnUse))
        );
        assert_eq!(
            MyUnits::parse("objectBoundingBox", ()),
            Ok(MyUnits(CoordUnits::ObjectBoundingBox))
        );
    }

    #[test]
    fn has_correct_default() {
        assert_eq!(MyUnits::default(), MyUnits(CoordUnits::ObjectBoundingBox));
    }

    #[test]
    fn converts_to_coord_units() {
        assert_eq!(
            CoordUnits::from(MyUnits(CoordUnits::ObjectBoundingBox)),
            CoordUnits::ObjectBoundingBox
        );
    }
}
<commit_msg>CoordUnits: use a cssparser<commit_after>use cssparser::{CowRcStr, Parser, ParserInput, Token};

use error::AttributeError;
use parsers::{Parse, ParseError};

/// Defines the units to be used for things that can consider a
/// coordinate system in terms of the current transformation, or in
/// terms of the current object's bounding box.
///
/// Keep in sync with rsvg-private.h:RsvgCoordUnits
#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum CoordUnits {
    UserSpaceOnUse,
    ObjectBoundingBox,
}

impl Parse for CoordUnits {
    type Data = ();
    type Err = AttributeError;

    fn parse(s: &str, _: ()) -> Result<CoordUnits, AttributeError> {
        let mut input = ParserInput::new(s);
        let mut parser = Parser::new(&mut input);
        let loc = parser.current_source_location();

        parser
            .expect_ident()
            .and_then(|cow| match cow.as_ref() {
                "userSpaceOnUse" => Ok(CoordUnits::UserSpaceOnUse),
                "objectBoundingBox" => Ok(CoordUnits::ObjectBoundingBox),
                _ => Err(
                    loc.new_basic_unexpected_token_error(Token::Ident(CowRcStr::from(
                        cow.as_ref().to_string(),
                    ))),
                ),
            })
            .map_err(|_| {
                AttributeError::Parse(ParseError::new(
                    "expected 'userSpaceOnUse' or 'objectBoundingBox'",
                ))
            })
    }
}

/// Creates a newtype around `CoordUnits`, with a default value.
///
/// SVG attributes that can take `userSpaceOnUse` or
/// `objectBoundingBox` values often have different default values
/// depending on the type of SVG element.  We use this macro to create
/// a newtype for each SVG element and attribute that requires values
/// of this type.  The newtype provides an `impl Default` with the
/// specified `$default` value.
#[macro_export]
macro_rules! coord_units {
    ($name:ident, $default:expr) => {
        #[derive(Debug, Copy, Clone, PartialEq, Eq)]
        pub struct $name(pub CoordUnits);

        impl Default for $name {
            fn default() -> Self {
                $name($default)
            }
        }

        impl From<$name> for CoordUnits {
            fn from(u: $name) -> Self {
                u.0
            }
        }

        impl $crate::parsers::Parse for $name {
            type Data = ();
            type Err = $crate::error::AttributeError;

            fn parse(s: &str, _: ()) -> Result<Self, $crate::error::AttributeError> {
                Ok($name($crate::coord_units::CoordUnits::parse(s, ())?))
            }
        }
    };
}

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

    coord_units!(MyUnits, CoordUnits::ObjectBoundingBox);

    #[test]
    fn parsing_invalid_strings_yields_error() {
        assert!(MyUnits::parse("", ()).is_err());
        assert!(MyUnits::parse("foo", ()).is_err());
    }

    #[test]
    fn parses_paint_server_units() {
        assert_eq!(
            MyUnits::parse("userSpaceOnUse", ()),
            Ok(MyUnits(CoordUnits::UserSpaceOnUse))
        );
        assert_eq!(
            MyUnits::parse("objectBoundingBox", ()),
            Ok(MyUnits(CoordUnits::ObjectBoundingBox))
        );
    }

    #[test]
    fn has_correct_default() {
        assert_eq!(MyUnits::default(), MyUnits(CoordUnits::ObjectBoundingBox));
    }

    #[test]
    fn converts_to_coord_units() {
        assert_eq!(
            CoordUnits::from(MyUnits(CoordUnits::ObjectBoundingBox)),
            CoordUnits::ObjectBoundingBox
        );
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.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.

use {LinkerFlavor, PanicStrategy};
use target::{Target, TargetOptions, TargetResult};

pub fn target() -> TargetResult {
    Ok(Target {
        llvm_target: "msp430-none-elf".to_string(),
        target_endian: "little".to_string(),
        target_pointer_width: "16".to_string(),
        target_c_int_width: "16".to_string(),
        data_layout: "e-m:e-p:16:16-i32:16-i64:16-f32:16-f64:16-a:8-n8:16-S16".to_string(),
        arch: "msp430".to_string(),
        target_os: "none".to_string(),
        target_env: "".to_string(),
        target_vendor: "".to_string(),
        linker_flavor: LinkerFlavor::Gcc,

        options: TargetOptions {
            executables: true,

            // The LLVM backend currently can't generate object files. To
            // workaround this LLVM generates assembly files which then we feed
            // to gcc to get object files. For this reason we have a hard
            // dependency on this specific gcc.
            asm_args: vec!["-mcpu=msp430".to_string()],
            linker: "msp430-elf-gcc".to_string(),
            no_integrated_as: true,

            // There are no atomic instructions available in the MSP430
            // instruction set
            max_atomic_width: Some(0),

            // Because these devices have very little resources having an
            // unwinder is too onerous so we default to "abort" because the
            // "unwind" strategy is very rare.
            panic_strategy: PanicStrategy::Abort,

            // Similarly, one almost always never wants to use relocatable
            // code because of the extra costs it involves.
            relocation_model: "static".to_string(),

            // Right now we invoke an external assembler and this isn't
            // compatible with multiple codegen units, and plus we probably
            // don't want to invoke that many gcc instances.
            default_codegen_units: Some(1),

            .. Default::default( )
        }
    })
}
<commit_msg>Rollup merge of #47890 - pftbest:no_trap, r=alexcrichton<commit_after>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.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.

use {LinkerFlavor, PanicStrategy};
use target::{Target, TargetOptions, TargetResult};

pub fn target() -> TargetResult {
    Ok(Target {
        llvm_target: "msp430-none-elf".to_string(),
        target_endian: "little".to_string(),
        target_pointer_width: "16".to_string(),
        target_c_int_width: "16".to_string(),
        data_layout: "e-m:e-p:16:16-i32:16-i64:16-f32:16-f64:16-a:8-n8:16-S16".to_string(),
        arch: "msp430".to_string(),
        target_os: "none".to_string(),
        target_env: "".to_string(),
        target_vendor: "".to_string(),
        linker_flavor: LinkerFlavor::Gcc,

        options: TargetOptions {
            executables: true,

            // The LLVM backend currently can't generate object files. To
            // workaround this LLVM generates assembly files which then we feed
            // to gcc to get object files. For this reason we have a hard
            // dependency on this specific gcc.
            asm_args: vec!["-mcpu=msp430".to_string()],
            linker: "msp430-elf-gcc".to_string(),
            no_integrated_as: true,

            // There are no atomic instructions available in the MSP430
            // instruction set
            max_atomic_width: Some(0),

            // Because these devices have very little resources having an
            // unwinder is too onerous so we default to "abort" because the
            // "unwind" strategy is very rare.
            panic_strategy: PanicStrategy::Abort,

            // Similarly, one almost always never wants to use relocatable
            // code because of the extra costs it involves.
            relocation_model: "static".to_string(),

            // Right now we invoke an external assembler and this isn't
            // compatible with multiple codegen units, and plus we probably
            // don't want to invoke that many gcc instances.
            default_codegen_units: Some(1),

            // Since MSP430 doesn't meaningfully support faulting on illegal
            // instructions, LLVM generates a call to abort() function instead
            // of a trap instruction. Such calls are 4 bytes long, and that is
            // too much overhead for such small target.
            trap_unreachable: false,

            .. Default::default( )
        }
    })
}
<|endoftext|>


<commit_before>use rustwlc::{input, Point, ResizeEdge, Geometry,
              RESIZE_LEFT, RESIZE_RIGHT, RESIZE_TOP, RESIZE_BOTTOM,
              RESIZE_TOPLEFT, RESIZE_TOPRIGHT, RESIZE_BOTTOMLEFT, RESIZE_BOTTOMRIGHT,
};

use super::super::{Action, Direction, LayoutTree, TreeError};
use super::super::commands::{CommandResult};
use super::super::core::container::{ContainerType, MIN_SIZE};
use uuid::Uuid;

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum ResizeErr {
    /// Expected the node associated with the UUID to be floating.
    ExpectedFloating(Uuid),
    /// Expected the node associated with the UUID to not be floating
    ExpectedNotFloating(Uuid)
}

impl LayoutTree {
    /// Updates pointer position for the node behind the id,
    /// in the same direction as the edge.
    pub fn update_pointer_pos(&self, id: Uuid, edge: ResizeEdge) -> CommandResult {
        let container = try!(self.lookup(id));
        let Geometry { mut origin, size } = container.get_geometry()
            .expect("Container had no geometry");
        drop(container);
        if edge.contains(RESIZE_TOPLEFT) {
            input::pointer::set_position(origin);
        } else if edge.contains(RESIZE_TOPRIGHT) {
            origin.x += size.w as i32;
            input::pointer::set_position(origin);
        } else if edge.contains(RESIZE_BOTTOMLEFT) {
            origin.y += size.h as i32;
            input::pointer::set_position(origin);
        } else if edge.contains(RESIZE_BOTTOMRIGHT) {
            origin.x += size.w as i32;
            origin.y += size.h as i32;
            input::pointer::set_position(origin);
        }
        Ok(())
    }

    /// Resizes a floating container. If the container was not floating, an Err is returned.
    pub fn resize_floating(&mut self, id: Uuid, edge: ResizeEdge, pointer: Point,
                           action: &mut Action) -> CommandResult {
        {
            let container = try!(self.lookup_mut(id));
            if !container.floating() {
                return Err(TreeError::Resize(ResizeErr::ExpectedFloating(container.get_id())))
            }
            match container.get_type() {
                ContainerType::View | ContainerType::Container => {},
                _ => return Err(TreeError::UuidWrongType(container.get_id(),
                                                        vec!(ContainerType::View)))
            }
            let geo = container.get_geometry()
                .expect("Could not get geometry of the container");

            let new_geo = calculate_resize(geo, edge, pointer, action.grab);
            action.grab = pointer;
            container.set_geometry(edge, new_geo);
        }
        self.update_pointer_pos(id, edge)
    }

    pub fn resize_tiled(&mut self, id: Uuid, edge: ResizeEdge, pointer: Point,
                        action: &mut Action) -> CommandResult {
        // This is the vector of operations we will perform, we do all geometry sets atomically.
        let mut resizing_ops: Vec<(Uuid, (ResizeEdge, Geometry))> = Vec::with_capacity(4);
        let dirs_moving_in = Direction::from_edge(edge);
        let next_containers: Vec<(Uuid, Uuid)> = dirs_moving_in.iter()
            .map(|dir| self.container_in_dir(id, *dir))
            .flat_map(|result| result.into_iter())
            .collect();
        for ancestor_id in next_containers.iter().map(|uuids| uuids.0) {
            let container = try!(self.lookup(ancestor_id));
            if container.floating() {
                return Err(TreeError::Resize(ResizeErr::ExpectedNotFloating(ancestor_id)))
            }
            match container.get_type() {
                ContainerType::View | ContainerType::Container => {},
                _ => return Err(TreeError::UuidWrongType(ancestor_id,
                                                        vec!(ContainerType::View)))
            }
            let geo = container.get_geometry()
                .expect("Could not get geometry of the container");
            if geometry_resize_too_small(geo, edge, pointer, action.grab) {
                return Ok(())
            }
            let new_geo = calculate_resize(geo, edge, pointer, action.grab);
            resizing_ops.push((ancestor_id, (edge, new_geo)));
        }
        let siblings: Vec<Uuid> = next_containers.into_iter()
            .map(|uuids| uuids.1).collect();
        // and now we mutate the siblings
        let reversed_dir: Vec<Direction> = dirs_moving_in.iter()
            .map(|dir| dir.reverse()).collect();
        let reversed_edge = Direction::to_edge(reversed_dir.as_slice());
        for sibling in siblings {
            let container = try!(self.lookup_mut(sibling));
            if container.floating() {
                return Err(TreeError::Resize(ResizeErr::ExpectedNotFloating(container.get_id())))
            }
            match container.get_type() {
                ContainerType::View | ContainerType::Container => {},
                _ => return Err(TreeError::UuidWrongType(container.get_id(),
                                                         vec!(ContainerType::View)))
            }
            let geo = container.get_geometry()
                .expect("Could not get geometry of the container");
            if geometry_resize_too_small(geo, reversed_edge, pointer, action.grab) {
                return Ok(())
            }
            let new_geo = calculate_resize(geo, reversed_edge, pointer, action.grab);
            if new_geo.size.w <= MIN_SIZE.w || new_geo.size.h <= MIN_SIZE.h {
                return Ok(())
            }
            resizing_ops.push((sibling, (reversed_edge, new_geo)));
        }
        action.grab = pointer;
        for (id, (edge, geo)) in resizing_ops {
            let container = self.lookup_mut(id)
                .expect("Id no longer points to node!");
            container.set_geometry(edge, geo);
        }
        action.grab = pointer;
        let node_ix = self.tree.lookup_id(id).unwrap();
        let workspace_ix = self.tree.ancestor_of_type(node_ix, ContainerType::Workspace).unwrap();
        self.layout(workspace_ix);
        self.update_pointer_pos(id, edge)
    }
}

/// Calculates what the new geometry is of a window.
/// Needs the geometry of the window, the edge direction the pointer is moving in,
/// the current position of the pointer, and the previous place the pointer was at.
fn calculate_resize(geo: Geometry, edge: ResizeEdge,
                    cur_pointer: Point, prev_pointer: Point) -> Geometry {
    let mut new_geo = geo.clone();
    let dx = cur_pointer.x - prev_pointer.x;
    let dy = cur_pointer.y - prev_pointer.y;
    if edge.contains(RESIZE_LEFT) {
        if dx < 0 {
            new_geo.size.w = geo.size.w.saturating_add(dx.abs() as u32);
        } else {
            new_geo.size.w = geo.size.w.saturating_sub(dx.abs() as u32);
        }
        new_geo.origin.x += dx;
    }
    else if edge.contains(RESIZE_RIGHT) {
        if dx < 0 {
            new_geo.size.w = geo.size.w.saturating_sub(dx.abs() as u32);
        } else {
            new_geo.size.w = geo.size.w.saturating_add(dx.abs() as u32);
        }
    }

    if edge.contains(RESIZE_TOP) {
        if dy < 0 {
            new_geo.size.h = geo.size.h.saturating_add(dy.abs() as u32);
        } else {
            new_geo.size.h = geo.size.h.saturating_sub(dy.abs() as u32);
        }
        new_geo.origin.y += dy;
    }
    else if edge.contains(RESIZE_BOTTOM) {
        if dy < 0 {
            new_geo.size.h = geo.size.h.saturating_sub(dy.abs() as u32);
        } else {
            new_geo.size.h = geo.size.h.saturating_add(dy.abs() as u32);
        }
    }

    if new_geo.size.w < MIN_SIZE.w {
        new_geo.origin.x = geo.origin.x;
        new_geo.size.w = MIN_SIZE.w;
    }

    if new_geo.size.h < MIN_SIZE.h {
        new_geo.origin.y = geo.origin.y;
        new_geo.size.h = MIN_SIZE.h;
    }
    new_geo
}

/// If the geometry is at the minimum size (in either the x or y plane)
/// and the pointer is trying to make it even smaller in that direction,
/// the it returns true (to indicate you should abandon all operations).
///
/// Otherwise returns false
fn geometry_resize_too_small(geo: Geometry, edge: ResizeEdge, cur_point: Point,
                       prev_point: Point) -> bool {
    if geo.size.w > MIN_SIZE.w && geo.size.h > MIN_SIZE.h {
        return false
    }
    if edge.contains(RESIZE_RIGHT) && cur_point.x - prev_point.x < 0 {
        true
    } else if edge.contains(RESIZE_LEFT) && cur_point.x - prev_point.x > 0 {
        true
    } else if edge.contains(RESIZE_TOP) && cur_point.y - prev_point.y > 0 {
        true
    } else if edge.contains(RESIZE_BOTTOM) && cur_point.y - prev_point.y < 0 {
        true
    } else {
        false
    }
}
<commit_msg>Fixes #180, prev pointer now uses window origin<commit_after>use rustwlc::{input, Point, ResizeEdge, Geometry,
              RESIZE_LEFT, RESIZE_RIGHT, RESIZE_TOP, RESIZE_BOTTOM,
              RESIZE_TOPLEFT, RESIZE_TOPRIGHT, RESIZE_BOTTOMLEFT, RESIZE_BOTTOMRIGHT,
};

use super::super::{Action, Direction, LayoutTree, TreeError};
use super::super::commands::{CommandResult};
use super::super::core::container::{ContainerType, MIN_SIZE};
use uuid::Uuid;

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum ResizeErr {
    /// Expected the node associated with the UUID to be floating.
    ExpectedFloating(Uuid),
    /// Expected the node associated with the UUID to not be floating
    ExpectedNotFloating(Uuid)
}

impl LayoutTree {
    /// Updates pointer position for the node behind the id,
    /// in the same direction as the edge.
    pub fn update_pointer_pos(&self, id: Uuid, edge: ResizeEdge)
        -> Result<Point, TreeError> {
        let container = try!(self.lookup(id));
        let Geometry { mut origin, size } = container.get_geometry()
            .expect("Container had no geometry");
        drop(container);
        if edge.contains(RESIZE_TOPLEFT) {
            input::pointer::set_position(origin);
        } else if edge.contains(RESIZE_TOPRIGHT) {
            origin.x += size.w as i32;
            input::pointer::set_position(origin);
        } else if edge.contains(RESIZE_BOTTOMLEFT) {
            origin.y += size.h as i32;
            input::pointer::set_position(origin);
        } else if edge.contains(RESIZE_BOTTOMRIGHT) {
            origin.x += size.w as i32;
            origin.y += size.h as i32;
            input::pointer::set_position(origin);
        }
        Ok((origin))
    }

    /// Resizes a floating container. If the container was not floating, an Err is returned.
    pub fn resize_floating(&mut self, id: Uuid, edge: ResizeEdge, pointer: Point,
                           action: &mut Action) -> CommandResult {
        {
            let container = try!(self.lookup_mut(id));
            if !container.floating() {
                return Err(TreeError::Resize(ResizeErr::ExpectedFloating(container.get_id())))
            }
            match container.get_type() {
                ContainerType::View | ContainerType::Container => {},
                _ => return Err(TreeError::UuidWrongType(container.get_id(),
                                                        vec!(ContainerType::View)))
            }
            let geo = container.get_geometry()
                .expect("Could not get geometry of the container");

            let new_geo = calculate_resize(geo, edge, pointer, action.grab);
            container.set_geometry(edge, new_geo);
        }
        action.grab = self.update_pointer_pos(id, edge)
            .expect("Could not update pointer position");
        Ok(())
    }

    pub fn resize_tiled(&mut self, id: Uuid, edge: ResizeEdge, pointer: Point,
                        action: &mut Action) -> CommandResult {
        // This is the vector of operations we will perform, we do all geometry sets atomically.
        let mut resizing_ops: Vec<(Uuid, (ResizeEdge, Geometry))> = Vec::with_capacity(4);
        let dirs_moving_in = Direction::from_edge(edge);
        let next_containers: Vec<(Uuid, Uuid)> = dirs_moving_in.iter()
            .map(|dir| self.container_in_dir(id, *dir))
            .flat_map(|result| result.into_iter())
            .collect();
        for ancestor_id in next_containers.iter().map(|uuids| uuids.0) {
            let container = try!(self.lookup(ancestor_id));
            if container.floating() {
                return Err(TreeError::Resize(ResizeErr::ExpectedNotFloating(ancestor_id)))
            }
            match container.get_type() {
                ContainerType::View | ContainerType::Container => {},
                _ => return Err(TreeError::UuidWrongType(ancestor_id,
                                                        vec!(ContainerType::View)))
            }
            let geo = container.get_geometry()
                .expect("Could not get geometry of the container");
            if geometry_resize_too_small(geo, edge, pointer, action.grab) {
                return Ok(())
            }
            let new_geo = calculate_resize(geo, edge, pointer, action.grab);
            resizing_ops.push((ancestor_id, (edge, new_geo)));
        }
        let siblings: Vec<Uuid> = next_containers.into_iter()
            .map(|uuids| uuids.1).collect();
        // and now we mutate the siblings
        let reversed_dir: Vec<Direction> = dirs_moving_in.iter()
            .map(|dir| dir.reverse()).collect();
        let reversed_edge = Direction::to_edge(reversed_dir.as_slice());
        for sibling in siblings {
            let container = try!(self.lookup_mut(sibling));
            if container.floating() {
                return Err(TreeError::Resize(ResizeErr::ExpectedNotFloating(container.get_id())))
            }
            match container.get_type() {
                ContainerType::View | ContainerType::Container => {},
                _ => return Err(TreeError::UuidWrongType(container.get_id(),
                                                         vec!(ContainerType::View)))
            }
            let geo = container.get_geometry()
                .expect("Could not get geometry of the container");
            if geometry_resize_too_small(geo, reversed_edge, pointer, action.grab) {
                return Ok(())
            }
            let new_geo = calculate_resize(geo, reversed_edge, pointer, action.grab);
            if new_geo.size.w <= MIN_SIZE.w || new_geo.size.h <= MIN_SIZE.h {
                return Ok(())
            }
            resizing_ops.push((sibling, (reversed_edge, new_geo)));
        }
        action.grab = pointer;
        for (id, (edge, geo)) in resizing_ops {
            let container = self.lookup_mut(id)
                .expect("Id no longer points to node!");
            container.set_geometry(edge, geo);
        }
        let node_ix = self.tree.lookup_id(id).unwrap();
        let workspace_ix = self.tree.ancestor_of_type(node_ix, ContainerType::Workspace).unwrap();
        self.layout(workspace_ix);
        action.grab = self.update_pointer_pos(id, edge)
            .expect("Could not update pointer position");
        Ok(())
    }
}

/// Calculates what the new geometry is of a window.
/// Needs the geometry of the window, the edge direction the pointer is moving in,
/// the current position of the pointer, and the previous place the pointer was at.
fn calculate_resize(geo: Geometry, edge: ResizeEdge,
                    cur_pointer: Point, prev_pointer: Point) -> Geometry {
    let mut new_geo = geo.clone();
    let dx = cur_pointer.x - prev_pointer.x;
    let dy = cur_pointer.y - prev_pointer.y;
    if edge.contains(RESIZE_LEFT) {
        if dx < 0 {
            new_geo.size.w = geo.size.w.saturating_add(dx.abs() as u32);
        } else {
            new_geo.size.w = geo.size.w.saturating_sub(dx.abs() as u32);
        }
        new_geo.origin.x += dx;
    }
    else if edge.contains(RESIZE_RIGHT) {
        if dx < 0 {
            new_geo.size.w = geo.size.w.saturating_sub(dx.abs() as u32);
        } else {
            new_geo.size.w = geo.size.w.saturating_add(dx.abs() as u32);
        }
    }

    if edge.contains(RESIZE_TOP) {
        if dy < 0 {
            new_geo.size.h = geo.size.h.saturating_add(dy.abs() as u32);
        } else {
            new_geo.size.h = geo.size.h.saturating_sub(dy.abs() as u32);
        }
        new_geo.origin.y += dy;
    }
    else if edge.contains(RESIZE_BOTTOM) {
        if dy < 0 {
            new_geo.size.h = geo.size.h.saturating_sub(dy.abs() as u32);
        } else {
            new_geo.size.h = geo.size.h.saturating_add(dy.abs() as u32);
        }
    }

    if new_geo.size.w < MIN_SIZE.w {
        new_geo.origin.x = geo.origin.x;
        new_geo.size.w = MIN_SIZE.w;
    }

    if new_geo.size.h < MIN_SIZE.h {
        new_geo.origin.y = geo.origin.y;
        new_geo.size.h = MIN_SIZE.h;
    }
    new_geo
}

/// If the geometry is at the minimum size (in either the x or y plane)
/// and the pointer is trying to make it even smaller in that direction,
/// the it returns true (to indicate you should abandon all operations).
///
/// Otherwise returns false
fn geometry_resize_too_small(geo: Geometry, edge: ResizeEdge, cur_point: Point,
                       prev_point: Point) -> bool {
    if geo.size.w > MIN_SIZE.w && geo.size.h > MIN_SIZE.h {
        return false
    }
    if edge.contains(RESIZE_RIGHT) && cur_point.x - prev_point.x < 0 {
        true
    } else if edge.contains(RESIZE_LEFT) && cur_point.x - prev_point.x > 0 {
        true
    } else if edge.contains(RESIZE_TOP) && cur_point.y - prev_point.y > 0 {
        true
    } else if edge.contains(RESIZE_BOTTOM) && cur_point.y - prev_point.y < 0 {
        true
    } else {
        false
    }
}
<|endoftext|>