<commit_before>pub use libc::dev_t;
pub use libc::stat as FileStat;

use {Errno, Result, NixPath};
use libc::{self, mode_t};
use std::mem;
use std::os::unix::io::RawFd;

mod ffi {
    use libc::{c_char, c_int, mode_t, dev_t};
    pub use libc::{stat, fstat, lstat};

    extern {
        pub fn mknod(pathname: *const c_char, mode: mode_t, dev: dev_t) -> c_int;
        pub fn umask(mask: mode_t) -> mode_t;
    }
}

bitflags!(
    flags SFlag: mode_t {
        const S_IFIFO = libc::S_IFIFO,
        const S_IFCHR = libc::S_IFCHR,
        const S_IFDIR = libc::S_IFDIR,
        const S_IFBLK = libc::S_IFBLK,
        const S_IFREG = libc::S_IFREG,
        const S_IFLNK = libc::S_IFLNK,
        const S_IFSOCK = libc::S_IFSOCK,
        const S_IFMT = libc::S_IFMT,
    }
);

bitflags! {
    flags Mode: mode_t {
        const S_IRWXU = libc::S_IRWXU,
        const S_IRUSR = libc::S_IRUSR,
        const S_IWUSR = libc::S_IWUSR,
        const S_IXUSR = libc::S_IXUSR,

        const S_IRWXG = libc::S_IRWXG,
        const S_IRGRP = libc::S_IRGRP,
        const S_IWGRP = libc::S_IWGRP,
        const S_IXGRP = libc::S_IXGRP,

        const S_IRWXO = libc::S_IRWXO,
        const S_IROTH = libc::S_IROTH,
        const S_IWOTH = libc::S_IWOTH,
        const S_IXOTH = libc::S_IXOTH,

        const S_ISUID = 0o4000,
        const S_ISGID = 0o2000,
        const S_ISVTX = 0o1000,
    }
}

pub fn mknod<P: ?Sized + NixPath>(path: &P, kind: SFlag, perm: Mode, dev: dev_t) -> Result<()> {
    let res = try!(path.with_nix_path(|cstr| {
        unsafe {
            ffi::mknod(cstr.as_ptr(), kind.bits | perm.bits() as mode_t, dev)
        }
    }));

    Errno::result(res).map(drop)
}

#[cfg(target_os = "linux")]
const MINORBITS: usize = 20;

#[cfg(target_os = "linux")]
pub fn mkdev(major: u64, minor: u64) -> dev_t {
    (major << MINORBITS) | minor
}

pub fn umask(mode: Mode) -> Mode {
    let prev = unsafe { ffi::umask(mode.bits() as mode_t) };
    Mode::from_bits(prev).expect("[BUG] umask returned invalid Mode")
}

pub fn stat<P: ?Sized + NixPath>(path: &P) -> Result<FileStat> {
    let mut dst = unsafe { mem::uninitialized() };
    let res = try!(path.with_nix_path(|cstr| {
        unsafe {
            ffi::stat(cstr.as_ptr(), &mut dst as *mut FileStat)
        }
    }));

    try!(Errno::result(res));

    Ok(dst)
}

pub fn lstat<P: ?Sized + NixPath>(path: &P) -> Result<FileStat> {
    let mut dst = unsafe { mem::uninitialized() };
    let res = try!(path.with_nix_path(|cstr| {
        unsafe {
            ffi::lstat(cstr.as_ptr(), &mut dst as *mut FileStat)
        }
    }));

    try!(Errno::result(res));

    Ok(dst)
}

pub fn fstat(fd: RawFd) -> Result<FileStat> {
    let mut dst = unsafe { mem::uninitialized() };
    let res = unsafe { ffi::fstat(fd, &mut dst as *mut FileStat) };

    try!(Errno::result(res));

    Ok(dst)
}
<commit_msg>Last few constants<commit_after>pub use libc::dev_t;
pub use libc::stat as FileStat;

use {Errno, Result, NixPath};
use libc::{self, mode_t};
use std::mem;
use std::os::unix::io::RawFd;

mod ffi {
    use libc::{c_char, c_int, mode_t, dev_t};
    pub use libc::{stat, fstat, lstat};

    extern {
        pub fn mknod(pathname: *const c_char, mode: mode_t, dev: dev_t) -> c_int;
        pub fn umask(mask: mode_t) -> mode_t;
    }
}

bitflags!(
    flags SFlag: mode_t {
        const S_IFIFO = libc::S_IFIFO,
        const S_IFCHR = libc::S_IFCHR,
        const S_IFDIR = libc::S_IFDIR,
        const S_IFBLK = libc::S_IFBLK,
        const S_IFREG = libc::S_IFREG,
        const S_IFLNK = libc::S_IFLNK,
        const S_IFSOCK = libc::S_IFSOCK,
        const S_IFMT = libc::S_IFMT,
    }
);

bitflags! {
    flags Mode: mode_t {
        const S_IRWXU = libc::S_IRWXU,
        const S_IRUSR = libc::S_IRUSR,
        const S_IWUSR = libc::S_IWUSR,
        const S_IXUSR = libc::S_IXUSR,

        const S_IRWXG = libc::S_IRWXG,
        const S_IRGRP = libc::S_IRGRP,
        const S_IWGRP = libc::S_IWGRP,
        const S_IXGRP = libc::S_IXGRP,

        const S_IRWXO = libc::S_IRWXO,
        const S_IROTH = libc::S_IROTH,
        const S_IWOTH = libc::S_IWOTH,
        const S_IXOTH = libc::S_IXOTH,

        const S_ISUID = libc::S_ISUID as mode_t,
        const S_ISGID = libc::S_ISGID as mode_t,
        const S_ISVTX = libc::S_ISVTX as mode_t,
    }
}

pub fn mknod<P: ?Sized + NixPath>(path: &P, kind: SFlag, perm: Mode, dev: dev_t) -> Result<()> {
    let res = try!(path.with_nix_path(|cstr| {
        unsafe {
            ffi::mknod(cstr.as_ptr(), kind.bits | perm.bits() as mode_t, dev)
        }
    }));

    Errno::result(res).map(drop)
}

#[cfg(target_os = "linux")]
const MINORBITS: usize = 20;

#[cfg(target_os = "linux")]
pub fn mkdev(major: u64, minor: u64) -> dev_t {
    (major << MINORBITS) | minor
}

pub fn umask(mode: Mode) -> Mode {
    let prev = unsafe { ffi::umask(mode.bits() as mode_t) };
    Mode::from_bits(prev).expect("[BUG] umask returned invalid Mode")
}

pub fn stat<P: ?Sized + NixPath>(path: &P) -> Result<FileStat> {
    let mut dst = unsafe { mem::uninitialized() };
    let res = try!(path.with_nix_path(|cstr| {
        unsafe {
            ffi::stat(cstr.as_ptr(), &mut dst as *mut FileStat)
        }
    }));

    try!(Errno::result(res));

    Ok(dst)
}

pub fn lstat<P: ?Sized + NixPath>(path: &P) -> Result<FileStat> {
    let mut dst = unsafe { mem::uninitialized() };
    let res = try!(path.with_nix_path(|cstr| {
        unsafe {
            ffi::lstat(cstr.as_ptr(), &mut dst as *mut FileStat)
        }
    }));

    try!(Errno::result(res));

    Ok(dst)
}

pub fn fstat(fd: RawFd) -> Result<FileStat> {
    let mut dst = unsafe { mem::uninitialized() };
    let res = unsafe { ffi::fstat(fd, &mut dst as *mut FileStat) };

    try!(Errno::result(res));

    Ok(dst)
}
<|endoftext|>


<commit_before>/* Copyright (C)  2016 nokaa <nokaa@cock.li>
 * This software is licensed under the terms of the
 * GNU Affero General Public License. You should have
 * received a copy of this license with this software.
 * The license may also be found at https://gnu.org/licenses/agpl.txt
 */

use rustty::{Terminal, Event, Color};
use editor;

/// Represents an `x, y` coordinate.
#[derive(Copy, Clone)]
struct Position {
    x: usize,
    y: usize,
}

/// Represents the cursor in our terminal
struct Cursor {
    pos: Position,
    lpos: Position,
    color: Color,
}

