<commit_before>//! TODO Documentation

use std::marker::PhantomData;
use std::panic;
use std::rc::{Rc, Weak};
use std::sync::atomic::{AtomicBool, Ordering};

use wlroots_sys::{wlr_cursor_attach_output_layout, wlr_output_effective_resolution,
                  wlr_output_layout, wlr_output_layout_add, wlr_output_layout_add_auto,
                  wlr_output_layout_create, wlr_output_layout_destroy, wlr_output_layout_move,
                  wlr_output_layout_output, wlr_output_layout_remove};

use errors::{UpgradeHandleErr, UpgradeHandleResult};

use {Cursor, CursorBuilder, Origin, Output, OutputHandle};

#[derive(Debug)]
pub struct OutputLayout {
    /// The structure that ensures weak handles to this structure are still alive.
    ///
    /// They contain weak handles, and will safely not use dead memory when this
    /// is freed by wlroots.
    ///
    /// If this is `None`, then this is from an upgraded `OutputHandle`, and
    /// the operations are **unchecked**.
    /// This is means safe operations might fail, but only if you use the unsafe
    /// marked function `upgrade` on a `OutputHandle`.
    liveliness: Option<Rc<AtomicBool>>,
    /// The output_layout ptr that refers to this `OutputLayout`
    layout: *mut wlr_output_layout,
    cursors: Vec<Cursor>
}

/// A handle to an `OutputLayout`.
///
/// Used internally by `Output` to gain access to the `OutputLayout`.
#[derive(Debug, Clone)]
pub(crate) struct OutputLayoutHandle {
    /// The Rc that ensures that this handle is still alive.
    ///
    /// When wlroots deallocates the pointer associated with this handle,
    /// this can no longer be used.
    handle: Weak<AtomicBool>,
    /// The output_layout ptr that refers to this `OutputLayout`
    layout: *mut wlr_output_layout
}

/// The coordinate information of an `Output` within an `OutputLayout`.
#[derive(Debug)]
pub struct OutputLayoutOutput<'output> {
    layout_output: *mut wlr_output_layout_output,
    phantom: PhantomData<&'output OutputLayout>
}

impl OutputLayout {
    /// Construct a new OuputLayout.
    pub fn new() -> Option<Self> {
        unsafe {
            let layout = wlr_output_layout_create();
            if layout.is_null() {
                None
            } else {
                Some(OutputLayout { liveliness: Some(Rc::new(AtomicBool::new(false))),
                                    layout,
                                    cursors: Vec::new() })
            }
        }
    }

    /// Get the outputs associated with this OutputLayout.
    ///
    /// Also returns their absolute position within the layout.
    pub fn outputs(&mut self) -> Vec<(OutputHandle, Origin)> {
        unsafe {
            let mut result = vec![];
            wl_list_for_each!((*self.layout).outputs, link, (pos: wlr_output_layout_output) => {
                result.push((OutputHandle::from_ptr((*pos).output),
                             Origin::new((*pos).x, (*pos).y)))
            });
            result
        }
    }

    /// Get the Outputs in the OutputLayout coupled with their output information.
    ///
    /// For a version that isn't bound by lifetimes, see `outputs`.
    pub fn outputs_layouts<'output>(&'output mut self) -> Vec<OutputLayoutOutput<'output>> {
        unsafe {
            let mut result = vec![];
            wl_list_for_each!((*self.layout).outputs, link,
                              (layout_output: wlr_output_layout_output) => {
                                  result.push(OutputLayoutOutput { layout_output,
                                                                   phantom: PhantomData
                                  })
                              });
            result
        }
    }

    pub fn cursors(&mut self) -> &mut [Cursor] {
        self.cursors.as_mut_slice()
    }

    /// Attach a cursor to this OutputLayout.
    pub fn attach_cursor(&mut self, cursor: CursorBuilder) {
        unsafe {
            let cursor = cursor.build(self.weak_reference());
            wlr_cursor_attach_output_layout(cursor.as_ptr(), self.layout);
            self.cursors.push(cursor);
        }
    }

    /// Adds an output to the layout at the given coordinates.
    pub fn add(&mut self, output: &mut Output, origin: Origin) {
        let (x, y) = (origin.x, origin.y);
        unsafe { wlr_output_layout_add(self.layout, output.as_ptr(), x, y) }
    }

    /// Adds an output to the layout, automatically positioning it with
    /// the others that are already there.
    pub fn add_auto(&mut self, output: &mut Output) {
        unsafe {
            let layout_handle = self.weak_reference();
            output.set_output_layout(Some(layout_handle));
            wlr_output_layout_add_auto(self.layout, output.as_ptr());
            wlr_log!(L_DEBUG, "Added {:?} to {:?}", output, self);
        }
    }

    /// Moves the output to the given coordinates.
    ///
    /// If the output is not part of this layout this does nothing.
    pub fn move_output(&mut self, output: &mut Output, origin: Origin) {
        let (x, y) = (origin.x, origin.y);
        unsafe { wlr_output_layout_move(self.layout, output.as_ptr(), x, y) }
    }

    /// Remove an output from this layout.
    ///
    /// If the output was not in the layout, does nothing.
    pub fn remove(&mut self, output: &mut Output) {
        wlr_log!(L_DEBUG, "Removing {:?} from {:?}", output, self);
        unsafe {
            output.clear_output_layout_data();
            wlr_output_layout_remove(self.layout, output.as_ptr());
        };
    }

    /// Creates a weak reference to an `OutputLayout`.
    ///
    /// # Panics
    /// If this `OutputLayout` is a previously upgraded `OutputLayoutHandle`,
    /// then this function will panic.
    pub(crate) fn weak_reference(&self) -> OutputLayoutHandle {
        let arc = self.liveliness.as_ref()
                      .expect("Cannot downgrade a previously upgraded OutputLayoutHandle");
        OutputLayoutHandle { handle: Rc::downgrade(arc),
                             layout: self.layout }
    }

    unsafe fn from_handle(handle: &OutputLayoutHandle) -> Self {
        OutputLayout { liveliness: None,
                       layout: handle.as_ptr(),
                       cursors: Vec::new() }
    }
}

impl Drop for OutputLayout {
    fn drop(&mut self) {
        match self.liveliness {
            None => {}
            Some(ref liveliness) => {
                if Rc::strong_count(liveliness) == 1 {
                    unsafe { wlr_output_layout_destroy(self.layout) }
                    wlr_log!(L_DEBUG, "Dropped {:?}", self);
                    let weak_count = Rc::weak_count(liveliness);
                    if weak_count > 0 {
                        wlr_log!(L_DEBUG,
                                 "Still {} weak pointers to OutputLayout {:?}",
                                 weak_count,
                                 self.layout);
                    }
                }
            }
        }
    }
}

impl OutputLayoutHandle {
    /// Upgrades the `OutputLayoutHandle` to a reference
    /// to the backing `OutputLayout`.
    ///
    /// # Unsafety
    /// This function is unsafe, because it creates an unbound `OutputLayout`
    /// which may live forever..
    /// But the actual lifetime of `OutputLayout` is determined by the user.
    pub(crate) unsafe fn upgrade(&self) -> UpgradeHandleResult<OutputLayout> {
        self.handle.upgrade()
            .ok_or(UpgradeHandleErr::AlreadyDropped)
            // NOTE
            // We drop the Rc here because having two would allow a dangling
            // pointer to exist!
            .and_then(|check| {
                let output_layout = OutputLayout::from_handle(self);
                if check.load(Ordering::Acquire) {
                    return Err(UpgradeHandleErr::AlreadyBorrowed)
                }
                check.store(true, Ordering::Release);
                Ok(output_layout)
            })
    }

    /// Run a function on the referenced OutputLayout, if it still exists
    ///
    /// Returns the result of the function, if successful.
    ///
    /// # Safety
    /// By enforcing a rather harsh limit on the lifetime of the OutputLayout
    /// to a short lived scope of an anonymous function,
    /// this function ensures the OutputLayout does not live longer
    /// than it exists (because the lifetime is controlled by the user).
    pub fn run<F, R>(&mut self, runner: F) -> UpgradeHandleResult<R>
        where F: FnOnce(&mut OutputLayout) -> R
    {
        let mut output_layout = unsafe { self.upgrade()? };
        let res = panic::catch_unwind(panic::AssertUnwindSafe(|| runner(&mut output_layout)));
        self.handle.upgrade().map(|check| {
                                      // Sanity check that it hasn't been tampered with.
                                      if !check.load(Ordering::Acquire) {
                                          wlr_log!(L_ERROR,
                                                   "After running OutputLayout callback, mutable \
                                                    lock was false for: {:?}",
                                                   output_layout);
                                          panic!("Lock in incorrect state!");
                                      }
                                      check.store(false, Ordering::Release);
                                  });
        match res {
            Ok(res) => Ok(res),
            Err(err) => panic::resume_unwind(err)
        }
    }

    unsafe fn as_ptr(&self) -> *mut wlr_output_layout {
        self.layout
    }
}

impl<'output> OutputLayoutOutput<'output> {
    /// Get the absolute top left edge coordinate of this output in the output
    /// layout.
    pub fn top_left_edge(&self) -> Origin {
        unsafe { Origin::new((*self.layout_output).x, (*self.layout_output).y) }
    }

    /// Get the absolute top right edge coordinate of this output in the output
    /// layout.
    pub fn top_right_edge(&self) -> Origin {
        unsafe {
            let (mut width, mut _height) = (0, 0);
            wlr_output_effective_resolution((*self.layout_output).output, &mut width, &mut _height);
            let (x, y) = ((*self.layout_output).x, (*self.layout_output).y);
            Origin::new(x + width, y)
        }
    }

    pub fn bottom_left_edge(&self) -> Origin {
        unsafe {
            let (mut _width, mut height) = (0, 0);
            wlr_output_effective_resolution((*self.layout_output).output, &mut _width, &mut height);
            let (x, y) = ((*self.layout_output).x, (*self.layout_output).y);
            Origin::new(x, y + height)
        }
    }

    pub fn bottom_right_edge(&self) -> Origin {
        unsafe {
            let (mut width, mut height) = (0, 0);
            wlr_output_effective_resolution((*self.layout_output).output, &mut width, &mut height);
            let (x, y) = ((*self.layout_output).x, (*self.layout_output).y);
            Origin::new(x + height, y + height)
        }
    }
}

impl PartialEq for OutputLayoutHandle {
    fn eq(&self, other: &OutputLayoutHandle) -> bool {
        self.layout == other.layout
    }
}

impl Eq for OutputLayoutHandle {}
<commit_msg>Wrape the rest of the OutputLayout functions<commit_after>//! TODO Documentation

use std::marker::PhantomData;
use std::{ptr, panic};
use std::rc::{Rc, Weak};
use std::sync::atomic::{AtomicBool, Ordering};

use wlroots_sys::{wlr_cursor_attach_output_layout, wlr_output_effective_resolution,
                  wlr_output_layout, wlr_output_layout_add, wlr_output_layout_add_auto,
                  wlr_output_layout_closest_point, wlr_output_layout_contains_point,
                  wlr_output_layout_create, wlr_output_layout_destroy, wlr_output_layout_get_box,
                  wlr_output_layout_get_center_output, wlr_output_layout_intersects,
                  wlr_output_layout_move, wlr_output_layout_output,
                  wlr_output_layout_output_coords, wlr_output_layout_remove};

use errors::{UpgradeHandleErr, UpgradeHandleResult};

use {Area, Cursor, CursorBuilder, Origin, Output, OutputHandle};

