<commit_before>mod set_like;
mod unused_slots;

use itertools::*;
use ordermap::OrderMap;
use rand::*;
use std::cmp::max;
use std::time::*;

use crusader::*;
use dps::*;
use formation::*;
use self::unused_slots::unused_slots;
use self::set_like::*;

pub struct FormationSearch<'a> {
    state: State<'a>,
    search_root: Node<'a>,
}

impl<'a> FormationSearch<'a> {
    pub fn new(formation: Formation<'a>, crusaders: &'a [Crusader]) -> Self {
        FormationSearch {
            state: State { formation, crusaders },
            search_root: Node::new(),
        }
    }

    pub fn perform_search(&mut self, mut max_time: Duration) {
        let empty_positions = self.state.formation.empty_positions().count();
        for _ in 0..empty_positions {
            let loop_start = Instant::now();
            // FIXME: Figure out the right multiple to utilize all available time.
            // Later search nodes are likely to take significantly less time than early
            max_time = max_time / 2;
            while loop_start.elapsed() < max_time && !self.search_root.is_complete() {
                self.search_root.expand(&mut self.state.clone())
            }
            let current_dps = self.state.formation.total_dps();
            let root_dps = self.search_root.highest_dps_seen;
            let best_option = self.search_root.children
                .drain(..)
                .filter(|&(_, ref c)| {
                    c.highest_dps_seen >= current_dps &&
                        c.highest_dps_seen >= root_dps
                }).nth(0); // FIXME: If multiple options have same, take most searched
            if let Some((placement, child)) = best_option {
                self.state.place(placement);
                self.search_root = child;
            } else {
                break;
            }
        }
    }

    pub fn formation(self) -> Formation<'a> {
        self.state.formation
    }
}

type Placement<'a> = (usize, &'a Crusader);

#[derive(Debug, PartialEq, Eq)]
enum Progress {
    Expandable,
    FullyExpanded,
    Complete,
}

struct Node<'a> {
    highest_dps_seen: Dps,
    progress: Progress,
    children: OrderMap<Placement<'a>, Node<'a>>,
}

impl<'a> Node<'a> {
    fn new() -> Self {
        Node {
            progress: Progress::Expandable,
            highest_dps_seen: Default::default(),
            children: Default::default(),
        }
    }

    fn expand(&mut self, state: &mut State<'a>) {
        match self.progress {
            Progress::Expandable => {
                if let Some(placement) = state.random_placement(&self.children) {
                    state.place(placement);
                    let mut child = Node::new();
                    child.check_single_formation(state);
                    self.maybe_assign_dps(child.highest_dps_seen);
                    self.children.insert(placement, child);
                } else {
                    self.progress = Progress::FullyExpanded;
                    self.expand(state);
                }
            }
            Progress::FullyExpanded => self.to_best_child(state, Self::expand),
            Progress::Complete => panic!("expand called on complete node"),
        }
    }

    fn best_child(&mut self) -> Option<(Placement<'a>, &mut Self)> {
        let mut rng = thread_rng();
        let possible_children = self.children.len()
            - self.children.values().filter(|c| !c.is_complete()).count();
        let selection = rng.gen_range(0, max(1, possible_children));
        self.children.iter_mut()
            .filter(|&(_, ref v)| !v.is_complete())
            .map(|(&k, v)| (k, v))
            .nth(selection)
    }

    fn check_single_formation(&mut self, state: &mut State<'a>) {
        let mut search_dps = None;
        match self.progress {
            Progress::Expandable => {
                if let Some(placement) = state.random_placement(&EmptySet) {
                    state.place(placement);
                    if let Some(child) = self.children.get_mut(&placement) {
                        child.check_single_formation(state);
                        search_dps = Some(child.highest_dps_seen);
                    } else {
                        state.fill_formation_randomly();
                        search_dps = Some(state.formation.total_dps())
                    }
                }
            }
            Progress::FullyExpanded => {
                if let Some((placement, child)) = self.best_child() {
                    state.place(placement);
                    child.check_single_formation(state);
                    search_dps = Some(state.formation.total_dps());
                }
            }
            Progress::Complete => {}
        }
        if let Some(dps) = search_dps {
            self.maybe_assign_dps(dps)
        } else {
            self.progress = Progress::Complete;
            self.maybe_assign_dps(state.formation.total_dps());
        }
    }

    fn maybe_assign_dps(&mut self, dps: Dps) {
        if dps > self.highest_dps_seen {
            self.highest_dps_seen = dps;
        }
    }

    fn is_complete(&self) -> bool {
        self.progress == Progress::Complete
    }

    fn to_best_child(
        &mut self,
        state: &mut State<'a>,
        action: fn(&mut Self, &mut State<'a>),
    ) {
        let mut search_dps = None;
        if let Some((placement, child)) = self.best_child() {
            state.place(placement);
            action(child, state);
            search_dps = Some(child.highest_dps_seen);
        }
        if let Some(dps) = search_dps {
            self.maybe_assign_dps(dps);
        } else {
            self.progress = Progress::Complete;
        }
    }
}

#[derive(Clone)]
struct State<'a> {
    formation: Formation<'a>,
    crusaders: &'a [Crusader],
}

impl<'a> State<'a> {
    fn fill_formation_randomly(&mut self) {
        while let Some(placement) = self.random_placement(&EmptySet) {
            self.place(placement)
        }
    }

    fn random_placement<T>(&self, exclude: &T) -> Option<Placement<'a>> where
        T: SetLike<Placement<'a>>,
    {
        let mut rng = thread_rng();
        let num_placements = self.formation.empty_positions().count()
            * unused_slots(&self.formation, self.crusaders).count()
            - exclude.len();
        let selection = rng.gen_range(0, max(1, num_placements));
        self.formation.empty_positions()
            .cartesian_product(unused_slots(&self.formation, self.crusaders))
            .filter(|p| !exclude.contains(p))
            .nth(selection)
    }

    fn place(&mut self, (position, crusader): Placement<'a>) {
        self.formation.place_crusader(position, crusader);
    }
}

#[test]
fn fill_formation_randomly_fills_formation_until_formation_is_full() {
    let crusaders = test_crusaders();
    let mut state = search_state(&crusaders);
    state.fill_formation_randomly();
    assert_eq!(0, state.formation.empty_positions().count());
}

#[test]
fn fill_formation_randomly_stops_when_no_more_crusaders() {
    let crusaders = vec![test_crusaders().pop().unwrap()];
    let mut state = search_state(&crusaders);
    state.fill_formation_randomly();
    assert_eq!(1, state.formation.empty_positions().count());
}

#[test]
fn expanding_expandable_node_always_adds_new_child() {
    let crusaders = test_crusaders();
    let state = search_state(&crusaders);
    let mut search = Node::new();
    for _ in 0..6 {
        assert_eq!(Progress::Expandable, search.progress);
        search.expand(&mut state.clone());
    }
    // Expand one more time. Whether determining FullyExpanded is eager when 1
    // choice left, or lazy when attempting to expand and no choices is an impl
    // detail
    search.expand(&mut state.clone());
    assert_eq!(Progress::FullyExpanded, search.progress);
}

#[test]
fn expanding_updates_highest_dps_seen() {
    let crusaders = test_crusaders();
    let mut state = search_state(&crusaders);
    let mut search = Node::new();
    assert_eq!(Dps(0.0), search.highest_dps_seen);
    search.expand(&mut state);
    // Sanity check, make sure the formation was filled
    assert_eq!(0, state.formation.empty_positions().count());
    assert_eq!(state.formation.total_dps(), search.highest_dps_seen);
}

#[test]
fn expand_eventually_completes_node() {
    let crusaders = test_crusaders();
    let state = search_state(&crusaders);
    let mut search = Node::new();
    for _ in 0..1000 {
        if search.progress == Progress::Complete {
            break;
        }
        search.expand(&mut state.clone());
    }
    assert_eq!(Progress::Complete, search.progress);
}

#[cfg(test)]
fn test_formation<'a>() -> Formation<'a> {
    let positions = vec![
        Coordinate::new(0, 0),
        Coordinate::new(0, 1),
    ];
    Formation::empty(positions)
}

#[cfg(test)]
fn test_crusaders() -> Vec<Crusader> {
    vec![
        Crusader::new(CrusaderName::JimTheLumberjack, Level(5000)),
        Crusader::new(CrusaderName::VeronicaTheAndroidArcher, Level(5000)),
        Crusader::new(CrusaderName::SallyTheSuccubus, Level(5000)),
    ]
}

#[cfg(test)]
fn search_state(crusaders: &[Crusader]) -> State {
    State {
        formation: test_formation(),
        crusaders,
    }
}
<commit_msg>Correctly choose from all incomplete children<commit_after>mod set_like;
mod unused_slots;

use itertools::*;
use ordermap::OrderMap;
use rand::*;
use std::cmp::max;
use std::time::*;

use crusader::*;
use dps::*;
use formation::*;
use self::unused_slots::unused_slots;
use self::set_like::*;

pub struct FormationSearch<'a> {
    state: State<'a>,
    search_root: Node<'a>,
}

impl<'a> FormationSearch<'a> {
    pub fn new(formation: Formation<'a>, crusaders: &'a [Crusader]) -> Self {
        FormationSearch {
            state: State { formation, crusaders },
            search_root: Node::new(),
        }
    }

    pub fn perform_search(&mut self, mut max_time: Duration) {
        let empty_positions = self.state.formation.empty_positions().count();
        for _ in 0..empty_positions {
            let loop_start = Instant::now();
            // FIXME: Figure out the right multiple to utilize all available time.
            // Later search nodes are likely to take significantly less time than early
            max_time = max_time / 2;
            while loop_start.elapsed() < max_time && !self.search_root.is_complete() {
                self.search_root.expand(&mut self.state.clone())
            }
            let current_dps = self.state.formation.total_dps();
            let root_dps = self.search_root.highest_dps_seen;
            let best_option = self.search_root.children
                .drain(..)
                .filter(|&(_, ref c)| {
                    c.highest_dps_seen >= current_dps &&
                        c.highest_dps_seen >= root_dps
                }).nth(0); // FIXME: If multiple options have same, take most searched
            if let Some((placement, child)) = best_option {
                self.state.place(placement);
                self.search_root = child;
            } else {
                break;
            }
        }
    }

