<commit_before>use common::*;
use gl::types::*;
use id_allocator::IdAllocator;
use nalgebra::{normalize, cross};
use nalgebra::{Pnt3, Vec3};
use ncollide::bounding_volume::{AABB, AABB3};
use noise::source::Perlin;
use noise::model::Plane;
use state::EntityId;
use std::collections::hash_map::HashMap;
use std::num::Float;
use stopwatch::TimerSet;
use terrain::{AMPLITUDE, TerrainType};

pub const BLOCK_WIDTH: i32 = 4;
// Number of samples in a single dimension per block.
pub const SAMPLES_PER_BLOCK: uint = 16;
pub const SAMPLE_WIDTH: f32 = BLOCK_WIDTH as f32 / SAMPLES_PER_BLOCK as f32;

#[deriving(Show, PartialEq, Eq, Hash, Copy, Clone)]
pub struct BlockPosition(Pnt3<i32>);

impl BlockPosition {
  #[inline(always)]
  pub fn new(x: i32, y: i32, z: i32) -> BlockPosition {
    BlockPosition(Pnt3::new(x, y, z))
  }

  #[inline(always)]
  pub fn as_pnt<'a>(&'a self) -> &'a Pnt3<i32> {
    let BlockPosition(ref pnt) = *self;
    pnt
  }

  #[inline(always)]
  pub fn as_mut_pnt3<'a>(&'a mut self) -> &'a mut Pnt3<i32> {
    let BlockPosition(ref mut pnt) = *self;
    pnt
  }

  pub fn from_world_position(world_position: &Pnt3<f32>) -> BlockPosition {
    macro_rules! convert_coordinate(
      ($x: expr) => ({
        let x = $x.floor() as i32;
        let x =
          if x < 0 {
            x - (BLOCK_WIDTH - 1)
          } else {
            x
          };
        x / BLOCK_WIDTH
      });
    );
    BlockPosition(
      Pnt3::new(
        convert_coordinate!(world_position.x),
        convert_coordinate!(world_position.y),
        convert_coordinate!(world_position.z),
      )
    )
  }

  pub fn to_world_position(&self) -> Pnt3<f32> {
    Pnt3::new(
      (self.as_pnt().x * BLOCK_WIDTH) as f32,
      (self.as_pnt().y * BLOCK_WIDTH) as f32,
      (self.as_pnt().z * BLOCK_WIDTH) as f32,
    )
  }
}

impl Add<Vec3<i32>, BlockPosition> for BlockPosition {
  fn add(mut self, rhs: Vec3<i32>) -> Self {
    self.as_mut_pnt3().x += rhs.x;
    self.as_mut_pnt3().y += rhs.y;
    self.as_mut_pnt3().z += rhs.z;
    self
  }
}

pub struct TerrainBlock {
  // These Vecs must all be ordered the same way.

  // vertex coordinates flattened into separate GLfloats (x, y, z order)
  pub vertex_coordinates: Vec<GLfloat>,
  // per-vertex normal vectors flattened into separate GLfloats (x, y, z order)
  pub normals: Vec<GLfloat>,
  // per-triangle terrain types
  pub typs: Vec<GLuint>,

  // per-triangle entity IDs
  pub ids: Vec<EntityId>,
  // per-triangle bounding boxes
  pub bounds: HashMap<EntityId, AABB3<GLfloat>>,
}

impl TerrainBlock {
  pub fn empty() -> TerrainBlock {
    TerrainBlock {
      vertex_coordinates: Vec::new(),
      normals: Vec::new(),
      typs: Vec::new(),
      ids: Vec::new(),
      bounds: HashMap::new(),
    }
  }

  pub fn generate(
    timers: &TimerSet,
    id_allocator: &mut IdAllocator<EntityId>,
    heightmap: &Perlin,
    position: &BlockPosition,
  ) -> TerrainBlock {
    timers.time("update.generate_block", || {
      let mut block = TerrainBlock::empty();
      let mut face_normals = [[None, ..SAMPLES_PER_BLOCK], ..SAMPLES_PER_BLOCK];

      let heightmap = Plane::new(heightmap);
      let x = (position.as_pnt().x * BLOCK_WIDTH) as f32;
      let y = (position.as_pnt().y * BLOCK_WIDTH) as f32;
      let z = (position.as_pnt().z * BLOCK_WIDTH) as f32;
      for dx in range(0, SAMPLES_PER_BLOCK) {
        let x = x + dx as f32 * SAMPLE_WIDTH;
        for dz in range(0, SAMPLES_PER_BLOCK) {
          let z = z + dz as f32 * SAMPLE_WIDTH;
          let position = Pnt3::new(x, y, z);
          TerrainBlock::add_square(
            timers,
            &heightmap,
            id_allocator,
            &mut block,
            &mut face_normals[dx][dz],
            &position,
          );
        }
      }

      if USE_LIGHTING {
        // This is all super mesh-specific.
        for x_square in range(0, SAMPLES_PER_BLOCK) {
          for z_square in range(0, SAMPLES_PER_BLOCK) {
            match face_normals[x_square][z_square] {
              None => {},
              Some(square_normals) => {
                let triangle_faces = [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3];

                for &triangle_face in triangle_faces.iter() {
                  block.normals.push_all(&[
                    square_normals[triangle_face].x,
                    square_normals[triangle_face].y,
                    square_normals[triangle_face].z,
                  ]);
                }
              },
            }
          }
        }
      }

      block
    })
  }

  #[inline]
  fn add_square<'a>(
    timers: &TimerSet,
    heightmap: &Plane<'a, Perlin>,
    id_allocator: &mut IdAllocator<EntityId>,
    block: &mut TerrainBlock,
    face_normals: &mut Option<[Vec3<f32>, ..4]>,
    position: &Pnt3<f32>,
  ) {
    macro_rules! at(
      ($x: expr, $z: expr) => ({
        let y = AMPLITUDE * (heightmap.get::<GLfloat>($x, $z) + 1.0) / 2.0;
        Pnt3::new($x, y, $z)
      });
    );

    let half_width = SAMPLE_WIDTH / 2.0;
    let center = at!(position.x + half_width, position.z + half_width);

    if position.y < center.y && center.y <= position.y + BLOCK_WIDTH as f32 {
      timers.time("update.generate_block.add_square", || {
        let mut new_face_normals = [Vec3::new(0.0, 0.0, 0.0), ..4];

        let x2 = position.x + SAMPLE_WIDTH;
        let z2 = position.z + SAMPLE_WIDTH;
        let v1 = at!(position.x, position.z);
        let v2 = at!(position.x, z2);
        let v3 = at!(x2, z2);
        let v4 = at!(x2, position.z);
        let mut center_lower_than: int = 0;
        for v in [v1, v2, v3, v4].iter() {
          if center.y < v.y {
            center_lower_than += 1;
          }
        }
        let terrain_type =
          if center_lower_than >= 3 {
            TerrainType::Dirt
          } else {
            TerrainType::Grass
          };

        macro_rules! place_terrain(
          ($normal_idx: expr, $v1: expr, $v2: expr, $minx: expr, $minz: expr, $maxx: expr, $maxz: expr) => ({
            let mut maxy = $v1.y;
            if $v2.y > $v1.y {
              maxy = $v2.y;
            }
            if center.y > maxy {
              maxy = center.y;
            }

            if USE_LIGHTING {
              let side1: Vec3<GLfloat> = center - *$v1;
              let side2: Vec3<GLfloat> = $v2.to_vec() - $v1.to_vec();
              let normal: Vec3<GLfloat> = normalize(&cross(&side1, &side2));
              new_face_normals[$normal_idx] = normal;
            }

            let id = id_allocator.allocate();
            block.vertex_coordinates.push_all(&[
              $v1.x, $v1.y, $v1.z,
              $v2.x, $v2.y, $v2.z,
              center.x, center.y, center.z,
            ]);
            block.typs.push(terrain_type as GLuint);
            block.ids.push(id);
            block.bounds.insert(
              id,
              AABB::new(
                Pnt3::new($minx, $v1.y, $minz),
                Pnt3::new($maxx, maxy, $maxz),
              ),
            );
          });
        );

        place_terrain!(0, &v1, &v2, v1.x, v1.z, center.x, v2.z);
        place_terrain!(1, &v2, &v3, v2.x, center.z, v3.x, v3.z);
        place_terrain!(2, &v3, &v4, center.x, center.z, v3.x, v3.z);
        place_terrain!(3, &v4, &v1, v1.x, v1.z, v4.x, center.z);

        *face_normals = Some(new_face_normals);
      })
    }
  }
}

<commit_msg>per-tile lighting smoothing<commit_after>use common::*;
use gl::types::*;
use id_allocator::IdAllocator;
use nalgebra::{normalize, cross};
use nalgebra::{Pnt3, Vec3};
use ncollide::bounding_volume::{AABB, AABB3};
use noise::source::Perlin;
use noise::model::Plane;
use state::EntityId;
use std::collections::hash_map::HashMap;
use std::num::Float;
use stopwatch::TimerSet;
use terrain::{AMPLITUDE, TerrainType};

pub const BLOCK_WIDTH: i32 = 4;
// Number of samples in a single dimension per block.
pub const SAMPLES_PER_BLOCK: uint = 16;
pub const SAMPLE_WIDTH: f32 = BLOCK_WIDTH as f32 / SAMPLES_PER_BLOCK as f32;

#[deriving(Show, PartialEq, Eq, Hash, Copy, Clone)]
pub struct BlockPosition(Pnt3<i32>);

impl BlockPosition {
  #[inline(always)]
  pub fn new(x: i32, y: i32, z: i32) -> BlockPosition {
    BlockPosition(Pnt3::new(x, y, z))
  }

  #[inline(always)]
  pub fn as_pnt<'a>(&'a self) -> &'a Pnt3<i32> {
    let BlockPosition(ref pnt) = *self;
    pnt
  }

  #[inline(always)]
  pub fn as_mut_pnt3<'a>(&'a mut self) -> &'a mut Pnt3<i32> {
    let BlockPosition(ref mut pnt) = *self;
    pnt
  }

  pub fn from_world_position(world_position: &Pnt3<f32>) -> BlockPosition {
    macro_rules! convert_coordinate(
      ($x: expr) => ({
        let x = $x.floor() as i32;
        let x =
          if x < 0 {
            x - (BLOCK_WIDTH - 1)
          } else {
            x
          };
        x / BLOCK_WIDTH
      });
    );
    BlockPosition(
      Pnt3::new(
        convert_coordinate!(world_position.x),
        convert_coordinate!(world_position.y),
        convert_coordinate!(world_position.z),
      )
    )
  }

  pub fn to_world_position(&self) -> Pnt3<f32> {
    Pnt3::new(
      (self.as_pnt().x * BLOCK_WIDTH) as f32,
      (self.as_pnt().y * BLOCK_WIDTH) as f32,
      (self.as_pnt().z * BLOCK_WIDTH) as f32,
    )
  }
}

impl Add<Vec3<i32>, BlockPosition> for BlockPosition {
  fn add(mut self, rhs: Vec3<i32>) -> Self {
    self.as_mut_pnt3().x += rhs.x;
    self.as_mut_pnt3().y += rhs.y;
    self.as_mut_pnt3().z += rhs.z;
    self
  }
}

