<commit_before>#![cfg(windows)]

extern crate libc;
extern crate winapi;
extern crate "kernel32-sys" as kernel32;

use self::kernel32::*;
use self::winapi::*;
use std::ptr;
use std::iter::repeat;
use std::cmp::min;
use screen::Key;
use screen::Key::*;

extern "system" {
    fn CreateFileA(lpFileName: LPCSTR, dwDesiredAccess: DWORD, dwShareMode: DWORD, lpSecurityAttributes: LPSECURITY_ATTRIBUTES, dwCreationDisposition: DWORD, dwFlagsAndAttributes: DWORD, hTemplateFile: HANDLE) -> HANDLE;
    fn ReadFile(hFile: HANDLE, lpBuffer: LPVOID, nNumberOfBytesToRead: DWORD, lpNumberOfBytesRead: LPDWORD, lpOverlapped: LPOVERLAPPED) -> BOOL;
    fn WriteFile(hFile: HANDLE, lpBuffer: LPVOID, nNumberOfBytesToWrite: DWORD, lpNumberOfBytesWritten: LPDWORD, lpOverlapped: LPOVERLAPPED) -> BOOL;
}

macro_rules! win32 {
    ($funcall:expr) => (
        if unsafe { $funcall } == 0 {
            panic!("Win32 call failed");
        }
    );
}

pub struct Screen {
    pub height: u16,
    pub width: u16,
    pub visible_choices: u16,
    pub start_line: u16,
    original_console_mode: DWORD,
    original_colors: WORD,
    conin: HANDLE,
    conout: HANDLE,
}

impl Drop for Screen {
    fn drop(&mut self) {
        win32!(SetConsoleMode(self.conout, self.original_console_mode));
    }
}

impl Screen {
    pub fn open_screen() -> Screen {
        let mut orig_mode;
        let conin: HANDLE;
        let conout: HANDLE;
        unsafe {
            // Unlike in Linux, we *could* use the default stdin/stdout file handles to talk
            // directly to the console. However, we need the default stdin handle for the list of
            // choices, and we need to write just the final selection to stdout. Therefore we have
            // to explicitly create new handles to talk to the console. The gory details are
            // available at:
            //
            // https://msdn.microsoft.com/en-us/library/windows/desktop/ms682075%28v=vs.85%29.aspx
            // http://stackoverflow.com/questions/377152/what-does-createfileconin-do
            const OPEN_EXISTING: DWORD = 3;
            let rw_access = GENERIC_READ | GENERIC_WRITE;
            conin = CreateFileA("CONIN$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            conout = CreateFileA("CONOUT$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            orig_mode = ::std::mem::uninitialized();
        }

        if conin == INVALID_HANDLE_VALUE || conout == INVALID_HANDLE_VALUE {
            panic!("Unable to open console");
        }

        let (cols, rows) = Screen::winsize(conout).unwrap();

        win32!(GetConsoleMode(conout, &mut orig_mode));
        let new_mode = orig_mode & !(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT);
        win32!(SetConsoleMode(conin, new_mode));

        let visible_choices = min(20, rows - 1);
        let start_line = rows - visible_choices - 1;
        let original_colors = Screen::get_original_colors(conout);
        Screen {
            height: rows,
            width: cols,
            visible_choices: visible_choices,
            start_line: start_line,
            original_console_mode: orig_mode,
            original_colors: original_colors,
            conin: conin,
            conout: conout,
        }
    }

    pub fn move_cursor(&mut self, line: u16, column: u16) {
        win32!(SetConsoleCursorPosition(self.conout, COORD { X: column as i16, Y: line as i16}));
    }

    pub fn move_cursor_to_bottom(&mut self) {
        let end_line = self.start_line + self.visible_choices;
        self.move_cursor(end_line, 0);
        self.write("\n");
    }

    pub fn blank_screen(&mut self) {
        let start_line = self.start_line;
        self.move_cursor(start_line, 0);
        let blank_line = repeat(' ').take(self.width as usize).collect::<String>();
        for _ in 0..self.height {
            self.write(&blank_line);
        }
        self.move_cursor(start_line, 0);
    }

    pub fn show_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: TRUE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn hide_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: FALSE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn write(&mut self, s: &str) {
        let mut copy = s.to_string();
        let len: DWORD = copy.len() as DWORD;
        copy.push('\0');
        let mut bytes_written: DWORD = 0;
        win32!(WriteFile(self.conout, copy.as_ptr() as PVOID, len, &mut bytes_written as LPDWORD, ptr::null_mut()));
    }

    pub fn write_inverted(&mut self, s: &str) {
        let orig = self.original_colors;
        const BLACK_ON_WHITE: WORD = (BACKGROUND_RED | BACKGROUND_GREEN | BACKGROUND_BLUE) as WORD;
        self.set_colors(BLACK_ON_WHITE);
        self.write(s);
        self.set_colors(orig);
    }

    fn set_colors(&mut self, colors: WORD) {
        win32!(SetConsoleTextAttribute(self.conout, colors));
    }

    // Currently the Windows implementation does not buffer, so this function just performs a
    // blocking read of a single key.
    pub fn get_buffered_keys(&mut self) -> Vec<Key> {
        let mut buf: Vec<u16> = repeat(0u16).take(0x1000).collect();
        let mut chars_read: DWORD = 0;
        win32!(ReadFile(self.conin, buf.as_mut_ptr() as LPVOID, 1, &mut chars_read as LPDWORD, ptr::null_mut()));

        let mut ret = Vec::new();
        for i in 0..chars_read {
            ret.push(Screen::translate_byte(buf[i as usize]));
        }
        ret
    }

    fn translate_byte(byte: u16) -> Key {
        if byte == '\r' as u16 {
            Enter
        } else if byte == 127 {
            Backspace
        } else if byte & 96 == 0 {
            Control(((byte + 96u16) as u8) as char)
        } else {
            Char((byte as u8) as char)
        }
    }

    fn get_original_colors(handle: HANDLE) -> WORD {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(handle, &mut buffer_info));
        buffer_info.wAttributes
    }

    fn winsize(conin: HANDLE) -> Option<(u16, u16)> {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        let result = unsafe { GetConsoleScreenBufferInfo(conin, &mut buffer_info) };
        if result != 0 {
            // This code specifically computes the size of the window,
            // *not* the size of the buffer (which is easily available
            // from dwSize). I got the algorithm from:
            //
            // http://stackoverflow.com/a/12642749
            let left = buffer_info.srWindow.Left;
            let top = buffer_info.srWindow.Top;
            let right = buffer_info.srWindow.Right;
            let bottom = buffer_info.srWindow.Bottom;
            let cols = right - left + 1;
            let rows = bottom - top + 1;
            Some((cols as u16, rows as u16))
        } else {
            None
        }
    }
}

#[test]
fn winsize_test() {
    let conout = unsafe { kernel32::GetStdHandle(-11) };
    let (cols, rows) = Screen::winsize(conout).expect("Failed to get window size!");
    // We don't know the window size a priori, but we can at least
    // assert that it is within some kind of sensible range.
    assert!(cols > 40);
    assert!(rows > 40);
    assert!(cols < 1000);
    assert!(rows < 1000);
}
<commit_msg>Remove now-redundant kernel32 bindings<commit_after>#![cfg(windows)]

extern crate libc;
extern crate winapi;
extern crate "kernel32-sys" as kernel32;

use self::kernel32::*;
use self::winapi::*;
use std::ptr;
use std::iter::repeat;
use std::cmp::min;
use screen::Key;
use screen::Key::*;

macro_rules! win32 {
    ($funcall:expr) => (
        if unsafe { $funcall } == 0 {
            panic!("Win32 call failed");
        }
    );
}

pub struct Screen {
    pub height: u16,
    pub width: u16,
    pub visible_choices: u16,
    pub start_line: u16,
    original_console_mode: DWORD,
    original_colors: WORD,
    conin: HANDLE,
    conout: HANDLE,
}

impl Drop for Screen {
    fn drop(&mut self) {
        win32!(SetConsoleMode(self.conout, self.original_console_mode));
    }
}

impl Screen {
    pub fn open_screen() -> Screen {
        let mut orig_mode;
        let conin: HANDLE;
        let conout: HANDLE;
        unsafe {
            // Unlike in Linux, we *could* use the default stdin/stdout file handles to talk
            // directly to the console. However, we need the default stdin handle for the list of
            // choices, and we need to write just the final selection to stdout. Therefore we have
            // to explicitly create new handles to talk to the console. The gory details are
            // available at:
            //
            // https://msdn.microsoft.com/en-us/library/windows/desktop/ms682075%28v=vs.85%29.aspx
            // http://stackoverflow.com/questions/377152/what-does-createfileconin-do
            const OPEN_EXISTING: DWORD = 3;
            let rw_access = GENERIC_READ | GENERIC_WRITE;
            conin = CreateFileA("CONIN$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            conout = CreateFileA("CONOUT$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            orig_mode = ::std::mem::uninitialized();
        }

        if conin == INVALID_HANDLE_VALUE || conout == INVALID_HANDLE_VALUE {
            panic!("Unable to open console");
        }

        let (cols, rows) = Screen::winsize(conout).unwrap();

        win32!(GetConsoleMode(conout, &mut orig_mode));
        let new_mode = orig_mode & !(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT);
        win32!(SetConsoleMode(conin, new_mode));

        let visible_choices = min(20, rows - 1);
        let start_line = rows - visible_choices - 1;
        let original_colors = Screen::get_original_colors(conout);
        Screen {
            height: rows,
            width: cols,
            visible_choices: visible_choices,
            start_line: start_line,
            original_console_mode: orig_mode,
            original_colors: original_colors,
            conin: conin,
            conout: conout,
        }
    }

    pub fn move_cursor(&mut self, line: u16, column: u16) {
        win32!(SetConsoleCursorPosition(self.conout, COORD { X: column as i16, Y: line as i16}));
    }

    pub fn move_cursor_to_bottom(&mut self) {
        let end_line = self.start_line + self.visible_choices;
        self.move_cursor(end_line, 0);
        self.write("\n");
    }

    pub fn blank_screen(&mut self) {
        let start_line = self.start_line;
        self.move_cursor(start_line, 0);
        let blank_line = repeat(' ').take(self.width as usize).collect::<String>();
        for _ in 0..self.height {
            self.write(&blank_line);
        }
        self.move_cursor(start_line, 0);
    }

    pub fn show_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: TRUE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn hide_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: FALSE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn write(&mut self, s: &str) {
        let mut copy = s.to_string();
        let len: DWORD = copy.len() as DWORD;
        copy.push('\0');
        let mut bytes_written: DWORD = 0;
        win32!(WriteFile(self.conout, copy.as_ptr() as PVOID, len, &mut bytes_written as LPDWORD, ptr::null_mut()));
    }

    pub fn write_inverted(&mut self, s: &str) {
        let orig = self.original_colors;
        const BLACK_ON_WHITE: WORD = (BACKGROUND_RED | BACKGROUND_GREEN | BACKGROUND_BLUE) as WORD;
        self.set_colors(BLACK_ON_WHITE);
        self.write(s);
        self.set_colors(orig);
    }

    fn set_colors(&mut self, colors: WORD) {
        win32!(SetConsoleTextAttribute(self.conout, colors));
    }

    // Currently the Windows implementation does not buffer, so this function just performs a
    // blocking read of a single key.
    pub fn get_buffered_keys(&mut self) -> Vec<Key> {
        let mut buf: Vec<u16> = repeat(0u16).take(0x1000).collect();
        let mut chars_read: DWORD = 0;
        win32!(ReadFile(self.conin, buf.as_mut_ptr() as LPVOID, 1, &mut chars_read as LPDWORD, ptr::null_mut()));

        let mut ret = Vec::new();
        for i in 0..chars_read {
            ret.push(Screen::translate_byte(buf[i as usize]));
        }
        ret
    }

    fn translate_byte(byte: u16) -> Key {
        if byte == '\r' as u16 {
            Enter
        } else if byte == 127 {
            Backspace
        } else if byte & 96 == 0 {
            Control(((byte + 96u16) as u8) as char)
        } else {
            Char((byte as u8) as char)
        }
    }

    fn get_original_colors(handle: HANDLE) -> WORD {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(handle, &mut buffer_info));
        buffer_info.wAttributes
    }