/// Launches the terminal using the given `editor`.
pub fn run(editor: &mut editor::Editor) {
    let mut cursor = Cursor {
        pos: Position {x: 0, y: 0},
        lpos: Position {x: 0, y: 0},
        color: Color::Red,
    };

    let mut term = Terminal::new().unwrap();
    term[(cursor.pos.x, cursor.pos.y)].set_bg(cursor.color);
    prompt(&mut term, editor.filename());
    term.swap_buffers().unwrap();

    loop {
        let evt = term.get_event(100).unwrap();
        if let Some(Event::Key(ch)) = evt {
            if editor.command_mode() {
                match ch {
                    'i' => {
                        editor.switch_mode();
                    }
                    'h' => {
                        if cursor.pos.x != 0 {
                            cursor.lpos = cursor.pos;
                            cursor.pos.x -= 1;
                        }
                    }
                    'j' => {
                        // TODO(nokaa): We don't want to go beyond
                        // the working area here.
                        cursor.lpos = cursor.pos;
                        cursor.pos.y += 1;

                    }
                    'k' => {
                        if cursor.pos.y != 0 {
                            cursor.lpos = cursor.pos;
                            cursor.pos.y -= 1;
                        }
                    }
                    'l' => {
                        // TODO(nokaa): We don't want to extend beyond the
                        // line length here, but we first need a way to
                        // determine a given line's length.
                        cursor.lpos = cursor.pos;
                        cursor.pos.x += 1;
                    }
                    'q' => {
                        break;
                    }
                    _ => { }
                }
            } else {
                match ch {
                    '\x1b' => { // Escape key
                        editor.switch_mode();
                    }
                    '\x7f' => { // Backspace key
                        cursor.lpos = cursor.pos;
                        if cursor.pos.x == 0 {
                            cursor.pos.y = cursor.pos.y.saturating_sub(1);
                        } else {
                            cursor.pos.x -= 1;
                        }
                        term[(cursor.pos.x, cursor.pos.y)].set_ch(' ');
                    }
                    '\r' => {
                        cursor.lpos = cursor.pos;
                        cursor.pos.x = 0;
                        cursor.pos.y += 1;
                    }
                    '\t' => {
                        term[(cursor.pos.x, cursor.pos.y)].set_ch(' ');
                        term[(cursor.pos.x + 1, cursor.pos.y)].set_ch(' ');
                        term[(cursor.pos.x + 2, cursor.pos.y)].set_ch(' ');
                        term[(cursor.pos.x + 3, cursor.pos.y)].set_ch(' ');
                        cursor.lpos = cursor.pos;
                        cursor.pos.x += 4;
                    }
                    c @ _ => {
                        term[(cursor.pos.x, cursor.pos.y)].set_ch(c);
                        cursor.lpos = cursor.pos;
                        cursor.pos.x += 1;
                    }
                }
            }

            if cursor.pos.x >= term.cols() - 1 {
                term[(cursor.lpos.x, cursor.lpos.y)].set_bg(Color::Default);
                cursor.lpos = cursor.pos;
                cursor.pos.x = 0;
                cursor.pos.y += 1;
            }
            if cursor.pos.y >= term.rows() - 1 {
                term[(cursor.lpos.x, cursor.lpos.y)].set_bg(Color::Default);
                cursor.lpos = cursor.pos;
                cursor.pos.x = 0;
                cursor.pos.y = 0;
            }

            term[(cursor.lpos.x, cursor.lpos.y)].set_bg(Color::Default);
            term[(cursor.pos.x, cursor.pos.y)].set_bg(cursor.color);
            term.swap_buffers().unwrap();
        }
    }
}

fn prompt(term: &mut Terminal, filename: String) {
    let w = term.cols();
    let h = term.rows();

    for i in 0..w {
        term[(i, h - 2)].set_bg(Color::Red);
    }

    for (i, c) in filename.chars().enumerate() {
        term[(i, h - 2)].set_ch(c);
    }
}
<commit_msg>Refactor tab handling<commit_after>/* Copyright (C)  2016 nokaa <nokaa@cock.li>
 * This software is licensed under the terms of the
 * GNU Affero General Public License. You should have
 * received a copy of this license with this software.
 * The license may also be found at https://gnu.org/licenses/agpl.txt
 */

use rustty::{Terminal, Event, Color};
use editor;

/// Represents an `x, y` coordinate.
#[derive(Copy, Clone)]
struct Position {
    x: usize,
    y: usize,
}

/// Represents the cursor in our terminal
struct Cursor {
    pos: Position,
    lpos: Position,
    color: Color,
}

/// Launches the terminal using the given `editor`.
pub fn run(editor: &mut editor::Editor) {
    let mut cursor = Cursor {
        pos: Position {x: 0, y: 0},
        lpos: Position {x: 0, y: 0},
        color: Color::Red,
    };

    let mut term = Terminal::new().unwrap();
    term[(cursor.pos.x, cursor.pos.y)].set_bg(cursor.color);
    prompt(&mut term, editor.filename());
    term.swap_buffers().unwrap();

    loop {
        let evt = term.get_event(100).unwrap();
        if let Some(Event::Key(ch)) = evt {
            if editor.command_mode() {
                match ch {
                    'i' => {
                        editor.switch_mode();
                    }
                    'h' => {
                        if cursor.pos.x != 0 {
                            cursor.lpos = cursor.pos;
                            cursor.pos.x -= 1;
                        }
                    }
                    'j' => {
                        // TODO(nokaa): We don't want to go beyond
                        // the working area here.
                        cursor.lpos = cursor.pos;
                        cursor.pos.y += 1;

                    }
                    'k' => {
                        if cursor.pos.y != 0 {
                            cursor.lpos = cursor.pos;
                            cursor.pos.y -= 1;
                        }
                    }
                    'l' => {
                        // TODO(nokaa): We don't want to extend beyond the
                        // line length here, but we first need a way to
                        // determine a given line's length.
                        cursor.lpos = cursor.pos;
                        cursor.pos.x += 1;
                    }
                    'q' => {
                        break;
                    }
                    _ => { }
                }
            } else {
                match ch {
                    '\x1b' => { // Escape key
                        editor.switch_mode();
                    }
                    '\x7f' => { // Backspace key
                        cursor.lpos = cursor.pos;
                        if cursor.pos.x == 0 {
                            cursor.pos.y = cursor.pos.y.saturating_sub(1);
                        } else {
                            cursor.pos.x -= 1;
                        }
                        term[(cursor.pos.x, cursor.pos.y)].set_ch(' ');
                    }
                    '\r' => {
                        cursor.lpos = cursor.pos;
                        cursor.pos.x = 0;
                        cursor.pos.y += 1;
                    }
                    '\t' => {
                        for i in 0..4 {
                            term[(cursor.pos.x + i, cursor.pos.y)].set_ch(' ');
                        }
                        cursor.lpos = cursor.pos;
                        cursor.pos.x += 4;
                    }
                    c @ _ => {
                        term[(cursor.pos.x, cursor.pos.y)].set_ch(c);
                        cursor.lpos = cursor.pos;
                        cursor.pos.x += 1;
                    }
                }
            }

            if cursor.pos.x >= term.cols() - 1 {
                term[(cursor.lpos.x, cursor.lpos.y)].set_bg(Color::Default);
                cursor.lpos = cursor.pos;
                cursor.pos.x = 0;
                cursor.pos.y += 1;
            }
            if cursor.pos.y >= term.rows() - 1 {
                term[(cursor.lpos.x, cursor.lpos.y)].set_bg(Color::Default);
                cursor.lpos = cursor.pos;
                cursor.pos.x = 0;
                cursor.pos.y = 0;
            }

            term[(cursor.lpos.x, cursor.lpos.y)].set_bg(Color::Default);
            term[(cursor.pos.x, cursor.pos.y)].set_bg(cursor.color);
            term.swap_buffers().unwrap();
        }
    }
}

fn prompt(term: &mut Terminal, filename: String) {
    let w = term.cols();
    let h = term.rows();

    for i in 0..w {
        term[(i, h - 2)].set_bg(Color::Red);
    }

    for (i, c) in filename.chars().enumerate() {
        term[(i, h - 2)].set_ch(c);
    }
}
<|endoftext|>


<commit_before>//! Filter

use std::{cell::RefCell, collections::VecDeque, rc::Rc};

/// Holder of global dispatch-related data
///
/// This struct serves as a dynamic container for the dispatch-time
/// global data that you gave to the dispatch method, and is given as
/// input to all your callbacks. It allows you to share global state
/// between your filters.
///
/// The main method of interest is the `get` method, which allows you to
/// access a `&mut _` reference to the global data itself. The other methods
/// are mostly used internally by the crate.
pub struct DispatchData<'a> {
    data: &'a mut dyn std::any::Any,
}

impl<'a> DispatchData<'a> {
    /// Access the dispatch data knowing its type
    ///
    /// Will return `None` if the provided type is not the correct
    /// inner type.
    pub fn get<T: std::any::Any>(&mut self) -> Option<&mut T> {
        self.data.downcast_mut()
    }

