<commit_before>//! https://arkada38.github.io/2018/01/14/creating-a-turtle/
//!
//! Inspired by child's drawing of turtle we are going to draw and paint figures
//! which consists of arcs and straight lines.
//! To draw arcs we need to use the loops which are determine tilt angle and the length.
//! The more sophisticated the figure is the more loops it requires to make.
extern crate turtle;

use turtle::{Turtle, Color, color};

const SIZE: f64 = 1.0;
const SHELL_COLOR: Color = Color {red: 62.0, green: 114.0, blue: 29.0, alpha: 1.0};
const BODY_COLOR: Color = Color {red: 119.0, green: 178.0, blue: 85.0, alpha: 1.0};
const EYE_COLOR: Color = color::BLACK;

fn main() {
    let mut turtle = Turtle::new();
    turtle.set_speed(8);

    turtle.pen_up();
    turtle.set_x(-280.0);
    turtle.set_y(-90.0);

    draw_shell(&mut turtle);

    draw_tail(&mut turtle);

    turtle.pen_up();
    turtle.right(55.0);
    turtle.forward(SIZE * 10.0);
    turtle.right(55.0);

    // Back leg
    draw_leg(&mut turtle);

    turtle.right(86.5);
    turtle.forward(SIZE * 55.0);
    turtle.right(86.0);

    // Front leg
    draw_leg(&mut turtle);

    turtle.right(54.5);
    turtle.forward(SIZE * 8.0);

    draw_neck(&mut turtle);

    turtle.right(172.6);
    turtle.forward(SIZE * 40.0);
    turtle.right(110.0);

    draw_head(&mut turtle);

    turtle.left(128.0);
    turtle.forward(SIZE * 15.0);

    draw_eye(&mut turtle);

    turtle.set_fill_color(BODY_COLOR);

    turtle.left(175.0);
    turtle.forward(SIZE * 43.0);

    draw_right_highlight(&mut turtle);

    turtle.right(18.0);
    turtle.forward(SIZE * 37.0);
    turtle.set_heading(180.0);

    draw_middle_highlight(&mut turtle);

    turtle.left(24.0);
    turtle.forward(SIZE * 36.0);
    turtle.set_heading(180.0);

    draw_left_highlight(&mut turtle);/**/
}

fn draw_shell(turtle: &mut Turtle) {
    turtle.set_fill_color(SHELL_COLOR);
    turtle.begin_fill();

    for _ in 0..180 {
        turtle.forward(SIZE);
        turtle.right(1.0);
    }

    for _ in 0..90 {
        turtle.forward(SIZE / 3.0);
        turtle.right(1.0);
    }

    turtle.set_speed(5);
    let d = SIZE * 360.0 / std::f64::consts::PI;
    turtle.forward(d - d / 3.0);
    turtle.set_speed(10);

    for _ in 0..90 {
        turtle.forward(SIZE / 3.0);
        turtle.right(1.0);
    }

    turtle.end_fill();
}

fn draw_tail(turtle: &mut Turtle) {
    turtle.set_fill_color(BODY_COLOR);
    turtle.begin_fill();

    turtle.left(90.0);
    turtle.forward(SIZE);

    for _ in 0..45 {
        turtle.forward(SIZE / 3.0);
        turtle.left(1.0);
    }

    turtle.forward(SIZE * 3.0);

    for _ in 0..25 {
        turtle.forward(SIZE / 3.0);
        turtle.left(6.0);
    }

    turtle.forward(SIZE * 9.0);
    turtle.right(6.0);
    turtle.forward(SIZE * 8.0);

    turtle.end_fill();
}

fn draw_leg(turtle: &mut Turtle) {
    turtle.begin_fill();

    for _ in 0..15 {
        turtle.forward(SIZE);
        turtle.left(0.5);
    }

    for _ in 0..90 {
        turtle.forward(SIZE / 6.0);
        turtle.left(1.0);
    }

    turtle.forward(SIZE * 3.0);

    for _ in 0..90 {
        turtle.forward(SIZE / 6.0);
        turtle.left(1.0);
    }

    turtle.forward(SIZE * 14.5);

    turtle.end_fill();
}

fn draw_neck(turtle: &mut Turtle) {
    turtle.begin_fill();

    for _ in 0..15 {
        turtle.forward(SIZE * 3.0);
        turtle.left(1.5);
    }

    turtle.end_fill();
    turtle.begin_fill();

    turtle.left(100.0);
    turtle.forward(SIZE * 20.0);
    turtle.left(80.0);

    for _ in 0..4 {
        turtle.forward(SIZE * 3.7);
        turtle.left(1.5);
    }

    turtle.left(30.0);

    for _ in 0..27 {
        turtle.forward(SIZE);
        turtle.right(1.0);
    }

    turtle.end_fill();
}

// In our case the most difficult part of the turtle is head. It requires six loops to draw.
// There was a way to draw it as one circle but
// I've decided to make it look a like a real turtle's head.
fn draw_head(turtle: &mut Turtle) {
    turtle.begin_fill();

    for _ in 0..20 {
        turtle.forward(SIZE * 1.2);
        turtle.left(1.0);
    }

    for _ in 0..10 {
        turtle.forward(SIZE * 1.2);
        turtle.left(4.0);
    }

    for _ in 0..10 {
        turtle.forward(SIZE / 1.5);
        turtle.left(7.0);
    }

    for _ in 0..10 {
        turtle.forward(SIZE);
        turtle.left(2.0);
    }

    for _ in 0..50 {
        turtle.forward(SIZE / 2.5);
        turtle.left(1.0);
    }

    for _ in 0..30 {
        turtle.forward(SIZE / 3.0);
        turtle.left(1.8);
    }

    for _ in 0..10 {
        turtle.forward(SIZE / 1.5);
        turtle.left(1.8);
    }

    turtle.end_fill();
}

fn draw_eye(turtle: &mut Turtle) {
    turtle.set_fill_color(EYE_COLOR);
    turtle.begin_fill();
    for _ in 0..90 {
        turtle.forward(SIZE / 3.0);
        turtle.right(4.0);
    }
    turtle.end_fill();
}

fn draw_right_highlight(turtle: &mut Turtle) {
    turtle.begin_fill();

    for _ in 0..39 {
        turtle.forward(SIZE / 5.0);
        turtle.left(2.0);
    }

    turtle.forward(SIZE * 36.0);

    for _ in 0..90 {
        turtle.forward(SIZE / 2.5);
        turtle.left(2.0);
    }

    for _ in 0..42 {
        turtle.forward(SIZE);
        turtle.left(0.9);
    }

    for _ in 0..26 {
        turtle.forward(SIZE / 5.0);
        turtle.left(2.0);
    }

    turtle.end_fill();
}

fn draw_middle_highlight(turtle: &mut Turtle) {
    turtle.begin_fill();

    for _ in 0..40 {
        turtle.forward(SIZE / 4.0);
        turtle.left(2.0);
    }

    turtle.forward(SIZE * 47.0);
    turtle.left(8.5);

    for _ in 0..90 {
        turtle.forward(SIZE / 2.0);
        turtle.left(2.0);
    }
    turtle.forward(SIZE / 2.0);

    turtle.left(8.5);
    turtle.forward(SIZE * 47.0);

    for _ in 0..40 {
        turtle.left(2.0);
        turtle.forward(SIZE / 4.0);
    }

    turtle.end_fill();
}

fn draw_left_highlight(turtle: &mut Turtle) {
    turtle.begin_fill();

    for _ in 0..26 {
        turtle.forward(SIZE / 5.0);
        turtle.left(2.0);
    }

    for _ in 0..42 {
        turtle.forward(SIZE);
        turtle.left(0.9);
    }

    for _ in 0..90 {
        turtle.forward(SIZE / 2.5);
        turtle.left(2.0);
    }

    turtle.forward(SIZE * 36.0);

    for _ in 0..39 {
        turtle.forward(SIZE / 5.0);
        turtle.left(2.0);
    }

    turtle.end_fill();
}
<commit_msg>Update draw_turtle.rs<commit_after>//! https://arkada38.github.io/2018/01/14/creating-a-turtle/
//!
//! This is inspired by a child's drawing of a turtle.
//! We are going to draw the turtle using arcs and straight lines.
//! To draw arcs, we use multiple Rust for-loops to create different tilt angles and
//! lengths. The more sophisticated the figure is, the more loops we need to make it.
extern crate turtle;

use turtle::{Turtle, Color, color};

const SIZE: f64 = 1.0;
const SHELL_COLOR: Color = Color {red: 62.0, green: 114.0, blue: 29.0, alpha: 1.0};
const BODY_COLOR: Color = Color {red: 119.0, green: 178.0, blue: 85.0, alpha: 1.0};
const EYE_COLOR: Color = color::BLACK;

fn main() {
    let mut turtle = Turtle::new();
    turtle.set_speed(8);

    turtle.pen_up();
    turtle.set_x(-280.0);
    turtle.set_y(-90.0);

    draw_shell(&mut turtle);

    draw_tail(&mut turtle);

    turtle.pen_up();
    turtle.right(55.0);
    turtle.forward(SIZE * 10.0);
    turtle.right(55.0);

    // Back leg
    draw_leg(&mut turtle);

    turtle.right(86.5);
    turtle.forward(SIZE * 55.0);
    turtle.right(86.0);

    // Front leg
    draw_leg(&mut turtle);

    turtle.right(54.5);
    turtle.forward(SIZE * 8.0);

    draw_neck(&mut turtle);

    turtle.right(172.6);
    turtle.forward(SIZE * 40.0);
    turtle.right(110.0);

    draw_head(&mut turtle);

    turtle.left(128.0);
    turtle.forward(SIZE * 15.0);

    draw_eye(&mut turtle);

    turtle.set_fill_color(BODY_COLOR);

    turtle.left(175.0);
    turtle.forward(SIZE * 43.0);

    // Here we start to draw highlights on the shell.
    // We have 3 highlights: on the right, in the middle and on the left.

    draw_right_highlight(&mut turtle);

    turtle.right(18.0);
    turtle.forward(SIZE * 37.0);
    turtle.set_heading(180.0);

    draw_middle_highlight(&mut turtle);

    turtle.left(24.0);
    turtle.forward(SIZE * 36.0);
    turtle.set_heading(180.0);

    draw_left_highlight(&mut turtle);/**/
}

fn draw_shell(turtle: &mut Turtle) {
    turtle.set_fill_color(SHELL_COLOR);
    turtle.begin_fill();

    for _ in 0..180 {
        turtle.forward(SIZE);
        turtle.right(1.0);
    }

    for _ in 0..90 {
        turtle.forward(SIZE / 3.0);
        turtle.right(1.0);
    }

    turtle.set_speed(5);
    let d = SIZE * 360.0 / std::f64::consts::PI;
    turtle.forward(d - d / 3.0);
    turtle.set_speed(10);

    for _ in 0..90 {
        turtle.forward(SIZE / 3.0);
        turtle.right(1.0);
    }

    turtle.end_fill();
}

fn draw_tail(turtle: &mut Turtle) {
    turtle.set_fill_color(BODY_COLOR);
    turtle.begin_fill();

    turtle.left(90.0);
    turtle.forward(SIZE);

    for _ in 0..45 {
        turtle.forward(SIZE / 3.0);
        turtle.left(1.0);
    }

    turtle.forward(SIZE * 3.0);

    for _ in 0..25 {
        turtle.forward(SIZE / 3.0);
        turtle.left(6.0);
    }

    turtle.forward(SIZE * 9.0);
    turtle.right(6.0);
    turtle.forward(SIZE * 8.0);

    turtle.end_fill();
}

fn draw_leg(turtle: &mut Turtle) {
    turtle.begin_fill();

    for _ in 0..15 {
        turtle.forward(SIZE);
        turtle.left(0.5);
    }

    for _ in 0..90 {
        turtle.forward(SIZE / 6.0);
        turtle.left(1.0);
    }

    turtle.forward(SIZE * 3.0);

    for _ in 0..90 {
        turtle.forward(SIZE / 6.0);
        turtle.left(1.0);
    }

    turtle.forward(SIZE * 14.5);

    turtle.end_fill();
}

fn draw_neck(turtle: &mut Turtle) {
    turtle.begin_fill();

    for _ in 0..15 {
        turtle.forward(SIZE * 3.0);
        turtle.left(1.5);
    }

    turtle.end_fill();
    turtle.begin_fill();

    turtle.left(100.0);
    turtle.forward(SIZE * 20.0);
    turtle.left(80.0);

    for _ in 0..4 {
        turtle.forward(SIZE * 3.7);
        turtle.left(1.5);
    }

    turtle.left(30.0);

    for _ in 0..27 {
        turtle.forward(SIZE);
        turtle.right(1.0);
    }

    turtle.end_fill();
}