pub struct TerrainBlock {
  // These Vecs must all be ordered the same way.

  // vertex coordinates flattened into separate GLfloats (x, y, z order)
  pub vertex_coordinates: Vec<GLfloat>,
  // per-vertex normal vectors flattened into separate GLfloats (x, y, z order)
  pub normals: Vec<GLfloat>,
  // per-triangle terrain types
  pub typs: Vec<GLuint>,

  // per-triangle entity IDs
  pub ids: Vec<EntityId>,
  // per-triangle bounding boxes
  pub bounds: HashMap<EntityId, AABB3<GLfloat>>,
}

impl TerrainBlock {
  pub fn empty() -> TerrainBlock {
    TerrainBlock {
      vertex_coordinates: Vec::new(),
      normals: Vec::new(),
      typs: Vec::new(),
      ids: Vec::new(),
      bounds: HashMap::new(),
    }
  }

  pub fn generate(
    timers: &TimerSet,
    id_allocator: &mut IdAllocator<EntityId>,
    heightmap: &Perlin,
    position: &BlockPosition,
  ) -> TerrainBlock {
    timers.time("update.generate_block", || {
      let mut block = TerrainBlock::empty();
      let mut face_normals = [[None, ..SAMPLES_PER_BLOCK], ..SAMPLES_PER_BLOCK];

      let heightmap = Plane::new(heightmap);
      let x = (position.as_pnt().x * BLOCK_WIDTH) as f32;
      let y = (position.as_pnt().y * BLOCK_WIDTH) as f32;
      let z = (position.as_pnt().z * BLOCK_WIDTH) as f32;
      for dx in range(0, SAMPLES_PER_BLOCK) {
        let x = x + dx as f32 * SAMPLE_WIDTH;
        for dz in range(0, SAMPLES_PER_BLOCK) {
          let z = z + dz as f32 * SAMPLE_WIDTH;
          let position = Pnt3::new(x, y, z);
          TerrainBlock::add_square(
            timers,
            &heightmap,
            id_allocator,
            &mut block,
            &mut face_normals[dx][dz],
            &position,
          );
        }
      }

      if USE_LIGHTING {
        // This is all super mesh-specific.
        for x_square in range(0, SAMPLES_PER_BLOCK) {
          for z_square in range(0, SAMPLES_PER_BLOCK) {
            match face_normals[x_square][z_square] {
              None => {},
              Some(square_normals) => {
                let mut mesh_vertex_normals = [Vec3::new(0.0, 0.0, 0.0), ..5];

                for n in square_normals.iter() {
                  mesh_vertex_normals[0] = mesh_vertex_normals[0] + *n;
                }
                mesh_vertex_normals[0] = mesh_vertex_normals[0] / 4.0;

                {
                  let mut square_normal = 3;

                  for mesh_vertex in range(1, 5 as uint) {
                    mesh_vertex_normals[mesh_vertex] = mesh_vertex_normals[mesh_vertex] + square_normals[square_normal];

                    // square_normal == 4 when mesh_vertex == 1.
                    if mesh_vertex == 1 {
                      square_normal = 0;
                    } else {
                      square_normal += 1;
                    }

                    mesh_vertex_normals[mesh_vertex] = mesh_vertex_normals[mesh_vertex] + square_normals[square_normal];
                    mesh_vertex_normals[mesh_vertex] = mesh_vertex_normals[mesh_vertex] / 2.0;
                  }
                }

                const N_INDICES: uint = 12;
                assert_eq!(N_INDICES, VERTICES_PER_TRIANGLE * square_normals.len());

                {
                  let mesh_vertex_indices: [uint, ..N_INDICES] = [1,2,0,2,3,0,3,4,0,4,1,0];
                  for &mesh_vertex_index in mesh_vertex_indices.iter() {
                    block.normals.push(mesh_vertex_normals[mesh_vertex_index].x);
                    block.normals.push(mesh_vertex_normals[mesh_vertex_index].y);
                    block.normals.push(mesh_vertex_normals[mesh_vertex_index].z);
                  }
                }

              },
            }
          }
        }
      }

      block
    })
  }

  #[inline]
  fn add_square<'a>(
    timers: &TimerSet,
    heightmap: &Plane<'a, Perlin>,
    id_allocator: &mut IdAllocator<EntityId>,
    block: &mut TerrainBlock,
    face_normals: &mut Option<[Vec3<f32>, ..4]>,
    position: &Pnt3<f32>,
  ) {
    macro_rules! at(
      ($x: expr, $z: expr) => ({
        let y = AMPLITUDE * (heightmap.get::<GLfloat>($x, $z) + 1.0) / 2.0;
        Pnt3::new($x, y, $z)
      });
    );

    let half_width = SAMPLE_WIDTH / 2.0;
    let center = at!(position.x + half_width, position.z + half_width);

    if position.y < center.y && center.y <= position.y + BLOCK_WIDTH as f32 {
      timers.time("update.generate_block.add_square", || {
        let mut new_face_normals = [Vec3::new(0.0, 0.0, 0.0), ..4];

        let x2 = position.x + SAMPLE_WIDTH;
        let z2 = position.z + SAMPLE_WIDTH;
        let v1 = at!(position.x, position.z);
        let v2 = at!(position.x, z2);
        let v3 = at!(x2, z2);
        let v4 = at!(x2, position.z);
        let mut center_lower_than: int = 0;
        for v in [v1, v2, v3, v4].iter() {
          if center.y < v.y {
            center_lower_than += 1;
          }
        }
        let terrain_type =
          if center_lower_than >= 3 {
            TerrainType::Dirt
          } else {
            TerrainType::Grass
          };

        macro_rules! place_terrain(
          ($normal_idx: expr, $v1: expr, $v2: expr, $minx: expr, $minz: expr, $maxx: expr, $maxz: expr) => ({
            let mut maxy = $v1.y;
            if $v2.y > $v1.y {
              maxy = $v2.y;
            }
            if center.y > maxy {
              maxy = center.y;
            }

            if USE_LIGHTING {
              let side1: Vec3<GLfloat> = center - *$v1;
              let side2: Vec3<GLfloat> = $v2.to_vec() - $v1.to_vec();
              let normal: Vec3<GLfloat> = normalize(&cross(&side1, &side2));
              new_face_normals[$normal_idx] = normal;
            }

            let id = id_allocator.allocate();
            block.vertex_coordinates.push_all(&[
              $v1.x, $v1.y, $v1.z,
              $v2.x, $v2.y, $v2.z,
              center.x, center.y, center.z,
            ]);
            block.typs.push(terrain_type as GLuint);
            block.ids.push(id);
            block.bounds.insert(
              id,
              AABB::new(
                Pnt3::new($minx, $v1.y, $minz),
                Pnt3::new($maxx, maxy, $maxz),
              ),
            );
          });
        );

        place_terrain!(0, &v1, &v2, v1.x, v1.z, center.x, v2.z);
        place_terrain!(1, &v2, &v3, v2.x, center.z, v3.x, v3.z);
        place_terrain!(2, &v3, &v4, center.x, center.z, v3.x, v3.z);
        place_terrain!(3, &v4, &v1, v1.x, v1.z, v4.x, center.z);

        *face_normals = Some(new_face_normals);
      })
    }
  }
}

<|endoftext|>


<commit_before>use crate::lib::*;

#[cfg(feature = "std")]
use parking_lot::Mutex;

#[cfg(not(feature = "std"))]
use spin::Mutex;

#[derive(Clone)]
/// Wraps the atomic operations on a Decider's state in a threadsafe
/// fashion.
pub(crate) struct ThreadsafeWrapper<T>
where
    T: fmt::Debug + Default + Clone + PartialEq + Eq,
{
    data: Arc<Mutex<T>>,
}

impl<T> Default for ThreadsafeWrapper<T>
where
    T: fmt::Debug + Default + Clone + PartialEq + Eq,
{
    fn default() -> Self {
        ThreadsafeWrapper {
            data: Arc::new(Mutex::new(T::default())),
        }
    }
}

impl<T> PartialEq<Self> for ThreadsafeWrapper<T>
where
    T: fmt::Debug + Default + Clone + PartialEq + Eq,
{
    fn eq(&self, other: &Self) -> bool {
        if self as *const _ == other as *const _ {
            return true;
        }
        let mine = self.data.lock();
        let other = other.data.lock();
        *other == *mine
    }
}

impl<T> Eq for ThreadsafeWrapper<T> where T: fmt::Debug + Default + Clone + PartialEq + Eq {}

impl<T> fmt::Debug for ThreadsafeWrapper<T>
where
    T: fmt::Debug + Default + Clone + PartialEq + Eq,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        let data = self.data.lock();
        data.fmt(f)
    }
}

#[cfg(feature = "std")]
mod std {
    use super::*;
    use evmap::ShallowCopy;

    impl<T> ShallowCopy for ThreadsafeWrapper<T>
    where
        T: fmt::Debug + Default + Clone + PartialEq + Eq,
    {
        unsafe fn shallow_copy(&mut self) -> Self {
            ThreadsafeWrapper {
                data: self.data.shallow_copy(),
            }
        }
    }
}

impl<T> ThreadsafeWrapper<T>
where
    T: fmt::Debug + Default + Clone + PartialEq + Eq,
{
    #[inline]
    /// Wraps retrieving a bucket's data, calls a function to make a
    /// decision and return a new state, and then tries to set the
    /// state on the bucket.
    ///
    /// This function can loop and call the decision closure again if
    /// the bucket state couldn't be set.
    ///
    /// # Panics
    /// Panics if an error occurs in acquiring any locks.
    pub(crate) fn measure_and_replace<F, E>(&self, f: F) -> Result<(), E>
    where
        F: Fn(&T) -> (Result<(), E>, Option<T>),
    {
        let mut data = self.data.lock();
        let (decision, new_data) = f(&*data);
        if let Some(new_data) = new_data {
            *data = new_data;
        }
        decision
    }

    /// Retrieves and returns a snapshot of the bucket state. This
    /// isn't thread safe, but can be used to restore an old copy of
    /// the bucket if necessary.
    ///
    /// # Thread safety
    /// This function operates threadsafely, but you're literally
    /// taking a copy of data that will change. Relying on the data
    /// that is returned *will* race.
    pub(crate) fn snapshot(&self) -> T {
        let data = self.data.lock();
        data.clone()
    }
}
<commit_msg>Add a test for equality of identical ThreadsafeWrappers<commit_after>use crate::lib::*;

#[cfg(feature = "std")]
use parking_lot::Mutex;

#[cfg(not(feature = "std"))]
use spin::Mutex;

#[derive(Clone)]
/// Wraps the atomic operations on a Decider's state in a threadsafe
/// fashion.
pub(crate) struct ThreadsafeWrapper<T>
where
    T: fmt::Debug + Default + Clone + PartialEq + Eq,
{
    data: Arc<Mutex<T>>,
}

impl<T> Default for ThreadsafeWrapper<T>
where
    T: fmt::Debug + Default + Clone + PartialEq + Eq,
{
    fn default() -> Self {
        ThreadsafeWrapper {
            data: Arc::new(Mutex::new(T::default())),
        }
    }
}