    /// Wrap a mutable reference
    ///
    /// This creates a new `DispatchData` from a mutable reference
    pub fn wrap<T: std::any::Any>(data: &'a mut T) -> DispatchData<'a> {
        DispatchData { data }
    }

    /// Reborrows this `DispatchData` to create a new one with the same content
    ///
    /// This is a quick and cheap way to propagate the `DispatchData` down a
    /// callback stack by value. It is basically a noop only there to ease
    /// work with the borrow checker.
    pub fn reborrow(&mut self) -> DispatchData {
        DispatchData {
            data: &mut *self.data,
        }
    }
}

struct Inner<E, F: ?Sized> {
    pending: RefCell<VecDeque<E>>,
    cb: RefCell<F>,
}

type DynInner<E> = Inner<E, dyn FnMut(E, &Filter<E>, DispatchData<'_>)>;

/// An event filter
///
/// Can be used in wayland-client and wayland-server to aggregate
/// messages from different objects into the same closure.
///
/// You need to provide it a closure of type `FnMut(E, &Filter<E>)`,
/// which will be called eny time a message is sent to the filter
/// via the `send(..)` method. Your closure also receives a handle
/// to the filter as argument, so that you can use it from within
/// the callback (to assign new wayland objects to this filter for
/// example).
///
/// The `Filter` can be cloned, and all clones send messages to the
/// same closure. However it is not threadsafe.
pub struct Filter<E> {
    inner: Rc<DynInner<E>>,
}

impl<E> Clone for Filter<E> {
    fn clone(&self) -> Filter<E> {
        Filter {
            inner: self.inner.clone(),
        }
    }
}

impl<E> Filter<E> {
    /// Create a new filter from given closure
    pub fn new<F: FnMut(E, &Filter<E>, DispatchData<'_>) + 'static>(f: F) -> Filter<E> {
        Filter {
            inner: Rc::new(Inner {
                pending: RefCell::new(VecDeque::new()),
                cb: RefCell::new(f),
            }),
        }
    }

    /// Send a message to this filter
    pub fn send(&self, evt: E, mut data: DispatchData) {
        // gracefully handle reentrancy
        if let Ok(mut guard) = self.inner.cb.try_borrow_mut() {
            (&mut *guard)(evt, self, data.reborrow());
            // process all events that might have been enqueued by the cb
            while let Some(evt) = self.inner.pending.borrow_mut().pop_front() {
                (&mut *guard)(evt, self, data.reborrow());
            }
        } else {
            self.inner.pending.borrow_mut().push_back(evt);
        }
    }
}
<commit_msg>Fix a documentation typo<commit_after>//! Filter

use std::{cell::RefCell, collections::VecDeque, rc::Rc};

/// Holder of global dispatch-related data
///
/// This struct serves as a dynamic container for the dispatch-time
/// global data that you gave to the dispatch method, and is given as
/// input to all your callbacks. It allows you to share global state
/// between your filters.
///
/// The main method of interest is the `get` method, which allows you to
/// access a `&mut _` reference to the global data itself. The other methods
/// are mostly used internally by the crate.
pub struct DispatchData<'a> {
    data: &'a mut dyn std::any::Any,
}

impl<'a> DispatchData<'a> {
    /// Access the dispatch data knowing its type
    ///
    /// Will return `None` if the provided type is not the correct
    /// inner type.
    pub fn get<T: std::any::Any>(&mut self) -> Option<&mut T> {
        self.data.downcast_mut()
    }

    /// Wrap a mutable reference
    ///
    /// This creates a new `DispatchData` from a mutable reference
    pub fn wrap<T: std::any::Any>(data: &'a mut T) -> DispatchData<'a> {
        DispatchData { data }
    }

    /// Reborrows this `DispatchData` to create a new one with the same content
    ///
    /// This is a quick and cheap way to propagate the `DispatchData` down a
    /// callback stack by value. It is basically a noop only there to ease
    /// work with the borrow checker.
    pub fn reborrow(&mut self) -> DispatchData {
        DispatchData {
            data: &mut *self.data,
        }
    }
}

struct Inner<E, F: ?Sized> {
    pending: RefCell<VecDeque<E>>,
    cb: RefCell<F>,
}

type DynInner<E> = Inner<E, dyn FnMut(E, &Filter<E>, DispatchData<'_>)>;

/// An event filter
///
/// Can be used in wayland-client and wayland-server to aggregate
/// messages from different objects into the same closure.
///
/// You need to provide it a closure of type `FnMut(E, &Filter<E>)`,
/// which will be called any time a message is sent to the filter
/// via the `send(..)` method. Your closure also receives a handle
/// to the filter as argument, so that you can use it from within
/// the callback (to assign new wayland objects to this filter for
/// example).
///
/// The `Filter` can be cloned, and all clones send messages to the
/// same closure. However it is not threadsafe.
pub struct Filter<E> {
    inner: Rc<DynInner<E>>,
}

impl<E> Clone for Filter<E> {
    fn clone(&self) -> Filter<E> {
        Filter {
            inner: self.inner.clone(),
        }
    }
}

impl<E> Filter<E> {
    /// Create a new filter from given closure
    pub fn new<F: FnMut(E, &Filter<E>, DispatchData<'_>) + 'static>(f: F) -> Filter<E> {
        Filter {
            inner: Rc::new(Inner {
                pending: RefCell::new(VecDeque::new()),
                cb: RefCell::new(f),
            }),
        }
    }

    /// Send a message to this filter
    pub fn send(&self, evt: E, mut data: DispatchData) {
        // gracefully handle reentrancy
        if let Ok(mut guard) = self.inner.cb.try_borrow_mut() {
            (&mut *guard)(evt, self, data.reborrow());
            // process all events that might have been enqueued by the cb
            while let Some(evt) = self.inner.pending.borrow_mut().pop_front() {
                (&mut *guard)(evt, self, data.reborrow());
            }
        } else {
            self.inner.pending.borrow_mut().push_back(evt);
        }
    }
}
<|endoftext|>


<commit_before>use alloc::borrow::Cow;
use alloc::vec::Vec;
use core::fmt::{LowerHex, Debug};

use scroll::{self, Pread};
use scroll::ctx::TryFromCtx;
use error;

use pe::section_table;
use pe::utils;
use pe::data_directories;

pub const IMPORT_BY_ORDINAL_32: u32 = 0x8000_0000;
pub const IMPORT_BY_ORDINAL_64: u64 = 0x8000_0000_0000_0000;
pub const IMPORT_RVA_MASK_32: u32 = 0x8fff_ffff;
pub const IMPORT_RVA_MASK_64: u64 = 0x0000_0000_8fff_ffff;

pub trait Bitfield<'a>: PartialEq + Eq + LowerHex + Debug + TryFromCtx<'a, scroll::Endian, Error=scroll::Error, Size=usize> {
    fn is_ordinal(&self) -> bool;
    fn to_ordinal(&self) -> u16;
    fn to_rva(&self) -> usize;
    fn size_of() -> usize;
    fn is_zero(&self) -> bool;
}

impl<'a> Bitfield<'a> for u64 {
    fn is_ordinal(&self) -> bool { self & IMPORT_BY_ORDINAL_64 == IMPORT_BY_ORDINAL_64 }
    fn to_ordinal(&self) -> u16 { (0xffff & self) as u16 }
    fn to_rva(&self) -> usize { (self & IMPORT_RVA_MASK_64) as usize }
    fn size_of() -> usize { 8 }
    fn is_zero(&self) -> bool { *self == 0}
}

impl<'a> Bitfield<'a> for u32 {
    fn is_ordinal(&self) -> bool { self & IMPORT_BY_ORDINAL_32 == IMPORT_BY_ORDINAL_32 }
    fn to_ordinal(&self) -> u16 { (0xffff & self) as u16 }
    fn to_rva(&self) -> usize { (self & IMPORT_RVA_MASK_32) as usize }
    fn size_of() -> usize { 4 }
    fn is_zero(&self) -> bool { *self == 0}
}

#[derive(Debug, Clone)]
pub struct HintNameTableEntry<'a> {
    pub hint: u16,
    pub name: &'a str,
}

impl<'a> HintNameTableEntry<'a> {
    fn parse(bytes: &'a [u8], mut offset: usize) -> error::Result<Self> {
        let offset = &mut offset;
        let hint = bytes.gread_with(offset, scroll::LE)?;
        let name = bytes.pread::<&'a str>(*offset)?;
        Ok(HintNameTableEntry { hint: hint, name: name })
    }
}

#[derive(Debug, Clone)]
pub enum SyntheticImportLookupTableEntry<'a> {
    OrdinalNumber(u16),
    HintNameTableRVA ((usize, HintNameTableEntry<'a>)), // [u8; 31] bitfield :/
}

#[derive(Debug)]
pub struct ImportLookupTableEntry<'a> {
    pub synthetic: SyntheticImportLookupTableEntry<'a>,
}

pub type ImportLookupTable<'a> = Vec<ImportLookupTableEntry<'a>>;

impl<'a> ImportLookupTableEntry<'a> {
    pub fn parse<T: Bitfield<'a>>(bytes: &'a [u8], mut offset: usize, sections: &[section_table::SectionTable])
                                                                      -> error::Result<ImportLookupTable<'a>> {
        let le = scroll::LE;
        let offset = &mut offset;
        let mut table = Vec::new();
        loop {
            let bitfield: T = bytes.gread_with(offset, le)?;
            if bitfield.is_zero() {
                debug!("imports done");
                break;
            } else {
                let synthetic = {
                    debug!("bitfield {:#x}", bitfield);
                    use self::SyntheticImportLookupTableEntry::*;
                    if bitfield.is_ordinal() {
                        let ordinal = bitfield.to_ordinal();
                        debug!("importing by ordinal {:#x}", ordinal);
                        OrdinalNumber(ordinal)
                    } else {
                        let rva = bitfield.to_rva();
                        let hentry = {
                            debug!("searching for RVA {:#x}", rva);
                            if let Some(offset) = utils::find_offset(rva as usize, sections) {
                                debug!("offset {:#x}", offset);
                                HintNameTableEntry::parse(bytes, offset)?
                            } else {
                                warn!("Entry {} has bad RVA: {:#x}", table.len(), rva);
                                continue
                            }
                        };
                        HintNameTableRVA ((rva, hentry))
                    }
                };
                let entry = ImportLookupTableEntry { synthetic: synthetic };
                table.push(entry);
            }
        }
        Ok(table)
    }
}

// get until entry is 0
pub type ImportAddressTable = Vec<usize>;

#[repr(C)]
#[derive(Debug)]
#[derive(Pread, Pwrite, SizeWith)]
pub struct ImportDirectoryEntry {
    pub import_lookup_table_rva: u32,
    pub time_date_stamp: u32,
    pub forwarder_chain: u32,
    pub name_rva: u32,
    pub import_address_table_rva: u32,
}

pub const SIZEOF_IMPORT_DIRECTORY_ENTRY: usize = 20;

impl ImportDirectoryEntry {
    pub fn is_null (&self) -> bool {
        (self.import_lookup_table_rva == 0) &&
            (self.time_date_stamp == 0) &&
            (self.forwarder_chain == 0) &&
            (self.name_rva == 0) &&
            (self.import_address_table_rva == 0)
    }
}

#[derive(Debug)]
pub struct SyntheticImportDirectoryEntry<'a> {
    pub import_directory_entry: ImportDirectoryEntry,
    /// Computed
    pub name: &'a str,
    /// The import lookup table is a vector of either ordinals, or RVAs + import names
    pub import_lookup_table: Option<ImportLookupTable<'a>>,
    /// Computed
    pub import_address_table: ImportAddressTable,
}

impl<'a> SyntheticImportDirectoryEntry<'a> {
    pub fn parse<T: Bitfield<'a>>(bytes: &'a [u8], import_directory_entry: ImportDirectoryEntry, sections: &[section_table::SectionTable]) -> error::Result<SyntheticImportDirectoryEntry<'a>> {
        const LE: scroll::Endian = scroll::LE;
        let name_rva = import_directory_entry.name_rva;
        let name = utils::try_name(bytes, name_rva as usize, sections)?;
        let import_lookup_table = {
            let import_lookup_table_rva = import_directory_entry.import_lookup_table_rva;
            debug!("Synthesizing lookup table imports for {} lib, with import lookup table rva: {:#x}", name, import_lookup_table_rva);
            if let Some(import_lookup_table_offset) = utils::find_offset(import_lookup_table_rva as usize, sections) {
                let import_lookup_table = ImportLookupTableEntry::parse::<T>(bytes, import_lookup_table_offset, sections)?;
                debug!("Successfully synthesized import lookup table entry: {:#?}", import_lookup_table);
                Some(import_lookup_table)
            } else {
                None
            }
        };

