<commit_before>// Copyright 2015, 2016 Ethcore (UK) Ltd.
// This file is part of Parity.

// Parity is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Parity is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Parity.  If not, see <http://www.gnu.org/licenses/>.

//! jsonrpc server over win named pipes
//!
//! ```no_run
//! extern crate jsonrpc_core;
//! extern crate jsonrpc_ipc_server;
//!
//! use std::sync::Arc;
//! use jsonrpc_core::*;
//! use jsonrpc_ipc_server::Server;
//!
//! fn main() {
//! 	let mut io = IoHandler::new();
//! 	io.add_method("say_hello", |_params| {
//!			Ok(Value::String("hello".into()))
//!		});
//! 	let server = Server::new("/tmp/json-ipc-test.ipc", io).unwrap();
//!     ::std::thread::spawn(move || server.run());
//! }
//! ```

//! Named pipes library

use miow::pipe::{NamedPipe, NamedPipeBuilder};
use std;
use std::io;
use std::io::{Read, Write};
use std::sync::atomic::*;
use std::sync::{Arc, Mutex};
use jsonrpc_core::{Metadata, MetaIoHandler};
use jsonrpc_core::reactor::{RpcHandler, RpcEventLoop, RpcEventLoopHandle};
use validator;

pub type Result<T> = std::result::Result<T, Error>;

const MAX_REQUEST_LEN: u32 = 65536;
const REQUEST_READ_BATCH: usize = 4096;

#[derive(Debug)]
pub enum Error {
	Io(std::io::Error),
	NotStarted,
	AlreadyStopping,
	NotStopped,
	IsStopping,
}

impl std::convert::From<std::io::Error> for Error {
	fn from(io_error: std::io::Error) -> Error {
		Error::Io(io_error)
	}
}

pub struct PipeHandler<M: Metadata = ()> {
	waiting_pipe: NamedPipe,
	io_handler: RpcHandler<M>,
}

impl<M: Metadata> PipeHandler<M> {
	/// start ipc rpc server
	pub fn start(addr: &str, io_handler: RpcHandler<M>) -> Result<PipeHandler<M>> {
		Ok(PipeHandler {
			waiting_pipe: try!(
				NamedPipeBuilder::new(addr)
					.first(true)
					.accept_remote(true)
					.max_instances(255)
					.inbound(true)
					.outbound(true)
					.out_buffer_size(MAX_REQUEST_LEN)
					.in_buffer_size(MAX_REQUEST_LEN)
					.create()
			),
			io_handler: io_handler.clone(),
		})
	}

	fn handle_incoming(&mut self, addr: &str, stop: Arc<AtomicBool>) -> io::Result<()> {
		trace!(target: "ipc", "Waiting for client: [{}]", addr);
		try!(self.waiting_pipe.connect());
		trace!(target: "ipc", "Received connection to address [{}]", addr);
		if stop.load(Ordering::Relaxed) {
			trace!(target: "ipc", "Stopped listening sequence [{}]", addr);
			return Ok(())
		}

		let mut connected_pipe = std::mem::replace::<NamedPipe>(&mut self.waiting_pipe,
			try!(NamedPipeBuilder::new(addr)
				.first(false)
				.accept_remote(true)
				.inbound(true)
				.outbound(true)
				.out_buffer_size(MAX_REQUEST_LEN)
				.in_buffer_size(MAX_REQUEST_LEN)
				.create()));

		let thread_handler = self.io_handler.clone();
		std::thread::spawn(move || {
			let mut buf = vec![0u8; MAX_REQUEST_LEN as usize];
			let mut fin = REQUEST_READ_BATCH;
			loop {
				let start = fin - REQUEST_READ_BATCH;
				trace!(target: "ipc", "Reading {} - {} of the buffer", start, fin);
				match connected_pipe.read(&mut buf[start..fin]) {
					Ok(size) => {
						let (requests, last_index) = {
							let effective = &buf[0..start + size];
							fin = fin + size;
							trace!(target: "ipc", "Received rpc data: {} bytes", effective.len());

							validator::extract_requests(effective)
						};
						if requests.len() > 0 {
							let mut response_buf = Vec::new();
							for rpc_msg in requests  {
								trace!(target: "ipc", "Request: {}", rpc_msg);
								let meta = Default::default();
								let response: Option<String> = thread_handler.handle_request_sync(&rpc_msg, meta);

								if let Some(response_str) = response {
									trace!(target: "ipc", "Response: {}", &response_str);
									response_buf.extend(response_str.into_bytes());
								}
							}

							if let Err(write_err) = connected_pipe.write_all(&response_buf[..]).and_then(|_| connected_pipe.flush()) {
								trace!(target: "ipc", "Response write error: {:?}", write_err);
							}
							else {
								trace!(target: "ipc", "Sent rpc response: {} bytes", response_buf.len());
							}

							let leftover_len = start + size - (last_index + 1);
							if leftover_len > 0 {
								let leftover = buf[last_index + 1..start + size].to_vec();
								buf[0..leftover_len].copy_from_slice(&leftover[..]);
							}
							fin = leftover_len + REQUEST_READ_BATCH;
						}
						else { continue; }
					},
					Err(e) => {
						// closed connection
						trace!(target: "ipc", "Dropped connection {:?}", e);
						break;
					}
				}
			}
		});

		Ok(())
	}
}

pub struct Server<M: Metadata = ()> {
	is_stopping: Arc<AtomicBool>,
	is_stopped: Arc<AtomicBool>,
	io_handler: RpcHandler<M>,
	rpc_event_loop: Mutex<Option<RpcEventLoopHandle>>,
	addr: String,
}

impl<M: Metadata> Server<M> {
	/// New server
	pub fn new<T>(socket_addr: &str, io_handler: T) -> Result<Server<M>> where
		T: Into<MetaIoHandler<M>>,
	{
		let rpc_loop = RpcEventLoop::spawn();
		let mut server = try!(Self::with_rpc_handler(socket_addr, rpc_loop.handler(Arc::new(io_handler.into()))));
		server.rpc_event_loop = Mutex::new(Some(rpc_loop.into()));
		Ok(server)
	}

	// New Server using RpcHandler
	pub fn with_rpc_handler(socket_addr: &str, io_handler: RpcHandler<M>) -> Result<Server<M>> {
		Ok(Server {
			io_handler: io_handler,
			is_stopping: Arc::new(AtomicBool::new(false)),
			is_stopped: Arc::new(AtomicBool::new(true)),
			rpc_event_loop: Mutex::new(None),
			addr: socket_addr.to_owned(),
		})
	}

	/// Run server (in this thread)
	pub fn run(&self) -> Result<()> {
		let mut pipe_handler = try!(PipeHandler::start(&self.addr, self.io_handler.clone()));
		loop  {
			try!(pipe_handler.handle_incoming(&self.addr, Arc::new(AtomicBool::new(false))));
		}
	}

	/// Run server (in separate thread)
	pub fn run_async(&self) -> Result<()> {
		if self.is_stopping.load(Ordering::Relaxed) { return Err(Error::IsStopping) }
		if !self.is_stopped.load(Ordering::Relaxed) { return Err(Error::NotStopped) }

		trace!(target: "ipc", "Started named pipes server [{}]", self.addr);

		let thread_stopping = self.is_stopping.clone();
		let thread_stopped = self.is_stopped.clone();
		let thread_handler = self.io_handler.clone();
		let addr = self.addr.clone();
		std::thread::spawn(move || {
			let mut pipe_handler = PipeHandler::start(&addr, thread_handler).unwrap();
			while !thread_stopping.load(Ordering::Relaxed) {
				trace!(target: "ipc", "Accepting pipe connection");
				if let Err(pipe_listener_error) = pipe_handler.handle_incoming(&addr, thread_stopping.clone()) {
					trace!(target: "ipc", "Pipe listening error: {:?}", pipe_listener_error);
				}
			}
			thread_stopped.store(true, Ordering::Relaxed);
		});

		self.is_stopped.store(false, Ordering::Relaxed);
		Ok(())
	}

	pub fn stop_async(&self) -> Result<()> {
		self.rpc_event_loop.lock().unwrap().take().map(|s| s.close());
		if self.is_stopped.load(Ordering::Relaxed) { return Err(Error::NotStarted) }
		if self.is_stopping.load(Ordering::Relaxed) { return Err(Error::AlreadyStopping)}
		self.is_stopping.store(true, Ordering::Relaxed);
		Ok(())
	}

	pub fn stop(&self) -> Result<()> {
		self.rpc_event_loop.lock().unwrap().take().map(|s| s.close());
		if self.is_stopped.load(Ordering::Relaxed) { return Err(Error::NotStarted) }
		if self.is_stopping.load(Ordering::Relaxed) { return Err(Error::AlreadyStopping)}
		self.is_stopping.store(true, Ordering::Relaxed);
		while !self.is_stopped.load(Ordering::Relaxed) { std::thread::park_timeout(std::time::Duration::new(0, 50)); }
		Ok(())
	}
}

impl<M: Metadata> Drop for Server<M> {
	fn drop(&mut self) {
		self.stop_async().unwrap_or_else(|_| {}); // ignore error - can be stopped already
		// todo : no stable logging for windows?
		trace!(target: "ipc", "IPC Server : shutdown");
	}
}
<commit_msg>add wait with timeout<commit_after>// Copyright 2015, 2016 Ethcore (UK) Ltd.
// This file is part of Parity.

// Parity is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Parity is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Parity.  If not, see <http://www.gnu.org/licenses/>.

//! jsonrpc server over win named pipes
//!
//! ```no_run
//! extern crate jsonrpc_core;
//! extern crate jsonrpc_ipc_server;
//!
//! use std::sync::Arc;
//! use jsonrpc_core::*;
//! use jsonrpc_ipc_server::Server;
//!
//! fn main() {
//! 	let mut io = IoHandler::new();
//! 	io.add_method("say_hello", |_params| {
//!			Ok(Value::String("hello".into()))
//!		});
//! 	let server = Server::new("/tmp/json-ipc-test.ipc", io).unwrap();
//!     ::std::thread::spawn(move || server.run());
//! }
//! ```

//! Named pipes library

use miow::pipe::{NamedPipe, NamedPipeBuilder};
use std;
use std::io;
use std::io::{Read, Write};
use std::sync::atomic::*;
use std::sync::{Arc, Mutex};
use jsonrpc_core::{Metadata, MetaIoHandler};
use jsonrpc_core::reactor::{RpcHandler, RpcEventLoop, RpcEventLoopHandle};
use validator;

pub type Result<T> = std::result::Result<T, Error>;

const MAX_REQUEST_LEN: u32 = 65536;
const REQUEST_READ_BATCH: usize = 4096;

#[derive(Debug)]
pub enum Error {
	Io(std::io::Error),
	NotStarted,
	AlreadyStopping,
	NotStopped,
	IsStopping,
}

impl std::convert::From<std::io::Error> for Error {
	fn from(io_error: std::io::Error) -> Error {
		Error::Io(io_error)
	}
}

pub struct PipeHandler<M: Metadata = ()> {
	waiting_pipe: NamedPipe,
	io_handler: RpcHandler<M>,
}

impl<M: Metadata> PipeHandler<M> {
	/// start ipc rpc server
	pub fn start(addr: &str, io_handler: RpcHandler<M>) -> Result<PipeHandler<M>> {
		Ok(PipeHandler {
			waiting_pipe: try!(
				NamedPipeBuilder::new(addr)
					.first(true)
					.accept_remote(true)
					.max_instances(255)
					.inbound(true)
					.outbound(true)
					.out_buffer_size(MAX_REQUEST_LEN)
					.in_buffer_size(MAX_REQUEST_LEN)
					.create()
			),
			io_handler: io_handler.clone(),
		})
	}