#[derive(Debug)]
pub struct OutputLayout {
    /// The structure that ensures weak handles to this structure are still alive.
    ///
    /// They contain weak handles, and will safely not use dead memory when this
    /// is freed by wlroots.
    ///
    /// If this is `None`, then this is from an upgraded `OutputHandle`, and
    /// the operations are **unchecked**.
    /// This is means safe operations might fail, but only if you use the unsafe
    /// marked function `upgrade` on a `OutputHandle`.
    liveliness: Option<Rc<AtomicBool>>,
    /// The output_layout ptr that refers to this `OutputLayout`
    layout: *mut wlr_output_layout,
    cursors: Vec<Cursor>
}

/// A handle to an `OutputLayout`.
///
/// Used internally by `Output` to gain access to the `OutputLayout`.
#[derive(Debug, Clone)]
pub(crate) struct OutputLayoutHandle {
    /// The Rc that ensures that this handle is still alive.
    ///
    /// When wlroots deallocates the pointer associated with this handle,
    /// this can no longer be used.
    handle: Weak<AtomicBool>,
    /// The output_layout ptr that refers to this `OutputLayout`
    layout: *mut wlr_output_layout
}

/// The coordinate information of an `Output` within an `OutputLayout`.
#[derive(Debug)]
pub struct OutputLayoutOutput<'output> {
    layout_output: *mut wlr_output_layout_output,
    phantom: PhantomData<&'output OutputLayout>
}

impl OutputLayout {
    /// Construct a new OuputLayout.
    pub fn new() -> Option<Self> {
        unsafe {
            let layout = wlr_output_layout_create();
            if layout.is_null() {
                None
            } else {
                Some(OutputLayout { liveliness: Some(Rc::new(AtomicBool::new(false))),
                                    layout,
                                    cursors: Vec::new() })
            }
        }
    }

    /// Get the outputs associated with this OutputLayout.
    ///
    /// Also returns their absolute position within the layout.
    pub fn outputs(&mut self) -> Vec<(OutputHandle, Origin)> {
        unsafe {
            let mut result = vec![];
            wl_list_for_each!((*self.layout).outputs, link, (pos: wlr_output_layout_output) => {
                result.push((OutputHandle::from_ptr((*pos).output),
                             Origin::new((*pos).x, (*pos).y)))
            });
            result
        }
    }

    /// Get the Outputs in the OutputLayout coupled with their output information.
    ///
    /// For a version that isn't bound by lifetimes, see `outputs`.
    pub fn outputs_layouts<'output>(&'output mut self) -> Vec<OutputLayoutOutput<'output>> {
        unsafe {
            let mut result = vec![];
            wl_list_for_each!((*self.layout).outputs, link,
                              (layout_output: wlr_output_layout_output) => {
                                  result.push(OutputLayoutOutput { layout_output,
                                                                   phantom: PhantomData
                                  })
                              });
            result
        }
    }

    pub fn cursors(&mut self) -> &mut [Cursor] {
        self.cursors.as_mut_slice()
    }

    /// Attach a cursor to this OutputLayout.
    pub fn attach_cursor(&mut self, cursor: CursorBuilder) {
        unsafe {
            let cursor = cursor.build(self.weak_reference());
            wlr_cursor_attach_output_layout(cursor.as_ptr(), self.layout);
            self.cursors.push(cursor);
        }
    }

    /// Adds an output to the layout at the given coordinates.
    pub fn add(&mut self, output: &mut Output, origin: Origin) {
        let (x, y) = (origin.x, origin.y);
        unsafe { wlr_output_layout_add(self.layout, output.as_ptr(), x, y) }
    }

    /// Adds an output to the layout, automatically positioning it with
    /// the others that are already there.
    pub fn add_auto(&mut self, output: &mut Output) {
        unsafe {
            let layout_handle = self.weak_reference();
            output.set_output_layout(Some(layout_handle));
            wlr_output_layout_add_auto(self.layout, output.as_ptr());
            wlr_log!(L_DEBUG, "Added {:?} to {:?}", output, self);
        }
    }

    /// Moves the output to the given coordinates.
    ///
    /// If the output is not part of this layout this does nothing.
    pub fn move_output(&mut self, output: &mut Output, origin: Origin) {
        let (x, y) = (origin.x, origin.y);
        unsafe { wlr_output_layout_move(self.layout, output.as_ptr(), x, y) }
    }

    /// Get the closest point on this layout from the given point from the reference
    /// output.
    ///
    /// If reference is None, gets the closest point from the entire layout.
    ///
    /// Returns the closest point in the format (x, y)
    pub fn closest_point<'this, O>(&mut self, reference: O, x: f64, y: f64) -> (f64, f64)
        where O: Into<Option<&'this mut Output>> {
        unsafe {
            let output_ptr = reference.into().map(|output| output.as_ptr()).unwrap_or(ptr::null_mut());
            let (ref mut out_x, ref mut out_y) = (0.0, 0.0);
            wlr_output_layout_closest_point(self.layout, output_ptr, x, y, out_x, out_y);
            (*out_x, *out_y)
        }
    }

    /// Determines if the `OutputLayout` contains the `Output` at the given point.
    pub fn contains_point(&mut self, output: &mut Output, origin: Origin) -> bool {
        unsafe {
            wlr_output_layout_contains_point(self.layout, output.as_ptr(), origin.x, origin.y)
        }
    }

    /// Get the box of the layout for the given reference output.
    ///
    /// If `reference` is None, the box will be for the extents of the entire layout.
    pub fn get_box<'this, O>(&mut self, reference: O) -> Area where O: Into<Option<&'this mut Output>> {
        unsafe {
            let output_ptr = reference.into().map(|output| output.as_ptr()).unwrap_or(ptr::null_mut());
            Area(*wlr_output_layout_get_box(self.layout, output_ptr))
        }
    }

    /// Get the output closest to the center of the layout extents, if one exists.
    pub fn get_center_output(&mut self) -> Option<OutputHandle> {
        unsafe {
            let output = wlr_output_layout_get_center_output(self.layout);
            if output.is_null() {
                None
            } else {
                Some(OutputHandle::from_ptr(output))
            }
        }
    }

    /// Determines if the `Output` in the `OutputLayout` intersects with
    /// the provided `Area`.
    pub fn intersects(&mut self, output: &mut Output, area: Area) -> bool {
        unsafe {
            wlr_output_layout_intersects(self.layout, output.as_ptr(), &area.0)
        }
    }

    /// Given x and y as pointers to global coordinates, adjusts them to local output
    /// coordinates relative to the given reference output.
    pub fn output_coords(&mut self, output: &mut Output, x: &mut f64, y: &mut f64) {
        unsafe {
            wlr_output_layout_output_coords(self.layout, output.as_ptr(), x, y)
        }
    }

    /// Remove an output from this layout.
    ///
    /// If the output was not in the layout, does nothing.
    pub fn remove(&mut self, output: &mut Output) {
        wlr_log!(L_DEBUG, "Removing {:?} from {:?}", output, self);
        unsafe {
            output.clear_output_layout_data();
            wlr_output_layout_remove(self.layout, output.as_ptr());
        };
    }

    /// Creates a weak reference to an `OutputLayout`.
    ///
    /// # Panics
    /// If this `OutputLayout` is a previously upgraded `OutputLayoutHandle`,
    /// then this function will panic.
    pub(crate) fn weak_reference(&self) -> OutputLayoutHandle {
        let arc = self.liveliness.as_ref()
                      .expect("Cannot downgrade a previously upgraded OutputLayoutHandle");
        OutputLayoutHandle { handle: Rc::downgrade(arc),
                             layout: self.layout }
    }

    unsafe fn from_handle(handle: &OutputLayoutHandle) -> Self {
        OutputLayout { liveliness: None,
                       layout: handle.as_ptr(),
                       cursors: Vec::new() }
    }
}

impl Drop for OutputLayout {
    fn drop(&mut self) {
        match self.liveliness {
            None => {}
            Some(ref liveliness) => {
                if Rc::strong_count(liveliness) == 1 {
                    unsafe { wlr_output_layout_destroy(self.layout) }
                    wlr_log!(L_DEBUG, "Dropped {:?}", self);
                    let weak_count = Rc::weak_count(liveliness);
                    if weak_count > 0 {
                        wlr_log!(L_DEBUG,
                                 "Still {} weak pointers to OutputLayout {:?}",
                                 weak_count,
                                 self.layout);
                    }
                }
            }
        }
    }
}

impl OutputLayoutHandle {
    /// Upgrades the `OutputLayoutHandle` to a reference
    /// to the backing `OutputLayout`.
    ///
    /// # Unsafety
    /// This function is unsafe, because it creates an unbound `OutputLayout`
    /// which may live forever..
    /// But the actual lifetime of `OutputLayout` is determined by the user.
    pub(crate) unsafe fn upgrade(&self) -> UpgradeHandleResult<OutputLayout> {
        self.handle.upgrade()
            .ok_or(UpgradeHandleErr::AlreadyDropped)
            // NOTE
            // We drop the Rc here because having two would allow a dangling
            // pointer to exist!
            .and_then(|check| {
                let output_layout = OutputLayout::from_handle(self);
                if check.load(Ordering::Acquire) {
                    return Err(UpgradeHandleErr::AlreadyBorrowed)
                }
                check.store(true, Ordering::Release);
                Ok(output_layout)
            })
    }

    /// Run a function on the referenced OutputLayout, if it still exists
    ///
    /// Returns the result of the function, if successful.
    ///
    /// # Safety
    /// By enforcing a rather harsh limit on the lifetime of the OutputLayout
    /// to a short lived scope of an anonymous function,
    /// this function ensures the OutputLayout does not live longer
    /// than it exists (because the lifetime is controlled by the user).
    pub fn run<F, R>(&mut self, runner: F) -> UpgradeHandleResult<R>
        where F: FnOnce(&mut OutputLayout) -> R
    {
        let mut output_layout = unsafe { self.upgrade()? };
        let res = panic::catch_unwind(panic::AssertUnwindSafe(|| runner(&mut output_layout)));
        self.handle.upgrade().map(|check| {
                                      // Sanity check that it hasn't been tampered with.
                                      if !check.load(Ordering::Acquire) {
                                          wlr_log!(L_ERROR,
                                                   "After running OutputLayout callback, mutable \
                                                    lock was false for: {:?}",
                                                   output_layout);
                                          panic!("Lock in incorrect state!");
                                      }
                                      check.store(false, Ordering::Release);
                                  });
        match res {
            Ok(res) => Ok(res),
            Err(err) => panic::resume_unwind(err)
        }
    }

    unsafe fn as_ptr(&self) -> *mut wlr_output_layout {
        self.layout
    }
}

impl<'output> OutputLayoutOutput<'output> {
    /// Get the absolute top left edge coordinate of this output in the output
    /// layout.
    pub fn top_left_edge(&self) -> Origin {
        unsafe { Origin::new((*self.layout_output).x, (*self.layout_output).y) }
    }

    /// Get the absolute top right edge coordinate of this output in the output
    /// layout.
    pub fn top_right_edge(&self) -> Origin {
        unsafe {
            let (mut width, mut _height) = (0, 0);
            wlr_output_effective_resolution((*self.layout_output).output, &mut width, &mut _height);
            let (x, y) = ((*self.layout_output).x, (*self.layout_output).y);
            Origin::new(x + width, y)
        }
    }

    pub fn bottom_left_edge(&self) -> Origin {
        unsafe {
            let (mut _width, mut height) = (0, 0);
            wlr_output_effective_resolution((*self.layout_output).output, &mut _width, &mut height);
            let (x, y) = ((*self.layout_output).x, (*self.layout_output).y);
            Origin::new(x, y + height)
        }
    }