    fn winsize(conin: HANDLE) -> Option<(u16, u16)> {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        let result = unsafe { GetConsoleScreenBufferInfo(conin, &mut buffer_info) };
        if result != 0 {
            // This code specifically computes the size of the window,
            // *not* the size of the buffer (which is easily available
            // from dwSize). I got the algorithm from:
            //
            // http://stackoverflow.com/a/12642749
            let left = buffer_info.srWindow.Left;
            let top = buffer_info.srWindow.Top;
            let right = buffer_info.srWindow.Right;
            let bottom = buffer_info.srWindow.Bottom;
            let cols = right - left + 1;
            let rows = bottom - top + 1;
            Some((cols as u16, rows as u16))
        } else {
            None
        }
    }
}

#[test]
fn winsize_test() {
    let conout = unsafe { kernel32::GetStdHandle(-11) };
    let (cols, rows) = Screen::winsize(conout).expect("Failed to get window size!");
    // We don't know the window size a priori, but we can at least
    // assert that it is within some kind of sensible range.
    assert!(cols > 40);
    assert!(rows > 40);
    assert!(cols < 1000);
    assert!(rows < 1000);
}
<|endoftext|>


<commit_before>use termion::terminal_size;

pub struct ScreenManager {
    current_width: usize,
    prev_width: usize
}

impl ScreenManager {
    pub fn new () -> Self {
        let w = terminal_size().unwrap().0 as usize;
        ScreenManager {
            current_width: w,
            prev_width: w,
        }
    }
    pub fn getWidth(&self) -> usize {
        self.current_width
    }
    pub fn isWidthChanged(&mut self) -> bool {
        let w = terminal_size().unwrap().0 as usize;
        if self.current_width != w {
            self.prev_width = self.current_width;
            self.current_width = w;
            return true
        }
        return false
    }
    pub fn getHeight(&self) -> usize {
        terminal_size().unwrap().1 as usize
    }
}
<commit_msg>screen manager - replace unwrap with expect<commit_after>use termion::terminal_size;

pub struct ScreenManager {
    current_width: usize,
    prev_width: usize
}

impl ScreenManager {
    pub fn new () -> Self {
        let w = terminal_size().expect("fail to get terminal size").0 as usize;
        ScreenManager {
            current_width: w,
            prev_width: w,
        }
    }
    pub fn getWidth(&self) -> usize {
        self.current_width
    }
    pub fn isWidthChanged(&mut self) -> bool {
        let w = terminal_size().expect("fail to get terminal size").0 as usize;
        if self.current_width != w {
            self.prev_width = self.current_width;
            self.current_width = w;
            return true
        }
        return false
    }
    pub fn getHeight(&self) -> usize {
        terminal_size().expect("fail to get terminal size").1 as usize
    }
}
<|endoftext|>


<commit_before>//! Plan in FFTW
//!
//! See also [Using Plans] in the original document
//! [Using Plans]: http://www.fftw.org/fftw3_doc/Using-Plans.html

use crate::array::{alignment_of, AlignedAllocable, AlignedVec, Alignment};
use crate::error::*;
use crate::ffi::*;
use crate::types::{c32, c64, Flag, R2RKind, Sign};

use std::marker::PhantomData;

pub type C2CPlan64 = Plan<c64, c64, Plan64>;
pub type C2CPlan32 = Plan<c32, c32, Plan32>;
pub type R2CPlan64 = Plan<f64, c64, Plan64>;
pub type R2CPlan32 = Plan<f32, c32, Plan32>;
pub type C2RPlan64 = Plan<c64, f64, Plan64>;
pub type C2RPlan32 = Plan<c32, f32, Plan32>;
pub type R2RPlan64 = Plan<f64, f64, Plan64>;
pub type R2RPlan32 = Plan<f32, f32, Plan32>;

/// Typed wrapper of `fftw_plan`
///
/// The plan in FFTW manages the contains all information necessary to compute the transform,
/// including the pointers to the input and output arrays.
/// However, this wrapper *does not modify* this pointer once after the plan is created
/// since it should be *unsafe* in terms of Rust.
/// Instead, this plan executes a transform for different arrays with [new-array execute functions]
/// with related associated functions, e.g. `C2CPlan::c2c`.
///
/// [new-array execute functions]: http://www.fftw.org/fftw3_doc/New_002darray-Execute-Functions.html
pub struct Plan<A, B, Plan: PlanSpec> {
    plan: Plan,
    input: (usize, Alignment),
    output: (usize, Alignment),
    phantom: PhantomData<(A, B)>,
}

impl<A, B, P: PlanSpec> Drop for Plan<A, B, P> {
    fn drop(&mut self) {
        self.plan.destroy();
    }
}

impl<A, B, P: PlanSpec> Plan<A, B, P> {
    fn check(&self, in_: &[A], out: &[B]) -> Result<()> {
        if self.input != slice_info(in_) {
            return Err(Error::InputArrayMismatch {
                expect: self.input,
                actual: slice_info(in_),
            });
        }
        if self.output != slice_info(out) {
            return Err(Error::OutputArrayMismatch {
                expect: self.output,
                actual: slice_info(out),
            });
        }
        Ok(())
    }
}

/// Trait for Plan makers
pub trait PlanSpec: Clone + Copy {
    fn validate(self) -> Result<Self>;
    fn destroy(self);
    fn print(self);
}

/// Marker for 64-bit floating point FFT
pub type Plan64 = fftw_plan;
/// Marker for 32-bit floating point FFT
pub type Plan32 = fftwf_plan;

/// Trait for the plan of Complex-to-Complex transformation
pub trait C2CPlan: Sized {
    type Complex: AlignedAllocable;

    /// Create new plan with aligned vector
    fn aligned(shape: &[usize], sign: Sign, flag: Flag) -> Result<Self> {
        let n: usize = shape.iter().product();
        let mut in_ = AlignedVec::new(n);
        let mut out = AlignedVec::new(n);
        Self::new(shape, &mut in_, &mut out, sign, flag)
    }

    /// Create new plan
    fn new(
        shape: &[usize],
        in_: &mut [Self::Complex],
        out: &mut [Self::Complex],
        sign: Sign,
        flag: Flag,
    ) -> Result<Self>;

    /// Execute complex-to-complex transform
    fn c2c(&mut self, in_: &mut [Self::Complex], out: &mut [Self::Complex]) -> Result<()>;
}

/// Trait for the plan of Real-to-Complex transformation
pub trait R2CPlan: Sized {
    type Real: AlignedAllocable;
    type Complex: AlignedAllocable;

    /// Create new plan with aligned vector
    fn aligned(shape: &[usize], flag: Flag) -> Result<Self> {
        let n: usize = shape.iter().product();
        let n_d = shape.last().unwrap();
        let n_sub = (n / n_d) * (n_d / 2 + 1);
        let mut in_ = AlignedVec::new(n);
        let mut out = AlignedVec::new(n_sub);
        Self::new(shape, &mut in_, &mut out, flag)
    }

    /// Create new plan
    fn new(
        shape: &[usize],
        in_: &mut [Self::Real],
        out: &mut [Self::Complex],
        flag: Flag,
    ) -> Result<Self>;

    /// Execute real-to-complex transform
    fn r2c(&mut self, in_: &mut [Self::Real], out: &mut [Self::Complex]) -> Result<()>;
}

/// Trait for the plan of Complex-to-Real transformation
pub trait C2RPlan: Sized {
    type Real: AlignedAllocable;
    type Complex: AlignedAllocable;

    /// Create new plan with aligned vector
    fn aligned(shape: &[usize], flag: Flag) -> Result<Self> {
        let n: usize = shape.iter().product();
        let n_d = shape.last().unwrap();
        let n_sub = (n / n_d) * (n_d / 2 + 1);
        let mut in_ = AlignedVec::new(n_sub);
        let mut out = AlignedVec::new(n);
        Self::new(shape, &mut in_, &mut out, flag)
    }

    /// Create new plan
    fn new(
        shape: &[usize],
        in_: &mut [Self::Complex],
        out: &mut [Self::Real],
        flag: Flag,
    ) -> Result<Self>;

    /// Execute complex-to-real transform
    fn c2r(&mut self, in_: &mut [Self::Complex], out: &mut [Self::Real]) -> Result<()>;
}

pub trait R2RPlan: Sized {
    type Real: AlignedAllocable;

    /// Create new plan with aligned vector
    fn aligned(shape: &[usize], kind: R2RKind, flag: Flag) -> Result<Self> {
        let n: usize = shape.iter().product();
        let mut in_ = AlignedVec::new(n);
        let mut out = AlignedVec::new(n);
        Self::new(shape, &mut in_, &mut out, kind, flag)
    }

    /// Create new plan
    fn new(
        shape: &[usize],
        in_: &mut [Self::Real],
        out: &mut [Self::Real],
        kind: R2RKind,
        flag: Flag,
    ) -> Result<Self>;

    /// Execute complex-to-complex transform
    fn r2r(&mut self, in_: &mut [Self::Real], out: &mut [Self::Real]) -> Result<()>;
}

macro_rules! impl_c2c {
    ($C:ty, $Plan:ty; $plan:ident, $exec:ident) => {
        impl C2CPlan for Plan<$C, $C, $Plan> {
            type Complex = $C;
            fn new(
                shape: &[usize],
                in_: &mut [Self::Complex],
                out: &mut [Self::Complex],
                sign: Sign,
                flag: Flag,
            ) -> Result<Self> {
                let plan = excall! { $plan(
                    shape.len() as i32,
                    shape.to_cint().as_mut_ptr() as *mut _,
                    in_.as_mut_ptr(),
                    out.as_mut_ptr(),
                    sign as i32, flag.bits())
                }
                .validate()?;
                Ok(Self {
                    plan,
                    input: slice_info(in_),
                    output: slice_info(out),
                    phantom: PhantomData,
                })
            }
            fn c2c(&mut self, in_: &mut [Self::Complex], out: &mut [Self::Complex]) -> Result<()> {
                unsafe { $exec(self.plan, in_.as_mut_ptr(), out.as_mut_ptr()) };
                Ok(())
            }
        }
    };
} // impl_c2c!

impl_c2c!(c64, Plan64; fftw_plan_dft, fftw_execute_dft);
impl_c2c!(c32, Plan32; fftwf_plan_dft, fftwf_execute_dft);

macro_rules! impl_r2c {
    ($R:ty, $C:ty, $Plan:ty; $plan:ident, $exec:ident) => {
        impl R2CPlan for Plan<$R, $C, $Plan> {
            type Real = $R;
            type Complex = $C;
            fn new(
                shape: &[usize],
                in_: &mut [Self::Real],
                out: &mut [Self::Complex],
                flag: Flag,
            ) -> Result<Self> {
                let plan = excall! { $plan(
                    shape.len() as i32,
                    shape.to_cint().as_mut_ptr() as *mut _,
                    in_.as_mut_ptr(),
                    out.as_mut_ptr(),
                    flag.bits())
                }
                .validate()?;
                Ok(Self {
                    plan,
                    input: slice_info(in_),
                    output: slice_info(out),
                    phantom: PhantomData,
                })
            }
            fn r2c(&mut self, in_: &mut [Self::Real], out: &mut [Self::Complex]) -> Result<()> {
                self.check(in_, out)?;
                unsafe { $exec(self.plan, in_.as_mut_ptr(), out.as_mut_ptr()) };
                Ok(())
            }
        }
    };
} // impl_r2c!

impl_r2c!(f64, c64, Plan64; fftw_plan_dft_r2c, fftw_execute_dft_r2c);
impl_r2c!(f32, c32, Plan32; fftwf_plan_dft_r2c, fftwf_execute_dft_r2c);