	fn handle_incoming(&mut self, addr: &str, stop: Arc<AtomicBool>) -> io::Result<()> {
		trace!(target: "ipc", "Waiting for client: [{}]", addr);

		// blocking wait with small timeouts
		// allows check if the server is actually stopped to quit gracefully
		// (`connect` does not allow that, it will block indefinitely)
		loop {
			if let Ok(_) = NamedPipe::wait(addr, Some(200)) {
				try!(self.waiting_pipe.connect());
				trace!(target: "ipc", "Received connection to address [{}]", addr);
				break;
			}
			if stop.load(Ordering::Relaxed) {
				trace!(target: "ipc", "Stopped listening sequence [{}]", addr);
				return Ok(())
			}
		}

		let mut connected_pipe = std::mem::replace::<NamedPipe>(&mut self.waiting_pipe,
			try!(NamedPipeBuilder::new(addr)
				.first(false)
				.accept_remote(true)
				.inbound(true)
				.outbound(true)
				.out_buffer_size(MAX_REQUEST_LEN)
				.in_buffer_size(MAX_REQUEST_LEN)
				.create()));

		let thread_handler = self.io_handler.clone();
		std::thread::spawn(move || {
			let mut buf = vec![0u8; MAX_REQUEST_LEN as usize];
			let mut fin = REQUEST_READ_BATCH;
			loop {
				let start = fin - REQUEST_READ_BATCH;
				trace!(target: "ipc", "Reading {} - {} of the buffer", start, fin);
				match connected_pipe.read(&mut buf[start..fin]) {
					Ok(size) => {
						let (requests, last_index) = {
							let effective = &buf[0..start + size];
							fin = fin + size;
							trace!(target: "ipc", "Received rpc data: {} bytes", effective.len());

							validator::extract_requests(effective)
						};
						if requests.len() > 0 {
							let mut response_buf = Vec::new();
							for rpc_msg in requests  {
								trace!(target: "ipc", "Request: {}", rpc_msg);
								let meta = Default::default();
								let response: Option<String> = thread_handler.handle_request_sync(&rpc_msg, meta);

								if let Some(response_str) = response {
									trace!(target: "ipc", "Response: {}", &response_str);
									response_buf.extend(response_str.into_bytes());
								}
							}

							if let Err(write_err) = connected_pipe.write_all(&response_buf[..]).and_then(|_| connected_pipe.flush()) {
								trace!(target: "ipc", "Response write error: {:?}", write_err);
							}
							else {
								trace!(target: "ipc", "Sent rpc response: {} bytes", response_buf.len());
							}

							let leftover_len = start + size - (last_index + 1);
							if leftover_len > 0 {
								let leftover = buf[last_index + 1..start + size].to_vec();
								buf[0..leftover_len].copy_from_slice(&leftover[..]);
							}
							fin = leftover_len + REQUEST_READ_BATCH;
						}
						else { continue; }
					},
					Err(e) => {
						// closed connection
						trace!(target: "ipc", "Dropped connection {:?}", e);
						break;
					}
				}
			}
		});

		Ok(())
	}
}

pub struct Server<M: Metadata = ()> {
	is_stopping: Arc<AtomicBool>,
	is_stopped: Arc<AtomicBool>,
	io_handler: RpcHandler<M>,
	rpc_event_loop: Mutex<Option<RpcEventLoopHandle>>,
	addr: String,
}

impl<M: Metadata> Server<M> {
	/// New server
	pub fn new<T>(socket_addr: &str, io_handler: T) -> Result<Server<M>> where
		T: Into<MetaIoHandler<M>>,
	{
		let rpc_loop = RpcEventLoop::spawn();
		let mut server = try!(Self::with_rpc_handler(socket_addr, rpc_loop.handler(Arc::new(io_handler.into()))));
		server.rpc_event_loop = Mutex::new(Some(rpc_loop.into()));
		Ok(server)
	}

	// New Server using RpcHandler
	pub fn with_rpc_handler(socket_addr: &str, io_handler: RpcHandler<M>) -> Result<Server<M>> {
		Ok(Server {
			io_handler: io_handler,
			is_stopping: Arc::new(AtomicBool::new(false)),
			is_stopped: Arc::new(AtomicBool::new(true)),
			rpc_event_loop: Mutex::new(None),
			addr: socket_addr.to_owned(),
		})
	}

	/// Run server (in this thread)
	pub fn run(&self) -> Result<()> {
		let mut pipe_handler = try!(PipeHandler::start(&self.addr, self.io_handler.clone()));
		loop  {
			try!(pipe_handler.handle_incoming(&self.addr, Arc::new(AtomicBool::new(false))));
		}
	}

	/// Run server (in separate thread)
	pub fn run_async(&self) -> Result<()> {
		if self.is_stopping.load(Ordering::Relaxed) { return Err(Error::IsStopping) }
		if !self.is_stopped.load(Ordering::Relaxed) { return Err(Error::NotStopped) }

		trace!(target: "ipc", "Started named pipes server [{}]", self.addr);

		let thread_stopping = self.is_stopping.clone();
		let thread_stopped = self.is_stopped.clone();
		let thread_handler = self.io_handler.clone();
		let addr = self.addr.clone();
		std::thread::spawn(move || {
			let mut pipe_handler = PipeHandler::start(&addr, thread_handler).unwrap();
			while !thread_stopping.load(Ordering::Relaxed) {
				trace!(target: "ipc", "Accepting pipe connection");
				if let Err(pipe_listener_error) = pipe_handler.handle_incoming(&addr, thread_stopping.clone()) {
					trace!(target: "ipc", "Pipe listening error: {:?}", pipe_listener_error);
				}
			}
			thread_stopped.store(true, Ordering::Relaxed);
		});

		self.is_stopped.store(false, Ordering::Relaxed);
		Ok(())
	}

	pub fn stop_async(&self) -> Result<()> {
		self.rpc_event_loop.lock().unwrap().take().map(|s| s.close());
		if self.is_stopped.load(Ordering::Relaxed) { return Err(Error::NotStarted) }
		if self.is_stopping.load(Ordering::Relaxed) { return Err(Error::AlreadyStopping)}
		self.is_stopping.store(true, Ordering::Relaxed);
		Ok(())
	}

	pub fn stop(&self) -> Result<()> {
		self.rpc_event_loop.lock().unwrap().take().map(|s| s.close());
		if self.is_stopped.load(Ordering::Relaxed) { return Err(Error::NotStarted) }
		if self.is_stopping.load(Ordering::Relaxed) { return Err(Error::AlreadyStopping)}
		self.is_stopping.store(true, Ordering::Relaxed);
		while !self.is_stopped.load(Ordering::Relaxed) { std::thread::park_timeout(std::time::Duration::new(0, 50)); }
		Ok(())
	}
}

impl<M: Metadata> Drop for Server<M> {
	fn drop(&mut self) {
		self.stop_async().unwrap_or_else(|_| {}); // ignore error - can be stopped already
		// todo : no stable logging for windows?
		trace!(target: "ipc", "IPC Server : shutdown");
	}
}
<|endoftext|>


<commit_before>// Copyright 2013-2015, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

//! Pack widgets in a rows and columns

// FIXME: Missings methods

use libc::{c_int, c_uint};

use {PositionType};
use cast::GTK_GRID;
use ffi;
use glib::{to_bool, to_gboolean};

/// Grid — Pack widgets in a rows and columns
struct_Widget!(Grid);

impl Grid {
    pub fn new() -> Option<Grid> {
        let tmp_pointer = unsafe { ffi::gtk_grid_new() };
        check_pointer!(tmp_pointer, Grid)
    }

    pub fn attach<T: ::WidgetTrait>(&self,
                                child: &T,
                                left: i32,
                                top: i32,
                                width: i32,
                                height: i32) -> () {
        unsafe {
            ffi::gtk_grid_attach(GTK_GRID(self.pointer),
                                 child.unwrap_widget(),
                                 left as c_int,
                                 top as c_int,
                                 width as c_int,
                                 height as c_int);
        }
    }

    pub fn attach_next_to<T: ::WidgetTrait, T2: ::WidgetTrait>(&self,
                                        child: &T,
                                        sibling: &T2,
                                        side: PositionType,
                                        width: i32,
                                        height: i32) -> () {
        unsafe {
            ffi::gtk_grid_attach_next_to(GTK_GRID(self.pointer),
                                         child.unwrap_widget(),
                                         sibling.unwrap_widget(),
                                         side,
                                         width as c_int,
                                         height as c_int);
        }
    }

    pub fn insert_row(&self, position: i32) -> () {
        unsafe {
            ffi::gtk_grid_insert_row(GTK_GRID(self.pointer), position as c_int);
        }
    }

     pub fn insert_column(&self, position: i32) -> () {
        unsafe {
            ffi::gtk_grid_insert_column(GTK_GRID(self.pointer), position as c_int);
        }
    }

    #[cfg(gtk_3_10)]
     pub fn remove_row(&self, position: i32) -> () {
        unsafe {
            ffi::gtk_grid_remove_row(GTK_GRID(self.pointer), position as c_int);
        }
    }

    #[cfg(gtk_3_10)]
     pub fn remove_column(&self, position: i32) -> () {
        unsafe {
            ffi::gtk_grid_remove_column(GTK_GRID(self.pointer), position as c_int);
        }
    }

    pub fn insert_next_to<T: ::WidgetTrait>(&self, sibling: &T, side: PositionType) -> () {
        unsafe {
            ffi::gtk_grid_insert_next_to(GTK_GRID(self.pointer), sibling.unwrap_widget(), side);
        }
    }

    pub fn set_row_homogeneous(&self, homogeneous: bool) -> () {
        unsafe { ffi::gtk_grid_set_row_homogeneous(GTK_GRID(self.pointer), to_gboolean(homogeneous)); }
    }

    pub fn get_row_homogeneous(&self) -> bool {
        unsafe { to_bool(ffi::gtk_grid_get_row_homogeneous(GTK_GRID(self.pointer))) }
    }

    pub fn set_row_spacing(&self, spacing: u32) -> () {
        unsafe {
            ffi::gtk_grid_set_row_spacing(GTK_GRID(self.pointer), spacing as c_uint);
        }
    }

    pub fn get_row_spacing(&self) -> u32 {
        unsafe {
            ffi::gtk_grid_get_row_spacing(GTK_GRID(self.pointer)) as u32
        }
    }

    pub fn set_column_homogeneous(&self, homogeneous: bool) -> () {
        unsafe { ffi::gtk_grid_set_column_homogeneous(GTK_GRID(self.pointer), to_gboolean(homogeneous)); }
    }

    pub fn get_column_homogeneous(&self) -> bool {
        unsafe { to_bool(ffi::gtk_grid_get_column_homogeneous(GTK_GRID(self.pointer))) }
    }

    pub fn set_column_spacing(&self, spacing: u32) -> () {
        unsafe {
            ffi::gtk_grid_set_column_spacing(GTK_GRID(self.pointer), spacing as c_uint);
        }
    }

    pub fn get_column_spacing(&self) -> u32 {
        unsafe {
            ffi::gtk_grid_get_column_spacing(GTK_GRID(self.pointer)) as u32
        }
    }

