<commit_before>// Copyright 2022 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::error::Error;
use std::fmt::{Display, Formatter};

use thiserror::Error as ThisError;

use base::MappedRegion;
use base::MemoryMappingArena;
use base::MmapError;
use ffmpeg::avcodec::AvBuffer;
use ffmpeg::avcodec::AvBufferSource;
use ffmpeg::avcodec::AvError;
use ffmpeg::avcodec::AvFrame;
use ffmpeg::avcodec::AvFrameError;
use ffmpeg::avcodec::AvPixelFormat;
use ffmpeg::avcodec::Dimensions;
use ffmpeg::avcodec::PlaneDescriptor;

use crate::virtio::video::format::Format;
use crate::virtio::video::resource::BufferHandle;
use crate::virtio::video::resource::GuestResource;

/// A simple wrapper that turns a [`MemoryMappingArena`] into an [`AvBufferSource`].
///
/// **Note:** Only use this if you can reasonably assume the mapped memory won't be written to from
/// another thread or the VM! The reason [`AvBufferSource`] is not directly implemented on
/// [`MemoryMappingArena`] is exactly due to this unsafety: If the guest or someone else writes to
/// the buffer FFmpeg is currently reading from or writing to, undefined behavior might occur. This
/// struct acts as an opt-in mechanism for this potential unsafety.
pub struct MemoryMappingAvBufferSource(MemoryMappingArena);

impl From<MemoryMappingArena> for MemoryMappingAvBufferSource {
    fn from(inner: MemoryMappingArena) -> Self {
        Self(inner)
    }
}
impl AvBufferSource for MemoryMappingAvBufferSource {
    fn as_ptr(&self) -> *const u8 {
        self.0.as_ptr() as _
    }

    fn len(&self) -> usize {
        self.0.size()
    }
}

pub trait TryAsAvFrameExt {
    type BufferSource;
    type Error: Error;

    /// Try to create an [`AvFrame`] from `self`.
    ///
    /// `wrapper` is a constructor for `T` that wraps `Self::BufferSource` and implement the
    /// `AvBufferSource` functionality. This can be used to provide a custom destructor
    /// implementation: for example, sending a virtio message signaling the buffer source is no
    /// longer used and it can be recycled by the guest.
    ///
    /// Note that `Self::BufferSource` might not always implement `AvBufferSource`. For instance,
    /// it's not guaranteed that a `MemoryMappingArena` created from `GuestResource` will be always
    /// immutable, and the backend need to explicit acknowledge the potential unsoundness by
    /// wrapping it in [`MemoryMappingAvBufferSource`].
    fn try_as_av_frame<T: AvBufferSource + 'static>(
        &self,
        wrapper: impl FnOnce(Self::BufferSource) -> T,
    ) -> Result<AvFrame, Self::Error>;
}

#[derive(Debug, ThisError)]
pub enum GuestResourceToAvFrameError {
    #[error("error while creating the AvFrame: {0}")]
    AvFrameError(#[from] AvFrameError),
    #[error("cannot mmap guest resource: {0}")]
    MappingResource(#[from] MmapError),
    #[error("virtio resource format is not convertible to AvPixelFormat")]
    InvalidFormat,
    #[error("out of memory while allocating AVBuffer")]
    CannotAllocateAvBuffer,
}

impl From<AvError> for GuestResourceToAvFrameError {
    fn from(e: AvError) -> Self {
        GuestResourceToAvFrameError::AvFrameError(AvFrameError::AvError(e))
    }
}

impl TryAsAvFrameExt for GuestResource {
    type BufferSource = MemoryMappingArena;
    type Error = GuestResourceToAvFrameError;

    fn try_as_av_frame<T: AvBufferSource + 'static>(
        &self,
        wrapper: impl FnOnce(MemoryMappingArena) -> T,
    ) -> Result<AvFrame, Self::Error> {
        let mut builder = AvFrame::builder()?;
        builder.set_dimensions(Dimensions {
            width: self.width,
            height: self.height,
        })?;
        let format = self
            .format
            .try_into()
            .map_err(|_| GuestResourceToAvFrameError::InvalidFormat)?;
        builder.set_format(format)?;
        let planes = &self.planes;
        // We have at least one plane, so we can unwrap below.
        let len = format.plane_sizes(planes.iter().map(|p| p.stride as _), self.height)?;
        let start = planes.iter().map(|p| p.offset).min().unwrap();
        let end = planes
            .iter()
            .enumerate()
            .map(|(i, p)| p.offset + len[i])
            .max()
            .unwrap();
        let mapping = self.handle.get_mapping(start, end - start)?;
        Ok(builder.build_owned(
            [AvBuffer::new(wrapper(mapping))
                .ok_or(GuestResourceToAvFrameError::CannotAllocateAvBuffer)?],
            planes.iter().map(|p| PlaneDescriptor {
                buffer_index: 0,
                offset: p.offset - start,
                stride: p.stride,
            }),
        )?)
    }
}

/// The error returned by `AvPixelFormat::try_from` when there's no applicable format.
// The empty field prevents constructing this and allows extending it in the future.
#[derive(Debug)]
pub struct TryFromFormatError(());

impl Display for TryFromFormatError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "No matching format to convert between AvPixelFormat and Format"
        )
    }
}

impl Error for TryFromFormatError {}

impl TryFrom<Format> for AvPixelFormat {
    type Error = TryFromFormatError;

    fn try_from(value: Format) -> Result<Self, Self::Error> {
        use ffmpeg::*;
        AvPixelFormat::try_from(match value {
            Format::NV12 => AVPixelFormat_AV_PIX_FMT_NV12,
            Format::YUV420 => AVPixelFormat_AV_PIX_FMT_YUV420P,
            _ => return Err(TryFromFormatError(())),
        })
        .map_err(|_|
            // The error case should never happen as long as we use valid constant values, but
            // don't panic in case something goes wrong.
            TryFromFormatError(()))
    }
}
<commit_msg>virtio: video: Add support for converting AvPixelFormat -> Format.<commit_after>// Copyright 2022 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::error::Error;
use std::fmt::{Display, Formatter};

use thiserror::Error as ThisError;

use base::MappedRegion;
use base::MemoryMappingArena;
use base::MmapError;
use ffmpeg::avcodec::AvBuffer;
use ffmpeg::avcodec::AvBufferSource;
use ffmpeg::avcodec::AvError;
use ffmpeg::avcodec::AvFrame;
use ffmpeg::avcodec::AvFrameError;
use ffmpeg::avcodec::AvPixelFormat;
use ffmpeg::avcodec::Dimensions;
use ffmpeg::avcodec::PlaneDescriptor;

use crate::virtio::video::format::Format;
use crate::virtio::video::resource::BufferHandle;
use crate::virtio::video::resource::GuestResource;

/// A simple wrapper that turns a [`MemoryMappingArena`] into an [`AvBufferSource`].
///
/// **Note:** Only use this if you can reasonably assume the mapped memory won't be written to from
/// another thread or the VM! The reason [`AvBufferSource`] is not directly implemented on
/// [`MemoryMappingArena`] is exactly due to this unsafety: If the guest or someone else writes to
/// the buffer FFmpeg is currently reading from or writing to, undefined behavior might occur. This
/// struct acts as an opt-in mechanism for this potential unsafety.
pub struct MemoryMappingAvBufferSource(MemoryMappingArena);

impl From<MemoryMappingArena> for MemoryMappingAvBufferSource {
    fn from(inner: MemoryMappingArena) -> Self {
        Self(inner)
    }
}
impl AvBufferSource for MemoryMappingAvBufferSource {
    fn as_ptr(&self) -> *const u8 {
        self.0.as_ptr() as _
    }

    fn len(&self) -> usize {
        self.0.size()
    }
}

pub trait TryAsAvFrameExt {
    type BufferSource;
    type Error: Error;

    /// Try to create an [`AvFrame`] from `self`.
    ///
    /// `wrapper` is a constructor for `T` that wraps `Self::BufferSource` and implement the
    /// `AvBufferSource` functionality. This can be used to provide a custom destructor
    /// implementation: for example, sending a virtio message signaling the buffer source is no
    /// longer used and it can be recycled by the guest.
    ///
    /// Note that `Self::BufferSource` might not always implement `AvBufferSource`. For instance,
    /// it's not guaranteed that a `MemoryMappingArena` created from `GuestResource` will be always
    /// immutable, and the backend need to explicit acknowledge the potential unsoundness by
    /// wrapping it in [`MemoryMappingAvBufferSource`].
    fn try_as_av_frame<T: AvBufferSource + 'static>(
        &self,
        wrapper: impl FnOnce(Self::BufferSource) -> T,
    ) -> Result<AvFrame, Self::Error>;
}

#[derive(Debug, ThisError)]
pub enum GuestResourceToAvFrameError {
    #[error("error while creating the AvFrame: {0}")]
    AvFrameError(#[from] AvFrameError),
    #[error("cannot mmap guest resource: {0}")]
    MappingResource(#[from] MmapError),
    #[error("virtio resource format is not convertible to AvPixelFormat")]
    InvalidFormat,
    #[error("out of memory while allocating AVBuffer")]
    CannotAllocateAvBuffer,
}

impl From<AvError> for GuestResourceToAvFrameError {
    fn from(e: AvError) -> Self {
        GuestResourceToAvFrameError::AvFrameError(AvFrameError::AvError(e))
    }
}

impl TryAsAvFrameExt for GuestResource {
    type BufferSource = MemoryMappingArena;
    type Error = GuestResourceToAvFrameError;

    fn try_as_av_frame<T: AvBufferSource + 'static>(
        &self,
        wrapper: impl FnOnce(MemoryMappingArena) -> T,
    ) -> Result<AvFrame, Self::Error> {
        let mut builder = AvFrame::builder()?;
        builder.set_dimensions(Dimensions {
            width: self.width,
            height: self.height,
        })?;
        let format = self
            .format
            .try_into()
            .map_err(|_| GuestResourceToAvFrameError::InvalidFormat)?;
        builder.set_format(format)?;
        let planes = &self.planes;
        // We have at least one plane, so we can unwrap below.
        let len = format.plane_sizes(planes.iter().map(|p| p.stride as _), self.height)?;
        let start = planes.iter().map(|p| p.offset).min().unwrap();
        let end = planes
            .iter()
            .enumerate()
            .map(|(i, p)| p.offset + len[i])
            .max()
            .unwrap();
        let mapping = self.handle.get_mapping(start, end - start)?;
        Ok(builder.build_owned(
            [AvBuffer::new(wrapper(mapping))
                .ok_or(GuestResourceToAvFrameError::CannotAllocateAvBuffer)?],
            planes.iter().map(|p| PlaneDescriptor {
                buffer_index: 0,
                offset: p.offset - start,
                stride: p.stride,
            }),
        )?)
    }
}

/// The error returned when converting between `AvPixelFormat` and `Format` and there's no
/// applicable format.
// The empty field prevents constructing this and allows extending it in the future.
#[derive(Debug)]
pub struct TryFromFormatError(());

impl Display for TryFromFormatError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "No matching format to convert between AvPixelFormat and Format"
        )
    }
}