macro_rules! impl_c2r {
    ($R:ty, $C:ty, $Plan:ty; $plan:ident, $exec:ident) => {
        impl C2RPlan for Plan<$C, $R, $Plan> {
            type Real = $R;
            type Complex = $C;
            fn new(
                shape: &[usize],
                in_: &mut [Self::Complex],
                out: &mut [Self::Real],
                flag: Flag,
            ) -> Result<Self> {
                let plan = excall! { $plan(
                    shape.len() as i32,
                    shape.to_cint().as_mut_ptr() as *mut _,
                    in_.as_mut_ptr(),
                    out.as_mut_ptr(),
                    flag.bits())
                }
                .validate()?;
                Ok(Self {
                    plan,
                    input: slice_info(in_),
                    output: slice_info(out),
                    phantom: PhantomData,
                })
            }
            fn c2r(&mut self, in_: &mut [Self::Complex], out: &mut [Self::Real]) -> Result<()> {
                self.check(in_, out)?;
                unsafe { $exec(self.plan, in_.as_mut_ptr(), out.as_mut_ptr()) };
                Ok(())
            }
        }
    };
} // impl_c2r!

impl_c2r!(f64, c64, Plan64; fftw_plan_dft_c2r, fftw_execute_dft_c2r);
impl_c2r!(f32, c32, Plan32; fftwf_plan_dft_c2r, fftwf_execute_dft_c2r);

macro_rules! impl_r2r {
    ($R:ty, $Plan:ty; $plan:ident, $exec:ident) => {
        impl R2RPlan for Plan<$R, $R, $Plan> {
            type Real = $R;
            fn new(
                shape: &[usize],
                in_: &mut [Self::Real],
                out: &mut [Self::Real],
                kind: R2RKind,
                flag: Flag,
            ) -> Result<Self> {
                let plan = excall! { $plan(
                    shape.len() as i32,
                    shape.to_cint().as_mut_ptr() as *mut _,
                    in_.as_mut_ptr(),
                    out.as_mut_ptr(),
                    &kind as *const _, flag.bits())
                }
                .validate()?;
                Ok(Self {
                    plan,
                    input: slice_info(in_),
                    output: slice_info(out),
                    phantom: PhantomData,
                })
            }
            fn r2r(&mut self, in_: &mut [Self::Real], out: &mut [Self::Real]) -> Result<()> {
                unsafe { $exec(self.plan, in_.as_mut_ptr(), out.as_mut_ptr()) };
                Ok(())
            }
        }
    };
} // impl_r2r!

impl_r2r!(f64, Plan64; fftw_plan_r2r, fftw_execute_r2r);
impl_r2r!(f32, Plan32; fftwf_plan_r2r, fftwf_execute_r2r);

macro_rules! impl_plan_spec {
    ($Plan:ty; $destroy_plan:ident, $print_plan:ident) => {
        impl PlanSpec for $Plan {
            fn validate(self) -> Result<Self> {
                if self.is_null() {
                    Err(Error::InvalidPlanError {})
                } else {
                    Ok(self)
                }
            }
            fn destroy(self) {
                excall! { $destroy_plan(self) }
            }
            fn print(self) {
                excall! { $print_plan(self) }
            }
        }
    };
} // impl_plan_spec!

impl_plan_spec!(Plan64; fftw_destroy_plan, fftw_print_plan);
impl_plan_spec!(Plan32; fftwf_destroy_plan, fftwf_print_plan);

// Convert [usize] -> [i32]
trait ToCInt {
    fn to_cint(&self) -> Vec<i32>;
}

impl ToCInt for [usize] {
    fn to_cint(&self) -> Vec<i32> {
        self.iter().map(|&x| x as i32).collect()
    }
}

fn slice_info<T>(a: &[T]) -> (usize, Alignment) {
    (a.len(), alignment_of(a))
}
<commit_msg>Validate input to C2CPlan::c2c and R2RPlan::r2r<commit_after>//! Plan in FFTW
//!
//! See also [Using Plans] in the original document
//! [Using Plans]: http://www.fftw.org/fftw3_doc/Using-Plans.html

use crate::array::{alignment_of, AlignedAllocable, AlignedVec, Alignment};
use crate::error::*;
use crate::ffi::*;
use crate::types::{c32, c64, Flag, R2RKind, Sign};

use std::marker::PhantomData;

pub type C2CPlan64 = Plan<c64, c64, Plan64>;
pub type C2CPlan32 = Plan<c32, c32, Plan32>;
pub type R2CPlan64 = Plan<f64, c64, Plan64>;
pub type R2CPlan32 = Plan<f32, c32, Plan32>;
pub type C2RPlan64 = Plan<c64, f64, Plan64>;
pub type C2RPlan32 = Plan<c32, f32, Plan32>;
pub type R2RPlan64 = Plan<f64, f64, Plan64>;
pub type R2RPlan32 = Plan<f32, f32, Plan32>;

/// Typed wrapper of `fftw_plan`
///
/// The plan in FFTW manages the contains all information necessary to compute the transform,
/// including the pointers to the input and output arrays.
/// However, this wrapper *does not modify* this pointer once after the plan is created
/// since it should be *unsafe* in terms of Rust.
/// Instead, this plan executes a transform for different arrays with [new-array execute functions]
/// with related associated functions, e.g. `C2CPlan::c2c`.
///
/// [new-array execute functions]: http://www.fftw.org/fftw3_doc/New_002darray-Execute-Functions.html
pub struct Plan<A, B, Plan: PlanSpec> {
    plan: Plan,
    input: (usize, Alignment),
    output: (usize, Alignment),
    phantom: PhantomData<(A, B)>,
}

impl<A, B, P: PlanSpec> Drop for Plan<A, B, P> {
    fn drop(&mut self) {
        self.plan.destroy();
    }
}

impl<A, B, P: PlanSpec> Plan<A, B, P> {
    fn check(&self, in_: &[A], out: &[B]) -> Result<()> {
        if self.input != slice_info(in_) {
            return Err(Error::InputArrayMismatch {
                expect: self.input,
                actual: slice_info(in_),
            });
        }
        if self.output != slice_info(out) {
            return Err(Error::OutputArrayMismatch {
                expect: self.output,
                actual: slice_info(out),
            });
        }
        Ok(())
    }
}

/// Trait for Plan makers
pub trait PlanSpec: Clone + Copy {
    fn validate(self) -> Result<Self>;
    fn destroy(self);
    fn print(self);
}

/// Marker for 64-bit floating point FFT
pub type Plan64 = fftw_plan;
/// Marker for 32-bit floating point FFT
pub type Plan32 = fftwf_plan;

/// Trait for the plan of Complex-to-Complex transformation
pub trait C2CPlan: Sized {
    type Complex: AlignedAllocable;

    /// Create new plan with aligned vector
    fn aligned(shape: &[usize], sign: Sign, flag: Flag) -> Result<Self> {
        let n: usize = shape.iter().product();
        let mut in_ = AlignedVec::new(n);
        let mut out = AlignedVec::new(n);
        Self::new(shape, &mut in_, &mut out, sign, flag)
    }

    /// Create new plan
    fn new(
        shape: &[usize],
        in_: &mut [Self::Complex],
        out: &mut [Self::Complex],
        sign: Sign,
        flag: Flag,
    ) -> Result<Self>;

    /// Execute complex-to-complex transform
    fn c2c(&mut self, in_: &mut [Self::Complex], out: &mut [Self::Complex]) -> Result<()>;
}

/// Trait for the plan of Real-to-Complex transformation
pub trait R2CPlan: Sized {
    type Real: AlignedAllocable;
    type Complex: AlignedAllocable;

    /// Create new plan with aligned vector
    fn aligned(shape: &[usize], flag: Flag) -> Result<Self> {
        let n: usize = shape.iter().product();
        let n_d = shape.last().unwrap();
        let n_sub = (n / n_d) * (n_d / 2 + 1);
        let mut in_ = AlignedVec::new(n);
        let mut out = AlignedVec::new(n_sub);
        Self::new(shape, &mut in_, &mut out, flag)
    }

    /// Create new plan
    fn new(
        shape: &[usize],
        in_: &mut [Self::Real],
        out: &mut [Self::Complex],
        flag: Flag,
    ) -> Result<Self>;

    /// Execute real-to-complex transform
    fn r2c(&mut self, in_: &mut [Self::Real], out: &mut [Self::Complex]) -> Result<()>;
}

/// Trait for the plan of Complex-to-Real transformation
pub trait C2RPlan: Sized {
    type Real: AlignedAllocable;
    type Complex: AlignedAllocable;

    /// Create new plan with aligned vector
    fn aligned(shape: &[usize], flag: Flag) -> Result<Self> {
        let n: usize = shape.iter().product();
        let n_d = shape.last().unwrap();
        let n_sub = (n / n_d) * (n_d / 2 + 1);
        let mut in_ = AlignedVec::new(n_sub);
        let mut out = AlignedVec::new(n);
        Self::new(shape, &mut in_, &mut out, flag)
    }

    /// Create new plan
    fn new(
        shape: &[usize],
        in_: &mut [Self::Complex],
        out: &mut [Self::Real],
        flag: Flag,
    ) -> Result<Self>;

    /// Execute complex-to-real transform
    fn c2r(&mut self, in_: &mut [Self::Complex], out: &mut [Self::Real]) -> Result<()>;
}

pub trait R2RPlan: Sized {
    type Real: AlignedAllocable;

    /// Create new plan with aligned vector
    fn aligned(shape: &[usize], kind: R2RKind, flag: Flag) -> Result<Self> {
        let n: usize = shape.iter().product();
        let mut in_ = AlignedVec::new(n);
        let mut out = AlignedVec::new(n);
        Self::new(shape, &mut in_, &mut out, kind, flag)
    }

    /// Create new plan
    fn new(
        shape: &[usize],
        in_: &mut [Self::Real],
        out: &mut [Self::Real],
        kind: R2RKind,
        flag: Flag,
    ) -> Result<Self>;

    /// Execute complex-to-complex transform
    fn r2r(&mut self, in_: &mut [Self::Real], out: &mut [Self::Real]) -> Result<()>;
}

macro_rules! impl_c2c {
    ($C:ty, $Plan:ty; $plan:ident, $exec:ident) => {
        impl C2CPlan for Plan<$C, $C, $Plan> {
            type Complex = $C;
            fn new(
                shape: &[usize],
                in_: &mut [Self::Complex],
                out: &mut [Self::Complex],
                sign: Sign,
                flag: Flag,
            ) -> Result<Self> {
                let plan = excall! { $plan(
                    shape.len() as i32,
                    shape.to_cint().as_mut_ptr() as *mut _,
                    in_.as_mut_ptr(),
                    out.as_mut_ptr(),
                    sign as i32, flag.bits())
                }
                .validate()?;
                Ok(Self {
                    plan,
                    input: slice_info(in_),
                    output: slice_info(out),
                    phantom: PhantomData,
                })
            }
            fn c2c(&mut self, in_: &mut [Self::Complex], out: &mut [Self::Complex]) -> Result<()> {
                self.check(in_, out)?;
                unsafe { $exec(self.plan, in_.as_mut_ptr(), out.as_mut_ptr()) };
                Ok(())
            }
        }
    };
} // impl_c2c!

impl_c2c!(c64, Plan64; fftw_plan_dft, fftw_execute_dft);
impl_c2c!(c32, Plan32; fftwf_plan_dft, fftwf_execute_dft);