    #[cfg(gtk_3_10)]
    pub fn get_baseline_row(&self) -> i32 {
        unsafe {
            ffi::gtk_grid_get_baseline_row(GTK_GRID(self.pointer)) as i32
        }
    }

    #[cfg(gtk_3_10)]
    pub fn set_baseline_row(&self, row: i32) -> () {
        unsafe {
            ffi::gtk_grid_set_baseline_row(GTK_GRID(self.pointer), row as c_int);
        }
    }
}

impl_drop!(Grid);
impl_TraitWidget!(Grid);

impl ::ContainerTrait for Grid {}
impl ::OrientableTrait for Grid {}
<commit_msg>Added missing get_child_at function for Grid<commit_after>// Copyright 2013-2015, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

//! Pack widgets in a rows and columns

// FIXME: Missings methods

use libc::{c_int, c_uint};

use {PositionType};
use cast::GTK_GRID;
use ffi;
use glib::{to_bool, to_gboolean};

/// Grid — Pack widgets in a rows and columns
struct_Widget!(Grid);

impl Grid {
    pub fn new() -> Option<Grid> {
        let tmp_pointer = unsafe { ffi::gtk_grid_new() };
        check_pointer!(tmp_pointer, Grid)
    }

    pub fn attach<T: ::WidgetTrait>(&self,
                                child: &T,
                                left: i32,
                                top: i32,
                                width: i32,
                                height: i32) -> () {
        unsafe {
            ffi::gtk_grid_attach(GTK_GRID(self.pointer),
                                 child.unwrap_widget(),
                                 left as c_int,
                                 top as c_int,
                                 width as c_int,
                                 height as c_int);
        }
    }

    pub fn attach_next_to<T: ::WidgetTrait, T2: ::WidgetTrait>(&self,
                                        child: &T,
                                        sibling: &T2,
                                        side: PositionType,
                                        width: i32,
                                        height: i32) -> () {
        unsafe {
            ffi::gtk_grid_attach_next_to(GTK_GRID(self.pointer),
                                         child.unwrap_widget(),
                                         sibling.unwrap_widget(),
                                         side,
                                         width as c_int,
                                         height as c_int);
        }
    }

    pub fn get_child_at<T: ::WidgetTrait>(&self, left: i32, top: i32) -> Option<T> {
        let tmp_pointer = unsafe {
            ffi::gtk_grid_get_child_at(GTK_GRID(self.pointer), left, top)
        };

        if tmp_pointer.is_null() {
            None
        } else {
            Some(::FFIWidget::wrap_widget(tmp_pointer))
        }
    }

    pub fn insert_row(&self, position: i32) -> () {
        unsafe {
            ffi::gtk_grid_insert_row(GTK_GRID(self.pointer), position as c_int);
        }
    }

     pub fn insert_column(&self, position: i32) -> () {
        unsafe {
            ffi::gtk_grid_insert_column(GTK_GRID(self.pointer), position as c_int);
        }
    }

    #[cfg(gtk_3_10)]
     pub fn remove_row(&self, position: i32) -> () {
        unsafe {
            ffi::gtk_grid_remove_row(GTK_GRID(self.pointer), position as c_int);
        }
    }

    #[cfg(gtk_3_10)]
     pub fn remove_column(&self, position: i32) -> () {
        unsafe {
            ffi::gtk_grid_remove_column(GTK_GRID(self.pointer), position as c_int);
        }
    }

    pub fn insert_next_to<T: ::WidgetTrait>(&self, sibling: &T, side: PositionType) -> () {
        unsafe {
            ffi::gtk_grid_insert_next_to(GTK_GRID(self.pointer), sibling.unwrap_widget(), side);
        }
    }

    pub fn set_row_homogeneous(&self, homogeneous: bool) -> () {
        unsafe { ffi::gtk_grid_set_row_homogeneous(GTK_GRID(self.pointer), to_gboolean(homogeneous)); }
    }

    pub fn get_row_homogeneous(&self) -> bool {
        unsafe { to_bool(ffi::gtk_grid_get_row_homogeneous(GTK_GRID(self.pointer))) }
    }

    pub fn set_row_spacing(&self, spacing: u32) -> () {
        unsafe {
            ffi::gtk_grid_set_row_spacing(GTK_GRID(self.pointer), spacing as c_uint);
        }
    }

    pub fn get_row_spacing(&self) -> u32 {
        unsafe {
            ffi::gtk_grid_get_row_spacing(GTK_GRID(self.pointer)) as u32
        }
    }

    pub fn set_column_homogeneous(&self, homogeneous: bool) -> () {
        unsafe { ffi::gtk_grid_set_column_homogeneous(GTK_GRID(self.pointer), to_gboolean(homogeneous)); }
    }

    pub fn get_column_homogeneous(&self) -> bool {
        unsafe { to_bool(ffi::gtk_grid_get_column_homogeneous(GTK_GRID(self.pointer))) }
    }

    pub fn set_column_spacing(&self, spacing: u32) -> () {
        unsafe {
            ffi::gtk_grid_set_column_spacing(GTK_GRID(self.pointer), spacing as c_uint);
        }
    }

    pub fn get_column_spacing(&self) -> u32 {
        unsafe {
            ffi::gtk_grid_get_column_spacing(GTK_GRID(self.pointer)) as u32
        }
    }

    #[cfg(gtk_3_10)]
    pub fn get_baseline_row(&self) -> i32 {
        unsafe {
            ffi::gtk_grid_get_baseline_row(GTK_GRID(self.pointer)) as i32
        }
    }

    #[cfg(gtk_3_10)]
    pub fn set_baseline_row(&self, row: i32) -> () {
        unsafe {
            ffi::gtk_grid_set_baseline_row(GTK_GRID(self.pointer), row as c_int);
        }
    }
}

impl_drop!(Grid);
impl_TraitWidget!(Grid);

impl ::ContainerTrait for Grid {}
impl ::OrientableTrait for Grid {}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::io::{self, Write};
use std::net::Shutdown;
use std::rc::Rc;
use mio::net::TcpStream;
use mio::{Event, PollOpt, Ready, Token};

use super::binary;
use super::close_listener::CloseListener;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::ipv4_packet_buffer::IPv4PacketBuffer;
use super::router::Router;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;

const TAG: &'static str = "Client";

pub struct Client {
    id: u32,
    stream: TcpStream,
    token: Token,
    client_to_network: IPv4PacketBuffer,
    network_to_client: StreamBuffer,
    router: Router,
    close_listener: Box<CloseListener<Client>>,
    closed: bool,
    // number of remaining bytes of "id" to send to the client before relaying any data
    pending_id_bytes: usize,
}

impl Client {
    pub fn new(id: u32, selector: &mut Selector, stream: TcpStream, close_listener: Box<CloseListener<Client>>) -> io::Result<Rc<RefCell<Self>>> {
        let rc = Rc::new(RefCell::new(Self {
            id: id,
            stream: stream,
            token: Token(0), // default value, will be set afterwards
            client_to_network: IPv4PacketBuffer::new(),
            network_to_client: StreamBuffer::new(16 * MAX_PACKET_LENGTH),
            router: Router::new(),
            closed: false,
            close_listener: close_listener,
            pending_id_bytes: 4,
        }));
        // set client as router owner
        rc.borrow_mut().router.set_client(Rc::downgrade(&rc));

        {
            let rc_clone = rc.clone();
            let handler = move |selector: &mut Selector, ready| {
                let mut self_ref = rc_clone.borrow_mut();
                self_ref.on_ready(selector, ready);
            };
            let mut self_ref = rc.borrow_mut();
            // on start, we are interested only in writing (we must first send the client id)
            let token = selector.register(&self_ref.stream, handler, Ready::writable(), PollOpt::level())?;
            self_ref.token = token;
        }
        Ok(rc)
    }

    pub fn id(&self) -> u32 {
        return self.id;
    }