    pub fn bottom_right_edge(&self) -> Origin {
        unsafe {
            let (mut width, mut height) = (0, 0);
            wlr_output_effective_resolution((*self.layout_output).output, &mut width, &mut height);
            let (x, y) = ((*self.layout_output).x, (*self.layout_output).y);
            Origin::new(x + height, y + height)
        }
    }
}

impl PartialEq for OutputLayoutHandle {
    fn eq(&self, other: &OutputLayoutHandle) -> bool {
        self.layout == other.layout
    }
}

impl Eq for OutputLayoutHandle {}
<|endoftext|>


<commit_before>// spell-checker:ignore (URL) appspot

use crate::numeric::*;

pub(crate) trait Basis {
    const BASIS: &'static [u64];
}

impl Basis for Montgomery<u64> {
    // Small set of bases for the Miller-Rabin prime test, valid for all 64b integers;
    //  discovered by Jim Sinclair on 2011-04-20, see miller-rabin.appspot.com
    #[allow(clippy::unreadable_literal)]
    const BASIS: &'static [u64] = &[2, 325, 9375, 28178, 450775, 9780504, 1795265022];
}

impl Basis for Montgomery<u32> {
    // Small set of bases for the Miller-Rabin prime test, valid for all 32b integers;
    //  discovered by Steve Worley on 2013-05-27, see miller-rabin.appspot.com
    #[allow(clippy::unreadable_literal)]
    const BASIS: &'static [u64] = &[
        4230279247111683200,
        14694767155120705706,
        16641139526367750375,
    ];
}

#[derive(Eq, PartialEq)]
pub(crate) enum Result {
    Prime,
    Pseudoprime,
    Composite(u64),
}

impl Result {
    pub(crate) fn is_prime(&self) -> bool {
        *self == Result::Prime
    }
}

// Deterministic Miller-Rabin primality-checking algorithm, adapted to extract
// (some) dividers; it will fail to factor strong pseudoprimes.
#[allow(clippy::many_single_char_names)]
pub(crate) fn test<A: Arithmetic + Basis>(m: A) -> Result {
    use self::Result::*;

    let n = m.modulus();
    debug_assert!(n > 1);
    debug_assert!(n % 2 != 0);

    // n-1 = r 2ⁱ
    let i = (n - 1).trailing_zeros();
    let r = (n - 1) >> i;

    let one = m.one();
    let minus_one = m.minus_one();

    'witness: for _a in A::BASIS.iter() {
        let _a = _a % n;
        if _a == 0 {
            continue;
        }

        let a = m.from_u64(_a);

        // x = a^r mod n
        let mut x = m.pow(a, r);

        if x == one || x == minus_one {
            continue;
        }

        for _ in 1..i {
            let y = m.mul(x, x);
            if y == one {
                return Composite(gcd(m.to_u64(x) - 1, m.modulus()));
            }
            if y == minus_one {
                // This basis element is not a witness of `n` being composite.
                // Keep looking.
                continue 'witness;
            }
            x = y;
        }

        return Pseudoprime;
    }

    Prime
}

// Used by build.rs' tests and debug assertions
#[allow(dead_code)]
pub(crate) fn is_prime(n: u64) -> bool {
    if n < 2 {
        false
    } else if n % 2 == 0 {
        n == 2
    } else {
        test::<Montgomery<u64>>(Montgomery::new(n)).is_prime()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::numeric::{Arithmetic, Montgomery};
    use quickcheck::quickcheck;
    use std::iter;
    const LARGEST_U64_PRIME: u64 = 0xFFFFFFFFFFFFFFC5;

    fn primes() -> impl Iterator<Item = u64> {
        iter::once(2).chain(odd_primes())
    }

    fn odd_primes() -> impl Iterator<Item = u64> {
        use crate::table::{NEXT_PRIME, P_INVS_U64};
        P_INVS_U64
            .iter()
            .map(|(p, _, _)| *p)
            .chain(iter::once(NEXT_PRIME))
    }

    #[test]
    fn largest_prime() {
        assert!(is_prime(LARGEST_U64_PRIME));
    }

    #[test]
    fn largest_composites() {
        for i in LARGEST_U64_PRIME + 1..=u64::MAX {
            assert!(!is_prime(i), "2⁶⁴ - {} reported prime", u64::MAX - i + 1);
        }
    }

    #[test]
    fn two() {
        assert!(is_prime(2));
    }

    // TODO: Deduplicate with macro in numeric.rs
    macro_rules! parametrized_check {
        ( $f:ident ) => {
            paste::item! {
                #[test]
                fn [< $f _ u32 >]() {
                    $f::<Montgomery<u32>>()
                }
                #[test]
                fn [< $f _ u64 >]() {
                    $f::<Montgomery<u64>>()
                }
            }
        };
    }

    fn first_primes<A: Arithmetic + Basis>() {
        for p in odd_primes() {
            assert!(test(A::new(p)).is_prime(), "{} reported composite", p);
        }
    }
    parametrized_check!(first_primes);

    #[test]
    fn one() {
        assert!(!is_prime(1));
    }
    #[test]
    fn zero() {
        assert!(!is_prime(0));
    }

    fn first_composites<A: Arithmetic + Basis>() {
        for (p, q) in primes().zip(odd_primes()) {
            for i in p + 1..q {
                assert!(!is_prime(i), "{} reported prime", i);
            }
        }
    }
    parametrized_check!(first_composites);

    #[test]
    fn issue_1556() {
        // 10 425 511 = 2441 × 4271
        assert!(!is_prime(10_425_511));
    }

    fn small_semiprimes<A: Arithmetic + Basis>() {
        for p in odd_primes() {
            for q in odd_primes().take_while(|q| *q <= p) {
                let n = p * q;
                let m = A::new(n);
                assert!(!test(m).is_prime(), "{} = {} × {} reported prime", n, p, q);
            }
        }
    }
    parametrized_check!(small_semiprimes);

    quickcheck! {
        fn composites(i: u64, j: u64) -> bool {
            i < 2 || j < 2 || !is_prime(i*j)
        }
    }
}
<commit_msg>factor::miller_rabin::test: Minor readability improvement<commit_after>// spell-checker:ignore (URL) appspot

use crate::numeric::*;

pub(crate) trait Basis {
    const BASIS: &'static [u64];
}

impl Basis for Montgomery<u64> {
    // Small set of bases for the Miller-Rabin prime test, valid for all 64b integers;
    //  discovered by Jim Sinclair on 2011-04-20, see miller-rabin.appspot.com
    #[allow(clippy::unreadable_literal)]
    const BASIS: &'static [u64] = &[2, 325, 9375, 28178, 450775, 9780504, 1795265022];
}

impl Basis for Montgomery<u32> {
    // Small set of bases for the Miller-Rabin prime test, valid for all 32b integers;
    //  discovered by Steve Worley on 2013-05-27, see miller-rabin.appspot.com
    #[allow(clippy::unreadable_literal)]
    const BASIS: &'static [u64] = &[
        4230279247111683200,
        14694767155120705706,
        16641139526367750375,
    ];
}

#[derive(Eq, PartialEq)]
pub(crate) enum Result {
    Prime,
    Pseudoprime,
    Composite(u64),
}

impl Result {
    pub(crate) fn is_prime(&self) -> bool {
        *self == Result::Prime
    }
}

// Deterministic Miller-Rabin primality-checking algorithm, adapted to extract
// (some) dividers; it will fail to factor strong pseudoprimes.
#[allow(clippy::many_single_char_names)]
pub(crate) fn test<A: Arithmetic + Basis>(m: A) -> Result {
    use self::Result::*;

    let n = m.modulus();
    debug_assert!(n > 1);
    debug_assert!(n % 2 != 0);

    // n-1 = r 2ⁱ
    let i = (n - 1).trailing_zeros();
    let r = (n - 1) >> i;

    let one = m.one();
    let minus_one = m.minus_one();

    'witness: for _a in A::BASIS.iter() {
        let _a = _a % n;
        if _a == 0 {
            continue;
        }

        let a = m.from_u64(_a);

        // x = a^r mod n
        let mut x = m.pow(a, r);

        if x == one || x == minus_one {
            continue;
        }

        for _ in 1..i {
            let y = m.mul(x, x);
            if y == one {
                return Composite(gcd(m.to_u64(x) - 1, m.modulus()));
            } else if y == minus_one {
                // This basis element is not a witness of `n` being composite.
                // Keep looking.
                continue 'witness;
            }
            x = y;
        }

        return Pseudoprime;
    }

    Prime
}