macro_rules! impl_r2c {
    ($R:ty, $C:ty, $Plan:ty; $plan:ident, $exec:ident) => {
        impl R2CPlan for Plan<$R, $C, $Plan> {
            type Real = $R;
            type Complex = $C;
            fn new(
                shape: &[usize],
                in_: &mut [Self::Real],
                out: &mut [Self::Complex],
                flag: Flag,
            ) -> Result<Self> {
                let plan = excall! { $plan(
                    shape.len() as i32,
                    shape.to_cint().as_mut_ptr() as *mut _,
                    in_.as_mut_ptr(),
                    out.as_mut_ptr(),
                    flag.bits())
                }
                .validate()?;
                Ok(Self {
                    plan,
                    input: slice_info(in_),
                    output: slice_info(out),
                    phantom: PhantomData,
                })
            }
            fn r2c(&mut self, in_: &mut [Self::Real], out: &mut [Self::Complex]) -> Result<()> {
                self.check(in_, out)?;
                unsafe { $exec(self.plan, in_.as_mut_ptr(), out.as_mut_ptr()) };
                Ok(())
            }
        }
    };
} // impl_r2c!

impl_r2c!(f64, c64, Plan64; fftw_plan_dft_r2c, fftw_execute_dft_r2c);
impl_r2c!(f32, c32, Plan32; fftwf_plan_dft_r2c, fftwf_execute_dft_r2c);

macro_rules! impl_c2r {
    ($R:ty, $C:ty, $Plan:ty; $plan:ident, $exec:ident) => {
        impl C2RPlan for Plan<$C, $R, $Plan> {
            type Real = $R;
            type Complex = $C;
            fn new(
                shape: &[usize],
                in_: &mut [Self::Complex],
                out: &mut [Self::Real],
                flag: Flag,
            ) -> Result<Self> {
                let plan = excall! { $plan(
                    shape.len() as i32,
                    shape.to_cint().as_mut_ptr() as *mut _,
                    in_.as_mut_ptr(),
                    out.as_mut_ptr(),
                    flag.bits())
                }
                .validate()?;
                Ok(Self {
                    plan,
                    input: slice_info(in_),
                    output: slice_info(out),
                    phantom: PhantomData,
                })
            }
            fn c2r(&mut self, in_: &mut [Self::Complex], out: &mut [Self::Real]) -> Result<()> {
                self.check(in_, out)?;
                unsafe { $exec(self.plan, in_.as_mut_ptr(), out.as_mut_ptr()) };
                Ok(())
            }
        }
    };
} // impl_c2r!

impl_c2r!(f64, c64, Plan64; fftw_plan_dft_c2r, fftw_execute_dft_c2r);
impl_c2r!(f32, c32, Plan32; fftwf_plan_dft_c2r, fftwf_execute_dft_c2r);

macro_rules! impl_r2r {
    ($R:ty, $Plan:ty; $plan:ident, $exec:ident) => {
        impl R2RPlan for Plan<$R, $R, $Plan> {
            type Real = $R;
            fn new(
                shape: &[usize],
                in_: &mut [Self::Real],
                out: &mut [Self::Real],
                kind: R2RKind,
                flag: Flag,
            ) -> Result<Self> {
                let plan = excall! { $plan(
                    shape.len() as i32,
                    shape.to_cint().as_mut_ptr() as *mut _,
                    in_.as_mut_ptr(),
                    out.as_mut_ptr(),
                    &kind as *const _, flag.bits())
                }
                .validate()?;
                Ok(Self {
                    plan,
                    input: slice_info(in_),
                    output: slice_info(out),
                    phantom: PhantomData,
                })
            }
            fn r2r(&mut self, in_: &mut [Self::Real], out: &mut [Self::Real]) -> Result<()> {
                self.check(in_, out)?;
                unsafe { $exec(self.plan, in_.as_mut_ptr(), out.as_mut_ptr()) };
                Ok(())
            }
        }
    };
} // impl_r2r!

impl_r2r!(f64, Plan64; fftw_plan_r2r, fftw_execute_r2r);
impl_r2r!(f32, Plan32; fftwf_plan_r2r, fftwf_execute_r2r);

macro_rules! impl_plan_spec {
    ($Plan:ty; $destroy_plan:ident, $print_plan:ident) => {
        impl PlanSpec for $Plan {
            fn validate(self) -> Result<Self> {
                if self.is_null() {
                    Err(Error::InvalidPlanError {})
                } else {
                    Ok(self)
                }
            }
            fn destroy(self) {
                excall! { $destroy_plan(self) }
            }
            fn print(self) {
                excall! { $print_plan(self) }
            }
        }
    };
} // impl_plan_spec!

impl_plan_spec!(Plan64; fftw_destroy_plan, fftw_print_plan);
impl_plan_spec!(Plan32; fftwf_destroy_plan, fftwf_print_plan);

// Convert [usize] -> [i32]
trait ToCInt {
    fn to_cint(&self) -> Vec<i32>;
}

impl ToCInt for [usize] {
    fn to_cint(&self) -> Vec<i32> {
        self.iter().map(|&x| x as i32).collect()
    }
}

fn slice_info<T>(a: &[T]) -> (usize, Alignment) {
    (a.len(), alignment_of(a))
}
<|endoftext|>


<commit_before>use gtk;
use gtk::prelude::*;
use std::cell::RefCell;
use portaudio as pa;
use std::rc::Rc;
use std::sync::Arc;
use std::sync::RwLock;
use std::io;
use std::io::Write;
use std::thread;
use std::sync::mpsc::*;

struct RustyUi {
    dropdown: gtk::ComboBoxText,
    pitch_label: gtk::Label,
    pitch_error_indicator: gtk::DrawingArea,
    oscilloscope_chart: gtk::DrawingArea,
    freq_chart: gtk::DrawingArea,
    correlation_chart: gtk::DrawingArea
}

struct ApplicationState {
    pa: pa::PortAudio,
    pa_stream: Option<pa::Stream<pa::NonBlocking, pa::Input<f32>>>,
    ui: RustyUi
}

struct CrossThreadState {
    pitch: String,
    error: f64,
    signal: Vec<f64>,
    freq_spectrum: Vec<::transforms::FrequencyBucket>,
    correlation: Vec<f64>
}

pub fn start_gui() -> Result<(), String> {
    let pa = try!(::audio::init().map_err(|e| e.to_string()));
    let microphones = try!(::audio::get_device_list(&pa).map_err(|e| e.to_string()));

    try!(gtk::init().map_err(|_| "Failed to initialize GTK."));

    let state = Rc::new(RefCell::new(ApplicationState {
        pa: pa,
        pa_stream: None,
        ui: create_window(microphones)
    }));

    let cross_thread_state = Arc::new(RwLock::new(CrossThreadState {
        pitch: String::new(),
        error: 0.0,
        signal: Vec::new(),
        freq_spectrum: Vec::new(),
        correlation: Vec::new()
    }));

    let (mic_sender, mic_receiver) = channel();

    connect_dropdown_choose_microphone(mic_sender, state.clone());

    start_processing_audio(mic_receiver, cross_thread_state.clone());
    setup_pitch_label_callbacks(state.clone(), cross_thread_state.clone());
    setup_pitch_error_indicator_callbacks(state.clone(), cross_thread_state.clone());
    setup_oscilloscope_drawing_area_callbacks(state.clone(), cross_thread_state.clone());
    setup_freq_drawing_area_callbacks(state.clone(), cross_thread_state.clone());
    setup_correlation_drawing_area_callbacks(state.clone(), cross_thread_state.clone());

    gtk::main();
    Ok(())
}

fn create_window(microphones: Vec<(u32, String)>) -> RustyUi {
    let window = gtk::Window::new(gtk::WindowType::Toplevel);
    window.set_title("Rusty Microphone");
    window.connect_delete_event(|_, _| {
        gtk::main_quit();
        Inhibit(false)
    });

    let layout_box = gtk::Box::new(gtk::Orientation::Vertical, 5);
    window.add(&layout_box);

    let dropdown = gtk::ComboBoxText::new();
    set_dropdown_items(&dropdown, microphones);
    layout_box.add(&dropdown);

    let pitch_label = gtk::Label::new(None);
    layout_box.add(&pitch_label);

    let pitch_error_indicator = gtk::DrawingArea::new();
    pitch_error_indicator.set_size_request(600, 50);
    layout_box.add(&pitch_error_indicator);

    let oscilloscope_chart = gtk::DrawingArea::new();
    oscilloscope_chart.set_size_request(600, 250);
    layout_box.add(&oscilloscope_chart);

    let freq_chart = gtk::DrawingArea::new();
    freq_chart.set_size_request(600, 250);
    layout_box.add(&freq_chart);

    let correlation_chart = gtk::DrawingArea::new();
    correlation_chart.set_size_request(600, 250);
    layout_box.add(&correlation_chart);

    window.show_all();

    RustyUi {
        dropdown: dropdown,
        pitch_label: pitch_label,
        pitch_error_indicator: pitch_error_indicator,
        oscilloscope_chart: oscilloscope_chart,
        freq_chart: freq_chart,
        correlation_chart: correlation_chart
    }
}

fn set_dropdown_items(dropdown: &gtk::ComboBoxText, microphones: Vec<(u32, String)>) {
    for (index, name) in microphones {
        dropdown.append(Some(format!("{}", index).as_ref()), name.as_ref());
    }
}

fn connect_dropdown_choose_microphone(mic_sender: Sender<Vec<f64>>, state: Rc<RefCell<ApplicationState>>) {
    let outer_state = state.clone();
    let ref dropdown = outer_state.borrow().ui.dropdown;
    dropdown.connect_changed(move |dropdown: &gtk::ComboBoxText| {
        match state.borrow_mut().pa_stream {
            Some(ref mut stream) => {stream.stop().ok();},
            _ => {}
        }
        let selected_mic = match dropdown.get_active_id().and_then(|id| id.parse().ok()) {
            Some(mic) => mic,
            None => {return;}
        };
        let stream = ::audio::start_listening(&state.borrow().pa, selected_mic, mic_sender.clone()).ok();
        if stream.is_none() {
            writeln!(io::stderr(), "Failed to open audio channel").ok();
        }
        state.borrow_mut().pa_stream = stream;
    });
}

fn start_processing_audio(mic_receiver: Receiver<Vec<f64>>, cross_thread_state: Arc<RwLock<CrossThreadState>>) {
    thread::spawn(move || {
        loop {
            let mut samples = None;
            loop {
                let next = mic_receiver.try_recv().ok();
                if next.is_none() {
                    break;
                }
                samples = next;
            }
            let samples = match samples {
                Some(samples) => samples,
                None => {continue;}
            };

            let signal = ::transforms::align_to_rising_edge(&samples);
            let frequency_domain = ::transforms::fft(&samples, 44100.0);
            let correlation = ::transforms::correlation(&samples);
            let fundamental = ::transforms::find_fundamental_frequency_correlation(&samples, 44100.0);
            let (pitch, error) = match fundamental {
                Some(fundamental) => (::transforms::hz_to_pitch(fundamental), ::transforms::hz_to_cents_error(fundamental)),
                None => ("".to_string(), 0.0)
            };

            match cross_thread_state.write() {
                Ok(mut state) => {
                    state.pitch = pitch;
                    state.signal = signal;
                    state.freq_spectrum = frequency_domain;
                    state.correlation = correlation;
                    state.error = error
                },
                Err(_) => {}
            };
        }
    });
}

fn setup_pitch_label_callbacks(state: Rc<RefCell<ApplicationState>>, cross_thread_state: Arc<RwLock<CrossThreadState>>) {
    gtk::timeout_add(100, move || {
        let ref pitch = cross_thread_state.read().unwrap().pitch;
        let ref ui = state.borrow().ui;
        ui.pitch_label.set_label(pitch.as_ref());
        ui.pitch_error_indicator.queue_draw();
        ui.oscilloscope_chart.queue_draw();
        ui.correlation_chart.queue_draw();
        ui.freq_chart.queue_draw();

        gtk::Continue(true)
    });
}