impl Error for TryFromFormatError {}

impl TryFrom<Format> for AvPixelFormat {
    type Error = TryFromFormatError;

    fn try_from(value: Format) -> Result<Self, Self::Error> {
        use ffmpeg::*;
        AvPixelFormat::try_from(match value {
            Format::NV12 => AVPixelFormat_AV_PIX_FMT_NV12,
            Format::YUV420 => AVPixelFormat_AV_PIX_FMT_YUV420P,
            _ => return Err(TryFromFormatError(())),
        })
        .map_err(|_|
            // The error case should never happen as long as we use valid constant values, but
            // don't panic in case something goes wrong.
            TryFromFormatError(()))
    }
}

impl TryFrom<AvPixelFormat> for Format {
    type Error = TryFromFormatError;

    fn try_from(fmt: AvPixelFormat) -> Result<Self, Self::Error> {
        // https://github.com/rust-lang/rust/issues/39371 Lint wrongly warns the consumer
        #![allow(non_upper_case_globals)]
        use ffmpeg::*;
        Ok(match fmt.pix_fmt() {
            AVPixelFormat_AV_PIX_FMT_NV12 => Format::NV12,
            AVPixelFormat_AV_PIX_FMT_YUV420P => Format::YUV420,
            _ => return Err(TryFromFormatError(())),
        })
    }
}
<|endoftext|>


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

//! Hash Function
//!
//! Primary hash function used in the protocol
//!

use byteorder::{BigEndian, ByteOrder};
use std::cmp::min;
use std::convert::AsRef;
use std::ops::Add;
use std::{fmt, ops};

use crate::blake2::blake2b::Blake2b;

use crate::ser::{self, AsFixedBytes, Error, FixedLength, Readable, Reader, Writeable, Writer};
use crate::util;

/// A hash consisting of all zeroes, used as a sentinel. No known preimage.
pub const ZERO_HASH: Hash = Hash([0; 32]);

/// A hash to uniquely (or close enough) identify one of the main blockchain
/// constructs. Used pervasively for blocks, transactions and outputs.
#[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash, Serialize, Deserialize)]
pub struct Hash([u8; 32]);

impl fmt::Debug for Hash {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		for i in self.0[..4].iter() {
			write!(f, "{:02x}", i)?;
		}
		Ok(())
	}
}

impl fmt::Display for Hash {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		let hash_hex = self.to_hex();
		const NUM_SHOW: usize = 12;

		write!(f, "{}...", &hash_hex[..NUM_SHOW])
	}
}

impl FixedLength for Hash {
	/// Size of a hash in bytes.
	const LEN: usize = 32;
}

impl Hash {
	/// Builds a Hash from a byte vector. If the vector is too short, it will be
	/// completed by zeroes. If it's too long, it will be truncated.
	pub fn from_vec(v: &[u8]) -> Hash {
		let mut h = [0; Hash::LEN];
		let copy_size = min(v.len(), Hash::LEN);
		h[..copy_size].copy_from_slice(&v[..copy_size]);
		Hash(h)
	}

	/// Converts the hash to a byte vector
	pub fn to_vec(&self) -> Vec<u8> {
		self.0.to_vec()
	}

	/// Returns a byte slice of the hash contents.
	pub fn as_bytes(&self) -> &[u8] {
		&self.0
	}

	/// Convert a hash to hex string format.
	pub fn to_hex(&self) -> String {
		util::to_hex(self.to_vec())
	}

	/// Convert hex string back to hash.
	pub fn from_hex(hex: &str) -> Result<Hash, Error> {
		let bytes = util::from_hex(hex.to_string())
			.map_err(|_| Error::HexError(format!("failed to decode {}", hex)))?;
		Ok(Hash::from_vec(&bytes))
	}

	/// Most significant 64 bits
	pub fn to_u64(&self) -> u64 {
		BigEndian::read_u64(&self.0)
	}
}

impl ops::Index<usize> for Hash {
	type Output = u8;

	fn index(&self, idx: usize) -> &u8 {
		&self.0[idx]
	}
}

impl ops::Index<ops::Range<usize>> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::Range<usize>) -> &[u8] {
		&self.0[idx]
	}
}

impl ops::Index<ops::RangeTo<usize>> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::RangeTo<usize>) -> &[u8] {
		&self.0[idx]
	}
}

impl ops::Index<ops::RangeFrom<usize>> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::RangeFrom<usize>) -> &[u8] {
		&self.0[idx]
	}
}

impl ops::Index<ops::RangeFull> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::RangeFull) -> &[u8] {
		&self.0[idx]
	}
}

impl AsRef<[u8]> for Hash {
	fn as_ref(&self) -> &[u8] {
		&self.0
	}
}

impl Readable for Hash {
	fn read(reader: &mut dyn Reader) -> Result<Hash, ser::Error> {
		let v = reader.read_fixed_bytes(32)?;
		let mut a = [0; 32];
		a.copy_from_slice(&v[..]);
		Ok(Hash(a))
	}
}

impl Writeable for Hash {
	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), Error> {
		writer.write_fixed_bytes(&self.0)
	}
}

impl Add for Hash {
	type Output = Hash;
	fn add(self, other: Hash) -> Hash {
		self.hash_with(other)
	}
}

impl Default for Hash {
	fn default() -> Hash {
		ZERO_HASH
	}
}

/// Serializer that outputs a hash of the serialized object
pub struct HashWriter {
	state: Blake2b,
}

impl HashWriter {
	/// Consume the `HashWriter`, outputting its current hash into a 32-byte
	/// array
	pub fn finalize(self, output: &mut [u8]) {
		output.copy_from_slice(self.state.finalize().as_bytes());
	}

	/// Consume the `HashWriter`, outputting a `Hash` corresponding to its
	/// current state
	pub fn into_hash(self) -> Hash {
		let mut res = [0; 32];
		res.copy_from_slice(self.state.finalize().as_bytes());
		Hash(res)
	}
}

impl Default for HashWriter {
	fn default() -> HashWriter {
		HashWriter {
			state: Blake2b::new(32),
		}
	}
}

impl ser::Writer for HashWriter {
	fn serialization_mode(&self) -> ser::SerializationMode {
		ser::SerializationMode::Hash
	}

	fn write_fixed_bytes<T: AsFixedBytes>(&mut self, b32: &T) -> Result<(), ser::Error> {
		self.state.update(b32.as_ref());
		Ok(())
	}
}

/// A trait for types that have a canonical hash
pub trait Hashed {
	/// Obtain the hash of the object
	fn hash(&self) -> Hash;
	/// Hash the object together with another writeable object
	fn hash_with<T: Writeable>(&self, other: T) -> Hash;
}

impl<W: ser::Writeable> Hashed for W {
	fn hash(&self) -> Hash {
		let mut hasher = HashWriter::default();
		ser::Writeable::write(self, &mut hasher).unwrap();
		let mut ret = [0; 32];
		hasher.finalize(&mut ret);
		Hash(ret)
	}

	fn hash_with<T: Writeable>(&self, other: T) -> Hash {
		let mut hasher = HashWriter::default();
		ser::Writeable::write(self, &mut hasher).unwrap();
		ser::Writeable::write(&other, &mut hasher).unwrap();
		let mut ret = [0; 32];
		hasher.finalize(&mut ret);
		Hash(ret)
	}
}
<commit_msg>Show full hex string of hash in fmt::Debug for core::Hash<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Hash Function
//!
//! Primary hash function used in the protocol
//!

use byteorder::{BigEndian, ByteOrder};
use std::cmp::min;
use std::convert::AsRef;
use std::ops::Add;
use std::{fmt, ops};

use crate::blake2::blake2b::Blake2b;

use crate::ser::{self, AsFixedBytes, Error, FixedLength, Readable, Reader, Writeable, Writer};
use crate::util;

/// A hash consisting of all zeroes, used as a sentinel. No known preimage.
pub const ZERO_HASH: Hash = Hash([0; 32]);

/// A hash to uniquely (or close enough) identify one of the main blockchain
/// constructs. Used pervasively for blocks, transactions and outputs.
#[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash, Serialize, Deserialize)]
pub struct Hash([u8; 32]);

impl fmt::Debug for Hash {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		write!(f, "{}", self.to_hex())
	}
}

impl fmt::Display for Hash {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		let hash_hex = self.to_hex();
		const NUM_SHOW: usize = 12;

		write!(f, "{}...", &hash_hex[..NUM_SHOW])
	}
}

impl FixedLength for Hash {
	/// Size of a hash in bytes.
	const LEN: usize = 32;
}

impl Hash {
	/// Builds a Hash from a byte vector. If the vector is too short, it will be
	/// completed by zeroes. If it's too long, it will be truncated.
	pub fn from_vec(v: &[u8]) -> Hash {
		let mut h = [0; Hash::LEN];
		let copy_size = min(v.len(), Hash::LEN);
		h[..copy_size].copy_from_slice(&v[..copy_size]);
		Hash(h)
	}

	/// Converts the hash to a byte vector
	pub fn to_vec(&self) -> Vec<u8> {
		self.0.to_vec()
	}

	/// Returns a byte slice of the hash contents.
	pub fn as_bytes(&self) -> &[u8] {
		&self.0
	}

	/// Convert a hash to hex string format.
	pub fn to_hex(&self) -> String {
		util::to_hex(self.to_vec())
	}

	/// Convert hex string back to hash.
	pub fn from_hex(hex: &str) -> Result<Hash, Error> {
		let bytes = util::from_hex(hex.to_string())
			.map_err(|_| Error::HexError(format!("failed to decode {}", hex)))?;
		Ok(Hash::from_vec(&bytes))
	}

	/// Most significant 64 bits
	pub fn to_u64(&self) -> u64 {
		BigEndian::read_u64(&self.0)
	}
}

impl ops::Index<usize> for Hash {
	type Output = u8;

	fn index(&self, idx: usize) -> &u8 {
		&self.0[idx]
	}
}

impl ops::Index<ops::Range<usize>> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::Range<usize>) -> &[u8] {
		&self.0[idx]
	}
}

impl ops::Index<ops::RangeTo<usize>> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::RangeTo<usize>) -> &[u8] {
		&self.0[idx]
	}
}

impl ops::Index<ops::RangeFrom<usize>> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::RangeFrom<usize>) -> &[u8] {
		&self.0[idx]
	}
}

impl ops::Index<ops::RangeFull> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::RangeFull) -> &[u8] {
		&self.0[idx]
	}
}

impl AsRef<[u8]> for Hash {
	fn as_ref(&self) -> &[u8] {
		&self.0
	}
}

impl Readable for Hash {
	fn read(reader: &mut dyn Reader) -> Result<Hash, ser::Error> {
		let v = reader.read_fixed_bytes(32)?;
		let mut a = [0; 32];
		a.copy_from_slice(&v[..]);
		Ok(Hash(a))
	}
}

impl Writeable for Hash {
	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), Error> {
		writer.write_fixed_bytes(&self.0)
	}
}

impl Add for Hash {
	type Output = Hash;
	fn add(self, other: Hash) -> Hash {
		self.hash_with(other)
	}
}