        let import_address_table_offset = &mut utils::find_offset(import_directory_entry.import_address_table_rva as usize, sections).ok_or(error::Error::Malformed(format!("Cannot map import_address_table_rva {:#x} into offset for {}", import_directory_entry.import_address_table_rva, name)))?;
        let mut import_address_table = Vec::new();
        loop {
            let import_address = bytes.gread_with(import_address_table_offset, LE)?;
            if import_address == 0 { break } else { import_address_table.push(import_address); }
        }
        Ok(SyntheticImportDirectoryEntry {
            import_directory_entry: import_directory_entry,
            name: name,
            import_lookup_table: import_lookup_table,
            import_address_table: import_address_table
        })
    }
}

#[derive(Debug)]
/// Contains a list of synthesized import data for this binary, e.g., which symbols from which libraries it is importing from
pub struct ImportData<'a> {
    pub import_data: Vec<SyntheticImportDirectoryEntry<'a>>,
}

impl<'a> ImportData<'a> {
    pub fn parse<T: Bitfield<'a>>(bytes: &'a[u8], dd: &data_directories::DataDirectory, sections: &[section_table::SectionTable]) -> error::Result<ImportData<'a>> {
        let import_directory_table_rva = dd.virtual_address as usize;
        debug!("import_directory_table_rva {:#x}", import_directory_table_rva);
        let offset = &mut utils::find_offset(import_directory_table_rva, sections).ok_or(error::Error::Malformed(format!("Cannot create ImportData; cannot map import_directory_table_rva {:#x} into offset", import_directory_table_rva)))?;;
        debug!("import data offset {:#x}", offset);
        let mut import_data = Vec::new();
        loop {
            let import_directory_entry: ImportDirectoryEntry = bytes.gread_with(offset, scroll::LE)?;
            debug!("{:#?}", import_directory_entry);
            if import_directory_entry.is_null() {
                break;
            } else {
                let entry = SyntheticImportDirectoryEntry::parse::<T>(bytes, import_directory_entry, sections)?;
                debug!("entry {:#?}", entry);
                import_data.push(entry);
            }
        }
        debug!("finished ImportData");
        Ok(ImportData { import_data: import_data})
    }
}

#[derive(Debug)]
/// A synthesized symbol import, the name is pre-indexed, and the binary offset is computed, as well as which dll it belongs to
pub struct Import<'a> {
    pub name: Cow<'a, str>,
    pub dll: &'a str,
    pub ordinal: u16,
    pub offset: usize,
    pub rva: usize,
    pub size: usize,
}

impl<'a> Import<'a> {
    pub fn parse<T: Bitfield<'a>>(_bytes: &'a [u8], import_data: &ImportData<'a>, _sections: &[section_table::SectionTable]) -> error::Result<Vec<Import<'a>>> {
        let mut imports = Vec::new();
        for data in &import_data.import_data {
            if let Some(ref import_lookup_table) = data.import_lookup_table {
                let dll = data.name;
                let import_base = data.import_directory_entry.import_address_table_rva as usize;
                debug!("Getting imports from {}", &dll);
                for (i, entry) in import_lookup_table.iter().enumerate() {
                    let offset = import_base + (i * T::size_of());
                    use self::SyntheticImportLookupTableEntry::*;
                    let (rva, name, ordinal) =
                        match &entry.synthetic {
                            &HintNameTableRVA ((rva, ref hint_entry)) => {
                                // if hint_entry.name = "" && hint_entry.hint = 0 {
                                //     println!("<PE.Import> warning hint/name table rva from {} without hint {:#x}", dll, rva);
                                // }
                                (rva, Cow::Borrowed(hint_entry.name), hint_entry.hint.clone())
                            },
                            &OrdinalNumber(ordinal) => {
                                let name = format!("ORDINAL {}", ordinal);
                                (0x0, Cow::Owned(name), ordinal)
                            },
                        };
                    let import =
                        Import {
                            name: name,
                            ordinal: ordinal, dll: dll,
                            size: T::size_of(), offset: offset, rva: rva as usize
                        };
                    imports.push(import);
                }
            }
        }
        Ok (imports)
    }
}
<commit_msg>Remove unnecessary struct<commit_after>use alloc::borrow::Cow;
use alloc::vec::Vec;
use core::fmt::{LowerHex, Debug};

use scroll::{self, Pread};
use scroll::ctx::TryFromCtx;
use error;

use pe::section_table;
use pe::utils;
use pe::data_directories;

pub const IMPORT_BY_ORDINAL_32: u32 = 0x8000_0000;
pub const IMPORT_BY_ORDINAL_64: u64 = 0x8000_0000_0000_0000;
pub const IMPORT_RVA_MASK_32: u32 = 0x8fff_ffff;
pub const IMPORT_RVA_MASK_64: u64 = 0x0000_0000_8fff_ffff;

pub trait Bitfield<'a>: PartialEq + Eq + LowerHex + Debug + TryFromCtx<'a, scroll::Endian, Error=scroll::Error, Size=usize> {
    fn is_ordinal(&self) -> bool;
    fn to_ordinal(&self) -> u16;
    fn to_rva(&self) -> usize;
    fn size_of() -> usize;
    fn is_zero(&self) -> bool;
}

impl<'a> Bitfield<'a> for u64 {
    fn is_ordinal(&self) -> bool { self & IMPORT_BY_ORDINAL_64 == IMPORT_BY_ORDINAL_64 }
    fn to_ordinal(&self) -> u16 { (0xffff & self) as u16 }
    fn to_rva(&self) -> usize { (self & IMPORT_RVA_MASK_64) as usize }
    fn size_of() -> usize { 8 }
    fn is_zero(&self) -> bool { *self == 0}
}

impl<'a> Bitfield<'a> for u32 {
    fn is_ordinal(&self) -> bool { self & IMPORT_BY_ORDINAL_32 == IMPORT_BY_ORDINAL_32 }
    fn to_ordinal(&self) -> u16 { (0xffff & self) as u16 }
    fn to_rva(&self) -> usize { (self & IMPORT_RVA_MASK_32) as usize }
    fn size_of() -> usize { 4 }
    fn is_zero(&self) -> bool { *self == 0}
}

#[derive(Debug, Clone)]
pub struct HintNameTableEntry<'a> {
    pub hint: u16,
    pub name: &'a str,
}

impl<'a> HintNameTableEntry<'a> {
    fn parse(bytes: &'a [u8], mut offset: usize) -> error::Result<Self> {
        let offset = &mut offset;
        let hint = bytes.gread_with(offset, scroll::LE)?;
        let name = bytes.pread::<&'a str>(*offset)?;
        Ok(HintNameTableEntry { hint: hint, name: name })
    }
}

#[derive(Debug, Clone)]
pub enum SyntheticImportLookupTableEntry<'a> {
    OrdinalNumber(u16),
    HintNameTableRVA ((usize, HintNameTableEntry<'a>)), // [u8; 31] bitfield :/
}

pub type ImportLookupTable<'a> = Vec<SyntheticImportLookupTableEntry<'a>>;

impl<'a> SyntheticImportLookupTableEntry<'a> {
    pub fn parse<T: Bitfield<'a>>(bytes: &'a [u8], mut offset: usize, sections: &[section_table::SectionTable])
                                                                      -> error::Result<ImportLookupTable<'a>> {
        let le = scroll::LE;
        let offset = &mut offset;
        let mut table = Vec::new();
        loop {
            let bitfield: T = bytes.gread_with(offset, le)?;
            if bitfield.is_zero() {
                debug!("imports done");
                break;
            } else {
                let entry = {
                    debug!("bitfield {:#x}", bitfield);
                    use self::SyntheticImportLookupTableEntry::*;
                    if bitfield.is_ordinal() {
                        let ordinal = bitfield.to_ordinal();
                        debug!("importing by ordinal {:#x}", ordinal);
                        OrdinalNumber(ordinal)
                    } else {
                        let rva = bitfield.to_rva();
                        let hentry = {
                            debug!("searching for RVA {:#x}", rva);
                            if let Some(offset) = utils::find_offset(rva as usize, sections) {
                                debug!("offset {:#x}", offset);
                                HintNameTableEntry::parse(bytes, offset)?
                            } else {
                                warn!("Entry {} has bad RVA: {:#x}", table.len(), rva);
                                continue
                            }
                        };
                        HintNameTableRVA ((rva, hentry))
                    }
                };
                table.push(entry);
            }
        }
        Ok(table)
    }
}

// get until entry is 0
pub type ImportAddressTable = Vec<usize>;

#[repr(C)]
#[derive(Debug)]
#[derive(Pread, Pwrite, SizeWith)]
pub struct ImportDirectoryEntry {
    pub import_lookup_table_rva: u32,
    pub time_date_stamp: u32,
    pub forwarder_chain: u32,
    pub name_rva: u32,
    pub import_address_table_rva: u32,
}

pub const SIZEOF_IMPORT_DIRECTORY_ENTRY: usize = 20;

impl ImportDirectoryEntry {
    pub fn is_null (&self) -> bool {
        (self.import_lookup_table_rva == 0) &&
            (self.time_date_stamp == 0) &&
            (self.forwarder_chain == 0) &&
            (self.name_rva == 0) &&
            (self.import_address_table_rva == 0)
    }
}