    pub fn formation(self) -> Formation<'a> {
        self.state.formation
    }
}

type Placement<'a> = (usize, &'a Crusader);

#[derive(Debug, PartialEq, Eq)]
enum Progress {
    Expandable,
    FullyExpanded,
    Complete,
}

struct Node<'a> {
    highest_dps_seen: Dps,
    progress: Progress,
    children: OrderMap<Placement<'a>, Node<'a>>,
}

impl<'a> Node<'a> {
    fn new() -> Self {
        Node {
            progress: Progress::Expandable,
            highest_dps_seen: Default::default(),
            children: Default::default(),
        }
    }

    fn expand(&mut self, state: &mut State<'a>) {
        match self.progress {
            Progress::Expandable => {
                if let Some(placement) = state.random_placement(&self.children) {
                    state.place(placement);
                    let mut child = Node::new();
                    child.check_single_formation(state);
                    self.maybe_assign_dps(child.highest_dps_seen);
                    self.children.insert(placement, child);
                } else {
                    self.progress = Progress::FullyExpanded;
                    self.expand(state);
                }
            }
            Progress::FullyExpanded => self.to_best_child(state, Self::expand),
            Progress::Complete => panic!("expand called on complete node"),
        }
    }

    fn best_child(&mut self) -> Option<(Placement<'a>, &mut Self)> {
        let mut rng = thread_rng();
        let possible_children = self.children.values().filter(|c| !c.is_complete()).count();
        let selection = rng.gen_range(0, max(1, possible_children));
        self.children.iter_mut()
            .filter(|&(_, ref v)| !v.is_complete())
            .map(|(&k, v)| (k, v))
            .nth(selection)
    }

    fn check_single_formation(&mut self, state: &mut State<'a>) {
        let mut search_dps = None;
        match self.progress {
            Progress::Expandable => {
                if let Some(placement) = state.random_placement(&EmptySet) {
                    state.place(placement);
                    if let Some(child) = self.children.get_mut(&placement) {
                        child.check_single_formation(state);
                        search_dps = Some(child.highest_dps_seen);
                    } else {
                        state.fill_formation_randomly();
                        search_dps = Some(state.formation.total_dps())
                    }
                }
            }
            Progress::FullyExpanded => {
                if let Some((placement, child)) = self.best_child() {
                    state.place(placement);
                    child.check_single_formation(state);
                    search_dps = Some(state.formation.total_dps());
                }
            }
            Progress::Complete => {}
        }
        if let Some(dps) = search_dps {
            self.maybe_assign_dps(dps)
        } else {
            self.progress = Progress::Complete;
            self.maybe_assign_dps(state.formation.total_dps());
        }
    }

    fn maybe_assign_dps(&mut self, dps: Dps) {
        if dps > self.highest_dps_seen {
            self.highest_dps_seen = dps;
        }
    }

    fn is_complete(&self) -> bool {
        self.progress == Progress::Complete
    }

    fn to_best_child(
        &mut self,
        state: &mut State<'a>,
        action: fn(&mut Self, &mut State<'a>),
    ) {
        let mut search_dps = None;
        if let Some((placement, child)) = self.best_child() {
            state.place(placement);
            action(child, state);
            search_dps = Some(child.highest_dps_seen);
        }
        if let Some(dps) = search_dps {
            self.maybe_assign_dps(dps);
        } else {
            self.progress = Progress::Complete;
        }
    }
}

#[derive(Clone)]
struct State<'a> {
    formation: Formation<'a>,
    crusaders: &'a [Crusader],
}

impl<'a> State<'a> {
    fn fill_formation_randomly(&mut self) {
        while let Some(placement) = self.random_placement(&EmptySet) {
            self.place(placement)
        }
    }

    fn random_placement<T>(&self, exclude: &T) -> Option<Placement<'a>> where
        T: SetLike<Placement<'a>>,
    {
        let mut rng = thread_rng();
        let num_placements = self.formation.empty_positions().count()
            * unused_slots(&self.formation, self.crusaders).count()
            - exclude.len();
        let selection = rng.gen_range(0, max(1, num_placements));
        self.formation.empty_positions()
            .cartesian_product(unused_slots(&self.formation, self.crusaders))
            .filter(|p| !exclude.contains(p))
            .nth(selection)
    }

    fn place(&mut self, (position, crusader): Placement<'a>) {
        self.formation.place_crusader(position, crusader);
    }
}

#[test]
fn fill_formation_randomly_fills_formation_until_formation_is_full() {
    let crusaders = test_crusaders();
    let mut state = search_state(&crusaders);
    state.fill_formation_randomly();
    assert_eq!(0, state.formation.empty_positions().count());
}

#[test]
fn fill_formation_randomly_stops_when_no_more_crusaders() {
    let crusaders = vec![test_crusaders().pop().unwrap()];
    let mut state = search_state(&crusaders);
    state.fill_formation_randomly();
    assert_eq!(1, state.formation.empty_positions().count());
}

#[test]
fn expanding_expandable_node_always_adds_new_child() {
    let crusaders = test_crusaders();
    let state = search_state(&crusaders);
    let mut search = Node::new();
    for _ in 0..6 {
        assert_eq!(Progress::Expandable, search.progress);
        search.expand(&mut state.clone());
    }
    // Expand one more time. Whether determining FullyExpanded is eager when 1
    // choice left, or lazy when attempting to expand and no choices is an impl
    // detail
    search.expand(&mut state.clone());
    assert_eq!(Progress::FullyExpanded, search.progress);
}

#[test]
fn expanding_updates_highest_dps_seen() {
    let crusaders = test_crusaders();
    let mut state = search_state(&crusaders);
    let mut search = Node::new();
    assert_eq!(Dps(0.0), search.highest_dps_seen);
    search.expand(&mut state);
    // Sanity check, make sure the formation was filled
    assert_eq!(0, state.formation.empty_positions().count());
    assert_eq!(state.formation.total_dps(), search.highest_dps_seen);
}

#[test]
fn expand_eventually_completes_node() {
    let crusaders = test_crusaders();
    let state = search_state(&crusaders);
    let mut search = Node::new();
    for _ in 0..1000 {
        if search.progress == Progress::Complete {
            break;
        }
        search.expand(&mut state.clone());
    }
    assert_eq!(Progress::Complete, search.progress);
}

#[cfg(test)]
fn test_formation<'a>() -> Formation<'a> {
    let positions = vec![
        Coordinate::new(0, 0),
        Coordinate::new(0, 1),
    ];
    Formation::empty(positions)
}

#[cfg(test)]
fn test_crusaders() -> Vec<Crusader> {
    vec![
        Crusader::new(CrusaderName::JimTheLumberjack, Level(5000)),
        Crusader::new(CrusaderName::VeronicaTheAndroidArcher, Level(5000)),
        Crusader::new(CrusaderName::SallyTheSuccubus, Level(5000)),
    ]
}

#[cfg(test)]
fn search_state(crusaders: &[Crusader]) -> State {
    State {
        formation: test_formation(),
        crusaders,
    }
}
<|endoftext|>


<commit_before>// This file is part of rgtk.
//
// rgtk is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// rgtk is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with rgtk.  If not, see <http://www.gnu.org/licenses/>.

//! A base class for menu objects

use gtk::{ffi, traits};
use gtk::cast::GTK_MENU_SHELL;

/// A base class for menu objects
pub trait MenuShell: traits::Widget + traits::Container {
    fn append(&mut self, widget: &traits::Widget) {
        unsafe {
            ffi::gtk_menu_shell_append(GTK_MENU_SHELL(self.get_widget()), widget.get_widget())
        }
    }

    fn prepend(&mut self, widget: &traits::Widget) {
        unsafe {
            ffi::gtk_menu_shell_prepend(GTK_MENU_SHELL(self.get_widget()), widget.get_widget())
        }
    }

    fn insert(&mut self, widget: &traits::Widget, position: i32) {
        unsafe {
            ffi::gtk_menu_shell_insert(GTK_MENU_SHELL(self.get_widget()),
                                       widget.get_widget(),
                                       position)
        }
    }

    fn deactivate(&mut self) {
        unsafe {
            ffi::gtk_menu_shell_deactivate(GTK_MENU_SHELL(self.get_widget()))
        }
    }

    // FIXME(jeremyletang): uncomment when gtk::MenuItem will be bind
    // fn select_item(&mut self, menu_item: gtk::MenuItem) {
    //     unsafe {
    //         ffi::gtk_menu_shell_select_item(GTK_MENU_SHELL(self.get_widget()),
    //                                         menu_item.get_widget())
    //     }
    // }

    fn deselect(&mut self) {
        unsafe {
            ffi::gtk_menu_shell_deselect(GTK_MENU_SHELL(self.get_widget()))
        }
    }

    // FIXME(jeremyletang): uncomment when gtk::MenuItem will be bind
    // fn activate_item(&mut self, menu_item: &gtk::MenuItem, force_deactivate: bool) {
    //     unsafe {
    //         ffi::gtk_menu_shell_activate_item(GTK_MENU_SHELL(self.get_widget()),
    //                                           menu_item.get_widget(),
    //                                           ffi::to_gboolean(force_deactivate))
    //     }
    // }

    fn select_first(&mut self, search_sensitive: bool) {
        unsafe {
            ffi::gtk_menu_shell_select_first(GTK_MENU_SHELL(self.get_widget()),
                                             ffi::to_gboolean(search_sensitive))
        }
    }

    fn cancel(&mut self) {
        unsafe {
            ffi::gtk_menu_shell_cancel(GTK_MENU_SHELL(self.get_widget()))
        }
    }

    fn get_take_focus(&self) -> bool {
        unsafe {
            ffi::to_bool(ffi::gtk_menu_shell_get_take_focus(GTK_MENU_SHELL(self.get_widget())))
        }
    }

    fn set_take_focus(&mut self, take_focus: bool) {
        unsafe {
            ffi::gtk_menu_shell_set_take_focus(GTK_MENU_SHELL(self.get_widget()),
                                               ffi::to_gboolean(take_focus))
        }
    }

    fn get_selected_item<T: traits::Widget>(&self) -> T {
        unsafe {
            ffi::FFIWidget::wrap(ffi::gtk_menu_shell_get_selected_item(GTK_MENU_SHELL(self.get_widget())))
        }
    }