impl<T> PartialEq<Self> for ThreadsafeWrapper<T>
where
    T: fmt::Debug + Default + Clone + PartialEq + Eq,
{
    fn eq(&self, other: &Self) -> bool {
        if self as *const _ == other as *const _ {
            return true;
        }
        let mine = self.data.lock();
        let other = other.data.lock();
        *other == *mine
    }
}

impl<T> Eq for ThreadsafeWrapper<T> where T: fmt::Debug + Default + Clone + PartialEq + Eq {}

impl<T> fmt::Debug for ThreadsafeWrapper<T>
where
    T: fmt::Debug + Default + Clone + PartialEq + Eq,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        let data = self.data.lock();
        data.fmt(f)
    }
}

#[cfg(feature = "std")]
mod std {
    use super::*;
    use evmap::ShallowCopy;

    impl<T> ShallowCopy for ThreadsafeWrapper<T>
    where
        T: fmt::Debug + Default + Clone + PartialEq + Eq,
    {
        unsafe fn shallow_copy(&mut self) -> Self {
            ThreadsafeWrapper {
                data: self.data.shallow_copy(),
            }
        }
    }
}

impl<T> ThreadsafeWrapper<T>
where
    T: fmt::Debug + Default + Clone + PartialEq + Eq,
{
    #[inline]
    /// Wraps retrieving a bucket's data, calls a function to make a
    /// decision and return a new state, and then tries to set the
    /// state on the bucket.
    ///
    /// This function can loop and call the decision closure again if
    /// the bucket state couldn't be set.
    ///
    /// # Panics
    /// Panics if an error occurs in acquiring any locks.
    pub(crate) fn measure_and_replace<F, E>(&self, f: F) -> Result<(), E>
    where
        F: Fn(&T) -> (Result<(), E>, Option<T>),
    {
        let mut data = self.data.lock();
        let (decision, new_data) = f(&*data);
        if let Some(new_data) = new_data {
            *data = new_data;
        }
        decision
    }

    /// Retrieves and returns a snapshot of the bucket state. This
    /// isn't thread safe, but can be used to restore an old copy of
    /// the bucket if necessary.
    ///
    /// # Thread safety
    /// This function operates threadsafely, but you're literally
    /// taking a copy of data that will change. Relying on the data
    /// that is returned *will* race.
    pub(crate) fn snapshot(&self) -> T {
        let data = self.data.lock();
        data.clone()
    }
}

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

    #[test]
    fn no_deadlock_on_eq() {
        let wrapper: ThreadsafeWrapper<u32> = ThreadsafeWrapper::default();
        assert_eq!(wrapper, wrapper);
    }
}
<|endoftext|>


<commit_before>#[cfg(feature = "chrono")]
mod date_and_time;

use byteorder::WriteBytesExt;
use mysql::{Mysql, MysqlType};
use std::error::Error as StdError;
use std::io::Write;
use types::{ToSql, ToSqlOutput, IsNull, FromSql, HasSqlType};

primitive_impls!(Tinyint -> (i8, mysql: (Tiny)));

impl ToSql<::types::Tinyint, Mysql> for i8 {
    fn to_sql<W: Write>(&self, out: &mut W) -> Result<IsNull, Box<StdError+Send+Sync>> {
        out.write_i8(*self)
            .map(|_| IsNull::No)
            .map_err(|e| Box::new(e) as Box<StdError+Send+Sync>)
    }
}

impl FromSql<::types::Tinyint, Mysql> for i8 {
    fn from_sql(bytes: Option<&[u8]>) -> Result<Self, Box<StdError+Send+Sync>> {
        let bytes = not_none!(bytes);
        Ok(bytes[0] as i8)
    }
}

impl ToSql<::types::Bool, Mysql> for bool {
    fn to_sql<W: Write>(&self, out: &mut ToSqlOutput<W, Mysql>) -> Result<IsNull, Box<StdError+Send+Sync>> {
        let int_value = if *self {
            1
        } else {
            0
        };
        <i32 as ToSql<::types::Integer, Mysql>>::to_sql(&int_value, out)
    }
}

impl FromSql<::types::Bool, Mysql> for bool {
    fn from_sql(bytes: Option<&[u8]>) -> Result<Self, Box<StdError+Send+Sync>> {
        Ok(not_none!(bytes).iter().any(|x| *x != 0))
    }
}

impl HasSqlType<::types::Date> for Mysql {
    fn metadata(_: &()) -> MysqlType {
        MysqlType::Date
    }
}

impl HasSqlType<::types::Time> for Mysql {
    fn metadata(_: &()) -> MysqlType {
        MysqlType::Time
    }
}

impl HasSqlType<::types::Timestamp> for Mysql {
    fn metadata(_: &()) -> MysqlType {
        MysqlType::Timestamp
    }
}
<commit_msg>Fix failing build<commit_after>#[cfg(feature = "chrono")]
mod date_and_time;

use byteorder::WriteBytesExt;
use mysql::{Mysql, MysqlType};
use std::error::Error as StdError;
use std::io::Write;
use types::{ToSql, ToSqlOutput, IsNull, FromSql, HasSqlType};

primitive_impls!(Tinyint -> (i8, mysql: (Tiny)));

impl ToSql<::types::Tinyint, Mysql> for i8 {
    fn to_sql<W: Write>(&self, out: &mut ToSqlOutput<W, Mysql>) -> Result<IsNull, Box<StdError+Send+Sync>> {
        out.write_i8(*self)
            .map(|_| IsNull::No)
            .map_err(|e| Box::new(e) as Box<StdError+Send+Sync>)
    }
}

impl FromSql<::types::Tinyint, Mysql> for i8 {
    fn from_sql(bytes: Option<&[u8]>) -> Result<Self, Box<StdError+Send+Sync>> {
        let bytes = not_none!(bytes);
        Ok(bytes[0] as i8)
    }
}

impl ToSql<::types::Bool, Mysql> for bool {
    fn to_sql<W: Write>(&self, out: &mut ToSqlOutput<W, Mysql>) -> Result<IsNull, Box<StdError+Send+Sync>> {
        let int_value = if *self {
            1
        } else {
            0
        };
        <i32 as ToSql<::types::Integer, Mysql>>::to_sql(&int_value, out)
    }
}

impl FromSql<::types::Bool, Mysql> for bool {
    fn from_sql(bytes: Option<&[u8]>) -> Result<Self, Box<StdError+Send+Sync>> {
        Ok(not_none!(bytes).iter().any(|x| *x != 0))
    }
}

impl HasSqlType<::types::Date> for Mysql {
    fn metadata(_: &()) -> MysqlType {
        MysqlType::Date
    }
}

impl HasSqlType<::types::Time> for Mysql {
    fn metadata(_: &()) -> MysqlType {
        MysqlType::Time
    }
}

impl HasSqlType<::types::Timestamp> for Mysql {
    fn metadata(_: &()) -> MysqlType {
        MysqlType::Timestamp
    }
}
<|endoftext|>


<commit_before>use dicom_core::dicom_value;
use dicom_core::{DataElement, PrimitiveValue, VR};
use dicom_dictionary_std::tags;
use dicom_object::{mem::InMemDicomObject, StandardDataDictionary};
use dicom_ul::pdu;
use dicom_ul::{
    association::client::ClientAssociationOptions,
    pdu::{PDataValueType, Pdu},
};
use pdu::PDataValue;
use snafu::{prelude::*, Whatever};
use structopt::StructOpt;
use tracing::{warn, Level};

/// DICOM C-ECHO SCU
#[derive(Debug, StructOpt)]
struct App {
    /// socket address to SCP,
    /// optionally with AE title
    /// (example: "QUERY-SCP@127.0.0.1:1045")
    addr: String,
    /// verbose mode
    #[structopt(short = "v", long = "verbose")]
    verbose: bool,
    /// the C-ECHO message ID
    #[structopt(short = "m", long = "message-id", default_value = "1")]
    message_id: u16,
    /// the calling AE title
    #[structopt(long = "calling-ae-title", default_value = "ECHOSCU")]
    calling_ae_title: String,
    /// the called Application Entity title,
    /// overrides AE title in address if present [default: ANY-SCP]
    #[structopt(long = "called-ae-title")]
    called_ae_title: Option<String>,
}

fn main() {
    run().unwrap_or_else(|e| {
        tracing::error!("{}", snafu::Report::from_error(e));
        std::process::exit(-2);
    })
}

fn run() -> Result<(), Whatever> {
    let App {
        addr,
        verbose,
        message_id,
        called_ae_title,
        calling_ae_title,
    } = App::from_args();

    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_max_level(if verbose { Level::DEBUG } else { Level::INFO })
            .finish(),
    )
    .whatever_context("Could not set up global logging subscriber")
    .unwrap_or_else(|e: Whatever| {
        eprintln!("[ERROR] {}", snafu::Report::from_error(e));
    });

    let mut association_opt = ClientAssociationOptions::new()
        .with_abstract_syntax("1.2.840.10008.1.1")
        .calling_ae_title(calling_ae_title);
    if let Some(called_ae_title) = called_ae_title {
        association_opt = association_opt.called_ae_title(called_ae_title);
    }
    let mut association = association_opt
        .establish_with(&addr)
        .whatever_context("Could not establish association with SCP")?;

    let pc = association
        .presentation_contexts()
        .first()
        .whatever_context("No presentation context accepted")?
        .clone();

    if verbose {
        println!("Association with {} successful", addr);
    }

    // commands are always in implict VR LE
    let ts = dicom_transfer_syntax_registry::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased();

    let obj = create_echo_command(message_id);

    let mut data = Vec::new();

    obj.write_dataset_with_ts(&mut data, &ts)
        .whatever_context("Failed to construct C-ECHO request")?;

    association
        .send(&Pdu::PData {
            data: vec![PDataValue {
                presentation_context_id: pc.id,
                value_type: PDataValueType::Command,
                is_last: true,
                data,
            }],
        })
        .whatever_context("Failed to send C-ECHO request")?;

    if verbose {
        println!(
            "Echo message sent (msg id {}), awaiting reply...",
            message_id
        );
    }

    let pdu = association
        .receive()
        .whatever_context("Could not receive response from SCP")?;

    match pdu {
        Pdu::PData { data } => {
            let data_value = &data[0];
            let v = &data_value.data;

            let obj = InMemDicomObject::read_dataset_with_ts(v.as_slice(), &ts)
                .whatever_context("Failed to read response dataset from SCP")?;
            if verbose {
                println!("{:?}", obj);
            }

            // check status
            let status_elem = obj
                .element(tags::STATUS)
                .whatever_context("Missing Status code in response")?;

            let status = status_elem
                .to_int::<u16>()
                .whatever_context("Status code in response is not a valid integer")?;
            if verbose {
                println!("Status: {:04X}H", status);
            }
            match status {
                // Success
                0 => {
                    if verbose {
                        println!("C-ECHO successful");
                    }
                }
                // Warning
                1 | 0x0107 | 0x0116 | 0xB000..=0xBFFF => {
                    warn!("Possible issue in C-ECHO (status code {:04X}H)", status);
                }
                0xFF00 | 0xFF01 => {
                    warn!(
                        "Possible issue in C-ECHO: status is pending (status code {:04X}H)",
                        status
                    );
                }
                0xFE00 => {
                    warn!("Operation cancelled");
                }
                _ => {
                    eprintln!("C-ECHO failed (status code {:04X}H)", status);
                }
            }

            // msg ID response, should be equal to sent msg ID
            let msg_id_elem = obj
                .element(tags::MESSAGE_ID_BEING_RESPONDED_TO)
                .whatever_context("Could not retrieve Message ID from response")?;

            if message_id
                != msg_id_elem
                    .to_int()
                    .whatever_context("Message ID is not a valid integer")?
            {
                whatever!("Message ID mismatch");
            }
        }
        pdu => whatever!("Unexpected PDU {:?}", pdu),
    }

    Ok(())
}