impl Default for Hash {
	fn default() -> Hash {
		ZERO_HASH
	}
}

/// Serializer that outputs a hash of the serialized object
pub struct HashWriter {
	state: Blake2b,
}

impl HashWriter {
	/// Consume the `HashWriter`, outputting its current hash into a 32-byte
	/// array
	pub fn finalize(self, output: &mut [u8]) {
		output.copy_from_slice(self.state.finalize().as_bytes());
	}

	/// Consume the `HashWriter`, outputting a `Hash` corresponding to its
	/// current state
	pub fn into_hash(self) -> Hash {
		let mut res = [0; 32];
		res.copy_from_slice(self.state.finalize().as_bytes());
		Hash(res)
	}
}

impl Default for HashWriter {
	fn default() -> HashWriter {
		HashWriter {
			state: Blake2b::new(32),
		}
	}
}

impl ser::Writer for HashWriter {
	fn serialization_mode(&self) -> ser::SerializationMode {
		ser::SerializationMode::Hash
	}

	fn write_fixed_bytes<T: AsFixedBytes>(&mut self, b32: &T) -> Result<(), ser::Error> {
		self.state.update(b32.as_ref());
		Ok(())
	}
}

/// A trait for types that have a canonical hash
pub trait Hashed {
	/// Obtain the hash of the object
	fn hash(&self) -> Hash;
	/// Hash the object together with another writeable object
	fn hash_with<T: Writeable>(&self, other: T) -> Hash;
}

impl<W: ser::Writeable> Hashed for W {
	fn hash(&self) -> Hash {
		let mut hasher = HashWriter::default();
		ser::Writeable::write(self, &mut hasher).unwrap();
		let mut ret = [0; 32];
		hasher.finalize(&mut ret);
		Hash(ret)
	}

	fn hash_with<T: Writeable>(&self, other: T) -> Hash {
		let mut hasher = HashWriter::default();
		ser::Writeable::write(self, &mut hasher).unwrap();
		ser::Writeable::write(&other, &mut hasher).unwrap();
		let mut ret = [0; 32];
		hasher.finalize(&mut ret);
		Hash(ret)
	}
}
<|endoftext|>


<commit_before>#[macro_use]
pub mod types;

use self::types::*;
use common::combinator::Parser;
use std::iter::FromIterator;

pub fn tokenize<T: Iterator<Item = char>>(input: T) -> Result<Vec<Token>> {
    let mut l = Lexer::new(input);

    let mut result = Vec::default();

    loop {
        let _ = skip_whitespaces(&mut l);
        match l.preview() {
            None => break,
            _ => {
                result.push(try!(lex_token(&mut l)));
            }
        }
    }

    Ok(result)
}

fn skip_whitespaces(l: &mut Lexer) -> Result<Vec<char>> {
    l.many(|l| l.predicate(|x| [' ', '\t', '\n', '\r'].contains(x)))
}

fn lex_token(l: &mut Lexer) -> Result<Token> {
    l.choose(&[&lex_operator, &lex_punctuation, &lex_string])
}

macro_rules! parse_map {
    ($pat:expr, $tok:ident) => {
        &|l| {
            let pos = l.current_pos();
            let lit: String = try!(l.string($pat.chars()));
            Ok(token!($tok, pos.0, pos.1, lit))
        }
    }
}

fn lex_operator(l: &mut Lexer) -> Result<Token> {
    l.choose(&[parse_map!("==", Eq),
               parse_map!("=", Assign),
               parse_map!("+", Plus),
               parse_map!("-", Minus),
               parse_map!("*", Multiply),
               parse_map!("/", Divide),
               parse_map!("!=", NotEq),
               parse_map!("!", Not),
               parse_map!(">", GreaterThan),
               parse_map!("<", LessThan)])
}

fn lex_punctuation(l: &mut Lexer) -> Result<Token> {
    l.choose(&[parse_map!(":", Colon),
               parse_map!(";", SemiColon),
               parse_map!(",", Comma),
               parse_map!("(", LParen),
               parse_map!(")", RParen),
               parse_map!("{", LBrace),
               parse_map!("}", RBrace),
               parse_map!("[", LBracket),
               parse_map!("]", RBracket)])
}

fn lex_string(l: &mut Lexer) -> Result<Token> {
    let pos = l.current_pos();
    let mut result = Vec::default();

    result.push(try!(l.atom('"')));
    loop {
        let c = try!(l.next());
        result.push(c);
        match c {
            '\\' => result.push(try!(l.next())),
            '"' => break,
            _ => continue,
        }
    }

    Ok(token!(StringLiteral, pos.0, pos.1, String::from_iter(result)))
}

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

    fn lex(input: &str) -> Vec<Token> {
        tokenize(String::from(input).chars()).unwrap()
    }

    fn vec_eq<T: PartialEq + Debug>(xs: &Vec<T>, ys: &Vec<T>) {
        if xs.len() != ys.len() {
            panic!("{:?} != {:?}", xs, ys);
        }
        for i in 0..xs.len() {
            if xs[i] != ys[i] {
                panic!("{:?} != {:?}", xs, ys);
            }
        }
    }

    #[test]
    fn special_chars() {
        vec_eq(&lex("=+(){},;"),
               &vec![token!(Assign, 1, 1, '='),
                     token!(Plus, 1, 2, '+'),
                     token!(LParen, 1, 3, '('),
                     token!(RParen, 1, 4, ')'),
                     token!(LBrace, 1, 5, '{'),
                     token!(RBrace, 1, 6, '}'),
                     token!(Comma, 1, 7, ','),
                     token!(SemiColon, 1, 8, ';')]);
    }

    #[test]
    fn complex_code() {
        vec_eq(&lex("
let five = 5;
let ten = 10;
let add = fn(x, y) {
    x + y;
};
let result = add(five, ten);
               "),
               &vec![token!(Let, 2, 1, "let"),
                     token!(Ident, 2, 5, "five"),
                     token!(Assign, 2, 10, "="),
                     token!(IntLiteral, 2, 12, "5"),
                     token!(SemiColon, 2, 13, ";"),
                     token!(Let, 3, 1, "let"),
                     token!(Ident, 3, 5, "ten"),
                     token!(Assign, 3, 9, "="),
                     token!(IntLiteral, 3, 11, "10"),
                     token!(SemiColon, 3, 13, ";"),
                     token!(Let, 4, 1, "let"),
                     token!(Ident, 4, 5, "add"),
                     token!(Assign, 4, 9, "="),
                     token!(Function, 4, 11, "fn"),
                     token!(LParen, 4, 13, "("),
                     token!(Ident, 4, 14, "x"),
                     token!(Comma, 4, 15, ","),
                     token!(Ident, 4, 17, "y"),
                     token!(RParen, 4, 18, ")"),
                     token!(LBrace, 4, 20, "{"),
                     token!(Ident, 5, 5, "x"),
                     token!(Plus, 5, 7, "+"),
                     token!(Ident, 5, 9, "y"),
                     token!(SemiColon, 5, 10, ";"),
                     token!(RBrace, 6, 1, "}"),
                     token!(SemiColon, 6, 2, ";"),
                     token!(Let, 7, 1, "let"),
                     token!(Ident, 7, 5, "result"),
                     token!(Assign, 7, 12, "="),
                     token!(Ident, 7, 14, "add"),
                     token!(LParen, 7, 17, "("),
                     token!(Ident, 7, 18, "five"),
                     token!(Comma, 7, 22, ","),
                     token!(Ident, 7, 24, "ten"),
                     token!(RParen, 7, 27, ")"),
                     token!(SemiColon, 7, 28, ";")]);
    }

    #[test]
    fn complex_code2() {
        vec_eq(&lex("
if (a == 10) {
  return a;
} else if (a != 20) {
  return !a;
} else if (a > 20) {
  return -30 / 40 * 50;
} else if (a < 30) {
  return true;
}
return false;
                    "),
               &vec![token!(If, 2, 1, "if"),
                     token!(LParen, 2, 3, "("),
                     token!(Ident, 2, 4, "a"),
                     token!(Eq, 2, 6, "=="),
                     token!(IntLiteral, 2, 9, "10"),
                     token!(RParen, 2, 11, ")"),
                     token!(LBrace, 2, 13, "{"),
                     token!(Return, 3, 3, "return"),
                     token!(Ident, 3, 10, "a"),
                     token!(SemiColon, 3, 11, ";"),
                     token!(RBrace, 4, 1, "}"),
                     token!(Else, 4, 3, "else"),
                     token!(If, 4, 8, "if"),
                     token!(LParen, 4, 11, "("),
                     token!(Ident, 4, 12, "a"),
                     token!(NotEq, 4, 14, "!="),
                     token!(IntLiteral, 4, 17, "20"),
                     token!(RParen, 4, 19, ")"),
                     token!(LBrace, 4, 21, "{"),
                     token!(Return, 5, 3, "return"),
                     token!(Not, 5, 10, "!"),
                     token!(Ident, 5, 11, "a"),
                     token!(SemiColon, 5, 13, ";"),
                     token!(RBrace, 6, 1, "}"),
                     token!(Else, 6, 3, "else"),
                     token!(If, 6, 8, "if"),
                     token!(LParen, 6, 11, "("),
                     token!(Ident, 6, 12, "a"),
                     token!(GreaterThan, 6, 14, ">"),
                     token!(IntLiteral, 6, 16, "20"),
                     token!(RParen, 6, 18, ")"),
                     token!(LBrace, 6, 20, "{"),
                     token!(Return, 7, 3, "return"),
                     token!(Minus, 7, 10, "-"),
                     token!(IntLiteral, 7, 11, "30"),
                     token!(Divide, 7, 14, "/"),
                     token!(IntLiteral, 7, 16, "40"),
                     token!(Multiply, 7, 19, "*"),
                     token!(IntLiteral, 7, 21, "50"),
                     token!(SemiColon, 7, 23, ";"),
                     token!(RBrace, 8, 1, "}"),
                     token!(Else, 8, 3, "else"),
                     token!(If, 8, 8, "if"),
                     token!(LParen, 8, 11, "("),
                     token!(Ident, 8, 12, "a"),
                     token!(LessThan, 8, 14, "<"),
                     token!(IntLiteral, 8, 16, "30"),
                     token!(RParen, 8, 18, ")"),
                     token!(LBrace, 8, 20, "{"),
                     token!(Return, 9, 3, "return"),
                     token!(BoolLiteral, 9, 10, "true"),
                     token!(SemiColon, 9, 14, ";"),
                     token!(RBrace, 10, 1, "}"),
                     token!(Return, 11, 1, "return"),
                     token!(BoolLiteral, 11, 8, "false"),
                     token!(SemiColon, 11, 13, ";")]);
    }

    #[test]
    fn id_with_numbers() {
        vec_eq(&lex("hello2 hel301oo120"),
               &vec![token!(Ident, 1, 1, "hello2"),
                     token!(Ident, 1, 8, "hel301oo120")]);
    }

    #[test]
    fn string_literals() {
        vec_eq(&lex("\"foobar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foobar\"")]);
        vec_eq(&lex("\"foo bar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo bar\"")]);
        vec_eq(&lex("\"foo\\nbar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo\\nbar\"")]);
        vec_eq(&lex("\"foo\\tbar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo\\tbar\"")]);
        vec_eq(&lex("\"foo\\\"bar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo\\\"bar\"")]);
    }

    #[test]
    fn array_tokens() {
        vec_eq(&lex("[1, 2];"),
               &vec![token!(LBracket, 1, 1, "["),
                     token!(IntLiteral, 1, 2, "1"),
                     token!(Comma, 1, 3, ","),
                     token!(IntLiteral, 1, 5, "2"),
                     token!(RBracket, 1, 6, "]"),
                     token!(SemiColon, 1, 7, ";")]);
    }

    #[test]
    fn hash_tokens() {
        vec_eq(&lex("{\"hello\": \"world\"}"),
               &vec![token!(LBrace, 1, 1, "{"),
                    token!(StringLiteral, 1, 2, "\"hello\""),
                    token!(Colon, 1, 9, ":"),
                    token!(StringLiteral, 1, 11, "\"world\""),
                    token!(RBrace, 1, 18, "}")]);
    }
}
<commit_msg>Implement lexer for reserved words and ident<commit_after>#[macro_use]
pub mod types;