// In our case the most difficult part of the turtle is head. It requires six loops to draw.
// There was a way to draw it as one circle but
// I've decided to make it look a like a real turtle's head.
fn draw_head(turtle: &mut Turtle) {
    turtle.begin_fill();

    for _ in 0..20 {
        turtle.forward(SIZE * 1.2);
        turtle.left(1.0);
    }

    for _ in 0..10 {
        turtle.forward(SIZE * 1.2);
        turtle.left(4.0);
    }

    for _ in 0..10 {
        turtle.forward(SIZE / 1.5);
        turtle.left(7.0);
    }

    for _ in 0..10 {
        turtle.forward(SIZE);
        turtle.left(2.0);
    }

    for _ in 0..50 {
        turtle.forward(SIZE / 2.5);
        turtle.left(1.0);
    }

    for _ in 0..30 {
        turtle.forward(SIZE / 3.0);
        turtle.left(1.8);
    }

    for _ in 0..10 {
        turtle.forward(SIZE / 1.5);
        turtle.left(1.8);
    }

    turtle.end_fill();
}

fn draw_eye(turtle: &mut Turtle) {
    turtle.set_fill_color(EYE_COLOR);
    turtle.begin_fill();
    for _ in 0..90 {
        turtle.forward(SIZE / 3.0);
        turtle.right(4.0);
    }
    turtle.end_fill();
}

fn draw_right_highlight(turtle: &mut Turtle) {
    turtle.begin_fill();

    for _ in 0..39 {
        turtle.forward(SIZE / 5.0);
        turtle.left(2.0);
    }

    turtle.forward(SIZE * 36.0);

    for _ in 0..90 {
        turtle.forward(SIZE / 2.5);
        turtle.left(2.0);
    }

    for _ in 0..42 {
        turtle.forward(SIZE);
        turtle.left(0.9);
    }

    for _ in 0..26 {
        turtle.forward(SIZE / 5.0);
        turtle.left(2.0);
    }

    turtle.end_fill();
}

fn draw_middle_highlight(turtle: &mut Turtle) {
    turtle.begin_fill();

    for _ in 0..40 {
        turtle.forward(SIZE / 4.0);
        turtle.left(2.0);
    }

    turtle.forward(SIZE * 47.0);
    turtle.left(8.5);

    for _ in 0..90 {
        turtle.forward(SIZE / 2.0);
        turtle.left(2.0);
    }
    turtle.forward(SIZE / 2.0);

    turtle.left(8.5);
    turtle.forward(SIZE * 47.0);

    for _ in 0..40 {
        turtle.left(2.0);
        turtle.forward(SIZE / 4.0);
    }

    turtle.end_fill();
}

fn draw_left_highlight(turtle: &mut Turtle) {
    turtle.begin_fill();

    for _ in 0..26 {
        turtle.forward(SIZE / 5.0);
        turtle.left(2.0);
    }

    for _ in 0..42 {
        turtle.forward(SIZE);
        turtle.left(0.9);
    }

    for _ in 0..90 {
        turtle.forward(SIZE / 2.5);
        turtle.left(2.0);
    }

    turtle.forward(SIZE * 36.0);

    for _ in 0..39 {
        turtle.forward(SIZE / 5.0);
        turtle.left(2.0);
    }

    turtle.end_fill();
}
<|endoftext|>


<commit_before>// "Tifflin" Kernel
// - By John Hodge (thePowersGang)
//
// Core/device_manager.rs
// - Core device manager

use _common::*;
use sync::Mutex;
use lib::Queue;

module_define!{DeviceManager, [arch], init}

/// A semi-arbitatry integer denoting how well a driver handles a specific device
pub type DriverHandleLevel = u32;

/// IO range binding
pub enum IOBinding
{
	/// Memory-mapped IO space
	Memory(::memory::virt::AllocHandle),
	/// x86 IO bus (Base and offset)
	IO(u16,u16),
}

/// Interface a bus manager instance
pub trait BusManager:
	Send + Sync
{
	/// Returns the textual name of the bus type (e.g. "pci")
	fn bus_type(&self) -> &str;
	/// Returns a list of valid attributes for BusDevice::get_attr
	fn get_attr_names(&self) -> &[&str];
}

/// Interface to a device on a bus
pub trait BusDevice:
	Send
{
	/// Returns the device's address on the parent bus
	fn addr(&self) -> u32;
	/// Retrurns the specified attribute (or 0, if invalid)
	fn get_attr(&self, name: &str) -> u32;
	/// Set the power state of this device
	fn set_power(&mut self, state: bool);	// TODO: Power state enum for Off,Standby,Low,On
	/// Bind to the specified IO block (meaning of `block_id` depends on the bus)
	fn bind_io(&mut self, block_id: usize) -> IOBinding;
	/// Obtain the specified interrupt vector
	fn get_irq(&mut self, idx: usize) -> u32;
}

// TODO: Change this to instead be a structure with a bound Fn reference
// - Structure defines bus type and a set of attribute names/values/masks
/// Abstract driver for a device (creates instances when passed a device)
pub trait Driver:
	Send + Sync
{
	/// Driver's name
	fn name(&self) -> &str;
	/// Bus type the driver binds against (matches value from `BusManager::bus_type`)
	fn bus_type(&self) -> &str;
	/// Return the handling level of this driver for the specified device
	fn handles(&self, bus_dev: &BusDevice) -> DriverHandleLevel;
	/// Requests that the driver bind itself to the specified device
	fn bind(&self, bus_dev: &mut BusDevice) -> Box<DriverInstance>;
}

/// Driver instance (maps directly to a device)
pub trait DriverInstance:
	Send
{
}

/// Internal representation of a device on a bus
struct Device
{
	bus_dev: Box<BusDevice>,
	driver: Option<(Box<DriverInstance>, DriverHandleLevel)>,
	//attribs: Vec<u32>,
}

/// Internal representation of a bus
struct Bus
{
	manager: &'static BusManager,
	devices: Vec<Device>,
}

/// List of registered busses on the system
#[allow(non_upper_case_globals)]
static s_root_busses: Mutex<Queue<Bus>> = mutex_init!(queue_init!());

/// List of registered drivers
#[allow(non_upper_case_globals)]
static s_driver_list: Mutex<Queue<&'static Driver>> = mutex_init!( queue_init!() );

fn init()
{
	// Do nothing!
}