    fn get_parent_shell<T: MenuShell>(&self) -> T {
        unsafe {
            ffi::FFIWidget::wrap(ffi::gtk_menu_shell_get_parent_shell(GTK_MENU_SHELL(self.get_widget())))
        }
    }
}
<commit_msg>uncomment functions in traits::MenuShell<commit_after>// This file is part of rgtk.
//
// rgtk is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// rgtk is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with rgtk.  If not, see <http://www.gnu.org/licenses/>.

//! A base class for menu objects

use gtk::{ffi, traits};
use gtk::cast::GTK_MENU_SHELL;

/// A base class for menu objects
pub trait MenuShell: traits::Widget + traits::Container {
    fn append(&mut self, widget: &traits::Widget) {
        unsafe {
            ffi::gtk_menu_shell_append(GTK_MENU_SHELL(self.get_widget()), widget.get_widget())
        }
    }

    fn prepend(&mut self, widget: &traits::Widget) {
        unsafe {
            ffi::gtk_menu_shell_prepend(GTK_MENU_SHELL(self.get_widget()), widget.get_widget())
        }
    }

    fn insert(&mut self, widget: &traits::Widget, position: i32) {
        unsafe {
            ffi::gtk_menu_shell_insert(GTK_MENU_SHELL(self.get_widget()),
                                       widget.get_widget(),
                                       position)
        }
    }

    fn deactivate(&mut self) {
        unsafe {
            ffi::gtk_menu_shell_deactivate(GTK_MENU_SHELL(self.get_widget()))
        }
    }

    fn select_item(&mut self, menu_item: &traits::MenuItem) {
        unsafe {
            ffi::gtk_menu_shell_select_item(GTK_MENU_SHELL(self.get_widget()),
                                            menu_item.get_widget())
        }
    }

    fn deselect(&mut self) {
        unsafe {
            ffi::gtk_menu_shell_deselect(GTK_MENU_SHELL(self.get_widget()))
        }
    }

    fn activate_item(&mut self, menu_item: &traits::MenuItem, force_deactivate: bool) {
        unsafe {
            ffi::gtk_menu_shell_activate_item(GTK_MENU_SHELL(self.get_widget()),
                                              menu_item.get_widget(),
                                              ffi::to_gboolean(force_deactivate))
        }
    }

    fn select_first(&mut self, search_sensitive: bool) {
        unsafe {
            ffi::gtk_menu_shell_select_first(GTK_MENU_SHELL(self.get_widget()),
                                             ffi::to_gboolean(search_sensitive))
        }
    }

    fn cancel(&mut self) {
        unsafe {
            ffi::gtk_menu_shell_cancel(GTK_MENU_SHELL(self.get_widget()))
        }
    }

    fn get_take_focus(&self) -> bool {
        unsafe {
            ffi::to_bool(ffi::gtk_menu_shell_get_take_focus(GTK_MENU_SHELL(self.get_widget())))
        }
    }

    fn set_take_focus(&mut self, take_focus: bool) {
        unsafe {
            ffi::gtk_menu_shell_set_take_focus(GTK_MENU_SHELL(self.get_widget()),
                                               ffi::to_gboolean(take_focus))
        }
    }

    fn get_selected_item<T: traits::Widget>(&self) -> T {
        unsafe {
            ffi::FFIWidget::wrap(ffi::gtk_menu_shell_get_selected_item(GTK_MENU_SHELL(self.get_widget())))
        }
    }

    fn get_parent_shell<T: MenuShell>(&self) -> T {
        unsafe {
            ffi::FFIWidget::wrap(ffi::gtk_menu_shell_get_parent_shell(GTK_MENU_SHELL(self.get_widget())))
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use std;
use std::marker::PhantomData;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::io;

use futures::{Async, Future, Poll};
use futures::stream::{Fuse, Peekable, Stream};
use futures::sync::mpsc::{unbounded, UnboundedReceiver};
use futures::task;
use rand::Rng;
use rand;
use tokio_core;
use tokio_core::reactor::Handle;

use BufStreamHandle;
use error::*;

/// A UDP stream of DNS binary packets
#[must_use = "futures do nothing unless polled"]
pub struct UdpStream {
    // FIXME: change UdpStream to always select a new Socket for every request
    socket: tokio_core::net::UdpSocket,
    outbound_messages: Peekable<Fuse<UnboundedReceiver<(Vec<u8>, SocketAddr)>>>,
}

impl UdpStream {
    /// This method is intended for client connections, see `with_bound` for a method better for
    ///  straight listening. It is expected that the resolver wrapper will be responsible for
    ///  creating and managing new UdpStreams such that each new client would have a random port
    ///  (reduce chance of cache poisoning). This will return a randomly assigned local port.
    ///
    /// # Arguments
    ///
    /// * `name_server`: socket address for the remote server (used to determine IPv4 or IPv6)
    /// * `loop_handle` - handle to the IO loop
    ///
    /// # Return
    ///
    /// a tuple of a Future Stream which will handle sending and receiving messsages, and a
    ///  handle which can be used to send messages into the stream.
    pub fn new<E>(
        name_server: SocketAddr,
        loop_handle: &Handle,
    ) -> (Box<Future<Item = UdpStream, Error = io::Error>>, BufStreamHandle<E>)
    where
        E: FromProtoError,
    {
        let (message_sender, outbound_messages) = unbounded();
        let message_sender = BufStreamHandle::<E> {
            sender: message_sender,
            phantom: PhantomData::<E>,
        };

        // TODO: allow the bind address to be specified...
        // constructs a future for getting the next randomly bound port to a UdpSocket
        let next_socket = Self::next_bound_local_address(&name_server);

        // This set of futures collapses the next udp socket into a stream which can be used for
        //  sending and receiving udp packets.
        let stream: Box<Future<Item = UdpStream, Error = io::Error>> = {
            let handle = loop_handle.clone();
            Box::new(
                next_socket
                    .map(move |socket| {
                        tokio_core::net::UdpSocket::from_socket(socket, &handle)
                            .expect("something wrong with the handle?")
                    })
                    .map(move |socket| {
                        UdpStream {
                            socket: socket,
                            outbound_messages: outbound_messages.fuse().peekable(),
                        }
                    }),
            )
        };

        (stream, message_sender)
    }

    /// Initialize the Stream with an already bound socket. Generally this should be only used for
    ///  server listening sockets. See `new` for a client oriented socket. Specifically, this there
    ///  is already a bound socket in this context, whereas `new` makes sure to randomize ports
    ///  for additional cache poison prevention.
    ///
    /// # Arguments
    ///
    /// * `socket` - an already bound UDP socket
    /// * `loop_handle` - handle to the IO loop
    ///
    /// # Return
    ///
    /// a tuple of a Future Stream which will handle sending and receiving messsages, and a
    ///  handle which can be used to send messages into the stream.
    pub fn with_bound<E>(
        socket: std::net::UdpSocket,
        loop_handle: &Handle,
    ) -> (Self, BufStreamHandle<E>)
    where
        E: FromProtoError + 'static,
    {
        let (message_sender, outbound_messages) = unbounded();
        let message_sender = BufStreamHandle::<E> {
            sender: message_sender,
            phantom: PhantomData::<E>,
        };

        // TODO: consider making this return a Result...
        let socket = tokio_core::net::UdpSocket::from_socket(socket, loop_handle)
            .expect("could not register socket to loop");

        let stream = UdpStream {
            socket: socket,
            outbound_messages: outbound_messages.fuse().peekable(),
        };

        (stream, message_sender)
    }

    /// Creates a future for randomly binding to a local socket address for client connections.
    fn next_bound_local_address(name_server: &SocketAddr) -> NextRandomUdpSocket {
        let zero_addr: IpAddr = match *name_server {
            SocketAddr::V4(..) => IpAddr::V4(Ipv4Addr::new(0,0,0,0)),
            SocketAddr::V6(..) => IpAddr::V6(Ipv6Addr::new(0,0,0,0,0,0,0,0)),
        };

        NextRandomUdpSocket { bind_address: zero_addr }
    }
}

impl Stream for UdpStream {
    type Item = (Vec<u8>, SocketAddr);
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        // this will not accept incoming data while there is data to send
        //  makes this self throttling.
        loop {
            // first try to send
            if let Async::Ready(Some(&(ref buffer, addr))) = try!(self.outbound_messages.peek().map_err(|()| io::Error::new(io::ErrorKind::Other, "unknown"))) {
          match self.socket.poll_write() {
            Async::NotReady => {
              return Ok(Async::NotReady)
            },
            Async::Ready(_) => {
              // will return if the socket will block
              try_nb!(self.socket.send_to(buffer, &addr));
            },
          }
        }

            // now pop the request and check if we should break or continue.
            match try!(self.outbound_messages.poll().map_err(|()| io::Error::new(io::ErrorKind::Other, "unknown"))) {
        // already handled above, here to make sure the poll() pops the next message
        Async::Ready(Some(_)) => (),
        // now we get to drop through to the receives...
        // TODO: should we also return None if there are no more messages to send?
        Async::NotReady | Async::Ready(None) => break,
      }
        }

        // For QoS, this will only accept one message and output that
        // recieve all inbound messages

        // TODO: this should match edns settings
        let mut buf = [0u8; 2048];

        // TODO: should we drop this packet if it's not from the same src as dest?
        let (len, src) = try_nb!(self.socket.recv_from(&mut buf));
        Ok(Async::Ready(
            Some((buf.iter().take(len).cloned().collect(), src)),
        ))
    }
}

#[must_use = "futures do nothing unless polled"]
struct NextRandomUdpSocket {
    bind_address: IpAddr,
}

impl Future for NextRandomUdpSocket {
    type Item = std::net::UdpSocket;
    type Error = io::Error;

    /// polls until there is an available next random UDP port.
    ///
    /// if there is no port available after 10 attempts, returns NotReady
    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut rand = rand::thread_rng();

        for attempt in 0..10 {
            let zero_addr = SocketAddr::new(
                self.bind_address,
                rand.gen_range(1025_u16, u16::max_value()),
            );

            match std::net::UdpSocket::bind(&zero_addr) {
                Ok(socket) => return Ok(Async::Ready(socket)),
                Err(err) => debug!("unable to bind port, attempt: {}: {}", attempt, err),
            }
        }

        warn!("could not get next random port, delaying");

        task::current().notify();
        // returning NotReady here, perhaps the next poll there will be some more socket available.
        Ok(Async::NotReady)
    }
}

#[test]
fn test_next_random_socket() {
    let mut io_loop = tokio_core::reactor::Core::new().unwrap();
    let (stream, _) = UdpStream::new::<ProtoError>(
        SocketAddr::new(
            std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)),
            52,
        ),
        &io_loop.handle(),
    );
    drop(io_loop.run(stream).ok().expect(
        "failed to get next socket address",
    ));
}