fn create_echo_command(message_id: u16) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::new_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UI,
        PrimitiveValue::from(8 + 18 + 8 + 2 + 8 + 2 + 8 + 2),
    ));

    // service
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        dicom_value!(Str, "1.2.840.10008.1.1\0"),
    ));
    // command
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        dicom_value!(U16, [0x0030]),
    ));
    // message ID
    obj.put(DataElement::new(
        tags::MESSAGE_ID,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));
    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0101]),
    ));

    obj
}
<commit_msg>[echoscu] Format imports<commit_after>use dicom_core::{dicom_value, DataElement, PrimitiveValue, VR};
use dicom_dictionary_std::tags;
use dicom_object::{mem::InMemDicomObject, StandardDataDictionary};
use dicom_ul::{
    association::client::ClientAssociationOptions,
    pdu::{self, PDataValueType, Pdu},
};
use pdu::PDataValue;
use snafu::{prelude::*, Whatever};
use structopt::StructOpt;
use tracing::{warn, Level};

/// DICOM C-ECHO SCU
#[derive(Debug, StructOpt)]
struct App {
    /// socket address to SCP,
    /// optionally with AE title
    /// (example: "QUERY-SCP@127.0.0.1:1045")
    addr: String,
    /// verbose mode
    #[structopt(short = "v", long = "verbose")]
    verbose: bool,
    /// the C-ECHO message ID
    #[structopt(short = "m", long = "message-id", default_value = "1")]
    message_id: u16,
    /// the calling AE title
    #[structopt(long = "calling-ae-title", default_value = "ECHOSCU")]
    calling_ae_title: String,
    /// the called Application Entity title,
    /// overrides AE title in address if present [default: ANY-SCP]
    #[structopt(long = "called-ae-title")]
    called_ae_title: Option<String>,
}

fn main() {
    run().unwrap_or_else(|e| {
        tracing::error!("{}", snafu::Report::from_error(e));
        std::process::exit(-2);
    })
}

fn run() -> Result<(), Whatever> {
    let App {
        addr,
        verbose,
        message_id,
        called_ae_title,
        calling_ae_title,
    } = App::from_args();

    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_max_level(if verbose { Level::DEBUG } else { Level::INFO })
            .finish(),
    )
    .whatever_context("Could not set up global logging subscriber")
    .unwrap_or_else(|e: Whatever| {
        eprintln!("[ERROR] {}", snafu::Report::from_error(e));
    });

    let mut association_opt = ClientAssociationOptions::new()
        .with_abstract_syntax("1.2.840.10008.1.1")
        .calling_ae_title(calling_ae_title);
    if let Some(called_ae_title) = called_ae_title {
        association_opt = association_opt.called_ae_title(called_ae_title);
    }
    let mut association = association_opt
        .establish_with(&addr)
        .whatever_context("Could not establish association with SCP")?;

    let pc = association
        .presentation_contexts()
        .first()
        .whatever_context("No presentation context accepted")?
        .clone();

    if verbose {
        println!("Association with {} successful", addr);
    }

    // commands are always in implict VR LE
    let ts = dicom_transfer_syntax_registry::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased();

    let obj = create_echo_command(message_id);

    let mut data = Vec::new();

    obj.write_dataset_with_ts(&mut data, &ts)
        .whatever_context("Failed to construct C-ECHO request")?;

    association
        .send(&Pdu::PData {
            data: vec![PDataValue {
                presentation_context_id: pc.id,
                value_type: PDataValueType::Command,
                is_last: true,
                data,
            }],
        })
        .whatever_context("Failed to send C-ECHO request")?;

    if verbose {
        println!(
            "Echo message sent (msg id {}), awaiting reply...",
            message_id
        );
    }

    let pdu = association
        .receive()
        .whatever_context("Could not receive response from SCP")?;

    match pdu {
        Pdu::PData { data } => {
            let data_value = &data[0];
            let v = &data_value.data;

            let obj = InMemDicomObject::read_dataset_with_ts(v.as_slice(), &ts)
                .whatever_context("Failed to read response dataset from SCP")?;
            if verbose {
                println!("{:?}", obj);
            }

            // check status
            let status_elem = obj
                .element(tags::STATUS)
                .whatever_context("Missing Status code in response")?;

            let status = status_elem
                .to_int::<u16>()
                .whatever_context("Status code in response is not a valid integer")?;
            if verbose {
                println!("Status: {:04X}H", status);
            }
            match status {
                // Success
                0 => {
                    if verbose {
                        println!("C-ECHO successful");
                    }
                }
                // Warning
                1 | 0x0107 | 0x0116 | 0xB000..=0xBFFF => {
                    warn!("Possible issue in C-ECHO (status code {:04X}H)", status);
                }
                0xFF00 | 0xFF01 => {
                    warn!(
                        "Possible issue in C-ECHO: status is pending (status code {:04X}H)",
                        status
                    );
                }
                0xFE00 => {
                    warn!("Operation cancelled");
                }
                _ => {
                    eprintln!("C-ECHO failed (status code {:04X}H)", status);
                }
            }

            // msg ID response, should be equal to sent msg ID
            let msg_id_elem = obj
                .element(tags::MESSAGE_ID_BEING_RESPONDED_TO)
                .whatever_context("Could not retrieve Message ID from response")?;

            if message_id
                != msg_id_elem
                    .to_int()
                    .whatever_context("Message ID is not a valid integer")?
            {
                whatever!("Message ID mismatch");
            }
        }
        pdu => whatever!("Unexpected PDU {:?}", pdu),
    }

    Ok(())
}

fn create_echo_command(message_id: u16) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::new_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UI,
        PrimitiveValue::from(8 + 18 + 8 + 2 + 8 + 2 + 8 + 2),
    ));

    // service
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        dicom_value!(Str, "1.2.840.10008.1.1\0"),
    ));
    // command
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        dicom_value!(U16, [0x0030]),
    ));
    // message ID
    obj.put(DataElement::new(
        tags::MESSAGE_ID,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));
    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0101]),
    ));

    obj
}
<|endoftext|>


<commit_before>use nom::{self, IResult, Needed};

use std::io::{self, Read};
use std::ptr;
use std::cmp;

use super::{Sample, StreamProducer};

/// Represent the different kinds of errors.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum ErrorKind {
  /// Error from I/O.
  IO(io::ErrorKind),
  /// A parser stopped midway and need more bytes to consume.
  Incomplete(usize),
  /// A parser has completes and there is still more bytes to consume.
  Continue,
  /// A parser has completes and there is no more bytes to consume.
  EndOfInput,
  /// A non-specified error.
  Unknown,
  // Parser Error
  /// Failed parsing the "fLaC" header token.
  HeaderParser,
  /// Failed parsing a metadata header.
  MetadataHeaderParser,
  /// Failed parsing the metadata block `StreamInfo`.
  StreamInfoParser,
  /// Failed parsing the metadata block `Padding`.
  PaddingParser,
  /// Failed parsing the metadata block `Application`.
  ApplicationParser,
  /// Failed parsing the metadata block `SeekTable`.
  SeekTableParser,
  /// Failed parsing the metadata block `VorbisComment`.
  VorbisCommentParser,
  /// Failed parsing the metadata block `CueSheet`.
  CueSheetParser,
  /// Failed parsing the metadata block `Picture`.
  PictureParser,
  /// Failed parsing the metadata block `Unknown`.
  UnknownParser,
  /// Failed parsing the blocking strategy inside the frame header.
  BlockingStrategyParser,
  /// Failed parsing the blocking sample inside the frame header.
  BlockingSampleParser,
  /// Failed parsing the channel bits inside the frame header.
  ChannelBitsParser,
  /// Failed parsing the UTF-8 header inside the frame header.
  UTF8HeaderParser,
  /// Failed parsing the UTF-8 body inside the frame header.
  UTF8BodyParser,
  /// Failed parsing the secondary block size inside the frame header.
  BlockSizeParser,
  /// Failed parsing the secondary sample rate inside the frame header.
  SampleRateParser,
  /// Failed parsing the CRC-8 inside the frame header.
  CRC8Parser,
  /// Failed parsing the frame footer, also known as the CRC-16.
  FrameFooterParser,
  /// Failed parsing the subframe header.
  SubframeHeaderParser,
  /// Failed parsing the leading zero for a unary value.
  LeadingZerosParser,
  /// Failed parsing a Constant subframe data.
  ConstantParser,
  /// Failed parsing a Verbatim subframe data.
  VerbatimParser,
  /// Failed parsing a Fixed subframe data.
  FixedParser,
  /// Failed parsing a LPC subframe data.
  LPCParser,
  // Invalid Error
  /// A block type, base on the number, that is outside the range (0-126).
  InvalidBlockType,
  /// An incorrect sync code with the frame header.
  InvalidSyncCode,
  /// A block sample that could cause sync-fooling.
  InvalidBlockSample,
  /// One or more bits are reserved values.
  InvalidChannelBits,
  /// An error occurred in building the UTF-8 value.
  InvalidUTF8,
  /// The stored CRC-8 doesn't match the one generated from the bytes within
  /// the frame header.
  InvalidCRC8,
  /// The stored CRC-16 doesn't match the one generated from the bytes
  /// within the entire frame.
  InvalidCRC16,
  /// A subframe header that could cause sync-fooling.
  InvalidSubframeHeader,
  // Not Found
  /// Some metadata block was not found with a specific filter.
  NotFound,
}

/// Structure that hold a slice of bytes.
pub struct ByteStream<'a> {
  offset: usize,
  bytes: &'a [u8],
}

impl<'a> ByteStream<'a> {
  /// Construct a `ByteStream` based on the passed in byte slice.
  pub fn new(bytes: &'a [u8]) -> Self {
    ByteStream {
      offset: 0,
      bytes: bytes,
    }
  }

  /// Return the number of bytes that haven't been consumed yet.
  #[inline]
  pub fn len(&self) -> usize {
    self.bytes.len() - self.offset
  }

  /// Return true if the stream contains no more bytes.
  #[inline]
  pub fn is_empty(&self) -> bool {
    self.len() == 0
  }
}