fn setup_pitch_error_indicator_callbacks(state: Rc<RefCell<ApplicationState>>, cross_thread_state: Arc<RwLock<CrossThreadState>>) {
    let outer_state = state.clone();
    let ref canvas = outer_state.borrow().ui.pitch_error_indicator;
    canvas.connect_draw(move |ref canvas, ref context| {
        let error = cross_thread_state.read().unwrap().error;
        let width = canvas.get_allocated_width() as f64;
        let midpoint = width / 2.0;
        let height = canvas.get_allocated_height() as f64;

        //flat on the left
        context.set_source_rgb(0.0, 0.0, if error < 0.0 {-error/50.0} else {0.0});
        context.rectangle(0.0, 0.0, midpoint, height);
        context.fill();

        //sharp on the right
        context.set_source_rgb(if error > 0.0 {error/50.0} else {0.0}, 0.0, 0.0);
        context.rectangle(midpoint, 0.0, width, height);
        context.fill();

        gtk::Inhibit(false)
    });
}

fn setup_oscilloscope_drawing_area_callbacks(state: Rc<RefCell<ApplicationState>>, cross_thread_state: Arc<RwLock<CrossThreadState>>) {
    let outer_state = state.clone();
    let ref canvas = outer_state.borrow().ui.oscilloscope_chart;
    canvas.connect_draw(move |ref canvas, ref context| {
        let ref signal = cross_thread_state.read().unwrap().signal;
        let width = canvas.get_allocated_width() as f64;
        let len = 512.0; //Set as a constant so signal won't change size based on zero point.
        let height = canvas.get_allocated_height() as f64;
        let mid_height = height / 2.0;
        let max = signal.iter().map(|x| x.abs()).fold(0.0, |max, x| if max > x { max } else { x });

        context.new_path();
        context.move_to(0.0, mid_height);

        for (i, intensity) in signal.iter().enumerate() {
            let x = i as f64 * width / len;
            let y = mid_height - (intensity * mid_height / max);
            context.line_to(x, y);
        }

        context.stroke();

        gtk::Inhibit(false)
    });
}

fn setup_freq_drawing_area_callbacks(state: Rc<RefCell<ApplicationState>>, cross_thread_state: Arc<RwLock<CrossThreadState>>) {
    let outer_state = state.clone();
    let ref canvas = outer_state.borrow().ui.freq_chart;
    canvas.connect_draw(move |ref canvas, ref context| {
        let ref spectrum = cross_thread_state.read().unwrap().freq_spectrum;
        let width = canvas.get_allocated_width() as f64;
        let height = canvas.get_allocated_height() as f64;
        let max = spectrum.iter().map(|x| x.intensity).fold(0.0, |max, x| if max > x { max } else { x });
        let len = spectrum.len() as f64;

        context.new_path();
        context.move_to(0.0, height);

        for (i, bucket) in spectrum.iter().enumerate() {
            let x = i as f64 * width / len;
            let y = height - (bucket.intensity * height / max);
            context.line_to(x, y);
        }

        context.stroke();

        gtk::Inhibit(false)
    });
}

fn setup_correlation_drawing_area_callbacks(state: Rc<RefCell<ApplicationState>>, cross_thread_state: Arc<RwLock<CrossThreadState>>) {
    let outer_state = state.clone();
    let ref canvas = outer_state.borrow().ui.correlation_chart;
    canvas.connect_draw(move |ref canvas, ref context| {
        let ref correlation = cross_thread_state.read().unwrap().correlation;
        if correlation.len() == 0 {
            return gtk::Inhibit(false);
        }

        let width = canvas.get_allocated_width() as f64;
        let height = canvas.get_allocated_height() as f64;
        let max = correlation[0];
        let len = correlation.len() as f64;

        context.new_path();
        context.move_to(0.0, height);

        for (i, val) in correlation.iter().enumerate() {
            let x = i as f64 * width / len;
            let y = height - (val * height / max);
            context.line_to(x, y);
        }

        context.stroke();

        gtk::Inhibit(false)
    });
}

<commit_msg>Added basic buttons to toggle chart visibility<commit_after>use gtk;
use gtk::prelude::*;
use std::cell::RefCell;
use portaudio as pa;
use std::rc::Rc;
use std::sync::Arc;
use std::sync::RwLock;
use std::io;
use std::io::Write;
use std::thread;
use std::sync::mpsc::*;

struct RustyUi {
    dropdown: gtk::ComboBoxText,
    pitch_label: gtk::Label,
    pitch_error_indicator: gtk::DrawingArea,
    oscilloscope_chart: gtk::DrawingArea,
    freq_chart: gtk::DrawingArea,
    correlation_chart: gtk::DrawingArea,
    oscilloscope_toggle_button: gtk::Button,
    freq_toggle_button: gtk::Button,
    correlation_toggle_button: gtk::Button
}

struct ApplicationState {
    pa: pa::PortAudio,
    pa_stream: Option<pa::Stream<pa::NonBlocking, pa::Input<f32>>>,
    ui: RustyUi
}

struct CrossThreadState {
    pitch: String,
    error: f64,
    signal: Vec<f64>,
    freq_spectrum: Vec<::transforms::FrequencyBucket>,
    correlation: Vec<f64>
}

pub fn start_gui() -> Result<(), String> {
    let pa = try!(::audio::init().map_err(|e| e.to_string()));
    let microphones = try!(::audio::get_device_list(&pa).map_err(|e| e.to_string()));

    try!(gtk::init().map_err(|_| "Failed to initialize GTK."));

    let state = Rc::new(RefCell::new(ApplicationState {
        pa: pa,
        pa_stream: None,
        ui: create_window(microphones)
    }));

    let cross_thread_state = Arc::new(RwLock::new(CrossThreadState {
        pitch: String::new(),
        error: 0.0,
        signal: Vec::new(),
        freq_spectrum: Vec::new(),
        correlation: Vec::new()
    }));

    let (mic_sender, mic_receiver) = channel();

    connect_dropdown_choose_microphone(mic_sender, state.clone());

    start_processing_audio(mic_receiver, cross_thread_state.clone());
    setup_pitch_label_callbacks(state.clone(), cross_thread_state.clone());
    setup_pitch_error_indicator_callbacks(state.clone(), cross_thread_state.clone());
    setup_oscilloscope_drawing_area_callbacks(state.clone(), cross_thread_state.clone());
    setup_freq_drawing_area_callbacks(state.clone(), cross_thread_state.clone());
    setup_correlation_drawing_area_callbacks(state.clone(), cross_thread_state.clone());

    setup_chart_visibility_callbacks(state.clone());

    gtk::main();
    Ok(())
}

fn create_window(microphones: Vec<(u32, String)>) -> RustyUi {
    let window = gtk::Window::new(gtk::WindowType::Toplevel);
    window.set_title("Rusty Microphone");
    window.connect_delete_event(|_, _| {
        gtk::main_quit();
        Inhibit(false)
    });

    let vbox = gtk::Box::new(gtk::Orientation::Vertical, 5);
    window.add(&vbox);

    let hbox = gtk::Box::new(gtk::Orientation::Horizontal, 2);
    vbox.add(&hbox);
    let dropdown = gtk::ComboBoxText::new();
    dropdown.set_hexpand(true);
    set_dropdown_items(&dropdown, microphones);
    hbox.add(&dropdown);

    let oscilloscope_toggle_button = gtk::Button::new_with_label("Osc");
    hbox.add(&oscilloscope_toggle_button);
    let freq_toggle_button = gtk::Button::new_with_label("FFT");
    hbox.add(&freq_toggle_button);
    let correlation_toggle_button = gtk::Button::new_with_label("Corr");
    hbox.add(&correlation_toggle_button);

    let pitch_label = gtk::Label::new(None);
    vbox.add(&pitch_label);

    let pitch_error_indicator = gtk::DrawingArea::new();
    pitch_error_indicator.set_size_request(600, 50);
    vbox.add(&pitch_error_indicator);

    let oscilloscope_chart = gtk::DrawingArea::new();
    oscilloscope_chart.set_size_request(600, 250);
    oscilloscope_chart.set_vexpand(true);
    vbox.add(&oscilloscope_chart);

    let freq_chart = gtk::DrawingArea::new();
    freq_chart.set_size_request(600, 250);
    freq_chart.set_vexpand(true);
    vbox.add(&freq_chart);

    let correlation_chart = gtk::DrawingArea::new();
    correlation_chart.set_size_request(600, 250);
    correlation_chart.set_vexpand(true);
    vbox.add(&correlation_chart);

    window.show_all();

    RustyUi {
        dropdown: dropdown,
        pitch_label: pitch_label,
        pitch_error_indicator: pitch_error_indicator,
        oscilloscope_chart: oscilloscope_chart,
        freq_chart: freq_chart,
        correlation_chart: correlation_chart,
        oscilloscope_toggle_button: oscilloscope_toggle_button,
        freq_toggle_button: freq_toggle_button,
        correlation_toggle_button: correlation_toggle_button
    }
}

fn set_dropdown_items(dropdown: &gtk::ComboBoxText, microphones: Vec<(u32, String)>) {
    for (index, name) in microphones {
        dropdown.append(Some(format!("{}", index).as_ref()), name.as_ref());
    }
}

fn connect_dropdown_choose_microphone(mic_sender: Sender<Vec<f64>>, state: Rc<RefCell<ApplicationState>>) {
    let outer_state = state.clone();
    let ref dropdown = outer_state.borrow().ui.dropdown;
    dropdown.connect_changed(move |dropdown: &gtk::ComboBoxText| {
        match state.borrow_mut().pa_stream {
            Some(ref mut stream) => {stream.stop().ok();},
            _ => {}
        }
        let selected_mic = match dropdown.get_active_id().and_then(|id| id.parse().ok()) {
            Some(mic) => mic,
            None => {return;}
        };
        let stream = ::audio::start_listening(&state.borrow().pa, selected_mic, mic_sender.clone()).ok();
        if stream.is_none() {
            writeln!(io::stderr(), "Failed to open audio channel").ok();
        }
        state.borrow_mut().pa_stream = stream;
    });
}

fn start_processing_audio(mic_receiver: Receiver<Vec<f64>>, cross_thread_state: Arc<RwLock<CrossThreadState>>) {
    thread::spawn(move || {
        loop {
            let mut samples = None;
            loop {
                let next = mic_receiver.try_recv().ok();
                if next.is_none() {
                    break;
                }
                samples = next;
            }
            let samples = match samples {
                Some(samples) => samples,
                None => {continue;}
            };

            let signal = ::transforms::align_to_rising_edge(&samples);
            let frequency_domain = ::transforms::fft(&samples, 44100.0);
            let correlation = ::transforms::correlation(&samples);
            let fundamental = ::transforms::find_fundamental_frequency_correlation(&samples, 44100.0);
            let (pitch, error) = match fundamental {
                Some(fundamental) => (::transforms::hz_to_pitch(fundamental), ::transforms::hz_to_cents_error(fundamental)),
                None => ("".to_string(), 0.0)
            };

            match cross_thread_state.write() {
                Ok(mut state) => {
                    state.pitch = pitch;
                    state.signal = signal;
                    state.freq_spectrum = frequency_domain;
                    state.correlation = correlation;
                    state.error = error
                },
                Err(_) => {}
            };
        }
    });
}

fn setup_pitch_label_callbacks(state: Rc<RefCell<ApplicationState>>, cross_thread_state: Arc<RwLock<CrossThreadState>>) {
    gtk::timeout_add(100, move || {
        let ref pitch = cross_thread_state.read().unwrap().pitch;
        let ref ui = state.borrow().ui;
        ui.pitch_label.set_label(pitch.as_ref());
        ui.pitch_error_indicator.queue_draw();
        ui.oscilloscope_chart.queue_draw();
        ui.correlation_chart.queue_draw();
        ui.freq_chart.queue_draw();

        gtk::Continue(true)
    });
}

fn setup_pitch_error_indicator_callbacks(state: Rc<RefCell<ApplicationState>>, cross_thread_state: Arc<RwLock<CrossThreadState>>) {
    let outer_state = state.clone();
    let ref canvas = outer_state.borrow().ui.pitch_error_indicator;
    canvas.connect_draw(move |ref canvas, ref context| {
        let error = cross_thread_state.read().unwrap().error;
        let width = canvas.get_allocated_width() as f64;
        let midpoint = width / 2.0;
        let height = canvas.get_allocated_height() as f64;

        //flat on the left
        context.set_source_rgb(0.0, 0.0, if error < 0.0 {-error/50.0} else {0.0});
        context.rectangle(0.0, 0.0, midpoint, height);
        context.fill();

        //sharp on the right
        context.set_source_rgb(if error > 0.0 {error/50.0} else {0.0}, 0.0, 0.0);
        context.rectangle(midpoint, 0.0, width, height);
        context.fill();

        gtk::Inhibit(false)
    });
}