use self::types::*;
use common::combinator::Parser;
use common::util::*;
use std::iter::FromIterator;

pub fn tokenize<T: Iterator<Item = char>>(input: T) -> Result<Vec<Token>> {
    let mut l = Lexer::new(input);

    let mut result = Vec::default();

    loop {
        let _ = skip_whitespaces(&mut l);
        match l.preview() {
            None => break,
            _ => {
                result.push(try!(lex_token(&mut l)));
            }
        }
    }

    Ok(result)
}

fn skip_whitespaces(l: &mut Lexer) -> Result<Vec<char>> {
    l.many(|l| l.predicate(|x| [' ', '\t', '\n', '\r'].contains(x)))
}

fn lex_token(l: &mut Lexer) -> Result<Token> {
    l.choose(&[&lex_operator,
               &lex_punctuation,
               &lex_string,
               &lex_reserved_or_ident])
}

macro_rules! parse_map {
    ($pat:expr, $tok:ident) => {
        &|l| {
            let pos = l.current_pos();
            let lit: String = try!(l.string($pat.chars()));
            Ok(token!($tok, pos.0, pos.1, lit))
        }
    }
}

fn lex_operator(l: &mut Lexer) -> Result<Token> {
    l.choose(&[parse_map!("==", Eq),
               parse_map!("=", Assign),
               parse_map!("+", Plus),
               parse_map!("-", Minus),
               parse_map!("*", Multiply),
               parse_map!("/", Divide),
               parse_map!("!=", NotEq),
               parse_map!("!", Not),
               parse_map!(">", GreaterThan),
               parse_map!("<", LessThan)])
}

fn lex_punctuation(l: &mut Lexer) -> Result<Token> {
    l.choose(&[parse_map!(":", Colon),
               parse_map!(";", SemiColon),
               parse_map!(",", Comma),
               parse_map!("(", LParen),
               parse_map!(")", RParen),
               parse_map!("{", LBrace),
               parse_map!("}", RBrace),
               parse_map!("[", LBracket),
               parse_map!("]", RBracket)])
}

fn lex_string(l: &mut Lexer) -> Result<Token> {
    let pos = l.current_pos();
    let mut result = Vec::default();

    result.push(try!(l.atom('"')));
    loop {
        let c = try!(l.next());
        result.push(c);
        match c {
            '\\' => result.push(try!(l.next())),
            '"' => break,
            _ => continue,
        }
    }

    Ok(token!(StringLiteral, pos.0, pos.1, String::from_iter(result)))
}

fn lex_reserved_or_ident(l: &mut Lexer) -> Result<Token> {
    let pos = l.current_pos();
    let fst = try!(l.predicate(is_letter)).to_string();
    let rest: String = try!(l.many(|l| l.predicate(is_letter_or_digit)));
    let name = fst + &rest;

    Ok(match name.as_ref() {
           "let" => token!(Let, pos.0, pos.1, name),
           "fn" => token!(Function, pos.0, pos.1, name),
           "if" => token!(If, pos.0, pos.1, name),
           "else" => token!(Else, pos.0, pos.1, name),
           "return" => token!(Return, pos.0, pos.1, name),
           "true" => token!(BoolLiteral, pos.0, pos.1, name),
           "false" => token!(BoolLiteral, pos.0, pos.1, name),
           _ => token!(Ident, pos.0, pos.1, name),
       })
}

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

    fn lex(input: &str) -> Vec<Token> {
        tokenize(String::from(input).chars()).unwrap()
    }

    fn vec_eq<T: PartialEq + Debug>(xs: &Vec<T>, ys: &Vec<T>) {
        if xs.len() != ys.len() {
            panic!("{:?} != {:?}", xs, ys);
        }
        for i in 0..xs.len() {
            if xs[i] != ys[i] {
                panic!("{:?} != {:?}", xs, ys);
            }
        }
    }

    #[test]
    fn special_chars() {
        vec_eq(&lex("=+(){},;"),
               &vec![token!(Assign, 1, 1, '='),
                     token!(Plus, 1, 2, '+'),
                     token!(LParen, 1, 3, '('),
                     token!(RParen, 1, 4, ')'),
                     token!(LBrace, 1, 5, '{'),
                     token!(RBrace, 1, 6, '}'),
                     token!(Comma, 1, 7, ','),
                     token!(SemiColon, 1, 8, ';')]);
    }

    #[test]
    fn complex_code() {
        vec_eq(&lex("
let five = 5;
let ten = 10;
let add = fn(x, y) {
    x + y;
};
let result = add(five, ten);
               "),
               &vec![token!(Let, 2, 1, "let"),
                     token!(Ident, 2, 5, "five"),
                     token!(Assign, 2, 10, "="),
                     token!(IntLiteral, 2, 12, "5"),
                     token!(SemiColon, 2, 13, ";"),
                     token!(Let, 3, 1, "let"),
                     token!(Ident, 3, 5, "ten"),
                     token!(Assign, 3, 9, "="),
                     token!(IntLiteral, 3, 11, "10"),
                     token!(SemiColon, 3, 13, ";"),
                     token!(Let, 4, 1, "let"),
                     token!(Ident, 4, 5, "add"),
                     token!(Assign, 4, 9, "="),
                     token!(Function, 4, 11, "fn"),
                     token!(LParen, 4, 13, "("),
                     token!(Ident, 4, 14, "x"),
                     token!(Comma, 4, 15, ","),
                     token!(Ident, 4, 17, "y"),
                     token!(RParen, 4, 18, ")"),
                     token!(LBrace, 4, 20, "{"),
                     token!(Ident, 5, 5, "x"),
                     token!(Plus, 5, 7, "+"),
                     token!(Ident, 5, 9, "y"),
                     token!(SemiColon, 5, 10, ";"),
                     token!(RBrace, 6, 1, "}"),
                     token!(SemiColon, 6, 2, ";"),
                     token!(Let, 7, 1, "let"),
                     token!(Ident, 7, 5, "result"),
                     token!(Assign, 7, 12, "="),
                     token!(Ident, 7, 14, "add"),
                     token!(LParen, 7, 17, "("),
                     token!(Ident, 7, 18, "five"),
                     token!(Comma, 7, 22, ","),
                     token!(Ident, 7, 24, "ten"),
                     token!(RParen, 7, 27, ")"),
                     token!(SemiColon, 7, 28, ";")]);
    }

    #[test]
    fn complex_code2() {
        vec_eq(&lex("
if (a == 10) {
  return a;
} else if (a != 20) {
  return !a;
} else if (a > 20) {
  return -30 / 40 * 50;
} else if (a < 30) {
  return true;
}
return false;
                    "),
               &vec![token!(If, 2, 1, "if"),
                     token!(LParen, 2, 3, "("),
                     token!(Ident, 2, 4, "a"),
                     token!(Eq, 2, 6, "=="),
                     token!(IntLiteral, 2, 9, "10"),
                     token!(RParen, 2, 11, ")"),
                     token!(LBrace, 2, 13, "{"),
                     token!(Return, 3, 3, "return"),
                     token!(Ident, 3, 10, "a"),
                     token!(SemiColon, 3, 11, ";"),
                     token!(RBrace, 4, 1, "}"),
                     token!(Else, 4, 3, "else"),
                     token!(If, 4, 8, "if"),
                     token!(LParen, 4, 11, "("),
                     token!(Ident, 4, 12, "a"),
                     token!(NotEq, 4, 14, "!="),
                     token!(IntLiteral, 4, 17, "20"),
                     token!(RParen, 4, 19, ")"),
                     token!(LBrace, 4, 21, "{"),
                     token!(Return, 5, 3, "return"),
                     token!(Not, 5, 10, "!"),
                     token!(Ident, 5, 11, "a"),
                     token!(SemiColon, 5, 13, ";"),
                     token!(RBrace, 6, 1, "}"),
                     token!(Else, 6, 3, "else"),
                     token!(If, 6, 8, "if"),
                     token!(LParen, 6, 11, "("),
                     token!(Ident, 6, 12, "a"),
                     token!(GreaterThan, 6, 14, ">"),
                     token!(IntLiteral, 6, 16, "20"),
                     token!(RParen, 6, 18, ")"),
                     token!(LBrace, 6, 20, "{"),
                     token!(Return, 7, 3, "return"),
                     token!(Minus, 7, 10, "-"),
                     token!(IntLiteral, 7, 11, "30"),
                     token!(Divide, 7, 14, "/"),
                     token!(IntLiteral, 7, 16, "40"),
                     token!(Multiply, 7, 19, "*"),
                     token!(IntLiteral, 7, 21, "50"),
                     token!(SemiColon, 7, 23, ";"),
                     token!(RBrace, 8, 1, "}"),
                     token!(Else, 8, 3, "else"),
                     token!(If, 8, 8, "if"),
                     token!(LParen, 8, 11, "("),
                     token!(Ident, 8, 12, "a"),
                     token!(LessThan, 8, 14, "<"),
                     token!(IntLiteral, 8, 16, "30"),
                     token!(RParen, 8, 18, ")"),
                     token!(LBrace, 8, 20, "{"),
                     token!(Return, 9, 3, "return"),
                     token!(BoolLiteral, 9, 10, "true"),
                     token!(SemiColon, 9, 14, ";"),
                     token!(RBrace, 10, 1, "}"),
                     token!(Return, 11, 1, "return"),
                     token!(BoolLiteral, 11, 8, "false"),
                     token!(SemiColon, 11, 13, ";")]);
    }

    #[test]
    fn id_with_numbers() {
        vec_eq(&lex("hello2 hel301oo120"),
               &vec![token!(Ident, 1, 1, "hello2"),
                     token!(Ident, 1, 8, "hel301oo120")]);
    }

    #[test]
    fn string_literals() {
        vec_eq(&lex("\"foobar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foobar\"")]);
        vec_eq(&lex("\"foo bar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo bar\"")]);
        vec_eq(&lex("\"foo\\nbar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo\\nbar\"")]);
        vec_eq(&lex("\"foo\\tbar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo\\tbar\"")]);
        vec_eq(&lex("\"foo\\\"bar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo\\\"bar\"")]);
    }

    #[test]
    fn array_tokens() {
        vec_eq(&lex("[1, 2];"),
               &vec![token!(LBracket, 1, 1, "["),
                     token!(IntLiteral, 1, 2, "1"),
                     token!(Comma, 1, 3, ","),
                     token!(IntLiteral, 1, 5, "2"),
                     token!(RBracket, 1, 6, "]"),
                     token!(SemiColon, 1, 7, ";")]);
    }

    #[test]
    fn hash_tokens() {
        vec_eq(&lex("{\"hello\": \"world\"}"),
               &vec![token!(LBrace, 1, 1, "{"),
                    token!(StringLiteral, 1, 2, "\"hello\""),
                    token!(Colon, 1, 9, ":"),
                    token!(StringLiteral, 1, 11, "\"world\""),
                    token!(RBrace, 1, 18, "}")]);
    }
}
<|endoftext|>