impl<'a> StreamProducer for ByteStream<'a> {
  fn parse<F, T>(&mut self, f: F) -> Result<T, ErrorKind>
   where F: FnOnce(&[u8]) -> IResult<&[u8], T, ErrorKind> {
    if self.is_empty() {
      return Err(ErrorKind::EndOfInput);
    }

    match f(&self.bytes[self.offset..]) {
      IResult::Done(i, o)    => {
        self.offset += self.len() - i.len();

        Ok(o)
      }
      IResult::Incomplete(n) => {
        let mut needed = self.len();

        if let Needed::Size(size) = n {
          needed = size;
        }

        Err(ErrorKind::Incomplete(needed))
      }
      IResult::Error(e)      => {
        match e {
          nom::Err::Code(k)               |
          nom::Err::Node(k, _)            |
          nom::Err::Position(k, _)        |
          nom::Err::NodePosition(k, _, _) => {
            if let nom::ErrorKind::Custom(kind) = k {
              Err(kind)
            } else {
              Err(ErrorKind::Unknown)
            }
          }
        }
      },
    }
  }
}

// Growable buffer of bytes.
//
// Mainly used to the `ReadStream` structure but can be used seperately for
// manually filling with some `Read` source.
pub struct Buffer {
  data: Vec<u8>,
  filled: usize,
  offset: usize,
}

impl Buffer {
  // Default constructor for `Buffer`
  pub fn new() -> Self {
    Self::with_capacity(1024)
  }

  // Explicitly set the buffer capacity.
  pub fn with_capacity(capacity: usize) -> Self {
    let mut buffer = Vec::with_capacity(capacity);

    unsafe {
      buffer.set_len(capacity);
    }

    Buffer {
      data: buffer,
      filled: 0,
      offset: 0,
    }
  }

  // Return the number of read bytes that haven't been consumed yet.
  #[inline]
  pub fn len(&self) -> usize {
    self.filled - self.offset
  }

  // Return true if buffer contains no more bytes.
  #[inline]
  pub fn is_empty(&self) -> bool {
    self.len() == 0
  }

  // The set length of the unlining buffer.
  #[inline]
  pub fn capacity(&self) -> usize {
    self.data.len()
  }

  // Return a reference to the slice of unread bytes.
  pub fn as_slice(&self) -> &[u8] {
    &self.data[self.offset..self.filled]
  }

  // Fill the buffer with bytes from a `Read` source.
  pub fn fill<R: Read>(&mut self, reader: &mut R) -> io::Result<usize> {
    reader.read(&mut self.data[self.filled..]).map(|consumed| {
      self.filled += consumed;

      consumed
    })
  }

  // Resize the current buffer
  //
  // This will only allocate data when the size requests is larger than the
  // current capacity of the buffer, otherwise it moves the currently filled
  // data to the beginning of the buffer.
  pub fn resize(&mut self, size: usize) {
    if size > self.data.capacity() {
      self.data.reserve(size);

      let capacity = self.data.capacity();

      unsafe {
        self.data.set_len(capacity);
      }
    }

    if self.data.len() - self.filled < size  {
      let length  = self.filled - self.offset;
      let mut_ptr = self.data.as_mut_ptr();

      unsafe {
        let offset_ptr  = self.data.as_ptr().offset(self.offset as isize);

        ptr::copy(offset_ptr, mut_ptr, length);
      }

      self.filled -= self.offset;
      self.offset  = 0;
    }
  }

  // Move the offset by the amount of consumed bytes.
  pub fn consume(&mut self, consumed: usize) {
    self.offset += consumed;
  }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ParserState {
  Incomplete,
  EndOfInput,
}

fn fill<R: Read>(buffer: &mut Buffer, reader: &mut R, needed: usize)
                 -> io::Result<usize> {
  let mut read = 0;

  if buffer.len() < needed {
    buffer.resize(needed);

    while buffer.len() < needed {
      let size_read = try!(buffer.fill(reader));

      if size_read > 0 {
        read += size_read;
      } else {
        break;
      }
    }
  }

  Ok(read)
}

/// Structure that hold a reader for a source of bytes.
pub struct ReadStream<R: Read> {
  reader: R,
  buffer: Buffer,
  needed: usize,
  state: ParserState,
}

impl<R> ReadStream<R> where R: Read {
  /// Constructor for `ReadStream` based on a `Read` source.
  pub fn new(reader: R) -> Self {
    ReadStream {
      reader: reader,
      buffer: Buffer::new(),
      needed: 0,
      state: ParserState::Incomplete,
    }
  }

  // Fill the stream with bytes from a `Read` source.
  fn fill(&mut self) -> io::Result<usize> {
    let needed = cmp::max(1, self.needed);

    fill(&mut self.buffer, &mut self.reader, needed).map(|consumed| {
      if self.buffer.len() < needed {
        self.state = ParserState::EndOfInput;
      }

      consumed
    })
  }
}

fn from_iresult<T>(buffer: &Buffer, result: IResult<&[u8], T, ErrorKind>)
                   -> Result<(usize, T), ErrorKind> {
  match result {
    IResult::Done(i, o)    => Ok((buffer.len() - i.len(), o)),
    IResult::Incomplete(n) => {
      let mut needed = buffer.capacity() + 1024;

      if let Needed::Size(size) = n {
        needed = size;
      }

      Err(ErrorKind::Incomplete(needed))
    }
    IResult::Error(e)      => {
      match e {
        nom::Err::Code(k)               |
        nom::Err::Node(k, _)            |
        nom::Err::Position(k, _)        |
        nom::Err::NodePosition(k, _, _) => {
          if let nom::ErrorKind::Custom(kind) = k {
            Err(kind)
          } else {
            Err(ErrorKind::Unknown)
          }
        }
      }
    },
  }
}

impl<R> StreamProducer for ReadStream<R> where R: Read {
  fn parse<F, T>(&mut self, f: F) -> Result<T, ErrorKind>
   where F: FnOnce(&[u8]) -> IResult<&[u8], T, ErrorKind> {
    if self.state != ParserState::EndOfInput {
      try!(self.fill().map_err(|e| ErrorKind::IO(e.kind())));
    }

    let mut buffer = &mut self.buffer;

    if buffer.is_empty() {
      self.state = ParserState::EndOfInput;

      return Err(ErrorKind::EndOfInput);
    }

    let result = {
      let iresult = f(buffer.as_slice());

      from_iresult(&buffer, iresult)
    };

    match result {
      Ok((consumed, o)) => {
        buffer.consume(consumed);

        Ok(o)
      }
      Err(kind)         => {
        if let ErrorKind::Incomplete(needed) = kind {
          self.needed = needed;

          Err(ErrorKind::Continue)
        } else {
          Err(kind)
        }
      }
    }
  }
}

macro_rules! sample (
  ($normal: ident, $extended: ident, $bits_per_sample: expr) => (
    impl Sample for $extended {
      type Normal = $normal;

      #[inline]
      fn size() -> usize { $bits_per_sample }

      fn to_normal(sample: Self) -> Option<Self::Normal> {
        use std::$normal;

        let min = $normal::min_value() as $extended;
        let max = $normal::max_value() as $extended;

        if sample > min && sample < max {
          Some(sample as $normal)
        } else {
          None
        }
      }
    }
  )
);

sample!(i8, i16, 8);
sample!(i16, i32, 16);
sample!(i32, i64, 32);

#[cfg(test)]
mod tests {
  use super::*;
  use utility::{Sample, StreamProducer};
  use nom::{self, IResult};

  use std::{i8, i16};

  #[inline]
  fn be_u32(input: &[u8]) -> IResult<&[u8], u32, ErrorKind> {
    to_custom_error!(input, nom::be_u32, Unknown)
  }

  #[test]
  fn test_buffer() {
    let mut buffer = Buffer::new();
    let bytes      = b"Hello World";
    let mut reader = &bytes[..];

    assert!(buffer.is_empty());
    assert_eq!(buffer.capacity(), 1024);

    let bytes_read = buffer.fill(&mut reader).unwrap_or(0);
    let bytes_len  = bytes.len();

    assert_eq!(bytes_read, bytes_len);
    assert_eq!(buffer.len(), bytes_len);
    assert_eq!(buffer.as_slice(), bytes);

    buffer.resize(512);
    assert_eq!(buffer.capacity(), 1024);
  }

  #[test]
  fn test_byte_stream() {
    let bytes      = b"Hello World";
    let mut stream = ByteStream::new(bytes);

    assert_eq!(stream.len(), bytes.len());

    let result = stream.parse(be_u32).unwrap_or(0);

    assert_eq!(result, 1214606444);
    assert_eq!(stream.len(), 7);
  }

  #[test]
  fn test_read_stream() {
    let bytes      = b"Hello World";
    let mut stream = ReadStream::new(&bytes[..]);

    let result = stream.parse(be_u32).unwrap_or(0);

    assert_eq!(result, 1214606444)
  }

  #[test]
  fn test_sample_to_normal_i8() {
    let min = i8::min_value();
    let max = i8::max_value();

    assert_eq!(Sample::to_normal(min as i16), Some(min));
    assert_eq!(Sample::to_normal(0 as i16), Some(0));
    assert_eq!(Sample::to_normal(max as i16), Some(max));

    assert_eq!(Sample::to_normal((min as i16) - 1), None);
    assert_eq!(Sample::to_normal((max as i16) + 1), None);
  }

  #[test]
  fn test_sample_to_normal_i16() {
    let min = i16::min_value();
    let max = i16::max_value();

    assert_eq!(Sample::to_normal(min as i32), Some(min));
    assert_eq!(Sample::to_normal(0 as i32), Some(0));
    assert_eq!(Sample::to_normal(max as i32), Some(max));

    assert_eq!(Sample::to_normal((min as i32) - 1), None);
    assert_eq!(Sample::to_normal((max as i32) + 1), None);
  }
}
<commit_msg>Add failing tests for Sample i32<commit_after>use nom::{self, IResult, Needed};

use std::io::{self, Read};
use std::ptr;
use std::cmp;

use super::{Sample, StreamProducer};

/// Represent the different kinds of errors.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum ErrorKind {
  /// Error from I/O.
  IO(io::ErrorKind),
  /// A parser stopped midway and need more bytes to consume.
  Incomplete(usize),
  /// A parser has completes and there is still more bytes to consume.
  Continue,
  /// A parser has completes and there is no more bytes to consume.
  EndOfInput,
  /// A non-specified error.
  Unknown,
  // Parser Error
  /// Failed parsing the "fLaC" header token.
  HeaderParser,
  /// Failed parsing a metadata header.
  MetadataHeaderParser,
  /// Failed parsing the metadata block `StreamInfo`.
  StreamInfoParser,
  /// Failed parsing the metadata block `Padding`.
  PaddingParser,
  /// Failed parsing the metadata block `Application`.
  ApplicationParser,
  /// Failed parsing the metadata block `SeekTable`.
  SeekTableParser,
  /// Failed parsing the metadata block `VorbisComment`.
  VorbisCommentParser,
  /// Failed parsing the metadata block `CueSheet`.
  CueSheetParser,
  /// Failed parsing the metadata block `Picture`.
  PictureParser,
  /// Failed parsing the metadata block `Unknown`.
  UnknownParser,
  /// Failed parsing the blocking strategy inside the frame header.
  BlockingStrategyParser,
  /// Failed parsing the blocking sample inside the frame header.
  BlockingSampleParser,
  /// Failed parsing the channel bits inside the frame header.
  ChannelBitsParser,
  /// Failed parsing the UTF-8 header inside the frame header.
  UTF8HeaderParser,
  /// Failed parsing the UTF-8 body inside the frame header.
  UTF8BodyParser,
  /// Failed parsing the secondary block size inside the frame header.
  BlockSizeParser,
  /// Failed parsing the secondary sample rate inside the frame header.
  SampleRateParser,
  /// Failed parsing the CRC-8 inside the frame header.
  CRC8Parser,
  /// Failed parsing the frame footer, also known as the CRC-16.
  FrameFooterParser,
  /// Failed parsing the subframe header.
  SubframeHeaderParser,
  /// Failed parsing the leading zero for a unary value.
  LeadingZerosParser,
  /// Failed parsing a Constant subframe data.
  ConstantParser,
  /// Failed parsing a Verbatim subframe data.
  VerbatimParser,
  /// Failed parsing a Fixed subframe data.
  FixedParser,
  /// Failed parsing a LPC subframe data.
  LPCParser,
  // Invalid Error
  /// A block type, base on the number, that is outside the range (0-126).
  InvalidBlockType,
  /// An incorrect sync code with the frame header.
  InvalidSyncCode,
  /// A block sample that could cause sync-fooling.
  InvalidBlockSample,
  /// One or more bits are reserved values.
  InvalidChannelBits,
  /// An error occurred in building the UTF-8 value.
  InvalidUTF8,
  /// The stored CRC-8 doesn't match the one generated from the bytes within
  /// the frame header.
  InvalidCRC8,
  /// The stored CRC-16 doesn't match the one generated from the bytes
  /// within the entire frame.
  InvalidCRC16,
  /// A subframe header that could cause sync-fooling.
  InvalidSubframeHeader,
  // Not Found
  /// Some metadata block was not found with a specific filter.
  NotFound,
}