// Used by build.rs' tests and debug assertions
#[allow(dead_code)]
pub(crate) fn is_prime(n: u64) -> bool {
    if n < 2 {
        false
    } else if n % 2 == 0 {
        n == 2
    } else {
        test::<Montgomery<u64>>(Montgomery::new(n)).is_prime()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::numeric::{Arithmetic, Montgomery};
    use quickcheck::quickcheck;
    use std::iter;
    const LARGEST_U64_PRIME: u64 = 0xFFFFFFFFFFFFFFC5;

    fn primes() -> impl Iterator<Item = u64> {
        iter::once(2).chain(odd_primes())
    }

    fn odd_primes() -> impl Iterator<Item = u64> {
        use crate::table::{NEXT_PRIME, P_INVS_U64};
        P_INVS_U64
            .iter()
            .map(|(p, _, _)| *p)
            .chain(iter::once(NEXT_PRIME))
    }

    #[test]
    fn largest_prime() {
        assert!(is_prime(LARGEST_U64_PRIME));
    }

    #[test]
    fn largest_composites() {
        for i in LARGEST_U64_PRIME + 1..=u64::MAX {
            assert!(!is_prime(i), "2⁶⁴ - {} reported prime", u64::MAX - i + 1);
        }
    }

    #[test]
    fn two() {
        assert!(is_prime(2));
    }

    // TODO: Deduplicate with macro in numeric.rs
    macro_rules! parametrized_check {
        ( $f:ident ) => {
            paste::item! {
                #[test]
                fn [< $f _ u32 >]() {
                    $f::<Montgomery<u32>>()
                }
                #[test]
                fn [< $f _ u64 >]() {
                    $f::<Montgomery<u64>>()
                }
            }
        };
    }

    fn first_primes<A: Arithmetic + Basis>() {
        for p in odd_primes() {
            assert!(test(A::new(p)).is_prime(), "{} reported composite", p);
        }
    }
    parametrized_check!(first_primes);

    #[test]
    fn one() {
        assert!(!is_prime(1));
    }
    #[test]
    fn zero() {
        assert!(!is_prime(0));
    }

    fn first_composites<A: Arithmetic + Basis>() {
        for (p, q) in primes().zip(odd_primes()) {
            for i in p + 1..q {
                assert!(!is_prime(i), "{} reported prime", i);
            }
        }
    }
    parametrized_check!(first_composites);

    #[test]
    fn issue_1556() {
        // 10 425 511 = 2441 × 4271
        assert!(!is_prime(10_425_511));
    }

    fn small_semiprimes<A: Arithmetic + Basis>() {
        for p in odd_primes() {
            for q in odd_primes().take_while(|q| *q <= p) {
                let n = p * q;
                let m = A::new(n);
                assert!(!test(m).is_prime(), "{} = {} × {} reported prime", n, p, q);
            }
        }
    }
    parametrized_check!(small_semiprimes);

    quickcheck! {
        fn composites(i: u64, j: u64) -> bool {
            i < 2 || j < 2 || !is_prime(i*j)
        }
    }
}
<|endoftext|>


<commit_before>use std::slice;
use std::ops::{Add, Sub};
use std::f32::consts::PI;
use std::cmp::max;
use num::{Integer, Float};
use geom::V2;

/// Hex grid geometry for vectors.
pub trait HexGeom {
    /// Hex distance represented by a vector.
    fn hex_dist(&self) -> i32;
}

impl HexGeom for V2<i32> {
    fn hex_dist(&self) -> i32 {
        let xd = self.0;
        let yd = self.1;
        if xd.signum() == yd.signum() {
            max(xd.abs(), yd.abs())
        } else {
            xd.abs() + yd.abs()
        }
    }
}

/// Hex grid directions.
#[derive(Copy, Eq, PartialEq, Clone, Debug, RustcEncodable, RustcDecodable)]
pub enum Dir6 {
    North = 0,
    NorthEast,
    SouthEast,
    South,
    SouthWest,
    NorthWest,
}

impl Dir6 {
    /// Convert a vector into the closest hex direction.
    ///
    /// ```notrust
    ///        *0*       *1*
    ///           \ 14 15 | 00 01
    ///           13\     |      02
    ///               \   |
    ///         12      \ |        03
    ///     *5* ----------O-X------- *2*
    ///         11        Y \      04
    ///                   |   \
    ///           10      |     \05
    ///             09 08 | 07 06 \
    ///                  *4*       *3*
    ///
    /// The hexadecants (00 to 15) and the hex
    /// directions (*0* to *5*) around the origin.
    /// ```
    ///
    /// Vectors that are in a space between two hex direction vectors are
    /// rounded to a hexadecant, then assigned the hex direction whose vector
    /// is nearest to that hexadecant.
    pub fn from_v2(v: V2<i32>) -> Dir6 {
        let hexadecant = {
            let width = PI / 8.0;
            let mut radian = (v.0 as f32).atan2(-v.1 as f32);
            if radian < 0.0 { radian += 2.0 * PI }
            (radian / width).floor() as i32
        };

        Dir6::from_int(match hexadecant {
            13 | 14 => 0,
            15 | 0 | 1 => 1,
            2 | 3 | 4 => 2,
            5 | 6 => 3,
            7 | 8 | 9 => 4,
            10 | 11 | 12 => 5,
            _ => panic!("Bad hexadecant")
        })
    }

    /// Convert an integer to a hex dir using modular arithmetic.
    pub fn from_int(i: i32) -> Dir6 { DIRS[i.mod_floor(&6) as usize] }

    /// Convert a hex dir into the corresponding unit vector.
    pub fn to_v2(&self) -> V2<i32> {
        [V2(-1, -1),
         V2( 0, -1),
         V2( 1,  0),
         V2( 1,  1),
         V2( 0,  1),
         V2(-1,  0)][*self as usize]
    }

    /// Iterate through the six hex dirs in the standard order.
    pub fn iter() -> slice::Iter<'static, Dir6> { DIRS.iter() }
}

impl Add<i32> for Dir6 {
    type Output = Dir6;
    fn add(self, other: i32) -> Dir6 { Dir6::from_int(self as i32 + other) }
}

impl Sub<i32> for Dir6 {
    type Output = Dir6;
    fn sub(self, other: i32) -> Dir6 { Dir6::from_int(self as i32 - other) }
}

static DIRS: [Dir6; 6] = [
    Dir6::North,
    Dir6::NorthEast,
    Dir6::SouthEast,
    Dir6::South,
    Dir6::SouthWest,
    Dir6::NorthWest];

/// Field of view iterator for a hexagonal map.
///
/// Takes a function that
/// indicates opaque cells and yields the visible locations from around the
/// origin up to the given radius.
pub struct HexFov<F> {
    /// Predicate for whether a given point will block the field of view.
    is_opaque: F,
    range: u32,
    stack: Vec<Sector>,
    fake_isometric_hack: bool,
    /// Extra values generated by special cases.
    side_channel: Vec<V2<i32>>,
}

impl<F> HexFov<F>
    where F: Fn(V2<i32>) -> bool
{
    pub fn new(is_opaque: F, range: u32) -> HexFov<F> {
        // The origin position V2(0, 0) is a special case for the traversal
        // algorithm, but it's also always present, so instead of adding ugly
        // branches to the actual iterator, we just chain it in right here.
        let init_group = is_opaque(Dir6::from_int(0).to_v2());
        HexFov {
            is_opaque: is_opaque,
            range: range,
            stack: vec![Sector {
                begin: PolarPoint::new(0.0, 1),
                pt: PolarPoint::new(0.0, 1),
                end: PolarPoint::new(6.0, 1),
                group_opaque: init_group,
            }],
            fake_isometric_hack: false,
            // The FOV algorithm will not generate the origin point, so we use
            // the side channel to explicitly add it in the beginning.
            side_channel: vec![V2(0, 0)],
        }
    }

    /// Make wall tiles in acute corners visible when running the algorithm so
    /// that the complete wall rectangle of fake-isometric rooms will appear
    /// in the FOV.
    pub fn fake_isometric(mut self) -> HexFov<F> {
        self.fake_isometric_hack = true;
        self
    }
}

impl<F> Iterator for HexFov<F>
    where F: Fn(V2<i32>) -> bool
{
    type Item = V2<i32>;
    fn next(&mut self) -> Option<V2<i32>> {
        if let Some(ret) = self.side_channel.pop() {
            return Some(ret);
        }

        if let Some(mut current) = self.stack.pop() {
            if current.pt.is_below(current.end) {
                let pos = current.pt.to_v2();
                let current_opaque = (self.is_opaque)(pos);

                // Terrain opacity changed, branch out.
                if current_opaque != current.group_opaque {
                    // Add the rest of this sector with the new opacity.
                    self.stack.push(Sector {
                        begin: current.pt,
                        pt: current.pt,
                        end: current.end,
                        group_opaque: current_opaque,
                    });

                    // If this was a visible sector and we're below range, branch
                    // out further.
                    if !current.group_opaque && current.begin.radius < self.range {
                        self.stack.push(Sector {
                            begin: current.begin.further(),
                            pt: current.begin.further(),
                            end: current.pt.further(),
                            group_opaque: (self.is_opaque)(current.begin.further().to_v2()),
                        });
                    }
                    return self.next();
                }

                // Hack for making acute corner tiles of fake-isometric rooms
                // visible.

                // XXX: Side cells should only be visible with wallform tiles,
                // but the FOV API can't currently distinguish between
                // wallform and blockform FOV blockers.
                if self.fake_isometric_hack {
                    if let Some(side_pt) = current.pt.side_point() {
                        // Only do this if both the front tiles and the target
                        // tile are opaque.
                        let next = current.pt.next();
                        if next.is_below(current.end)
                            && current.group_opaque
                            && (self.is_opaque)(next.to_v2())
                            && (self.is_opaque)(side_pt)
                            && current.begin.radius < self.range {
                                self.side_channel.push(side_pt);
                        }
                    }
                }

                // Proceed along the current sector.
                current.pt = current.pt.next();
                self.stack.push(current);
                return Some(pos);
            } else {
                // Hit the end of the sector.

                // If this was a visible sector and we're below range, branch
                // out further.
                if !current.group_opaque && current.begin.radius < self.range {
                    self.stack.push(Sector {
                        begin: current.begin.further(),
                        pt: current.begin.further(),
                        end: current.end.further(),
                        group_opaque: (self.is_opaque)(current.begin.further().to_v2()),
                    });
                }

                self.next()
            }
        } else {
            None
        }
    }
}

struct Sector {
    /// Start point of current sector.
    begin: PolarPoint,
    /// Point currently being processed.
    pt: PolarPoint,
    /// End point of current sector.
    end: PolarPoint,
    /// Currently iterating through a sequence of opaque cells.
    group_opaque: bool,
}

/// Points on a hex circle expressed in polar coordinates.
#[derive(Copy, Clone, PartialEq)]
struct PolarPoint {
    pos: f32,
    radius: u32
}

impl PolarPoint {
    pub fn new(pos: f32, radius: u32) -> PolarPoint { PolarPoint { pos: pos, radius: radius } }
    /// Index of the discrete hex cell along the circle that corresponds to this point.
    fn winding_index(self) -> i32 { (self.pos + 0.5).floor() as i32 }

    pub fn is_below(self, other: PolarPoint) -> bool { self.winding_index() < other.end_index() }
    fn end_index(self) -> i32 { (self.pos + 0.5).ceil() as i32 }

    pub fn to_v2(self) -> V2<i32> {
        if self.radius == 0 { return V2(0, 0); }
        let index = self.winding_index();
        let sector = index.mod_floor(&(self.radius as i32 * 6)) / self.radius as i32;
        let offset = index.mod_floor(&(self.radius as i32));
        let rod = Dir6::from_int(sector).to_v2() * (self.radius as i32);
        let tangent = Dir6::from_int((sector + 2) % 6).to_v2() * offset;
        rod + tangent
    }

    /// If this point and the next point are adjacent vertically (along the xy
    /// axis), return a tuple of the point outside of the circle between the
    /// two points.
    ///
    /// This is a helper function for the FOV special case where acute corners
    /// of fake isometric rooms are marked visible even though strict hex FOV
    /// logic would keep them unseen.
    pub fn side_point(self) -> Option<V2<i32>> {
        let next = self.next();
        let V2(x1, y1) = self.to_v2();
        let V2(x2, y2) = next.to_v2();

        if x2 == x1 + 1 && y2 == y1 + 1 {
            // Going down the right rim.
            Some(V2(x1 + 1, y1))
        } else if x2 == x1 - 1 && y2 == y1 - 1 {
            // Going up the left rim.
            Some(V2(x1 - 1, y1))
        } else {
            None
        }
    }

    /// The point corresponding to this one on the hex circle with radius +1.
    pub fn further(self) -> PolarPoint {
        PolarPoint::new(
            self.pos * (self.radius + 1) as f32 / self.radius as f32,
            self.radius + 1)
    }

    /// The point next to this one along the hex circle.
    pub fn next(self) -> PolarPoint {
        PolarPoint::new((self.pos + 0.5).floor() + 0.5, self.radius)
    }
}

#[cfg(test)]
mod test {
    use geom::V2;
    // XXX: Why doesn't super::* work here?
    use super::{Dir6};
    use super::Dir6::*;

    #[test]
    fn test_dir6() {
        assert_eq!(North, Dir6::from_int(0));
        assert_eq!(NorthWest, Dir6::from_int(-1));
        assert_eq!(NorthWest, Dir6::from_int(5));
        assert_eq!(NorthEast, Dir6::from_int(1));

        assert_eq!(NorthEast, Dir6::from_v2(V2(20, -21)));
        assert_eq!(SouthEast, Dir6::from_v2(V2(20, -10)));
        assert_eq!(North, Dir6::from_v2(V2(-10, -10)));
        assert_eq!(South, Dir6::from_v2(V2(1, 1)));

        for i in 0..6 {
            let d = Dir6::from_int(i);
            let v = d.to_v2();
            let v1 = Dir6::from_int(i - 1).to_v2();
            let v2 = Dir6::from_int(i + 1).to_v2();

            // Test static iter
            assert_eq!(Some(d), Dir6::iter().nth(i as usize).map(|&x| x));

            // Test vector mapping.
            assert_eq!(d, Dir6::from_v2(v));

            // Test opposite dir vector mapping.
            assert_eq!(Dir6::from_int(i + 3), Dir6::from_v2(-v));

            // Test approximation of longer vectors.
            assert_eq!(d, Dir6::from_v2(v * 3));
            assert_eq!(d, Dir6::from_v2(v * 3 + v1));
            assert_eq!(d, Dir6::from_v2(v * 3 + v2));
        }
    }
}
<commit_msg>impl Rand for Dir6<commit_after>use std::slice;
use std::ops::{Add, Sub};
use std::f32::consts::PI;
use std::cmp::max;
use rand::{Rand, Rng};
use num::{Integer, Float};
use geom::V2;

/// Hex grid geometry for vectors.
pub trait HexGeom {
    /// Hex distance represented by a vector.
    fn hex_dist(&self) -> i32;
}

impl HexGeom for V2<i32> {
    fn hex_dist(&self) -> i32 {
        let xd = self.0;
        let yd = self.1;
        if xd.signum() == yd.signum() {
            max(xd.abs(), yd.abs())
        } else {
            xd.abs() + yd.abs()
        }
    }
}

/// Hex grid directions.
#[derive(Copy, Eq, PartialEq, Clone, Debug, RustcEncodable, RustcDecodable)]
pub enum Dir6 {
    North = 0,
    NorthEast,
    SouthEast,
    South,
    SouthWest,
    NorthWest,
}

impl Dir6 {
    /// Convert a vector into the closest hex direction.
    ///
    /// ```notrust
    ///        *0*       *1*
    ///           \ 14 15 | 00 01
    ///           13\     |      02
    ///               \   |
    ///         12      \ |        03
    ///     *5* ----------O-X------- *2*
    ///         11        Y \      04
    ///                   |   \
    ///           10      |     \05
    ///             09 08 | 07 06 \
    ///                  *4*       *3*
    ///
    /// The hexadecants (00 to 15) and the hex
    /// directions (*0* to *5*) around the origin.
    /// ```
    ///
    /// Vectors that are in a space between two hex direction vectors are
    /// rounded to a hexadecant, then assigned the hex direction whose vector
    /// is nearest to that hexadecant.
    pub fn from_v2(v: V2<i32>) -> Dir6 {
        let hexadecant = {
            let width = PI / 8.0;
            let mut radian = (v.0 as f32).atan2(-v.1 as f32);
            if radian < 0.0 { radian += 2.0 * PI }
            (radian / width).floor() as i32
        };

        Dir6::from_int(match hexadecant {
            13 | 14 => 0,
            15 | 0 | 1 => 1,
            2 | 3 | 4 => 2,
            5 | 6 => 3,
            7 | 8 | 9 => 4,
            10 | 11 | 12 => 5,
            _ => panic!("Bad hexadecant")
        })
    }

    /// Convert an integer to a hex dir using modular arithmetic.
    pub fn from_int(i: i32) -> Dir6 { DIRS[i.mod_floor(&6) as usize] }

    /// Convert a hex dir into the corresponding unit vector.
    pub fn to_v2(&self) -> V2<i32> {
        [V2(-1, -1),
         V2( 0, -1),
         V2( 1,  0),
         V2( 1,  1),
         V2( 0,  1),
         V2(-1,  0)][*self as usize]
    }

    /// Iterate through the six hex dirs in the standard order.
    pub fn iter() -> slice::Iter<'static, Dir6> { DIRS.iter() }
}

impl Add<i32> for Dir6 {
    type Output = Dir6;
    fn add(self, other: i32) -> Dir6 { Dir6::from_int(self as i32 + other) }
}

impl Sub<i32> for Dir6 {
    type Output = Dir6;
    fn sub(self, other: i32) -> Dir6 { Dir6::from_int(self as i32 - other) }
}

impl Rand for Dir6 {
    fn rand<R: Rng>(rng: &mut R) -> Dir6 { Dir6::from_int(rng.gen_range(0, 6)) }
}

static DIRS: [Dir6; 6] = [
    Dir6::North,
    Dir6::NorthEast,
    Dir6::SouthEast,
    Dir6::South,
    Dir6::SouthWest,
    Dir6::NorthWest];

/// Field of view iterator for a hexagonal map.
///
/// Takes a function that
/// indicates opaque cells and yields the visible locations from around the
/// origin up to the given radius.
pub struct HexFov<F> {
    /// Predicate for whether a given point will block the field of view.
    is_opaque: F,
    range: u32,
    stack: Vec<Sector>,
    fake_isometric_hack: bool,
    /// Extra values generated by special cases.
    side_channel: Vec<V2<i32>>,
}

impl<F> HexFov<F>
    where F: Fn(V2<i32>) -> bool
{
    pub fn new(is_opaque: F, range: u32) -> HexFov<F> {
        // The origin position V2(0, 0) is a special case for the traversal
        // algorithm, but it's also always present, so instead of adding ugly
        // branches to the actual iterator, we just chain it in right here.
        let init_group = is_opaque(Dir6::from_int(0).to_v2());
        HexFov {
            is_opaque: is_opaque,
            range: range,
            stack: vec![Sector {
                begin: PolarPoint::new(0.0, 1),
                pt: PolarPoint::new(0.0, 1),
                end: PolarPoint::new(6.0, 1),
                group_opaque: init_group,
            }],
            fake_isometric_hack: false,
            // The FOV algorithm will not generate the origin point, so we use
            // the side channel to explicitly add it in the beginning.
            side_channel: vec![V2(0, 0)],
        }
    }

    /// Make wall tiles in acute corners visible when running the algorithm so
    /// that the complete wall rectangle of fake-isometric rooms will appear
    /// in the FOV.
    pub fn fake_isometric(mut self) -> HexFov<F> {
        self.fake_isometric_hack = true;
        self
    }
}

impl<F> Iterator for HexFov<F>
    where F: Fn(V2<i32>) -> bool
{
    type Item = V2<i32>;
    fn next(&mut self) -> Option<V2<i32>> {
        if let Some(ret) = self.side_channel.pop() {
            return Some(ret);
        }

        if let Some(mut current) = self.stack.pop() {
            if current.pt.is_below(current.end) {
                let pos = current.pt.to_v2();
                let current_opaque = (self.is_opaque)(pos);

                // Terrain opacity changed, branch out.
                if current_opaque != current.group_opaque {
                    // Add the rest of this sector with the new opacity.
                    self.stack.push(Sector {
                        begin: current.pt,
                        pt: current.pt,
                        end: current.end,
                        group_opaque: current_opaque,
                    });

                    // If this was a visible sector and we're below range, branch
                    // out further.
                    if !current.group_opaque && current.begin.radius < self.range {
                        self.stack.push(Sector {
                            begin: current.begin.further(),
                            pt: current.begin.further(),
                            end: current.pt.further(),
                            group_opaque: (self.is_opaque)(current.begin.further().to_v2()),
                        });
                    }
                    return self.next();
                }

                // Hack for making acute corner tiles of fake-isometric rooms
                // visible.

                // XXX: Side cells should only be visible with wallform tiles,
                // but the FOV API can't currently distinguish between
                // wallform and blockform FOV blockers.
                if self.fake_isometric_hack {
                    if let Some(side_pt) = current.pt.side_point() {
                        // Only do this if both the front tiles and the target
                        // tile are opaque.
                        let next = current.pt.next();
                        if next.is_below(current.end)
                            && current.group_opaque
                            && (self.is_opaque)(next.to_v2())
                            && (self.is_opaque)(side_pt)
                            && current.begin.radius < self.range {
                                self.side_channel.push(side_pt);
                        }
                    }
                }

                // Proceed along the current sector.
                current.pt = current.pt.next();
                self.stack.push(current);
                return Some(pos);
            } else {
                // Hit the end of the sector.

                // If this was a visible sector and we're below range, branch
                // out further.
                if !current.group_opaque && current.begin.radius < self.range {
                    self.stack.push(Sector {
                        begin: current.begin.further(),
                        pt: current.begin.further(),
                        end: current.end.further(),
                        group_opaque: (self.is_opaque)(current.begin.further().to_v2()),
                    });
                }

                self.next()
            }
        } else {
            None
        }
    }
}

struct Sector {
    /// Start point of current sector.
    begin: PolarPoint,
    /// Point currently being processed.
    pt: PolarPoint,
    /// End point of current sector.
    end: PolarPoint,
    /// Currently iterating through a sequence of opaque cells.
    group_opaque: bool,
}

/// Points on a hex circle expressed in polar coordinates.
#[derive(Copy, Clone, PartialEq)]
struct PolarPoint {
    pos: f32,
    radius: u32
}

impl PolarPoint {
    pub fn new(pos: f32, radius: u32) -> PolarPoint { PolarPoint { pos: pos, radius: radius } }
    /// Index of the discrete hex cell along the circle that corresponds to this point.
    fn winding_index(self) -> i32 { (self.pos + 0.5).floor() as i32 }

    pub fn is_below(self, other: PolarPoint) -> bool { self.winding_index() < other.end_index() }
    fn end_index(self) -> i32 { (self.pos + 0.5).ceil() as i32 }

    pub fn to_v2(self) -> V2<i32> {
        if self.radius == 0 { return V2(0, 0); }
        let index = self.winding_index();
        let sector = index.mod_floor(&(self.radius as i32 * 6)) / self.radius as i32;
        let offset = index.mod_floor(&(self.radius as i32));
        let rod = Dir6::from_int(sector).to_v2() * (self.radius as i32);
        let tangent = Dir6::from_int((sector + 2) % 6).to_v2() * offset;
        rod + tangent
    }

    /// If this point and the next point are adjacent vertically (along the xy
    /// axis), return a tuple of the point outside of the circle between the
    /// two points.
    ///
    /// This is a helper function for the FOV special case where acute corners
    /// of fake isometric rooms are marked visible even though strict hex FOV
    /// logic would keep them unseen.
    pub fn side_point(self) -> Option<V2<i32>> {
        let next = self.next();
        let V2(x1, y1) = self.to_v2();
        let V2(x2, y2) = next.to_v2();

        if x2 == x1 + 1 && y2 == y1 + 1 {
            // Going down the right rim.
            Some(V2(x1 + 1, y1))
        } else if x2 == x1 - 1 && y2 == y1 - 1 {
            // Going up the left rim.
            Some(V2(x1 - 1, y1))
        } else {
            None
        }
    }

    /// The point corresponding to this one on the hex circle with radius +1.
    pub fn further(self) -> PolarPoint {
        PolarPoint::new(
            self.pos * (self.radius + 1) as f32 / self.radius as f32,
            self.radius + 1)
    }

    /// The point next to this one along the hex circle.
    pub fn next(self) -> PolarPoint {
        PolarPoint::new((self.pos + 0.5).floor() + 0.5, self.radius)
    }
}

#[cfg(test)]
mod test {
    use geom::V2;
    // XXX: Why doesn't super::* work here?
    use super::{Dir6};
    use super::Dir6::*;

    #[test]
    fn test_dir6() {
        assert_eq!(North, Dir6::from_int(0));
        assert_eq!(NorthWest, Dir6::from_int(-1));
        assert_eq!(NorthWest, Dir6::from_int(5));
        assert_eq!(NorthEast, Dir6::from_int(1));

        assert_eq!(NorthEast, Dir6::from_v2(V2(20, -21)));
        assert_eq!(SouthEast, Dir6::from_v2(V2(20, -10)));
        assert_eq!(North, Dir6::from_v2(V2(-10, -10)));
        assert_eq!(South, Dir6::from_v2(V2(1, 1)));

        for i in 0..6 {
            let d = Dir6::from_int(i);
            let v = d.to_v2();
            let v1 = Dir6::from_int(i - 1).to_v2();
            let v2 = Dir6::from_int(i + 1).to_v2();

            // Test static iter
            assert_eq!(Some(d), Dir6::iter().nth(i as usize).map(|&x| x));

            // Test vector mapping.
            assert_eq!(d, Dir6::from_v2(v));

            // Test opposite dir vector mapping.
            assert_eq!(Dir6::from_int(i + 3), Dir6::from_v2(-v));

            // Test approximation of longer vectors.
            assert_eq!(d, Dir6::from_v2(v * 3));
            assert_eq!(d, Dir6::from_v2(v * 3 + v1));
            assert_eq!(d, Dir6::from_v2(v * 3 + v2));
        }
    }
}
<|endoftext|>


<commit_before>
//
// FSUIPC library
// Copyright (c) 2015 Alvaro Polo
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::io;
use std::io::{Read, Write};

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

use super::raw::RawBytes;

/// The header of a message sent to FSUIPC module via IPC
#[allow(raw_pointer_derive)]
#[derive(Debug, PartialEq)]
pub enum MsgHeader {
    /// Read state data message header
    /// Read `len` bytes from given offset and prepar to store `data` in `target`.
    ReadStateData {
        offset: u16,
        len: usize,
        target: *mut u8,
    },
    /// Write state data message header
    /// Write `len` bytes from given `source` to given offset.
    WriteStateData {
        offset: u16,
        len: usize,
    },
    TerminationMark
}

pub trait MsgRead : Read {
    /// Read a IPC message header from the given `Read` object.
    /// It returns the read message header and the number of bytes processed.
    fn read_header(&mut self) -> io::Result<(MsgHeader, usize)> {
        match try!(self.read_u32::<LittleEndian>()) {
            FS6IPC_READSTATEDATA_ID => {
                let offset = try!(self.read_u32::<LittleEndian>()) as u16;
                let len = try!(self.read_u32::<LittleEndian>()) as usize;
                let target = try!(self.read_u32::<LittleEndian>()) as *mut u8;
                Ok((MsgHeader::ReadStateData {
                    offset: offset,
                    len: len,
                    target: target,
                }, 16))
            },
            FS6IPC_WRITESTATEDATA_ID => {
                let offset = try!(self.read_u32::<LittleEndian>()) as u16;
                let len = try!(self.read_u32::<LittleEndian>()) as usize;
                Ok((MsgHeader::WriteStateData {
                    offset: offset,
                    len: len,
                }, 12))
            },
            FS6IPC_TERMINATIONMARK_ID => return Ok((MsgHeader::TerminationMark, 4)),
            unexpected => Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("unexpected double word 0x{} while reading IPC message header",
                    unexpected))),
        }
    }

    fn read_body<W: Write>(&mut self, header: &MsgHeader, output: &mut W) -> io::Result<usize> {
        match header {
            &MsgHeader::ReadStateData { offset: _, len, target: _ } => {
                for _ in 0..len { try!(output.write_u8(try!(self.read_u8()))); }
                Ok(len)
            },
            &MsgHeader::WriteStateData { offset: _, len } => {
                for _ in 0..len { try!(output.write_u8(try!(self.read_u8()))); }
                Ok(len)
            },
            &MsgHeader::TerminationMark => Ok(0),
        }
    }
}