    pub fn router(&mut self) -> &mut Router {
        &mut self.router
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        selector.deregister(&self.stream, self.token);
        // shutdown only (there is no close), the socket will be closed on drop
        self.stream.shutdown(Shutdown::Both);
        self.router.clear();
        self.close_listener.on_closed(self);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        if self.must_send_id() {
            match self.send_id() {
                Ok(_) => {
                    if self.pending_id_bytes == 0 {
                        debug!("Client id #{} sent to client", self.id);
                    }
                }
                Err(_) => {
                    error!(target: TAG, "Cannot write client id #{}", self.id);
                    self.close(selector);
                }
            }
        } else {
            match self.write() {
                Ok(_) => self.process_pending(),
                Err(_) => {
                    error!(target: TAG, "Cannot write");
                    self.close(selector);
                }
            }
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        match self.read() {
            Ok(_) => self.push_to_network(selector),
            Err(_) => {
                error!(target: TAG, "Cannot read");
                self.close(selector);
            }
        }
    }

    pub fn send_to_client(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        if ipv4_packet.length() as usize <= self.network_to_client.remaining() {
            self.network_to_client.read_from(ipv4_packet.raw());
            self.update_interests(selector);
            Ok(())
        } else {
            warn!(target: TAG, "Client buffer full, delaying packet processing");
            Err(io::Error::new(io::ErrorKind::WouldBlock, "Client buffer full"))
        }
    }

    fn send_id(&mut self) -> io::Result<()> {
        assert!(self.must_send_id());
        let raw_id = binary::to_byte_array(self.id);
        let w = self.stream.write(&raw_id[4 - self.pending_id_bytes..])?;
        self.pending_id_bytes -= w;
        Ok(())
    }

    fn update_interests(&mut self, selector: &mut Selector) -> io::Result<()> {
        let ready = if self.network_to_client.is_empty() {
            Ready::readable()
        } else {
            Ready::readable() | Ready::writable()
        };
        selector.reregister(&self.stream, self.token, ready, PollOpt::level())
    }

    fn read(&mut self) -> io::Result<()> {
        self.client_to_network.read_from(&mut self.stream)
    }

    fn write(&mut self) -> io::Result<()> {
        self.network_to_client.write_to(&mut self.stream)?;
        Ok(())
    }

    fn push_to_network(&mut self, selector: &mut Selector) {
        while self.push_one_packet_to_network(selector) {
            self.client_to_network.next();
        }
    }

    fn push_one_packet_to_network(&mut self, selector: &mut Selector) -> bool {
        match self.client_to_network.as_ipv4_packet() {
            Some(ref packet) => {
                self.router.send_to_network(selector, packet);
                true
            }
            None => false
        }
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn process_pending(&mut self) {
        // TODO
    }

    fn clean_expired_connections() {
        // TODO
    }

    fn must_send_id(&self) -> bool{
        self.pending_id_bytes > 0
    }
}
<commit_msg>Add target in log<commit_after>use std::cell::RefCell;
use std::io::{self, Write};
use std::net::Shutdown;
use std::rc::Rc;
use mio::net::TcpStream;
use mio::{Event, PollOpt, Ready, Token};

use super::binary;
use super::close_listener::CloseListener;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::ipv4_packet_buffer::IPv4PacketBuffer;
use super::router::Router;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;

const TAG: &'static str = "Client";

pub struct Client {
    id: u32,
    stream: TcpStream,
    token: Token,
    client_to_network: IPv4PacketBuffer,
    network_to_client: StreamBuffer,
    router: Router,
    close_listener: Box<CloseListener<Client>>,
    closed: bool,
    // number of remaining bytes of "id" to send to the client before relaying any data
    pending_id_bytes: usize,
}

impl Client {
    pub fn new(id: u32, selector: &mut Selector, stream: TcpStream, close_listener: Box<CloseListener<Client>>) -> io::Result<Rc<RefCell<Self>>> {
        let rc = Rc::new(RefCell::new(Self {
            id: id,
            stream: stream,
            token: Token(0), // default value, will be set afterwards
            client_to_network: IPv4PacketBuffer::new(),
            network_to_client: StreamBuffer::new(16 * MAX_PACKET_LENGTH),
            router: Router::new(),
            closed: false,
            close_listener: close_listener,
            pending_id_bytes: 4,
        }));
        // set client as router owner
        rc.borrow_mut().router.set_client(Rc::downgrade(&rc));

        {
            let rc_clone = rc.clone();
            let handler = move |selector: &mut Selector, ready| {
                let mut self_ref = rc_clone.borrow_mut();
                self_ref.on_ready(selector, ready);
            };
            let mut self_ref = rc.borrow_mut();
            // on start, we are interested only in writing (we must first send the client id)
            let token = selector.register(&self_ref.stream, handler, Ready::writable(), PollOpt::level())?;
            self_ref.token = token;
        }
        Ok(rc)
    }

    pub fn id(&self) -> u32 {
        return self.id;
    }

    pub fn router(&mut self) -> &mut Router {
        &mut self.router
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        selector.deregister(&self.stream, self.token);
        // shutdown only (there is no close), the socket will be closed on drop
        self.stream.shutdown(Shutdown::Both);
        self.router.clear();
        self.close_listener.on_closed(self);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        if self.must_send_id() {
            match self.send_id() {
                Ok(_) => {
                    if self.pending_id_bytes == 0 {
                        debug!(target:TAG, "Client id #{} sent to client", self.id);
                    }
                }
                Err(_) => {
                    error!(target: TAG, "Cannot write client id #{}", self.id);
                    self.close(selector);
                }
            }
        } else {
            match self.write() {
                Ok(_) => self.process_pending(),
                Err(_) => {
                    error!(target: TAG, "Cannot write");
                    self.close(selector);
                }
            }
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        match self.read() {
            Ok(_) => self.push_to_network(selector),
            Err(_) => {
                error!(target: TAG, "Cannot read");
                self.close(selector);
            }
        }
    }

    pub fn send_to_client(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        if ipv4_packet.length() as usize <= self.network_to_client.remaining() {
            self.network_to_client.read_from(ipv4_packet.raw());
            self.update_interests(selector);
            Ok(())
        } else {
            warn!(target: TAG, "Client buffer full, delaying packet processing");
            Err(io::Error::new(io::ErrorKind::WouldBlock, "Client buffer full"))
        }
    }

    fn send_id(&mut self) -> io::Result<()> {
        assert!(self.must_send_id());
        let raw_id = binary::to_byte_array(self.id);
        let w = self.stream.write(&raw_id[4 - self.pending_id_bytes..])?;
        self.pending_id_bytes -= w;
        Ok(())
    }

    fn update_interests(&mut self, selector: &mut Selector) -> io::Result<()> {
        let ready = if self.network_to_client.is_empty() {
            Ready::readable()
        } else {
            Ready::readable() | Ready::writable()
        };
        selector.reregister(&self.stream, self.token, ready, PollOpt::level())
    }

    fn read(&mut self) -> io::Result<()> {
        self.client_to_network.read_from(&mut self.stream)
    }

    fn write(&mut self) -> io::Result<()> {
        self.network_to_client.write_to(&mut self.stream)?;
        Ok(())
    }

    fn push_to_network(&mut self, selector: &mut Selector) {
        while self.push_one_packet_to_network(selector) {
            self.client_to_network.next();
        }
    }

    fn push_one_packet_to_network(&mut self, selector: &mut Selector) -> bool {
        match self.client_to_network.as_ipv4_packet() {
            Some(ref packet) => {
                self.router.send_to_network(selector, packet);
                true
            }
            None => false
        }
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn process_pending(&mut self) {
        // TODO
    }

    fn clean_expired_connections() {
        // TODO
    }

    fn must_send_id(&self) -> bool{
        self.pending_id_bytes > 0
    }
}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// 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.

mod common;

use anyhow::Result;
use common::{
    create_partition_entry, create_primary_header, create_secondary_header,
};
use gpt_disk_io::{BlockIo, Disk, DiskError, MutSliceBlockIo};
use gpt_disk_types::{BlockSize, GptPartitionEntryArray};
use std::fs;
use std::path::Path;
use std::process::Command;
use tempfile::TempDir;
#[cfg(feature = "std")]
use {
    gpt_disk_io::StdBlockIo,
    std::fs::{File, OpenOptions},
};

fn test_disk_read<Io>(block_io: Io) -> Result<(), DiskError<Io::Error>>
where
    Io: BlockIo,
{
    let mut block_buf = vec![0u8; 512];
    let mut disk = Disk::new(block_io)?;

    let primary_header = disk.read_primary_gpt_header(&mut block_buf)?;
    assert_eq!(primary_header, create_primary_header());

    let secondary_header = disk.read_secondary_gpt_header(&mut block_buf)?;
    assert_eq!(secondary_header, create_secondary_header());

    let expected_partition_entry = create_partition_entry();

    // Check the primary partition entry array.
    let primary_partition_entry = disk
        .gpt_partition_entry_array_iter(
            primary_header.get_partition_entry_array_layout().unwrap(),
            &mut block_buf,
        )?
        .next()
        .unwrap()?;
    assert_eq!(primary_partition_entry, expected_partition_entry);

    // Check the secondary partition entry array.
    let second_partition_entry = disk
        .gpt_partition_entry_array_iter(
            primary_header.get_partition_entry_array_layout().unwrap(),
            &mut block_buf,
        )?
        .next()
        .unwrap()?;
    assert_eq!(second_partition_entry, expected_partition_entry);

    Ok(())
}

fn test_disk_write<Io>(block_io: Io) -> Result<(), DiskError<Io::Error>>
where
    Io: BlockIo,
{
    let bs = BlockSize::B512;
    let mut block_buf = vec![0u8; bs.to_usize().unwrap()];
    let mut disk = Disk::new(block_io)?;

    let primary_header = create_primary_header();
    let secondary_header = create_secondary_header();
    let partition_entry = create_partition_entry();

    disk.write_protective_mbr(&mut block_buf)?;
    disk.write_primary_gpt_header(&primary_header, &mut block_buf)?;
    disk.write_secondary_gpt_header(&secondary_header, &mut block_buf)?;

    let layout = primary_header.get_partition_entry_array_layout().unwrap();
    let mut bytes =
        vec![0; layout.num_bytes_rounded_to_block_as_usize(bs).unwrap()];
    let mut entry_array =
        GptPartitionEntryArray::new(layout, bs, &mut bytes).unwrap();
    *entry_array.get_partition_entry_mut(0).unwrap() = partition_entry;
    disk.write_gpt_partition_entry_array(&entry_array)?;

    entry_array.set_start_lba(secondary_header.partition_entry_lba.into());
    disk.write_gpt_partition_entry_array(&entry_array)?;

    disk.flush()?;

    Ok(())
}

fn run_cmd(cmd: &mut Command) -> Result<()> {
    let o = cmd.output()?;
    assert!(o.status.success());
    Ok(())
}

fn create_empty_file(path: &Path, size: &str) -> Result<()> {
    run_cmd(Command::new("truncate").args(&["--size", size]).arg(path))
}

fn test_with_mut_slice(sgdisk_path: &Path) -> Result<()> {
    // Test read.
    let mut contents = fs::read(&sgdisk_path)?;
    test_disk_read(MutSliceBlockIo::new(&mut contents, BlockSize::B512))
        .unwrap();

    // Test write.
    let mut new_contents = vec![0; contents.len()];
    test_disk_write(MutSliceBlockIo::new(&mut new_contents, BlockSize::B512))
        .unwrap();
    assert_eq!(contents, new_contents);

    Ok(())
}

#[cfg(feature = "std")]
fn test_with_file(tmp_path: &Path, sgdisk_path: &Path) -> Result<()> {
    // Test read.
    let mut file = File::open(&sgdisk_path)?;
    test_disk_read(StdBlockIo::new(&mut file, BlockSize::B512))?;

    // Test write.
    let new_disk_path = tmp_path.join("new_disk.bin");
    create_empty_file(&new_disk_path, "4MiB")?;
    let mut new_file = OpenOptions::new()
        .read(true)
        .write(true)
        .truncate(false)
        .open(&new_disk_path)?;
    test_disk_write(StdBlockIo::new(&mut new_file, BlockSize::B512)).unwrap();
    let expected_bytes = fs::read(&sgdisk_path)?;
    let actual_bytes = fs::read(&new_disk_path)?;
    assert_eq!(expected_bytes, actual_bytes);

    Ok(())
}

#[test]
#[cfg_attr(miri, ignore)]
fn test_disk() -> Result<()> {
    // If sgdisk is not available, silently skip the test. TODO: ensure
    // that the test is not skipped in CI.
    if Command::new("sgdisk").arg("--version").status().is_err() {
        return Ok(());
    }

    let tmp_dir = TempDir::new()?;
    let tmp_path = tmp_dir.path();

    let sgdisk_path = tmp_path.join("disk.bin");
    create_empty_file(&sgdisk_path, "4MiB")?;
    run_cmd(Command::new("sgdisk").arg(&sgdisk_path).args(&[
        "--disk-guid=57a7feb6-8cd5-4922-b7bd-c78b0914e870",
        "--new=1:2048:4096",
        "--change-name=1:hello world!",
        "--partition-guid=1:37c75ffd-8932-467a-9c56-8cf1f0456b12",
        "--typecode=1:ccf0994f-f7e0-4e26-a011-843e38aa2eac",
    ]))?;

    test_with_mut_slice(&sgdisk_path)?;

    #[cfg(feature = "std")]
    test_with_file(&tmp_path, &sgdisk_path)?;

    Ok(())
}
<commit_msg>Require sgdisk for running the tests<commit_after>// Copyright 2022 Google LLC
//
// 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.

mod common;

use anyhow::Result;
use common::{
    create_partition_entry, create_primary_header, create_secondary_header,
};
use gpt_disk_io::{BlockIo, Disk, DiskError, MutSliceBlockIo};
use gpt_disk_types::{BlockSize, GptPartitionEntryArray};
use std::fs;
use std::path::Path;
use std::process::Command;
use tempfile::TempDir;
#[cfg(feature = "std")]
use {
    gpt_disk_io::StdBlockIo,
    std::fs::{File, OpenOptions},
};

fn test_disk_read<Io>(block_io: Io) -> Result<(), DiskError<Io::Error>>
where
    Io: BlockIo,
{
    let mut block_buf = vec![0u8; 512];
    let mut disk = Disk::new(block_io)?;

    let primary_header = disk.read_primary_gpt_header(&mut block_buf)?;
    assert_eq!(primary_header, create_primary_header());

    let secondary_header = disk.read_secondary_gpt_header(&mut block_buf)?;
    assert_eq!(secondary_header, create_secondary_header());

    let expected_partition_entry = create_partition_entry();

    // Check the primary partition entry array.
    let primary_partition_entry = disk
        .gpt_partition_entry_array_iter(
            primary_header.get_partition_entry_array_layout().unwrap(),
            &mut block_buf,
        )?
        .next()
        .unwrap()?;
    assert_eq!(primary_partition_entry, expected_partition_entry);

    // Check the secondary partition entry array.
    let second_partition_entry = disk
        .gpt_partition_entry_array_iter(
            primary_header.get_partition_entry_array_layout().unwrap(),
            &mut block_buf,
        )?
        .next()
        .unwrap()?;
    assert_eq!(second_partition_entry, expected_partition_entry);

    Ok(())
}

fn test_disk_write<Io>(block_io: Io) -> Result<(), DiskError<Io::Error>>
where
    Io: BlockIo,
{
    let bs = BlockSize::B512;
    let mut block_buf = vec![0u8; bs.to_usize().unwrap()];
    let mut disk = Disk::new(block_io)?;

    let primary_header = create_primary_header();
    let secondary_header = create_secondary_header();
    let partition_entry = create_partition_entry();

    disk.write_protective_mbr(&mut block_buf)?;
    disk.write_primary_gpt_header(&primary_header, &mut block_buf)?;
    disk.write_secondary_gpt_header(&secondary_header, &mut block_buf)?;

    let layout = primary_header.get_partition_entry_array_layout().unwrap();
    let mut bytes =
        vec![0; layout.num_bytes_rounded_to_block_as_usize(bs).unwrap()];
    let mut entry_array =
        GptPartitionEntryArray::new(layout, bs, &mut bytes).unwrap();
    *entry_array.get_partition_entry_mut(0).unwrap() = partition_entry;
    disk.write_gpt_partition_entry_array(&entry_array)?;

    entry_array.set_start_lba(secondary_header.partition_entry_lba.into());
    disk.write_gpt_partition_entry_array(&entry_array)?;

    disk.flush()?;

    Ok(())
}

fn run_cmd(cmd: &mut Command) -> Result<()> {
    let o = cmd.output()?;
    assert!(o.status.success());
    Ok(())
}

fn create_empty_file(path: &Path, size: &str) -> Result<()> {
    run_cmd(Command::new("truncate").args(&["--size", size]).arg(path))
}

fn test_with_mut_slice(sgdisk_path: &Path) -> Result<()> {
    // Test read.
    let mut contents = fs::read(&sgdisk_path)?;
    test_disk_read(MutSliceBlockIo::new(&mut contents, BlockSize::B512))
        .unwrap();

    // Test write.
    let mut new_contents = vec![0; contents.len()];
    test_disk_write(MutSliceBlockIo::new(&mut new_contents, BlockSize::B512))
        .unwrap();
    assert_eq!(contents, new_contents);

    Ok(())
}

#[cfg(feature = "std")]
fn test_with_file(tmp_path: &Path, sgdisk_path: &Path) -> Result<()> {
    // Test read.
    let mut file = File::open(&sgdisk_path)?;
    test_disk_read(StdBlockIo::new(&mut file, BlockSize::B512))?;

    // Test write.
    let new_disk_path = tmp_path.join("new_disk.bin");
    create_empty_file(&new_disk_path, "4MiB")?;
    let mut new_file = OpenOptions::new()
        .read(true)
        .write(true)
        .truncate(false)
        .open(&new_disk_path)?;
    test_disk_write(StdBlockIo::new(&mut new_file, BlockSize::B512)).unwrap();
    let expected_bytes = fs::read(&sgdisk_path)?;
    let actual_bytes = fs::read(&new_disk_path)?;
    assert_eq!(expected_bytes, actual_bytes);

    Ok(())
}

#[test]
#[cfg_attr(miri, ignore)]
fn test_disk() -> Result<()> {
    if Command::new("sgdisk").arg("--version").status().is_err() {
        panic!("failed to run sgdisk, is it installed?");
    }

    let tmp_dir = TempDir::new()?;
    let tmp_path = tmp_dir.path();

    let sgdisk_path = tmp_path.join("disk.bin");
    create_empty_file(&sgdisk_path, "4MiB")?;
    run_cmd(Command::new("sgdisk").arg(&sgdisk_path).args(&[
        "--disk-guid=57a7feb6-8cd5-4922-b7bd-c78b0914e870",
        "--new=1:2048:4096",
        "--change-name=1:hello world!",
        "--partition-guid=1:37c75ffd-8932-467a-9c56-8cf1f0456b12",
        "--typecode=1:ccf0994f-f7e0-4e26-a011-843e38aa2eac",
    ]))?;

    test_with_mut_slice(&sgdisk_path)?;

    #[cfg(feature = "std")]
    test_with_file(&tmp_path, &sgdisk_path)?;

    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2013 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.

//! Composable iterator objects

use prelude::*;

pub trait Iterator<A> {
    /// Advance the iterator and return the next value. Return `None` when the end is reached.
    fn next(&mut self) -> Option<A>;
}

pub trait IteratorUtil<A> {
    fn chain(self, other: Self) -> ChainIterator<Self>;
    fn zip<B, U: Iterator<B>>(self, other: U) -> ZipIterator<Self, U>;
    // FIXME: #5898: should be called map
    fn transform<'r, B>(self, f: &'r fn(A) -> B) -> MapIterator<'r, A, B, Self>;
    fn filter<'r>(self, predicate: &'r fn(&A) -> bool) -> FilterIterator<'r, A, Self>;
    fn enumerate(self) -> EnumerateIterator<Self>;
    fn skip_while<'r>(self, predicate: &'r fn(&A) -> bool) -> SkipWhileIterator<'r, A, Self>;
    fn take_while<'r>(self, predicate: &'r fn(&A) -> bool) -> TakeWhileIterator<'r, A, Self>;
    fn skip(self, n: uint) -> SkipIterator<Self>;
    fn take(self, n: uint) -> TakeIterator<Self>;
    fn advance(&mut self, f: &fn(A) -> bool);
}

impl<A, T: Iterator<A>> IteratorUtil<A> for T {
    #[inline(always)]
    fn chain(self, other: T) -> ChainIterator<T> {
        ChainIterator{a: self, b: other, flag: false}
    }