<commit_before>fn on_osmain(f: fn~()) {
    let builder = task::builder();
    task::set_opts(builder, {
        sched: some({
            mode: task::osmain,
            foreign_stack_size: none
        })
        with task::get_opts(builder)
    });
    let po = comm::port();
    let ch = comm::chan(po);
    task::run(builder, || {
        f();
        comm::send(ch, ());
    });
    comm::recv(po);
}

// FIXME: Needs additional cocoa setup on OS X. rust-cocoa should probably just
// be a dependency
#[test]
#[ignore(cfg(target_os = "macos"))]
fn test_everything() {

    on_osmain(|| {
        init([init_video, init_timer]);
        run_tests([
            general::test_was_init,
            general::test_set_error,
            general::test_error,
            general::test_clear_error,
            video::test_set_video_mode,
            // FIXME: Doesn't work when called from a directory that
            // doesn't contain the test image file
            //video::test_blit,
            test_event::test_poll_event_none
            // FIXME: This test is interactive
            //test_event::test_keyboard
        ]);
        quit();
    })
}

fn run_tests(tests: [fn()]) {
    vec::iter(tests, |test| test());
}

mod general {
    fn test_was_init() {
        assert vec::contains(was_init([init_timer]), init_timer);
    }

    fn test_set_error() {
        set_error("test");
        assert get_error() == "test";
    }

    fn test_error() {
        clear_error();
        assert str::is_empty(get_error());
        error(enomem);
        assert str::is_not_empty(get_error());
    }

    fn test_clear_error() {
        set_error("test");
        clear_error();
        assert str::is_empty(get_error());
    }
}

mod test_event {
    fn test_poll_event_none() {
        ::event::poll_event(|event| assert event == ::event::no_event);
    }

    fn test_keyboard() {
        io::println("press a key in the window");
        let surface = ::video::set_video_mode(320, 200, 32,
            [::video::swsurface], [::video::doublebuf, ::video::resizable]);
        let mut keydown = false;
        let mut keyup = false;
        while !keydown || !keyup {
            ::event::poll_event(|event| {
                alt event {
                  event::keyup_event(_) { keyup = true; }
                  event::keydown_event(_) { keydown = true; }
                  _ { }
                }
            })
        }
        ::video::free_surface(surface);
    }
}

mod video {

    fn test_set_video_mode() {
        let surface = ::video::set_video_mode(320, 200, 32,
            [::video::hwsurface], [::video::doublebuf]);
        assert surface != ptr::null();
        ::video::free_surface(surface);
    }

    fn test_blit() {
        let screen = ::video::set_video_mode(320, 200, 32,
            [::video::swsurface], [::video::doublebuf]);
        assert screen != ptr::null();

        let image = {
            // FIXME: We need to load this from the crate instead of
            // off the filesystem
            let tmp = ::video::load_bmp("rust-logo-128x128-blk.bmp");
            assert tmp != ptr::null();
            let image = ::video::display_format(tmp);
            assert image != ptr::null();
            ::video::free_surface(tmp);
            image
        };

        iter::repeat(1u, || {
            ::video::blit_surface(image, ptr::null(),
                                  screen, ptr::null());
            ::video::flip(screen);
            ::event::poll_event(|_event| {})
        });

        ::video::free_surface(image);
        ::video::free_surface(screen);
    }
}
<commit_msg>Update to build with latest rustc syntax changes.<commit_after>fn on_osmain(f: fn~()) {
    let builder = task::builder();
    task::set_opts(builder, {
        sched: some({
            mode: task::osmain,
            foreign_stack_size: none
        })
        with task::get_opts(builder)
    });
    let po = comm::port();
    let ch = comm::chan(po);
    task::run(builder, || {
        f();
        comm::send(ch, ());
    });
    comm::recv(po);
}

// FIXME: Needs additional cocoa setup on OS X. rust-cocoa should probably just
// be a dependency
#[test]
#[ignore(cfg(target_os = "macos"))]
fn test_everything() {

    on_osmain(|| {
        init([init_video, init_timer]);
        run_tests([
            general::test_was_init,
            general::test_set_error,
            general::test_error,
            general::test_clear_error,
            video::test_set_video_mode,
            // FIXME: Doesn't work when called from a directory that
            // doesn't contain the test image file
            //video::test_blit,
            test_event::test_poll_event_none
            // FIXME: This test is interactive
            //test_event::test_keyboard
        ]);
        quit();
    })
}

fn run_tests(tests: [fn()]) {
    vec::iter(tests, |test| test());
}

mod general {
    fn test_was_init() {
        assert vec::contains(was_init([init_timer]), init_timer);
    }

    fn test_set_error() {
        set_error("test");
        assert get_error() == "test";
    }

    fn test_error() {
        clear_error();
        assert str::is_empty(get_error());
        error(enomem);
        assert str::is_not_empty(get_error());
    }

    fn test_clear_error() {
        set_error("test");
        clear_error();
        assert str::is_empty(get_error());
    }
}

mod test_event {
    fn test_poll_event_none() {
        ::event::poll_event(|event| assert event == ::event::no_event);
    }

    fn test_keyboard() {
        io::println("press a key in the window");
        let surface = ::video::set_video_mode(320, 200, 32,
            [::video::swsurface], [::video::doublebuf, ::video::resizable]);
        let mut keydown = false;
        let mut keyup = false;
        while !keydown || !keyup {
            ::event::poll_event(|event| {
                alt event {
                  event::keyup_event(_) { keyup = true; }
                  event::keydown_event(_) { keydown = true; }
                  _ { }
                }
            })
        }
        ::video::free_surface(surface);
    }
}

mod video {

    fn test_set_video_mode() {
        let surface = ::video::set_video_mode(320, 200, 32,
            [::video::hwsurface], [::video::doublebuf]);
        assert surface != ptr::null();
        ::video::free_surface(surface);
    }

    fn test_blit() {
        let screen = ::video::set_video_mode(320, 200, 32,
            [::video::swsurface], [::video::doublebuf]);
        assert screen != ptr::null();

        let image = {
            // FIXME: We need to load this from the crate instead of
            // off the filesystem
            let tmp = ::video::load_bmp("rust-logo-128x128-blk.bmp");
            assert tmp != ptr::null();
            let image = ::video::display_format(tmp);
            assert image != ptr::null();
            ::video::free_surface(tmp);
            image
        };

        for iter::repeat(1u) || {
            ::video::blit_surface(image, ptr::null(),
                                  screen, ptr::null());
            ::video::flip(screen);
            ::event::poll_event(|_event| {})
        };

        ::video::free_surface(image);
        ::video::free_surface(screen);
    }
}
<|endoftext|>


<commit_before>#![feature(macro_rules)]
extern crate opencl = "OpenCL#0.2";

use opencl::hl;

macro_rules! expect (
    ($test: expr, $expected: expr) => ({
            let test     = $test;
            let expected = $expected;
            if test != expected {
                fail!(format!("Test failure in {:s}: expected {:?}, got {:?}",
                              stringify!($test),
                              expected, test))
            }
        })
        )

fn main () {
    /*let src = "__kernel void test(__global int *i, long int k) {
                   *i += k;
                   }";
    let prog = ctx.create_program_from_source(src);
    prog.build(&device).unwrap();

    let k = prog.create_kernel("test");
    let v = ctx.create_buffer_from(&[1], CL_MEM_READ_WRITE);
    k.set_arg(0, &v);
    k.set_arg(1, &42);
    queue.enqueue_async_kernel(&k, 1, None, ()).wait();
    let v: ~[int] = queue.get(&v, ());

    expect!(v[0], 43);*/

    let platforms = hl::get_platforms();
    for p in platforms.iter() {
        println!("platform found");
        println!("platform: {}", p.name());
        let devices = p.get_devices();
        println!("got the devices!!")
        for d in devices.iter() {
            println!("device found");
            let context = d.create_context();
            let queue = context.create_command_queue(d);
            let name = d.name();
            println!("done with the device");
        }
    }

}
<commit_msg>Fix test code<commit_after>#![feature(macro_rules)]
extern crate opencl = "OpenCL#0.2";

use opencl::hl;
use opencl::hl::{GPU, EventList};
use opencl::CL::CL_MEM_READ_WRITE;

macro_rules! expect (
    ($test: expr, $expected: expr) => ({
      let test     = $test;
      let expected = $expected;
      if test != expected {
      fail!(format!("Test failure in {:s}: expected {:?}, got {:?}",
          stringify!($test),
          expected, test))
      }
      })
    )

fn main () {

  let platforms = hl::get_platforms();
  for p in platforms.iter() {
    println!("platform found");
    println!("platform: {}", p.name());
    let devices = p.get_devices_by_types(&[GPU]);
    println!("got the devices!!")
      for d in devices.iter() {
        println!("device found");
        let ctx = d.create_context();
        let queue = ctx.create_command_queue(d);
        let name = d.name();
        let src = "__kernel void test(__global int *i, long int k) {
          *i += k;
        }";
        let prog = ctx.create_program_from_source(src);
        prog.build(d).unwrap();

        let k = prog.create_kernel("test");
        let v = ctx.create_buffer_from(&[1], CL_MEM_READ_WRITE);
        k.set_arg(0, &v);
        k.set_arg(1, &42);
        queue.enqueue_async_kernel(&k, 1, None, ()).wait();
        let v: Vec<int> = queue.get(&v, ());

        expect!(*v.get(0), 43);

        println!("done with the device");
      }
  }

}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use common::{Diagnostic, Location, NamedItem, WithLocation};
use graphql_ir::{
    Argument, ConstantValue, Directive, Program, ScalarField, Selection, Transformed, Transformer,
    ValidationMessage, ValidationResult, Value,
};
use interner::{Intern, StringKey};
use lazy_static::lazy_static;
use schema::{AstValue, Field, FieldID, Type};
use std::sync::Arc;