impl<R: Read + ?Sized> MsgRead for R {}

pub trait MsgWrite : Write {
    /// Write a IPC message header into the given `Write` object.
    fn write_header(&mut self, msg: &MsgHeader) -> io::Result<usize> {
        match msg {
            &MsgHeader::ReadStateData { offset, len, target } => {
                try!(self.write_u32::<LittleEndian>(FS6IPC_READSTATEDATA_ID));
                try!(self.write_u32::<LittleEndian>(offset as u32));
                try!(self.write_u32::<LittleEndian>(len as u32));
                try!(self.write_u32::<LittleEndian>(target as u32));
                Ok(16)
            },
            &MsgHeader::WriteStateData { offset, len } => {
                try!(self.write_u32::<LittleEndian>(FS6IPC_WRITESTATEDATA_ID));
                try!(self.write_u32::<LittleEndian>(offset as u32));
                try!(self.write_u32::<LittleEndian>(len as u32));
                Ok(12)
            },
            &MsgHeader::TerminationMark => {
                try!(self.write_u32::<LittleEndian>(FS6IPC_TERMINATIONMARK_ID));
                Ok(4)
            },
        }
    }

    fn write_body<R: Read>(&mut self, header: &MsgHeader, input: &mut R) -> io::Result<usize> {
        match header {
            &MsgHeader::ReadStateData { offset: _, len, target: _ } => {
                for _ in 0..len { try!(self.write_u8(try!(input.read_u8()))); }
                Ok(len)
            },
            &MsgHeader::WriteStateData { offset: _, len } => {
                for _ in 0..len { try!(self.write_u8(try!(input.read_u8()))); }
                Ok(len)
            },
            &MsgHeader::TerminationMark => Ok(0),
        }
    }