#[test]
fn test_udp_stream_ipv4() {
    udp_stream_test(std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)))
}

#[test]
#[cfg(not(target_os = "linux"))] // ignored until Travis-CI fixes IPv6
fn test_udp_stream_ipv6() {
    udp_stream_test(std::net::IpAddr::V6(
        std::net::Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1),
    ))
}

#[cfg(test)]
fn udp_stream_test(server_addr: std::net::IpAddr) {
    use tokio_core::reactor::Core;

    use std;
    let succeeded = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
    let succeeded_clone = succeeded.clone();
    std::thread::Builder::new()
        .name("thread_killer".to_string())
        .spawn(move || {
            let succeeded = succeeded_clone.clone();
            for _ in 0..15 {
                std::thread::sleep(std::time::Duration::from_secs(1));
                if succeeded.load(std::sync::atomic::Ordering::Relaxed) {
                    return;
                }
            }

            panic!("timeout");
        })
        .unwrap();

    let server = std::net::UdpSocket::bind(SocketAddr::new(server_addr, 0)).unwrap();
    server
        .set_read_timeout(Some(std::time::Duration::from_secs(5)))
        .unwrap(); // should recieve something within 5 seconds...
    server
        .set_write_timeout(Some(std::time::Duration::from_secs(5)))
        .unwrap(); // should recieve something within 5 seconds...
    let server_addr = server.local_addr().unwrap();

    let test_bytes: &'static [u8; 8] = b"DEADBEEF";
    let send_recv_times = 4;

    // an in and out server
    let server_handle = std::thread::Builder::new()
        .name("test_udp_stream_ipv4:server".to_string())
        .spawn(move || {
            let mut buffer = [0_u8; 512];

            for _ in 0..send_recv_times {
                // wait for some bytes...
                let (len, addr) = server.recv_from(&mut buffer).expect("receive failed");

                assert_eq!(&buffer[0..len], test_bytes);

                // bounce them right back...
                assert_eq!(
                    server.send_to(&buffer[0..len], addr).expect("send failed"),
                    len
                );
            }
        })
        .unwrap();

    // setup the client, which is going to run on the testing thread...
    let mut io_loop = Core::new().unwrap();

    // the tests should run within 5 seconds... right?
    // TODO: add timeout here, so that test never hangs...
    let client_addr = match server_addr {
        std::net::SocketAddr::V4(_) => "127.0.0.1:0",
        std::net::SocketAddr::V6(_) => "[::1]:0",
    };

    let socket = std::net::UdpSocket::bind(client_addr).expect("could not create socket"); // some random address...
    let (mut stream, sender) = UdpStream::with_bound::<ProtoError>(socket, &io_loop.handle());
    //let mut stream: UdpStream = io_loop.run(stream).ok().unwrap();

    for _ in 0..send_recv_times {
        // test once
        sender.sender
            .unbounded_send((test_bytes.to_vec(), server_addr))
            .unwrap();
        let (buffer_and_addr, stream_tmp) = io_loop.run(stream.into_future()).ok().unwrap();
        stream = stream_tmp;
        let (buffer, addr) = buffer_and_addr.expect("no buffer received");
        assert_eq!(&buffer, test_bytes);
        assert_eq!(addr, server_addr);
    }

    succeeded.store(true, std::sync::atomic::Ordering::Relaxed);
    server_handle.join().expect("server thread failed");
}
<commit_msg>also fix rand generation for UDP src port<commit_after>// Copyright 2015-2016 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use std;
use std::marker::PhantomData;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::io;

use futures::{Async, Future, Poll};
use futures::stream::{Fuse, Peekable, Stream};
use futures::sync::mpsc::{unbounded, UnboundedReceiver};
use futures::task;
use rand;
use rand::distributions::{IndependentSample, Range};
use tokio_core;
use tokio_core::reactor::Handle;

use BufStreamHandle;
use error::*;

/// A UDP stream of DNS binary packets
#[must_use = "futures do nothing unless polled"]
pub struct UdpStream {
    // FIXME: change UdpStream to always select a new Socket for every request
    socket: tokio_core::net::UdpSocket,
    outbound_messages: Peekable<Fuse<UnboundedReceiver<(Vec<u8>, SocketAddr)>>>,
}

impl UdpStream {
    /// This method is intended for client connections, see `with_bound` for a method better for
    ///  straight listening. It is expected that the resolver wrapper will be responsible for
    ///  creating and managing new UdpStreams such that each new client would have a random port
    ///  (reduce chance of cache poisoning). This will return a randomly assigned local port.
    ///
    /// # Arguments
    ///
    /// * `name_server`: socket address for the remote server (used to determine IPv4 or IPv6)
    /// * `loop_handle` - handle to the IO loop
    ///
    /// # Return
    ///
    /// a tuple of a Future Stream which will handle sending and receiving messsages, and a
    ///  handle which can be used to send messages into the stream.
    pub fn new<E>(
        name_server: SocketAddr,
        loop_handle: &Handle,
    ) -> (Box<Future<Item = UdpStream, Error = io::Error>>, BufStreamHandle<E>)
    where
        E: FromProtoError,
    {
        let (message_sender, outbound_messages) = unbounded();
        let message_sender = BufStreamHandle::<E> {
            sender: message_sender,
            phantom: PhantomData::<E>,
        };

        // TODO: allow the bind address to be specified...
        // constructs a future for getting the next randomly bound port to a UdpSocket
        let next_socket = Self::next_bound_local_address(&name_server);

        // This set of futures collapses the next udp socket into a stream which can be used for
        //  sending and receiving udp packets.
        let stream: Box<Future<Item = UdpStream, Error = io::Error>> = {
            let handle = loop_handle.clone();
            Box::new(
                next_socket
                    .map(move |socket| {
                        tokio_core::net::UdpSocket::from_socket(socket, &handle)
                            .expect("something wrong with the handle?")
                    })
                    .map(move |socket| {
                        UdpStream {
                            socket: socket,
                            outbound_messages: outbound_messages.fuse().peekable(),
                        }
                    }),
            )
        };

        (stream, message_sender)
    }

    /// Initialize the Stream with an already bound socket. Generally this should be only used for
    ///  server listening sockets. See `new` for a client oriented socket. Specifically, this there
    ///  is already a bound socket in this context, whereas `new` makes sure to randomize ports
    ///  for additional cache poison prevention.
    ///
    /// # Arguments
    ///
    /// * `socket` - an already bound UDP socket
    /// * `loop_handle` - handle to the IO loop
    ///
    /// # Return
    ///
    /// a tuple of a Future Stream which will handle sending and receiving messsages, and a
    ///  handle which can be used to send messages into the stream.
    pub fn with_bound<E>(
        socket: std::net::UdpSocket,
        loop_handle: &Handle,
    ) -> (Self, BufStreamHandle<E>)
    where
        E: FromProtoError + 'static,
    {
        let (message_sender, outbound_messages) = unbounded();
        let message_sender = BufStreamHandle::<E> {
            sender: message_sender,
            phantom: PhantomData::<E>,
        };

        // TODO: consider making this return a Result...
        let socket = tokio_core::net::UdpSocket::from_socket(socket, loop_handle)
            .expect("could not register socket to loop");

        let stream = UdpStream {
            socket: socket,
            outbound_messages: outbound_messages.fuse().peekable(),
        };

        (stream, message_sender)
    }

    /// Creates a future for randomly binding to a local socket address for client connections.
    fn next_bound_local_address(name_server: &SocketAddr) -> NextRandomUdpSocket {
        let zero_addr: IpAddr = match *name_server {
            SocketAddr::V4(..) => IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
            SocketAddr::V6(..) => IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0)),
        };

        NextRandomUdpSocket { bind_address: zero_addr }
    }
}

impl Stream for UdpStream {
    type Item = (Vec<u8>, SocketAddr);
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        // this will not accept incoming data while there is data to send
        //  makes this self throttling.
        loop {
            // first try to send
            if let Async::Ready(Some(&(ref buffer, addr))) = try!(self.outbound_messages.peek().map_err(|()| io::Error::new(io::ErrorKind::Other, "unknown"))) {
          match self.socket.poll_write() {
            Async::NotReady => {
              return Ok(Async::NotReady)
            },
            Async::Ready(_) => {
              // will return if the socket will block
              try_nb!(self.socket.send_to(buffer, &addr));
            },
          }
        }

            // now pop the request and check if we should break or continue.
            match try!(self.outbound_messages.poll().map_err(|()| io::Error::new(io::ErrorKind::Other, "unknown"))) {
        // already handled above, here to make sure the poll() pops the next message
        Async::Ready(Some(_)) => (),
        // now we get to drop through to the receives...
        // TODO: should we also return None if there are no more messages to send?
        Async::NotReady | Async::Ready(None) => break,
      }
        }

        // For QoS, this will only accept one message and output that
        // recieve all inbound messages

        // TODO: this should match edns settings
        let mut buf = [0u8; 2048];

        // TODO: should we drop this packet if it's not from the same src as dest?
        let (len, src) = try_nb!(self.socket.recv_from(&mut buf));
        Ok(Async::Ready(
            Some((buf.iter().take(len).cloned().collect(), src)),
        ))
    }
}

#[must_use = "futures do nothing unless polled"]
struct NextRandomUdpSocket {
    bind_address: IpAddr,
}

impl Future for NextRandomUdpSocket {
    type Item = std::net::UdpSocket;
    type Error = io::Error;

    /// polls until there is an available next random UDP port.
    ///
    /// if there is no port available after 10 attempts, returns NotReady
    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let between = Range::new(1025_u32, u32::from(u16::max_value()) + 1);
        let mut rand = rand::thread_rng();

        for attempt in 0..10 {
            let port = between.ind_sample(&mut rand) as u16; // the range is [0 ... u16::max] aka [0 .. u16::max + 1)
            let zero_addr = SocketAddr::new(
                self.bind_address,
                port,
            );

            match std::net::UdpSocket::bind(&zero_addr) {
                Ok(socket) => return Ok(Async::Ready(socket)),
                Err(err) => debug!("unable to bind port, attempt: {}: {}", attempt, err),
            }
        }