/// Structure that hold a slice of bytes.
pub struct ByteStream<'a> {
  offset: usize,
  bytes: &'a [u8],
}

impl<'a> ByteStream<'a> {
  /// Construct a `ByteStream` based on the passed in byte slice.
  pub fn new(bytes: &'a [u8]) -> Self {
    ByteStream {
      offset: 0,
      bytes: bytes,
    }
  }

  /// Return the number of bytes that haven't been consumed yet.
  #[inline]
  pub fn len(&self) -> usize {
    self.bytes.len() - self.offset
  }

  /// Return true if the stream contains no more bytes.
  #[inline]
  pub fn is_empty(&self) -> bool {
    self.len() == 0
  }
}

impl<'a> StreamProducer for ByteStream<'a> {
  fn parse<F, T>(&mut self, f: F) -> Result<T, ErrorKind>
   where F: FnOnce(&[u8]) -> IResult<&[u8], T, ErrorKind> {
    if self.is_empty() {
      return Err(ErrorKind::EndOfInput);
    }

    match f(&self.bytes[self.offset..]) {
      IResult::Done(i, o)    => {
        self.offset += self.len() - i.len();

        Ok(o)
      }
      IResult::Incomplete(n) => {
        let mut needed = self.len();

        if let Needed::Size(size) = n {
          needed = size;
        }

        Err(ErrorKind::Incomplete(needed))
      }
      IResult::Error(e)      => {
        match e {
          nom::Err::Code(k)               |
          nom::Err::Node(k, _)            |
          nom::Err::Position(k, _)        |
          nom::Err::NodePosition(k, _, _) => {
            if let nom::ErrorKind::Custom(kind) = k {
              Err(kind)
            } else {
              Err(ErrorKind::Unknown)
            }
          }
        }
      },
    }
  }
}

// Growable buffer of bytes.
//
// Mainly used to the `ReadStream` structure but can be used seperately for
// manually filling with some `Read` source.
pub struct Buffer {
  data: Vec<u8>,
  filled: usize,
  offset: usize,
}

impl Buffer {
  // Default constructor for `Buffer`
  pub fn new() -> Self {
    Self::with_capacity(1024)
  }

  // Explicitly set the buffer capacity.
  pub fn with_capacity(capacity: usize) -> Self {
    let mut buffer = Vec::with_capacity(capacity);

    unsafe {
      buffer.set_len(capacity);
    }

    Buffer {
      data: buffer,
      filled: 0,
      offset: 0,
    }
  }

  // Return the number of read bytes that haven't been consumed yet.
  #[inline]
  pub fn len(&self) -> usize {
    self.filled - self.offset
  }

  // Return true if buffer contains no more bytes.
  #[inline]
  pub fn is_empty(&self) -> bool {
    self.len() == 0
  }

  // The set length of the unlining buffer.
  #[inline]
  pub fn capacity(&self) -> usize {
    self.data.len()
  }

  // Return a reference to the slice of unread bytes.
  pub fn as_slice(&self) -> &[u8] {
    &self.data[self.offset..self.filled]
  }

  // Fill the buffer with bytes from a `Read` source.
  pub fn fill<R: Read>(&mut self, reader: &mut R) -> io::Result<usize> {
    reader.read(&mut self.data[self.filled..]).map(|consumed| {
      self.filled += consumed;

      consumed
    })
  }

  // Resize the current buffer
  //
  // This will only allocate data when the size requests is larger than the
  // current capacity of the buffer, otherwise it moves the currently filled
  // data to the beginning of the buffer.
  pub fn resize(&mut self, size: usize) {
    if size > self.data.capacity() {
      self.data.reserve(size);

      let capacity = self.data.capacity();

      unsafe {
        self.data.set_len(capacity);
      }
    }

    if self.data.len() - self.filled < size  {
      let length  = self.filled - self.offset;
      let mut_ptr = self.data.as_mut_ptr();

      unsafe {
        let offset_ptr  = self.data.as_ptr().offset(self.offset as isize);

        ptr::copy(offset_ptr, mut_ptr, length);
      }

      self.filled -= self.offset;
      self.offset  = 0;
    }
  }

  // Move the offset by the amount of consumed bytes.
  pub fn consume(&mut self, consumed: usize) {
    self.offset += consumed;
  }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ParserState {
  Incomplete,
  EndOfInput,
}

fn fill<R: Read>(buffer: &mut Buffer, reader: &mut R, needed: usize)
                 -> io::Result<usize> {
  let mut read = 0;

  if buffer.len() < needed {
    buffer.resize(needed);

    while buffer.len() < needed {
      let size_read = try!(buffer.fill(reader));

      if size_read > 0 {
        read += size_read;
      } else {
        break;
      }
    }
  }

  Ok(read)
}

/// Structure that hold a reader for a source of bytes.
pub struct ReadStream<R: Read> {
  reader: R,
  buffer: Buffer,
  needed: usize,
  state: ParserState,
}

impl<R> ReadStream<R> where R: Read {
  /// Constructor for `ReadStream` based on a `Read` source.
  pub fn new(reader: R) -> Self {
    ReadStream {
      reader: reader,
      buffer: Buffer::new(),
      needed: 0,
      state: ParserState::Incomplete,
    }
  }

  // Fill the stream with bytes from a `Read` source.
  fn fill(&mut self) -> io::Result<usize> {
    let needed = cmp::max(1, self.needed);

    fill(&mut self.buffer, &mut self.reader, needed).map(|consumed| {
      if self.buffer.len() < needed {
        self.state = ParserState::EndOfInput;
      }

      consumed
    })
  }
}

fn from_iresult<T>(buffer: &Buffer, result: IResult<&[u8], T, ErrorKind>)
                   -> Result<(usize, T), ErrorKind> {
  match result {
    IResult::Done(i, o)    => Ok((buffer.len() - i.len(), o)),
    IResult::Incomplete(n) => {
      let mut needed = buffer.capacity() + 1024;

      if let Needed::Size(size) = n {
        needed = size;
      }

      Err(ErrorKind::Incomplete(needed))
    }
    IResult::Error(e)      => {
      match e {
        nom::Err::Code(k)               |
        nom::Err::Node(k, _)            |
        nom::Err::Position(k, _)        |
        nom::Err::NodePosition(k, _, _) => {
          if let nom::ErrorKind::Custom(kind) = k {
            Err(kind)
          } else {
            Err(ErrorKind::Unknown)
          }
        }
      }
    },
  }
}

impl<R> StreamProducer for ReadStream<R> where R: Read {
  fn parse<F, T>(&mut self, f: F) -> Result<T, ErrorKind>
   where F: FnOnce(&[u8]) -> IResult<&[u8], T, ErrorKind> {
    if self.state != ParserState::EndOfInput {
      try!(self.fill().map_err(|e| ErrorKind::IO(e.kind())));
    }

    let mut buffer = &mut self.buffer;

    if buffer.is_empty() {
      self.state = ParserState::EndOfInput;

      return Err(ErrorKind::EndOfInput);
    }

    let result = {
      let iresult = f(buffer.as_slice());

      from_iresult(&buffer, iresult)
    };

    match result {
      Ok((consumed, o)) => {
        buffer.consume(consumed);

        Ok(o)
      }
      Err(kind)         => {
        if let ErrorKind::Incomplete(needed) = kind {
          self.needed = needed;

          Err(ErrorKind::Continue)
        } else {
          Err(kind)
        }
      }
    }
  }
}

macro_rules! sample (
  ($normal: ident, $extended: ident, $bits_per_sample: expr) => (
    impl Sample for $extended {
      type Normal = $normal;

      #[inline]
      fn size() -> usize { $bits_per_sample }

      fn to_normal(sample: Self) -> Option<Self::Normal> {
        use std::$normal;

        let min = $normal::min_value() as $extended;
        let max = $normal::max_value() as $extended;

        if sample > min && sample < max {
          Some(sample as $normal)
        } else {
          None
        }
      }
    }
  )
);

sample!(i8, i16, 8);
sample!(i16, i32, 16);
sample!(i32, i64, 32);

#[cfg(test)]
mod tests {
  use super::*;
  use utility::{Sample, StreamProducer};
  use nom::{self, IResult};

  use std::{i8, i16, i32};

  #[inline]
  fn be_u32(input: &[u8]) -> IResult<&[u8], u32, ErrorKind> {
    to_custom_error!(input, nom::be_u32, Unknown)
  }

  #[test]
  fn test_buffer() {
    let mut buffer = Buffer::new();
    let bytes      = b"Hello World";
    let mut reader = &bytes[..];

    assert!(buffer.is_empty());
    assert_eq!(buffer.capacity(), 1024);

    let bytes_read = buffer.fill(&mut reader).unwrap_or(0);
    let bytes_len  = bytes.len();

    assert_eq!(bytes_read, bytes_len);
    assert_eq!(buffer.len(), bytes_len);
    assert_eq!(buffer.as_slice(), bytes);

    buffer.resize(512);
    assert_eq!(buffer.capacity(), 1024);
  }

  #[test]
  fn test_byte_stream() {
    let bytes      = b"Hello World";
    let mut stream = ByteStream::new(bytes);

    assert_eq!(stream.len(), bytes.len());

    let result = stream.parse(be_u32).unwrap_or(0);

    assert_eq!(result, 1214606444);
    assert_eq!(stream.len(), 7);
  }

  #[test]
  fn test_read_stream() {
    let bytes      = b"Hello World";
    let mut stream = ReadStream::new(&bytes[..]);

    let result = stream.parse(be_u32).unwrap_or(0);

    assert_eq!(result, 1214606444)
  }