    fn write_rsd(&mut self, offset: u16, dest: *mut u8, len: usize) -> io::Result<usize> {
        let header = MsgHeader::ReadStateData {
            offset: offset, len: len, target: dest,
        };
        let hdr_bytes = try!(self.write_header(&header));
        let body_bytes = try!(self.write_body(&header, &mut io::repeat(0)));
        Ok(hdr_bytes + body_bytes)
    }

    fn write_wsd(&mut self, offset: u16, src: *const u8, len: usize) -> io::Result<usize> {
        let header = MsgHeader::WriteStateData {
            offset: offset, len: len,
        };
        let hdr_bytes = try!(self.write_header(&header));
        let body_bytes = try!(self.write_body(&header, &mut RawBytes::new(src, len)));
        Ok(hdr_bytes + body_bytes)
    }
}

impl<W: Write + ?Sized> MsgWrite for W {}

const FS6IPC_TERMINATIONMARK_ID: u32 = 0;
const FS6IPC_READSTATEDATA_ID: u32 = 1;
const FS6IPC_WRITESTATEDATA_ID: u32 = 2;

#[cfg(test)]
mod test {

    use std::io::{Cursor, ErrorKind};

    use byteorder::{LittleEndian, ReadBytesExt};

    use super::*;

    #[test]
    fn should_read_rsd_header() {
        let mut buff: &[u8] = &[
            0x01, 0x00, 0x00, 0x00,
            0x00, 0x10, 0x00, 0x00,
            0x04, 0x00, 0x00, 0x00,
            0x00, 0x20, 0x00, 0x00,
        ];
        let expected = MsgHeader::ReadStateData {
            offset: 0x1000,
            len: 4,
            target: 0x2000 as *mut u8,
        };
        assert_eq!(buff.read_header().unwrap(), (expected, 16));
    }

    #[test]
    fn should_read_rsd_body() {
        let mut buff: &[u8] = &[ 0x01, 0x02, 0x03, 0x04 ];
        let header = MsgHeader::ReadStateData {
            offset: 0x1000,
            len: 4,
            target: 0x2000 as *mut u8,
        };
        let mut data = Vec::new();
        assert_eq!(buff.read_body(&header, &mut data).unwrap(), 4);
        assert_eq!(data[0], 1);
        assert_eq!(data[1], 2);
        assert_eq!(data[2], 3);
        assert_eq!(data[3], 4);
    }

    #[test]
    fn should_read_wsd_header() {
        let mut buff: &[u8] = &[
            0x02, 0x00, 0x00, 0x00,
            0x00, 0x10, 0x00, 0x00,
            0x04, 0x00, 0x00, 0x00,
        ];
        let expected = MsgHeader::WriteStateData {
            offset: 0x1000,
            len: 4,
        };
        assert_eq!(buff.read_header().unwrap(), (expected, 12));
    }

    #[test]
    fn should_read_wsd_body() {
        let mut buff: &[u8] = &[ 0x01, 0x02, 0x03, 0x04 ];
        let header = MsgHeader::WriteStateData {
            offset: 0x1000,
            len: 4,
        };
        let mut data = Vec::new();
        assert_eq!(buff.read_body(&header, &mut data).unwrap(), 4);
        assert_eq!(data[0], 1);
        assert_eq!(data[1], 2);
        assert_eq!(data[2], 3);
        assert_eq!(data[3], 4);
    }

    #[test]
    fn should_read_tm_header() {
        let mut buff: &[u8] = &[0x00, 0x00, 0x00, 0x00];
        assert_eq!(buff.read_header().unwrap(), ( MsgHeader::TerminationMark, 4));
    }

    #[test]
    fn should_read_tm_body() {
        let mut buff: &[u8] = &[ 0x01, 0x02, 0x03, 0x04 ];
        let mut data = Vec::new();
        assert_eq!(buff.read_body(&MsgHeader::TerminationMark, &mut data).unwrap(), 0);
        assert_eq!(data.len(), 0);
    }

    #[test]
    fn should_fail_to_read_from_invalid_stream() {
        let mut buff: &[u8] = &[0x01, 0x02, 0x03, 0x04];
        let expected_error = ErrorKind::InvalidInput;
        let actual_error = buff.read_header().err().unwrap().kind();
        assert_eq!(actual_error, expected_error);
    }

    #[test]
    fn should_write_rsd_header() {
        let mut buff = Cursor::new(Vec::new());
        let msg = MsgHeader::ReadStateData {
            offset: 0x1000,
            len: 4,
            target: 0x2000 as *mut u8,
        };
        assert_eq!(buff.write_header(&msg).unwrap(), 16);
        buff.set_position(0);
        assert_eq!(buff.get_ref().len(), 16);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 1);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x1000);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 4);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x2000);
    }

    #[test]
    fn should_write_rsd_body() {
        let mut buff = Vec::new();
        let mut input = Cursor::new(vec![ 0x01u8, 0x02, 0x03, 0x04 ]);
        let msg = MsgHeader::ReadStateData {
            offset: 0x1000,
            len: 4,
            target: 0x2000 as *mut u8,
        };

        assert_eq!(buff.write_body(&msg, &mut input).unwrap(), 4);
        assert_eq!(buff[0], 1);
        assert_eq!(buff[1], 2);
        assert_eq!(buff[2], 3);
        assert_eq!(buff[3], 4);
    }

    #[test]
    fn should_write_rsd() {
        let mut buff = Cursor::new(Vec::new());
        assert_eq!(buff.write_rsd(0x1000, 0x2000 as *mut u8, 4).unwrap(), 20);
        buff.set_position(0);
        assert_eq!(buff.get_ref().len(), 20);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 1);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x1000);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 4);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x2000);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0);
    }

    #[test]
    fn should_write_wsd_header() {
        let mut buff = Cursor::new(Vec::new());
        let msg = MsgHeader::WriteStateData {
            offset: 0x1000,
            len: 4,
        };
        assert_eq!(buff.write_header(&msg).unwrap(), 12);
        buff.set_position(0);
        assert_eq!(buff.get_ref().len(), 12);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 2);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x1000);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 4);
    }

    #[test]
    fn should_write_wsd_body() {
        let mut buff = Vec::new();
        let mut input = Cursor::new(vec![ 0x01u8, 0x02, 0x03, 0x04 ]);
        let msg = MsgHeader::WriteStateData {
            offset: 0x1000,
            len: 4,
        };

        assert_eq!(buff.write_body(&msg, &mut input).unwrap(), 4);
        assert_eq!(buff[0], 1);
        assert_eq!(buff[1], 2);
        assert_eq!(buff[2], 3);
        assert_eq!(buff[3], 4);
    }

    #[test]
    fn should_write_wsd() {
        let mut buff = Cursor::new(Vec::new());
        let data = &0x01020304u32 as *const u32 as *const u8;
        assert_eq!(buff.write_wsd(0x1000, data, 4).unwrap(), 16);
        buff.set_position(0);
        assert_eq!(buff.get_ref().len(), 16);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 2);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x1000);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 4);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x01020304);
    }

    #[test]
    fn should_write_tm_header() {
        let mut buff = Cursor::new(Vec::new());
        assert_eq!(buff.write_header(&MsgHeader::TerminationMark).unwrap(),4);
        buff.set_position(0);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0);
    }

    #[test]
    fn should_write_tm_body() {
        let mut buff = Vec::new();
        let mut input = Cursor::new(vec![ 0x01u8, 0x02, 0x03, 0x04 ]);
        assert_eq!(buff.write_body(&MsgHeader::TerminationMark, &mut input).unwrap(), 0);
        assert_eq!(buff.len(), 0);
    }
}
<commit_msg>Fix invalid error kind while reading hdr from invalid bytes<commit_after>
//
// FSUIPC library
// Copyright (c) 2015 Alvaro Polo
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::io;
use std::io::{Read, Write};

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