lazy_static! {
    pub static ref REACT_FLIGHT_DIRECTIVE_NAME: StringKey = "__ReactFlightComponent".intern();
    static ref REACT_FLIGHT_COMPONENT_ARGUMENT_NAME: StringKey = "component".intern();
    static ref REACT_FLIGHT_PROPS_ARGUMENT_NAME: StringKey = "props".intern();
    static ref REACT_FLIGHT_PROPS_TYPE: StringKey = "ReactFlightProps".intern();
    static ref REACT_FLIGHT_COMPONENT_TYPE: StringKey = "ReactFlightComponent".intern();
    static ref REACT_FLIGHT_FIELD_NAME: StringKey = "flight".intern();
    static ref REACT_FLIGHT_EXTENSION_DIRECTIVE_NAME: StringKey = "react_flight_component".intern();
    static ref NAME: StringKey = "name".intern();
}

/// Transform to find calls to React Flight schema extension fields and rewrite them into calls
/// to a generic `flight(component, props)` field.
pub fn react_flight(program: &Program) -> ValidationResult<Program> {
    // No-op unless the special props/component types and flight directive are defined
    let props_type = program.schema.get_type(*REACT_FLIGHT_PROPS_TYPE);
    let component_type = program.schema.get_type(*REACT_FLIGHT_COMPONENT_TYPE);
    let has_flight_directive = program
        .schema
        .has_directive(*REACT_FLIGHT_EXTENSION_DIRECTIVE_NAME);
    let (props_type, component_type) = match (props_type, component_type) {
        (Some(props_type), Some(component_type)) if has_flight_directive => {
            (props_type, component_type)
        }
        _ => return Ok(program.clone()),
    };
    let mut transform = ReactFlightTransform::new(program, props_type, component_type);
    if transform.errors.is_empty() {
        Ok(transform
            .transform_program(program)
            .replace_or_else(|| program.clone()))
    } else {
        Err(transform.errors)
    }
}

struct ReactFlightTransform<'s> {
    component_type: Type,
    errors: Vec<Diagnostic>,
    program: &'s Program,
    props_type: Type,
}

impl<'s> ReactFlightTransform<'s> {
    fn new(program: &'s Program, props_type: Type, component_type: Type) -> Self {
        Self {
            component_type,
            errors: Default::default(),
            program,
            props_type,
        }
    }

    fn get_component_name_for_field(
        &mut self,
        field_definition: &Field,
        location: Location,
    ) -> Result<StringKey, ()> {
        // the field definition must specify the backing component's module name
        let component_directive = match field_definition
            .directives
            .named(*REACT_FLIGHT_EXTENSION_DIRECTIVE_NAME)
        {
            Some(component_directive) => component_directive,
            None => {
                self.errors.push(Diagnostic::error(
                    ValidationMessage::InvalidFlightFieldMissingModuleDirective,
                    location,
                ));
                return Err(());
            }
        };
        // extract the component name
        let value = component_directive
            .arguments
            .iter()
            .cloned()
            .find(|arg| arg.name == *NAME)
            .unwrap()
            .value;
        match value {
            AstValue::String(name) => Ok(name.intern()),
            _ => {
                self.errors.push(Diagnostic::error(
                    ValidationMessage::InvalidFlightFieldExpectedModuleNameString,
                    location,
                ));
                Err(())
            }
        }
    }

    // validates that the field's parent type also has a field conforming to the
    // following specification:
    //
    // ```
    // flight(
    //   component: String
    //   props: ReactFlightProps
    // ): ReactFlightComponent
    // ```
    fn validate_flight_field(
        &mut self,
        field_definition: &Field,
        location: Location,
    ) -> Result<FieldID, ()> {
        // not a built-in field, so there must be a parent type or the schema is
        // invalid (which is a compiler error not a user error)
        let parent_type = field_definition
            .parent_type
            .unwrap_or_else(|| panic!("Expected field to have a parent type"));

        // the parent type must have the generic `flight` field
        let flight_field_id = match self
            .program
            .schema
            .named_field(parent_type, *REACT_FLIGHT_FIELD_NAME)
        {
            Some(flight_field_id) => flight_field_id,
            None => {
                self.errors.push(Diagnostic::error(
                    ValidationMessage::InvalidFlightFieldNotDefinedOnType {
                        field_name: field_definition.name,
                    },
                    location,
                ));
                return Err(());
            }
        };
        let flight_field_definition = self.program.schema.field(flight_field_id);

        // flight field must have `props: ReactFlightProps` arg
        let props_argument = flight_field_definition.arguments.iter().find(|arg| {
            arg.name == *REACT_FLIGHT_PROPS_ARGUMENT_NAME && arg.type_.inner() == self.props_type
        });
        if props_argument.is_none() {
            self.errors.push(Diagnostic::error(
                ValidationMessage::InvalidFlightFieldPropsArgument,
                location,
            ));
            return Err(());
        }
        // flight field must have `component: String` arg
        let component_argument = flight_field_definition.arguments.iter().find(|arg| {
            arg.name == *REACT_FLIGHT_COMPONENT_ARGUMENT_NAME
                && Some(arg.type_.inner()) == self.program.schema.get_type("String".intern())
        });
        if component_argument.is_none() {
            self.errors.push(Diagnostic::error(
                ValidationMessage::InvalidFlightFieldComponentArgument,
                location,
            ));
            return Err(());
        }
        // flight field must return `ReactFlightComponent`
        if flight_field_definition.type_.inner() != self.component_type {
            self.errors.push(Diagnostic::error(
                ValidationMessage::InvalidFlightFieldReturnType,
                location,
            ));
            return Err(());
        }
        Ok(flight_field_id)
    }
}

impl<'s> Transformer for ReactFlightTransform<'s> {
    const NAME: &'static str = "ReactFlightTransform";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_scalar_field(&mut self, field: &ScalarField) -> Transformed<Selection> {
        let field_definition = self.program.schema.field(field.definition.item);
        // Activate the transform based on the return type since this is a fast check
        if field_definition.type_.inner() != self.component_type {
            return Transformed::Keep;
        }

        // Extract the backing component's name from the field definition
        let component_name =
            match self.get_component_name_for_field(field_definition, field.definition.location) {
                Ok(value) => value,
                Err(_) => return Transformed::Keep,
            };

        // Determine the type's `flight` field
        let flight_field_id =
            match self.validate_flight_field(field_definition, field.definition.location) {
                Ok(value) => value,
                Err(_) => return Transformed::Keep,
            };

        // Rewrite into a call to the `flight` field, passing the original arguments
        // as values of the `props` argument:
        let alias = field
            .alias
            .unwrap_or_else(|| WithLocation::generated(field_definition.name));
        let mut directives = Vec::with_capacity(field.directives.len() + 1);
        directives.extend(field.directives.iter().cloned());
        directives.push(Directive {
            name: WithLocation::generated(*REACT_FLIGHT_DIRECTIVE_NAME),
            arguments: vec![],
        });
        Transformed::Replace(Selection::ScalarField(Arc::new(ScalarField {
            alias: Some(alias),
            arguments: vec![
                Argument {
                    name: WithLocation::generated(*REACT_FLIGHT_COMPONENT_ARGUMENT_NAME),
                    value: WithLocation::generated(Value::Constant(ConstantValue::String(
                        component_name,
                    ))),
                },
                Argument {
                    name: WithLocation::generated(*REACT_FLIGHT_PROPS_ARGUMENT_NAME),
                    value: WithLocation::generated(Value::Object(field.arguments.clone())),
                },
            ],
            definition: WithLocation::generated(flight_field_id),
            directives,
        })))
    }
}
<commit_msg>Remove unnecessary check in react_flight transform<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use common::{Diagnostic, Location, NamedItem, WithLocation};
use graphql_ir::{
    Argument, ConstantValue, Directive, Program, ScalarField, Selection, Transformed, Transformer,
    ValidationMessage, ValidationResult, Value,
};
use interner::{Intern, StringKey};
use lazy_static::lazy_static;
use schema::{AstValue, Field, FieldID, Type};
use std::sync::Arc;

lazy_static! {
    pub static ref REACT_FLIGHT_DIRECTIVE_NAME: StringKey = "__ReactFlightComponent".intern();
    static ref REACT_FLIGHT_COMPONENT_ARGUMENT_NAME: StringKey = "component".intern();
    static ref REACT_FLIGHT_PROPS_ARGUMENT_NAME: StringKey = "props".intern();
    static ref REACT_FLIGHT_PROPS_TYPE: StringKey = "ReactFlightProps".intern();
    static ref REACT_FLIGHT_COMPONENT_TYPE: StringKey = "ReactFlightComponent".intern();
    static ref REACT_FLIGHT_FIELD_NAME: StringKey = "flight".intern();
    static ref REACT_FLIGHT_EXTENSION_DIRECTIVE_NAME: StringKey = "react_flight_component".intern();
    static ref NAME: StringKey = "name".intern();
}

/// Transform to find calls to React Flight schema extension fields and rewrite them into calls
/// to a generic `flight(component, props)` field.
pub fn react_flight(program: &Program) -> ValidationResult<Program> {
    // No-op unless the special props/component types and flight directive are defined
    let props_type = program.schema.get_type(*REACT_FLIGHT_PROPS_TYPE);
    let component_type = program.schema.get_type(*REACT_FLIGHT_COMPONENT_TYPE);
    let (props_type, component_type) = match (props_type, component_type) {
        (Some(props_type), Some(component_type)) => (props_type, component_type),
        _ => return Ok(program.clone()),
    };
    let mut transform = ReactFlightTransform::new(program, props_type, component_type);
    if transform.errors.is_empty() {
        Ok(transform
            .transform_program(program)
            .replace_or_else(|| program.clone()))
    } else {
        Err(transform.errors)
    }
}

struct ReactFlightTransform<'s> {
    component_type: Type,
    errors: Vec<Diagnostic>,
    program: &'s Program,
    props_type: Type,
}

impl<'s> ReactFlightTransform<'s> {
    fn new(program: &'s Program, props_type: Type, component_type: Type) -> Self {
        Self {
            component_type,
            errors: Default::default(),
            program,
            props_type,
        }
    }

    fn get_component_name_for_field(
        &mut self,
        field_definition: &Field,
        location: Location,
    ) -> Result<StringKey, ()> {
        // the field definition must specify the backing component's module name
        let component_directive = match field_definition
            .directives
            .named(*REACT_FLIGHT_EXTENSION_DIRECTIVE_NAME)
        {
            Some(component_directive) => component_directive,
            None => {
                self.errors.push(Diagnostic::error(
                    ValidationMessage::InvalidFlightFieldMissingModuleDirective,
                    location,
                ));
                return Err(());
            }
        };
        // extract the component name
        let value = component_directive
            .arguments
            .iter()
            .cloned()
            .find(|arg| arg.name == *NAME)
            .unwrap()
            .value;
        match value {
            AstValue::String(name) => Ok(name.intern()),
            _ => {
                self.errors.push(Diagnostic::error(
                    ValidationMessage::InvalidFlightFieldExpectedModuleNameString,
                    location,
                ));
                Err(())
            }
        }
    }