    #[inline(always)]
    fn zip<B, U: Iterator<B>>(self, other: U) -> ZipIterator<T, U> {
        ZipIterator{a: self, b: other}
    }

    // FIXME: #5898: should be called map
    #[inline(always)]
    fn transform<'r, B>(self, f: &'r fn(A) -> B) -> MapIterator<'r, A, B, T> {
        MapIterator{iter: self, f: f}
    }

    #[inline(always)]
    fn filter<'r>(self, predicate: &'r fn(&A) -> bool) -> FilterIterator<'r, A, T> {
        FilterIterator{iter: self, predicate: predicate}
    }

    #[inline(always)]
    fn enumerate(self) -> EnumerateIterator<T> {
        EnumerateIterator{iter: self, count: 0}
    }

    #[inline(always)]
    fn skip_while<'r>(self, predicate: &'r fn(&A) -> bool) -> SkipWhileIterator<'r, A, T> {
        SkipWhileIterator{iter: self, flag: false, predicate: predicate}
    }

    #[inline(always)]
    fn take_while<'r>(self, predicate: &'r fn(&A) -> bool) -> TakeWhileIterator<'r, A, T> {
        TakeWhileIterator{iter: self, flag: false, predicate: predicate}
    }

    #[inline(always)]
    fn skip(self, n: uint) -> SkipIterator<T> {
        SkipIterator{iter: self, n: n}
    }

    #[inline(always)]
    fn take(self, n: uint) -> TakeIterator<T> {
        TakeIterator{iter: self, n: n}
    }

    /// A shim implementing the `for` loop iteration protocol for iterator objects
    #[inline]
    fn advance(&mut self, f: &fn(A) -> bool) {
        loop {
            match self.next() {
                Some(x) => {
                    if !f(x) { return }
                }
                None => return
            }
        }
    }
}

pub struct ChainIterator<T> {
    priv a: T,
    priv b: T,
    priv flag: bool
}

impl<A, T: Iterator<A>> Iterator<A> for ChainIterator<T> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        if self.flag {
            self.b.next()
        } else {
            match self.a.next() {
                Some(x) => return Some(x),
                _ => ()
            }
            self.flag = true;
            self.b.next()
        }
    }
}

pub struct ZipIterator<T, U> {
    priv a: T,
    priv b: U
}

impl<A, B, T: Iterator<A>, U: Iterator<B>> Iterator<(A, B)> for ZipIterator<T, U> {
    #[inline]
    fn next(&mut self) -> Option<(A, B)> {
        match (self.a.next(), self.b.next()) {
            (Some(x), Some(y)) => Some((x, y)),
            _ => None
        }
    }
}

pub struct MapIterator<'self, A, B, T> {
    priv iter: T,
    priv f: &'self fn(A) -> B
}

impl<'self, A, B, T: Iterator<A>> Iterator<B> for MapIterator<'self, A, B, T> {
    #[inline]
    fn next(&mut self) -> Option<B> {
        match self.iter.next() {
            Some(a) => Some((self.f)(a)),
            _ => None
        }
    }
}

pub struct FilterIterator<'self, A, T> {
    priv iter: T,
    priv predicate: &'self fn(&A) -> bool
}

impl<'self, A, T: Iterator<A>> Iterator<A> for FilterIterator<'self, A, T> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        for self.iter.advance |x| {
            if (self.predicate)(&x) {
                return Some(x);
            } else {
                loop
            }
        }
        None
    }
}

pub struct EnumerateIterator<T> {
    priv iter: T,
    priv count: uint
}

impl<A, T: Iterator<A>> Iterator<(uint, A)> for EnumerateIterator<T> {
    #[inline]
    fn next(&mut self) -> Option<(uint, A)> {
        match self.iter.next() {
            Some(a) => {
                let ret = Some((self.count, a));
                self.count += 1;
                ret
            }
            _ => None
        }
    }
}

pub struct SkipWhileIterator<'self, A, T> {
    priv iter: T,
    priv flag: bool,
    priv predicate: &'self fn(&A) -> bool
}

impl<'self, A, T: Iterator<A>> Iterator<A> for SkipWhileIterator<'self, A, T> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        let mut next = self.iter.next();
        if self.flag {
            next
        } else {
            loop {
                match next {
                    Some(x) => {
                        if (self.predicate)(&x) {
                            next = self.iter.next();
                            loop
                        } else {
                            self.flag = true;
                            return Some(x)
                        }
                    }
                    None => return None
                }
            }
        }
    }
}

pub struct TakeWhileIterator<'self, A, T> {
    priv iter: T,
    priv flag: bool,
    priv predicate: &'self fn(&A) -> bool
}

impl<'self, A, T: Iterator<A>> Iterator<A> for TakeWhileIterator<'self, A, T> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        if self.flag {
            None
        } else {
            match self.iter.next() {
                Some(x) => {
                    if (self.predicate)(&x) {
                        Some(x)
                    } else {
                        self.flag = true;
                        None
                    }
                }
                None => None
            }
        }
    }
}

pub struct SkipIterator<T> {
    priv iter: T,
    priv n: uint
}

impl<A, T: Iterator<A>> Iterator<A> for SkipIterator<T> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        let mut next = self.iter.next();
        if self.n == 0 {
            next
        } else {
            let n = self.n;
            for n.times {
                match next {
                    Some(_) => {
                        next = self.iter.next();
                        loop
                    }
                    None => {
                        self.n = 0;
                        return None
                    }
                }
            }
            self.n = 0;
            next
        }
    }
}

pub struct TakeIterator<T> {
    priv iter: T,
    priv n: uint
}

impl<A, T: Iterator<A>> Iterator<A> for TakeIterator<T> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        let next = self.iter.next();
        if self.n != 0 {
            self.n -= 1;
            next
        } else {
            None
        }
    }
}