use super::raw::RawBytes;

/// The header of a message sent to FSUIPC module via IPC
#[allow(raw_pointer_derive)]
#[derive(Debug, PartialEq)]
pub enum MsgHeader {
    /// Read state data message header
    /// Read `len` bytes from given offset and prepar to store `data` in `target`.
    ReadStateData {
        offset: u16,
        len: usize,
        target: *mut u8,
    },
    /// Write state data message header
    /// Write `len` bytes from given `source` to given offset.
    WriteStateData {
        offset: u16,
        len: usize,
    },
    TerminationMark
}

pub trait MsgRead : Read {
    /// Read a IPC message header from the given `Read` object.
    /// It returns the read message header and the number of bytes processed.
    fn read_header(&mut self) -> io::Result<(MsgHeader, usize)> {
        match try!(self.read_u32::<LittleEndian>()) {
            FS6IPC_READSTATEDATA_ID => {
                let offset = try!(self.read_u32::<LittleEndian>()) as u16;
                let len = try!(self.read_u32::<LittleEndian>()) as usize;
                let target = try!(self.read_u32::<LittleEndian>()) as *mut u8;
                Ok((MsgHeader::ReadStateData {
                    offset: offset,
                    len: len,
                    target: target,
                }, 16))
            },
            FS6IPC_WRITESTATEDATA_ID => {
                let offset = try!(self.read_u32::<LittleEndian>()) as u16;
                let len = try!(self.read_u32::<LittleEndian>()) as usize;
                Ok((MsgHeader::WriteStateData {
                    offset: offset,
                    len: len,
                }, 12))
            },
            FS6IPC_TERMINATIONMARK_ID => return Ok((MsgHeader::TerminationMark, 4)),
            unexpected => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("unexpected double word 0x{} while reading IPC message header",
                    unexpected))),
        }
    }

    fn read_body<W: Write>(&mut self, header: &MsgHeader, output: &mut W) -> io::Result<usize> {
        match header {
            &MsgHeader::ReadStateData { offset: _, len, target: _ } => {
                for _ in 0..len { try!(output.write_u8(try!(self.read_u8()))); }
                Ok(len)
            },
            &MsgHeader::WriteStateData { offset: _, len } => {
                for _ in 0..len { try!(output.write_u8(try!(self.read_u8()))); }
                Ok(len)
            },
            &MsgHeader::TerminationMark => Ok(0),
        }
    }
}

impl<R: Read + ?Sized> MsgRead for R {}

pub trait MsgWrite : Write {
    /// Write a IPC message header into the given `Write` object.
    fn write_header(&mut self, msg: &MsgHeader) -> io::Result<usize> {
        match msg {
            &MsgHeader::ReadStateData { offset, len, target } => {
                try!(self.write_u32::<LittleEndian>(FS6IPC_READSTATEDATA_ID));
                try!(self.write_u32::<LittleEndian>(offset as u32));
                try!(self.write_u32::<LittleEndian>(len as u32));
                try!(self.write_u32::<LittleEndian>(target as u32));
                Ok(16)
            },
            &MsgHeader::WriteStateData { offset, len } => {
                try!(self.write_u32::<LittleEndian>(FS6IPC_WRITESTATEDATA_ID));
                try!(self.write_u32::<LittleEndian>(offset as u32));
                try!(self.write_u32::<LittleEndian>(len as u32));
                Ok(12)
            },
            &MsgHeader::TerminationMark => {
                try!(self.write_u32::<LittleEndian>(FS6IPC_TERMINATIONMARK_ID));
                Ok(4)
            },
        }
    }

    fn write_body<R: Read>(&mut self, header: &MsgHeader, input: &mut R) -> io::Result<usize> {
        match header {
            &MsgHeader::ReadStateData { offset: _, len, target: _ } => {
                for _ in 0..len { try!(self.write_u8(try!(input.read_u8()))); }
                Ok(len)
            },
            &MsgHeader::WriteStateData { offset: _, len } => {
                for _ in 0..len { try!(self.write_u8(try!(input.read_u8()))); }
                Ok(len)
            },
            &MsgHeader::TerminationMark => Ok(0),
        }
    }

    fn write_rsd(&mut self, offset: u16, dest: *mut u8, len: usize) -> io::Result<usize> {
        let header = MsgHeader::ReadStateData {
            offset: offset, len: len, target: dest,
        };
        let hdr_bytes = try!(self.write_header(&header));
        let body_bytes = try!(self.write_body(&header, &mut io::repeat(0)));
        Ok(hdr_bytes + body_bytes)
    }

    fn write_wsd(&mut self, offset: u16, src: *const u8, len: usize) -> io::Result<usize> {
        let header = MsgHeader::WriteStateData {
            offset: offset, len: len,
        };
        let hdr_bytes = try!(self.write_header(&header));
        let body_bytes = try!(self.write_body(&header, &mut RawBytes::new(src, len)));
        Ok(hdr_bytes + body_bytes)
    }
}

impl<W: Write + ?Sized> MsgWrite for W {}

const FS6IPC_TERMINATIONMARK_ID: u32 = 0;
const FS6IPC_READSTATEDATA_ID: u32 = 1;
const FS6IPC_WRITESTATEDATA_ID: u32 = 2;

#[cfg(test)]
mod test {

    use std::io::{Cursor, ErrorKind};

    use byteorder::{LittleEndian, ReadBytesExt};

    use super::*;

    #[test]
    fn should_read_rsd_header() {
        let mut buff: &[u8] = &[
            0x01, 0x00, 0x00, 0x00,
            0x00, 0x10, 0x00, 0x00,
            0x04, 0x00, 0x00, 0x00,
            0x00, 0x20, 0x00, 0x00,
        ];
        let expected = MsgHeader::ReadStateData {
            offset: 0x1000,
            len: 4,
            target: 0x2000 as *mut u8,
        };
        assert_eq!(buff.read_header().unwrap(), (expected, 16));
    }

    #[test]
    fn should_read_rsd_body() {
        let mut buff: &[u8] = &[ 0x01, 0x02, 0x03, 0x04 ];
        let header = MsgHeader::ReadStateData {
            offset: 0x1000,
            len: 4,
            target: 0x2000 as *mut u8,
        };
        let mut data = Vec::new();
        assert_eq!(buff.read_body(&header, &mut data).unwrap(), 4);
        assert_eq!(data[0], 1);
        assert_eq!(data[1], 2);
        assert_eq!(data[2], 3);
        assert_eq!(data[3], 4);
    }

    #[test]
    fn should_read_wsd_header() {
        let mut buff: &[u8] = &[
            0x02, 0x00, 0x00, 0x00,
            0x00, 0x10, 0x00, 0x00,
            0x04, 0x00, 0x00, 0x00,
        ];
        let expected = MsgHeader::WriteStateData {
            offset: 0x1000,
            len: 4,
        };
        assert_eq!(buff.read_header().unwrap(), (expected, 12));
    }

    #[test]
    fn should_read_wsd_body() {
        let mut buff: &[u8] = &[ 0x01, 0x02, 0x03, 0x04 ];
        let header = MsgHeader::WriteStateData {
            offset: 0x1000,
            len: 4,
        };
        let mut data = Vec::new();
        assert_eq!(buff.read_body(&header, &mut data).unwrap(), 4);
        assert_eq!(data[0], 1);
        assert_eq!(data[1], 2);
        assert_eq!(data[2], 3);
        assert_eq!(data[3], 4);
    }

    #[test]
    fn should_read_tm_header() {
        let mut buff: &[u8] = &[0x00, 0x00, 0x00, 0x00];
        assert_eq!(buff.read_header().unwrap(), ( MsgHeader::TerminationMark, 4));
    }

    #[test]
    fn should_read_tm_body() {
        let mut buff: &[u8] = &[ 0x01, 0x02, 0x03, 0x04 ];
        let mut data = Vec::new();
        assert_eq!(buff.read_body(&MsgHeader::TerminationMark, &mut data).unwrap(), 0);
        assert_eq!(data.len(), 0);
    }

    #[test]
    fn should_fail_to_read_from_invalid_stream() {
        let mut buff: &[u8] = &[0x01, 0x02, 0x03, 0x04];
        let expected_error = ErrorKind::InvalidData;
        let actual_error = buff.read_header().err().unwrap().kind();
        assert_eq!(actual_error, expected_error);
    }

    #[test]
    fn should_write_rsd_header() {
        let mut buff = Cursor::new(Vec::new());
        let msg = MsgHeader::ReadStateData {
            offset: 0x1000,
            len: 4,
            target: 0x2000 as *mut u8,
        };
        assert_eq!(buff.write_header(&msg).unwrap(), 16);
        buff.set_position(0);
        assert_eq!(buff.get_ref().len(), 16);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 1);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x1000);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 4);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x2000);
    }

    #[test]
    fn should_write_rsd_body() {
        let mut buff = Vec::new();
        let mut input = Cursor::new(vec![ 0x01u8, 0x02, 0x03, 0x04 ]);
        let msg = MsgHeader::ReadStateData {
            offset: 0x1000,
            len: 4,
            target: 0x2000 as *mut u8,
        };

        assert_eq!(buff.write_body(&msg, &mut input).unwrap(), 4);
        assert_eq!(buff[0], 1);
        assert_eq!(buff[1], 2);
        assert_eq!(buff[2], 3);
        assert_eq!(buff[3], 4);
    }

    #[test]
    fn should_write_rsd() {
        let mut buff = Cursor::new(Vec::new());
        assert_eq!(buff.write_rsd(0x1000, 0x2000 as *mut u8, 4).unwrap(), 20);
        buff.set_position(0);
        assert_eq!(buff.get_ref().len(), 20);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 1);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x1000);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 4);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x2000);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0);
    }

    #[test]
    fn should_write_wsd_header() {
        let mut buff = Cursor::new(Vec::new());
        let msg = MsgHeader::WriteStateData {
            offset: 0x1000,
            len: 4,
        };
        assert_eq!(buff.write_header(&msg).unwrap(), 12);
        buff.set_position(0);
        assert_eq!(buff.get_ref().len(), 12);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 2);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x1000);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 4);
    }

    #[test]
    fn should_write_wsd_body() {
        let mut buff = Vec::new();
        let mut input = Cursor::new(vec![ 0x01u8, 0x02, 0x03, 0x04 ]);
        let msg = MsgHeader::WriteStateData {
            offset: 0x1000,
            len: 4,
        };

        assert_eq!(buff.write_body(&msg, &mut input).unwrap(), 4);
        assert_eq!(buff[0], 1);
        assert_eq!(buff[1], 2);
        assert_eq!(buff[2], 3);
        assert_eq!(buff[3], 4);
    }

    #[test]
    fn should_write_wsd() {
        let mut buff = Cursor::new(Vec::new());
        let data = &0x01020304u32 as *const u32 as *const u8;
        assert_eq!(buff.write_wsd(0x1000, data, 4).unwrap(), 16);
        buff.set_position(0);
        assert_eq!(buff.get_ref().len(), 16);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 2);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x1000);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 4);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0x01020304);
    }

    #[test]
    fn should_write_tm_header() {
        let mut buff = Cursor::new(Vec::new());
        assert_eq!(buff.write_header(&MsgHeader::TerminationMark).unwrap(),4);
        buff.set_position(0);
        assert_eq!(buff.read_u32::<LittleEndian>().unwrap(), 0);
    }

    #[test]
    fn should_write_tm_body() {
        let mut buff = Vec::new();
        let mut input = Cursor::new(vec![ 0x01u8, 0x02, 0x03, 0x04 ]);
        assert_eq!(buff.write_body(&MsgHeader::TerminationMark, &mut input).unwrap(), 0);
        assert_eq!(buff.len(), 0);
    }
}
<|endoftext|>