    // validates that the field's parent type also has a field conforming to the
    // following specification:
    //
    // ```
    // flight(
    //   component: String
    //   props: ReactFlightProps
    // ): ReactFlightComponent
    // ```
    fn validate_flight_field(
        &mut self,
        field_definition: &Field,
        location: Location,
    ) -> Result<FieldID, ()> {
        // not a built-in field, so there must be a parent type or the schema is
        // invalid (which is a compiler error not a user error)
        let parent_type = field_definition
            .parent_type
            .unwrap_or_else(|| panic!("Expected field to have a parent type"));

        // the parent type must have the generic `flight` field
        let flight_field_id = match self
            .program
            .schema
            .named_field(parent_type, *REACT_FLIGHT_FIELD_NAME)
        {
            Some(flight_field_id) => flight_field_id,
            None => {
                self.errors.push(Diagnostic::error(
                    ValidationMessage::InvalidFlightFieldNotDefinedOnType {
                        field_name: field_definition.name,
                    },
                    location,
                ));
                return Err(());
            }
        };
        let flight_field_definition = self.program.schema.field(flight_field_id);

        // flight field must have `props: ReactFlightProps` arg
        let props_argument = flight_field_definition.arguments.iter().find(|arg| {
            arg.name == *REACT_FLIGHT_PROPS_ARGUMENT_NAME && arg.type_.inner() == self.props_type
        });
        if props_argument.is_none() {
            self.errors.push(Diagnostic::error(
                ValidationMessage::InvalidFlightFieldPropsArgument,
                location,
            ));
            return Err(());
        }
        // flight field must have `component: String` arg
        let component_argument = flight_field_definition.arguments.iter().find(|arg| {
            arg.name == *REACT_FLIGHT_COMPONENT_ARGUMENT_NAME
                && Some(arg.type_.inner()) == self.program.schema.get_type("String".intern())
        });
        if component_argument.is_none() {
            self.errors.push(Diagnostic::error(
                ValidationMessage::InvalidFlightFieldComponentArgument,
                location,
            ));
            return Err(());
        }
        // flight field must return `ReactFlightComponent`
        if flight_field_definition.type_.inner() != self.component_type {
            self.errors.push(Diagnostic::error(
                ValidationMessage::InvalidFlightFieldReturnType,
                location,
            ));
            return Err(());
        }
        Ok(flight_field_id)
    }
}

impl<'s> Transformer for ReactFlightTransform<'s> {
    const NAME: &'static str = "ReactFlightTransform";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_scalar_field(&mut self, field: &ScalarField) -> Transformed<Selection> {
        let field_definition = self.program.schema.field(field.definition.item);
        // Activate the transform based on the return type since this is a fast check
        if field_definition.type_.inner() != self.component_type {
            return Transformed::Keep;
        }

        // Extract the backing component's name from the field definition
        let component_name =
            match self.get_component_name_for_field(field_definition, field.definition.location) {
                Ok(value) => value,
                Err(_) => return Transformed::Keep,
            };

        // Determine the type's `flight` field
        let flight_field_id =
            match self.validate_flight_field(field_definition, field.definition.location) {
                Ok(value) => value,
                Err(_) => return Transformed::Keep,
            };

        // Rewrite into a call to the `flight` field, passing the original arguments
        // as values of the `props` argument:
        let alias = field
            .alias
            .unwrap_or_else(|| WithLocation::generated(field_definition.name));
        let mut directives = Vec::with_capacity(field.directives.len() + 1);
        directives.extend(field.directives.iter().cloned());
        directives.push(Directive {
            name: WithLocation::generated(*REACT_FLIGHT_DIRECTIVE_NAME),
            arguments: vec![],
        });
        Transformed::Replace(Selection::ScalarField(Arc::new(ScalarField {
            alias: Some(alias),
            arguments: vec![
                Argument {
                    name: WithLocation::generated(*REACT_FLIGHT_COMPONENT_ARGUMENT_NAME),
                    value: WithLocation::generated(Value::Constant(ConstantValue::String(
                        component_name,
                    ))),
                },
                Argument {
                    name: WithLocation::generated(*REACT_FLIGHT_PROPS_ARGUMENT_NAME),
                    value: WithLocation::generated(Value::Object(field.arguments.clone())),
                },
            ],
            definition: WithLocation::generated(flight_field_id),
            directives,
        })))
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// 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.

use std::sync::{mpsc, Arc, RwLock};
use std::time::{Duration, Instant};
use std::thread::{self, JoinHandle};

use linked_hash_map::LinkedHashMap;
use hab_net::server::ZMQ_CONTEXT;
use protobuf::{parse_from_bytes, Message};
use protocol::jobsrv;
use zmq;

use config::Config;
use data_store::DataStore;
use error::Result;

const WORKER_MGR_ADDR: &'static str = "inproc://work-manager";
const WORKER_TIMEOUT_MS: u64 = 33_000;

pub struct WorkerMgrClient {
    socket: zmq::Socket,
}

impl WorkerMgrClient {
    pub fn connect(&mut self) -> Result<()> {
        try!(self.socket.connect(WORKER_MGR_ADDR));
        Ok(())
    }

    pub fn notify_work(&mut self) -> Result<()> {
        try!(self.socket.send(&[1], 0));
        Ok(())
    }
}

impl Default for WorkerMgrClient {
    fn default() -> WorkerMgrClient {
        let socket = (**ZMQ_CONTEXT).as_mut().socket(zmq::DEALER).unwrap();
        socket.set_sndhwm(1).unwrap();
        socket.set_linger(0).unwrap();
        socket.set_immediate(true).unwrap();
        WorkerMgrClient { socket: socket }
    }
}

pub struct WorkerMgr {
    config: Arc<RwLock<Config>>,
    datastore: DataStore,
    hb_sock: zmq::Socket,
    rq_sock: zmq::Socket,
    work_mgr_sock: zmq::Socket,
    msg: zmq::Message,
    workers: LinkedHashMap<String, Instant>,
}

impl WorkerMgr {
    pub fn new(config: Arc<RwLock<Config>>, datastore: DataStore) -> Result<Self> {
        let hb_sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::SUB));
        let rq_sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::ROUTER));
        let work_mgr_sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::DEALER));
        try!(rq_sock.set_router_mandatory(true));
        try!(hb_sock.set_subscribe(&[]));
        try!(work_mgr_sock.set_rcvhwm(1));
        try!(work_mgr_sock.set_linger(0));
        try!(work_mgr_sock.set_immediate(true));
        let msg = try!(zmq::Message::new());
        Ok(WorkerMgr {
            config: config,
            datastore: datastore,
            hb_sock: hb_sock,
            rq_sock: rq_sock,
            work_mgr_sock: work_mgr_sock,
            msg: msg,
            workers: LinkedHashMap::new(),
        })
    }

    pub fn start(cfg: Arc<RwLock<Config>>, ds: DataStore) -> Result<JoinHandle<()>> {
        let (tx, rx) = mpsc::sync_channel(1);
        let handle = thread::Builder::new()
            .name("worker-manager".to_string())
            .spawn(move || {
                let mut manager = Self::new(cfg, ds).unwrap();
                manager.run(tx).unwrap();
            })
            .unwrap();
        match rx.recv() {
            Ok(()) => Ok(handle),
            Err(e) => panic!("worker-manager thread startup error, err={}", e),
        }
    }

    fn run(&mut self, rz: mpsc::SyncSender<()>) -> Result<()> {
        try!(self.work_mgr_sock.bind(WORKER_MGR_ADDR));
        {
            let cfg = self.config.read().unwrap();
            let worker_command = cfg.net.worker_command_addr();
            let worker_heartbeat = cfg.net.worker_heartbeat_addr();
            println!("Listening for commands on {}", worker_command);
            self.rq_sock.bind(&worker_command)?;
            println!("Listening for heartbeats on {}", worker_heartbeat);
            self.hb_sock.bind(&worker_heartbeat)?;
        }
        let mut hb_sock = false;
        let mut rq_sock = false;
        let mut work_mgr_sock = false;
        let mut process_work = false;
        rz.send(()).unwrap();

        // Reset any Dispatched jobs to Pending
        self.datastore.reset_jobs()?;

        loop {
            {
                let timeout = self.poll_timeout();
                let mut items = [
                    self.hb_sock.as_poll_item(1),
                    self.rq_sock.as_poll_item(1),
                    self.work_mgr_sock.as_poll_item(1),
                ];
                // Poll until timeout or message is received. Checking for the zmq::POLLIN flag on
                // a poll item's revents will let you know if you have received a message or not
                // on that socket.
                try!(zmq::poll(&mut items, timeout));
                if (items[0].get_revents() & zmq::POLLIN) > 0 {
                    hb_sock = true;
                }
                if (items[1].get_revents() & zmq::POLLIN) > 0 {
                    rq_sock = true;
                }
                if (items[2].get_revents() & zmq::POLLIN) > 0 {
                    work_mgr_sock = true;
                }
            }
            if hb_sock {
                process_work = try!(self.process_heartbeat());
                hb_sock = false;
            }
            self.expire_workers();
            if rq_sock {
                try!(self.process_job_status());
                rq_sock = false;
            }
            if work_mgr_sock {
                process_work = true;
                work_mgr_sock = false;
                try!(self.work_mgr_sock.recv(&mut self.msg, 0));
            }

            // Handle potential work in pending_jobs queue
            if process_work {
                try!(self.process_work());
            }
        }
    }

    fn poll_timeout(&self) -> i64 {
        if let Some((_, expiry)) = self.workers.front() {
            let timeout = *expiry - Instant::now();
            (timeout.as_secs() as i64 * 1000) + (timeout.subsec_nanos() as i64 / 1000 / 1000)
        } else {
            -1
        }
    }

    fn process_work(&mut self) -> Result<()> {
        loop {
            // Take one job from the pending list
            let mut jobs = self.datastore.pending_jobs(1)?;
            // 0 means there are no pending jobs, so we can exit
            if jobs.len() == 0 {
                debug!("process_work, no pending jobs");
                break;
            }
            // This unwrap is fine, because we just checked our length
            let mut job = jobs.pop().unwrap();

            match self.workers.pop_front() {
                Some((worker, _)) => {
                    debug!("sending work, worker={:?}, job={:?}", worker, job);
                    if self.rq_sock.send_str(&worker, zmq::SNDMORE).is_err() {
                        debug!("failed to send, worker went away, worker={:?}", worker);
                        job.set_state(jobsrv::JobState::Pending);
                        self.datastore.update_job(&job)?;
                        continue;
                    }
                    if self.rq_sock.send(&[], zmq::SNDMORE).is_err() {
                        debug!("failed to send, worker went away, worker={:?}", worker);
                        job.set_state(jobsrv::JobState::Pending);
                        self.datastore.update_job(&job)?;
                        continue;
                    }
                    if self.rq_sock
                        .send(&job.write_to_bytes().unwrap(), 0)
                        .is_err()
                    {
                        debug!("failed to send, worker went away, worker={:?}", worker);
                        job.set_state(jobsrv::JobState::Pending);
                        self.datastore.update_job(&job)?;
                        continue;
                    }
                }
                None => {
                    debug!("no workers available - bailing for now");
                    job.set_state(jobsrv::JobState::Pending);
                    self.datastore.update_job(&job)?;
                    return Ok(());
                }
            }
        }
        Ok(())
    }

    fn expire_workers(&mut self) {
        let now = Instant::now();
        loop {
            if let Some((_, expiry)) = self.workers.front() {
                if expiry >= &now {
                    break;
                }
            } else {
                break;
            }
            let worker = self.workers.pop_front();
            debug!("expiring worker due to inactivity, worker={:?}", worker);
        }
    }

    fn process_heartbeat(&mut self) -> Result<bool> {
        try!(self.hb_sock.recv(&mut self.msg, 0));
        let heartbeat: jobsrv::Heartbeat = try!(parse_from_bytes(&self.msg));
        debug!("heartbeat={:?}", heartbeat);
        let result = match heartbeat.get_state() {
            jobsrv::WorkerState::Ready => {
                let now = Instant::now();
                let expiry = now + Duration::from_millis(WORKER_TIMEOUT_MS);
                self.workers.insert(
                    heartbeat.get_endpoint().to_string(),
                    expiry,
                );
                true
            }
            jobsrv::WorkerState::Busy => {
                self.workers.remove(heartbeat.get_endpoint());
                false
            }
        };
        Ok(result)
    }

    fn process_job_status(&mut self) -> Result<()> {
        // Pop message delimiter
        try!(self.rq_sock.recv(&mut self.msg, 0));
        // Pop message body
        try!(self.rq_sock.recv(&mut self.msg, 0));
        let job: jobsrv::Job = try!(parse_from_bytes(&self.msg));
        debug!("job_status={:?}", job);
        try!(self.datastore.update_job(&job));

        Ok(())
    }
}
<commit_msg>Check expiry before doing calculations with it to avoid panicking<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// 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.