        warn!("could not get next random port, delaying");

        task::current().notify();
        // returning NotReady here, perhaps the next poll there will be some more socket available.
        Ok(Async::NotReady)
    }
}

#[test]
fn test_next_random_socket() {
    let mut io_loop = tokio_core::reactor::Core::new().unwrap();
    let (stream, _) = UdpStream::new::<ProtoError>(
        SocketAddr::new(
            std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)),
            52,
        ),
        &io_loop.handle(),
    );
    drop(io_loop.run(stream).ok().expect(
        "failed to get next socket address",
    ));
}

#[test]
fn test_udp_stream_ipv4() {
    udp_stream_test(std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)))
}

#[test]
#[cfg(not(target_os = "linux"))] // ignored until Travis-CI fixes IPv6
fn test_udp_stream_ipv6() {
    udp_stream_test(std::net::IpAddr::V6(
        std::net::Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1),
    ))
}

#[cfg(test)]
fn udp_stream_test(server_addr: std::net::IpAddr) {
    use tokio_core::reactor::Core;

    use std;
    let succeeded = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
    let succeeded_clone = succeeded.clone();
    std::thread::Builder::new()
        .name("thread_killer".to_string())
        .spawn(move || {
            let succeeded = succeeded_clone.clone();
            for _ in 0..15 {
                std::thread::sleep(std::time::Duration::from_secs(1));
                if succeeded.load(std::sync::atomic::Ordering::Relaxed) {
                    return;
                }
            }

            panic!("timeout");
        })
        .unwrap();

    let server = std::net::UdpSocket::bind(SocketAddr::new(server_addr, 0)).unwrap();
    server
        .set_read_timeout(Some(std::time::Duration::from_secs(5)))
        .unwrap(); // should recieve something within 5 seconds...
    server
        .set_write_timeout(Some(std::time::Duration::from_secs(5)))
        .unwrap(); // should recieve something within 5 seconds...
    let server_addr = server.local_addr().unwrap();

    let test_bytes: &'static [u8; 8] = b"DEADBEEF";
    let send_recv_times = 4;

    // an in and out server
    let server_handle = std::thread::Builder::new()
        .name("test_udp_stream_ipv4:server".to_string())
        .spawn(move || {
            let mut buffer = [0_u8; 512];

            for _ in 0..send_recv_times {
                // wait for some bytes...
                let (len, addr) = server.recv_from(&mut buffer).expect("receive failed");

                assert_eq!(&buffer[0..len], test_bytes);

                // bounce them right back...
                assert_eq!(
                    server.send_to(&buffer[0..len], addr).expect("send failed"),
                    len
                );
            }
        })
        .unwrap();

    // setup the client, which is going to run on the testing thread...
    let mut io_loop = Core::new().unwrap();

    // the tests should run within 5 seconds... right?
    // TODO: add timeout here, so that test never hangs...
    let client_addr = match server_addr {
        std::net::SocketAddr::V4(_) => "127.0.0.1:0",
        std::net::SocketAddr::V6(_) => "[::1]:0",
    };

    let socket = std::net::UdpSocket::bind(client_addr).expect("could not create socket"); // some random address...
    let (mut stream, sender) = UdpStream::with_bound::<ProtoError>(socket, &io_loop.handle());
    //let mut stream: UdpStream = io_loop.run(stream).ok().unwrap();

    for _ in 0..send_recv_times {
        // test once
        sender
            .sender
            .unbounded_send((test_bytes.to_vec(), server_addr))
            .unwrap();
        let (buffer_and_addr, stream_tmp) = io_loop.run(stream.into_future()).ok().unwrap();
        stream = stream_tmp;
        let (buffer, addr) = buffer_and_addr.expect("no buffer received");
        assert_eq!(&buffer, test_bytes);
        assert_eq!(addr, server_addr);
    }

    succeeded.store(true, std::sync::atomic::Ordering::Relaxed);
    server_handle.join().expect("server thread failed");
}
<|endoftext|>


<commit_before>// See LICENSE file for copyright and license details.

use std;
use gl;
use gl::types::{
  GLfloat,
  GLint,
  GLuint,
  GLchar,
  GLenum,
};
use cgmath::matrix::{
  Matrix,
  Mat4,
  Mat3,
  ToMat4,
};
use cgmath::vector::Vec3;
use cgmath::angle;
use misc::{
  c_str,
  deg_to_rad,
};
use color::Color3;

pub fn compile_shader(src: &str, shader_type: GLenum) -> GLuint {
  let shader = gl::CreateShader(shader_type);
  unsafe {
    gl::ShaderSource(shader, 1, &c_str(src), std::ptr::null());
    gl::CompileShader(shader);
    let mut status = gl::FALSE as GLint;
    gl::GetShaderiv(shader, gl::COMPILE_STATUS, &mut status);
    if status != (gl::TRUE as GLint) {
      let mut len = 0;
      gl::GetShaderiv(shader, gl::INFO_LOG_LENGTH, &mut len);
      // subtract 1 to skip the trailing null character
      let mut buf = std::vec::from_elem(len as uint - 1, 0u8);
      gl::GetShaderInfoLog(shader, len, std::ptr::mut_null(),
        buf.as_mut_ptr() as *mut GLchar
      );
      fail!("compile_shader(): " + std::str::raw::from_utf8(buf));
    }
  }
  shader
}

pub fn link_program(vertex_shader: GLuint, fragment_shader: GLuint) -> GLuint {
  let program = gl::CreateProgram();
  gl::AttachShader(program, vertex_shader);
  gl::AttachShader(program, fragment_shader);
  gl::LinkProgram(program);
  unsafe {
    let mut status = gl::FALSE as GLint;
    gl::GetProgramiv(program, gl::LINK_STATUS, &mut status);
    if status != (gl::TRUE as GLint) {
      let mut len: GLint = 0;
      gl::GetProgramiv(program, gl::INFO_LOG_LENGTH, &mut len);
      // subtract 1 to skip the trailing null character
      let mut buf = std::vec::from_elem(len as uint - 1, 0u8);
      gl::GetProgramInfoLog(program, len, std::ptr::mut_null(),
        buf.as_mut_ptr() as *mut GLchar
      );
      fail!("link_program(): " + std::str::raw::from_utf8(buf));
    }
  }
  program
}

pub fn compile_program(
  vertex_shader_src: &str,
  frag_shader_src: &str,
) -> GLuint {
  let vertex_shader = compile_shader(
    vertex_shader_src, gl::VERTEX_SHADER);
  let fragment_shader = compile_shader(
    frag_shader_src, gl::FRAGMENT_SHADER);
  let program = link_program(vertex_shader, fragment_shader);
  // mark shaders for deletion after program deletion
  gl::DeleteShader(fragment_shader);
  gl::DeleteShader(vertex_shader);
  program
}

pub fn get_attr(program_id: GLuint, name: &str) -> GLuint {
  unsafe {
    gl::GetAttribLocation(program_id, c_str(name)) as GLuint
  }
}

pub fn get_uniform(program: GLuint, name: &str) -> GLint {
  unsafe {
    gl::GetUniformLocation(program, c_str(name))
  }
}

pub fn draw_mesh(faces_count: int) {
  let starting_index = 0;
  let vertices_count = faces_count as i32 * 3;
  gl::DrawArrays(gl::TRIANGLES, starting_index, vertices_count);
}

pub fn uniform_mat4f(mat_id: GLint, mat: &Mat4<GLfloat>) {
  unsafe {
    gl::UniformMatrix4fv(mat_id, 1, gl::FALSE, mat.cr(0, 0));
  }
}

pub fn tr(m: Mat4<f32>, v: Vec3<f32>) -> Mat4<f32> {
  let mut t = Mat4::<f32>::identity();
  *t.mut_cr(3, 0) = v.x;
  *t.mut_cr(3, 1) = v.y;
  *t.mut_cr(3, 2) = v.z;
  m.mul_m(&t)
}

pub fn rot_x(m: Mat4<f32>, angle: f32) -> Mat4<f32> {
  let rad = angle::rad(deg_to_rad(angle));
  let r = Mat3::from_angle_x(rad).to_mat4();
  m.mul_m(&r)
}

pub fn rot_z(m: Mat4<f32>, angle: f32) -> Mat4<f32> {
  let rad = angle::rad(deg_to_rad(angle));
  let r = Mat3::from_angle_z(rad).to_mat4();
  m.mul_m(&r)
}

pub fn fill_current_coord_vbo(data: &[Vec3<GLfloat>]) {
  let glfloat_size = std::mem::size_of::<GLfloat>();
  let buffer_size = (data.len() * 3 * glfloat_size) as gl::types::GLsizeiptr;
  unsafe {
    gl::BufferData(
      gl::ARRAY_BUFFER,
      buffer_size,
      std::cast::transmute(&data[0]),
      gl::STATIC_DRAW,
    );
  }
}

pub fn gen_buffer() -> GLuint {
  let mut n = 0 as GLuint;
  unsafe {
    gl::GenBuffers(1, &mut n);
  }
  n
}

pub fn delete_buffer(buffer: GLuint) {
  unsafe {
    gl::DeleteBuffers(1, &buffer);
  }
}

pub fn fill_current_color_vbo(data: &[Color3]) {
  let color3_size = std::mem::size_of::<Color3>();
  let buffer_size = (data.len() * color3_size) as gl::types::GLsizeiptr;
  unsafe {
    gl::BufferData(
      gl::ARRAY_BUFFER,
      buffer_size,
      std::cast::transmute(&data[0]),
      gl::STATIC_DRAW,
    );
  }
}

pub fn vertex_attrib_pointer(attr: GLuint) {
  let components_count = 3;
  let normalized = gl::FALSE;
  let stride = 0;
  unsafe {
    gl::VertexAttribPointer(
      attr,
      components_count,
      gl::FLOAT,
      normalized,
      stride,
      std::ptr::null(),
    );
  }
}

// vim: set tabstop=2 shiftwidth=2 softtabstop=2 expandtab:
<commit_msg>gl_helpers.rs: link_program: Removed ': GLint'<commit_after>// See LICENSE file for copyright and license details.