#[derive(Debug)]
pub struct SyntheticImportDirectoryEntry<'a> {
    pub import_directory_entry: ImportDirectoryEntry,
    /// Computed
    pub name: &'a str,
    /// The import lookup table is a vector of either ordinals, or RVAs + import names
    pub import_lookup_table: Option<ImportLookupTable<'a>>,
    /// Computed
    pub import_address_table: ImportAddressTable,
}

impl<'a> SyntheticImportDirectoryEntry<'a> {
    pub fn parse<T: Bitfield<'a>>(bytes: &'a [u8], import_directory_entry: ImportDirectoryEntry, sections: &[section_table::SectionTable]) -> error::Result<SyntheticImportDirectoryEntry<'a>> {
        const LE: scroll::Endian = scroll::LE;
        let name_rva = import_directory_entry.name_rva;
        let name = utils::try_name(bytes, name_rva as usize, sections)?;
        let import_lookup_table = {
            let import_lookup_table_rva = import_directory_entry.import_lookup_table_rva;
            debug!("Synthesizing lookup table imports for {} lib, with import lookup table rva: {:#x}", name, import_lookup_table_rva);
            if let Some(import_lookup_table_offset) = utils::find_offset(import_lookup_table_rva as usize, sections) {
                let import_lookup_table = SyntheticImportLookupTableEntry::parse::<T>(bytes, import_lookup_table_offset, sections)?;
                debug!("Successfully synthesized import lookup table entry: {:#?}", import_lookup_table);
                Some(import_lookup_table)
            } else {
                None
            }
        };

        let import_address_table_offset = &mut utils::find_offset(import_directory_entry.import_address_table_rva as usize, sections).ok_or(error::Error::Malformed(format!("Cannot map import_address_table_rva {:#x} into offset for {}", import_directory_entry.import_address_table_rva, name)))?;
        let mut import_address_table = Vec::new();
        loop {
            let import_address = bytes.gread_with(import_address_table_offset, LE)?;
            if import_address == 0 { break } else { import_address_table.push(import_address); }
        }
        Ok(SyntheticImportDirectoryEntry {
            import_directory_entry: import_directory_entry,
            name: name,
            import_lookup_table: import_lookup_table,
            import_address_table: import_address_table
        })
    }
}

#[derive(Debug)]
/// Contains a list of synthesized import data for this binary, e.g., which symbols from which libraries it is importing from
pub struct ImportData<'a> {
    pub import_data: Vec<SyntheticImportDirectoryEntry<'a>>,
}

impl<'a> ImportData<'a> {
    pub fn parse<T: Bitfield<'a>>(bytes: &'a[u8], dd: &data_directories::DataDirectory, sections: &[section_table::SectionTable]) -> error::Result<ImportData<'a>> {
        let import_directory_table_rva = dd.virtual_address as usize;
        debug!("import_directory_table_rva {:#x}", import_directory_table_rva);
        let offset = &mut utils::find_offset(import_directory_table_rva, sections).ok_or(error::Error::Malformed(format!("Cannot create ImportData; cannot map import_directory_table_rva {:#x} into offset", import_directory_table_rva)))?;;
        debug!("import data offset {:#x}", offset);
        let mut import_data = Vec::new();
        loop {
            let import_directory_entry: ImportDirectoryEntry = bytes.gread_with(offset, scroll::LE)?;
            debug!("{:#?}", import_directory_entry);
            if import_directory_entry.is_null() {
                break;
            } else {
                let entry = SyntheticImportDirectoryEntry::parse::<T>(bytes, import_directory_entry, sections)?;
                debug!("entry {:#?}", entry);
                import_data.push(entry);
            }
        }
        debug!("finished ImportData");
        Ok(ImportData { import_data: import_data})
    }
}

#[derive(Debug)]
/// A synthesized symbol import, the name is pre-indexed, and the binary offset is computed, as well as which dll it belongs to
pub struct Import<'a> {
    pub name: Cow<'a, str>,
    pub dll: &'a str,
    pub ordinal: u16,
    pub offset: usize,
    pub rva: usize,
    pub size: usize,
}

impl<'a> Import<'a> {
    pub fn parse<T: Bitfield<'a>>(_bytes: &'a [u8], import_data: &ImportData<'a>, _sections: &[section_table::SectionTable]) -> error::Result<Vec<Import<'a>>> {
        let mut imports = Vec::new();
        for data in &import_data.import_data {
            if let Some(ref import_lookup_table) = data.import_lookup_table {
                let dll = data.name;
                let import_base = data.import_directory_entry.import_address_table_rva as usize;
                debug!("Getting imports from {}", &dll);
                for (i, entry) in import_lookup_table.iter().enumerate() {
                    let offset = import_base + (i * T::size_of());
                    use self::SyntheticImportLookupTableEntry::*;
                    let (rva, name, ordinal) =
                        match entry {
                            &HintNameTableRVA ((rva, ref hint_entry)) => {
                                // if hint_entry.name = "" && hint_entry.hint = 0 {
                                //     println!("<PE.Import> warning hint/name table rva from {} without hint {:#x}", dll, rva);
                                // }
                                (rva, Cow::Borrowed(hint_entry.name), hint_entry.hint.clone())
                            },
                            &OrdinalNumber(ordinal) => {
                                let name = format!("ORDINAL {}", ordinal);
                                (0x0, Cow::Owned(name), ordinal)
                            },
                        };
                    let import =
                        Import {
                            name: name,
                            ordinal: ordinal, dll: dll,
                            size: T::size_of(), offset: offset, rva: rva as usize
                        };
                    imports.push(import);
                }
            }
        }
        Ok (imports)
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;

mod from_str;
// mod relabel;

pub mod cli;

/// Variable or terminal symbol in an MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
pub enum VarT<T> {
    /// `Var(i, j)` represents the `j`th component of the `i`th successor.
    /// Indexing starts from `0`.
    Var(usize, usize),
    T(T),
}

/// Composition function in an MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
pub struct Composition<T> {
    pub composition: Vec<Vec<VarT<T>>>,
}

impl<T> From<Vec<Vec<VarT<T>>>> for Composition<T> {
    fn from(encapsulated_value: Vec<Vec<VarT<T>>>) -> Self {
        Composition { composition: encapsulated_value }
    }
}

/// Rule of a weighted MCFG.
#[derive(Debug, PartialOrd, Ord, Clone)]
pub struct PMCFGRule<N, T, W> {
    pub head: N,
    pub tail: Vec<N>,
    pub composition: Composition<T>,
    pub weight: W,
}

impl<N, T, W> PMCFGRule<N, T, W>
    where T: Clone,
          W: Clone,
{
    pub fn map_nonterminals<F, M>(&self, f: F) -> PMCFGRule<M, T, W>
        where F: Fn(&N) -> M,
    {
        PMCFGRule {
            head: f(&self.head),
            tail: self.tail.iter().map(f).collect(),
            composition: self.composition.clone(),
            weight: self.weight.clone()
        }
    }
}

/// A weighted MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
pub struct PMCFG<N, T, W> {
    pub _dummy: PhantomData<T>,
    pub initial: Vec<N>,
    pub rules: Vec<PMCFGRule<N, T, W>>,
}

impl<N: Hash, T: Hash, W> Hash for PMCFGRule<N, T, W> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.head.hash(state);
        self.tail.hash(state);
        self.composition.hash(state);
    }
}

impl<N: PartialEq, T: PartialEq, W> PartialEq for PMCFGRule<N, T, W> {
    fn eq(&self, other: &Self) -> bool {
        self.head == other.head && self.tail == other.tail && self.composition == other.composition
    }
}

impl<N: Eq, T: Eq, W> Eq for PMCFGRule<N, T, W> {}

impl<T: fmt::Display> fmt::Display for Composition<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = "".to_string();

        let mut iter0 = self.composition.iter().peekable();
        let mut iter1;

        buffer.push_str("[");
        while let Some(proj) = iter0.next() {
            iter1 = proj.into_iter().peekable();
            buffer.push_str("[");
            while let Some(vart) = iter1.next() {
                buffer.push_str(format!("{}", vart).as_str());
                if iter1.peek().is_some() {
                    buffer.push_str(", ");
                }
            }
            buffer.push_str("]");
            if iter0.peek().is_some() {
                buffer.push_str(", ");
            }
        }
        buffer.push_str("]");

        write!(f, "{}", buffer)
    }
}

impl<T: fmt::Display> fmt::Display for VarT<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            VarT::Var(i, j) => {
                write!(f, "Var {} {}", i, j)
            },
            VarT::T(ref x) => {
                write!(f, "T \"{}\"", x)
            },
        }
    }
}

impl<N: fmt::Display, T: fmt::Display, W: fmt::Display> fmt::Display for PMCFGRule<N, T, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
           let mut buffer = "".to_string();

           let mut iter = self.tail.iter().peekable();

           buffer.push_str("(");
           while let Some(nt) = iter.next() {
               buffer.push_str(format!("\"{}\"", nt).as_str());
               if iter.peek().is_some() {
                   buffer.push_str(", ");
               }
           }
           buffer.push_str(")");

        write!(f, "\"{}\" → {} {}  # {}", self.head, self.composition, buffer, self.weight)
        // write!(f, "\"{}\" → {}", self.head, self.composition)
    }
}

impl<N: fmt::Display, T: fmt::Display, W: fmt::Display> fmt::Display for PMCFG<N, T, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = "".to_string();

        let mut iter = self.initial.iter().peekable();

        buffer.push_str("initial: [");
        while let Some(nt) = iter.next() {
            buffer.push_str(format!("\"{}\"", nt).as_str());
            if iter.peek().is_some() {
                buffer.push_str(", ");
            }
        }
        buffer.push_str("]\n\n");

        for r in &self.rules {
            buffer.push_str(format!("{}\n", r).as_str());
        }

        write!(f, "{}", buffer)
    }
}

pub fn evaluate<T: Clone + fmt::Debug>(term_map: &BTreeMap<Vec<usize>, Composition<T>>) -> Composition<T> {
    evaluate_pos(term_map, vec![])
}