use std::sync::{mpsc, Arc, RwLock};
use std::time::{Duration, Instant};
use std::thread::{self, JoinHandle};

use linked_hash_map::LinkedHashMap;
use hab_net::server::ZMQ_CONTEXT;
use protobuf::{parse_from_bytes, Message};
use protocol::jobsrv;
use zmq;

use config::Config;
use data_store::DataStore;
use error::Result;

const WORKER_MGR_ADDR: &'static str = "inproc://work-manager";
const WORKER_TIMEOUT_MS: u64 = 33_000;

pub struct WorkerMgrClient {
    socket: zmq::Socket,
}

impl WorkerMgrClient {
    pub fn connect(&mut self) -> Result<()> {
        try!(self.socket.connect(WORKER_MGR_ADDR));
        Ok(())
    }

    pub fn notify_work(&mut self) -> Result<()> {
        try!(self.socket.send(&[1], 0));
        Ok(())
    }
}

impl Default for WorkerMgrClient {
    fn default() -> WorkerMgrClient {
        let socket = (**ZMQ_CONTEXT).as_mut().socket(zmq::DEALER).unwrap();
        socket.set_sndhwm(1).unwrap();
        socket.set_linger(0).unwrap();
        socket.set_immediate(true).unwrap();
        WorkerMgrClient { socket: socket }
    }
}

pub struct WorkerMgr {
    config: Arc<RwLock<Config>>,
    datastore: DataStore,
    hb_sock: zmq::Socket,
    rq_sock: zmq::Socket,
    work_mgr_sock: zmq::Socket,
    msg: zmq::Message,
    workers: LinkedHashMap<String, Instant>,
}

impl WorkerMgr {
    pub fn new(config: Arc<RwLock<Config>>, datastore: DataStore) -> Result<Self> {
        let hb_sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::SUB));
        let rq_sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::ROUTER));
        let work_mgr_sock = try!((**ZMQ_CONTEXT).as_mut().socket(zmq::DEALER));
        try!(rq_sock.set_router_mandatory(true));
        try!(hb_sock.set_subscribe(&[]));
        try!(work_mgr_sock.set_rcvhwm(1));
        try!(work_mgr_sock.set_linger(0));
        try!(work_mgr_sock.set_immediate(true));
        let msg = try!(zmq::Message::new());
        Ok(WorkerMgr {
            config: config,
            datastore: datastore,
            hb_sock: hb_sock,
            rq_sock: rq_sock,
            work_mgr_sock: work_mgr_sock,
            msg: msg,
            workers: LinkedHashMap::new(),
        })
    }

    pub fn start(cfg: Arc<RwLock<Config>>, ds: DataStore) -> Result<JoinHandle<()>> {
        let (tx, rx) = mpsc::sync_channel(1);
        let handle = thread::Builder::new()
            .name("worker-manager".to_string())
            .spawn(move || {
                let mut manager = Self::new(cfg, ds).unwrap();
                manager.run(tx).unwrap();
            })
            .unwrap();
        match rx.recv() {
            Ok(()) => Ok(handle),
            Err(e) => panic!("worker-manager thread startup error, err={}", e),
        }
    }

    fn run(&mut self, rz: mpsc::SyncSender<()>) -> Result<()> {
        try!(self.work_mgr_sock.bind(WORKER_MGR_ADDR));
        {
            let cfg = self.config.read().unwrap();
            let worker_command = cfg.net.worker_command_addr();
            let worker_heartbeat = cfg.net.worker_heartbeat_addr();
            println!("Listening for commands on {}", worker_command);
            self.rq_sock.bind(&worker_command)?;
            println!("Listening for heartbeats on {}", worker_heartbeat);
            self.hb_sock.bind(&worker_heartbeat)?;
        }
        let mut hb_sock = false;
        let mut rq_sock = false;
        let mut work_mgr_sock = false;
        let mut process_work = false;
        rz.send(()).unwrap();

        // Reset any Dispatched jobs to Pending
        self.datastore.reset_jobs()?;

        loop {
            {
                let timeout = self.poll_timeout();
                let mut items = [
                    self.hb_sock.as_poll_item(1),
                    self.rq_sock.as_poll_item(1),
                    self.work_mgr_sock.as_poll_item(1),
                ];
                // Poll until timeout or message is received. Checking for the zmq::POLLIN flag on
                // a poll item's revents will let you know if you have received a message or not
                // on that socket.
                try!(zmq::poll(&mut items, timeout));
                if (items[0].get_revents() & zmq::POLLIN) > 0 {
                    hb_sock = true;
                }
                if (items[1].get_revents() & zmq::POLLIN) > 0 {
                    rq_sock = true;
                }
                if (items[2].get_revents() & zmq::POLLIN) > 0 {
                    work_mgr_sock = true;
                }
            }
            if hb_sock {
                process_work = try!(self.process_heartbeat());
                hb_sock = false;
            }
            self.expire_workers();
            if rq_sock {
                try!(self.process_job_status());
                rq_sock = false;
            }
            if work_mgr_sock {
                process_work = true;
                work_mgr_sock = false;
                try!(self.work_mgr_sock.recv(&mut self.msg, 0));
            }

            // Handle potential work in pending_jobs queue
            if process_work {
                try!(self.process_work());
            }
        }
    }

    fn poll_timeout(&self) -> i64 {
        let now = Instant::now();
        let timeout;

        if let Some((_, expiry)) = self.workers.front() {
            // uh-oh. our expiration date is in the past. it's supposed to be in the
            // future. blindly subtracting now from this will panic the current
            // thread, since Instant's are supposed to monotonically increase.
            // let's just timeout immediately instead.
            if expiry < &now {
                return 0;
            } else {
                timeout = *expiry - now;
            }

            (timeout.as_secs() as i64 * 1000) + (timeout.subsec_nanos() as i64 / 1000 / 1000)
        } else {
            -1
        }
    }

    fn process_work(&mut self) -> Result<()> {
        loop {
            // Take one job from the pending list
            let mut jobs = self.datastore.pending_jobs(1)?;
            // 0 means there are no pending jobs, so we can exit
            if jobs.len() == 0 {
                debug!("process_work, no pending jobs");
                break;
            }
            // This unwrap is fine, because we just checked our length
            let mut job = jobs.pop().unwrap();

            match self.workers.pop_front() {
                Some((worker, _)) => {
                    debug!("sending work, worker={:?}, job={:?}", worker, job);
                    if self.rq_sock.send_str(&worker, zmq::SNDMORE).is_err() {
                        debug!("failed to send, worker went away, worker={:?}", worker);
                        job.set_state(jobsrv::JobState::Pending);
                        self.datastore.update_job(&job)?;
                        continue;
                    }
                    if self.rq_sock.send(&[], zmq::SNDMORE).is_err() {
                        debug!("failed to send, worker went away, worker={:?}", worker);
                        job.set_state(jobsrv::JobState::Pending);
                        self.datastore.update_job(&job)?;
                        continue;
                    }
                    if self.rq_sock
                        .send(&job.write_to_bytes().unwrap(), 0)
                        .is_err()
                    {
                        debug!("failed to send, worker went away, worker={:?}", worker);
                        job.set_state(jobsrv::JobState::Pending);
                        self.datastore.update_job(&job)?;
                        continue;
                    }
                }
                None => {
                    debug!("no workers available - bailing for now");
                    job.set_state(jobsrv::JobState::Pending);
                    self.datastore.update_job(&job)?;
                    return Ok(());
                }
            }
        }
        Ok(())
    }

    fn expire_workers(&mut self) {
        let now = Instant::now();
        loop {
            if let Some((_, expiry)) = self.workers.front() {
                if expiry >= &now {
                    break;
                }
            } else {
                break;
            }
            let worker = self.workers.pop_front();
            debug!("expiring worker due to inactivity, worker={:?}", worker);
        }
    }

    fn process_heartbeat(&mut self) -> Result<bool> {
        try!(self.hb_sock.recv(&mut self.msg, 0));
        let heartbeat: jobsrv::Heartbeat = try!(parse_from_bytes(&self.msg));
        debug!("heartbeat={:?}", heartbeat);
        let result = match heartbeat.get_state() {
            jobsrv::WorkerState::Ready => {
                let now = Instant::now();
                let expiry = now + Duration::from_millis(WORKER_TIMEOUT_MS);
                self.workers.insert(
                    heartbeat.get_endpoint().to_string(),
                    expiry,
                );
                true
            }
            jobsrv::WorkerState::Busy => {
                self.workers.remove(heartbeat.get_endpoint());
                false
            }
        };
        Ok(result)
    }

    fn process_job_status(&mut self) -> Result<()> {
        // Pop message delimiter
        try!(self.rq_sock.recv(&mut self.msg, 0));
        // Pop message body
        try!(self.rq_sock.recv(&mut self.msg, 0));
        let job: jobsrv::Job = try!(parse_from_bytes(&self.msg));
        debug!("job_status={:?}", job);
        try!(self.datastore.update_job(&job));

        Ok(())
    }
}
<|endoftext|>


<commit_before>// run-pass

fn main () {
    let n = (1, (2, 3)).1.1;
    assert_eq!(n, 3);
}
<commit_msg>Test even deeper nested indexing<commit_after>// run-pass

fn main () {
    let n = (1, (2, 3)).1.1;
    assert_eq!(n, 3);

    let n = (1, (2, (3, 4))).1.1.1;
    assert_eq!(n, 4);
}
<|endoftext|>