pub struct UnfoldrIterator<'self, A, St> {
    priv f: &'self fn(&mut St) -> Option<A>,
    priv state: St
}

pub impl<'self, A, St> UnfoldrIterator<'self, A, St> {
    #[inline]
    fn new(f: &'self fn(&mut St) -> Option<A>, initial_state: St) -> UnfoldrIterator<'self, A, St> {
        UnfoldrIterator {
            f: f,
            state: initial_state
        }
    }
}

impl<'self, A, St> Iterator<A> for UnfoldrIterator<'self, A, St> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        (self.f)(&mut self.state)
    }
}

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

    #[test]
    fn test_iterator_chain() {
        let xs = [0u, 1, 2, 3, 4, 5];
        let ys = [30, 40, 50, 60];
        let expected = [0, 1, 2, 3, 4, 5, 30, 40, 50, 60];
        let mut it = xs.iter().chain(ys.iter());
        let mut i = 0;
        for it.advance |&x: &uint| {
            assert_eq!(x, expected[i]);
            i += 1;
        }
        assert_eq!(i, expected.len());
    }

    #[test]
    fn test_iterator_enumerate() {
        let xs = [0u, 1, 2, 3, 4, 5];
        let mut it = xs.iter().enumerate();
        for it.advance |(i, &x): (uint, &uint)| {
            assert_eq!(i, x);
        }
    }

    #[test]
    fn test_iterator_take_while() {
        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19];
        let ys = [0u, 1, 2, 3, 5, 13];
        let mut it = xs.iter().take_while(|&x| *x < 15u);
        let mut i = 0;
        for it.advance |&x: &uint| {
            assert_eq!(x, ys[i]);
            i += 1;
        }
        assert_eq!(i, ys.len());
    }

    #[test]
    fn test_iterator_skip_while() {
        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19];
        let ys = [15, 16, 17, 19];
        let mut it = xs.iter().skip_while(|&x| *x < 15u);
        let mut i = 0;
        for it.advance |&x: &uint| {
            assert_eq!(x, ys[i]);
            i += 1;
        }
        assert_eq!(i, ys.len());
    }

    #[test]
    fn test_iterator_skip() {
        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19, 20, 30];
        let ys = [13, 15, 16, 17, 19, 20, 30];
        let mut it = xs.iter().skip(5);
        let mut i = 0;
        for it.advance |&x: &uint| {
            assert_eq!(x, ys[i]);
            i += 1;
        }
        assert_eq!(i, ys.len());
    }

    #[test]
    fn test_iterator_take() {
        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19];
        let ys = [0u, 1, 2, 3, 5];
        let mut it = xs.iter().take(5);
        let mut i = 0;
        for it.advance |&x: &uint| {
            assert_eq!(x, ys[i]);
            i += 1;
        }
        assert_eq!(i, ys.len());
    }

    #[test]
    fn test_unfoldr() {
        fn count(st: &mut uint) -> Option<uint> {
            if *st < 10 {
                let ret = Some(*st);
                *st += 1;
                ret
            } else {
                None
            }
        }

        let mut it = UnfoldrIterator::new(count, 0);
        let mut i = 0;
        for it.advance |counted| {
            assert_eq!(counted, i);
            i += 1;
        }
        assert_eq!(i, 10);
    }
}
<commit_msg>iterator: add a bit of documentation<commit_after>// Copyright 2013 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.

/*! Composable external iterators

The `Iterator` trait defines an interface for objects which implement iteration as a state machine.

Algorithms like `zip` are provided as `Iterator` implementations which wrap other objects
implementing the `Iterator` trait.

*/

use prelude::*;

pub trait Iterator<A> {
    /// Advance the iterator and return the next value. Return `None` when the end is reached.
    fn next(&mut self) -> Option<A>;
}

/// Iterator adaptors provided for every `Iterator` implementation. The adaptor objects are also
/// implementations of the `Iterator` trait.
///
/// In the future these will be default methods instead of a utility trait.
pub trait IteratorUtil<A> {
    fn chain(self, other: Self) -> ChainIterator<Self>;
    fn zip<B, U: Iterator<B>>(self, other: U) -> ZipIterator<Self, U>;
    // FIXME: #5898: should be called map
    fn transform<'r, B>(self, f: &'r fn(A) -> B) -> MapIterator<'r, A, B, Self>;
    fn filter<'r>(self, predicate: &'r fn(&A) -> bool) -> FilterIterator<'r, A, Self>;
    fn enumerate(self) -> EnumerateIterator<Self>;
    fn skip_while<'r>(self, predicate: &'r fn(&A) -> bool) -> SkipWhileIterator<'r, A, Self>;
    fn take_while<'r>(self, predicate: &'r fn(&A) -> bool) -> TakeWhileIterator<'r, A, Self>;
    fn skip(self, n: uint) -> SkipIterator<Self>;
    fn take(self, n: uint) -> TakeIterator<Self>;
    fn advance(&mut self, f: &fn(A) -> bool);
}

/// Iterator adaptors provided for every `Iterator` implementation. The adaptor objects are also
/// implementations of the `Iterator` trait.
///
/// In the future these will be default methods instead of a utility trait.
impl<A, T: Iterator<A>> IteratorUtil<A> for T {
    #[inline(always)]
    fn chain(self, other: T) -> ChainIterator<T> {
        ChainIterator{a: self, b: other, flag: false}
    }

    #[inline(always)]
    fn zip<B, U: Iterator<B>>(self, other: U) -> ZipIterator<T, U> {
        ZipIterator{a: self, b: other}
    }

    // FIXME: #5898: should be called map
    #[inline(always)]
    fn transform<'r, B>(self, f: &'r fn(A) -> B) -> MapIterator<'r, A, B, T> {
        MapIterator{iter: self, f: f}
    }

    #[inline(always)]
    fn filter<'r>(self, predicate: &'r fn(&A) -> bool) -> FilterIterator<'r, A, T> {
        FilterIterator{iter: self, predicate: predicate}
    }

    #[inline(always)]
    fn enumerate(self) -> EnumerateIterator<T> {
        EnumerateIterator{iter: self, count: 0}
    }

    #[inline(always)]
    fn skip_while<'r>(self, predicate: &'r fn(&A) -> bool) -> SkipWhileIterator<'r, A, T> {
        SkipWhileIterator{iter: self, flag: false, predicate: predicate}
    }

    #[inline(always)]
    fn take_while<'r>(self, predicate: &'r fn(&A) -> bool) -> TakeWhileIterator<'r, A, T> {
        TakeWhileIterator{iter: self, flag: false, predicate: predicate}
    }

    #[inline(always)]
    fn skip(self, n: uint) -> SkipIterator<T> {
        SkipIterator{iter: self, n: n}
    }

    #[inline(always)]
    fn take(self, n: uint) -> TakeIterator<T> {
        TakeIterator{iter: self, n: n}
    }

    /// A shim implementing the `for` loop iteration protocol for iterator objects
    #[inline]
    fn advance(&mut self, f: &fn(A) -> bool) {
        loop {
            match self.next() {
                Some(x) => {
                    if !f(x) { return }
                }
                None => return
            }
        }
    }
}

pub struct ChainIterator<T> {
    priv a: T,
    priv b: T,
    priv flag: bool
}

impl<A, T: Iterator<A>> Iterator<A> for ChainIterator<T> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        if self.flag {
            self.b.next()
        } else {
            match self.a.next() {
                Some(x) => return Some(x),
                _ => ()
            }
            self.flag = true;
            self.b.next()
        }
    }
}

pub struct ZipIterator<T, U> {
    priv a: T,
    priv b: U
}

impl<A, B, T: Iterator<A>, U: Iterator<B>> Iterator<(A, B)> for ZipIterator<T, U> {
    #[inline]
    fn next(&mut self) -> Option<(A, B)> {
        match (self.a.next(), self.b.next()) {
            (Some(x), Some(y)) => Some((x, y)),
            _ => None
        }
    }
}

pub struct MapIterator<'self, A, B, T> {
    priv iter: T,
    priv f: &'self fn(A) -> B
}

impl<'self, A, B, T: Iterator<A>> Iterator<B> for MapIterator<'self, A, B, T> {
    #[inline]
    fn next(&mut self) -> Option<B> {
        match self.iter.next() {
            Some(a) => Some((self.f)(a)),
            _ => None
        }
    }
}

pub struct FilterIterator<'self, A, T> {
    priv iter: T,
    priv predicate: &'self fn(&A) -> bool
}

impl<'self, A, T: Iterator<A>> Iterator<A> for FilterIterator<'self, A, T> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        for self.iter.advance |x| {
            if (self.predicate)(&x) {
                return Some(x);
            } else {
                loop
            }
        }
        None
    }
}

pub struct EnumerateIterator<T> {
    priv iter: T,
    priv count: uint
}

impl<A, T: Iterator<A>> Iterator<(uint, A)> for EnumerateIterator<T> {
    #[inline]
    fn next(&mut self) -> Option<(uint, A)> {
        match self.iter.next() {
            Some(a) => {
                let ret = Some((self.count, a));
                self.count += 1;
                ret
            }
            _ => None
        }
    }
}

pub struct SkipWhileIterator<'self, A, T> {
    priv iter: T,
    priv flag: bool,
    priv predicate: &'self fn(&A) -> bool
}

impl<'self, A, T: Iterator<A>> Iterator<A> for SkipWhileIterator<'self, A, T> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        let mut next = self.iter.next();
        if self.flag {
            next
        } else {
            loop {
                match next {
                    Some(x) => {
                        if (self.predicate)(&x) {
                            next = self.iter.next();
                            loop
                        } else {
                            self.flag = true;
                            return Some(x)
                        }
                    }
                    None => return None
                }
            }
        }
    }
}

pub struct TakeWhileIterator<'self, A, T> {
    priv iter: T,
    priv flag: bool,
    priv predicate: &'self fn(&A) -> bool
}

impl<'self, A, T: Iterator<A>> Iterator<A> for TakeWhileIterator<'self, A, T> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        if self.flag {
            None
        } else {
            match self.iter.next() {
                Some(x) => {
                    if (self.predicate)(&x) {
                        Some(x)
                    } else {
                        self.flag = true;
                        None
                    }
                }
                None => None
            }
        }
    }
}

pub struct SkipIterator<T> {
    priv iter: T,
    priv n: uint
}

impl<A, T: Iterator<A>> Iterator<A> for SkipIterator<T> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        let mut next = self.iter.next();
        if self.n == 0 {
            next
        } else {
            let n = self.n;
            for n.times {
                match next {
                    Some(_) => {
                        next = self.iter.next();
                        loop
                    }
                    None => {
                        self.n = 0;
                        return None
                    }
                }
            }
            self.n = 0;
            next
        }
    }
}

pub struct TakeIterator<T> {
    priv iter: T,
    priv n: uint
}

impl<A, T: Iterator<A>> Iterator<A> for TakeIterator<T> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        let next = self.iter.next();
        if self.n != 0 {
            self.n -= 1;
            next
        } else {
            None
        }
    }
}

pub struct UnfoldrIterator<'self, A, St> {
    priv f: &'self fn(&mut St) -> Option<A>,
    priv state: St
}

pub impl<'self, A, St> UnfoldrIterator<'self, A, St> {
    #[inline]
    fn new(f: &'self fn(&mut St) -> Option<A>, initial_state: St) -> UnfoldrIterator<'self, A, St> {
        UnfoldrIterator {
            f: f,
            state: initial_state
        }
    }
}

impl<'self, A, St> Iterator<A> for UnfoldrIterator<'self, A, St> {
    #[inline]
    fn next(&mut self) -> Option<A> {
        (self.f)(&mut self.state)
    }
}

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