  #[test]
  fn test_sample_to_normal_i8() {
    let min = i8::min_value();
    let max = i8::max_value();

    assert_eq!(Sample::to_normal(min as i16), Some(min));
    assert_eq!(Sample::to_normal(0 as i16), Some(0));
    assert_eq!(Sample::to_normal(max as i16), Some(max));

    assert_eq!(Sample::to_normal((min as i16) - 1), None);
    assert_eq!(Sample::to_normal((max as i16) + 1), None);
  }

  #[test]
  fn test_sample_to_normal_i16() {
    let min = i16::min_value();
    let max = i16::max_value();

    assert_eq!(Sample::to_normal(min as i32), Some(min));
    assert_eq!(Sample::to_normal(0 as i32), Some(0));
    assert_eq!(Sample::to_normal(max as i32), Some(max));

    assert_eq!(Sample::to_normal((min as i32) - 1), None);
    assert_eq!(Sample::to_normal((max as i32) + 1), None);
  }

  #[test]
  fn test_sample_to_normal_i32() {
    let min = i32::min_value();
    let max = i32::max_value();

    assert_eq!(Sample::to_normal(min as i64), Some(min));
    assert_eq!(Sample::to_normal(0 as i64), Some(0));
    assert_eq!(Sample::to_normal(max as i64), Some(max));

    assert_eq!(Sample::to_normal((min as i64) - 1), None);
    assert_eq!(Sample::to_normal((max as i64) + 1), None);
  }
}
<|endoftext|>


<commit_before>use wire::peer_info::PeerState;
use metainfo::MetaInfo;
use wire::action::PeerAction;
use wire::msg::PeerMsg;
use wire::action::PeerStreamAction;
use std::collections::HashMap;
use bit_vec::BitVec;
use file::PartialFile;
use file::PartialFileTrait;
use wire::action::PeerId;
use file::PeerFile;

pub trait Strategy {
    fn query(&mut self, pending: Vec<&Order>, done: Vec<OrderResult>, peers: Vec<&PeerState>, file: &PartialFile) -> Vec<Order>;
    fn on_handshake(&mut self, id: PeerId) -> Vec<Order>;
    fn on_choke(&mut self, id: PeerId) -> Vec<Order>;
    fn on_unchoke(&mut self, id: PeerId) -> Vec<Order>;
    fn on_interested(&mut self, id: PeerId) -> Vec<Order>;
    fn on_not_interested(&mut self, id: PeerId) -> Vec<Order>;
    fn on_have(&mut self, id: PeerId, piece_index: usize) -> Vec<Order>;
    fn on_bitfield(&mut self, id: PeerId, bitfield: BitVec) ->  Vec<Order>;
    fn on_request(&mut self, id: PeerId, index: u32, begin: u32, length: u32) -> Vec<Order>;
    fn on_piece(&mut self, id: PeerId, index: u32, begin: u32, block: Vec<u8>) -> Vec<Order>;
    fn on_cancel(&mut self, id: PeerId, index: u32, begin: u32, block: Vec<u8>) -> Vec<Order>;
    fn on_port(&mut self, id: PeerId, port: u16) -> Vec<Order>;
}

const MAX_BLOCK_SIZE: u64 = 2 << 14;
const MAX_BYTES_PER_REQUEST: u64 = 1024 * 512;
const MAX_PIECES_PEER: usize = 10;

pub struct NormalStrategy {
    orders: HashMap<OrderId, OrderInfo>,
    peers: HashMap<PeerId, PeerState>,
    pieces_to_request: BitVec,
    next_order_id: usize,
    num_pieces: usize,
    piece_length: u64,
    partial_file: PartialFile
}

impl NormalStrategy {
    pub fn new (metainfo: MetaInfo) -> NormalStrategy {
        let num_pieces = metainfo.info.pieces.len();
        let partial_file = PartialFile::new(&metainfo.info);
        let piece_length = partial_file.piece_length();

        NormalStrategy {
            orders: HashMap::new(),
            peers: HashMap::new(),
            pieces_to_request: BitVec::from_elem(num_pieces, true),
            next_order_id: 0,
            partial_file: partial_file,
            num_pieces: num_pieces,
            piece_length: piece_length,
        }
    }

    pub fn request_pieces<F: PartialFileTrait>(&mut self, peers: Vec<&PeerState>, partial_file: &F) -> Vec<PeerAction> {
        let mut actions = Vec::new();
        let ready_peers = peers.iter().filter(|peer| !peer.peer_choking && !peer.am_choking && peer.has_handshake);
        for peer in ready_peers {
            let id = peer.peer_id;
            let missing = self._get_missing(&peer.file);
            //now have eligible pieces to request
            //let mut pieces = missing.into_iter().enumerate().filter(|&(_, x)| x).map(|(i, _)| i as u64).collect::<Vec<u64>>();
            //let piece_len = self.piece_length;
            //let mut bytes_requested = 0;
            //let mut pieces_request = 0;

            let mut msgs = Vec::new();
            let missing = self._get_missing(&peer.file);
            for (piece_index, _) in missing.into_iter().enumerate().filter(|&(_, val)| val){
                let mut pieces = NormalStrategy::piece_request(piece_index as u64, self.piece_length, MAX_BLOCK_SIZE);
                self.pieces_to_request.set(piece_index as usize, false);
                msgs.append(&mut pieces);
            }

            let req_actions = PeerAction(id, PeerStreamAction::SendMessages(msgs));
            actions.push(req_actions);
        }
        actions
    }

    fn piece_request(piece_index: u64, piece_len: u64, max_block_size: u64) -> Vec<PeerMsg> {
        let mut msgs = Vec::new();
        let mut bytes_requested = 0;

        if piece_len > max_block_size {
            let mut offset = 0;
            while offset < piece_len {
                if (offset + max_block_size > piece_len) {
                    msgs.push(PeerMsg::Request(piece_index as u32, offset as u32,
                        max_block_size as u32));
                } else {
                    msgs.push(PeerMsg::Request(piece_index as u32, offset as u32,
                        max_block_size as u32));
                }
                offset += max_block_size;
            }
        } else {
            msgs.push(PeerMsg::Request(piece_index as u32, 0, piece_len as u32));
        }

        msgs
    }

    fn _get_piece_from_req(&mut self, index: usize, begin: u32, offset: u32) -> Option<PeerMsg> {
        if self.partial_file.has_piece(index as usize) {
            let piece = self.partial_file.get_piece_mut(index as usize);
            return match piece.get_offset(begin as usize, offset as usize) {
                Some(piece_data) => {
                    let msg = PeerMsg::Piece(begin, offset, Vec::from(piece_data));
                    Some(msg)
                }
                _ => None,
            };
        }
        None
    }

    fn _make_order(&mut self, action: PeerAction) -> Order {
        self.next_order_id += 1;
        Order {order_id: self.next_order_id, action: action, status: OrderStatus::NotStarted}
    }

    fn _make_msg_order(&mut self, id: PeerId, msgs: Vec<PeerMsg>) -> Order {
        self.next_order_id += 1;
        let action = PeerAction(id, PeerStreamAction::SendMessages(msgs));
        Order {order_id: self.next_order_id, action: action, status: OrderStatus::NotStarted}
    }

    fn _get_missing(&self, peer_file: &PeerFile) -> BitVec {
        let mut missing = self.partial_file.bit_array();
        missing.negate();
        let mut them = peer_file.bit_array();
        missing.intersect(&them);
        missing.intersect(&self.pieces_to_request);
        missing
    }
}

impl Strategy for NormalStrategy {
    fn on_handshake(&mut self, id: PeerId) -> Vec<Order> {
        let peer = PeerState::new(self.num_pieces, id);
        self.peers.insert(id, peer);
        vec![]

    }
    fn on_choke(&mut self, id: PeerId) -> Vec<Order> {
        match self.peers.get_mut(&id) {
            Some(ref mut peer) => peer.peer_choking = true,
            None => {}
        };
        vec![]
    }
    fn on_unchoke(&mut self, id: PeerId) -> Vec<Order> {
        match self.peers.get_mut(&id) {
            Some(ref mut peer) => peer.peer_choking = false,
            None => {}
        };
        vec![]

    }
    fn on_interested(&mut self, id: PeerId) -> Vec<Order>  {
        match self.peers.get_mut(&id) {
            Some(ref mut peer) => peer.peer_interested = true,
            None => {}
        };
        vec![]

    }
    fn on_not_interested(&mut self, id: PeerId) -> Vec<Order> {
        match self.peers.get_mut(&id) {
            Some(ref mut peer) => peer.peer_interested = false,
            None => {}
        };
        vec![]
    }

    fn on_have(&mut self, id: PeerId, piece_index: usize) -> Vec<Order> {
        let msgs = {
            match self.peers.get_mut(&id) {
                Some(ref mut peer) => {
                    peer.file.set(piece_index as usize, true);
                    let file_vec = peer.file.bit_array();
                    if let Some(true) = file_vec.get(piece_index) {
                        let p_msgs = NormalStrategy::piece_request(piece_index as u64, self.piece_length, MAX_BLOCK_SIZE);
                        p_msgs
                    } else {
                        return Vec::new();
                    }
                }
                None => return Vec::new()
            }
        };
        let p_orders = self._make_msg_order(id, msgs);
        return vec![p_orders];
    }

    fn on_bitfield(&mut self, id: PeerId, bitfield: BitVec) ->  Vec<Order> {
        match self.peers.get_mut(&id) {
            Some(ref mut peer) => {
                    let limit = peer.file.bit_array().len();
                    for (i, bit) in bitfield.iter().enumerate()  {
                        if i >= limit {
                            break;
                        }
                        peer.file.set(i, bit);
                    }
                },
            None => {}
        };
        vec![]
    }
    fn on_request(&mut self, id: PeerId, index: u32, begin: u32, length: u32) -> Vec<Order> {
        let piece_result = self._get_piece_from_req(index as usize, begin, length);
        match self.peers.get_mut(&id) {
            Some(ref mut peer) => {
                let response = piece_result;
                match response {
                    Some(r) => unimplemented!(),
                    _ => Vec::new(),
                }
            },
            _ => Vec::new()
        }
    }

    fn on_piece(&mut self, id: PeerId, index: u32, begin: u32, block: Vec<u8>) -> Vec<Order> {
        self.partial_file.add_piece(index as usize, begin as usize, block);
        Vec::new()
    }

    fn on_cancel(&mut self, id: PeerId, index: u32, begin: u32, block: Vec<u8>) -> Vec<Order> {
        Vec::new()
    }

    fn on_port(&mut self, id: PeerId, port: u16) -> Vec<Order> {
        Vec::new()
    }

    fn query(&mut self, pending: Vec<&Order>, done: Vec<OrderResult>, peers: Vec<&PeerState>, file: &PartialFile)
        -> Vec<Order> {
        let actions = self.request_pieces(peers, file);
        let mut orders = Vec::new();

        for action in actions {
            orders.push(self._make_order(action));
            self.next_order_id += 1;
        }

        orders
    }

}

struct OrderInfo;
type OrderId = usize;

pub struct OrderResult {
    status: OrderStatus,
    order_id: OrderId
}

pub struct Order {
    status: OrderStatus,
    order_id: OrderId,
    pub action: PeerAction
}

impl Order {
    pub fn complete(&self, status: OrderStatus) -> OrderResult {
        OrderResult { order_id: self.order_id, status: status}
    }