fn setup_oscilloscope_drawing_area_callbacks(state: Rc<RefCell<ApplicationState>>, cross_thread_state: Arc<RwLock<CrossThreadState>>) {
    let outer_state = state.clone();
    let ref canvas = outer_state.borrow().ui.oscilloscope_chart;
    canvas.connect_draw(move |ref canvas, ref context| {
        let ref signal = cross_thread_state.read().unwrap().signal;
        let width = canvas.get_allocated_width() as f64;
        let len = 512.0; //Set as a constant so signal won't change size based on zero point.
        let height = canvas.get_allocated_height() as f64;
        let mid_height = height / 2.0;
        let max = signal.iter().map(|x| x.abs()).fold(0.0, |max, x| if max > x { max } else { x });

        context.new_path();
        context.move_to(0.0, mid_height);

        for (i, intensity) in signal.iter().enumerate() {
            let x = i as f64 * width / len;
            let y = mid_height - (intensity * mid_height / max);
            context.line_to(x, y);
        }

        context.stroke();

        gtk::Inhibit(false)
    });
}

fn setup_freq_drawing_area_callbacks(state: Rc<RefCell<ApplicationState>>, cross_thread_state: Arc<RwLock<CrossThreadState>>) {
    let outer_state = state.clone();
    let ref canvas = outer_state.borrow().ui.freq_chart;
    canvas.connect_draw(move |ref canvas, ref context| {
        let ref spectrum = cross_thread_state.read().unwrap().freq_spectrum;
        let width = canvas.get_allocated_width() as f64;
        let height = canvas.get_allocated_height() as f64;
        let max = spectrum.iter().map(|x| x.intensity).fold(0.0, |max, x| if max > x { max } else { x });
        let len = spectrum.len() as f64;

        context.new_path();
        context.move_to(0.0, height);

        for (i, bucket) in spectrum.iter().enumerate() {
            let x = i as f64 * width / len;
            let y = height - (bucket.intensity * height / max);
            context.line_to(x, y);
        }

        context.stroke();

        gtk::Inhibit(false)
    });
}

fn setup_correlation_drawing_area_callbacks(state: Rc<RefCell<ApplicationState>>, cross_thread_state: Arc<RwLock<CrossThreadState>>) {
    let outer_state = state.clone();
    let ref canvas = outer_state.borrow().ui.correlation_chart;
    canvas.connect_draw(move |ref canvas, ref context| {
        let ref correlation = cross_thread_state.read().unwrap().correlation;
        if correlation.len() == 0 {
            return gtk::Inhibit(false);
        }

        let width = canvas.get_allocated_width() as f64;
        let height = canvas.get_allocated_height() as f64;
        let max = correlation[0];
        let len = correlation.len() as f64;

        context.new_path();
        context.move_to(0.0, height);

        for (i, val) in correlation.iter().enumerate() {
            let x = i as f64 * width / len;
            let y = height - (val * height / max);
            context.line_to(x, y);
        }

        context.stroke();

        gtk::Inhibit(false)
    });
}

fn setup_chart_visibility_callbacks(state: Rc<RefCell<ApplicationState>>) {
    let outer_state = state.clone();
    let ref oscilloscope_toggle_button = outer_state.borrow().ui.oscilloscope_toggle_button;
    let ref freq_toggle_button = outer_state.borrow().ui.freq_toggle_button;
    let ref correlation_toggle_button = outer_state.borrow().ui.correlation_toggle_button;

    let oscilloscope_state = state.clone();
    oscilloscope_toggle_button.connect_clicked(move |_| {
        let ref chart = oscilloscope_state.borrow().ui.oscilloscope_chart;
        chart.set_visible(!chart.get_visible());
    });

    let freq_state = state.clone();
    freq_toggle_button.connect_clicked(move |_| {
        let ref chart = freq_state.borrow().ui.freq_chart;
        chart.set_visible(!chart.get_visible());
    });

    let correlation_state = state.clone();
    correlation_toggle_button.connect_clicked(move |_| {
        let ref chart = correlation_state.borrow().ui.correlation_chart;
        chart.set_visible(!chart.get_visible());
    });
}
<|endoftext|>


<commit_before>//!
//! A binding for the library `SDL2_image`
//!
//!
//! Note that you need to build with the
//! feature `image` for this module to be enabled,
//! like so:
//!
//! ```bash
//! $ cargo build --features "image"
//! ```
//!
//! If you want to use this with from inside your own
//! crate, you will need to add this in your Cargo.toml
//!
//! ```toml
//! [dependencies.sdl2]
//! version = ...
//! default-features = false
//! features = ["image"]
//! ```

use std::os::raw::{c_int, c_char};
use std::ffi::CString;
use std::path::Path;
use ::surface::Surface;
use ::render::{Renderer, Texture};
use ::rwops::RWops;
use ::version::Version;
use ::get_error;
use sys;

// Setup linking for all targets.
#[cfg(target_os="macos")]
mod mac {
    #[cfg(any(mac_framework, feature="use_mac_framework"))]
    #[link(kind="framework", name="SDL2_image")]
    extern {}

    #[cfg(not(any(mac_framework, feature="use_mac_framework")))]
    #[link(name="SDL2_image")]
    extern {}
}

#[cfg(any(target_os="windows", target_os="linux", target_os="freebsd"))]
mod others {
    #[link(name="SDL2_image")]
    extern {}
}

#[allow(non_camel_case_types, dead_code)]
mod ffi;

/// InitFlags are passed to init() to control which subsystem
/// functionality to load.
bitflags! {
    pub flags InitFlag : u32 {
        const INIT_JPG  = ffi::IMG_INIT_JPG as u32,
        const INIT_PNG  = ffi::IMG_INIT_PNG as u32,
        const INIT_TIF  = ffi::IMG_INIT_TIF as u32,
        const INIT_WEBP = ffi::IMG_INIT_WEBP as u32
    }
}

// This is used for error message for init
impl ToString for InitFlag {
    fn to_string(&self) -> String {
        let mut string = "".to_string();
        if self.contains(INIT_JPG) {
            string = string + &"INIT_JPG ".to_string();
        }
        if self.contains(INIT_PNG) {
            string = string + &"INIT_PNG ".to_string();
        }
        if self.contains(INIT_TIF) {
            string = string + &"INIT_TIF ".to_string();
        }
        if self.contains(INIT_WEBP) {
            string = string + &"INIT_WEBP ".to_string();
        }
        string
    }
}

/// Static method extensions for creating Surfaces
pub trait LoadSurface: Sized {
    // Self is only returned here to type hint to the compiler.
    // The syntax for type hinting in this case is not yet defined.
    // The intended return value is Result<~Surface, String>.
    fn from_file<P: AsRef<Path>>(filename: P) -> Result<Self, String>;
    fn from_xpm_array(xpm: *const *const i8) -> Result<Self, String>;
}

/// Method extensions to Surface for saving to disk
pub trait SaveSurface {
    fn save<P: AsRef<Path>>(&self, filename: P) -> Result<(), String>;
    fn save_rw(&self, dst: &mut RWops) -> Result<(), String>;
}

impl<'a> LoadSurface for Surface<'a> {
    fn from_file<P: AsRef<Path>>(filename: P) -> Result<Surface<'a>, String> {
        //! Loads an SDL Surface from a file
        unsafe {
            let c_filename = CString::new(filename.as_ref().to_str().unwrap()).unwrap();
            let raw = ffi::IMG_Load(c_filename.as_ptr() as *const _);
            if (raw as *mut ()).is_null() {
                Err(get_error())
            } else {
                Ok(Surface::from_ll(raw))
            }
        }
    }

    fn from_xpm_array(xpm: *const *const i8) -> Result<Surface<'a>, String> {
        //! Loads an SDL Surface from XPM data
        unsafe {
            let raw = ffi::IMG_ReadXPMFromArray(xpm as *const *const c_char);
            if (raw as *mut ()).is_null() {
                Err(get_error())
            } else {
                Ok(Surface::from_ll(raw))
            }
        }
    }
}

impl<'a> SaveSurface for Surface<'a> {
    fn save<P: AsRef<Path>>(&self, filename: P) -> Result<(), String> {
        //! Saves an SDL Surface to a file
        unsafe {
            let c_filename = CString::new(filename.as_ref().to_str().unwrap()).unwrap();
            let status = ffi::IMG_SavePNG(self.raw(), c_filename.as_ptr() as *const _);
            if status != 0 {
                Err(get_error())
            } else {
                Ok(())
            }
        }
    }

    fn save_rw(&self, dst: &mut RWops) -> Result<(), String> {
        //! Saves an SDL Surface to an RWops
        unsafe {
            let status = ffi::IMG_SavePNG_RW(self.raw(), dst.raw(), 0);

            if status != 0 {
                Err(get_error())
            } else {
                Ok(())
            }
        }
    }
}

/// Method extensions for creating Textures from a Renderer
pub trait LoadTexture {
    fn load_texture<P: AsRef<Path>>(&self, filename: P) -> Result<Texture, String>;
}

impl<'a> LoadTexture for Renderer<'a> {
    fn load_texture<P: AsRef<Path>>(&self, filename: P) -> Result<Texture, String> {
        //! Loads an SDL Texture from a file
        unsafe {
            let c_filename = CString::new(filename.as_ref().to_str().unwrap()).unwrap();
            let raw = ffi::IMG_LoadTexture(self.raw(), c_filename.as_ptr() as *const _);
            if (raw as *mut ()).is_null() {
                Err(get_error())
            } else {
                Ok(Texture::from_ll(self, raw))
            }
        }
    }
}

/// Context manager for `sdl2_image` to manage quiting. Can't do much with it but
/// keep it alive while you are using it.
pub struct Sdl2ImageContext;

impl Drop for Sdl2ImageContext {
    fn drop(&mut self) {
        unsafe { ffi::IMG_Quit(); }
    }
}

/// Initializes `SDL2_image` with `InitFlags`.
/// If not every flag is set it returns an error
pub fn init(flags: InitFlag) -> Result<Sdl2ImageContext, String> {
    let return_flags = unsafe {
        let used = ffi::IMG_Init(flags.bits() as c_int);
        InitFlag::from_bits_truncate(used as u32)
    };
    if !flags.intersects(return_flags) {
        // According to docs, error message text is not always set
        if get_error() == "" {
            let un_init_flags = return_flags ^ flags;
            let error_str = &("Could not init: ".to_string() + &un_init_flags.to_string());
            let _ = ::set_error(error_str);
        }
        Err(get_error())
    } else {
        Ok(Sdl2ImageContext)
    }
}

/// Returns the version of the dynamically linked `SDL_image` library
pub fn get_linked_version() -> Version {
    unsafe {
        Version::from_ll(*ffi::IMG_Linked_Version())
    }
}

#[inline]
fn to_surface_result<'a>(raw: *mut sys::surface::SDL_Surface) -> Result<Surface<'a>, String> {
    if (raw as *mut ()).is_null() {
        Err(get_error())
    } else {
        unsafe { Ok(Surface::from_ll(raw)) }
    }
}

pub trait ImageRWops {
    /// load as a surface. except TGA
    fn load(&self) -> Result<Surface, String>;
    /// load as a surface. This can load all supported image formats.
    fn load_typed(&self, _type: &str) -> Result<Surface, String>;