    #[test]
    fn test_iterator_chain() {
        let xs = [0u, 1, 2, 3, 4, 5];
        let ys = [30, 40, 50, 60];
        let expected = [0, 1, 2, 3, 4, 5, 30, 40, 50, 60];
        let mut it = xs.iter().chain(ys.iter());
        let mut i = 0;
        for it.advance |&x: &uint| {
            assert_eq!(x, expected[i]);
            i += 1;
        }
        assert_eq!(i, expected.len());
    }

    #[test]
    fn test_iterator_enumerate() {
        let xs = [0u, 1, 2, 3, 4, 5];
        let mut it = xs.iter().enumerate();
        for it.advance |(i, &x): (uint, &uint)| {
            assert_eq!(i, x);
        }
    }

    #[test]
    fn test_iterator_take_while() {
        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19];
        let ys = [0u, 1, 2, 3, 5, 13];
        let mut it = xs.iter().take_while(|&x| *x < 15u);
        let mut i = 0;
        for it.advance |&x: &uint| {
            assert_eq!(x, ys[i]);
            i += 1;
        }
        assert_eq!(i, ys.len());
    }

    #[test]
    fn test_iterator_skip_while() {
        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19];
        let ys = [15, 16, 17, 19];
        let mut it = xs.iter().skip_while(|&x| *x < 15u);
        let mut i = 0;
        for it.advance |&x: &uint| {
            assert_eq!(x, ys[i]);
            i += 1;
        }
        assert_eq!(i, ys.len());
    }

    #[test]
    fn test_iterator_skip() {
        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19, 20, 30];
        let ys = [13, 15, 16, 17, 19, 20, 30];
        let mut it = xs.iter().skip(5);
        let mut i = 0;
        for it.advance |&x: &uint| {
            assert_eq!(x, ys[i]);
            i += 1;
        }
        assert_eq!(i, ys.len());
    }

    #[test]
    fn test_iterator_take() {
        let xs = [0u, 1, 2, 3, 5, 13, 15, 16, 17, 19];
        let ys = [0u, 1, 2, 3, 5];
        let mut it = xs.iter().take(5);
        let mut i = 0;
        for it.advance |&x: &uint| {
            assert_eq!(x, ys[i]);
            i += 1;
        }
        assert_eq!(i, ys.len());
    }

    #[test]
    fn test_unfoldr() {
        fn count(st: &mut uint) -> Option<uint> {
            if *st < 10 {
                let ret = Some(*st);
                *st += 1;
                ret
            } else {
                None
            }
        }

        let mut it = UnfoldrIterator::new(count, 0);
        let mut i = 0;
        for it.advance |counted| {
            assert_eq!(counted, i);
            i += 1;
        }
        assert_eq!(i, 10);
    }
}
<|endoftext|>


<commit_before>//! Abstractions for default-sized and huge virtual memory pages.

use crate::structures::paging::page_table::PageTableLevel;
use crate::structures::paging::PageTableIndex;
use crate::VirtAddr;
use core::fmt;
#[cfg(feature = "step_trait")]
use core::iter::Step;
use core::marker::PhantomData;
use core::ops::{Add, AddAssign, Sub, SubAssign};

/// Trait for abstracting over the three possible page sizes on x86_64, 4KiB, 2MiB, 1GiB.
pub trait PageSize: Copy + Eq + PartialOrd + Ord {
    /// The page size in bytes.
    const SIZE: u64;

    /// A string representation of the page size for debug output.
    const SIZE_AS_DEBUG_STR: &'static str;
}

/// This trait is implemented for 4KiB and 2MiB pages, but not for 1GiB pages.
pub trait NotGiantPageSize: PageSize {}

/// A standard 4KiB page.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size4KiB {}

/// A “huge” 2MiB page.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size2MiB {}

/// A “giant” 1GiB page.
///
/// (Only available on newer x86_64 CPUs.)
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size1GiB {}

impl PageSize for Size4KiB {
    const SIZE: u64 = 4096;
    const SIZE_AS_DEBUG_STR: &'static str = "4KiB";
}

impl NotGiantPageSize for Size4KiB {}

impl PageSize for Size2MiB {
    const SIZE: u64 = Size4KiB::SIZE * 512;
    const SIZE_AS_DEBUG_STR: &'static str = "2MiB";
}

impl NotGiantPageSize for Size2MiB {}

impl PageSize for Size1GiB {
    const SIZE: u64 = Size2MiB::SIZE * 512;
    const SIZE_AS_DEBUG_STR: &'static str = "1GiB";
}

/// A virtual memory page.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(C)]
pub struct Page<S: PageSize = Size4KiB> {
    start_address: VirtAddr,
    size: PhantomData<S>,
}

impl<S: PageSize> Page<S> {
    /// The page size in bytes.
    pub const SIZE: u64 = S::SIZE;

    /// Returns the page that starts at the given virtual address.
    ///
    /// Returns an error if the address is not correctly aligned (i.e. is not a valid page start).
    #[inline]
    pub fn from_start_address(address: VirtAddr) -> Result<Self, AddressNotAligned> {
        if !address.is_aligned(S::SIZE) {
            return Err(AddressNotAligned);
        }
        Ok(Page::containing_address(address))
    }

    const_fn! {
        /// Returns the page that starts at the given virtual address.
        ///
        /// ## Safety
        ///
        /// The address must be correctly aligned.
        #[inline]
        pub unsafe fn from_start_address_unchecked(start_address: VirtAddr) -> Self {
            Page {
                start_address,
                size: PhantomData,
            }
        }
    }

    /// Returns the page that contains the given virtual address.
    #[inline]
    pub fn containing_address(address: VirtAddr) -> Self {
        Page {
            start_address: address.align_down(S::SIZE),
            size: PhantomData,
        }
    }

    const_fn! {
        /// Returns the start address of the page.
        #[inline]
        pub fn start_address(self) -> VirtAddr {
            self.start_address
        }
    }

    const_fn! {
        /// Returns the size the page (4KB, 2MB or 1GB).
        #[inline]
        pub fn size(self) -> u64 {
            S::SIZE
        }
    }

    const_fn! {
        /// Returns the level 4 page table index of this page.
        #[inline]
        pub fn p4_index(self) -> PageTableIndex {
            self.start_address().p4_index()
        }
    }

    const_fn! {
        /// Returns the level 3 page table index of this page.
        #[inline]
        pub fn p3_index(self) -> PageTableIndex {
            self.start_address().p3_index()
        }
    }

    const_fn! {
        /// Returns the table index of this page at the specified level.
        #[inline]
        pub fn page_table_index(self, level: PageTableLevel) -> PageTableIndex {
            self.start_address().page_table_index(level)
        }
    }

    const_fn! {
        /// Returns a range of pages, exclusive `end`.
        #[inline]
        pub fn range(start: Self, end: Self) -> PageRange<S> {
            PageRange { start, end }
        }
    }

    const_fn! {
        /// Returns a range of pages, inclusive `end`.
        #[inline]
        pub fn range_inclusive(start: Self, end: Self) -> PageRangeInclusive<S> {
            PageRangeInclusive { start, end }
        }
    }
}

impl<S: NotGiantPageSize> Page<S> {
    const_fn! {
        /// Returns the level 2 page table index of this page.
        #[inline]
        pub fn p2_index(self) -> PageTableIndex {
            self.start_address().p2_index()
        }
    }
}

impl Page<Size1GiB> {
    /// Returns the 1GiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices_1gib(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        Page::containing_address(VirtAddr::new(addr))
    }
}

impl Page<Size2MiB> {
    /// Returns the 2MiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices_2mib(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
        p2_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        addr.set_bits(21..30, u64::from(p2_index));
        Page::containing_address(VirtAddr::new(addr))
    }
}

impl Page<Size4KiB> {
    /// Returns the 4KiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
        p2_index: PageTableIndex,
        p1_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        addr.set_bits(21..30, u64::from(p2_index));
        addr.set_bits(12..21, u64::from(p1_index));
        Page::containing_address(VirtAddr::new(addr))
    }

    /// Returns the level 1 page table index of this page.
    #[inline]
    pub const fn p1_index(self) -> PageTableIndex {
        self.start_address.p1_index()
    }
}

impl<S: PageSize> fmt::Debug for Page<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_fmt(format_args!(
            "Page[{}]({:#x})",
            S::SIZE_AS_DEBUG_STR,
            self.start_address().as_u64()
        ))
    }
}

impl<S: PageSize> Add<u64> for Page<S> {
    type Output = Self;
    #[inline]
    fn add(self, rhs: u64) -> Self::Output {
        Page::containing_address(self.start_address() + rhs * S::SIZE)
    }
}

impl<S: PageSize> AddAssign<u64> for Page<S> {
    #[inline]
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl<S: PageSize> Sub<u64> for Page<S> {
    type Output = Self;
    #[inline]
    fn sub(self, rhs: u64) -> Self::Output {
        Page::containing_address(self.start_address() - rhs * S::SIZE)
    }
}

impl<S: PageSize> SubAssign<u64> for Page<S> {
    #[inline]
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl<S: PageSize> Sub<Self> for Page<S> {
    type Output = u64;
    #[inline]
    fn sub(self, rhs: Self) -> Self::Output {
        (self.start_address - rhs.start_address) / S::SIZE
    }
}

#[cfg(feature = "step_trait")]
impl<S: PageSize> Step for Page<S> {
    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
        Step::steps_between(&start.start_address, &end.start_address)
            .map(|steps| steps / S::SIZE as usize)
    }

    fn forward_checked(start: Self, count: usize) -> Option<Self> {
        let count = count.checked_mul(S::SIZE as usize)?;
        let start_address = Step::forward_checked(start.start_address, count)?;
        Some(Self {
            start_address,
            size: PhantomData,
        })
    }

    fn backward_checked(start: Self, count: usize) -> Option<Self> {
        let count = count.checked_mul(S::SIZE as usize)?;
        let start_address = Step::backward_checked(start.start_address, count)?;
        Some(Self {
            start_address,
            size: PhantomData,
        })
    }
}

/// A range of pages with exclusive upper bound.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct PageRange<S: PageSize = Size4KiB> {
    /// The start of the range, inclusive.
    pub start: Page<S>,
    /// The end of the range, exclusive.
    pub end: Page<S>,
}

impl<S: PageSize> PageRange<S> {
    /// Returns wether this range contains no pages.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.start >= self.end
    }
}

impl<S: PageSize> Iterator for PageRange<S> {
    type Item = Page<S>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.start < self.end {
            let page = self.start;
            self.start += 1;
            Some(page)
        } else {
            None
        }
    }
}

impl PageRange<Size2MiB> {
    /// Converts the range of 2MiB pages to a range of 4KiB pages.
    #[inline]
    pub fn as_4kib_page_range(self) -> PageRange<Size4KiB> {
        PageRange {
            start: Page::containing_address(self.start.start_address()),
            end: Page::containing_address(self.end.start_address()),
        }
    }
}

impl<S: PageSize> fmt::Debug for PageRange<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("PageRange")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

/// A range of pages with inclusive upper bound.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct PageRangeInclusive<S: PageSize = Size4KiB> {
    /// The start of the range, inclusive.
    pub start: Page<S>,
    /// The end of the range, inclusive.
    pub end: Page<S>,
}