    pub fn id(&self) -> OrderId {
        self.order_id
    }

    pub fn status(&self) -> OrderStatus {
        self.status
    }

    pub fn set_status(&mut self, status: OrderStatus) {
        self.status = status;
    }
}

#[derive(Eq, PartialEq, Clone, Copy)]
pub enum OrderStatus {
    NotStarted,
    InProgress,
    Failed,
    Done
}
<commit_msg>Make handling of Request correct<commit_after>use wire::peer_info::PeerState;
use metainfo::MetaInfo;
use wire::action::PeerAction;
use wire::msg::PeerMsg;
use wire::action::PeerStreamAction;
use std::collections::HashMap;
use bit_vec::BitVec;
use file::PartialFile;
use file::PartialFileTrait;
use wire::action::PeerId;
use file::PeerFile;

pub trait Strategy {
    fn query(&mut self, pending: Vec<&Order>, done: Vec<OrderResult>, peers: Vec<&PeerState>, file: &PartialFile) -> Vec<Order>;
    fn on_handshake(&mut self, id: PeerId) -> Vec<Order>;
    fn on_choke(&mut self, id: PeerId) -> Vec<Order>;
    fn on_unchoke(&mut self, id: PeerId) -> Vec<Order>;
    fn on_interested(&mut self, id: PeerId) -> Vec<Order>;
    fn on_not_interested(&mut self, id: PeerId) -> Vec<Order>;
    fn on_have(&mut self, id: PeerId, piece_index: usize) -> Vec<Order>;
    fn on_bitfield(&mut self, id: PeerId, bitfield: BitVec) ->  Vec<Order>;
    fn on_request(&mut self, id: PeerId, index: u32, begin: u32, length: u32) -> Vec<Order>;
    fn on_piece(&mut self, id: PeerId, index: u32, begin: u32, block: Vec<u8>) -> Vec<Order>;
    fn on_cancel(&mut self, id: PeerId, index: u32, begin: u32, block: Vec<u8>) -> Vec<Order>;
    fn on_port(&mut self, id: PeerId, port: u16) -> Vec<Order>;
}

const MAX_BLOCK_SIZE: u64 = 2 << 14;
const MAX_BYTES_PER_REQUEST: u64 = 1024 * 512;
const MAX_PIECES_PEER: usize = 10;

pub struct NormalStrategy {
    orders: HashMap<OrderId, OrderInfo>,
    peers: HashMap<PeerId, PeerState>,
    pieces_to_request: BitVec,
    next_order_id: usize,
    num_pieces: usize,
    piece_length: u64,
    partial_file: PartialFile
}

impl NormalStrategy {
    pub fn new (metainfo: MetaInfo) -> NormalStrategy {
        let num_pieces = metainfo.info.pieces.len();
        let partial_file = PartialFile::new(&metainfo.info);
        let piece_length = partial_file.piece_length();

        NormalStrategy {
            orders: HashMap::new(),
            peers: HashMap::new(),
            pieces_to_request: BitVec::from_elem(num_pieces, true),
            next_order_id: 0,
            partial_file: partial_file,
            num_pieces: num_pieces,
            piece_length: piece_length,
        }
    }

    pub fn request_pieces<F: PartialFileTrait>(&mut self, peers: Vec<&PeerState>, partial_file: &F) -> Vec<PeerAction> {
        let mut actions = Vec::new();
        let ready_peers = peers.iter().filter(|peer| !peer.peer_choking && !peer.am_choking && peer.has_handshake);
        for peer in ready_peers {
            let id = peer.peer_id;
            let missing = self._get_missing(&peer.file);
            //now have eligible pieces to request
            //let mut pieces = missing.into_iter().enumerate().filter(|&(_, x)| x).map(|(i, _)| i as u64).collect::<Vec<u64>>();
            //let piece_len = self.piece_length;
            //let mut bytes_requested = 0;
            //let mut pieces_request = 0;

            let mut msgs = Vec::new();
            let missing = self._get_missing(&peer.file);
            for (piece_index, _) in missing.into_iter().enumerate().filter(|&(_, val)| val){
                let mut pieces = NormalStrategy::piece_request(piece_index as u64, self.piece_length, MAX_BLOCK_SIZE);
                self.pieces_to_request.set(piece_index as usize, false);
                msgs.append(&mut pieces);
            }

            let req_actions = PeerAction(id, PeerStreamAction::SendMessages(msgs));
            actions.push(req_actions);
        }
        actions
    }

    fn piece_request(piece_index: u64, piece_len: u64, max_block_size: u64) -> Vec<PeerMsg> {
        let mut msgs = Vec::new();
        let mut bytes_requested = 0;

        if piece_len > max_block_size {
            let mut offset = 0;
            while offset < piece_len {
                if (offset + max_block_size > piece_len) {
                    msgs.push(PeerMsg::Request(piece_index as u32, offset as u32,
                        max_block_size as u32));
                } else {
                    msgs.push(PeerMsg::Request(piece_index as u32, offset as u32,
                        max_block_size as u32));
                }
                offset += max_block_size;
            }
        } else {
            msgs.push(PeerMsg::Request(piece_index as u32, 0, piece_len as u32));
        }

        msgs
    }

    fn _get_piece_from_req(&mut self, index: usize, begin: u32, offset: u32) -> Option<PeerMsg> {
        if self.partial_file.has_piece(index as usize) {
            let piece = self.partial_file.get_piece_mut(index as usize);
            return match piece.get_offset(begin as usize, offset as usize) {
                Some(piece_data) => {
                    let msg = PeerMsg::Piece(begin, offset, Vec::from(piece_data));
                    Some(msg)
                }
                _ => None,
            };
        }
        None
    }

    fn _make_order(&mut self, action: PeerAction) -> Order {
        self.next_order_id += 1;
        Order {order_id: self.next_order_id, action: action, status: OrderStatus::NotStarted}
    }

    fn _make_msg_order(&mut self, id: PeerId, msgs: Vec<PeerMsg>) -> Order {
        self.next_order_id += 1;
        let action = PeerAction(id, PeerStreamAction::SendMessages(msgs));
        Order {order_id: self.next_order_id, action: action, status: OrderStatus::NotStarted}
    }

    fn _get_missing(&self, peer_file: &PeerFile) -> BitVec {
        let mut missing = self.partial_file.bit_array();
        missing.negate();
        let mut them = peer_file.bit_array();
        missing.intersect(&them);
        missing.intersect(&self.pieces_to_request);
        missing
    }
}

impl Strategy for NormalStrategy {
    fn on_handshake(&mut self, id: PeerId) -> Vec<Order> {
        let peer = PeerState::new(self.num_pieces, id);
        self.peers.insert(id, peer);
        vec![]

    }
    fn on_choke(&mut self, id: PeerId) -> Vec<Order> {
        match self.peers.get_mut(&id) {
            Some(ref mut peer) => peer.peer_choking = true,
            None => {}
        };
        vec![]
    }
    fn on_unchoke(&mut self, id: PeerId) -> Vec<Order> {
        match self.peers.get_mut(&id) {
            Some(ref mut peer) => peer.peer_choking = false,
            None => {}
        };
        vec![]

    }
    fn on_interested(&mut self, id: PeerId) -> Vec<Order>  {
        match self.peers.get_mut(&id) {
            Some(ref mut peer) => peer.peer_interested = true,
            None => {}
        };
        vec![]

    }
    fn on_not_interested(&mut self, id: PeerId) -> Vec<Order> {
        match self.peers.get_mut(&id) {
            Some(ref mut peer) => peer.peer_interested = false,
            None => {}
        };
        vec![]
    }

    fn on_have(&mut self, id: PeerId, piece_index: usize) -> Vec<Order> {
        let msgs = {
            match self.peers.get_mut(&id) {
                Some(ref mut peer) => {
                    peer.file.set(piece_index as usize, true);
                    let file_vec = peer.file.bit_array();
                    if let Some(true) = file_vec.get(piece_index) {
                        let p_msgs = NormalStrategy::piece_request(piece_index as u64, self.piece_length, MAX_BLOCK_SIZE);
                        p_msgs
                    } else {
                        return Vec::new();
                    }
                }
                None => return Vec::new()
            }
        };
        let p_orders = self._make_msg_order(id, msgs);
        return vec![p_orders];
    }

    fn on_bitfield(&mut self, id: PeerId, bitfield: BitVec) ->  Vec<Order> {
        match self.peers.get_mut(&id) {
            Some(ref mut peer) => {
                    let limit = peer.file.bit_array().len();
                    for (i, bit) in bitfield.iter().enumerate()  {
                        if i >= limit {
                            break;
                        }
                        peer.file.set(i, bit);
                    }
                },
            None => {}
        };
        vec![]
    }
    fn on_request(&mut self, id: PeerId, index: u32, begin: u32, length: u32) -> Vec<Order> {
        let piece_result = self._get_piece_from_req(index as usize, begin, length);
        match self.peers.get_mut(&id) {
            Some(ref mut peer) => {
                let response = piece_result;
                match response {
                    Some(r) => {
                        if piece_result.is_ok() {
                            return self._make_msg_order(piece_result.unwrap());
                        }
                    },
                    _ => Vec::new(),
                }
            },
            _ => Vec::new()
        }
    }

    fn on_piece(&mut self, id: PeerId, index: u32, begin: u32, block: Vec<u8>) -> Vec<Order> {
        self.partial_file.add_piece(index as usize, begin as usize, block);
        Vec::new()
    }

    fn on_cancel(&mut self, id: PeerId, index: u32, begin: u32, block: Vec<u8>) -> Vec<Order> {
        Vec::new()
    }

    fn on_port(&mut self, id: PeerId, port: u16) -> Vec<Order> {
        Vec::new()
    }

    fn query(&mut self, pending: Vec<&Order>, done: Vec<OrderResult>, peers: Vec<&PeerState>, file: &PartialFile)
        -> Vec<Order> {
        let actions = self.request_pieces(peers, file);
        let mut orders = Vec::new();

        for action in actions {
            orders.push(self._make_order(action));
            self.next_order_id += 1;
        }

        orders
    }

}

struct OrderInfo;
type OrderId = usize;

pub struct OrderResult {
    status: OrderStatus,
    order_id: OrderId
}

pub struct Order {
    status: OrderStatus,
    order_id: OrderId,
    pub action: PeerAction
}

impl Order {
    pub fn complete(&self, status: OrderStatus) -> OrderResult {
        OrderResult { order_id: self.order_id, status: status}
    }

    pub fn id(&self) -> OrderId {
        self.order_id
    }

    pub fn status(&self) -> OrderStatus {
        self.status
    }

    pub fn set_status(&mut self, status: OrderStatus) {
        self.status = status;
    }
}

#[derive(Eq, PartialEq, Clone, Copy)]
pub enum OrderStatus {
    NotStarted,
    InProgress,
    Failed,
    Done
}
<|endoftext|>