use std;
use gl;
use gl::types::{
  GLfloat,
  GLint,
  GLuint,
  GLchar,
  GLenum,
};
use cgmath::matrix::{
  Matrix,
  Mat4,
  Mat3,
  ToMat4,
};
use cgmath::vector::Vec3;
use cgmath::angle;
use misc::{
  c_str,
  deg_to_rad,
};
use color::Color3;

pub fn compile_shader(src: &str, shader_type: GLenum) -> GLuint {
  let shader = gl::CreateShader(shader_type);
  unsafe {
    gl::ShaderSource(shader, 1, &c_str(src), std::ptr::null());
    gl::CompileShader(shader);
    let mut status = gl::FALSE as GLint;
    gl::GetShaderiv(shader, gl::COMPILE_STATUS, &mut status);
    if status != (gl::TRUE as GLint) {
      let mut len = 0;
      gl::GetShaderiv(shader, gl::INFO_LOG_LENGTH, &mut len);
      // subtract 1 to skip the trailing null character
      let mut buf = std::vec::from_elem(len as uint - 1, 0u8);
      gl::GetShaderInfoLog(shader, len, std::ptr::mut_null(),
        buf.as_mut_ptr() as *mut GLchar
      );
      fail!("compile_shader(): " + std::str::raw::from_utf8(buf));
    }
  }
  shader
}

pub fn link_program(vertex_shader: GLuint, fragment_shader: GLuint) -> GLuint {
  let program = gl::CreateProgram();
  gl::AttachShader(program, vertex_shader);
  gl::AttachShader(program, fragment_shader);
  gl::LinkProgram(program);
  unsafe {
    let mut status = gl::FALSE as GLint;
    gl::GetProgramiv(program, gl::LINK_STATUS, &mut status);
    if status != (gl::TRUE as GLint) {
      let mut len = 0;
      gl::GetProgramiv(program, gl::INFO_LOG_LENGTH, &mut len);
      // subtract 1 to skip the trailing null character
      let mut buf = std::vec::from_elem(len as uint - 1, 0u8);
      gl::GetProgramInfoLog(program, len, std::ptr::mut_null(),
        buf.as_mut_ptr() as *mut GLchar
      );
      fail!("link_program(): " + std::str::raw::from_utf8(buf));
    }
  }
  program
}

pub fn compile_program(
  vertex_shader_src: &str,
  frag_shader_src: &str,
) -> GLuint {
  let vertex_shader = compile_shader(
    vertex_shader_src, gl::VERTEX_SHADER);
  let fragment_shader = compile_shader(
    frag_shader_src, gl::FRAGMENT_SHADER);
  let program = link_program(vertex_shader, fragment_shader);
  // mark shaders for deletion after program deletion
  gl::DeleteShader(fragment_shader);
  gl::DeleteShader(vertex_shader);
  program
}

pub fn get_attr(program_id: GLuint, name: &str) -> GLuint {
  unsafe {
    gl::GetAttribLocation(program_id, c_str(name)) as GLuint
  }
}

pub fn get_uniform(program: GLuint, name: &str) -> GLint {
  unsafe {
    gl::GetUniformLocation(program, c_str(name))
  }
}

pub fn draw_mesh(faces_count: int) {
  let starting_index = 0;
  let vertices_count = faces_count as i32 * 3;
  gl::DrawArrays(gl::TRIANGLES, starting_index, vertices_count);
}

pub fn uniform_mat4f(mat_id: GLint, mat: &Mat4<GLfloat>) {
  unsafe {
    gl::UniformMatrix4fv(mat_id, 1, gl::FALSE, mat.cr(0, 0));
  }
}

pub fn tr(m: Mat4<f32>, v: Vec3<f32>) -> Mat4<f32> {
  let mut t = Mat4::<f32>::identity();
  *t.mut_cr(3, 0) = v.x;
  *t.mut_cr(3, 1) = v.y;
  *t.mut_cr(3, 2) = v.z;
  m.mul_m(&t)
}

pub fn rot_x(m: Mat4<f32>, angle: f32) -> Mat4<f32> {
  let rad = angle::rad(deg_to_rad(angle));
  let r = Mat3::from_angle_x(rad).to_mat4();
  m.mul_m(&r)
}

pub fn rot_z(m: Mat4<f32>, angle: f32) -> Mat4<f32> {
  let rad = angle::rad(deg_to_rad(angle));
  let r = Mat3::from_angle_z(rad).to_mat4();
  m.mul_m(&r)
}

pub fn fill_current_coord_vbo(data: &[Vec3<GLfloat>]) {
  let glfloat_size = std::mem::size_of::<GLfloat>();
  let buffer_size = (data.len() * 3 * glfloat_size) as gl::types::GLsizeiptr;
  unsafe {
    gl::BufferData(
      gl::ARRAY_BUFFER,
      buffer_size,
      std::cast::transmute(&data[0]),
      gl::STATIC_DRAW,
    );
  }
}

pub fn gen_buffer() -> GLuint {
  let mut n = 0 as GLuint;
  unsafe {
    gl::GenBuffers(1, &mut n);
  }
  n
}

pub fn delete_buffer(buffer: GLuint) {
  unsafe {
    gl::DeleteBuffers(1, &buffer);
  }
}

pub fn fill_current_color_vbo(data: &[Color3]) {
  let color3_size = std::mem::size_of::<Color3>();
  let buffer_size = (data.len() * color3_size) as gl::types::GLsizeiptr;
  unsafe {
    gl::BufferData(
      gl::ARRAY_BUFFER,
      buffer_size,
      std::cast::transmute(&data[0]),
      gl::STATIC_DRAW,
    );
  }
}

pub fn vertex_attrib_pointer(attr: GLuint) {
  let components_count = 3;
  let normalized = gl::FALSE;
  let stride = 0;
  unsafe {
    gl::VertexAttribPointer(
      attr,
      components_count,
      gl::FLOAT,
      normalized,
      stride,
      std::ptr::null(),
    );
  }
}

// vim: set tabstop=2 shiftwidth=2 softtabstop=2 expandtab:
<|endoftext|>


<commit_before>use std::ffi::CString;
use libc::{c_double, c_int};
use vector::Defn;
use utils::{_string, _last_null_pointer_err};
use gdal_sys::{self, OGRErr, OGRFeatureH, OGRFieldType};
use vector::geometry::Geometry;
use vector::layer::Layer;
use chrono::{Date, FixedOffset, DateTime, TimeZone, Datelike, Timelike};

use errors::*;

/// OGR Feature
pub struct Feature<'a> {
    _defn: &'a Defn,
    c_feature: OGRFeatureH,
    geometry: Vec<Geometry>,
}

impl<'a> Feature<'a> {
    pub fn new(defn: &'a Defn) -> Result<Feature> {
        let c_feature = unsafe { gdal_sys::OGR_F_Create(defn.c_defn()) };
        if c_feature.is_null() {
            Err(_last_null_pointer_err("OGR_F_Create"))?;
        };
        Ok(Feature {
                 _defn: defn,
                 c_feature,
                 geometry: Feature::_lazy_feature_geometries(defn),
             })
    }

    pub unsafe fn _with_c_feature(defn: &'a Defn, c_feature: OGRFeatureH) -> Feature {
        Feature {
            _defn: defn,
            c_feature,
            geometry: Feature::_lazy_feature_geometries(defn),
        }
    }

    pub fn _lazy_feature_geometries(defn: &'a Defn) -> Vec<Geometry> {
        let geom_field_count = unsafe { gdal_sys::OGR_FD_GetGeomFieldCount(defn.c_defn()) } as isize;
        (0..geom_field_count).map(|_| unsafe { Geometry::lazy_feature_geometry() }).collect()
    }

    /// Get the value of a named field. If the field exists, it returns a
    /// `FieldValue` wrapper, that you need to unpack to a base type
    /// (string, float, etc). If the field is missing, returns `None`.
    pub fn field(&self, name: &str) -> Result<FieldValue> {
        let c_name = CString::new(name)?;
        let field_id = unsafe { gdal_sys::OGR_F_GetFieldIndex(self.c_feature, c_name.as_ptr()) };
        if field_id == -1 {
            Err(ErrorKind::InvalidFieldName{field_name: name.to_string(), method_name: "OGR_F_GetFieldIndex"})?;
        }
        let field_defn = unsafe { gdal_sys::OGR_F_GetFieldDefnRef(self.c_feature, field_id) };
        let field_type = unsafe { gdal_sys::OGR_Fld_GetType(field_defn) };
        match field_type {
            OGRFieldType::OFTString => {
                let rv = unsafe { gdal_sys::OGR_F_GetFieldAsString(self.c_feature, field_id) };
                Ok(FieldValue::StringValue(_string(rv)))
            },
            OGRFieldType::OFTReal => {
                let rv = unsafe { gdal_sys::OGR_F_GetFieldAsDouble(self.c_feature, field_id) };
                Ok(FieldValue::RealValue(rv as f64))
            },
            OGRFieldType::OFTInteger => {
                let rv = unsafe { gdal_sys::OGR_F_GetFieldAsInteger(self.c_feature, field_id) };
                Ok(FieldValue::IntegerValue(rv as i32))
            },
            OGRFieldType::OFTDateTime => {
                Ok(FieldValue::DateTimeValue(self.get_field_datetime(field_id)?))
            },
            OGRFieldType::OFTDate => {
                Ok(FieldValue::DateValue(self.get_field_datetime(field_id)?.date()))
            },
            _ => Err(ErrorKind::UnhandledFieldType{field_type, method_name: "OGR_Fld_GetType"})?
        }
    }

    fn get_field_datetime(&self, field_id: c_int) -> Result<DateTime<FixedOffset>> {
        let mut year: c_int = 0;
        let mut month: c_int = 0;
        let mut day: c_int = 0;
        let mut hour: c_int = 0;
        let mut minute: c_int = 0;
        let mut second: c_int = 0;
        let mut tzflag: c_int = 0;

        let success = unsafe {
            gdal_sys::OGR_F_GetFieldAsDateTime(
                self.c_feature, field_id,
                &mut year, &mut month, &mut day, &mut hour, &mut minute, &mut second, &mut tzflag
            )
        };
        if success == 0 {
            Err(ErrorKind::OgrError { err: OGRErr::OGRERR_FAILURE, method_name: "OGR_F_GetFieldAsDateTime" })?;
        }

        // from https://github.com/OSGeo/gdal/blob/33a8a0edc764253b582e194d330eec3b83072863/gdal/ogr/ogrutils.cpp#L1309
        let tzoffset_secs = if tzflag == 0 || tzflag == 100 {
            0
        } else {
            (tzflag as i32 - 100) * 15 * 60
        };
        let rv = FixedOffset::east(tzoffset_secs)
            .ymd(year as i32, month as u32, day as u32)
            .and_hms(hour as u32, minute as u32, second as u32);
        Ok(rv)
    }