pub fn evaluate_pos<T: Clone + fmt::Debug>(term_map: &BTreeMap<Vec<usize>, Composition<T>>, address: Vec<usize>) -> Composition<T> {
    let unexpanded_composition = &term_map.get(&address).unwrap().composition;
    let mut expanded_nonterminals: BTreeMap<_, Vec<Vec<VarT<T>>>> = BTreeMap::new();
    let mut expanded_composition = Vec::new();

    for component in unexpanded_composition {
        let mut expanded_component = Vec::new();

        for variable in component {
            match variable {
                &VarT::Var(num_nonter, num_compon) => {
                    if !expanded_nonterminals.contains_key(&num_nonter) {
                        let mut child_address = address.clone();
                        child_address.push(num_nonter);
                        let nonter_compos = evaluate_pos(term_map, child_address).composition;
                        expanded_nonterminals.insert(num_nonter, nonter_compos);
                    }

                    let nonter_compos = expanded_nonterminals.get(&num_nonter).unwrap();

                    if let Some(compon) = nonter_compos.get(num_compon) {
                        for terminal in compon {
                            expanded_component.push(terminal.clone());
                        }
                    } else {
                        panic!("{:?}: use of {}-th component of nonterminal {} that has only {} components!",
                               unexpanded_composition, num_compon, num_nonter, nonter_compos.len());
                    }

                },
                &VarT::T(ref terminal) => {
                    expanded_component.push(VarT::T(terminal.clone()));
                },
            }
        }

        expanded_composition.push(expanded_component);
    }

    Composition::from(expanded_composition)
}

#[cfg(test)]
mod tests {
    use super::*;
    use VarT::{Var, T};

    #[test]
    fn test_evaluate() {
        let mut term_map: BTreeMap<Vec<usize>, _> = BTreeMap::new();
        term_map.insert(vec![], Composition::from(vec![
                vec![Var(0,0), Var(1,0), Var(0,1), Var(1,1)]
        ]));
        term_map.insert(vec![0], Composition::from(vec![
                vec![Var(1,0), Var(0,0)],
                vec![Var(2,0), Var(0,1)]
        ]));
        term_map.insert(vec![0,0], Composition::from(vec![
                vec![Var(1,0), Var(0,0)],
                vec![Var(2,0), Var(0,1)]
        ]));
        term_map.insert(vec![0,0,0], Composition::from(vec![
                vec![],
                vec![]
        ]));
        term_map.insert(vec![0,0,1], Composition::from(vec![
                vec![T('a')]
        ]));
        term_map.insert(vec![0,0,2], Composition::from(vec![
                vec![T('c')]
        ]));
        term_map.insert(vec![0,1], Composition::from(vec![
                vec![T('a')]
        ]));
        term_map.insert(vec![0,2], Composition::from(vec![
                vec![T('c')]
        ]));
        term_map.insert(vec![1], Composition::from(vec![
                vec![Var(1,0), Var(0,0)],
                vec![Var(2,0), Var(0,1)]
        ]));
        term_map.insert(vec![1,0], Composition::from(vec![
                vec![],
                vec![]
        ]));
        term_map.insert(vec![1,1], Composition::from(vec![
                vec![T('b')]
        ]));
        term_map.insert(vec![1,2], Composition::from(vec![
                vec![T('d')]
        ]));

        let expanded_compos = Composition::from(vec![
            vec![T('a'), T('a'), T('b'), T('c'), T('c'), T('d')]
        ]);

        assert_eq!(expanded_compos, evaluate(&term_map));
    }

    #[test]
    #[should_panic(expected =
        "[[Var(0, 0), Var(0, 1)]]: use of 1-th component of nonterminal 0 that has only 1 components!"
    )]
    fn test_evaluate_invalid_composition() {
        let mut term_map: BTreeMap<Vec<usize>, _> = BTreeMap::new();
        term_map.insert(vec![], Composition::from(vec![
                vec![Var(0,0), Var(0,1)]
        ]));
        term_map.insert(vec![0], Composition::from(vec![
                vec![T('a')]
        ]));

        evaluate(&term_map);
    }
}
<commit_msg>src/pmcfg/mod.rs: added identify_terminals()<commit_after>use std::collections::BTreeMap;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;

mod from_str;
// mod relabel;

pub mod cli;

/// Variable or terminal symbol in an MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
pub enum VarT<T> {
    /// `Var(i, j)` represents the `j`th component of the `i`th successor.
    /// Indexing starts from `0`.
    Var(usize, usize),
    T(T),
}

/// Composition function in an MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
pub struct Composition<T> {
    pub composition: Vec<Vec<VarT<T>>>,
}

impl<T> From<Vec<Vec<VarT<T>>>> for Composition<T> {
    fn from(encapsulated_value: Vec<Vec<VarT<T>>>) -> Self {
        Composition { composition: encapsulated_value }
    }
}

/// Rule of a weighted MCFG.
#[derive(Debug, PartialOrd, Ord, Clone)]
pub struct PMCFGRule<N, T, W> {
    pub head: N,
    pub tail: Vec<N>,
    pub composition: Composition<T>,
    pub weight: W,
}

impl<N, T, W> PMCFGRule<N, T, W>
    where T: Clone,
          W: Clone,
{
    pub fn map_nonterminals<F, M>(&self, f: F) -> PMCFGRule<M, T, W>
        where F: Fn(&N) -> M,
    {
        PMCFGRule {
            head: f(&self.head),
            tail: self.tail.iter().map(f).collect(),
            composition: self.composition.clone(),
            weight: self.weight.clone()
        }
    }
}

/// A weighted MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
pub struct PMCFG<N, T, W> {
    pub _dummy: PhantomData<T>,
    pub initial: Vec<N>,
    pub rules: Vec<PMCFGRule<N, T, W>>,
}

impl<N: Hash, T: Hash, W> Hash for PMCFGRule<N, T, W> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.head.hash(state);
        self.tail.hash(state);
        self.composition.hash(state);
    }
}

impl<N: PartialEq, T: PartialEq, W> PartialEq for PMCFGRule<N, T, W> {
    fn eq(&self, other: &Self) -> bool {
        self.head == other.head && self.tail == other.tail && self.composition == other.composition
    }
}

impl<N: Eq, T: Eq, W> Eq for PMCFGRule<N, T, W> {}

impl<T: fmt::Display> fmt::Display for Composition<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = "".to_string();

        let mut iter0 = self.composition.iter().peekable();
        let mut iter1;

        buffer.push_str("[");
        while let Some(proj) = iter0.next() {
            iter1 = proj.into_iter().peekable();
            buffer.push_str("[");
            while let Some(vart) = iter1.next() {
                buffer.push_str(format!("{}", vart).as_str());
                if iter1.peek().is_some() {
                    buffer.push_str(", ");
                }
            }
            buffer.push_str("]");
            if iter0.peek().is_some() {
                buffer.push_str(", ");
            }
        }
        buffer.push_str("]");

        write!(f, "{}", buffer)
    }
}

impl<T: fmt::Display> fmt::Display for VarT<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            VarT::Var(i, j) => {
                write!(f, "Var {} {}", i, j)
            },
            VarT::T(ref x) => {
                write!(f, "T \"{}\"", x)
            },
        }
    }
}

impl<N: fmt::Display, T: fmt::Display, W: fmt::Display> fmt::Display for PMCFGRule<N, T, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
           let mut buffer = "".to_string();

           let mut iter = self.tail.iter().peekable();

           buffer.push_str("(");
           while let Some(nt) = iter.next() {
               buffer.push_str(format!("\"{}\"", nt).as_str());
               if iter.peek().is_some() {
                   buffer.push_str(", ");
               }
           }
           buffer.push_str(")");

        write!(f, "\"{}\" → {} {}  # {}", self.head, self.composition, buffer, self.weight)
        // write!(f, "\"{}\" → {}", self.head, self.composition)
    }
}

impl<N: fmt::Display, T: fmt::Display, W: fmt::Display> fmt::Display for PMCFG<N, T, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = "".to_string();

        let mut iter = self.initial.iter().peekable();

        buffer.push_str("initial: [");
        while let Some(nt) = iter.next() {
            buffer.push_str(format!("\"{}\"", nt).as_str());
            if iter.peek().is_some() {
                buffer.push_str(", ");
            }
        }
        buffer.push_str("]\n\n");

        for r in &self.rules {
            buffer.push_str(format!("{}\n", r).as_str());
        }

        write!(f, "{}", buffer)
    }
}

pub fn evaluate<T: Clone + fmt::Debug>(term_map: &BTreeMap<Vec<usize>, Composition<T>>) -> Composition<T> {
    evaluate_pos(term_map, vec![])
}

pub fn evaluate_pos<T: Clone + fmt::Debug>(term_map: &BTreeMap<Vec<usize>, Composition<T>>, address: Vec<usize>) -> Composition<T> {
    let unexpanded_composition = &term_map.get(&address).unwrap().composition;
    let mut expanded_nonterminals: BTreeMap<_, Vec<Vec<VarT<T>>>> = BTreeMap::new();
    let mut expanded_composition = Vec::new();

    for component in unexpanded_composition {
        let mut expanded_component = Vec::new();

        for variable in component {
            match variable {
                &VarT::Var(num_nonter, num_compon) => {
                    if !expanded_nonterminals.contains_key(&num_nonter) {
                        let mut child_address = address.clone();
                        child_address.push(num_nonter);
                        let nonter_compos = evaluate_pos(term_map, child_address).composition;
                        expanded_nonterminals.insert(num_nonter, nonter_compos);
                    }

                    let nonter_compos = expanded_nonterminals.get(&num_nonter).unwrap();

                    if let Some(compon) = nonter_compos.get(num_compon) {
                        for terminal in compon {
                            expanded_component.push(terminal.clone());
                        }
                    } else {
                        panic!("{:?}: use of {}-th component of nonterminal {} that has only {} components!",
                               unexpanded_composition, num_compon, num_nonter, nonter_compos.len());
                    }

                },
                &VarT::T(ref terminal) => {
                    expanded_component.push(VarT::T(terminal.clone()));
                },
            }
        }

        expanded_composition.push(expanded_component);
    }

    Composition::from(expanded_composition)
}