/// Register a bus with the device manager
///
/// Creates a new internal representation of the bus, containg the passed set of devices.
pub fn register_bus(manager: &'static BusManager, devices: Vec<Box<BusDevice>>)
{
	let bus = Bus {
		manager: manager,
		// For each device, locate a driver
		devices: devices.into_iter().map(|mut d| Device {
			driver: find_driver(manager, &mut *d),
			//attribs: Vec::new(),
			bus_dev: d,
			}).collect(),
		};
	s_root_busses.lock().push(bus);
}

/// Registers a driver with the device manger
pub fn register_driver(driver: &'static Driver)
{
	s_driver_list.lock().push(driver);
	log_debug!("Registering driver {}", driver.name());
	// Iterate known devices and spin up instances if needed
	for bus in s_root_busses.lock().iter_mut()
	{
		log_trace!("bus type {}", bus.manager.bus_type());
		for dev in bus.devices.iter_mut()
		{
			let rank = driver.handles(&*dev.bus_dev);
			log_debug!("rank = {:?}", rank);
			if rank == 0
			{
				// SKIP!
			}
			else if dev.driver.is_some()
			{
				let bind = dev.driver.as_ref().unwrap();
				let cur_rank = bind.1;
				if cur_rank > rank
				{
					// Existing driver is better
				}
				else if cur_rank == rank
				{
					// Fight!
				}
				else
				{
					// New driver is better
					panic!("TODO: Unbind driver and bind in new one");
				}
			}
			else
			{
				// Bind new driver
				dev.driver = Some( (driver.bind(&mut *dev.bus_dev), rank) );
			}
		}
	}
}

/**
 * Locate the best registered driver for this device and instanciate it
 */
fn find_driver(bus: &BusManager, bus_dev: &mut BusDevice) -> Option<(Box<DriverInstance>,DriverHandleLevel)>
{
	log_debug!("Finding driver for {}:{:x}", bus.bus_type(), bus_dev.addr());
	let mut best_ranking = 0;
	let mut best_driver = None;
	for driver in s_driver_list.lock().iter()
	{
		if bus.bus_type() == driver.bus_type()
		{
			let ranking = driver.handles(bus_dev);
			if ranking == 0
			{
				// Doesn't handle this device
			}
			else if ranking > best_ranking
			{
				// Best so far
				best_driver = Some( *driver );
				best_ranking = ranking;
			}
			else if ranking == best_ranking
			{
				// A tie, this is not very good
				//log_warning!("Tie for device {}:{:x} between {} and {}",
				//	bus.bus_type(), bus_dev.addr(), driver, best_driver.unwrap());
			}
			else
			{
				// Not as good as current, move along
			}
		}
	}
	best_driver.map(|d| (d.bind(bus_dev), best_ranking))
}

impl IOBinding
{
	/// Returns the x86 IO space base
	pub fn io_base(&self) -> u16 {
		match *self
		{
		IOBinding::IO(base, _size) => base,
		IOBinding::Memory(_) => panic!("Called IOBinding::io_base on IOBinding::Memory"),
		}
	}
	/// Writes a single u8 to the binding
	pub unsafe fn write_8(&self, ofs: usize, val: u8)
	{
		match *self
		{
		IOBinding::IO(base, s) => {
			assert!( ofs < s as usize );
			::arch::x86_io::outb(base + ofs as u16, val);
			},
		IOBinding::Memory(ref h) => {
			*h.as_mut::<u8>(ofs) = val;
			},
		}
	}
	/// Write a single u32 to the binding
	pub unsafe fn write_32(&self, ofs: usize, val: u32)
	{
		match *self
		{
		IOBinding::IO(base, s) => {
			assert!( ofs < s as usize );
			::arch::x86_io::outl(base + ofs as u16, val);
			},
		IOBinding::Memory(ref h) => {
			*h.as_mut::<u32>(ofs) = val;
			},
		}
	}
}

impl ::core::fmt::Debug for IOBinding
{
	fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
		match *self
		{
		IOBinding::IO(b, s) => write!(f, "IO({:#x}+{:#x})", b, s),
		IOBinding::Memory(ref h) => write!(f, "Memory({:?})", h),
		}
	}
}

// vim: ft=rust
<commit_msg>Core/devmgr - Use volatile stores in MMIO handling<commit_after>// "Tifflin" Kernel
// - By John Hodge (thePowersGang)
//
// Core/device_manager.rs
// - Core device manager

use _common::*;
use sync::Mutex;
use lib::Queue;

module_define!{DeviceManager, [arch], init}

/// A semi-arbitatry integer denoting how well a driver handles a specific device
pub type DriverHandleLevel = u32;

/// IO range binding
pub enum IOBinding
{
	/// Memory-mapped IO space
	Memory(::memory::virt::AllocHandle),
	/// x86 IO bus (Base and offset)
	IO(u16,u16),
}

/// Interface a bus manager instance
pub trait BusManager:
	Send + Sync
{
	/// Returns the textual name of the bus type (e.g. "pci")
	fn bus_type(&self) -> &str;
	/// Returns a list of valid attributes for BusDevice::get_attr
	fn get_attr_names(&self) -> &[&str];
}

/// Interface to a device on a bus
pub trait BusDevice:
	Send
{
	/// Returns the device's address on the parent bus
	fn addr(&self) -> u32;
	/// Retrurns the specified attribute (or 0, if invalid)
	fn get_attr(&self, name: &str) -> u32;
	/// Set the power state of this device
	fn set_power(&mut self, state: bool);	// TODO: Power state enum for Off,Standby,Low,On
	/// Bind to the specified IO block (meaning of `block_id` depends on the bus)
	fn bind_io(&mut self, block_id: usize) -> IOBinding;
	/// Obtain the specified interrupt vector
	fn get_irq(&mut self, idx: usize) -> u32;
}

// TODO: Change this to instead be a structure with a bound Fn reference
// - Structure defines bus type and a set of attribute names/values/masks
/// Abstract driver for a device (creates instances when passed a device)
pub trait Driver:
	Send + Sync
{
	/// Driver's name
	fn name(&self) -> &str;
	/// Bus type the driver binds against (matches value from `BusManager::bus_type`)
	fn bus_type(&self) -> &str;
	/// Return the handling level of this driver for the specified device
	fn handles(&self, bus_dev: &BusDevice) -> DriverHandleLevel;
	/// Requests that the driver bind itself to the specified device
	fn bind(&self, bus_dev: &mut BusDevice) -> Box<DriverInstance>;
}

/// Driver instance (maps directly to a device)
pub trait DriverInstance:
	Send
{
}

/// Internal representation of a device on a bus
struct Device
{
	bus_dev: Box<BusDevice>,
	driver: Option<(Box<DriverInstance>, DriverHandleLevel)>,
	//attribs: Vec<u32>,
}

/// Internal representation of a bus
struct Bus
{
	manager: &'static BusManager,
	devices: Vec<Device>,
}

/// List of registered busses on the system
#[allow(non_upper_case_globals)]
static s_root_busses: Mutex<Queue<Bus>> = mutex_init!(queue_init!());

/// List of registered drivers
#[allow(non_upper_case_globals)]
static s_driver_list: Mutex<Queue<&'static Driver>> = mutex_init!( queue_init!() );

fn init()
{
	// Do nothing!
}

/// Register a bus with the device manager
///
/// Creates a new internal representation of the bus, containg the passed set of devices.
pub fn register_bus(manager: &'static BusManager, devices: Vec<Box<BusDevice>>)
{
	let bus = Bus {
		manager: manager,
		// For each device, locate a driver
		devices: devices.into_iter().map(|mut d| Device {
			driver: find_driver(manager, &mut *d),
			//attribs: Vec::new(),
			bus_dev: d,
			}).collect(),
		};
	s_root_busses.lock().push(bus);
}

/// Registers a driver with the device manger
pub fn register_driver(driver: &'static Driver)
{
	s_driver_list.lock().push(driver);
	log_debug!("Registering driver {}", driver.name());
	// Iterate known devices and spin up instances if needed
	for bus in s_root_busses.lock().iter_mut()
	{
		log_trace!("bus type {}", bus.manager.bus_type());
		for dev in bus.devices.iter_mut()
		{
			let rank = driver.handles(&*dev.bus_dev);
			log_debug!("rank = {:?}", rank);
			if rank == 0
			{
				// SKIP!
			}
			else if dev.driver.is_some()
			{
				let bind = dev.driver.as_ref().unwrap();
				let cur_rank = bind.1;
				if cur_rank > rank
				{
					// Existing driver is better
				}
				else if cur_rank == rank
				{
					// Fight!
				}
				else
				{
					// New driver is better
					panic!("TODO: Unbind driver and bind in new one");
				}
			}
			else
			{
				// Bind new driver
				dev.driver = Some( (driver.bind(&mut *dev.bus_dev), rank) );
			}
		}
	}
}

/**
 * Locate the best registered driver for this device and instanciate it
 */
fn find_driver(bus: &BusManager, bus_dev: &mut BusDevice) -> Option<(Box<DriverInstance>,DriverHandleLevel)>
{
	log_debug!("Finding driver for {}:{:x}", bus.bus_type(), bus_dev.addr());
	let mut best_ranking = 0;
	let mut best_driver = None;
	for driver in s_driver_list.lock().iter()
	{
		if bus.bus_type() == driver.bus_type()
		{
			let ranking = driver.handles(bus_dev);
			if ranking == 0
			{
				// Doesn't handle this device
			}
			else if ranking > best_ranking
			{
				// Best so far
				best_driver = Some( *driver );
				best_ranking = ranking;
			}
			else if ranking == best_ranking
			{
				// A tie, this is not very good
				//log_warning!("Tie for device {}:{:x} between {} and {}",
				//	bus.bus_type(), bus_dev.addr(), driver, best_driver.unwrap());
			}
			else
			{
				// Not as good as current, move along
			}
		}
	}
	best_driver.map(|d| (d.bind(bus_dev), best_ranking))
}

impl IOBinding
{
	/// Returns the x86 IO space base
	pub fn io_base(&self) -> u16 {
		match *self
		{
		IOBinding::IO(base, _size) => base,
		IOBinding::Memory(_) => panic!("Called IOBinding::io_base on IOBinding::Memory"),
		}
	}
	/// Writes a single u8 to the binding
	pub unsafe fn write_8(&self, ofs: usize, val: u8)
	{
		match *self
		{
		IOBinding::IO(base, s) => {
			assert!( ofs < s as usize );
			::arch::x86_io::outb(base + ofs as u16, val);
			},
		IOBinding::Memory(ref h) => {
			::core::intrinsics::volatile_store( h.as_mut::<u8>(ofs), val );
			//*h.as_mut::<u8>(ofs) = val;
			},
		}
	}
	/// Write a single u32 to the binding
	pub unsafe fn write_32(&self, ofs: usize, val: u32)
	{
		match *self
		{
		IOBinding::IO(base, s) => {
			assert!( ofs < s as usize );
			::arch::x86_io::outl(base + ofs as u16, val);
			},
		IOBinding::Memory(ref h) => {
			::core::intrinsics::volatile_store( h.as_mut::<u32>(ofs), val );
			//*h.as_mut::<u32>(ofs) = val;
			},
		}
	}
}

impl ::core::fmt::Debug for IOBinding
{
	fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
		match *self
		{
		IOBinding::IO(b, s) => write!(f, "IO({:#x}+{:#x})", b, s),
		IOBinding::Memory(ref h) => write!(f, "Memory({:?})", h),
		}
	}
}

// vim: ft=rust
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::rc::Rc;

use board::Board;
use game_elements::{SnakeElement, new_snake_element};

pub struct Snake {
    head  : Rc<RefCell<SnakeElement>>,
    tail  : Rc<RefCell<SnakeElement>>,
    board : Rc<RefCell<Board>>
}

impl Snake {
    pub fn move_up(&mut self) -> bool {
        let new_x;
        let new_y;
        {
            let head  = self.head.borrow();
            let board = self.board.borrow();
            new_x = head.get_x();
            let mut y = head.get_y();
            if y == 0 { y = board.get_height(); }
            new_y = y - 1;
        }
        self.do_move(new_snake_element(new_x, new_y))
    }

    pub fn move_down(&mut self) -> bool {
        let new_x;
        let new_y;
        {
            let head  = self.head.borrow();
            let board = self.board.borrow();
            new_x = head.get_x();
            let mut y = head.get_y() + 1;
            if y == board.get_height() { y = 0; }
            new_y = y;
        }
        self.do_move(new_snake_element(new_x, new_y))
    }

    pub fn move_left(&mut self) -> bool {
        let new_x;
        let new_y;
        {
            let head  = self.head.borrow();
            let board = self.board.borrow();
            let mut x = head.get_x();
            if x == 0 { x = board.get_width() }
            new_x = x - 1;
            new_y = head.get_y();
        }
        self.do_move(new_snake_element(new_x, new_y))
    }

    pub fn move_right(&mut self) -> bool {
        let new_x;
        let new_y;
        {
            let head  = self.head.borrow();
            let board = self.board.borrow();
            let mut x = head.get_x() + 1;
            if x == board.get_width() { x = 0; }
            new_x = x;
            new_y = head.get_y();
        }
        self.do_move(new_snake_element(new_x, new_y))
    }

    fn do_move(&mut self, new_head: Rc<RefCell<SnakeElement>>) -> bool {
        let x;
        let y;
        {
            let mut nh = new_head.borrow_mut();
            x = nh.get_x();
            y = nh.get_y();
        }
        {
            let mut head = self.head.borrow_mut();
            head.set_prev(new_head.clone());
        }
        self.head = new_head.clone();

        let mut board = self.board.borrow_mut();

        if board.is_apple(x, y) {
            board.add_apple();
        } else {
            if board.is_snake(x, y) {
                return false;
            }

            board.remove(self.tail.clone());
            let prev;
            {
                let tail = self.tail.borrow();
                prev = tail.get_prev().unwrap();
            }
            self.tail = prev;
        }

        board.add(new_head.clone());

        true
    }
}

pub fn new(x: usize, y: usize, board: Rc<RefCell<Board>>) -> Snake {
    let snake_start = new_snake_element(x, y);
    let snake = Snake {head:  snake_start.clone(),
                       tail:  snake_start.clone(),
                       board: board.clone()};
    let b = board.borrow();
    b.add(snake_start);
    snake
}
<commit_msg>new_head doesn't need to be read mutably<commit_after>// THINGS to complain about: `in` keyword, `move` keyword WTF!!!

use std::cell::RefCell;
use std::rc::Rc;

use board::Board;
use game_elements::{SnakeElement, new_snake_element};

pub struct Snake {
    head  : Rc<RefCell<SnakeElement>>,
    tail  : Rc<RefCell<SnakeElement>>,
    board : Rc<RefCell<Board>>
}

impl Snake {
    pub fn move_up(&mut self) -> bool {
        let new_x;
        let new_y;
        {
            let head  = self.head.borrow();
            let board = self.board.borrow();
            new_x = head.get_x();
            let mut y = head.get_y();
            if y == 0 { y = board.get_height(); }
            new_y = y - 1;
        }
        self.do_move(new_snake_element(new_x, new_y))
    }

    pub fn move_down(&mut self) -> bool {
        let new_x;
        let new_y;
        {
            let head  = self.head.borrow();
            let board = self.board.borrow();
            new_x = head.get_x();
            let mut y = head.get_y() + 1;
            if y == board.get_height() { y = 0; }
            new_y = y;
        }
        self.do_move(new_snake_element(new_x, new_y))
    }

    pub fn move_left(&mut self) -> bool {
        let new_x;
        let new_y;
        {
            let head  = self.head.borrow();
            let board = self.board.borrow();
            let mut x = head.get_x();
            if x == 0 { x = board.get_width() }
            new_x = x - 1;
            new_y = head.get_y();
        }
        self.do_move(new_snake_element(new_x, new_y))
    }

    pub fn move_right(&mut self) -> bool {
        let new_x;
        let new_y;
        {
            let head  = self.head.borrow();
            let board = self.board.borrow();
            let mut x = head.get_x() + 1;
            if x == board.get_width() { x = 0; }
            new_x = x;
            new_y = head.get_y();
        }
        self.do_move(new_snake_element(new_x, new_y))
    }

    fn do_move(&mut self, new_head: Rc<RefCell<SnakeElement>>) -> bool {
        let x;
        let y;
        {
            let nh = new_head.borrow();
            x = nh.get_x();
            y = nh.get_y();
        }
        {
            let mut head = self.head.borrow_mut();
            head.set_prev(new_head.clone());
        }
        self.head = new_head.clone();

        let mut board = self.board.borrow_mut();

        if board.is_apple(x, y) {
            board.add_apple();
        } else {
            if board.is_snake(x, y) {
                return false;
            }

            board.remove(self.tail.clone());
            let prev;
            {
                let tail = self.tail.borrow();
                prev = tail.get_prev().unwrap();
            }
            self.tail = prev;
        }

        board.add(new_head.clone());

        true
    }
}

pub fn new(x: usize, y: usize, board: Rc<RefCell<Board>>) -> Snake {
    let snake_start = new_snake_element(x, y);
    let snake = Snake {head:  snake_start.clone(),
                       tail:  snake_start.clone(),
                       board: board.clone()};
    let b = board.borrow();
    b.add(snake_start);
    snake
}
<|endoftext|>


<commit_before>// Implements http://rosettacode.org/wiki/99_Bottles_of_Beer
use std::string::String;

#[cfg(not(test))]
fn main() {
	for num_bottles in std::iter::range_step_inclusive(99i, 1, -1) {
        let num_bottles=num_bottles as uint;

		println!("{}",bottles_line(num_bottles, true));
		println!("{}",bottles_line(num_bottles, false));
		println!("Take one down, pass it around...");
		println!("{}",bottles_line(num_bottles - 1, true));
		println!("-----------------------------------");
	}
}

fn bottles_line(num_bottles: uint, on_the_wall: bool) -> String {
    let mut ln = match num_bottles {
        0 => "No bottles".to_string(),
        1 => "One bottle".to_string(),
        n => format!("{} bottles", n)
    };

    ln.push_str(" of beer");

	if on_the_wall {
		ln.push_str(" on the wall!");
	}

	ln.push_str("\n");
    ln
}

#[test]
fn gen_bottle_line() {
    let ln=bottles_line(42, false);
    let ln2=bottles_line(42, true);

    assert_eq!(ln.as_slice(),"42 bottles of beer\n");
    assert_eq!(ln2.as_slice(),"42 bottles of beer on the wall!\n");
}
<commit_msg>99 bottles: style<commit_after>// Implements http://rosettacode.org/wiki/99_Bottles_of_Beer
use std::string::String;

#[cfg(not(test))]
fn main() {
	for num_bottles in std::iter::range_inclusive(1u, 99).rev() {
        println!("{}", bottles_line(num_bottles, true));
		println!("{}", bottles_line(num_bottles, false));
		println!("Take one down, pass it around...");
		println!("{}", bottles_line(num_bottles - 1, true));
		println!("-----------------------------------");
	}
}

fn bottles_line(num_bottles: uint, on_the_wall: bool) -> String {
    let tail = match on_the_wall {
        true => "of beer on the wall!\n",
        false => "of beer\n"
    };

    match num_bottles {
        0 => format!("No bottles {}", tail),
        1 => format!("One bottle {}", tail),
        n => format!("{} bottles {}", n, tail)
    }
}

#[test]
fn gen_bottle_line() {
    let ln = bottles_line(42, false);
    let ln2 = bottles_line(42, true);

    assert_eq!(ln.as_slice(), "42 bottles of beer\n");
    assert_eq!(ln2.as_slice(), "42 bottles of beer on the wall!\n");
}
<|endoftext|>


<commit_before>//! Thread-local storage for JavaScript threads.
//!
//! At runtime, an instance of a Node.js addon can contain its own local storage,
//! which can then be shared and accessed as needed from Rust in a Neon module. This can
//! be useful for setting up long-lived state that needs to be shared between calls
//! of an addon's APIs.
//!
//! For example, an addon may wish to track the [thread ID][threadId] of each of its
//! instances:
//!
//! ```
//! # use neon::prelude::*;
//! # use neon::thread::LocalKey;
//! static THREAD_ID: LocalKey<u32> = LocalKey::new();
//!
//! pub fn thread_id<'cx, C: Context<'cx>>(cx: &mut C) -> NeonResult<u32> {
//!     THREAD_ID.get_or_try_init(cx, |cx| {
//!         let global = cx.global();
//!         let require: Handle<JsFunction> = global.get(cx, "require")?;
//!         let worker: Handle<JsObject> = require.call_with(cx)
//!             .arg(cx.string("node:worker_threads"))
//!             .apply(cx)?;
//!         let thread_id: Handle<JsNumber> = worker.get(cx, "threadId")?;
//!         Ok(thread_id.value(cx) as u32)
//!     }).cloned()
//! }
//! ```
//!
//! ### The Addon Lifecycle
//!
//! For some use cases, a single shared global constant stored in a `static` variable
//! might be sufficient:
//!
//! ```
//! static MY_CONSTANT: &'static str = "hello Neon";
//! ```
//!
//! This variable will be allocated when the addon is first loaded into the Node.js
//! process. This works fine for single-threaded applications, or global immutable
//! data.
//!
//! However, since the addition of [worker threads][workers] in Node v10,
//! modules can be instantiated multiple times in a single Node process. This means
//! that while the dynamically-loaded binary library (i.e., the Rust implementation of
//! the addon) is only loaded once in the running process, but its `main()` function
//! is executed multiple times with distinct module objects, once per application thread:
//!
//! ![The Node.js addon lifecycle, described in detail below.][lifecycle]
//!
//! This means that any thread-local data needs to be initialized separately for each
//! instance of the addon. This module provides a simple container type, [`LocalKey`](LocalKey),
//! for allocating and initializing thread-local data. (Technically, this data is stored in the
//! addon's [module instance][environment], which is equivalent to being thread-local.)
//!
//! A common example is when an addon defines a custom datatype, which cannot be used across
//! separate threads. The constructor for the datatype can be [rooted](crate::handle::Root) and
//! saved in thread-local storage:
//!
//! ```
//! # use neon::prelude::*;
//! # use neon::thread::LocalKey;
//! # fn initialize_my_datatype<'cx, C: Context<'cx>>(cx: &mut C) -> JsResult<'cx, JsFunction> { unimplemented!() }
//! static MY_CONSTRUCTOR: LocalKey<Root<JsFunction>> = LocalKey::new();
//!
//! pub fn my_constructor<'cx, C: Context<'cx>>(cx: &mut C) -> JsResult<'cx, JsFunction> {
//!     let constructor = MY_CONSTRUCTOR.get_or_try_init(cx, |cx| {
//!         let constructor: Handle<JsFunction> = initialize_my_datatype(cx)?;
//!         Ok(constructor.root(cx))
//!     })?;
//!     Ok(constructor.to_inner(cx))
//! }
//! ```
//!
//! Notice that if this code were implemented without a `LocalKey`, it would fail whenever
//! one thread stores an instance of the constructor and a different thread attempts to
//! access it with the call to [`to_inner()`](crate::handle::Root::to_inner).
//!
//! ### When to Use Thread-Local Storage
//!
//! Single-threaded applications don't generally need to worry about thread-local data.
//! There are two cases where Neon apps should consider storing static data in a
//! `LocalKey` storage cell:
//!
//! - **Multi-threaded applications:** If your Node application uses the `Worker`
//!   API, you'll want to store any static data that might get access from multiple
//!   threads in thread-local data.
//! - **Libraries:** If your addon is part of a library that could be used by multiple
//!   applications, you'll want to store static data in thread-local data in case the
//!   addon ends up instantiated by multiple threads in some future application.
//!
//! ### Why Not Use Standard TLS?
//!
//! Because the JavaScript engine may not tie JavaScript threads 1:1 to system threads,
//! it is recommended to use this module instead of the Rust standard thread-local storage
//! when tying data to a JavaScript thread.
//!
//! [environment]: https://nodejs.org/api/n-api.html#environment-life-cycle-apis
//! [lifecycle]: https://raw.githubusercontent.com/neon-bindings/neon/main/doc/lifecycle.png
//! [workers]: https://nodejs.org/api/worker_threads.html
//! [threadId]: https://nodejs.org/api/worker_threads.html#workerthreadid

use std::any::Any;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicUsize, Ordering};

use once_cell::sync::OnceCell;

use crate::context::Context;
use crate::lifecycle::LocalCell;

static COUNTER: AtomicUsize = AtomicUsize::new(0);

fn next_id() -> usize {
    COUNTER.fetch_add(1, Ordering::SeqCst)
}

/// A JavaScript thread-local container that owns its contents, similar to
/// [`std::thread::LocalKey`](std::thread::LocalKey) but tied to a JavaScript thread rather
/// than a system thread.
///
/// ### Initialization and Destruction
///
/// Initialization is dynamically performed on the first call to one of the `init` methods
/// of `LocalKey`, and values that implement [`Drop`](std::ops::Drop) get destructed when
/// the JavaScript thread exits, i.e. when a worker thread terminates or the main thread
/// terminates on process exit.
#[derive(Default)]
pub struct LocalKey<T> {
    _type: PhantomData<T>,
    id: OnceCell<usize>,
}

impl<T> LocalKey<T> {
    /// Creates a new local value. This method is `const`, so it can be assigned to
    /// static variables.
    pub const fn new() -> Self {
        Self {
            _type: PhantomData,
            id: OnceCell::new(),
        }
    }

    fn id(&self) -> usize {
        *self.id.get_or_init(next_id)
    }
}

impl<T: Any + Send + 'static> LocalKey<T> {
    /// Gets the current value of the cell. Returns `None` if the cell has not
    /// yet been initialized.
    pub fn get<'cx, 'a, C>(&self, cx: &'a mut C) -> Option<&'cx T>
    where
        C: Context<'cx>,
    {
        // Unwrap safety: The type bound LocalKey<T> and the fact that every LocalKey has a unique
        // id guarantees that the cell is only ever assigned instances of type T.
        let r: Option<&T> =
            LocalCell::get(cx, self.id()).map(|value| value.downcast_ref().unwrap());

        // Safety: Since the Box is immutable and heap-allocated, it's guaranteed not to
        // move or change for the duration of the context.
        unsafe { std::mem::transmute::<Option<&'a T>, Option<&'cx T>>(r) }
    }

    /// Gets the current value of the cell, initializing it with the result of
    /// calling `f` if it has not yet been initialized.
    pub fn get_or_init<'cx, 'a, C, F>(&self, cx: &'a mut C, f: F) -> &'cx T
    where
        C: Context<'cx>,
        F: FnOnce() -> T,
    {
        // Unwrap safety: The type bound LocalKey<T> and the fact that every LocalKey has a unique
        // id guarantees that the cell is only ever assigned instances of type T.
        let r: &T = LocalCell::get_or_init(cx, self.id(), || Box::new(f()))
            .downcast_ref()
            .unwrap();

        // Safety: Since the Box is immutable and heap-allocated, it's guaranteed not to
        // move or change for the duration of the context.
        unsafe { std::mem::transmute::<&'a T, &'cx T>(r) }
    }

    /// Gets the current value of the cell, initializing it with the result of
    /// calling `f` if it has not yet been initialized. Returns `Err` if the
    /// callback triggers a JavaScript exception.
    ///
    /// # Panics
    ///
    /// During the execution of `f`, calling any methods on this `LocalKey` that
    /// attempt to initialize it will panic.
    pub fn get_or_try_init<'cx, 'a, C, E, F>(&self, cx: &'a mut C, f: F) -> Result<&'cx T, E>
    where
        C: Context<'cx>,
        F: FnOnce(&mut C) -> Result<T, E>,
    {
        // Unwrap safety: The type bound LocalKey<T> and the fact that every LocalKey has a unique
        // id guarantees that the cell is only ever assigned instances of type T.
        let r: &T = LocalCell::get_or_try_init(cx, self.id(), |cx| Ok(Box::new(f(cx)?)))?
            .downcast_ref()
            .unwrap();

        // Safety: Since the Box is immutable and heap-allocated, it's guaranteed not to
        // move or change for the duration of the context.
        Ok(unsafe { std::mem::transmute::<&'a T, &'cx T>(r) })
    }
}

impl<T: Any + Send + Default + 'static> LocalKey<T> {
    /// Gets the current value of the cell, initializing it with the default value
    /// if it has not yet been initialized.
    pub fn get_or_init_default<'cx, 'a, C>(&self, cx: &'a mut C) -> &'cx T
    where
        C: Context<'cx>,
    {
        self.get_or_init(cx, Default::default)
    }
}
<commit_msg>Link to `neon::main` docs in doc comment<commit_after>//! Thread-local storage for JavaScript threads.
//!
//! At runtime, an instance of a Node.js addon can contain its own local storage,
//! which can then be shared and accessed as needed from Rust in a Neon module. This can
//! be useful for setting up long-lived state that needs to be shared between calls
//! of an addon's APIs.
//!
//! For example, an addon may wish to track the [thread ID][threadId] of each of its
//! instances:
//!
//! ```
//! # use neon::prelude::*;
//! # use neon::thread::LocalKey;
//! static THREAD_ID: LocalKey<u32> = LocalKey::new();
//!
//! pub fn thread_id<'cx, C: Context<'cx>>(cx: &mut C) -> NeonResult<u32> {
//!     THREAD_ID.get_or_try_init(cx, |cx| {
//!         let global = cx.global();
//!         let require: Handle<JsFunction> = global.get(cx, "require")?;
//!         let worker: Handle<JsObject> = require.call_with(cx)
//!             .arg(cx.string("node:worker_threads"))
//!             .apply(cx)?;
//!         let thread_id: Handle<JsNumber> = worker.get(cx, "threadId")?;
//!         Ok(thread_id.value(cx) as u32)
//!     }).cloned()
//! }
//! ```
//!
//! ### The Addon Lifecycle
//!
//! For some use cases, a single shared global constant stored in a `static` variable
//! might be sufficient:
//!
//! ```
//! static MY_CONSTANT: &'static str = "hello Neon";
//! ```
//!
//! This variable will be allocated when the addon is first loaded into the Node.js
//! process. This works fine for single-threaded applications, or global immutable
//! data.
//!
//! However, since the addition of [worker threads][workers] in Node v10,
//! modules can be instantiated multiple times in a single Node process. This means
//! that while the dynamically-loaded binary library (i.e., the Rust implementation of
//! the addon) is only loaded once in the running process, but its [`#[main]`](neon::main) function
//! is executed multiple times with distinct module objects, once per application thread:
//!
//! ![The Node.js addon lifecycle, described in detail below.][lifecycle]
//!
//! This means that any thread-local data needs to be initialized separately for each
//! instance of the addon. This module provides a simple container type, [`LocalKey`](LocalKey),
//! for allocating and initializing thread-local data. (Technically, this data is stored in the
//! addon's [module instance][environment], which is equivalent to being thread-local.)
//!
//! A common example is when an addon defines a custom datatype, which cannot be used across
//! separate threads. The constructor for the datatype can be [rooted](crate::handle::Root) and
//! saved in thread-local storage:
//!
//! ```
//! # use neon::prelude::*;
//! # use neon::thread::LocalKey;
//! # fn initialize_my_datatype<'cx, C: Context<'cx>>(cx: &mut C) -> JsResult<'cx, JsFunction> { unimplemented!() }
//! static MY_CONSTRUCTOR: LocalKey<Root<JsFunction>> = LocalKey::new();
//!
//! pub fn my_constructor<'cx, C: Context<'cx>>(cx: &mut C) -> JsResult<'cx, JsFunction> {
//!     let constructor = MY_CONSTRUCTOR.get_or_try_init(cx, |cx| {
//!         let constructor: Handle<JsFunction> = initialize_my_datatype(cx)?;
//!         Ok(constructor.root(cx))
//!     })?;
//!     Ok(constructor.to_inner(cx))
//! }
//! ```
//!
//! Notice that if this code were implemented without a `LocalKey`, it would fail whenever
//! one thread stores an instance of the constructor and a different thread attempts to
//! access it with the call to [`to_inner()`](crate::handle::Root::to_inner).
//!
//! ### When to Use Thread-Local Storage
//!
//! Single-threaded applications don't generally need to worry about thread-local data.
//! There are two cases where Neon apps should consider storing static data in a
//! `LocalKey` storage cell:
//!
//! - **Multi-threaded applications:** If your Node application uses the `Worker`
//!   API, you'll want to store any static data that might get access from multiple
//!   threads in thread-local data.
//! - **Libraries:** If your addon is part of a library that could be used by multiple
//!   applications, you'll want to store static data in thread-local data in case the
//!   addon ends up instantiated by multiple threads in some future application.
//!
//! ### Why Not Use Standard TLS?
//!
//! Because the JavaScript engine may not tie JavaScript threads 1:1 to system threads,
//! it is recommended to use this module instead of the Rust standard thread-local storage
//! when tying data to a JavaScript thread.
//!
//! [environment]: https://nodejs.org/api/n-api.html#environment-life-cycle-apis
//! [lifecycle]: https://raw.githubusercontent.com/neon-bindings/neon/main/doc/lifecycle.png
//! [workers]: https://nodejs.org/api/worker_threads.html
//! [threadId]: https://nodejs.org/api/worker_threads.html#workerthreadid

use std::any::Any;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicUsize, Ordering};

use once_cell::sync::OnceCell;

use crate::context::Context;
use crate::lifecycle::LocalCell;

static COUNTER: AtomicUsize = AtomicUsize::new(0);

fn next_id() -> usize {
    COUNTER.fetch_add(1, Ordering::SeqCst)
}

/// A JavaScript thread-local container that owns its contents, similar to
/// [`std::thread::LocalKey`](std::thread::LocalKey) but tied to a JavaScript thread rather
/// than a system thread.
///
/// ### Initialization and Destruction
///
/// Initialization is dynamically performed on the first call to one of the `init` methods
/// of `LocalKey`, and values that implement [`Drop`](std::ops::Drop) get destructed when
/// the JavaScript thread exits, i.e. when a worker thread terminates or the main thread
/// terminates on process exit.
#[derive(Default)]
pub struct LocalKey<T> {
    _type: PhantomData<T>,
    id: OnceCell<usize>,
}

impl<T> LocalKey<T> {
    /// Creates a new local value. This method is `const`, so it can be assigned to
    /// static variables.
    pub const fn new() -> Self {
        Self {
            _type: PhantomData,
            id: OnceCell::new(),
        }
    }

    fn id(&self) -> usize {
        *self.id.get_or_init(next_id)
    }
}

impl<T: Any + Send + 'static> LocalKey<T> {
    /// Gets the current value of the cell. Returns `None` if the cell has not
    /// yet been initialized.
    pub fn get<'cx, 'a, C>(&self, cx: &'a mut C) -> Option<&'cx T>
    where
        C: Context<'cx>,
    {
        // Unwrap safety: The type bound LocalKey<T> and the fact that every LocalKey has a unique
        // id guarantees that the cell is only ever assigned instances of type T.
        let r: Option<&T> =
            LocalCell::get(cx, self.id()).map(|value| value.downcast_ref().unwrap());

        // Safety: Since the Box is immutable and heap-allocated, it's guaranteed not to
        // move or change for the duration of the context.
        unsafe { std::mem::transmute::<Option<&'a T>, Option<&'cx T>>(r) }
    }

    /// Gets the current value of the cell, initializing it with the result of
    /// calling `f` if it has not yet been initialized.
    pub fn get_or_init<'cx, 'a, C, F>(&self, cx: &'a mut C, f: F) -> &'cx T
    where
        C: Context<'cx>,
        F: FnOnce() -> T,
    {
        // Unwrap safety: The type bound LocalKey<T> and the fact that every LocalKey has a unique
        // id guarantees that the cell is only ever assigned instances of type T.
        let r: &T = LocalCell::get_or_init(cx, self.id(), || Box::new(f()))
            .downcast_ref()
            .unwrap();

        // Safety: Since the Box is immutable and heap-allocated, it's guaranteed not to
        // move or change for the duration of the context.
        unsafe { std::mem::transmute::<&'a T, &'cx T>(r) }
    }

    /// Gets the current value of the cell, initializing it with the result of
    /// calling `f` if it has not yet been initialized. Returns `Err` if the
    /// callback triggers a JavaScript exception.
    ///
    /// # Panics
    ///
    /// During the execution of `f`, calling any methods on this `LocalKey` that
    /// attempt to initialize it will panic.
    pub fn get_or_try_init<'cx, 'a, C, E, F>(&self, cx: &'a mut C, f: F) -> Result<&'cx T, E>
    where
        C: Context<'cx>,
        F: FnOnce(&mut C) -> Result<T, E>,
    {
        // Unwrap safety: The type bound LocalKey<T> and the fact that every LocalKey has a unique
        // id guarantees that the cell is only ever assigned instances of type T.
        let r: &T = LocalCell::get_or_try_init(cx, self.id(), |cx| Ok(Box::new(f(cx)?)))?
            .downcast_ref()
            .unwrap();

        // Safety: Since the Box is immutable and heap-allocated, it's guaranteed not to
        // move or change for the duration of the context.
        Ok(unsafe { std::mem::transmute::<&'a T, &'cx T>(r) })
    }
}

impl<T: Any + Send + Default + 'static> LocalKey<T> {
    /// Gets the current value of the cell, initializing it with the default value
    /// if it has not yet been initialized.
    pub fn get_or_init_default<'cx, 'a, C>(&self, cx: &'a mut C) -> &'cx T
    where
        C: Context<'cx>,
    {
        self.get_or_init(cx, Default::default)
    }
}
<|endoftext|>


<commit_before>extern crate getopts;
#[macro_use]
extern crate modparse_derive;
#[macro_use]
extern crate derive_builder;

#[macro_use]
extern crate torchrs;

use torchrs::autograd::{Variable, VariableArgs, VarAccess};
use torchrs::tensor::Tensor;
use torchrs::optim;

use torchrs::macros::*;
use torchrs::nn;
use torchrs::nn::{ModuleStruct, ModIntf, ModDelegate, Module};
use torchrs::nn::functional as F;
use torchrs::utils::data as D;

use getopts::Options;
use std::env;

fn MNIST() -> D::Dataset<D::Batch<f32, i64>> {
    unimplemented!()
}

#[derive(Clone, Builder)]
struct NetArgs {
    #[builder(default="64")]
    batch_size: usize,
    #[builder(default="1000")]
    test_batch_size: usize,
    #[builder(default="10")]
    epochs: u32,
    #[builder(default="0.01")]
    lr: f32,
    #[builder(default="0.5")]
    momentum: f32,
    #[builder(default="true")]
    cuda: bool,
    #[builder(default="1")]
    seed: usize,
    #[builder(default="10")]
    log_interval: usize,
}

impl Default for NetArgs {
    fn default() -> Self {
        NetArgsBuilder::default().build().unwrap()
    }
}

fn parse_args() -> NetArgs {
    let cmd_args: Vec<String> = env::args().collect();
    let mut args = NetArgs::default();
    let mut opts = Options::new();
    unimplemented!();
    /*
    	* XXX do parsey stuff
    	*/
    args
}

#[derive(ModParse)]
struct Net<T: Default + Copy> {
    delegate: nn::Module<T>,
    conv1: nn::Conv2d<T>,
    conv2: nn::Conv2d<T>,
    fc1: nn::Linear<T>,
    fc2: nn::Linear<T>,
}

impl<T: Default + Copy> Net<T> {
    pub fn new() -> Net<T> {
        Net {
                delegate: nn::Module::new(),
                conv1: nn::Conv2d::build(1, 10, (5, 5)).done(),
                conv2: nn::Conv2d::build(10, 20, (5, 5)).done(),
                fc1: nn::Linear::build(320, 50).done(),
                fc2: nn::Linear::build(50, 10).done(),
            }
            .init_module()
    }
}
impl_mod_delegate!(Net);

// The forward operations could take on one of two implementations.
// The first supporting a near verbatim version of the python
// implementation, and the second supporting a slightly more
// idiomatic to Rust method chaining.

// a) as a near verbatim implementation of the python version
impl ModIntf<f32> for Net<f32> {
    fn forward(&mut self, args: &mut Variable<f32>) -> Variable<f32> {
        let training = self.delegate.training;
        let pool_val = F::MaxPool2dArgs::default();
        let dropout_val = F::DropoutArgs::default();
        let mut x = F::relu(&F::max_pool2d(&self.conv1.f(args), (2, 2), &pool_val));
        let mut x = F::relu(&F::max_pool2d(&F::dropout2d(&self.conv2.f(&mut x), &dropout_val),
                                           (2, 2),
                                           &pool_val));
        let mut x = x.view(&[-1, 320]);
        let x = F::relu(&self.fc1.f(&mut x));
        let mut x = F::dropout(&x, &dropout_val);
        let x = self.fc2.f(&mut x);
        F::log_softmax(&x)
    }
}

fn train(model: &mut Net<f32>,
         args: &NetArgs,
         train_loader: &D::BatchLoader<f32, i64>,
         epoch: u32,
         optimizer: &mut optim::OptIntf) {
    model.train();
    for (batch_idx, (ref data, ref target)) in train_loader.iter().enumerate() {
        let (mut data, target) = if args.cuda {
            (Variable::new(data.cuda(None).clone()), Variable::new(target.cuda(None).clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        optimizer.zero_grad();
        let output = model.f(&mut data);
        let mut loss = F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        loss.backward();
        optimizer.step();
        if batch_idx % args.log_interval == 0 {
            println!("Train Epoch: {} [{}/{} ({:.0}%)]\tLoss: {:.6}",
                     epoch,
                     batch_idx * data.data().len(),
                     train_loader.dataset.len(),
                     100. * (batch_idx as f32) / train_loader.len() as f32,
                     loss.data()[0]);
        }
    }
}

fn test(model: &mut Net<f32>, args: &NetArgs, test_loader: &D::BatchLoader<f32, i64>) {
    model.eval();
    let mut test_loss = 0.;
    let mut correct = 0;
    for (ref data, ref target) in test_loader.iter() {
        let varargs = VariableArgs {
            volatile: true,
            ..Default::default()
        };
        let (mut data, mut target) = if args.cuda {
            (Variable::new_args(data.cuda(None).clone(), &varargs),
             Variable::new(target.cuda(None).clone()))
        } else {
            (Variable::new_args(data.clone(), &varargs), Variable::new(target.clone()))
        };
        let mut output = model.f(&mut data);
        test_loss += F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        let pred = output.data().max_reduce(1).1;
        correct += pred.eq_tensor(&*target.data()).cpu().sum() as u32;
    }
    test_loss /= test_loader.len() as f32;
    println!("\nTest set: Average loss: {:.4}, Accuracy: {}/{} ({:.0}%)\n",
             test_loss,
             correct,
             test_loader.dataset.len(),
             100. * (correct as f32) / (test_loader.dataset.len() as f32))

}

fn main() {
    // no data loaders yet so demonstrate with Vec placeholder
    let train_loader: D::BatchLoader<f32, i64> = D::DataLoader::new(MNIST(),
                                                                    D::DataLoaderArgs::default());
    let test_loader: D::BatchLoader<f32, i64> = D::DataLoader::new(MNIST(),
                                                                   D::DataLoaderArgs::default());
    let args = parse_args();
    let mut optimizer = optim::SGD::new();

    let mut model = Net::new();
    for epoch in 1..args.epochs + 1 {
        train(&mut model, &args, &train_loader, epoch, &mut optimizer);
        test(&mut model, &args, &test_loader);
    }
}
<commit_msg>fix a few warnings from excess mut use, set training in dropout<commit_after>extern crate getopts;
#[macro_use]
extern crate modparse_derive;
#[macro_use]
extern crate derive_builder;

#[macro_use]
extern crate torchrs;

use torchrs::autograd::{Variable, VariableArgs, VarAccess};
use torchrs::tensor::Tensor;
use torchrs::optim;
use torchrs::macros::*;
use torchrs::nn;
use torchrs::nn::{ModuleStruct, ModIntf, ModDelegate, Module};
use torchrs::nn::functional as F;
use torchrs::utils::data as D;

use getopts::Options;
use std::env;

fn MNIST() -> D::Dataset<D::Batch<f32, i64>> {
    unimplemented!()
}

#[derive(Clone, Builder)]
struct NetArgs {
    #[builder(default="64")]
    batch_size: usize,
    #[builder(default="1000")]
    test_batch_size: usize,
    #[builder(default="10")]
    epochs: u32,
    #[builder(default="0.01")]
    lr: f32,
    #[builder(default="0.5")]
    momentum: f32,
    #[builder(default="true")]
    cuda: bool,
    #[builder(default="1")]
    seed: usize,
    #[builder(default="10")]
    log_interval: usize,
}

impl Default for NetArgs {
    fn default() -> Self {
        NetArgsBuilder::default().build().unwrap()
    }
}

fn parse_args() -> NetArgs {
    unimplemented!();
    let cmd_args: Vec<String> = env::args().collect();
    let mut args = NetArgs::default();
    let mut opts = Options::new();
    /*
    	* XXX do parsey stuff
    	*/
    args
}

#[derive(ModParse)]
struct Net<T: Default + Copy> {
    delegate: nn::Module<T>,
    conv1: nn::Conv2d<T>,
    conv2: nn::Conv2d<T>,
    fc1: nn::Linear<T>,
    fc2: nn::Linear<T>,
}

impl<T: Default + Copy> Net<T> {
    pub fn new() -> Net<T> {
        Net {
                delegate: nn::Module::new(),
                conv1: nn::Conv2d::build(1, 10, (5, 5)).done(),
                conv2: nn::Conv2d::build(10, 20, (5, 5)).done(),
                fc1: nn::Linear::build(320, 50).done(),
                fc2: nn::Linear::build(50, 10).done(),
            }
            .init_module()
    }
}
impl_mod_delegate!(Net);

// The forward operations could take on one of two implementations.
// The first supporting a near verbatim version of the python
// implementation, and the second supporting a slightly more
// idiomatic to Rust method chaining.

// a) as a near verbatim implementation of the python version
impl ModIntf<f32> for Net<f32> {
    fn forward(&mut self, args: &mut Variable<f32>) -> Variable<f32> {
        let pool_val = F::MaxPool2dArgs::default();
        let mut dropout_val = F::DropoutArgs::default();
        dropout_val.training = self.delegate.training;
        let mut x = F::relu(&F::max_pool2d(&self.conv1.f(args), (2, 2), &pool_val));
        x = F::relu(&F::max_pool2d(&F::dropout2d(&self.conv2.f(&mut x), &dropout_val),
                                           (2, 2),
                                           &pool_val));
        x = x.view(&[-1, 320]);
        x = F::relu(&self.fc1.f(&mut x));
        x = F::dropout(&x, &dropout_val);
        x = self.fc2.f(&mut x);
        F::log_softmax(&x)
    }
}

fn train(model: &mut Net<f32>,
         args: &NetArgs,
         train_loader: &D::BatchLoader<f32, i64>,
         epoch: u32,
         optimizer: &mut optim::OptIntf) {
    model.train();
    for (batch_idx, (ref data, ref target)) in train_loader.iter().enumerate() {
        let (mut data, target) = if args.cuda {
            (Variable::new(data.cuda(None).clone()), Variable::new(target.cuda(None).clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        optimizer.zero_grad();
        let output = model.f(&mut data);
        let mut loss = F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        loss.backward();
        optimizer.step();
        if batch_idx % args.log_interval == 0 {
            println!("Train Epoch: {} [{}/{} ({:.0}%)]\tLoss: {:.6}",
                     epoch,
                     batch_idx * data.data().len(),
                     train_loader.dataset.len(),
                     100. * (batch_idx as f32) / train_loader.len() as f32,
                     loss.data()[0]);
        }
    }
}

fn test(model: &mut Net<f32>, args: &NetArgs, test_loader: &D::BatchLoader<f32, i64>) {
    model.eval();
    let mut test_loss = 0.;
    let mut correct = 0;
    for (ref data, ref target) in test_loader.iter() {
        let varargs = VariableArgs {
            volatile: true,
            ..Default::default()
        };
        let (mut data, mut target) = if args.cuda {
            (Variable::new_args(data.cuda(None).clone(), &varargs),
             Variable::new(target.cuda(None).clone()))
        } else {
            (Variable::new_args(data.clone(), &varargs), Variable::new(target.clone()))
        };
        let mut output = model.f(&mut data);
        test_loss += F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        let pred = output.data().max_reduce(1).1;
        correct += pred.eq_tensor(&*target.data()).cpu().sum() as u32;
    }
    test_loss /= test_loader.len() as f32;
    println!("\nTest set: Average loss: {:.4}, Accuracy: {}/{} ({:.0}%)\n",
             test_loss,
             correct,
             test_loader.dataset.len(),
             100. * (correct as f32) / (test_loader.dataset.len() as f32))

}

fn main() {
    // no data loaders yet so demonstrate with Vec placeholder
    let train_loader: D::BatchLoader<f32, i64> = D::DataLoader::new(MNIST(),
                                                                    D::DataLoaderArgs::default());
    let test_loader: D::BatchLoader<f32, i64> = D::DataLoader::new(MNIST(),
                                                                   D::DataLoaderArgs::default());
    let args = parse_args();
    let mut optimizer = optim::SGD::new();

    let mut model = Net::new();
    for epoch in 1..args.epochs + 1 {
        train(&mut model, &args, &train_loader, epoch, &mut optimizer);
        test(&mut model, &args, &test_loader);
    }
}
<|endoftext|>


<commit_before>extern crate torchrs;
extern crate getopts;
#[macro_use]
extern crate modparse_derive;
#[macro_use]
extern crate derive_builder;

use torchrs::autograd::{Variable, BackwardArgs};
use torchrs::tensor::Tensor;
use torchrs::optim;

use torchrs::nn;
use torchrs::nn::{ModuleStruct, ModIntf};
use torchrs::nn::functional as F;
use torchrs::utils::data as D;

use getopts::Options;
use std::env;

fn MNIST() -> D::Dataset<D::Batch<f32, i64>> {
    panic!("just a stub")
}

#[derive(ModParse)]
struct Net {
    delegate: nn::Module<f32>,
    conv1: nn::Conv2d<f32>,
    conv2: nn::Conv2d<f32>,
    fc1: nn::Linear<f32>,
    fc2: nn::Linear<f32>,
}

impl Net {
    pub fn new() -> Net {
        let mut t = Net {
            delegate: nn::Module::new(),
            conv1: nn::Conv2d::build(1, 10, 5).done(),
            conv2: nn::Conv2d::build(10, 20, 5).done(),
            fc1: nn::Linear::build(320, 50).done(),
            fc2: nn::Linear::build(50, 10).done(),
        };
        t.init_module();
        t
    }
}
// The forward operations could take on one of two implementations.
// The first supporting a near verbatim version of the python
// implementation, and the second supporting a slightly more
// idiomatic to Rust method chaining.

// a) as a near verbatim implementation of the python version
impl ModIntf<f32> for Net {
    fn forward(&mut self, args: &mut Variable<f32>) -> Variable<f32> {
        let training = self.delegate.training;
        let pool_val = F::MaxPool2dArgs::default();
        let dropout_val = F::DropoutArgs::default();
        let mut x = F::relu(&F::max_pool2d(&self.conv1.f(args), (2, 2), &pool_val));
        let mut x = F::relu(&F::max_pool2d(&F::dropout2d(&self.conv2.f(&mut x), &dropout_val),
                                           (2, 2),
                                           &pool_val));
        let mut x = x.view(&[-1, 320]);
        let x = F::relu(&self.fc1.f(&mut x));
        let mut x = F::dropout(&x, &dropout_val);
        let x = self.fc2.f(&mut x);
        F::log_softmax(&x)
    }
    fn forwardv(&mut self, input: &mut Vec<Variable<f32>>) -> Vec<Variable<f32>> {
        panic!("not valid")
    }

    fn delegate(&mut self) -> &mut nn::Module<f32> {
        &mut self.delegate
    }
}

fn train(model: &mut Net,
         args: &NetArgs,
         train_loader: &D::BatchLoader<f32, i64>,
         epoch: u32,
         optimizer: &mut optim::OptIntf) {
    model.train();
    for (batch_idx, (ref data, ref target)) in train_loader.iter().enumerate() {
        let (mut data, target) = if args.cuda {
            (Variable::new(data.cuda().clone()), Variable::new(target.cuda().clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        optimizer.zero_grad();
        let output = model.f(&mut data);
        let mut loss = F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        loss.backward(&BackwardArgs::default());
        optimizer.step();
        if batch_idx % args.log_interval == 0 {
            println!("Train Epoch: {} [{}/{} ({:.0}%)]\tLoss: {:.6}",
                     epoch,
                     batch_idx * data.data().len(),
                     train_loader.dataset.len(),
                     100. * (batch_idx as f32) / train_loader.len() as f32,
                     loss.data()[0]);
        }
    }
}

fn test(model: &mut Net, args: &NetArgs, test_loader: &D::BatchLoader<f32, i64>, epoch: u32) {
    model.eval();
    let mut test_loss = 0.;
    let mut correct = 0;
    for (ref data, ref target) in test_loader.iter() {
        let (mut data, mut target) = if args.cuda {
            (Variable::new_volatile(data.cuda().clone()), Variable::new(target.cuda().clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        let mut output = model.f(&mut data);
        test_loss += F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        let pred = output.data().reduce_max(1).1;
        correct += pred.tensor_eq(target.data()).cpu().sum();
    }
    test_loss /= test_loader.len() as f32;
    println!("\nTest set: Average loss: {:.4}, Accuracy: {}/{} ({:.0}%)\n",
             test_loss,
             correct,
             test_loader.dataset.len(),
             100. * (correct as f32) / (test_loader.dataset.len() as f32))

}

#[derive(Clone, Builder)]
struct NetArgs {
    #[builder(default="64")]
    batch_size: usize,
    #[builder(default="1000")]
    test_batch_size: usize,
    #[builder(default="10")]
    epochs: u32,
    #[builder(default="0.01")]
    lr: f32,
    #[builder(default="0.5")]
    momentum: f32,
    #[builder(default="true")]
    cuda: bool,
    #[builder(default="1")]
    seed: usize,
    #[builder(default="10")]
    log_interval: usize,
}

impl Default for NetArgs {
    fn default() -> Self {
        NetArgsBuilder::default().build().unwrap()
    }
}

fn parse_args() -> NetArgs {
    let cmd_args: Vec<String> = env::args().collect();
    let mut args = NetArgs::default();
    let mut opts = Options::new();
    /*
    	* XXX do parsey stuff
    	*/
    args
}

fn main() {
    // no data loaders yet so demonstrate with Vec placeholder
    let train_loader: D::BatchLoader<f32, i64> = D::DataLoader::new(MNIST(),
                                                                    D::DataLoaderArgs::default());
    let test_loader: D::BatchLoader<f32, i64> = D::DataLoader::new(MNIST(),
                                                                   D::DataLoaderArgs::default());
    let args = parse_args();
    let mut optimizer = optim::SGD::new();

    let mut model = Net::new();
    for epoch in 1..args.epochs + 1 {
        train(&mut model, &args, &train_loader, epoch, &mut optimizer);
        test(&mut model, &args, &test_loader, epoch);
    }
}
<commit_msg>shuffle up args, factor out boilerplate<commit_after>extern crate getopts;
#[macro_use]
extern crate modparse_derive;
#[macro_use]
extern crate derive_builder;

#[macro_use]
extern crate torchrs;

use torchrs::autograd::{Variable, BackwardArgs};
use torchrs::tensor::Tensor;
use torchrs::optim;

use torchrs::macros::*;
use torchrs::nn;
use torchrs::nn::{ModuleStruct, ModIntf, ModDelegate, Module};
use torchrs::nn::functional as F;
use torchrs::utils::data as D;

use getopts::Options;
use std::env;

fn MNIST() -> D::Dataset<D::Batch<f32, i64>> {
    panic!("just a stub")
}

#[derive(Clone, Builder)]
struct NetArgs {
    #[builder(default="64")]
    batch_size: usize,
    #[builder(default="1000")]
    test_batch_size: usize,
    #[builder(default="10")]
    epochs: u32,
    #[builder(default="0.01")]
    lr: f32,
    #[builder(default="0.5")]
    momentum: f32,
    #[builder(default="true")]
    cuda: bool,
    #[builder(default="1")]
    seed: usize,
    #[builder(default="10")]
    log_interval: usize,
}

impl Default for NetArgs {
    fn default() -> Self {
        NetArgsBuilder::default().build().unwrap()
    }
}

fn parse_args() -> NetArgs {
    let cmd_args: Vec<String> = env::args().collect();
    let mut args = NetArgs::default();
    let mut opts = Options::new();
    /*
    	* XXX do parsey stuff
    	*/
    args
}

#[derive(ModParse)]
struct Net<T: Default> {
    delegate: nn::Module<T>,
    conv1: nn::Conv2d<T>,
    conv2: nn::Conv2d<T>,
    fc1: nn::Linear<T>,
    fc2: nn::Linear<T>,
}

impl<T: Default> Net<T> {
    pub fn new() -> Net<T> {
        let mut t = Net {
            delegate: nn::Module::new(),
            conv1: nn::Conv2d::build(1, 10, 5).done(),
            conv2: nn::Conv2d::build(10, 20, 5).done(),
            fc1: nn::Linear::build(320, 50).done(),
            fc2: nn::Linear::build(50, 10).done(),
        };
        t.init_module();
        t
    }
}
impl_mod_delegate!(Net);

// The forward operations could take on one of two implementations.
// The first supporting a near verbatim version of the python
// implementation, and the second supporting a slightly more
// idiomatic to Rust method chaining.

// a) as a near verbatim implementation of the python version
impl ModIntf<f32> for Net<f32> {
    fn forward(&mut self, args: &mut Variable<f32>) -> Variable<f32> {
        let training = self.delegate.training;
        let pool_val = F::MaxPool2dArgs::default();
        let dropout_val = F::DropoutArgs::default();
        let mut x = F::relu(&F::max_pool2d(&self.conv1.f(args), (2, 2), &pool_val));
        let mut x = F::relu(&F::max_pool2d(&F::dropout2d(&self.conv2.f(&mut x), &dropout_val),
                                           (2, 2),
                                           &pool_val));
        let mut x = x.view(&[-1, 320]);
        let x = F::relu(&self.fc1.f(&mut x));
        let mut x = F::dropout(&x, &dropout_val);
        let x = self.fc2.f(&mut x);
        F::log_softmax(&x)
    }
}

fn train(model: &mut Net<f32>,
         args: &NetArgs,
         train_loader: &D::BatchLoader<f32, i64>,
         epoch: u32,
         optimizer: &mut optim::OptIntf) {
    model.train();
    for (batch_idx, (ref data, ref target)) in train_loader.iter().enumerate() {
        let (mut data, target) = if args.cuda {
            (Variable::new(data.cuda().clone()), Variable::new(target.cuda().clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        optimizer.zero_grad();
        let output = model.f(&mut data);
        let mut loss = F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        loss.backward(&BackwardArgs::default());
        optimizer.step();
        if batch_idx % args.log_interval == 0 {
            println!("Train Epoch: {} [{}/{} ({:.0}%)]\tLoss: {:.6}",
                     epoch,
                     batch_idx * data.data().len(),
                     train_loader.dataset.len(),
                     100. * (batch_idx as f32) / train_loader.len() as f32,
                     loss.data()[0]);
        }
    }
}

fn test(model: &mut Net<f32>, args: &NetArgs, test_loader: &D::BatchLoader<f32, i64>, epoch: u32) {
    model.eval();
    let mut test_loss = 0.;
    let mut correct = 0;
    for (ref data, ref target) in test_loader.iter() {
        let (mut data, mut target) = if args.cuda {
            (Variable::new_volatile(data.cuda().clone()), Variable::new(target.cuda().clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        let mut output = model.f(&mut data);
        test_loss += F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        let pred = output.data().reduce_max(1).1;
        correct += pred.tensor_eq(target.data()).cpu().sum();
    }
    test_loss /= test_loader.len() as f32;
    println!("\nTest set: Average loss: {:.4}, Accuracy: {}/{} ({:.0}%)\n",
             test_loss,
             correct,
             test_loader.dataset.len(),
             100. * (correct as f32) / (test_loader.dataset.len() as f32))

}

fn main() {
    // no data loaders yet so demonstrate with Vec placeholder
    let train_loader: D::BatchLoader<f32, i64> = D::DataLoader::new(MNIST(),
                                                                    D::DataLoaderArgs::default());
    let test_loader: D::BatchLoader<f32, i64> = D::DataLoader::new(MNIST(),
                                                                   D::DataLoaderArgs::default());
    let args = parse_args();
    let mut optimizer = optim::SGD::new();

    let mut model = Net::new();
    for epoch in 1..args.epochs + 1 {
        train(&mut model, &args, &train_loader, epoch, &mut optimizer);
        test(&mut model, &args, &test_loader, epoch);
    }
}
<|endoftext|>


<commit_before>//
// Copyright 2018 The Project Oak Authors
//
// 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.
//

extern crate byteorder;
extern crate fmt;
#[macro_use]
extern crate log;
extern crate protobuf;

use byteorder::WriteBytesExt;
use proto::oak_api::OakStatus;
use protobuf::{Message, ProtobufEnum};
use std::io;
use std::io::Write;

pub mod proto;
pub mod storage;
#[cfg(test)]
mod tests;

#[cfg(test)]
#[macro_use]
extern crate assert_matches;

pub type GrpcResult<T> = Result<T, proto::status::Status>;

/// Trait to allow repeated writing of responses for server-streaming gRPC methods.
pub trait ResponseWriter<T: protobuf::Message> {
    fn write(&mut self, rsp: T);
}

/// Implementation of ResponseWriter that encapsulates response messages into
/// GrpcResponse wrapper messages and writes serialized versions to a (mutably
/// borrowed) send channel.
pub struct ChannelResponseWriter<'a> {
    pub channel: &'a mut SendChannelHalf,
}

impl<'a, T> ResponseWriter<T> for ChannelResponseWriter<'a>
where
    T: protobuf::Message,
{
    fn write(&mut self, rsp: T) {
        // Put the serialized response into a GrpcResponse message wrapper.
        let mut grpc_rsp = proto::grpc_encap::GrpcResponse::new();
        let mut any = protobuf::well_known_types::Any::new();
        rsp.write_to_writer(&mut any.value).unwrap();
        grpc_rsp.set_rsp_msg(any);
        // Serialize the GrpcResponse into the send channel.
        grpc_rsp.write_to_writer(&mut self.channel).unwrap();
    }
}

pub type Handle = u64;

/// Map an OakStatus to the nearest available std::io::Result.
fn result_from_status<T>(status: Option<OakStatus>, val: T) -> std::io::Result<T> {
    match status {
        Some(OakStatus::OAK_STATUS_UNSPECIFIED) => Err(io::Error::new(
            io::ErrorKind::Other,
            "Unspecified Oak status value",
        )),
        Some(OakStatus::OK) => Ok(val),
        Some(OakStatus::ERR_BAD_HANDLE) => {
            Err(io::Error::new(io::ErrorKind::NotConnected, "Bad handle"))
        }
        Some(OakStatus::ERR_INVALID_ARGS) => Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "Invalid arguments",
        )),
        Some(OakStatus::ERR_CHANNEL_CLOSED) => Err(io::Error::new(
            io::ErrorKind::ConnectionReset,
            "Channel closed",
        )),
        Some(OakStatus::ERR_BUFFER_TOO_SMALL) => Err(io::Error::new(
            io::ErrorKind::UnexpectedEof,
            "Buffer too small",
        )),
        Some(OakStatus::ERR_OUT_OF_RANGE) => {
            Err(io::Error::new(io::ErrorKind::NotConnected, "Out of range"))
        }
        Some(OakStatus::ERR_INTERNAL) => {
            Err(io::Error::new(io::ErrorKind::Other, "Internal error"))
        }
        None => Err(io::Error::new(
            io::ErrorKind::Other,
            "Unknown Oak status value",
        )),
    }
}

mod wasm {
    // See https://rustwasm.github.io/book/reference/js-ffi.html
    #[link(wasm_import_module = "oak")]
    extern "C" {
        pub fn wait_on_channels(buf: *mut u8, count: u32) -> i32;
        pub fn channel_read(handle: u64, buf: *mut u8, size: usize, actual_size: *mut u32) -> i32;
        pub fn channel_write(handle: u64, buf: *const u8, size: usize) -> i32;
        pub fn channel_find(buf: *const u8, len: usize) -> u64;
    }
}

pub fn channel_find(port_name: &str) -> Handle {
    unsafe { wasm::channel_find(port_name.as_ptr(), port_name.len()) }
}

pub const SPACE_BYTES_PER_HANDLE: usize = 9;

// Build a chunk of memory that is suitable for passing to wasm::wait_on_channels,
// holding the given collection of channel handles.
fn new_handle_space(handles: &[Handle]) -> Vec<u8> {
    let mut space = Vec::with_capacity(SPACE_BYTES_PER_HANDLE * handles.len());
    for handle in handles {
        space.write_u64::<byteorder::LittleEndian>(*handle).unwrap();
        space.push(0x00);
    }
    space
}

// Prepare a handle space for polling by clearing all of the message-pending
// indicator bytes.
fn prep_handle_space(space: &mut [u8]) {
    let count = space.len() / 8;
    for i in 0..count {
        space[i * SPACE_BYTES_PER_HANDLE + (SPACE_BYTES_PER_HANDLE - 1)] = 0;
    }
}

// Convenience wrapper around wait_on_channels host function. This version is
// easier to use in Rust but is less efficient (because the notification space
// is re-created on each invocation).
pub fn wait_on_channels(handles: &[Handle]) -> Result<Vec<Handle>, OakStatus> {
    let mut space = new_handle_space(handles);
    unsafe {
        let status = wasm::wait_on_channels(space.as_mut_ptr(), handles.len() as u32);
        match OakStatus::from_i32(status) {
            Some(OakStatus::OK) => (),
            Some(err) => return Err(err),
            None => return Err(OakStatus::OAK_STATUS_UNSPECIFIED),
        }
        let mut results = Vec::with_capacity(handles.len());
        for i in 0..handles.len() {
            if space[i * SPACE_BYTES_PER_HANDLE + (SPACE_BYTES_PER_HANDLE - 1)] != 0 {
                results.push(handles[i]);
            }
        }
        Ok(results)
    }
}

pub struct SendChannelHalf {
    handle: Handle,
}

impl SendChannelHalf {
    pub fn new(handle: Handle) -> SendChannelHalf {
        SendChannelHalf { handle }
    }

    pub fn write_message(&mut self, buf: &[u8]) -> std::io::Result<()> {
        result_from_status(
            OakStatus::from_i32(unsafe {
                wasm::channel_write(self.handle, buf.as_ptr(), buf.len())
            }),
            (),
        )
    }
}

// Implement the Write trait on a send channel for convenience, particularly for
// the logging channel.
impl Write for SendChannelHalf {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        match self.write_message(buf) {
            Ok(_) => Ok(buf.len()),
            Err(e) => Err(e),
        }
    }
    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

pub fn logging_channel() -> impl Write {
    let logging_channel = SendChannelHalf::new(channel_find("log"));
    // Only flush logging channel on newlines.
    std::io::LineWriter::new(logging_channel)
}

pub struct ReceiveChannelHalf {
    handle: Handle,
}

impl ReceiveChannelHalf {
    pub fn new(handle: Handle) -> ReceiveChannelHalf {
        ReceiveChannelHalf { handle }
    }

    pub fn read_message(&mut self, buf: &mut Vec<u8>) -> std::io::Result<usize> {
        // Try reading from the channel twice: first with provided vector,
        // then with a vector that's been resized to meet size requirements.
        for resized in &[false, true] {
            let mut actual_size: u32 = 0;
            let status = OakStatus::from_i32(unsafe {
                wasm::channel_read(
                    self.handle,
                    buf.as_mut_ptr(),
                    buf.capacity(),
                    &mut actual_size,
                )
            });
            match status {
                Some(OakStatus::OK) => {
                    unsafe {
                        buf.set_len(actual_size as usize);
                    };
                    return Ok(actual_size as usize);
                }
                Some(OakStatus::ERR_BUFFER_TOO_SMALL) => {
                    if *resized {
                        return result_from_status(status, 0);
                    }
                    // Can escape the match if buffer is too small and !resized.
                }
                _ => return result_from_status(status, 0),
            }

            // Extend the vector to be large enough for the message
            debug!("Got {}, need {}", buf.capacity(), actual_size);
            if (actual_size as usize) < buf.len() {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    format!(
                        "Internal error: provided {} bytes for receive, asked for {}",
                        buf.len(),
                        actual_size
                    ),
                ));
            }
            let extra = (actual_size as usize) - buf.len();
            buf.reserve(extra);
        }
        Err(io::Error::new(io::ErrorKind::Other, "Unreachable reached!"))
    }
}

/// Convenience wrapper for a send/receive pair of channels.
pub struct ChannelPair {
    pub receive: ReceiveChannelHalf,
    pub send: SendChannelHalf,
}

impl ChannelPair {
    pub fn new(in_handle: Handle, out_handle: Handle) -> ChannelPair {
        ChannelPair {
            receive: ReceiveChannelHalf::new(in_handle),
            send: SendChannelHalf::new(out_handle),
        }
    }
}

/// Trait encapsulating the operations required for an Oak Node.
pub trait OakNode {
    fn new() -> Self
    where
        Self: Sized;
    fn invoke(&mut self, method: &str, req: &[u8], out: &mut SendChannelHalf);
}

/// Perform a gRPC event loop on the given channels, invoking the given node.
pub fn grpc_event_loop<T: OakNode>(
    mut node: T,
    grpc_in_handle: Handle,
    grpc_out_handle: Handle,
) -> i32 {
    info!(
        "start event loop for node with handles in:{} out:{}",
        grpc_in_handle, grpc_out_handle
    );
    if grpc_in_handle == 0 || grpc_out_handle == 0 {
        return OakStatus::ERR_CHANNEL_CLOSED.value();
    }
    set_panic_hook();

    let read_handles = vec![grpc_in_handle];
    let mut space = new_handle_space(&read_handles);

    let mut grpc_in_channel = ReceiveChannelHalf::new(grpc_in_handle);
    let mut grpc_out_channel = SendChannelHalf::new(grpc_out_handle);
    loop {
        // Block until there is a message to read on an input channel.
        prep_handle_space(&mut space);
        unsafe {
            let status = wasm::wait_on_channels(space.as_mut_ptr(), read_handles.len() as u32);
            match OakStatus::from_i32(status) {
                Some(OakStatus::OK) => (),
                Some(err) => return err as i32,
                None => return OakStatus::OAK_STATUS_UNSPECIFIED.value(),
            }
        }

        let mut buf = Vec::<u8>::with_capacity(1024);
        grpc_in_channel.read_message(&mut buf).unwrap();
        if buf.is_empty() {
            info!("no pending message; poll again");
            continue;
        }
        let req: proto::grpc_encap::GrpcRequest = protobuf::parse_from_bytes(&buf).unwrap();
        if !req.last {
            panic!("Support for streaming requests not yet implemented");
        }
        node.invoke(
            &req.method_name,
            req.get_req_msg().value.as_slice(),
            &mut grpc_out_channel,
        );
    }
}

/// Install a panic hook so that panics are logged to the logging channel, if one is set.
/// See https://doc.rust-lang.org/std/panic/struct.PanicInfo.html.
fn set_panic_hook() {
    std::panic::set_hook(Box::new(|panic_info| {
        let msg = panic_info
            .payload()
            .downcast_ref::<&str>()
            .unwrap_or(&"<UNKWOWN MESSAGE>");
        let (file, line) = match panic_info.location() {
            Some(location) => (location.file(), location.line()),
            None => ("<UNKNOWN FILE>", 0),
        };
        error!(
            "panic occurred in file '{}' at line {}: {}",
            file, line, msg
        );
    }));
}
<commit_msg>Make oak::set_panic_hook public<commit_after>//
// Copyright 2018 The Project Oak Authors
//
// 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.
//

extern crate byteorder;
extern crate fmt;
#[macro_use]
extern crate log;
extern crate protobuf;

use byteorder::WriteBytesExt;
use proto::oak_api::OakStatus;
use protobuf::{Message, ProtobufEnum};
use std::io;
use std::io::Write;

pub mod proto;
pub mod storage;
#[cfg(test)]
mod tests;

#[cfg(test)]
#[macro_use]
extern crate assert_matches;

pub type GrpcResult<T> = Result<T, proto::status::Status>;

/// Trait to allow repeated writing of responses for server-streaming gRPC methods.
pub trait ResponseWriter<T: protobuf::Message> {
    fn write(&mut self, rsp: T);
}

/// Implementation of ResponseWriter that encapsulates response messages into
/// GrpcResponse wrapper messages and writes serialized versions to a (mutably
/// borrowed) send channel.
pub struct ChannelResponseWriter<'a> {
    pub channel: &'a mut SendChannelHalf,
}

impl<'a, T> ResponseWriter<T> for ChannelResponseWriter<'a>
where
    T: protobuf::Message,
{
    fn write(&mut self, rsp: T) {
        // Put the serialized response into a GrpcResponse message wrapper.
        let mut grpc_rsp = proto::grpc_encap::GrpcResponse::new();
        let mut any = protobuf::well_known_types::Any::new();
        rsp.write_to_writer(&mut any.value).unwrap();
        grpc_rsp.set_rsp_msg(any);
        // Serialize the GrpcResponse into the send channel.
        grpc_rsp.write_to_writer(&mut self.channel).unwrap();
    }
}

pub type Handle = u64;

/// Map an OakStatus to the nearest available std::io::Result.
fn result_from_status<T>(status: Option<OakStatus>, val: T) -> std::io::Result<T> {
    match status {
        Some(OakStatus::OAK_STATUS_UNSPECIFIED) => Err(io::Error::new(
            io::ErrorKind::Other,
            "Unspecified Oak status value",
        )),
        Some(OakStatus::OK) => Ok(val),
        Some(OakStatus::ERR_BAD_HANDLE) => {
            Err(io::Error::new(io::ErrorKind::NotConnected, "Bad handle"))
        }
        Some(OakStatus::ERR_INVALID_ARGS) => Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "Invalid arguments",
        )),
        Some(OakStatus::ERR_CHANNEL_CLOSED) => Err(io::Error::new(
            io::ErrorKind::ConnectionReset,
            "Channel closed",
        )),
        Some(OakStatus::ERR_BUFFER_TOO_SMALL) => Err(io::Error::new(
            io::ErrorKind::UnexpectedEof,
            "Buffer too small",
        )),
        Some(OakStatus::ERR_OUT_OF_RANGE) => {
            Err(io::Error::new(io::ErrorKind::NotConnected, "Out of range"))
        }
        Some(OakStatus::ERR_INTERNAL) => {
            Err(io::Error::new(io::ErrorKind::Other, "Internal error"))
        }
        None => Err(io::Error::new(
            io::ErrorKind::Other,
            "Unknown Oak status value",
        )),
    }
}

mod wasm {
    // See https://rustwasm.github.io/book/reference/js-ffi.html
    #[link(wasm_import_module = "oak")]
    extern "C" {
        pub fn wait_on_channels(buf: *mut u8, count: u32) -> i32;
        pub fn channel_read(handle: u64, buf: *mut u8, size: usize, actual_size: *mut u32) -> i32;
        pub fn channel_write(handle: u64, buf: *const u8, size: usize) -> i32;
        pub fn channel_find(buf: *const u8, len: usize) -> u64;
    }
}

pub fn channel_find(port_name: &str) -> Handle {
    unsafe { wasm::channel_find(port_name.as_ptr(), port_name.len()) }
}

pub const SPACE_BYTES_PER_HANDLE: usize = 9;

// Build a chunk of memory that is suitable for passing to wasm::wait_on_channels,
// holding the given collection of channel handles.
fn new_handle_space(handles: &[Handle]) -> Vec<u8> {
    let mut space = Vec::with_capacity(SPACE_BYTES_PER_HANDLE * handles.len());
    for handle in handles {
        space.write_u64::<byteorder::LittleEndian>(*handle).unwrap();
        space.push(0x00);
    }
    space
}

// Prepare a handle space for polling by clearing all of the message-pending
// indicator bytes.
fn prep_handle_space(space: &mut [u8]) {
    let count = space.len() / 8;
    for i in 0..count {
        space[i * SPACE_BYTES_PER_HANDLE + (SPACE_BYTES_PER_HANDLE - 1)] = 0;
    }
}

// Convenience wrapper around wait_on_channels host function. This version is
// easier to use in Rust but is less efficient (because the notification space
// is re-created on each invocation).
pub fn wait_on_channels(handles: &[Handle]) -> Result<Vec<Handle>, OakStatus> {
    let mut space = new_handle_space(handles);
    unsafe {
        let status = wasm::wait_on_channels(space.as_mut_ptr(), handles.len() as u32);
        match OakStatus::from_i32(status) {
            Some(OakStatus::OK) => (),
            Some(err) => return Err(err),
            None => return Err(OakStatus::OAK_STATUS_UNSPECIFIED),
        }
        let mut results = Vec::with_capacity(handles.len());
        for i in 0..handles.len() {
            if space[i * SPACE_BYTES_PER_HANDLE + (SPACE_BYTES_PER_HANDLE - 1)] != 0 {
                results.push(handles[i]);
            }
        }
        Ok(results)
    }
}

pub struct SendChannelHalf {
    handle: Handle,
}

impl SendChannelHalf {
    pub fn new(handle: Handle) -> SendChannelHalf {
        SendChannelHalf { handle }
    }

    pub fn write_message(&mut self, buf: &[u8]) -> std::io::Result<()> {
        result_from_status(
            OakStatus::from_i32(unsafe {
                wasm::channel_write(self.handle, buf.as_ptr(), buf.len())
            }),
            (),
        )
    }
}

// Implement the Write trait on a send channel for convenience, particularly for
// the logging channel.
impl Write for SendChannelHalf {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        match self.write_message(buf) {
            Ok(_) => Ok(buf.len()),
            Err(e) => Err(e),
        }
    }
    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

pub fn logging_channel() -> impl Write {
    let logging_channel = SendChannelHalf::new(channel_find("log"));
    // Only flush logging channel on newlines.
    std::io::LineWriter::new(logging_channel)
}

pub struct ReceiveChannelHalf {
    handle: Handle,
}

impl ReceiveChannelHalf {
    pub fn new(handle: Handle) -> ReceiveChannelHalf {
        ReceiveChannelHalf { handle }
    }

    pub fn read_message(&mut self, buf: &mut Vec<u8>) -> std::io::Result<usize> {
        // Try reading from the channel twice: first with provided vector,
        // then with a vector that's been resized to meet size requirements.
        for resized in &[false, true] {
            let mut actual_size: u32 = 0;
            let status = OakStatus::from_i32(unsafe {
                wasm::channel_read(
                    self.handle,
                    buf.as_mut_ptr(),
                    buf.capacity(),
                    &mut actual_size,
                )
            });
            match status {
                Some(OakStatus::OK) => {
                    unsafe {
                        buf.set_len(actual_size as usize);
                    };
                    return Ok(actual_size as usize);
                }
                Some(OakStatus::ERR_BUFFER_TOO_SMALL) => {
                    if *resized {
                        return result_from_status(status, 0);
                    }
                    // Can escape the match if buffer is too small and !resized.
                }
                _ => return result_from_status(status, 0),
            }

            // Extend the vector to be large enough for the message
            debug!("Got {}, need {}", buf.capacity(), actual_size);
            if (actual_size as usize) < buf.len() {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    format!(
                        "Internal error: provided {} bytes for receive, asked for {}",
                        buf.len(),
                        actual_size
                    ),
                ));
            }
            let extra = (actual_size as usize) - buf.len();
            buf.reserve(extra);
        }
        Err(io::Error::new(io::ErrorKind::Other, "Unreachable reached!"))
    }
}

/// Convenience wrapper for a send/receive pair of channels.
pub struct ChannelPair {
    pub receive: ReceiveChannelHalf,
    pub send: SendChannelHalf,
}

impl ChannelPair {
    pub fn new(in_handle: Handle, out_handle: Handle) -> ChannelPair {
        ChannelPair {
            receive: ReceiveChannelHalf::new(in_handle),
            send: SendChannelHalf::new(out_handle),
        }
    }
}

/// Trait encapsulating the operations required for an Oak Node.
pub trait OakNode {
    fn new() -> Self
    where
        Self: Sized;
    fn invoke(&mut self, method: &str, req: &[u8], out: &mut SendChannelHalf);
}

/// Perform a gRPC event loop on the given channels, invoking the given node.
pub fn grpc_event_loop<T: OakNode>(
    mut node: T,
    grpc_in_handle: Handle,
    grpc_out_handle: Handle,
) -> i32 {
    info!(
        "start event loop for node with handles in:{} out:{}",
        grpc_in_handle, grpc_out_handle
    );
    if grpc_in_handle == 0 || grpc_out_handle == 0 {
        return OakStatus::ERR_CHANNEL_CLOSED.value();
    }
    set_panic_hook();

    let read_handles = vec![grpc_in_handle];
    let mut space = new_handle_space(&read_handles);

    let mut grpc_in_channel = ReceiveChannelHalf::new(grpc_in_handle);
    let mut grpc_out_channel = SendChannelHalf::new(grpc_out_handle);
    loop {
        // Block until there is a message to read on an input channel.
        prep_handle_space(&mut space);
        unsafe {
            let status = wasm::wait_on_channels(space.as_mut_ptr(), read_handles.len() as u32);
            match OakStatus::from_i32(status) {
                Some(OakStatus::OK) => (),
                Some(err) => return err as i32,
                None => return OakStatus::OAK_STATUS_UNSPECIFIED.value(),
            }
        }

        let mut buf = Vec::<u8>::with_capacity(1024);
        grpc_in_channel.read_message(&mut buf).unwrap();
        if buf.is_empty() {
            info!("no pending message; poll again");
            continue;
        }
        let req: proto::grpc_encap::GrpcRequest = protobuf::parse_from_bytes(&buf).unwrap();
        if !req.last {
            panic!("Support for streaming requests not yet implemented");
        }
        node.invoke(
            &req.method_name,
            req.get_req_msg().value.as_slice(),
            &mut grpc_out_channel,
        );
    }
}

/// Install a panic hook so that panics are logged to the logging channel, if one is set.
/// See https://doc.rust-lang.org/std/panic/struct.PanicInfo.html.
pub fn set_panic_hook() {
    std::panic::set_hook(Box::new(|panic_info| {
        let msg = panic_info
            .payload()
            .downcast_ref::<&str>()
            .unwrap_or(&"<UNKNOWN MESSAGE>");
        let (file, line) = match panic_info.location() {
            Some(location) => (location.file(), location.line()),
            None => ("<UNKNOWN FILE>", 0),
        };
        error!(
            "panic occurred in file '{}' at line {}: {}",
            file, line, msg
        );
    }));
}
<|endoftext|>