    /// Get the field's geometry.
    pub fn geometry(&self) -> &Geometry {
        if !self.geometry[0].has_gdal_ptr() {
            let c_geom = unsafe { gdal_sys::OGR_F_GetGeometryRef(self.c_feature) };
            unsafe { self.geometry[0].set_c_geometry(c_geom) };
        }
        &self.geometry[0]
    }

    pub fn geometry_by_name(&self, field_name: &str) -> Result<&Geometry> {
        let c_str_field_name = CString::new(field_name)?;
        let idx = unsafe { gdal_sys::OGR_F_GetGeomFieldIndex(self.c_feature, c_str_field_name.as_ptr())};
        if idx == -1 {
            Err(ErrorKind::InvalidFieldName{field_name: field_name.to_string(), method_name: "geometry_by_name"})?
        } else {
            self.geometry_by_index(idx as usize)
        }
    }

    pub fn geometry_by_index(&self, idx: usize) -> Result<&Geometry> {
        if idx >= self.geometry.len() {
            Err(ErrorKind::InvalidFieldIndex{index: idx, method_name: "geometry_by_name"})?;
        }
        if ! self.geometry[idx].has_gdal_ptr() {
            let c_geom = unsafe { gdal_sys::OGR_F_GetGeomFieldRef(self.c_feature, idx as i32) };
            if c_geom.is_null() {
                Err(_last_null_pointer_err("OGR_F_GetGeomFieldRef"))?;
            }
            unsafe { self.geometry[idx].set_c_geometry(c_geom) };
        }
        Ok(&self.geometry[idx])
    }

    pub fn create(&self, lyr: &Layer) -> Result<()> {
        let rv = unsafe { gdal_sys::OGR_L_CreateFeature(lyr.c_layer(), self.c_feature) };
        if rv != OGRErr::OGRERR_NONE {
            Err(ErrorKind::OgrError{err: rv, method_name: "OGR_L_CreateFeature"})?;
        }
        Ok(())
    }

    pub fn set_field_string(&self, field_name: &str, value: &str) -> Result<()> {
        let c_str_field_name = CString::new(field_name)?;
        let c_str_value = CString::new(value)?;
        let idx = unsafe { gdal_sys::OGR_F_GetFieldIndex(self.c_feature, c_str_field_name.as_ptr())};
        if idx == -1 {
            Err(ErrorKind::InvalidFieldName{field_name: field_name.to_string(), method_name: "OGR_F_GetFieldIndex"})?;
        }
        unsafe { gdal_sys::OGR_F_SetFieldString(self.c_feature, idx, c_str_value.as_ptr()) };
        Ok(())
    }

    pub fn set_field_double(&self, field_name: &str, value: f64) -> Result<()> {
        let c_str_field_name = CString::new(field_name)?;
        let idx = unsafe { gdal_sys::OGR_F_GetFieldIndex(self.c_feature, c_str_field_name.as_ptr())};
        if idx == -1 {
            Err(ErrorKind::InvalidFieldName{field_name: field_name.to_string(), method_name: "OGR_F_GetFieldIndex"})?;
        }
        unsafe { gdal_sys::OGR_F_SetFieldDouble(self.c_feature, idx, value as c_double) };
        Ok(())
    }

    pub fn set_field_integer(&self, field_name: &str, value: i32) -> Result<()> {
        let c_str_field_name = CString::new(field_name)?;
        let idx = unsafe { gdal_sys::OGR_F_GetFieldIndex(self.c_feature, c_str_field_name.as_ptr())};
        if idx == -1 {
            Err(ErrorKind::InvalidFieldName{field_name: field_name.to_string(), method_name: "OGR_F_GetFieldIndex"})?;
        }
        unsafe { gdal_sys::OGR_F_SetFieldInteger(self.c_feature, idx, value as c_int) };
        Ok(())
    }

    pub fn set_field_datetime(&self, field_name: &str, value: DateTime<FixedOffset>) -> Result<()> {
        let c_str_field_name = CString::new(field_name)?;
        let idx = unsafe { gdal_sys::OGR_F_GetFieldIndex(self.c_feature, c_str_field_name.as_ptr())};
        if idx == -1 {
            Err(ErrorKind::InvalidFieldName{field_name: field_name.to_string(), method_name: "OGR_F_GetFieldIndex"})?;
        }

        let year = value.year() as c_int;
        let month = value.month() as c_int;
        let day = value.day() as c_int;
        let hour= value.hour() as c_int;
        let minute = value.minute() as c_int;
        let second = value.second() as c_int;
        let tzflag: c_int = if value.offset().local_minus_utc() == 0 {
            0
        } else {
            100 + (value.offset().local_minus_utc() / (15 * 60))
        };

        unsafe { gdal_sys::OGR_F_SetFieldDateTime(self.c_feature, idx, year, month, day, hour, minute, second, tzflag) };
        Ok(())
    }

    pub fn set_field(&self, field_name: &str,  value: &FieldValue) -> Result<()> {
          match *value {
             FieldValue::RealValue(value) => self.set_field_double(field_name, value),
             FieldValue::StringValue(ref value) => self.set_field_string(field_name, value.as_str()),
             FieldValue::IntegerValue(value) => self.set_field_integer(field_name, value),
             FieldValue::DateTimeValue(value) => self.set_field_datetime(field_name, value),
             FieldValue::DateValue(value) => self.set_field_datetime(field_name, value.and_hms(0, 0, 0)),
         }
     }

    pub fn set_geometry(&mut self, geom: Geometry) -> Result<()> {
        let rv = unsafe { gdal_sys::OGR_F_SetGeometry(self.c_feature, geom.c_geometry()) };
        if rv != OGRErr::OGRERR_NONE {
            Err(ErrorKind::OgrError{err: rv, method_name: "OGR_G_SetGeometry"})?;
        }
        self.geometry[0] = geom;
        Ok(())
    }
}

impl<'a> Drop for Feature<'a> {
    fn drop(&mut self) {
        unsafe { gdal_sys::OGR_F_Destroy(self.c_feature); }
    }
}

pub enum FieldValue {
    IntegerValue(i32),
    StringValue(String),
    RealValue(f64),
    DateValue(Date<FixedOffset>),
    DateTimeValue(DateTime<FixedOffset>),
}

impl FieldValue {
    /// Interpret the value as `String`. Panics if the value is something else.
    pub fn into_string(self) -> Option<String> {
        match self {
            FieldValue::StringValue(rv) => Some(rv),
            _ => None
        }
    }

    /// Interpret the value as `f64`. Panics if the value is something else.
    pub fn into_real(self) -> Option<f64> {
        match self {
            FieldValue::RealValue(rv) => Some(rv),
            _ => None
        }
    }

    /// Interpret the value as `i32`. Panics if the value is something else.
    pub fn into_int(self) -> Option<i32> {
        match self {
            FieldValue::IntegerValue(rv) => Some(rv),
            _ => None
        }
    }

    /// Interpret the value as `Date`.
    pub fn into_date(self) -> Option<Date<FixedOffset>> {
        match self {
            FieldValue::DateValue(rv) => Some(rv),
            FieldValue::DateTimeValue(rv) => Some(rv.date()),
            _ => None
        }
    }

    /// Interpret the value as `DateTime`.
    pub fn into_datetime(self) -> Option<DateTime<FixedOffset>> {
        match self {
            FieldValue::DateTimeValue(rv) => Some(rv),
            _ => None
        }
    }
}
<commit_msg>feature gate for datetime support<commit_after>use std::ffi::CString;
use libc::{c_double, c_int};
use vector::Defn;
use utils::{_string, _last_null_pointer_err};
use gdal_sys::{self, OGRErr, OGRFeatureH, OGRFieldType};
use vector::geometry::Geometry;
use vector::layer::Layer;

#[cfg(feature = "datetime")]
use chrono::{Date, FixedOffset, DateTime, TimeZone, Datelike, Timelike};

use errors::*;

/// OGR Feature
pub struct Feature<'a> {
    _defn: &'a Defn,
    c_feature: OGRFeatureH,
    geometry: Vec<Geometry>,
}

impl<'a> Feature<'a> {
    pub fn new(defn: &'a Defn) -> Result<Feature> {
        let c_feature = unsafe { gdal_sys::OGR_F_Create(defn.c_defn()) };
        if c_feature.is_null() {
            Err(_last_null_pointer_err("OGR_F_Create"))?;
        };
        Ok(Feature {
                 _defn: defn,
                 c_feature,
                 geometry: Feature::_lazy_feature_geometries(defn),
             })
    }

    pub unsafe fn _with_c_feature(defn: &'a Defn, c_feature: OGRFeatureH) -> Feature {
        Feature {
            _defn: defn,
            c_feature,
            geometry: Feature::_lazy_feature_geometries(defn),
        }
    }

    pub fn _lazy_feature_geometries(defn: &'a Defn) -> Vec<Geometry> {
        let geom_field_count = unsafe { gdal_sys::OGR_FD_GetGeomFieldCount(defn.c_defn()) } as isize;
        (0..geom_field_count).map(|_| unsafe { Geometry::lazy_feature_geometry() }).collect()
    }