pub fn identify_terminals<A: Clone>(tree_map: &BTreeMap<Vec<usize>, Composition<A>>) -> (BTreeMap<Vec<usize>, Composition<(Vec<usize>, usize)>>, BTreeMap<(Vec<usize>, usize), A>) {
    let mut identified_tree_map = BTreeMap::new();
    let mut terminal_map = BTreeMap::new();

    for (address, composition) in tree_map {
        let vec_compos = &composition.composition;
        let mut identified_compos = Vec::new();
        let mut compos_var_pos = 0;

        for component in vec_compos {
            let mut identified_compon = Vec::new();

            for variable in component {
                match variable {
                    &VarT::Var(x, y) => {
                        identified_compon.push(VarT::Var(x, y));
                    },
                    &VarT::T(ref terminal) => {
                        let terminal_id = (address.clone(), compos_var_pos);
                        identified_compon.push(VarT::T(terminal_id.clone()));
                        terminal_map.insert(terminal_id, terminal.clone());
                    },
                };

                compos_var_pos += 1;
            }

            identified_compos.push(identified_compon);
        }

        identified_tree_map.insert(address.clone(), Composition::from(identified_compos));
    }

    (identified_tree_map, terminal_map)
}

#[cfg(test)]
mod tests {
    use super::*;
    use VarT::{Var, T};

    #[test]
    fn test_evaluate() {
        let mut term_map: BTreeMap<Vec<usize>, _> = BTreeMap::new();
        term_map.insert(vec![], Composition::from(vec![
                vec![Var(0,0), Var(1,0), Var(0,1), Var(1,1)]
        ]));
        term_map.insert(vec![0], Composition::from(vec![
                vec![Var(1,0), Var(0,0)],
                vec![Var(2,0), Var(0,1)]
        ]));
        term_map.insert(vec![0,0], Composition::from(vec![
                vec![Var(1,0), Var(0,0)],
                vec![Var(2,0), Var(0,1)]
        ]));
        term_map.insert(vec![0,0,0], Composition::from(vec![
                vec![],
                vec![]
        ]));
        term_map.insert(vec![0,0,1], Composition::from(vec![
                vec![T('a')]
        ]));
        term_map.insert(vec![0,0,2], Composition::from(vec![
                vec![T('c')]
        ]));
        term_map.insert(vec![0,1], Composition::from(vec![
                vec![T('a')]
        ]));
        term_map.insert(vec![0,2], Composition::from(vec![
                vec![T('c')]
        ]));
        term_map.insert(vec![1], Composition::from(vec![
                vec![Var(1,0), Var(0,0)],
                vec![Var(2,0), Var(0,1)]
        ]));
        term_map.insert(vec![1,0], Composition::from(vec![
                vec![],
                vec![]
        ]));
        term_map.insert(vec![1,1], Composition::from(vec![
                vec![T('b')]
        ]));
        term_map.insert(vec![1,2], Composition::from(vec![
                vec![T('d')]
        ]));

        let expanded_compos = Composition::from(vec![
            vec![T('a'), T('a'), T('b'), T('c'), T('c'), T('d')]
        ]);

        assert_eq!(expanded_compos, evaluate(&term_map));
    }

    #[test]
    #[should_panic(expected =
        "[[Var(0, 0), Var(0, 1)]]: use of 1-th component of nonterminal 0 that has only 1 components!"
    )]
    fn test_evaluate_invalid_composition() {
        let mut term_map: BTreeMap<Vec<usize>, _> = BTreeMap::new();
        term_map.insert(vec![], Composition::from(vec![
                vec![Var(0,0), Var(0,1)]
        ]));
        term_map.insert(vec![0], Composition::from(vec![
                vec![T('a')]
        ]));

        evaluate(&term_map);
    }
}
<|endoftext|>


<commit_before>use rand;

use std::rc::Rc;

use sdl2_mixer::Music;

use super::player::Player;
use sprite::Sprite;
use renderer::Renderer;
use mixer::Mixer;

use unit;
use unit::{Unit, Vec2, AABB, Point, MS};

static BALL_MOVE_SPEED: Unit = Unit(0.5);

pub struct Ball {
    sprite: Sprite,

    pos: Vec2,
    v: Vec2,

    aabb: AABB,

    beeep: Rc<~Music>,
    peeeeeep: Rc<~Music>,
    plop: Rc<~Music>,
}

impl Ball {
    pub fn new(renderer: &mut Renderer, mixer: &mut Mixer) -> Ball {
        let mut sprite = Sprite::new(renderer, ~"assets/ball.png");
        let size = sprite.size();
        sprite.set_pivot(Point::new(size.w / 2, size.h / 2));

        Ball {
            sprite: sprite,
            pos: unit::vec2::ZERO,
            v: unit::vec2::ZERO,
            aabb: AABB::new(
                0.0, 0.0, size.w, size.h
            ),

            beeep: mixer.load_music(~"assets/beeep.ogg"),
            peeeeeep: mixer.load_music(~"assets/peeeeeep.ogg"),
            plop: mixer.load_music(~"assets/plop.ogg"),
        }
    }

    pub fn emit(&mut self) {
        if self.v.length() == unit::ZERO {
            let dx = if rand::random::<f32>() > 0.5 {
                0.5
            } else {
                -0.5
            };
            self.set_direction(Vec2::new(dx,
                                         rand::random::<f32>() * 2.0 - 1.0));
        }
    }

    pub fn offset(&mut self, offset: Vec2) {
        self.pos = self.pos + offset;
    }

    pub fn update(&mut self, dt: MS,
                  top_wall_aabb: &AABB, bottom_wall_aabb: &AABB,
                  left_wall_aabb: &AABB, right_wall_aabb: &AABB,
                  player1: &mut Player, player2: &mut Player) {
        let dp = self.v * dt;

        // check collision of x direction
        let aabb = self.aabb.transform(self.pos + Vec2::new(dp.x, 0.0));
        if dp.x < unit::ZERO {
            if aabb.is_collided_with(left_wall_aabb) {
                player2.win();
                println!("{} : {}", player1.score(), player2.score());
                self.reset();

                self.peeeeeep.play(1);
            } if aabb.is_collided_with(&player1.aabb()) {
                let dy = self.reflection(player1);
                self.set_direction(Vec2::new(1.0, dy));

                self.beeep.play(1);
            } else {
                self.pos.x = self.pos.x + dp.x;
            }
        } else {
            if aabb.is_collided_with(right_wall_aabb) {
                player1.win();
                println!("{} : {}", player1.score(), player2.score());
                self.reset();

                self.peeeeeep.play(1);
            } else if aabb.is_collided_with(&player2.aabb()) {
                let dy = self.reflection(player2);
                self.set_direction(Vec2::new(-1.0, dy));

                self.beeep.play(1);
            } else {
                self.pos.x = self.pos.x + dp.x;
            }
        }

        // check collision of y direction
        let aabb = self.aabb.transform(self.pos + Vec2::new(0.0, dp.y));
        if dp.y < unit::ZERO {
            if aabb.is_collided_with(top_wall_aabb) {
                self.pos.y = top_wall_aabb.bottom() + self.aabb.size().y / 2.0;
                self.v.y = -self.v.y;

                self.plop.play(1);
            } else {
                self.pos.y = self.pos.y + dp.y;
            }
        } else {
            if aabb.is_collided_with(bottom_wall_aabb) {
                self.pos.y = bottom_wall_aabb.top() - self.aabb.size().y / 2.0;
                self.v.y = -self.v.y;

                self.plop.play(1);
            } else {
                self.pos.y = self.pos.y + dp.y;
            }
        }
    }

    pub fn render(&self, renderer: &Renderer) {
        self.sprite.render(renderer, self.pos);
    }

    pub fn reset(&mut self) {
        self.v = unit::vec2::ZERO;
        self.pos = Vec2::new(240.0, 160.0);
    }

    fn set_direction(&mut self, dir: Vec2) {
        self.v = dir.norm() * BALL_MOVE_SPEED;
    }

    fn reflection(&self, player: &Player) -> Unit {
        let mut y = self.pos.y - player.position().y;
        if y > player.aabb().size.y / 2.0 {
            y = player.aabb().size.y / 2.0;
        } else if y < -player.aabb().size.y / 2.0 {
            y = -player.aabb().size.y / 2.0;
        }

        y * 4.0 / player.aabb().size.y
    }
}
<commit_msg>Add random reflection.<commit_after>use rand;

use std::rc::Rc;

use sdl2_mixer::Music;

use super::player::Player;
use sprite::Sprite;
use renderer::Renderer;
use mixer::Mixer;

use unit;
use unit::{Unit, Vec2, AABB, Point, MS};

static BALL_MOVE_SPEED: Unit = Unit(0.5);

pub struct Ball {
    sprite: Sprite,

    pos: Vec2,
    v: Vec2,

    aabb: AABB,

    beeep: Rc<~Music>,
    peeeeeep: Rc<~Music>,
    plop: Rc<~Music>,
}

impl Ball {
    pub fn new(renderer: &mut Renderer, mixer: &mut Mixer) -> Ball {
        let mut sprite = Sprite::new(renderer, ~"assets/ball.png");
        let size = sprite.size();
        sprite.set_pivot(Point::new(size.w / 2, size.h / 2));

        Ball {
            sprite: sprite,
            pos: unit::vec2::ZERO,
            v: unit::vec2::ZERO,
            aabb: AABB::new(
                0.0, 0.0, size.w, size.h
            ),

            beeep: mixer.load_music(~"assets/beeep.ogg"),
            peeeeeep: mixer.load_music(~"assets/peeeeeep.ogg"),
            plop: mixer.load_music(~"assets/plop.ogg"),
        }
    }