    fn load_cur(&self) -> Result<Surface, String>;
    fn load_ico(&self) -> Result<Surface, String>;
    fn load_bmp(&self) -> Result<Surface, String>;
    fn load_pnm(&self) -> Result<Surface, String>;
    fn load_xpm(&self) -> Result<Surface, String>;
    fn load_xcf(&self) -> Result<Surface, String>;
    fn load_pcx(&self) -> Result<Surface, String>;
    fn load_gif(&self) -> Result<Surface, String>;
    fn load_jpg(&self) -> Result<Surface, String>;
    fn load_tif(&self) -> Result<Surface, String>;
    fn load_png(&self) -> Result<Surface, String>;
    fn load_tga(&self) -> Result<Surface, String>;
    fn load_lbm(&self) -> Result<Surface, String>;
    fn load_xv(&self)  -> Result<Surface, String>;
    fn load_webp(&self) -> Result<Surface, String>;

    fn is_cur(&self) -> bool;
    fn is_ico(&self) -> bool;
    fn is_bmp(&self) -> bool;
    fn is_pnm(&self) -> bool;
    fn is_xpm(&self) -> bool;
    fn is_xcf(&self) -> bool;
    fn is_pcx(&self) -> bool;
    fn is_gif(&self) -> bool;
    fn is_jpg(&self) -> bool;
    fn is_tif(&self) -> bool;
    fn is_png(&self) -> bool;
    fn is_lbm(&self) -> bool;
    fn is_xv(&self)  -> bool;
    fn is_webp(&self) -> bool;
}

impl<'a> ImageRWops for RWops<'a> {
    fn load(&self) -> Result<Surface, String> {
        let raw = unsafe {
            ffi::IMG_Load_RW(self.raw(), 0)
        };
        to_surface_result(raw)
    }
    fn load_typed(&self, _type: &str) -> Result<Surface, String> {
        let raw = unsafe {
            let c_type = CString::new(_type.as_bytes()).unwrap();
            ffi::IMG_LoadTyped_RW(self.raw(), 0, c_type.as_ptr() as *const _)
        };
        to_surface_result(raw)
    }