    /// Get the value of a named field. If the field exists, it returns a
    /// `FieldValue` wrapper, that you need to unpack to a base type
    /// (string, float, etc). If the field is missing, returns `None`.
    pub fn field(&self, name: &str) -> Result<FieldValue> {
        let c_name = CString::new(name)?;
        let field_id = unsafe { gdal_sys::OGR_F_GetFieldIndex(self.c_feature, c_name.as_ptr()) };
        if field_id == -1 {
            Err(ErrorKind::InvalidFieldName{field_name: name.to_string(), method_name: "OGR_F_GetFieldIndex"})?;
        }
        let field_defn = unsafe { gdal_sys::OGR_F_GetFieldDefnRef(self.c_feature, field_id) };
        let field_type = unsafe { gdal_sys::OGR_Fld_GetType(field_defn) };
        match field_type {
            OGRFieldType::OFTString => {
                let rv = unsafe { gdal_sys::OGR_F_GetFieldAsString(self.c_feature, field_id) };
                Ok(FieldValue::StringValue(_string(rv)))
            },
            OGRFieldType::OFTReal => {
                let rv = unsafe { gdal_sys::OGR_F_GetFieldAsDouble(self.c_feature, field_id) };
                Ok(FieldValue::RealValue(rv as f64))
            },
            OGRFieldType::OFTInteger => {
                let rv = unsafe { gdal_sys::OGR_F_GetFieldAsInteger(self.c_feature, field_id) };
                Ok(FieldValue::IntegerValue(rv as i32))
            },
            #[cfg(feature = "datetime")]
            OGRFieldType::OFTDateTime => {
                Ok(FieldValue::DateTimeValue(self.get_field_datetime(field_id)?))
            },
            #[cfg(feature = "datetime")]
            OGRFieldType::OFTDate => {
                Ok(FieldValue::DateValue(self.get_field_datetime(field_id)?.date()))
            },
            _ => Err(ErrorKind::UnhandledFieldType{field_type, method_name: "OGR_Fld_GetType"})?
        }
    }

    #[cfg(feature = "datetime")]
    fn get_field_datetime(&self, field_id: c_int) -> Result<DateTime<FixedOffset>> {
        let mut year: c_int = 0;
        let mut month: c_int = 0;
        let mut day: c_int = 0;
        let mut hour: c_int = 0;
        let mut minute: c_int = 0;
        let mut second: c_int = 0;
        let mut tzflag: c_int = 0;

        let success = unsafe {
            gdal_sys::OGR_F_GetFieldAsDateTime(
                self.c_feature, field_id,
                &mut year, &mut month, &mut day, &mut hour, &mut minute, &mut second, &mut tzflag
            )
        };
        if success == 0 {
            Err(ErrorKind::OgrError { err: OGRErr::OGRERR_FAILURE, method_name: "OGR_F_GetFieldAsDateTime" })?;
        }

        // from https://github.com/OSGeo/gdal/blob/33a8a0edc764253b582e194d330eec3b83072863/gdal/ogr/ogrutils.cpp#L1309
        let tzoffset_secs = if tzflag == 0 || tzflag == 100 {
            0
        } else {
            (tzflag as i32 - 100) * 15 * 60
        };
        let rv = FixedOffset::east(tzoffset_secs)
            .ymd(year as i32, month as u32, day as u32)
            .and_hms(hour as u32, minute as u32, second as u32);
        Ok(rv)
    }

    /// Get the field's geometry.
    pub fn geometry(&self) -> &Geometry {
        if !self.geometry[0].has_gdal_ptr() {
            let c_geom = unsafe { gdal_sys::OGR_F_GetGeometryRef(self.c_feature) };
            unsafe { self.geometry[0].set_c_geometry(c_geom) };
        }
        &self.geometry[0]
    }

    pub fn geometry_by_name(&self, field_name: &str) -> Result<&Geometry> {
        let c_str_field_name = CString::new(field_name)?;
        let idx = unsafe { gdal_sys::OGR_F_GetGeomFieldIndex(self.c_feature, c_str_field_name.as_ptr())};
        if idx == -1 {
            Err(ErrorKind::InvalidFieldName{field_name: field_name.to_string(), method_name: "geometry_by_name"})?
        } else {
            self.geometry_by_index(idx as usize)
        }
    }

    pub fn geometry_by_index(&self, idx: usize) -> Result<&Geometry> {
        if idx >= self.geometry.len() {
            Err(ErrorKind::InvalidFieldIndex{index: idx, method_name: "geometry_by_name"})?;
        }
        if ! self.geometry[idx].has_gdal_ptr() {
            let c_geom = unsafe { gdal_sys::OGR_F_GetGeomFieldRef(self.c_feature, idx as i32) };
            if c_geom.is_null() {
                Err(_last_null_pointer_err("OGR_F_GetGeomFieldRef"))?;
            }
            unsafe { self.geometry[idx].set_c_geometry(c_geom) };
        }
        Ok(&self.geometry[idx])
    }

    pub fn create(&self, lyr: &Layer) -> Result<()> {
        let rv = unsafe { gdal_sys::OGR_L_CreateFeature(lyr.c_layer(), self.c_feature) };
        if rv != OGRErr::OGRERR_NONE {
            Err(ErrorKind::OgrError{err: rv, method_name: "OGR_L_CreateFeature"})?;
        }
        Ok(())
    }

    pub fn set_field_string(&self, field_name: &str, value: &str) -> Result<()> {
        let c_str_field_name = CString::new(field_name)?;
        let c_str_value = CString::new(value)?;
        let idx = unsafe { gdal_sys::OGR_F_GetFieldIndex(self.c_feature, c_str_field_name.as_ptr())};
        if idx == -1 {
            Err(ErrorKind::InvalidFieldName{field_name: field_name.to_string(), method_name: "OGR_F_GetFieldIndex"})?;
        }
        unsafe { gdal_sys::OGR_F_SetFieldString(self.c_feature, idx, c_str_value.as_ptr()) };
        Ok(())
    }

    pub fn set_field_double(&self, field_name: &str, value: f64) -> Result<()> {
        let c_str_field_name = CString::new(field_name)?;
        let idx = unsafe { gdal_sys::OGR_F_GetFieldIndex(self.c_feature, c_str_field_name.as_ptr())};
        if idx == -1 {
            Err(ErrorKind::InvalidFieldName{field_name: field_name.to_string(), method_name: "OGR_F_GetFieldIndex"})?;
        }
        unsafe { gdal_sys::OGR_F_SetFieldDouble(self.c_feature, idx, value as c_double) };
        Ok(())
    }

    pub fn set_field_integer(&self, field_name: &str, value: i32) -> Result<()> {
        let c_str_field_name = CString::new(field_name)?;
        let idx = unsafe { gdal_sys::OGR_F_GetFieldIndex(self.c_feature, c_str_field_name.as_ptr())};
        if idx == -1 {
            Err(ErrorKind::InvalidFieldName{field_name: field_name.to_string(), method_name: "OGR_F_GetFieldIndex"})?;
        }
        unsafe { gdal_sys::OGR_F_SetFieldInteger(self.c_feature, idx, value as c_int) };
        Ok(())
    }

    #[cfg(feature = "datetime")]
    pub fn set_field_datetime(&self, field_name: &str, value: DateTime<FixedOffset>) -> Result<()> {
        let c_str_field_name = CString::new(field_name)?;
        let idx = unsafe { gdal_sys::OGR_F_GetFieldIndex(self.c_feature, c_str_field_name.as_ptr())};
        if idx == -1 {
            Err(ErrorKind::InvalidFieldName{field_name: field_name.to_string(), method_name: "OGR_F_GetFieldIndex"})?;
        }

        let year = value.year() as c_int;
        let month = value.month() as c_int;
        let day = value.day() as c_int;
        let hour= value.hour() as c_int;
        let minute = value.minute() as c_int;
        let second = value.second() as c_int;
        let tzflag: c_int = if value.offset().local_minus_utc() == 0 {
            0
        } else {
            100 + (value.offset().local_minus_utc() / (15 * 60))
        };

        unsafe { gdal_sys::OGR_F_SetFieldDateTime(self.c_feature, idx, year, month, day, hour, minute, second, tzflag) };
        Ok(())
    }

    pub fn set_field(&self, field_name: &str,  value: &FieldValue) -> Result<()> {
          match *value {
             FieldValue::RealValue(value) => self.set_field_double(field_name, value),
             FieldValue::StringValue(ref value) => self.set_field_string(field_name, value.as_str()),
             FieldValue::IntegerValue(value) => self.set_field_integer(field_name, value),

             #[cfg(feature = "datetime")]
             FieldValue::DateTimeValue(value) => self.set_field_datetime(field_name, value),

              #[cfg(feature = "datetime")]
             FieldValue::DateValue(value) => self.set_field_datetime(field_name, value.and_hms(0, 0, 0)),
         }
     }

    pub fn set_geometry(&mut self, geom: Geometry) -> Result<()> {
        let rv = unsafe { gdal_sys::OGR_F_SetGeometry(self.c_feature, geom.c_geometry()) };
        if rv != OGRErr::OGRERR_NONE {
            Err(ErrorKind::OgrError{err: rv, method_name: "OGR_G_SetGeometry"})?;
        }
        self.geometry[0] = geom;
        Ok(())
    }
}

impl<'a> Drop for Feature<'a> {
    fn drop(&mut self) {
        unsafe { gdal_sys::OGR_F_Destroy(self.c_feature); }
    }
}

pub enum FieldValue {
    IntegerValue(i32),
    StringValue(String),
    RealValue(f64),

    #[cfg(feature = "datetime")]
    DateValue(Date<FixedOffset>),

    #[cfg(feature = "datetime")]
    DateTimeValue(DateTime<FixedOffset>),
}

impl FieldValue {
    /// Interpret the value as `String`. Panics if the value is something else.
    pub fn into_string(self) -> Option<String> {
        match self {
            FieldValue::StringValue(rv) => Some(rv),
            _ => None
        }
    }

    /// Interpret the value as `f64`. Panics if the value is something else.
    pub fn into_real(self) -> Option<f64> {
        match self {
            FieldValue::RealValue(rv) => Some(rv),
            _ => None
        }
    }

    /// Interpret the value as `i32`. Panics if the value is something else.
    pub fn into_int(self) -> Option<i32> {
        match self {
            FieldValue::IntegerValue(rv) => Some(rv),
            _ => None
        }
    }

    /// Interpret the value as `Date`.
    #[cfg(feature = "datetime")]
    pub fn into_date(self) -> Option<Date<FixedOffset>> {
        match self {
            FieldValue::DateValue(rv) => Some(rv),
            FieldValue::DateTimeValue(rv) => Some(rv.date()),
            _ => None
        }
    }

    /// Interpret the value as `DateTime`.
    #[cfg(feature = "datetime")]
    pub fn into_datetime(self) -> Option<DateTime<FixedOffset>> {
        match self {
            FieldValue::DateTimeValue(rv) => Some(rv),
            _ => None
        }
    }
}
<|endoftext|>