    pub fn emit(&mut self) {
        if self.v.length() == unit::ZERO {
            let dx = if rand::random::<f32>() > 0.5 {
                0.5
            } else {
                -0.5
            };
            self.set_direction(Vec2::new(dx,
                                         rand::random::<f32>() * 2.0 - 1.0));
        }
    }

    pub fn update(&mut self, dt: MS,
                  top_wall_aabb: &AABB, bottom_wall_aabb: &AABB,
                  left_wall_aabb: &AABB, right_wall_aabb: &AABB,
                  player1: &mut Player, player2: &mut Player) {
        let dp = self.v * dt;

        // check collision of x direction
        let aabb = self.aabb.transform(self.pos + Vec2::new(dp.x, 0.0));
        if dp.x < unit::ZERO {
            if aabb.is_collided_with(left_wall_aabb) {
                player2.win();
                println!("{} : {}", player1.score(), player2.score());
                self.reset();

                self.peeeeeep.play(1).unwrap();
            } if aabb.is_collided_with(&player1.aabb()) {
                let dy = self.reflection(player1);
                self.set_direction(Vec2::new(1.0, dy));

                self.beeep.play(1).unwrap();
            } else {
                self.pos.x = self.pos.x + dp.x;
            }
        } else {
            if aabb.is_collided_with(right_wall_aabb) {
                player1.win();
                println!("{} : {}", player1.score(), player2.score());
                self.reset();

                self.peeeeeep.play(1).unwrap();
            } else if aabb.is_collided_with(&player2.aabb()) {
                let dy = self.reflection(player2);
                self.set_direction(Vec2::new(-1.0, dy));

                self.beeep.play(1).unwrap();
            } else {
                self.pos.x = self.pos.x + dp.x;
            }
        }

        // check collision of y direction
        let aabb = self.aabb.transform(self.pos + Vec2::new(0.0, dp.y));
        if dp.y < unit::ZERO {
            if aabb.is_collided_with(top_wall_aabb) {
                self.pos.y = top_wall_aabb.bottom() + self.aabb.size().y / 2.0;
                self.v.y = -self.v.y;

                self.plop.play(1).unwrap();
            } else {
                self.pos.y = self.pos.y + dp.y;
            }
        } else {
            if aabb.is_collided_with(bottom_wall_aabb) {
                self.pos.y = bottom_wall_aabb.top() - self.aabb.size().y / 2.0;
                self.v.y = -self.v.y;

                self.plop.play(1).unwrap();
            } else {
                self.pos.y = self.pos.y + dp.y;
            }
        }
    }

    pub fn render(&self, renderer: &Renderer) {
        self.sprite.render(renderer, self.pos);
    }

    pub fn reset(&mut self) {
        self.v = unit::vec2::ZERO;
        self.pos = Vec2::new(240.0, 160.0);
    }

    fn set_direction(&mut self, dir: Vec2) {
        self.v = dir.norm() * BALL_MOVE_SPEED;
    }

    fn reflection(&self, player: &Player) -> Unit {
        let mut y = self.pos.y - player.position().y;

        y = y * rand::random::<f32>() * 0.4 + 0.8;

        if y > player.aabb().size.y / 2.0 {
            y = player.aabb().size.y / 2.0;
        } else if y < -player.aabb().size.y / 2.0 {
            y = -player.aabb().size.y / 2.0;
        }

        y * 4.0 / player.aabb().size.y
    }
}
<|endoftext|>


<commit_before>//use num::Complex;

use ndarray::{Array3, Zip};
use ndarray_parallel::prelude::*;
use config::*;

//struct PotentialValue<T>(T); //Some kind of generic like this may work.

pub fn generate(config: &Config) -> Array3<f64> {
    let num = &config.grid.size;
    let init_size: [usize; 3] = [(num.x + 5) as usize, (num.y + 5) as usize, (num.z + 5) as usize];
    let mut v = Array3::<f64>::zeros(init_size);

    Zip::indexed(&mut v)
        .par_apply(|(i, j, k), x| *x = potential(config, &Index3 { x: i, y: j, z: k }));
    v
}

//TODO: Maybe have the OPTION to be complex rather than force it.
//For now we're dropping complex all together here.
pub fn potential(config: &Config, idx: &Index3) -> f64 {
    let num = &config.grid.size;
    match config.potential {
        Potential::NoPotential => 0.0,
        Potential::Cube => {
            if (idx.x > num.x / 4 && idx.x <= 3 * num.x / 4) &&
               (idx.y > num.y / 4 && idx.y <= 3 * num.y / 4) &&
               (idx.z > num.z / 4 && idx.z <= 3 * num.z / 4) {
                -10.0
            } else {
                0.0
            }
        }
        Potential::QuadWell => {
            if (idx.x > num.x / 4 && idx.x <= 3 * num.x / 4) &&
               (idx.y > num.y / 4 && idx.y <= 3 * num.y / 4) &&
               (idx.z > 3 * num.z / 8 && idx.z <= 5 * num.z / 8) {
                -10.0
            } else {
                0.0
            }
        }
        Potential::ComplexCoulomb => {
            let r = calculate_r(idx, &config.grid);
            if r < config.grid.dn {
                -1. / config.grid.dn
            } else {
                -1. / r
            }
        }
        Potential::FromFile => 0.0,
        Potential::FromScript => 0.0,
    }
}

fn calculate_r(idx: &Index3, grid: &Grid) -> f64 {
    let dx = (idx.x as f64) - ((grid.size.x as f64) + 1.) / 2.;
    let dy = (idx.y as f64) - ((grid.size.y as f64) + 1.) / 2.;
    let dz = (idx.z as f64) - ((grid.size.z as f64) + 1.) / 2.; //TODO: DISTNUMZ (if needed)
    grid.dn * (dx * dx + dy * dy + dz * dz).sqrt()
}
<commit_msg>Initial error handling for potentials, placeholder functions for file and script read-in<commit_after>//use num::Complex;
use ndarray::{Array3, Zip};
use ndarray_parallel::prelude::*;
use std::error::Error;
use std::fmt;

use config::*;

//TODO: Add failure modes for file read and scripting.
#[derive(Debug)]
enum PotentialError {
    NotAvailable,
}

impl fmt::Display for PotentialError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            PotentialError::NotAvailable => write!(f, "Not able to calculate potential value at an index for this potential type."),
        }
    }
}

impl Error for PotentialError {
    fn description(&self) -> &str {
        match *self {
            PotentialError::NotAvailable => "not available",
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            PotentialError::NotAvailable => None,
        }
    }
}

//struct PotentialValue<T>(T); //Some kind of generic like this may work.

pub fn generate(config: &Config) -> Array3<f64> {
    let num = &config.grid.size;
    let init_size: [usize; 3] = [(num.x + 5) as usize, (num.y + 5) as usize, (num.z + 5) as usize];
    let mut v = Array3::<f64>::zeros(init_size);

    Zip::indexed(&mut v)
        .par_apply(|(i, j, k), x| {
            match potential(config, &Index3 { x: i, y: j, z: k }) {
                Ok(result) => *x = result,
                Err(err) => panic!("Error: {}", err),
            }
        });
    v
}

/// Loads a pre-calculated potential from a user defined file.
pub fn from_file() -> Result<Array3<f64>, PotentialError> {
    //TODO: This is currently just a placeholder.
    let mut v = Array3::<f64>::zeros((2,2,2));

    Zip::from(&mut v).par_apply(|x| *x = 7.);
    Ok(v)
}

/// Loads a pre-calculated potential from a user defined script.
pub fn from_script() -> Result<Array3<f64>, PotentialError> {
    //TODO: This is currently just a placeholder.
    let ret = 2.3;
    let mut v = Array3::<f64>::zeros((2,2,2));

    Zip::from(&mut v).par_apply(|x| *x = ret);
    Ok(v)
}

//TODO: Maybe have the OPTION to be complex rather than force it.
//For now we're dropping complex all together here.
fn potential(config: &Config, idx: &Index3) -> Result<f64, PotentialError> {
    let num = &config.grid.size;
    match config.potential {
        Potential::NoPotential => Ok(0.0),
        Potential::Cube => {
            if (idx.x > num.x / 4 && idx.x <= 3 * num.x / 4) &&
               (idx.y > num.y / 4 && idx.y <= 3 * num.y / 4) &&
               (idx.z > num.z / 4 && idx.z <= 3 * num.z / 4) {
                Ok(-10.0)
            } else {
                Ok(0.0)
            }
        }
        Potential::QuadWell => {
            if (idx.x > num.x / 4 && idx.x <= 3 * num.x / 4) &&
               (idx.y > num.y / 4 && idx.y <= 3 * num.y / 4) &&
               (idx.z > 3 * num.z / 8 && idx.z <= 5 * num.z / 8) {
                Ok(-10.0)
            } else {
                Ok(0.0)
            }
        }
        Potential::ComplexCoulomb => {
            let r = calculate_r(idx, &config.grid);
            if r < config.grid.dn {
                Ok(-1. / config.grid.dn)
            } else {
                Ok(-1. / r)
            }
        }
        Potential::FromFile => Err(PotentialError::NotAvailable),
        Potential::FromScript => Err(PotentialError::NotAvailable), //TODO: This may not need to be an error.
    }
}

fn calculate_r(idx: &Index3, grid: &Grid) -> f64 {
    let dx = (idx.x as f64) - ((grid.size.x as f64) + 1.) / 2.;
    let dy = (idx.y as f64) - ((grid.size.y as f64) + 1.) / 2.;
    let dz = (idx.z as f64) - ((grid.size.z as f64) + 1.) / 2.; //TODO: DISTNUMZ (if needed)
    grid.dn * (dx * dx + dy * dy + dz * dz).sqrt()
}
<|endoftext|>