<commit_before>/*
Tristen Horton
tristen@tristenhorton.com
2017-04-03 23:50
*/
<commit_msg>added some irc parsing functions that determine some status about an irc message<commit_after>/*
Tristen Horton
tristen@tristenhorton.com
2017-04-03 23:50
*/

use channel;
use regex::Regex;

pub fn detect_connected(message: &String) -> bool {

	match message.split(":").nth(2) {
		Some("You are in a maze of twisty passages, all alike.") => true,
		_ => false
	}
}

pub fn detect_new_subscriber(
	irc: &String,
	channels: &Vec<channel::Channel>
) -> Option<String> {

	for chan in channels {
		match read_type(irc, chan.name()) {
			Some(rtype) => {

				if irc.split("!").nth(0).unwrap() == ":twitchnotify" &&
					rtype == "PRIVMSG" &&
					(irc.contains("just subscribed!") ||
						irc.to_lowercase().contains("just subscriber with twitch prime!")
					)
				{
					return Some(chan.name().clone());
				}
			}
			_ => {}
		}
	}

	None
}

pub fn read_type(irc: &String, channel: &String) -> Option<String> {

	if irc.contains(" ") {

		let mut found = false;
		for word in irc.split(" ") {
			if word.chars().nth(0).unwrap() == '#' {
				if word == format!("#{}", channel) {
					found = true;
				}
			}
		}
		if found {
			// TODO: Clean this up because this is messy as heck
			let reg = Regex::new(format!(" #{}", channel).as_ref()).unwrap();
			let mut split_a = reg.split(irc);
			let mut split_a_0 = split_a.nth(0).unwrap().to_owned();
			let mut type_wrap = split_a_0.split(" ");
			let mut type_copy = type_wrap.clone();
			let wrap_count = type_wrap.count();
			let ret = type_copy.nth(wrap_count - 1).unwrap().to_owned();
			return Some(ret);
		}
		else {
			if irc.split(" ").count() > 1 && irc.split(" ").nth(1).unwrap() == String::from("NOTICE") {
				return Some(String::from("NOTICE"));
			}
		}
	}

	None
}
<|endoftext|>


<commit_before>pub struct RepeatingKey {
    key: Vec<u8>,
    iter_val: u16,
    key_size: u16,
}

impl Iterator for RepeatingKey {
	type Item = u8;

    fn next(&mut self) -> Option<u8> {
		let result = Some(self.key[self.iter_val as usize]);
    	self.iter_val = (self.iter_val + 1) % self.key_size;

    	result
    }
}

impl RepeatingKey {
	pub fn new_bytes(key_bytes: Vec<u8>) -> RepeatingKey {
		let key_size: u16 = key_bytes.len() as u16;

		RepeatingKey {
			key: key_bytes,
			iter_val: 0,
			key_size: key_size,
		}
	}

	pub fn new(key: &str) -> RepeatingKey {
		let key_bytes: Vec<u8> = key.bytes().collect();
		RepeatingKey::new_bytes(key_bytes)
	}

	pub fn of_length(&mut self, length: usize) -> Vec<u8> {
		let mut result = vec![];
		for x in self.take(length) {
			result.push(x);
		}
		self.iter_val = 0;
		result
	}

	pub fn encrypt_bytes(&mut self, bytes: &Vec<u8>) -> Vec<u8> {
		use byte_manipulation::*;

		let length = bytes.len();

		xor_byte_streams(&bytes, &self.of_length(length))
	}
}

//-----------------------------------------------------------------------------

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

	#[test]
	fn RepeatingKey_short() {
		let rk = RepeatingKey::new(&"ICE");
		let mut actual: Vec<u8> = vec![];

		for x in rk.take(2) {
			actual.push(x);
		}

		let expected: Vec<u8> = "IC".to_string().bytes().collect();

		assert_eq!(expected, actual);
	}

	#[test]
	fn RepeatingKey_long() {
		let rk = RepeatingKey::new(&"ICE");
		let mut actual: Vec<u8> = vec![];

		for x in rk.take(7) {
			actual.push(x);
		}

		let expected: Vec<u8> = "ICEICEI".to_string().bytes().collect();

		assert_eq!(expected, actual);
	}

	#[test]
	fn RepeatingKey_problem5() {
		// http://cryptopals.com/sets/1/challenges/5/
		use byte_conversion::*;

		let bytes = readable_text_to_bytes(&"Burning 'em, if you ain't quick and nimble\nI go crazy when I hear a cymbal");
		let output = RepeatingKey::new(&"ICE").encrypt_bytes(&bytes);

		let expected = hex_to_bytes(&"0b3637272a2b2e63622c2e69692a23693a2a3c6324202d623d63343c2a26226324272765272a282b2f20430a652e2c652a3124333a653e2b2027630c692b20283165286326302e27282f");

		assert_eq!(expected, output);
	}

	#[test]
	fn RepeatingKey_of_length() {
		let actual = RepeatingKey::new(&"ICE").of_length(7);
		let expected = vec![73, 67, 69, 73, 67, 69, 73];

		assert_eq!(expected, actual);
	}

	#[test]
	fn RepeatingKey_encrypt_bytes() {
		use byte_conversion::*;

		let bytes = readable_text_to_bytes(&"Burning 'em, if you ain't quick and nimble\nI go crazy when I hear a cymbal");
		let output = RepeatingKey::new(&"ICE").encrypt_bytes(&bytes);

		let expected = hex_to_bytes(&"0b3637272a2b2e63622c2e69692a23693a2a3c6324202d623d63343c2a26226324272765272a282b2f20430a652e2c652a3124333a653e2b2027630c692b20283165286326302e27282f");

		assert_eq!(expected, output);
	}

	#[test]
	fn RepeatingKey_encrypt_bytes_twice() {
		use byte_conversion::*;

		let word = "Burning 'em, if you ain't quick and nimble\nI go crazy when I hear a cymbal";

		let bytes = readable_text_to_bytes(&word);
		let mut key = RepeatingKey::new(&"ICE");

		let encrypted = key.encrypt_bytes(&bytes);
		let decrypted = key.encrypt_bytes(&encrypted);

		assert_eq!(bytes, decrypted);
	}
}<commit_msg>RepeatingKey takes &Vec<commit_after>pub struct RepeatingKey {
    key: Vec<u8>,
    iter_val: u16,
    key_size: u16,
}

impl Iterator for RepeatingKey {
	type Item = u8;

    fn next(&mut self) -> Option<u8> {
		let result = Some(self.key[self.iter_val as usize]);
    	self.iter_val = (self.iter_val + 1) % self.key_size;

    	result
    }
}

impl RepeatingKey {
	pub fn new_bytes(key_bytes: &Vec<u8>) -> RepeatingKey {
		let key_size: u16 = key_bytes.len() as u16;

		RepeatingKey {
			key: key_bytes.clone(),
			iter_val: 0,
			key_size: key_size,
		}
	}

	pub fn new(key: &str) -> RepeatingKey {
		let key_bytes: Vec<u8> = key.bytes().collect();
		RepeatingKey::new_bytes(&key_bytes)
	}

	pub fn of_length(&mut self, length: usize) -> Vec<u8> {
		let mut result = vec![];
		for x in self.take(length) {
			result.push(x);
		}
		self.iter_val = 0;
		result
	}

	pub fn encrypt_bytes(&mut self, bytes: &Vec<u8>) -> Vec<u8> {
		use byte_manipulation::*;

		let length = bytes.len();

		xor_byte_streams(&bytes, &self.of_length(length))
	}
}

//-----------------------------------------------------------------------------

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

	#[test]
	fn RepeatingKey_short() {
		let rk = RepeatingKey::new(&"ICE");
		let mut actual: Vec<u8> = vec![];

		for x in rk.take(2) {
			actual.push(x);
		}

		let expected: Vec<u8> = "IC".to_string().bytes().collect();

		assert_eq!(expected, actual);
	}

	#[test]
	fn RepeatingKey_long() {
		let rk = RepeatingKey::new(&"ICE");
		let mut actual: Vec<u8> = vec![];

		for x in rk.take(7) {
			actual.push(x);
		}

		let expected: Vec<u8> = "ICEICEI".to_string().bytes().collect();

		assert_eq!(expected, actual);
	}

	#[test]
	fn RepeatingKey_problem5() {
		// http://cryptopals.com/sets/1/challenges/5/
		use byte_conversion::*;

		let bytes = readable_text_to_bytes(&"Burning 'em, if you ain't quick and nimble\nI go crazy when I hear a cymbal");
		let output = RepeatingKey::new(&"ICE").encrypt_bytes(&bytes);

		let expected = hex_to_bytes(&"0b3637272a2b2e63622c2e69692a23693a2a3c6324202d623d63343c2a26226324272765272a282b2f20430a652e2c652a3124333a653e2b2027630c692b20283165286326302e27282f");

		assert_eq!(expected, output);
	}

	#[test]
	fn RepeatingKey_of_length() {
		let actual = RepeatingKey::new(&"ICE").of_length(7);
		let expected = vec![73, 67, 69, 73, 67, 69, 73];

		assert_eq!(expected, actual);
	}

	#[test]
	fn RepeatingKey_encrypt_bytes() {
		use byte_conversion::*;

		let bytes = readable_text_to_bytes(&"Burning 'em, if you ain't quick and nimble\nI go crazy when I hear a cymbal");
		let output = RepeatingKey::new(&"ICE").encrypt_bytes(&bytes);

		let expected = hex_to_bytes(&"0b3637272a2b2e63622c2e69692a23693a2a3c6324202d623d63343c2a26226324272765272a282b2f20430a652e2c652a3124333a653e2b2027630c692b20283165286326302e27282f");

		assert_eq!(expected, output);
	}

	#[test]
	fn RepeatingKey_encrypt_bytes_twice() {
		use byte_conversion::*;

		let word = "Burning 'em, if you ain't quick and nimble\nI go crazy when I hear a cymbal";

		let bytes = readable_text_to_bytes(&word);
		let mut key = RepeatingKey::new(&"ICE");

		let encrypted = key.encrypt_bytes(&bytes);
		let decrypted = key.encrypt_bytes(&encrypted);

		assert_eq!(bytes, decrypted);
	}
}<|endoftext|>