    fn load_cur(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadCUR_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_ico(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadICO_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_bmp(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadBMP_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_pnm(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadPNM_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_xpm(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadXPM_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_xcf(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadXCF_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_pcx(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadPCX_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_gif(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadGIF_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_jpg(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadJPG_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_tif(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadTIF_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_png(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadPNG_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_tga(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadTGA_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_lbm(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadLBM_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_xv(&self)  -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadXV_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_webp(&self)  -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadWEBP_RW(self.raw()) };
        to_surface_result(raw)
    }

    fn is_cur(&self) -> bool {
        unsafe { ffi::IMG_isCUR(self.raw()) == 1 }
    }
    fn is_ico(&self) -> bool {
        unsafe { ffi::IMG_isICO(self.raw()) == 1 }
    }
    fn is_bmp(&self) -> bool {
        unsafe { ffi::IMG_isBMP(self.raw()) == 1 }
    }
    fn is_pnm(&self) -> bool {
        unsafe { ffi::IMG_isPNM(self.raw()) == 1 }
    }
    fn is_xpm(&self) -> bool {
        unsafe { ffi::IMG_isXPM(self.raw()) == 1 }
    }
    fn is_xcf(&self) -> bool {
        unsafe { ffi::IMG_isXCF(self.raw()) == 1 }
    }
    fn is_pcx(&self) -> bool {
        unsafe { ffi::IMG_isPCX(self.raw()) == 1 }
    }
    fn is_gif(&self) -> bool {
        unsafe { ffi::IMG_isGIF(self.raw()) == 1 }
    }
    fn is_jpg(&self) -> bool {
        unsafe { ffi::IMG_isJPG(self.raw()) == 1 }
    }
    fn is_tif(&self) -> bool {
        unsafe { ffi::IMG_isTIF(self.raw()) == 1 }
    }
    fn is_png(&self) -> bool {
        unsafe { ffi::IMG_isPNG(self.raw()) == 1 }
    }
    fn is_lbm(&self) -> bool {
        unsafe { ffi::IMG_isLBM(self.raw()) == 1 }
    }
    fn is_xv(&self)  -> bool {
        unsafe { ffi::IMG_isXV(self.raw())  == 1 }
    }
    fn is_webp(&self) -> bool {
        unsafe { ffi::IMG_isWEBP(self.raw())  == 1 }
    }
}
<commit_msg>Reduce allocation churn in init error handling<commit_after>//!
//! A binding for the library `SDL2_image`
//!
//!
//! Note that you need to build with the
//! feature `image` for this module to be enabled,
//! like so:
//!
//! ```bash
//! $ cargo build --features "image"
//! ```
//!
//! If you want to use this with from inside your own
//! crate, you will need to add this in your Cargo.toml
//!
//! ```toml
//! [dependencies.sdl2]
//! version = ...
//! default-features = false
//! features = ["image"]
//! ```

use std::os::raw::{c_int, c_char};
use std::ffi::CString;
use std::path::Path;
use ::surface::Surface;
use ::render::{Renderer, Texture};
use ::rwops::RWops;
use ::version::Version;
use ::get_error;
use sys;

// Setup linking for all targets.
#[cfg(target_os="macos")]
mod mac {
    #[cfg(any(mac_framework, feature="use_mac_framework"))]
    #[link(kind="framework", name="SDL2_image")]
    extern {}

    #[cfg(not(any(mac_framework, feature="use_mac_framework")))]
    #[link(name="SDL2_image")]
    extern {}
}

#[cfg(any(target_os="windows", target_os="linux", target_os="freebsd"))]
mod others {
    #[link(name="SDL2_image")]
    extern {}
}

#[allow(non_camel_case_types, dead_code)]
mod ffi;

/// InitFlags are passed to init() to control which subsystem
/// functionality to load.
bitflags! {
    pub flags InitFlag : u32 {
        const INIT_JPG  = ffi::IMG_INIT_JPG as u32,
        const INIT_PNG  = ffi::IMG_INIT_PNG as u32,
        const INIT_TIF  = ffi::IMG_INIT_TIF as u32,
        const INIT_WEBP = ffi::IMG_INIT_WEBP as u32
    }
}

// This is used for error message for init
impl ::std::fmt::Display for InitFlag {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        if self.contains(INIT_JPG) {
            try!(f.write_str("INIT_JPG "));
        }
        if self.contains(INIT_PNG) {
            try!(f.write_str("INIT_PNG "));
        }
        if self.contains(INIT_TIF) {
            try!(f.write_str("INIT_TIF "));
        }
        if self.contains(INIT_WEBP) {
            try!(f.write_str("INIT_WEBP "));
        }
        Ok(())
    }
}

/// Static method extensions for creating Surfaces
pub trait LoadSurface: Sized {
    // Self is only returned here to type hint to the compiler.
    // The syntax for type hinting in this case is not yet defined.
    // The intended return value is Result<~Surface, String>.
    fn from_file<P: AsRef<Path>>(filename: P) -> Result<Self, String>;
    fn from_xpm_array(xpm: *const *const i8) -> Result<Self, String>;
}

/// Method extensions to Surface for saving to disk
pub trait SaveSurface {
    fn save<P: AsRef<Path>>(&self, filename: P) -> Result<(), String>;
    fn save_rw(&self, dst: &mut RWops) -> Result<(), String>;
}

impl<'a> LoadSurface for Surface<'a> {
    fn from_file<P: AsRef<Path>>(filename: P) -> Result<Surface<'a>, String> {
        //! Loads an SDL Surface from a file
        unsafe {
            let c_filename = CString::new(filename.as_ref().to_str().unwrap()).unwrap();
            let raw = ffi::IMG_Load(c_filename.as_ptr() as *const _);
            if (raw as *mut ()).is_null() {
                Err(get_error())
            } else {
                Ok(Surface::from_ll(raw))
            }
        }
    }

    fn from_xpm_array(xpm: *const *const i8) -> Result<Surface<'a>, String> {
        //! Loads an SDL Surface from XPM data
        unsafe {
            let raw = ffi::IMG_ReadXPMFromArray(xpm as *const *const c_char);
            if (raw as *mut ()).is_null() {
                Err(get_error())
            } else {
                Ok(Surface::from_ll(raw))
            }
        }
    }
}

impl<'a> SaveSurface for Surface<'a> {
    fn save<P: AsRef<Path>>(&self, filename: P) -> Result<(), String> {
        //! Saves an SDL Surface to a file
        unsafe {
            let c_filename = CString::new(filename.as_ref().to_str().unwrap()).unwrap();
            let status = ffi::IMG_SavePNG(self.raw(), c_filename.as_ptr() as *const _);
            if status != 0 {
                Err(get_error())
            } else {
                Ok(())
            }
        }
    }

    fn save_rw(&self, dst: &mut RWops) -> Result<(), String> {
        //! Saves an SDL Surface to an RWops
        unsafe {
            let status = ffi::IMG_SavePNG_RW(self.raw(), dst.raw(), 0);

            if status != 0 {
                Err(get_error())
            } else {
                Ok(())
            }
        }
    }
}

/// Method extensions for creating Textures from a Renderer
pub trait LoadTexture {
    fn load_texture<P: AsRef<Path>>(&self, filename: P) -> Result<Texture, String>;
}

impl<'a> LoadTexture for Renderer<'a> {
    fn load_texture<P: AsRef<Path>>(&self, filename: P) -> Result<Texture, String> {
        //! Loads an SDL Texture from a file
        unsafe {
            let c_filename = CString::new(filename.as_ref().to_str().unwrap()).unwrap();
            let raw = ffi::IMG_LoadTexture(self.raw(), c_filename.as_ptr() as *const _);
            if (raw as *mut ()).is_null() {
                Err(get_error())
            } else {
                Ok(Texture::from_ll(self, raw))
            }
        }
    }
}

/// Context manager for `sdl2_image` to manage quiting. Can't do much with it but
/// keep it alive while you are using it.
pub struct Sdl2ImageContext;

impl Drop for Sdl2ImageContext {
    fn drop(&mut self) {
        unsafe { ffi::IMG_Quit(); }
    }
}

/// Initializes `SDL2_image` with `InitFlags`.
/// If not every flag is set it returns an error
pub fn init(flags: InitFlag) -> Result<Sdl2ImageContext, String> {
    let return_flags = unsafe {
        let used = ffi::IMG_Init(flags.bits() as c_int);
        InitFlag::from_bits_truncate(used as u32)
    };
    if !flags.intersects(return_flags) {
        // According to docs, error message text is not always set
        let mut error = get_error();
        if error.is_empty() {
            let un_init_flags = return_flags ^ flags;
            error = format!("Could not init: {}", un_init_flags);
            let _ = ::set_error(&error);
        }
        Err(error)
    } else {
        Ok(Sdl2ImageContext)
    }
}

/// Returns the version of the dynamically linked `SDL_image` library
pub fn get_linked_version() -> Version {
    unsafe {
        Version::from_ll(*ffi::IMG_Linked_Version())
    }
}

#[inline]
fn to_surface_result<'a>(raw: *mut sys::surface::SDL_Surface) -> Result<Surface<'a>, String> {
    if (raw as *mut ()).is_null() {
        Err(get_error())
    } else {
        unsafe { Ok(Surface::from_ll(raw)) }
    }
}

pub trait ImageRWops {
    /// load as a surface. except TGA
    fn load(&self) -> Result<Surface, String>;
    /// load as a surface. This can load all supported image formats.
    fn load_typed(&self, _type: &str) -> Result<Surface, String>;

    fn load_cur(&self) -> Result<Surface, String>;
    fn load_ico(&self) -> Result<Surface, String>;
    fn load_bmp(&self) -> Result<Surface, String>;
    fn load_pnm(&self) -> Result<Surface, String>;
    fn load_xpm(&self) -> Result<Surface, String>;
    fn load_xcf(&self) -> Result<Surface, String>;
    fn load_pcx(&self) -> Result<Surface, String>;
    fn load_gif(&self) -> Result<Surface, String>;
    fn load_jpg(&self) -> Result<Surface, String>;
    fn load_tif(&self) -> Result<Surface, String>;
    fn load_png(&self) -> Result<Surface, String>;
    fn load_tga(&self) -> Result<Surface, String>;
    fn load_lbm(&self) -> Result<Surface, String>;
    fn load_xv(&self)  -> Result<Surface, String>;
    fn load_webp(&self) -> Result<Surface, String>;

    fn is_cur(&self) -> bool;
    fn is_ico(&self) -> bool;
    fn is_bmp(&self) -> bool;
    fn is_pnm(&self) -> bool;
    fn is_xpm(&self) -> bool;
    fn is_xcf(&self) -> bool;
    fn is_pcx(&self) -> bool;
    fn is_gif(&self) -> bool;
    fn is_jpg(&self) -> bool;
    fn is_tif(&self) -> bool;
    fn is_png(&self) -> bool;
    fn is_lbm(&self) -> bool;
    fn is_xv(&self)  -> bool;
    fn is_webp(&self) -> bool;
}

impl<'a> ImageRWops for RWops<'a> {
    fn load(&self) -> Result<Surface, String> {
        let raw = unsafe {
            ffi::IMG_Load_RW(self.raw(), 0)
        };
        to_surface_result(raw)
    }
    fn load_typed(&self, _type: &str) -> Result<Surface, String> {
        let raw = unsafe {
            let c_type = CString::new(_type.as_bytes()).unwrap();
            ffi::IMG_LoadTyped_RW(self.raw(), 0, c_type.as_ptr() as *const _)
        };
        to_surface_result(raw)
    }

    fn load_cur(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadCUR_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_ico(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadICO_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_bmp(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadBMP_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_pnm(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadPNM_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_xpm(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadXPM_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_xcf(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadXCF_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_pcx(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadPCX_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_gif(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadGIF_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_jpg(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadJPG_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_tif(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadTIF_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_png(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadPNG_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_tga(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadTGA_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_lbm(&self) -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadLBM_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_xv(&self)  -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadXV_RW(self.raw()) };
        to_surface_result(raw)
    }
    fn load_webp(&self)  -> Result<Surface, String> {
        let raw = unsafe { ffi::IMG_LoadWEBP_RW(self.raw()) };
        to_surface_result(raw)
    }

    fn is_cur(&self) -> bool {
        unsafe { ffi::IMG_isCUR(self.raw()) == 1 }
    }
    fn is_ico(&self) -> bool {
        unsafe { ffi::IMG_isICO(self.raw()) == 1 }
    }
    fn is_bmp(&self) -> bool {
        unsafe { ffi::IMG_isBMP(self.raw()) == 1 }
    }
    fn is_pnm(&self) -> bool {
        unsafe { ffi::IMG_isPNM(self.raw()) == 1 }
    }
    fn is_xpm(&self) -> bool {
        unsafe { ffi::IMG_isXPM(self.raw()) == 1 }
    }
    fn is_xcf(&self) -> bool {
        unsafe { ffi::IMG_isXCF(self.raw()) == 1 }
    }
    fn is_pcx(&self) -> bool {
        unsafe { ffi::IMG_isPCX(self.raw()) == 1 }
    }
    fn is_gif(&self) -> bool {
        unsafe { ffi::IMG_isGIF(self.raw()) == 1 }
    }
    fn is_jpg(&self) -> bool {
        unsafe { ffi::IMG_isJPG(self.raw()) == 1 }
    }
    fn is_tif(&self) -> bool {
        unsafe { ffi::IMG_isTIF(self.raw()) == 1 }
    }
    fn is_png(&self) -> bool {
        unsafe { ffi::IMG_isPNG(self.raw()) == 1 }
    }
    fn is_lbm(&self) -> bool {
        unsafe { ffi::IMG_isLBM(self.raw()) == 1 }
    }
    fn is_xv(&self)  -> bool {
        unsafe { ffi::IMG_isXV(self.raw())  == 1 }
    }
    fn is_webp(&self) -> bool {
        unsafe { ffi::IMG_isWEBP(self.raw())  == 1 }
    }
}
<|endoftext|>


<commit_before>//! https://en.wikipedia.org/wiki/Address_Resolution_Protocol

use alloc::prelude::v1::*;

use crate::{EtherType, Ipv4Addr, MacAddr};

/// Only supports Ethernet with MAC addresses and IPv4
/// https://en.wikipedia.org/wiki/Address_Resolution_Protocol#Packet_structure
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Packet {
    pub ptype: EtherType,
    pub operation: Operation,
    pub sender_hw: MacAddr,
    pub sender_ip: Ipv4Addr,
    pub target_hw: MacAddr,
    pub target_ip: Ipv4Addr,
}
impl Packet {
    pub fn from_bytes(input: &[u8]) -> Self {
        assert!(&input[..2] == [0, 1]); // Ethernet only

        let hlen = input[4] as usize;
        let plen = input[5] as usize;

        assert!(hlen == 6); // Mac addresses only
        assert!(plen == 4); // Ipv4

        Self {
            ptype: EtherType::from_bytes(&input[2..4]),
            operation: Operation::from_bytes(&input[6..8]),
            sender_hw: MacAddr::from_bytes(&input[8..8 + hlen]),
            sender_ip: Ipv4Addr::from_bytes(&input[8 + hlen..8 + hlen + plen]),
            target_hw: MacAddr::from_bytes(&input[8 + hlen + plen..8 + hlen * 2 + plen]),
            target_ip: Ipv4Addr::from_bytes(&input[8 + hlen * 2 + plen..8 + hlen * 2 + plen * 2]),
        }
    }

    /// https://en.wikipedia.org/wiki/Address_Resolution_Protocol#Packet_structure
    pub fn to_bytes(self) -> Vec<u8> {
        let mut result = Vec::new();
        // HTYPE: Ethernet
        result.push(0);
        result.push(1);
        // PTYPE: Ipv4
        result.push(8);
        result.push(0);
        // HLEN: 6 (MacAddr)
        result.push(6);
        // PLEN: 4 (Ipv4)
        result.push(4);
        // Operation
        result.push(0);
        result.push(match self.operation {
            Operation::Request => 1,
            Operation::Reply => 2,
        });
        // Sender MacAddr
        result.extend(&self.sender_hw.0);
        // Sender Ipv4
        result.extend(&self.sender_ip.0);
        // Target MacAddr
        result.extend(&self.target_hw.0);
        // Target Ipv4
        result.extend(&self.target_ip.0);
        // Return
        result
    }

    pub fn is_request(&self) -> bool {
        self.operation == Operation::Request
    }

    pub fn to_reply(mut self, mac: MacAddr, ip: Ipv4Addr) -> Self {
        assert!(self.is_request());

        self.operation == Operation::Reply;
        self.target_hw = mac;
        self.target_ip = ip;

        self
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Operation {
    Request,
    Reply,
}
impl Operation {
    pub fn from_bytes(bytes: &[u8]) -> Self {
        assert!(bytes[0] == 0);
        match bytes[1] {
            1 => Self::Request,
            2 => Self::Reply,
            other => panic!("Unknown ARP operation {}", other),
        }
    }
}

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

    #[test]
    fn test_parse() {
        let example: Vec<u8> = vec![
            0, 1, 8, 0, 6, 4, 0, 1, 1, 2, 3, 4, 5, 6, 10, 0, 2, 2, 0, 0, 0, 0, 0, 0, 10, 0, 2, 15,
        ];

        let packet = Packet::from_bytes(&example);

        assert_eq!(packet, Packet {
            ptype: EtherType::Ipv4,
            operation: Operation::Request,
            sender_hw: MacAddr::from_bytes(&[1, 2, 3, 4, 5, 6]),
            sender_ip: Ipv4Addr::from_bytes(&[10, 0, 2, 2]),
            target_hw: MacAddr::from_bytes(&[0, 0, 0, 0, 0, 0]),
            target_ip: Ipv4Addr::from_bytes(&[10, 0, 2, 15]),
        });

        assert_eq!(packet.to_bytes(), example);
    }
}
<commit_msg>Fix invalid implementation of arp::Packet::reply<commit_after>//! https://en.wikipedia.org/wiki/Address_Resolution_Protocol

use alloc::prelude::v1::*;

use crate::{EtherType, Ipv4Addr, MacAddr};

/// Only supports Ethernet with MAC addresses and IPv4
/// https://en.wikipedia.org/wiki/Address_Resolution_Protocol#Packet_structure
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Packet {
    pub ptype: EtherType,
    pub operation: Operation,
    pub sender_hw: MacAddr,
    pub sender_ip: Ipv4Addr,
    pub target_hw: MacAddr,
    pub target_ip: Ipv4Addr,
}
impl Packet {
    pub fn from_bytes(input: &[u8]) -> Self {
        assert!(&input[..2] == [0, 1]); // Ethernet only

        let hlen = input[4] as usize;
        let plen = input[5] as usize;

        assert!(hlen == 6); // Mac addresses only
        assert!(plen == 4); // Ipv4

        Self {
            ptype: EtherType::from_bytes(&input[2..4]),
            operation: Operation::from_bytes(&input[6..8]),
            sender_hw: MacAddr::from_bytes(&input[8..8 + hlen]),
            sender_ip: Ipv4Addr::from_bytes(&input[8 + hlen..8 + hlen + plen]),
            target_hw: MacAddr::from_bytes(&input[8 + hlen + plen..8 + hlen * 2 + plen]),
            target_ip: Ipv4Addr::from_bytes(&input[8 + hlen * 2 + plen..8 + hlen * 2 + plen * 2]),
        }
    }

    /// https://en.wikipedia.org/wiki/Address_Resolution_Protocol#Packet_structure
    pub fn to_bytes(self) -> Vec<u8> {
        let mut result = Vec::new();
        // HTYPE: Ethernet
        result.push(0);
        result.push(1);
        // PTYPE: Ipv4
        result.push(8);
        result.push(0);
        // HLEN: 6 (MacAddr)
        result.push(6);
        // PLEN: 4 (Ipv4)
        result.push(4);
        // Operation
        result.push(0);
        result.push(match self.operation {
            Operation::Request => 1,
            Operation::Reply => 2,
        });
        // Sender MacAddr
        result.extend(&self.sender_hw.0);
        // Sender Ipv4
        result.extend(&self.sender_ip.0);
        // Target MacAddr
        result.extend(&self.target_hw.0);
        // Target Ipv4
        result.extend(&self.target_ip.0);
        // Return
        result
    }

    pub fn is_request(&self) -> bool {
        self.operation == Operation::Request
    }

    pub fn to_reply(mut self, mac: MacAddr, ip: Ipv4Addr) -> Self {
        assert!(self.is_request());

        self.operation = Operation::Reply;
        self.target_hw = self.sender_hw;
        self.target_ip = self.sender_ip;
        self.sender_hw = mac;
        self.sender_ip = ip;

        self
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Operation {
    Request,
    Reply,
}
impl Operation {
    pub fn from_bytes(bytes: &[u8]) -> Self {
        assert!(bytes[0] == 0);
        match bytes[1] {
            1 => Self::Request,
            2 => Self::Reply,
            other => panic!("Unknown ARP operation {}", other),
        }
    }
}

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

    #[test]
    fn test_parse() {
        let example: Vec<u8> = vec![
            0, 1, 8, 0, 6, 4, 0, 1, 1, 2, 3, 4, 5, 6, 10, 0, 2, 2, 0, 0, 0, 0, 0, 0, 10, 0, 2, 15,
        ];

        let packet = Packet::from_bytes(&example);

        assert_eq!(packet, Packet {
            ptype: EtherType::Ipv4,
            operation: Operation::Request,
            sender_hw: MacAddr::from_bytes(&[1, 2, 3, 4, 5, 6]),
            sender_ip: Ipv4Addr::from_bytes(&[10, 0, 2, 2]),
            target_hw: MacAddr::from_bytes(&[0, 0, 0, 0, 0, 0]),
            target_ip: Ipv4Addr::from_bytes(&[10, 0, 2, 15]),
        });

        assert_eq!(packet.to_bytes(), example);
    }
}
<|endoftext|>