impl<S: PageSize> PageRangeInclusive<S> {
    /// Returns wether this range contains no pages.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.start > self.end
    }
}

impl<S: PageSize> Iterator for PageRangeInclusive<S> {
    type Item = Page<S>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.start <= self.end {
            let page = self.start;
            self.start += 1;
            Some(page)
        } else {
            None
        }
    }
}

impl<S: PageSize> fmt::Debug for PageRangeInclusive<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("PageRangeInclusive")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

/// The given address was not sufficiently aligned.
#[derive(Debug)]
pub struct AddressNotAligned;

impl fmt::Display for AddressNotAligned {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "the given address was not sufficiently aligned")
    }
}

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

    #[test]
    pub fn test_page_ranges() {
        let page_size = Size4KiB::SIZE;
        let number = 1000;

        let start_addr = VirtAddr::new(0xdead_beaf);
        let start: Page = Page::containing_address(start_addr);
        let end = start + number;

        let mut range = Page::range(start, end);
        for i in 0..number {
            assert_eq!(
                range.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range.next(), None);

        let mut range_inclusive = Page::range_inclusive(start, end);
        for i in 0..=number {
            assert_eq!(
                range_inclusive.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range_inclusive.next(), None);
    }
}
<commit_msg>Fixed overflow bug in PageRangeInclusive<commit_after>//! Abstractions for default-sized and huge virtual memory pages.

use crate::structures::paging::page_table::PageTableLevel;
use crate::structures::paging::PageTableIndex;
use crate::VirtAddr;
use core::fmt;
#[cfg(feature = "step_trait")]
use core::iter::Step;
use core::marker::PhantomData;
use core::ops::{Add, AddAssign, Sub, SubAssign};

/// Trait for abstracting over the three possible page sizes on x86_64, 4KiB, 2MiB, 1GiB.
pub trait PageSize: Copy + Eq + PartialOrd + Ord {
    /// The page size in bytes.
    const SIZE: u64;

    /// A string representation of the page size for debug output.
    const SIZE_AS_DEBUG_STR: &'static str;
}

/// This trait is implemented for 4KiB and 2MiB pages, but not for 1GiB pages.
pub trait NotGiantPageSize: PageSize {}

/// A standard 4KiB page.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size4KiB {}

/// A “huge” 2MiB page.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size2MiB {}

/// A “giant” 1GiB page.
///
/// (Only available on newer x86_64 CPUs.)
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size1GiB {}

impl PageSize for Size4KiB {
    const SIZE: u64 = 4096;
    const SIZE_AS_DEBUG_STR: &'static str = "4KiB";
}

impl NotGiantPageSize for Size4KiB {}

impl PageSize for Size2MiB {
    const SIZE: u64 = Size4KiB::SIZE * 512;
    const SIZE_AS_DEBUG_STR: &'static str = "2MiB";
}

impl NotGiantPageSize for Size2MiB {}

impl PageSize for Size1GiB {
    const SIZE: u64 = Size2MiB::SIZE * 512;
    const SIZE_AS_DEBUG_STR: &'static str = "1GiB";
}

/// A virtual memory page.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(C)]
pub struct Page<S: PageSize = Size4KiB> {
    start_address: VirtAddr,
    size: PhantomData<S>,
}

impl<S: PageSize> Page<S> {
    /// The page size in bytes.
    pub const SIZE: u64 = S::SIZE;

    /// Returns the page that starts at the given virtual address.
    ///
    /// Returns an error if the address is not correctly aligned (i.e. is not a valid page start).
    #[inline]
    pub fn from_start_address(address: VirtAddr) -> Result<Self, AddressNotAligned> {
        if !address.is_aligned(S::SIZE) {
            return Err(AddressNotAligned);
        }
        Ok(Page::containing_address(address))
    }

    const_fn! {
        /// Returns the page that starts at the given virtual address.
        ///
        /// ## Safety
        ///
        /// The address must be correctly aligned.
        #[inline]
        pub unsafe fn from_start_address_unchecked(start_address: VirtAddr) -> Self {
            Page {
                start_address,
                size: PhantomData,
            }
        }
    }

    /// Returns the page that contains the given virtual address.
    #[inline]
    pub fn containing_address(address: VirtAddr) -> Self {
        Page {
            start_address: address.align_down(S::SIZE),
            size: PhantomData,
        }
    }

    const_fn! {
        /// Returns the start address of the page.
        #[inline]
        pub fn start_address(self) -> VirtAddr {
            self.start_address
        }
    }

    const_fn! {
        /// Returns the size the page (4KB, 2MB or 1GB).
        #[inline]
        pub fn size(self) -> u64 {
            S::SIZE
        }
    }

    const_fn! {
        /// Returns the level 4 page table index of this page.
        #[inline]
        pub fn p4_index(self) -> PageTableIndex {
            self.start_address().p4_index()
        }
    }

    const_fn! {
        /// Returns the level 3 page table index of this page.
        #[inline]
        pub fn p3_index(self) -> PageTableIndex {
            self.start_address().p3_index()
        }
    }

    const_fn! {
        /// Returns the table index of this page at the specified level.
        #[inline]
        pub fn page_table_index(self, level: PageTableLevel) -> PageTableIndex {
            self.start_address().page_table_index(level)
        }
    }

    const_fn! {
        /// Returns a range of pages, exclusive `end`.
        #[inline]
        pub fn range(start: Self, end: Self) -> PageRange<S> {
            PageRange { start, end }
        }
    }

    const_fn! {
        /// Returns a range of pages, inclusive `end`.
        #[inline]
        pub fn range_inclusive(start: Self, end: Self) -> PageRangeInclusive<S> {
            PageRangeInclusive { start, end }
        }
    }
}

impl<S: NotGiantPageSize> Page<S> {
    const_fn! {
        /// Returns the level 2 page table index of this page.
        #[inline]
        pub fn p2_index(self) -> PageTableIndex {
            self.start_address().p2_index()
        }
    }
}

impl Page<Size1GiB> {
    /// Returns the 1GiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices_1gib(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        Page::containing_address(VirtAddr::new(addr))
    }
}

impl Page<Size2MiB> {
    /// Returns the 2MiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices_2mib(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
        p2_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        addr.set_bits(21..30, u64::from(p2_index));
        Page::containing_address(VirtAddr::new(addr))
    }
}

impl Page<Size4KiB> {
    /// Returns the 4KiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
        p2_index: PageTableIndex,
        p1_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        addr.set_bits(21..30, u64::from(p2_index));
        addr.set_bits(12..21, u64::from(p1_index));
        Page::containing_address(VirtAddr::new(addr))
    }

    /// Returns the level 1 page table index of this page.
    #[inline]
    pub const fn p1_index(self) -> PageTableIndex {
        self.start_address.p1_index()
    }
}

impl<S: PageSize> fmt::Debug for Page<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_fmt(format_args!(
            "Page[{}]({:#x})",
            S::SIZE_AS_DEBUG_STR,
            self.start_address().as_u64()
        ))
    }
}

impl<S: PageSize> Add<u64> for Page<S> {
    type Output = Self;
    #[inline]
    fn add(self, rhs: u64) -> Self::Output {
        Page::containing_address(self.start_address() + rhs * S::SIZE)
    }
}

impl<S: PageSize> AddAssign<u64> for Page<S> {
    #[inline]
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl<S: PageSize> Sub<u64> for Page<S> {
    type Output = Self;
    #[inline]
    fn sub(self, rhs: u64) -> Self::Output {
        Page::containing_address(self.start_address() - rhs * S::SIZE)
    }
}

impl<S: PageSize> SubAssign<u64> for Page<S> {
    #[inline]
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl<S: PageSize> Sub<Self> for Page<S> {
    type Output = u64;
    #[inline]
    fn sub(self, rhs: Self) -> Self::Output {
        (self.start_address - rhs.start_address) / S::SIZE
    }
}

#[cfg(feature = "step_trait")]
impl<S: PageSize> Step for Page<S> {
    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
        Step::steps_between(&start.start_address, &end.start_address)
            .map(|steps| steps / S::SIZE as usize)
    }

    fn forward_checked(start: Self, count: usize) -> Option<Self> {
        let count = count.checked_mul(S::SIZE as usize)?;
        let start_address = Step::forward_checked(start.start_address, count)?;
        Some(Self {
            start_address,
            size: PhantomData,
        })
    }

    fn backward_checked(start: Self, count: usize) -> Option<Self> {
        let count = count.checked_mul(S::SIZE as usize)?;
        let start_address = Step::backward_checked(start.start_address, count)?;
        Some(Self {
            start_address,
            size: PhantomData,
        })
    }
}

/// A range of pages with exclusive upper bound.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct PageRange<S: PageSize = Size4KiB> {
    /// The start of the range, inclusive.
    pub start: Page<S>,
    /// The end of the range, exclusive.
    pub end: Page<S>,
}

impl<S: PageSize> PageRange<S> {
    /// Returns wether this range contains no pages.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.start >= self.end
    }
}

impl<S: PageSize> Iterator for PageRange<S> {
    type Item = Page<S>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.start < self.end {
            let page = self.start;
            self.start += 1;
            Some(page)
        } else {
            None
        }
    }
}

impl PageRange<Size2MiB> {
    /// Converts the range of 2MiB pages to a range of 4KiB pages.
    #[inline]
    pub fn as_4kib_page_range(self) -> PageRange<Size4KiB> {
        PageRange {
            start: Page::containing_address(self.start.start_address()),
            end: Page::containing_address(self.end.start_address()),
        }
    }
}

impl<S: PageSize> fmt::Debug for PageRange<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("PageRange")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

/// A range of pages with inclusive upper bound.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct PageRangeInclusive<S: PageSize = Size4KiB> {
    /// The start of the range, inclusive.
    pub start: Page<S>,
    /// The end of the range, inclusive.
    pub end: Page<S>,
}

impl<S: PageSize> PageRangeInclusive<S> {
    /// Returns wether this range contains no pages.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.start > self.end
    }
}

impl<S: PageSize> Iterator for PageRangeInclusive<S> {
    type Item = Page<S>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.start <= self.end {
            let page = self.start;

            // If the end of the inclusive range is the maximum page possible for size S,
            // incrementing start until it is greater than the end will cause an integer overflow.
            // So instead, in that case we decrement end rather than incrementing start.
            let max_page_addr = VirtAddr::new(u64::MAX) - (S::SIZE - 1);
            if self.start.start_address() < max_page_addr {
                self.start += 1;
            } else {
                self.end -= 1;
            }
            Some(page)
        } else {
            None
        }
    }
}

impl<S: PageSize> fmt::Debug for PageRangeInclusive<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("PageRangeInclusive")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

/// The given address was not sufficiently aligned.
#[derive(Debug)]
pub struct AddressNotAligned;

impl fmt::Display for AddressNotAligned {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "the given address was not sufficiently aligned")
    }
}

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

    #[test]
    pub fn test_page_ranges() {
        let page_size = Size4KiB::SIZE;
        let number = 1000;

        let start_addr = VirtAddr::new(0xdead_beaf);
        let start: Page = Page::containing_address(start_addr);
        let end = start + number;

        let mut range = Page::range(start, end);
        for i in 0..number {
            assert_eq!(
                range.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range.next(), None);

        let mut range_inclusive = Page::range_inclusive(start, end);
        for i in 0..=number {
            assert_eq!(
                range_inclusive.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range_inclusive.next(), None);
    }
}
<|endoftext|>