<commit_before>use calx_grid::Dir6;
use display;
use euclid::{Point2D, Rect};
use scancode::Scancode;
use std::fs::File;
use std::io::prelude::*;
use vitral::{Context, FracPoint2D, FracSize2D, FracRect, Align};
use world::{Command, CommandResult, Event, ItemType, Location, Query, Slot, World};

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
enum State {
    Main,
    Inventory(InventoryAction),
    Console,
    Aim(AimAction),
}

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
enum InventoryAction {
    Drop,
    Equip,
    Use,
}

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
enum AimAction {
    Zap(Slot),
    // Maybe add intrinsic abilities not tied to a specific entity later
}

pub struct GameLoop {
    pub world: World,
    pub console: display::Console,
    camera_loc: Location,
    state: State,
}

impl GameLoop {
    pub fn new(world: World) -> GameLoop {
        GameLoop {
            world,
            console: display::Console::default(),
            camera_loc: Location::new(0, 0, 0),
            state: State::Main,
        }
    }

    /// Step command that turns into melee attack if an enemy is in the way.
    fn smart_step(&mut self, dir: Dir6) -> CommandResult {
        let player = self.world.player().ok_or(())?;
        let loc = self.world.location(player).ok_or(())?;
        let destination = loc.jump(&self.world, dir);

        if let Some(mob) = self.world.mob_at(destination) {
            if self.world.is_hostile_to(player, mob) {
                // Fight on!
                self.world.melee(dir)
            } else {
                // Do we want to do something smarter than walk into friendlies?
                // The world might treat this as a displace action so keep it like this for now.
                self.world.step(dir)
            }
        } else {
            self.world.step(dir)
        }
    }

    fn game_input(&mut self, scancode: Scancode) -> CommandResult {
        use scancode::Scancode::*;
        match scancode {
            Tab => {
                self.enter_state(State::Console);
                Ok(Vec::new())
            }
            Q => self.smart_step(Dir6::Northwest),
            W => self.smart_step(Dir6::North),
            E => self.smart_step(Dir6::Northeast),
            A => self.smart_step(Dir6::Southwest),
            S => self.smart_step(Dir6::South),
            D => self.smart_step(Dir6::Southeast),
            I => {
                self.enter_state(State::Inventory(InventoryAction::Equip));
                Ok(Vec::new())
            }
            // XXX: Key mnemonics, bit awkward when D is taken by movement.
            B => {
                self.enter_state(State::Inventory(InventoryAction::Drop));
                Ok(Vec::new())
            }
            U => {
                self.enter_state(State::Inventory(InventoryAction::Use));
                Ok(Vec::new())
            }
            G => self.world.take(),
            Space => self.world.pass(),
            F5 => {
                self.world
                    .save(&mut File::create("save.gam").unwrap())
                    .unwrap();
                Ok(Vec::new())
            }
            F9 => {
                let mut savefile = File::open("save.gam").unwrap();
                self.world = World::load(&mut savefile).unwrap();
                Ok(Vec::new())
            }
            _ => Ok(Vec::new()),
        }
    }

    fn zap(&mut self, slot: Slot, dir: Dir6) -> CommandResult {
        let ret = self.world.zap_item(slot, dir)?;
        self.enter_state(State::Main);
        Ok(ret)
    }

    fn aim_input(&mut self, slot: Slot, scancode: Scancode) -> CommandResult {
        use scancode::Scancode::*;
        match scancode {
            Q => self.zap(slot, Dir6::Northwest),
            W => self.zap(slot, Dir6::North),
            E => self.zap(slot, Dir6::Northeast),
            A => self.zap(slot, Dir6::Southwest),
            S => self.zap(slot, Dir6::South),
            D => self.zap(slot, Dir6::Southeast),
            Escape => {
                self.enter_state(State::Main);
                Ok(Vec::new())
            }
            _ => Ok(Vec::new()),
        }
    }

    fn inventory_input(&mut self, scancode: Scancode) -> CommandResult {
        use scancode::Scancode::*;
        for slot in SLOT_DATA.iter() {
            if scancode == slot.code {
                if let State::Inventory(action) = self.state {
                    let ret = self.inventory_action(slot.slot, action);
                    if ret.is_ok() {
                        return ret;
                    }
                }
            }
        }

        match scancode {
            Escape => {
                self.enter_state(State::Main);
                Ok(Vec::new())
            }
            _ => Ok(Vec::new()),
        }
    }

    fn enter_state(&mut self, new_state: State) {
        if self.state == new_state {
            return;
        }

        if let State::Aim(_) = new_state {
            let _ = writeln!(&mut self.console, "Press direction to aim, Esc to cancel");
        }

        self.state = new_state;
    }

    fn inventory_action(&mut self, slot: Slot, action: InventoryAction) -> CommandResult {
        match action {
            InventoryAction::Drop => {
                let ret = self.world.drop(slot);
                // After succesful drop, go back to main state.
                if ret.is_ok() {
                    self.enter_state(State::Main);
                }
                ret
            }
            // Can equip multiple items in one go, wait for ESC to return to main state.
            InventoryAction::Equip => self.world.equip(slot),
            InventoryAction::Use => {
                let player = self.world.player().ok_or(())?;

                if let Some(item) = self.world.entity_equipped(player, slot) {
                    match self.world.item_type(item) {
                        Some(ItemType::UntargetedUsable(_)) => {
                            let ret = self.world.use_item(slot)?;
                            self.enter_state(State::Main);
                            return Ok(ret);
                        }
                        Some(ItemType::TargetedUsable(_)) => {
                            // If we need to aim, switch to aim state before calling world.
                            self.enter_state(State::Aim(AimAction::Zap(slot)));
                            return Ok(Vec::new());
                        }
                        _ => {}
                    }
                }
                Err(())
            }
        }
    }

    fn console_input(&mut self, scancode: Scancode) -> CommandResult {
        use scancode::Scancode::*;
        match scancode {
            Tab => {
                self.enter_state(State::Main);
                Ok(Vec::new())
            }
            Enter | PadEnter => {
                let input = self.console.get_input();
                let _ = writeln!(&mut self.console, "{}", input);
                if let Err(e) = self.parse(&input) {
                    let _ = writeln!(&mut self.console, "{}", e);
                }
                Ok(Vec::new())
            }
            _ => Ok(Vec::new()),
        }
    }

    fn todo(&mut self) {
        // TODO: Bring back some debug commands
    }

    command_parser!{
        fn todo(&mut self);
    }

    fn draw_inventory(&mut self, c: &mut display::Backend) -> Result<(), ()> {
        let player = self.world.player().ok_or(())?;

        // Start with hardcoded invetory data to test the UI logic.
        c.fill_rect(
            FracRect::new(FracPoint2D::new(0.0, 0.0), FracSize2D::new(1.0, 1.0)),
            [0.0, 0.0, 0.0, 0.99],
        );

        let mut letter_pos = Point2D::new(0.0, 0.0);
        let mut slot_name_pos = Point2D::new(20.0, 0.0);
        let mut item_name_pos = Point2D::new(80.0, 0.0);
        let text_color = [1.0, 1.0, 1.0, 1.0];

        for slot in SLOT_DATA.iter() {
            // TODO: Bounding box for these is a button...
            letter_pos = c.draw_text(
                letter_pos,
                Align::Left,
                text_color,
                &format!("{})", slot.key),
            );
            slot_name_pos = c.draw_text(slot_name_pos, Align::Left, text_color, slot.name);
            let item_name = if let Some(item) = self.world.entity_equipped(player, slot.slot) {
                self.world.entity_name(item)
            } else {
                "".to_string()
            };

            item_name_pos = c.draw_text(item_name_pos, Align::Left, text_color, &item_name);
        }

        Ok(())
    }

    pub fn draw(&mut self, context: &mut display::Backend, screen_area: &Rect<f32>) {
        // Ugh
        self.world.player().map(|x| {
            self.world.location(x).map(|l| self.camera_loc = l)
        });

        let mut view = display::WorldView::new(self.camera_loc, *screen_area);
        view.show_cursor = true;

        view.draw(&self.world, context);

        match self.state {
            State::Inventory(_) => {
                let _ = self.draw_inventory(context);
            }
            State::Console => {
                let mut console_area = *screen_area;
                console_area.size.height = 184.0;
                self.console.draw_large(context, &console_area);
            }
            _ => {
                let mut console_area = *screen_area;
                console_area.size.height = 32.0;
                self.console.draw_small(context, &console_area);
            }
        }

        if let Some(scancode) = context.poll_key().and_then(|k| Scancode::new(k.scancode)) {
            let ret = match self.state {
                State::Inventory(_) => self.inventory_input(scancode),
                State::Console => self.console_input(scancode),
                State::Aim(AimAction::Zap(slot)) => self.aim_input(slot, scancode),
                _ => self.game_input(scancode),
            };

            if let Ok(events) = ret {
                // Input event caused a successful world step and we got an event sequence out.
                // Convert events into UI display effects.
                for e in events {
                    match e {
                        Event::Msg(text) => {
                            let _ = writeln!(&mut self.console, "{}", text);
                        }
                    }
                }
            }
        }
    }
}

struct SlotData {
    key: char,
    code: Scancode,
    slot: Slot,
    name: &'static str,
}

#[cfg_attr(rustfmt, rustfmt_skip)]
static SLOT_DATA: [SlotData; 34] = [
    SlotData { key: '1', code: Scancode::Num1, slot: Slot::Spell1,     name: "Ability" },
    SlotData { key: '2', code: Scancode::Num2, slot: Slot::Spell2,     name: "Ability" },
    SlotData { key: '3', code: Scancode::Num3, slot: Slot::Spell3,     name: "Ability" },
    SlotData { key: '4', code: Scancode::Num4, slot: Slot::Spell4,     name: "Ability" },
    SlotData { key: '5', code: Scancode::Num5, slot: Slot::Spell5,     name: "Ability" },
    SlotData { key: '6', code: Scancode::Num6, slot: Slot::Spell6,     name: "Ability" },
    SlotData { key: '7', code: Scancode::Num7, slot: Slot::Spell7,     name: "Ability" },
    SlotData { key: '8', code: Scancode::Num8, slot: Slot::Spell8,     name: "Ability" },
    SlotData { key: 'a', code: Scancode::A,    slot: Slot::Melee,      name: "Weapon" },
    SlotData { key: 'b', code: Scancode::B,    slot: Slot::Ranged,     name: "Ranged" },
    SlotData { key: 'c', code: Scancode::C,    slot: Slot::Head,       name: "Head" },
    SlotData { key: 'd', code: Scancode::D,    slot: Slot::Body,       name: "Body" },
    SlotData { key: 'e', code: Scancode::E,    slot: Slot::Feet,       name: "Feet" },
    SlotData { key: 'f', code: Scancode::F,    slot: Slot::TrinketF,   name: "Trinket" },
    SlotData { key: 'g', code: Scancode::G,    slot: Slot::TrinketG,   name: "Trinket" },
    SlotData { key: 'h', code: Scancode::H,    slot: Slot::TrinketH,   name: "Trinket" },
    SlotData { key: 'i', code: Scancode::I,    slot: Slot::TrinketI,   name: "Trinket" },
    SlotData { key: 'j', code: Scancode::J,    slot: Slot::InventoryJ, name: "" },
    SlotData { key: 'k', code: Scancode::K,    slot: Slot::InventoryK, name: "" },
    SlotData { key: 'l', code: Scancode::L,    slot: Slot::InventoryL, name: "" },
    SlotData { key: 'm', code: Scancode::M,    slot: Slot::InventoryM, name: "" },
    SlotData { key: 'n', code: Scancode::N,    slot: Slot::InventoryN, name: "" },
    SlotData { key: 'o', code: Scancode::O,    slot: Slot::InventoryO, name: "" },
    SlotData { key: 'p', code: Scancode::P,    slot: Slot::InventoryP, name: "" },
    SlotData { key: 'q', code: Scancode::Q,    slot: Slot::InventoryQ, name: "" },
    SlotData { key: 'r', code: Scancode::R,    slot: Slot::InventoryR, name: "" },
    SlotData { key: 's', code: Scancode::S,    slot: Slot::InventoryS, name: "" },
    SlotData { key: 't', code: Scancode::T,    slot: Slot::InventoryT, name: "" },
    SlotData { key: 'u', code: Scancode::U,    slot: Slot::InventoryU, name: "" },
    SlotData { key: 'v', code: Scancode::V,    slot: Slot::InventoryV, name: "" },
    SlotData { key: 'w', code: Scancode::W,    slot: Slot::InventoryW, name: "" },
    SlotData { key: 'x', code: Scancode::X,    slot: Slot::InventoryX, name: "" },
    SlotData { key: 'y', code: Scancode::Y,    slot: Slot::InventoryY, name: "" },
    SlotData { key: 'z', code: Scancode::Z,    slot: Slot::InventoryZ, name: "" },
];
<commit_msg>Reactivate screenshotting<commit_after>use calx_grid::Dir6;
use display;
use euclid::{Point2D, Rect};
use scancode::Scancode;
use std::fs::File;
use std::io::prelude::*;
use vitral::{Context, FracPoint2D, FracSize2D, FracRect, Align};
use world::{Command, CommandResult, Event, ItemType, Location, Query, Slot, World};

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
enum State {
    Main,
    Inventory(InventoryAction),
    Console,
    Aim(AimAction),
}

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
enum InventoryAction {
    Drop,
    Equip,
    Use,
}

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
enum AimAction {
    Zap(Slot),
    // Maybe add intrinsic abilities not tied to a specific entity later
}

pub struct GameLoop {
    pub world: World,
    pub console: display::Console,
    camera_loc: Location,
    state: State,
}

impl GameLoop {
    pub fn new(world: World) -> GameLoop {
        GameLoop {
            world,
            console: display::Console::default(),
            camera_loc: Location::new(0, 0, 0),
            state: State::Main,
        }
    }

    /// Step command that turns into melee attack if an enemy is in the way.
    fn smart_step(&mut self, dir: Dir6) -> CommandResult {
        let player = self.world.player().ok_or(())?;
        let loc = self.world.location(player).ok_or(())?;
        let destination = loc.jump(&self.world, dir);

        if let Some(mob) = self.world.mob_at(destination) {
            if self.world.is_hostile_to(player, mob) {
                // Fight on!
                self.world.melee(dir)
            } else {
                // Do we want to do something smarter than walk into friendlies?
                // The world might treat this as a displace action so keep it like this for now.
                self.world.step(dir)
            }
        } else {
            self.world.step(dir)
        }
    }

    fn game_input(&mut self, context: &mut display::Backend, scancode: Scancode) -> CommandResult {
        use scancode::Scancode::*;
        match scancode {
            Tab => {
                self.enter_state(State::Console);
                Ok(Vec::new())
            }
            Q => self.smart_step(Dir6::Northwest),
            W => self.smart_step(Dir6::North),
            E => self.smart_step(Dir6::Northeast),
            A => self.smart_step(Dir6::Southwest),
            S => self.smart_step(Dir6::South),
            D => self.smart_step(Dir6::Southeast),
            I => {
                self.enter_state(State::Inventory(InventoryAction::Equip));
                Ok(Vec::new())
            }
            // XXX: Key mnemonics, bit awkward when D is taken by movement.
            B => {
                self.enter_state(State::Inventory(InventoryAction::Drop));
                Ok(Vec::new())
            }
            U => {
                self.enter_state(State::Inventory(InventoryAction::Use));
                Ok(Vec::new())
            }
            G => self.world.take(),
            Space => self.world.pass(),
            F5 => {
                self.world
                    .save(&mut File::create("save.gam").unwrap())
                    .unwrap();
                Ok(Vec::new())
            }
            F9 => {
                let mut savefile = File::open("save.gam").unwrap();
                self.world = World::load(&mut savefile).unwrap();
                Ok(Vec::new())
            }
            F12 => {
                context.save_screenshot("magog");
                Ok(Vec::new())
            }
            _ => Ok(Vec::new()),
        }
    }

    fn zap(&mut self, slot: Slot, dir: Dir6) -> CommandResult {
        let ret = self.world.zap_item(slot, dir)?;
        self.enter_state(State::Main);
        Ok(ret)
    }

    fn aim_input(&mut self, slot: Slot, scancode: Scancode) -> CommandResult {
        use scancode::Scancode::*;
        match scancode {
            Q => self.zap(slot, Dir6::Northwest),
            W => self.zap(slot, Dir6::North),
            E => self.zap(slot, Dir6::Northeast),
            A => self.zap(slot, Dir6::Southwest),
            S => self.zap(slot, Dir6::South),
            D => self.zap(slot, Dir6::Southeast),
            Escape => {
                self.enter_state(State::Main);
                Ok(Vec::new())
            }
            _ => Ok(Vec::new()),
        }
    }

    fn inventory_input(&mut self, scancode: Scancode) -> CommandResult {
        use scancode::Scancode::*;
        for slot in SLOT_DATA.iter() {
            if scancode == slot.code {
                if let State::Inventory(action) = self.state {
                    let ret = self.inventory_action(slot.slot, action);
                    if ret.is_ok() {
                        return ret;
                    }
                }
            }
        }

        match scancode {
            Escape => {
                self.enter_state(State::Main);
                Ok(Vec::new())
            }
            _ => Ok(Vec::new()),
        }
    }

    fn enter_state(&mut self, new_state: State) {
        if self.state == new_state {
            return;
        }

        if let State::Aim(_) = new_state {
            let _ = writeln!(&mut self.console, "Press direction to aim, Esc to cancel");
        }

        self.state = new_state;
    }

    fn inventory_action(&mut self, slot: Slot, action: InventoryAction) -> CommandResult {
        match action {
            InventoryAction::Drop => {
                let ret = self.world.drop(slot);
                // After succesful drop, go back to main state.
                if ret.is_ok() {
                    self.enter_state(State::Main);
                }
                ret
            }
            // Can equip multiple items in one go, wait for ESC to return to main state.
            InventoryAction::Equip => self.world.equip(slot),
            InventoryAction::Use => {
                let player = self.world.player().ok_or(())?;

                if let Some(item) = self.world.entity_equipped(player, slot) {
                    match self.world.item_type(item) {
                        Some(ItemType::UntargetedUsable(_)) => {
                            let ret = self.world.use_item(slot)?;
                            self.enter_state(State::Main);
                            return Ok(ret);
                        }
                        Some(ItemType::TargetedUsable(_)) => {
                            // If we need to aim, switch to aim state before calling world.
                            self.enter_state(State::Aim(AimAction::Zap(slot)));
                            return Ok(Vec::new());
                        }
                        _ => {}
                    }
                }
                Err(())
            }
        }
    }

    fn console_input(&mut self, scancode: Scancode) -> CommandResult {
        use scancode::Scancode::*;
        match scancode {
            Tab => {
                self.enter_state(State::Main);
                Ok(Vec::new())
            }
            Enter | PadEnter => {
                let input = self.console.get_input();
                let _ = writeln!(&mut self.console, "{}", input);
                if let Err(e) = self.parse(&input) {
                    let _ = writeln!(&mut self.console, "{}", e);
                }
                Ok(Vec::new())
            }
            _ => Ok(Vec::new()),
        }
    }

    fn todo(&mut self) {
        // TODO: Bring back some debug commands
    }

    command_parser!{
        fn todo(&mut self);
    }

    fn draw_inventory(&mut self, c: &mut display::Backend) -> Result<(), ()> {
        let player = self.world.player().ok_or(())?;

        // Start with hardcoded invetory data to test the UI logic.
        c.fill_rect(
            FracRect::new(FracPoint2D::new(0.0, 0.0), FracSize2D::new(1.0, 1.0)),
            [0.0, 0.0, 0.0, 0.99],
        );

        let mut letter_pos = Point2D::new(0.0, 0.0);
        let mut slot_name_pos = Point2D::new(20.0, 0.0);
        let mut item_name_pos = Point2D::new(80.0, 0.0);
        let text_color = [1.0, 1.0, 1.0, 1.0];

        for slot in SLOT_DATA.iter() {
            // TODO: Bounding box for these is a button...
            letter_pos = c.draw_text(
                letter_pos,
                Align::Left,
                text_color,
                &format!("{})", slot.key),
            );
            slot_name_pos = c.draw_text(slot_name_pos, Align::Left, text_color, slot.name);
            let item_name = if let Some(item) = self.world.entity_equipped(player, slot.slot) {
                self.world.entity_name(item)
            } else {
                "".to_string()
            };

            item_name_pos = c.draw_text(item_name_pos, Align::Left, text_color, &item_name);
        }

        Ok(())
    }

    pub fn draw(&mut self, context: &mut display::Backend, screen_area: &Rect<f32>) {
        // Ugh
        self.world.player().map(|x| {
            self.world.location(x).map(|l| self.camera_loc = l)
        });

        let mut view = display::WorldView::new(self.camera_loc, *screen_area);
        view.show_cursor = true;

        view.draw(&self.world, context);

        match self.state {
            State::Inventory(_) => {
                let _ = self.draw_inventory(context);
            }
            State::Console => {
                let mut console_area = *screen_area;
                console_area.size.height = 184.0;
                self.console.draw_large(context, &console_area);
            }
            _ => {
                let mut console_area = *screen_area;
                console_area.size.height = 32.0;
                self.console.draw_small(context, &console_area);
            }
        }

        if let Some(scancode) = context.poll_key().and_then(|k| Scancode::new(k.scancode)) {
            let ret = match self.state {
                State::Inventory(_) => self.inventory_input(scancode),
                State::Console => self.console_input(scancode),
                State::Aim(AimAction::Zap(slot)) => self.aim_input(slot, scancode),
                _ => self.game_input(context, scancode),
            };

            if let Ok(events) = ret {
                // Input event caused a successful world step and we got an event sequence out.
                // Convert events into UI display effects.
                for e in events {
                    match e {
                        Event::Msg(text) => {
                            let _ = writeln!(&mut self.console, "{}", text);
                        }
                    }
                }
            }
        }
    }
}

struct SlotData {
    key: char,
    code: Scancode,
    slot: Slot,
    name: &'static str,
}

#[cfg_attr(rustfmt, rustfmt_skip)]
static SLOT_DATA: [SlotData; 34] = [
    SlotData { key: '1', code: Scancode::Num1, slot: Slot::Spell1,     name: "Ability" },
    SlotData { key: '2', code: Scancode::Num2, slot: Slot::Spell2,     name: "Ability" },
    SlotData { key: '3', code: Scancode::Num3, slot: Slot::Spell3,     name: "Ability" },
    SlotData { key: '4', code: Scancode::Num4, slot: Slot::Spell4,     name: "Ability" },
    SlotData { key: '5', code: Scancode::Num5, slot: Slot::Spell5,     name: "Ability" },
    SlotData { key: '6', code: Scancode::Num6, slot: Slot::Spell6,     name: "Ability" },
    SlotData { key: '7', code: Scancode::Num7, slot: Slot::Spell7,     name: "Ability" },
    SlotData { key: '8', code: Scancode::Num8, slot: Slot::Spell8,     name: "Ability" },
    SlotData { key: 'a', code: Scancode::A,    slot: Slot::Melee,      name: "Weapon" },
    SlotData { key: 'b', code: Scancode::B,    slot: Slot::Ranged,     name: "Ranged" },
    SlotData { key: 'c', code: Scancode::C,    slot: Slot::Head,       name: "Head" },
    SlotData { key: 'd', code: Scancode::D,    slot: Slot::Body,       name: "Body" },
    SlotData { key: 'e', code: Scancode::E,    slot: Slot::Feet,       name: "Feet" },
    SlotData { key: 'f', code: Scancode::F,    slot: Slot::TrinketF,   name: "Trinket" },
    SlotData { key: 'g', code: Scancode::G,    slot: Slot::TrinketG,   name: "Trinket" },
    SlotData { key: 'h', code: Scancode::H,    slot: Slot::TrinketH,   name: "Trinket" },
    SlotData { key: 'i', code: Scancode::I,    slot: Slot::TrinketI,   name: "Trinket" },
    SlotData { key: 'j', code: Scancode::J,    slot: Slot::InventoryJ, name: "" },
    SlotData { key: 'k', code: Scancode::K,    slot: Slot::InventoryK, name: "" },
    SlotData { key: 'l', code: Scancode::L,    slot: Slot::InventoryL, name: "" },
    SlotData { key: 'm', code: Scancode::M,    slot: Slot::InventoryM, name: "" },
    SlotData { key: 'n', code: Scancode::N,    slot: Slot::InventoryN, name: "" },
    SlotData { key: 'o', code: Scancode::O,    slot: Slot::InventoryO, name: "" },
    SlotData { key: 'p', code: Scancode::P,    slot: Slot::InventoryP, name: "" },
    SlotData { key: 'q', code: Scancode::Q,    slot: Slot::InventoryQ, name: "" },
    SlotData { key: 'r', code: Scancode::R,    slot: Slot::InventoryR, name: "" },
    SlotData { key: 's', code: Scancode::S,    slot: Slot::InventoryS, name: "" },
    SlotData { key: 't', code: Scancode::T,    slot: Slot::InventoryT, name: "" },
    SlotData { key: 'u', code: Scancode::U,    slot: Slot::InventoryU, name: "" },
    SlotData { key: 'v', code: Scancode::V,    slot: Slot::InventoryV, name: "" },
    SlotData { key: 'w', code: Scancode::W,    slot: Slot::InventoryW, name: "" },
    SlotData { key: 'x', code: Scancode::X,    slot: Slot::InventoryX, name: "" },
    SlotData { key: 'y', code: Scancode::Y,    slot: Slot::InventoryY, name: "" },
    SlotData { key: 'z', code: Scancode::Z,    slot: Slot::InventoryZ, name: "" },
];
<|endoftext|>


<commit_before>// This file is part of libfringe, a low-level green threading library.
// Copyright (c) whitequark <whitequark@whitequark.org>
// 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.

//! Generators.
//!
//! Generators allow repeatedly suspending the execution of a function,
//! returning a value to the caller, and resuming the suspended function
//! afterwards.

use core::marker::PhantomData;
use core::{ptr, mem};

use stack;
use context::Context;

#[derive(Debug, Clone, Copy)]
pub enum State {
  /// Generator can be resumed. This is the initial state.
  Runnable,
  /// Generator cannot be resumed. This is the state of the generator after
  /// the generator function has returned or panicked.
  Unavailable
}

/// Generator wraps a function and allows suspending its execution more than once, returning
/// a value each time.
///
/// The first time `resume(input0)` is called, the function is called as `f(yielder, input0)`.
/// It runs until it suspends its execution through `yielder.suspend(output0)`, after which
/// `resume(input0)` returns `output0`. The function can be resumed again using `resume(input1)`,
/// after which `yielder.suspend(output0)` returns `input1`, and so on. Once the function returns,
/// the `resume()` call will return `None`, and it will return `None` every time it is called
/// after that.
///
/// If the generator function panics, the panic is propagated through the `resume()` call as usual.
///
/// After the generator function returns or panics, it is safe to reclaim the generator stack
/// using `unwrap()`.
///
/// `state()` can be used to determine whether the generator function has returned;
/// the state is `State::Runnable` after creation and suspension, and `State::Unavailable`
/// once the generator function returns or panics.
///
/// When the input type is `()`, a generator implements the Iterator trait.
///
/// # Example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut add_one = Generator::new(stack, move |yielder, mut input| {
///   loop {
///     if input == 0 { break }
///     input = yielder.suspend(input + 1)
///   }
/// });
/// println!("{:?}", add_one.resume(2)); // prints Some(3)
/// println!("{:?}", add_one.resume(3)); // prints Some(4)
/// println!("{:?}", add_one.resume(0)); // prints None
/// ```
///
/// # Iterator example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut nat = Generator::new(stack, move |yielder, ()| {
///   for i in 1.. { yielder.suspend(i) }
/// });
/// println!("{:?}", nat.next()); // prints Some(0)
/// println!("{:?}", nat.next()); // prints Some(1)
/// println!("{:?}", nat.next()); // prints Some(2)
/// ```
#[derive(Debug)]
pub struct Generator<Input: Send, Output: Send, Stack: stack::Stack> {
  state:   State,
  context: Context<Stack>,
  phantom: (PhantomData<*const Input>, PhantomData<*const Output>)
}

impl<Input, Output, Stack> Generator<Input, Output, Stack>
    where Input: Send, Output: Send, Stack: stack::Stack {
  /// Creates a new generator.
  pub fn new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
      where Stack: stack::GuardedStack,
            F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) + Send {
    unsafe { Generator::unsafe_new(stack, f) }
  }

  /// Same as `new`, but does not require `stack` to have a guard page.
  ///
  /// This function is unsafe because the generator function can easily violate
  /// memory safety by overflowing the stack. It is useful in environments where
  /// guarded stacks do not exist, e.g. in absence of an MMU.
  pub unsafe fn unsafe_new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
      where F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) + Send {
    unsafe extern "C" fn generator_wrapper<Input, Output, Stack, F>(env: usize) -> !
        where Input: Send, Output: Send, Stack: stack::Stack,
              F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) {
      // Retrieve our environment from the callee and return control to it.
      let (mut yielder, f) = ptr::read(env as *mut (Yielder<Input, Output, Stack>, F));
      let data = Context::swap(yielder.context, yielder.context, 0);
      // See the second half of Yielder::suspend_bare.
      let (new_context, input) = ptr::read(data as *mut (*mut Context<Stack>, Input));
      yielder.context = new_context as *mut Context<Stack>;
      // Run the body of the generator.
      f(&mut yielder, input);
      // Past this point, the generator has dropped everything it has held.
      loop { yielder.suspend_bare(None); }
    }

    let mut generator = Generator {
      state:   State::Runnable,
      context: Context::new(stack, generator_wrapper::<Input, Output, Stack, F>),
      phantom: (PhantomData, PhantomData)
    };

    // Transfer environment to the callee.
    let mut env = (Yielder::new(&mut generator.context), f);
    Context::swap(&mut generator.context, &generator.context,
                  &mut env as *mut (Yielder<Input, Output, Stack>, F) as usize);
    mem::forget(env);

    generator
  }

  /// Resumes the generator and return the next value it yields.
  /// If the generator function has returned, returns `None`.
  #[inline]
  pub fn resume(&mut self, input: Input) -> Option<Output> {
    match self.state {
      State::Runnable => {
        // Set the state to Unavailable. Since we have exclusive access to the generator,
        // the only case where this matters is the generator function panics, after which
        // it must not be invocable again.
        self.state = State::Unavailable;

        // Switch to the generator function, and retrieve the yielded value.
        let val = unsafe {
          let mut data_in = (&mut self.context as *mut Context<Stack>, input);
          let data_out =
            ptr::read(Context::swap(&mut self.context, &self.context,
                                    &mut data_in as *mut (*mut Context<Stack>, Input)  as usize)
                      as *mut Option<Output>);
          mem::forget(data_in);
          data_out
        };

        // Unless the generator function has returned, it can be switched to again, so
        // set the state to Runnable.
        if val.is_some() { self.state = State::Runnable }

        val
      }
      State::Unavailable => None
    }
  }

  /// Returns the state of the generator.
  #[inline]
  pub fn state(&self) -> State { self.state }

  /// Extracts the stack from a generator when the generator function has returned.
  /// If the generator function has not returned
  /// (i.e. `self.state() == State::Runnable`), panics.
  pub fn unwrap(self) -> Stack {
    match self.state {
      State::Runnable    => panic!("Argh! Bastard! Don't touch that!"),
      State::Unavailable => unsafe { self.context.unwrap() }
    }
  }
}

/// Yielder is an interface provided to every generator through which it
/// returns a value.
#[derive(Debug)]
pub struct Yielder<Input: Send, Output: Send, Stack: stack::Stack> {
  context: *mut Context<Stack>,
  phantom: (PhantomData<*const Input>, PhantomData<*const Output>)
}

impl<Input, Output, Stack> Yielder<Input, Output, Stack>
    where Input: Send, Output: Send, Stack: stack::Stack {
  fn new(context: *mut Context<Stack>) -> Yielder<Input, Output, Stack> {
    Yielder {
      context: context,
      phantom: (PhantomData, PhantomData)
    }
  }

  #[inline(always)]
  fn suspend_bare(&mut self, mut val: Option<Output>) -> Input {
    unsafe {
      let data = Context::swap(self.context, self.context,
                               &mut val as *mut Option<Output> as usize);
      let (new_context, input) = ptr::read(data as *mut (*mut Context<Stack>, Input));
      // The generator can be moved (and with it, the context).
      // This changes the address of the context.
      // Thus, we update it after each swap.
      self.context = new_context;
      // However, between this point and the next time we enter suspend_bare
      // the generator cannot be moved, as a &mut Generator is necessary
      // to resume the generator function.
      input
    }
  }

  /// Suspends the generator and returns `Some(item)` from the `resume()`
  /// invocation that resumed the generator.
  #[inline(always)]
  pub fn suspend(&mut self, item: Output) -> Input {
    self.suspend_bare(Some(item))
  }
}

impl<Output, Stack> Iterator for Generator<(), Output, Stack>
    where Output: Send, Stack: stack::Stack {
  type Item = Output;

  fn next(&mut self) -> Option<Self::Item> { self.resume(()) }
}
<commit_msg>don't drop yielded values in the generator<commit_after>// This file is part of libfringe, a low-level green threading library.
// Copyright (c) whitequark <whitequark@whitequark.org>
// 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.

//! Generators.
//!
//! Generators allow repeatedly suspending the execution of a function,
//! returning a value to the caller, and resuming the suspended function
//! afterwards.

use core::marker::PhantomData;
use core::{ptr, mem};

use stack;
use context::Context;

#[derive(Debug, Clone, Copy)]
pub enum State {
  /// Generator can be resumed. This is the initial state.
  Runnable,
  /// Generator cannot be resumed. This is the state of the generator after
  /// the generator function has returned or panicked.
  Unavailable
}

/// Generator wraps a function and allows suspending its execution more than once, returning
/// a value each time.
///
/// The first time `resume(input0)` is called, the function is called as `f(yielder, input0)`.
/// It runs until it suspends its execution through `yielder.suspend(output0)`, after which
/// `resume(input0)` returns `output0`. The function can be resumed again using `resume(input1)`,
/// after which `yielder.suspend(output0)` returns `input1`, and so on. Once the function returns,
/// the `resume()` call will return `None`, and it will return `None` every time it is called
/// after that.
///
/// If the generator function panics, the panic is propagated through the `resume()` call as usual.
///
/// After the generator function returns or panics, it is safe to reclaim the generator stack
/// using `unwrap()`.
///
/// `state()` can be used to determine whether the generator function has returned;
/// the state is `State::Runnable` after creation and suspension, and `State::Unavailable`
/// once the generator function returns or panics.
///
/// When the input type is `()`, a generator implements the Iterator trait.
///
/// # Example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut add_one = Generator::new(stack, move |yielder, mut input| {
///   loop {
///     if input == 0 { break }
///     input = yielder.suspend(input + 1)
///   }
/// });
/// println!("{:?}", add_one.resume(2)); // prints Some(3)
/// println!("{:?}", add_one.resume(3)); // prints Some(4)
/// println!("{:?}", add_one.resume(0)); // prints None
/// ```
///
/// # Iterator example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut nat = Generator::new(stack, move |yielder, ()| {
///   for i in 1.. { yielder.suspend(i) }
/// });
/// println!("{:?}", nat.next()); // prints Some(0)
/// println!("{:?}", nat.next()); // prints Some(1)
/// println!("{:?}", nat.next()); // prints Some(2)
/// ```
#[derive(Debug)]
pub struct Generator<Input: Send, Output: Send, Stack: stack::Stack> {
  state:   State,
  context: Context<Stack>,
  phantom: (PhantomData<*const Input>, PhantomData<*const Output>)
}

impl<Input, Output, Stack> Generator<Input, Output, Stack>
    where Input: Send, Output: Send, Stack: stack::Stack {
  /// Creates a new generator.
  pub fn new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
      where Stack: stack::GuardedStack,
            F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) + Send {
    unsafe { Generator::unsafe_new(stack, f) }
  }

  /// Same as `new`, but does not require `stack` to have a guard page.
  ///
  /// This function is unsafe because the generator function can easily violate
  /// memory safety by overflowing the stack. It is useful in environments where
  /// guarded stacks do not exist, e.g. in absence of an MMU.
  pub unsafe fn unsafe_new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
      where F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) + Send {
    unsafe extern "C" fn generator_wrapper<Input, Output, Stack, F>(env: usize) -> !
        where Input: Send, Output: Send, Stack: stack::Stack,
              F: FnOnce(&mut Yielder<Input, Output, Stack>, Input) {
      // Retrieve our environment from the callee and return control to it.
      let (mut yielder, f) = ptr::read(env as *mut (Yielder<Input, Output, Stack>, F));
      let data = Context::swap(yielder.context, yielder.context, 0);
      // See the second half of Yielder::suspend_bare.
      let (new_context, input) = ptr::read(data as *mut (*mut Context<Stack>, Input));
      yielder.context = new_context as *mut Context<Stack>;
      // Run the body of the generator.
      f(&mut yielder, input);
      // Past this point, the generator has dropped everything it has held.
      loop { yielder.suspend_bare(None); }
    }

    let mut generator = Generator {
      state:   State::Runnable,
      context: Context::new(stack, generator_wrapper::<Input, Output, Stack, F>),
      phantom: (PhantomData, PhantomData)
    };

    // Transfer environment to the callee.
    let mut env = (Yielder::new(&mut generator.context), f);
    Context::swap(&mut generator.context, &generator.context,
                  &mut env as *mut (Yielder<Input, Output, Stack>, F) as usize);
    mem::forget(env);

    generator
  }

  /// Resumes the generator and return the next value it yields.
  /// If the generator function has returned, returns `None`.
  #[inline]
  pub fn resume(&mut self, input: Input) -> Option<Output> {
    match self.state {
      State::Runnable => {
        // Set the state to Unavailable. Since we have exclusive access to the generator,
        // the only case where this matters is the generator function panics, after which
        // it must not be invocable again.
        self.state = State::Unavailable;

        // Switch to the generator function, and retrieve the yielded value.
        let val = unsafe {
          let mut data_in = (&mut self.context as *mut Context<Stack>, input);
          let data_out =
            ptr::read(Context::swap(&mut self.context, &self.context,
                                    &mut data_in as *mut (*mut Context<Stack>, Input)  as usize)
                      as *mut Option<Output>);
          mem::forget(data_in);
          data_out
        };

        // Unless the generator function has returned, it can be switched to again, so
        // set the state to Runnable.
        if val.is_some() { self.state = State::Runnable }

        val
      }
      State::Unavailable => None
    }
  }

  /// Returns the state of the generator.
  #[inline]
  pub fn state(&self) -> State { self.state }

  /// Extracts the stack from a generator when the generator function has returned.
  /// If the generator function has not returned
  /// (i.e. `self.state() == State::Runnable`), panics.
  pub fn unwrap(self) -> Stack {
    match self.state {
      State::Runnable    => panic!("Argh! Bastard! Don't touch that!"),
      State::Unavailable => unsafe { self.context.unwrap() }
    }
  }
}

/// Yielder is an interface provided to every generator through which it
/// returns a value.
#[derive(Debug)]
pub struct Yielder<Input: Send, Output: Send, Stack: stack::Stack> {
  context: *mut Context<Stack>,
  phantom: (PhantomData<*const Input>, PhantomData<*const Output>)
}

impl<Input, Output, Stack> Yielder<Input, Output, Stack>
    where Input: Send, Output: Send, Stack: stack::Stack {
  fn new(context: *mut Context<Stack>) -> Yielder<Input, Output, Stack> {
    Yielder {
      context: context,
      phantom: (PhantomData, PhantomData)
    }
  }

  #[inline(always)]
  fn suspend_bare(&mut self, mut val: Option<Output>) -> Input {
    unsafe {
      let data = Context::swap(self.context, self.context,
                               &mut val as *mut Option<Output> as usize);
      mem::forget(val);
      let (new_context, input) = ptr::read(data as *mut (*mut Context<Stack>, Input));
      // The generator can be moved (and with it, the context).
      // This changes the address of the context.
      // Thus, we update it after each swap.
      self.context = new_context;
      // However, between this point and the next time we enter suspend_bare
      // the generator cannot be moved, as a &mut Generator is necessary
      // to resume the generator function.
      input
    }
  }

  /// Suspends the generator and returns `Some(item)` from the `resume()`
  /// invocation that resumed the generator.
  #[inline(always)]
  pub fn suspend(&mut self, item: Output) -> Input {
    self.suspend_bare(Some(item))
  }
}

impl<Output, Stack> Iterator for Generator<(), Output, Stack>
    where Output: Send, Stack: stack::Stack {
  type Item = Output;

  fn next(&mut self) -> Option<Self::Item> { self.resume(()) }
}
<|endoftext|>


<commit_before>use socket::Socket;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum TriggerSource {
    EXT_PE,
    EXT_NE,
    INT,
    GATED,
}

impl ::std::fmt::Display for TriggerSource {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &TriggerSource::EXT_PE => "EXT_PE",
            &TriggerSource::EXT_NE => "EXT_NE",
            &TriggerSource::INT => "INT",
            &TriggerSource::GATED => "GATED",
        };

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

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Form {
    SINE,
    SQUARE,
    TRIANGLE,
    SAWU,
    SAWD,
    PWM,
    ARBITRARY,
}

impl ::std::fmt::Display for Form {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &Form::SINE => "SINE",
            &Form::SQUARE => "SQUARE",
            &Form::TRIANGLE => "TRIANGLE",
            &Form::SAWU => "SAWU",
            &Form::SAWD => "SAWD",
            &Form::PWM => "PWM",
            &Form::ARBITRARY => "ARBITRARY",
        };

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

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Source {
    OUT1,
    OUT2,
}

impl ::std::fmt::Display for Source {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &Source::OUT1 => "SOUR1",
            &Source::OUT2 => "SOUR2",
        };

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

// @TODO impl std::slice::SliceIndex<generator::State> instead
impl ::std::convert::Into<usize> for Source {
    fn into(self) -> usize {
        match self {
            Source::OUT1 => 0,
            Source::OUT2 => 1,
        }
    }
}

struct State {
    started: bool,
    form: Form,
    amplitude: f32,
    frequency: u32,
    dcyc: u32,
}

impl State {
    pub fn new() -> Self {
        State {
            started: false,
            form: Form::SINE,
            amplitude: 1.0,
            frequency: 1000,
            dcyc: 50,
        }
    }
}

pub struct Generator {
    socket: Socket,
    states: [State; 2],
}

impl Generator {
    pub fn new(socket: Socket) -> Self {
        Generator {
            socket: socket,
            states: [
                State::new(),
                State::new(),
            ],
        }
    }

    /**
     * Enable fast analog outputs.
     */
    pub fn start(&mut self, source: Source) {
        self.set_state(&source, "ON");
        self.states[source as usize].started = true;
    }

    /**
     * Disable fast analog outputs.
     */
    pub fn stop(&mut self, source: Source) {
        self.set_state(&source, "OFF");
        self.states[source as usize].started = false;
    }

    fn set_state(&mut self, source: &Source, state: &str) {
        let output = match *source {
            Source::OUT1 => "OUTPUT1",
            Source::OUT2 => "OUTPUT2",
        };

        self.socket.send(format!("{}:STATE {}", output, state));
    }

    pub fn is_started(&self, source: Source) -> bool {
        self.states[source as usize].started
    }

    /**
     * Set frequency of fast analog outputs.
     */
    pub fn set_frequency(&mut self, source: Source, frequency: u32) {
        self.socket.send(format!("{}:FREQ:FIX {}", source, frequency));
        self.states[source as usize].frequency = frequency;
    }

    pub fn get_frequency(&self, source: Source) -> u32 {
        self.states[source as usize].frequency
    }

    /**
     * Set waveform of fast analog outputs.
     */
    pub fn set_form(&mut self, source: Source, form: Form) {
        self.socket.send(format!("{}:FUNC {}", source, form));
        self.states[source as usize].form = form;
    }

    pub fn get_form(&self, source: Source) -> Form {
        self.states[source as usize].form.clone()
    }

    /**
     * Set amplitude voltage of fast analog outputs.
     *
     * Amplitude + offset value must be less than maximum output range ± 1V
     */
    pub fn set_amplitude(&mut self, source: Source, amplitude: f32) {
        self.socket.send(format!("{}:VOLT {}", source, amplitude));
        self.states[source as usize].amplitude = amplitude;
    }

    pub fn get_amplitude(&self, source: Source) -> f32 {
        self.states[source as usize].amplitude
    }

    /**
     * Set offset voltage of fast analog outputs.
     *
     * Amplitude + offset value must be less than maximum output range ± 1V
     */
    pub fn set_offset(&mut self, source: Source, offset: f32) {
        self.socket.send(format!("{}:VOLT:OFFS {}", source, offset));
    }

    /**
     * Set phase of fast analog outputs.
     */
    pub fn set_phase(&mut self, source: Source, phase: i32) {
        self.socket.send(format!("{}:PHAS {}", source, phase));
    }

    /**
     * Set duty cycle of PWM waveform.
     */
    pub fn set_duty_cicle(&mut self, source: Source, dcyc: u32) {
        self.socket.send(format!("{}:DCYC {}", source, dcyc));
        self.states[source as usize].dcyc = dcyc;
    }

    pub fn get_duty_cicle(&self, source: Source) -> u32 {
        self.states[source as usize].dcyc
    }

    /**
     * Import data for arbitrary waveform generation.
     */
    pub fn arbitrary_waveform(&mut self, source: Source, data: Vec<f32>) {
        let mut data = data.iter()
            .fold(String::new(), |acc, e| {
                format!("{}{},", acc, e)
            });
        data.pop();

        self.socket.send(format!("{}:TRAC:DATA:DATA {}", source, data));
    }

    /**
     * Set trigger source for selected signal.
     */
    pub fn set_trigger_source(&mut self, source: Source, trigger: TriggerSource) {
        self.socket.send(format!("{}:TRIG:SOUR {}", source, trigger));
    }

    /**
     * Triggers selected source immediately.
     */
    pub fn trigger(&mut self, source: Source) {
        self.socket.send(format!("{}:TRIG:IMM", source));
    }

    /**
     * Triggers both sources immediately.
     */
    pub fn trigger_all(&mut self) {
        self.socket.send("TRIG:IMM");
    }

    /**
     * Reset generator to default settings.
     */
    pub fn reset(&mut self) {
        self.socket.send("GEN:RST");
    }
}

#[cfg(test)]
mod test {
    macro_rules! generator_assert {
        ($f:ident, $e:expr) => {
            let (rx, mut generator) = create_generator();

            generator.$f();
            assert_eq!($e, rx.recv().unwrap());
        }
    }

    #[test]
    fn test_start() {
        let (rx, mut generator) = create_generator();

        generator.start(::generator::Source::OUT2);
        assert_eq!("OUTPUT2:STATE ON\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_stop() {
        let (rx, mut generator) = create_generator();

        generator.stop(::generator::Source::OUT2);
        assert_eq!("OUTPUT2:STATE OFF\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_is_started() {
        let (_, mut generator) = create_generator();

        assert_eq!(generator.is_started(::generator::Source::OUT1), false);
        generator.start(::generator::Source::OUT1);
        assert_eq!(generator.is_started(::generator::Source::OUT1), true);
        generator.stop(::generator::Source::OUT1);
        assert_eq!(generator.is_started(::generator::Source::OUT1), false);
    }

    #[test]
    fn test_set_frequency() {
        let (rx, mut generator) = create_generator();

        generator.set_frequency(::generator::Source::OUT1, 500);
        assert_eq!("SOUR1:FREQ:FIX 500\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_frequency() {
        let (_, generator) = create_generator();

        assert_eq!(generator.get_frequency(::generator::Source::OUT1), 1000);
    }

    #[test]
    fn test_set_form() {
        let (rx, mut generator) = create_generator();

        generator.set_form(::generator::Source::OUT1, ::generator::Form::SINE);
        assert_eq!("SOUR1:FUNC SINE\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_amplitude() {
        let (rx, mut generator) = create_generator();

        generator.set_amplitude(::generator::Source::OUT1, -0.9);
        assert_eq!("SOUR1:VOLT -0.9\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_set_offset() {
        let (rx, mut generator) = create_generator();

        generator.set_offset(::generator::Source::OUT1, -1.0);
        assert_eq!("SOUR1:VOLT:OFFS -1\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_set_phase() {
        let (rx, mut generator) = create_generator();

        generator.set_phase(::generator::Source::OUT1, -360);
        assert_eq!("SOUR1:PHAS -360\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_set_duty_cicle() {
        let (rx, mut generator) = create_generator();

        generator.set_duty_cicle(::generator::Source::OUT1, 100);
        assert_eq!("SOUR1:DCYC 100\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_arbitrary_waveform() {
        let (rx, mut generator) = create_generator();

        generator.arbitrary_waveform(::generator::Source::OUT1, vec![1.0, 0.5, 0.2]);
        assert_eq!("SOUR1:TRAC:DATA:DATA 1,0.5,0.2\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_set_trigger_source() {
        let (rx, mut generator) = create_generator();

        generator.set_trigger_source(::generator::Source::OUT1, ::generator::TriggerSource::EXT_PE);
        assert_eq!("SOUR1:TRIG:SOUR EXT_PE\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_trigger() {
        let (rx, mut generator) = create_generator();

        generator.trigger(::generator::Source::OUT1);
        assert_eq!("SOUR1:TRIG:IMM\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_trigger_all() {
        generator_assert!(trigger_all, "TRIG:IMM\r\n");
    }

    #[test]
    fn test_reset() {
        generator_assert!(reset, "GEN:RST\r\n");
    }

    fn create_generator() -> (::std::sync::mpsc::Receiver<String>, ::generator::Generator) {
        let (addr, rx) = ::test::launch_server();
        let socket = ::socket::Socket::new(addr);

        (rx, ::generator::Generator::new(socket))
    }
}
<commit_msg>Fix typo in method name<commit_after>use socket::Socket;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum TriggerSource {
    EXT_PE,
    EXT_NE,
    INT,
    GATED,
}

impl ::std::fmt::Display for TriggerSource {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &TriggerSource::EXT_PE => "EXT_PE",
            &TriggerSource::EXT_NE => "EXT_NE",
            &TriggerSource::INT => "INT",
            &TriggerSource::GATED => "GATED",
        };

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

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Form {
    SINE,
    SQUARE,
    TRIANGLE,
    SAWU,
    SAWD,
    PWM,
    ARBITRARY,
}

impl ::std::fmt::Display for Form {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &Form::SINE => "SINE",
            &Form::SQUARE => "SQUARE",
            &Form::TRIANGLE => "TRIANGLE",
            &Form::SAWU => "SAWU",
            &Form::SAWD => "SAWD",
            &Form::PWM => "PWM",
            &Form::ARBITRARY => "ARBITRARY",
        };

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

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Source {
    OUT1,
    OUT2,
}

impl ::std::fmt::Display for Source {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &Source::OUT1 => "SOUR1",
            &Source::OUT2 => "SOUR2",
        };

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

// @TODO impl std::slice::SliceIndex<generator::State> instead
impl ::std::convert::Into<usize> for Source {
    fn into(self) -> usize {
        match self {
            Source::OUT1 => 0,
            Source::OUT2 => 1,
        }
    }
}

struct State {
    started: bool,
    form: Form,
    amplitude: f32,
    frequency: u32,
    dcyc: u32,
}

impl State {
    pub fn new() -> Self {
        State {
            started: false,
            form: Form::SINE,
            amplitude: 1.0,
            frequency: 1000,
            dcyc: 50,
        }
    }
}

pub struct Generator {
    socket: Socket,
    states: [State; 2],
}

impl Generator {
    pub fn new(socket: Socket) -> Self {
        Generator {
            socket: socket,
            states: [
                State::new(),
                State::new(),
            ],
        }
    }

    /**
     * Enable fast analog outputs.
     */
    pub fn start(&mut self, source: Source) {
        self.set_state(&source, "ON");
        self.states[source as usize].started = true;
    }

    /**
     * Disable fast analog outputs.
     */
    pub fn stop(&mut self, source: Source) {
        self.set_state(&source, "OFF");
        self.states[source as usize].started = false;
    }

    fn set_state(&mut self, source: &Source, state: &str) {
        let output = match *source {
            Source::OUT1 => "OUTPUT1",
            Source::OUT2 => "OUTPUT2",
        };

        self.socket.send(format!("{}:STATE {}", output, state));
    }

    pub fn is_started(&self, source: Source) -> bool {
        self.states[source as usize].started
    }

    /**
     * Set frequency of fast analog outputs.
     */
    pub fn set_frequency(&mut self, source: Source, frequency: u32) {
        self.socket.send(format!("{}:FREQ:FIX {}", source, frequency));
        self.states[source as usize].frequency = frequency;
    }

    pub fn get_frequency(&self, source: Source) -> u32 {
        self.states[source as usize].frequency
    }

    /**
     * Set waveform of fast analog outputs.
     */
    pub fn set_form(&mut self, source: Source, form: Form) {
        self.socket.send(format!("{}:FUNC {}", source, form));
        self.states[source as usize].form = form;
    }

    pub fn get_form(&self, source: Source) -> Form {
        self.states[source as usize].form.clone()
    }

    /**
     * Set amplitude voltage of fast analog outputs.
     *
     * Amplitude + offset value must be less than maximum output range ± 1V
     */
    pub fn set_amplitude(&mut self, source: Source, amplitude: f32) {
        self.socket.send(format!("{}:VOLT {}", source, amplitude));
        self.states[source as usize].amplitude = amplitude;
    }

    pub fn get_amplitude(&self, source: Source) -> f32 {
        self.states[source as usize].amplitude
    }

    /**
     * Set offset voltage of fast analog outputs.
     *
     * Amplitude + offset value must be less than maximum output range ± 1V
     */
    pub fn set_offset(&mut self, source: Source, offset: f32) {
        self.socket.send(format!("{}:VOLT:OFFS {}", source, offset));
    }

    /**
     * Set phase of fast analog outputs.
     */
    pub fn set_phase(&mut self, source: Source, phase: i32) {
        self.socket.send(format!("{}:PHAS {}", source, phase));
    }

    /**
     * Set duty cycle of PWM waveform.
     */
    pub fn set_duty_cycle(&mut self, source: Source, dcyc: u32) {
        self.socket.send(format!("{}:DCYC {}", source, dcyc));
        self.states[source as usize].dcyc = dcyc;
    }

    pub fn get_duty_cycle(&self, source: Source) -> u32 {
        self.states[source as usize].dcyc
    }

    /**
     * Import data for arbitrary waveform generation.
     */
    pub fn arbitrary_waveform(&mut self, source: Source, data: Vec<f32>) {
        let mut data = data.iter()
            .fold(String::new(), |acc, e| {
                format!("{}{},", acc, e)
            });
        data.pop();

        self.socket.send(format!("{}:TRAC:DATA:DATA {}", source, data));
    }

    /**
     * Set trigger source for selected signal.
     */
    pub fn set_trigger_source(&mut self, source: Source, trigger: TriggerSource) {
        self.socket.send(format!("{}:TRIG:SOUR {}", source, trigger));
    }

    /**
     * Triggers selected source immediately.
     */
    pub fn trigger(&mut self, source: Source) {
        self.socket.send(format!("{}:TRIG:IMM", source));
    }

    /**
     * Triggers both sources immediately.
     */
    pub fn trigger_all(&mut self) {
        self.socket.send("TRIG:IMM");
    }

    /**
     * Reset generator to default settings.
     */
    pub fn reset(&mut self) {
        self.socket.send("GEN:RST");
    }
}

#[cfg(test)]
mod test {
    macro_rules! generator_assert {
        ($f:ident, $e:expr) => {
            let (rx, mut generator) = create_generator();

            generator.$f();
            assert_eq!($e, rx.recv().unwrap());
        }
    }

    #[test]
    fn test_start() {
        let (rx, mut generator) = create_generator();

        generator.start(::generator::Source::OUT2);
        assert_eq!("OUTPUT2:STATE ON\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_stop() {
        let (rx, mut generator) = create_generator();

        generator.stop(::generator::Source::OUT2);
        assert_eq!("OUTPUT2:STATE OFF\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_is_started() {
        let (_, mut generator) = create_generator();

        assert_eq!(generator.is_started(::generator::Source::OUT1), false);
        generator.start(::generator::Source::OUT1);
        assert_eq!(generator.is_started(::generator::Source::OUT1), true);
        generator.stop(::generator::Source::OUT1);
        assert_eq!(generator.is_started(::generator::Source::OUT1), false);
    }

    #[test]
    fn test_set_frequency() {
        let (rx, mut generator) = create_generator();

        generator.set_frequency(::generator::Source::OUT1, 500);
        assert_eq!("SOUR1:FREQ:FIX 500\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_frequency() {
        let (_, generator) = create_generator();

        assert_eq!(generator.get_frequency(::generator::Source::OUT1), 1000);
    }

    #[test]
    fn test_set_form() {
        let (rx, mut generator) = create_generator();

        generator.set_form(::generator::Source::OUT1, ::generator::Form::SINE);
        assert_eq!("SOUR1:FUNC SINE\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_amplitude() {
        let (rx, mut generator) = create_generator();

        generator.set_amplitude(::generator::Source::OUT1, -0.9);
        assert_eq!("SOUR1:VOLT -0.9\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_set_offset() {
        let (rx, mut generator) = create_generator();

        generator.set_offset(::generator::Source::OUT1, -1.0);
        assert_eq!("SOUR1:VOLT:OFFS -1\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_set_phase() {
        let (rx, mut generator) = create_generator();

        generator.set_phase(::generator::Source::OUT1, -360);
        assert_eq!("SOUR1:PHAS -360\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_set_duty_cycle() {
        let (rx, mut generator) = create_generator();

        generator.set_duty_cycle(::generator::Source::OUT1, 100);
        assert_eq!("SOUR1:DCYC 100\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_arbitrary_waveform() {
        let (rx, mut generator) = create_generator();

        generator.arbitrary_waveform(::generator::Source::OUT1, vec![1.0, 0.5, 0.2]);
        assert_eq!("SOUR1:TRAC:DATA:DATA 1,0.5,0.2\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_set_trigger_source() {
        let (rx, mut generator) = create_generator();

        generator.set_trigger_source(::generator::Source::OUT1, ::generator::TriggerSource::EXT_PE);
        assert_eq!("SOUR1:TRIG:SOUR EXT_PE\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_trigger() {
        let (rx, mut generator) = create_generator();

        generator.trigger(::generator::Source::OUT1);
        assert_eq!("SOUR1:TRIG:IMM\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_trigger_all() {
        generator_assert!(trigger_all, "TRIG:IMM\r\n");
    }

    #[test]
    fn test_reset() {
        generator_assert!(reset, "GEN:RST\r\n");
    }

    fn create_generator() -> (::std::sync::mpsc::Receiver<String>, ::generator::Generator) {
        let (addr, rx) = ::test::launch_server();
        let socket = ::socket::Socket::new(addr);

        (rx, ::generator::Generator::new(socket))
    }
}
<|endoftext|>


<commit_before>// (C)opyleft 2013 Frank Denis

/*!
 * Bindings for the GeoIP library
 */
#![desc = "Bindings for the GeoIP library."]
#![license = "BSD"]
#![crate_id = "geoip#0.1"]
#![crate_type = "rlib"]

#![warn(non_camel_case_types,
        non_uppercase_statics,
        unnecessary_qualification,
        managed_heap_memory)]

use std::c_str::CString;
use std::fmt;
use std::io::net::ip::{IpAddr,Ipv4Addr,Ipv6Addr};
use std::libc::{c_void, c_char, c_int, c_ulong};

type GeoIP_ = *c_void;
type In6Addr = [u8, ..16];

struct GeoIPLookup {
    netmask: c_int
}

impl GeoIPLookup {
    fn new() -> GeoIPLookup {
        GeoIPLookup {
            netmask: 0
        }
    }
}

#[link(name = "GeoIP")]
extern {
    fn GeoIP_open(dbtype: *c_char, flags: c_int) -> GeoIP_;
    fn GeoIP_delete(db: GeoIP_);
    fn GeoIP_name_by_ipnum_gl(db: GeoIP_, ipnum: c_ulong, gl: &GeoIPLookup) -> *c_char;
    fn GeoIP_name_by_ipnum_v6_gl(db: GeoIP_, ipnum: In6Addr, gl: &GeoIPLookup) -> *c_char;
}

pub enum Options {
    Standard = 0,
    MemoryCache = 1,
    CheckCache = 2,
    IndexCache = 4,
    MmapCache = 8
}

pub enum DBType {
    CountryEdition = 1,
    RegionEditionRev0 = 7,
    CityEditionRev0 = 6,
    ORGEdition = 5,
    ISPEdition = 4,
    CityEditionRev1 = 2,
    RegionEditionRev1 = 3,
    ProxyEdition = 8,
    ASNUMEdition = 9,
    NetSpeedEdition = 10,
    DomainEdition = 11,
    CountryEditionV6 = 12,
    LocationAEdition = 13,
    AccuracyRadiusEdition = 14,
    LargeCountryEdition = 17,
    LargeCountryEditionV6 = 18,
    ASNumEditionV6 = 21,
    ISPEditionV6 = 22,
    ORGEditionV6 = 23,
    DomainEditionV6 = 24,
    LoctionAEditionV6 = 25,
    RegistrarEdition = 26,
    RegistrarEditionV6 = 27,
    UserTypeEdition = 28,
    UserTypeEditionV6 = 29,
    CityEditionRev1V6 = 30,
    CityEditionRev0V6 = 31,
    NetSpeedEditionRev1 = 32,
    NetSpeedEditionRev1V6 = 33,
    CountryConfEdition = 34,
    CityConfEdition = 35,
    RegionConfEdition = 36,
    PostalConfEdition = 37,
    AccuracyRadiusEditionV6 = 38
}

pub struct GeoIP {
    db: GeoIP_
}

pub struct ASInfo {
    pub asn: uint,
    pub name: ~str,
    pub netmask: uint
}

impl fmt::Show for ASInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f.buf, "{}\t{}", self.asn, self.name)
    }
}

impl GeoIP {
    pub fn open(path: &Path, options: Options) -> Result<GeoIP, ~str> {
        let file = match path.as_str() {
            None => return Err(format!("Invalid path {}", path.display())),
            Some(file) => file
        };
        let db = unsafe {
            GeoIP_open(file.to_c_str().unwrap(), options as c_int)
        };
        match db.is_null() {
            true => Err(format!("Can't open {}", file)),
            false => Ok(GeoIP {
                    db: db
                })
        }
    }

    pub fn as_info_by_ip(&self, ip: IpAddr) -> Option<ASInfo> {
        let gl = GeoIPLookup::new();
        let cres = match ip {
            Ipv4Addr(a, b, c, d) => {
                let ipnum: c_ulong =
                    (a as c_ulong << 24) | (b as c_ulong << 16) |
                    (c as c_ulong << 8)  | (d as c_ulong);
                unsafe {
                    GeoIP_name_by_ipnum_gl(self.db, ipnum, &gl)
                }
            },
            Ipv6Addr(a, b, c, d, e, f, g, h) => {
                let in6_addr: In6Addr = [(a >> 8) as u8, a as u8,
                                         (b >> 8) as u8, b as u8,
                                         (c >> 8) as u8, c as u8,
                                         (d >> 8) as u8, d as u8,
                                         (e >> 8) as u8, e as u8,
                                         (f >> 8) as u8, f as u8,
                                         (g >> 8) as u8, g as u8,
                                         (h >> 8) as u8, h as u8];
                unsafe {
                    GeoIP_name_by_ipnum_v6_gl(self.db, in6_addr, &gl)
                }
            }
        };
        if cres.is_null() {
            return None;
        }
        let description_cstr = unsafe { CString::new(cres, true) };
        let description = match description_cstr.as_str() {
            None => return None,
            Some(description) => description
        };
        let mut di = description.splitn(' ', 1);
        let asn = match di.next() {
            None => return None,
            Some(asn) => {
                if ! asn.starts_with("AS") {
                    return None
                } else {
                    from_str::<uint>(asn.slice_from(2)).unwrap()
                }
            }
        };
        let name = di.next().unwrap_or("(none)");
        let asinfo = ASInfo {
            asn: asn,
            name: name.to_owned(),
            netmask: gl.netmask as uint
        };
        Some(asinfo)
    }
}

impl Drop for GeoIP {
    fn drop(&mut self) {
        unsafe {
            GeoIP_delete(self.db);
        }
    }
}

#[test]
fn geoip_test_basic() {
    let geoip = match GeoIP::open(&from_str("/opt/geoip/GeoIPASNum.dat").unwrap(), MemoryCache) {
        Err(err) => fail!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("91.203.184.192").unwrap();
    let res = geoip.as_info_by_ip(ip).unwrap();
    assert!(res.asn == 41064);
    assert!(res.name.contains("Telefun"));
    assert!(res.netmask == 22);
}
<commit_msg>libc is now in its own crate<commit_after>// (C)opyleft 2013 Frank Denis

/*!
 * Bindings for the GeoIP library
 */
#![desc = "Bindings for the GeoIP library."]
#![license = "BSD"]
#![crate_id = "geoip#0.1"]
#![crate_type = "rlib"]

#![warn(non_camel_case_types,
        non_uppercase_statics,
        unnecessary_qualification,
        managed_heap_memory)]

extern crate libc;

use libc::{c_void, c_char, c_int, c_ulong};
use std::c_str::CString;
use std::fmt;
use std::io::net::ip::{IpAddr,Ipv4Addr,Ipv6Addr};

type GeoIP_ = *c_void;
type In6Addr = [u8, ..16];

struct GeoIPLookup {
    netmask: c_int
}

impl GeoIPLookup {
    fn new() -> GeoIPLookup {
        GeoIPLookup {
            netmask: 0
        }
    }
}

#[link(name = "GeoIP")]
extern {
    fn GeoIP_open(dbtype: *c_char, flags: c_int) -> GeoIP_;
    fn GeoIP_delete(db: GeoIP_);
    fn GeoIP_name_by_ipnum_gl(db: GeoIP_, ipnum: c_ulong, gl: &GeoIPLookup) -> *c_char;
    fn GeoIP_name_by_ipnum_v6_gl(db: GeoIP_, ipnum: In6Addr, gl: &GeoIPLookup) -> *c_char;
}

pub enum Options {
    Standard = 0,
    MemoryCache = 1,
    CheckCache = 2,
    IndexCache = 4,
    MmapCache = 8
}

pub enum DBType {
    CountryEdition = 1,
    RegionEditionRev0 = 7,
    CityEditionRev0 = 6,
    ORGEdition = 5,
    ISPEdition = 4,
    CityEditionRev1 = 2,
    RegionEditionRev1 = 3,
    ProxyEdition = 8,
    ASNUMEdition = 9,
    NetSpeedEdition = 10,
    DomainEdition = 11,
    CountryEditionV6 = 12,
    LocationAEdition = 13,
    AccuracyRadiusEdition = 14,
    LargeCountryEdition = 17,
    LargeCountryEditionV6 = 18,
    ASNumEditionV6 = 21,
    ISPEditionV6 = 22,
    ORGEditionV6 = 23,
    DomainEditionV6 = 24,
    LoctionAEditionV6 = 25,
    RegistrarEdition = 26,
    RegistrarEditionV6 = 27,
    UserTypeEdition = 28,
    UserTypeEditionV6 = 29,
    CityEditionRev1V6 = 30,
    CityEditionRev0V6 = 31,
    NetSpeedEditionRev1 = 32,
    NetSpeedEditionRev1V6 = 33,
    CountryConfEdition = 34,
    CityConfEdition = 35,
    RegionConfEdition = 36,
    PostalConfEdition = 37,
    AccuracyRadiusEditionV6 = 38
}

pub struct GeoIP {
    db: GeoIP_
}

pub struct ASInfo {
    pub asn: uint,
    pub name: ~str,
    pub netmask: uint
}

impl fmt::Show for ASInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f.buf, "{}\t{}", self.asn, self.name)
    }
}

impl GeoIP {
    pub fn open(path: &Path, options: Options) -> Result<GeoIP, ~str> {
        let file = match path.as_str() {
            None => return Err(format!("Invalid path {}", path.display())),
            Some(file) => file
        };
        let db = unsafe {
            GeoIP_open(file.to_c_str().unwrap(), options as c_int)
        };
        match db.is_null() {
            true => Err(format!("Can't open {}", file)),
            false => Ok(GeoIP {
                    db: db
                })
        }
    }

    pub fn as_info_by_ip(&self, ip: IpAddr) -> Option<ASInfo> {
        let gl = GeoIPLookup::new();
        let cres = match ip {
            Ipv4Addr(a, b, c, d) => {
                let ipnum: c_ulong =
                    (a as c_ulong << 24) | (b as c_ulong << 16) |
                    (c as c_ulong << 8)  | (d as c_ulong);
                unsafe {
                    GeoIP_name_by_ipnum_gl(self.db, ipnum, &gl)
                }
            },
            Ipv6Addr(a, b, c, d, e, f, g, h) => {
                let in6_addr: In6Addr = [(a >> 8) as u8, a as u8,
                                         (b >> 8) as u8, b as u8,
                                         (c >> 8) as u8, c as u8,
                                         (d >> 8) as u8, d as u8,
                                         (e >> 8) as u8, e as u8,
                                         (f >> 8) as u8, f as u8,
                                         (g >> 8) as u8, g as u8,
                                         (h >> 8) as u8, h as u8];
                unsafe {
                    GeoIP_name_by_ipnum_v6_gl(self.db, in6_addr, &gl)
                }
            }
        };
        if cres.is_null() {
            return None;
        }
        let description_cstr = unsafe { CString::new(cres, true) };
        let description = match description_cstr.as_str() {
            None => return None,
            Some(description) => description
        };
        let mut di = description.splitn(' ', 1);
        let asn = match di.next() {
            None => return None,
            Some(asn) => {
                if ! asn.starts_with("AS") {
                    return None
                } else {
                    from_str::<uint>(asn.slice_from(2)).unwrap()
                }
            }
        };
        let name = di.next().unwrap_or("(none)");
        let asinfo = ASInfo {
            asn: asn,
            name: name.to_owned(),
            netmask: gl.netmask as uint
        };
        Some(asinfo)
    }
}

impl Drop for GeoIP {
    fn drop(&mut self) {
        unsafe {
            GeoIP_delete(self.db);
        }
    }
}

#[test]
fn geoip_test_basic() {
    let geoip = match GeoIP::open(&from_str("/opt/geoip/GeoIPASNum.dat").unwrap(), MemoryCache) {
        Err(err) => fail!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("91.203.184.192").unwrap();
    let res = geoip.as_info_by_ip(ip).unwrap();
    assert!(res.asn == 41064);
    assert!(res.name.contains("Telefun"));
    assert!(res.netmask == 22);
}
<|endoftext|>


<commit_before>//! SSS provides Shamir's secret sharing with raw data.

use rand::{OsRng, Rng};
use merkle_sigs::sign_data_vec;

use errors::*;
use sss::{Share, HASH_ALGO};
use sss::format::format_share_for_signing;
use share::validation::{validate_share_count, validate_signed_shares};
use lagrange::interpolate_at;

use super::encode::encode_secret_byte;

/// SSS provides Shamir's secret sharing with raw data.
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
pub(crate) struct SSS;

impl SSS {
    /// Performs threshold k-out-of-n Shamir's secret sharing.
    pub fn split_secret(
        &self,
        threshold: u8,
        shares_count: u8,
        secret: &[u8],
        sign_shares: bool,
    ) -> Result<Vec<Share>> {
        let (threshold, shares_count) = validate_share_count(threshold, shares_count)?;
        let shares = Self::secret_share(secret, threshold, shares_count)?;

        let signatures = if sign_shares {
            let shares_to_sign = shares
                .iter()
                .enumerate()
                .map(|(i, x)| format_share_for_signing(threshold, (i + 1) as u8, x))
                .collect::<Vec<_>>();

            let sign = sign_data_vec(&shares_to_sign, HASH_ALGO)
                .unwrap()
                .into_iter()
                .map(Some)
                .collect::<Vec<_>>();

            Some(sign)
        } else {
            None
        };

        let sig_pairs = signatures
            .unwrap_or_else(|| vec![None; shares_count as usize])
            .into_iter()
            .map(|sig_pair| sig_pair.map(From::from));

        let shares_and_sigs = shares.into_iter().enumerate().zip(sig_pairs);

        let result = shares_and_sigs.map(|((index, data), signature_pair)| {
            // This is actually safe since we alwaays generate less than 256 shares.
            let id = (index + 1) as u8;

            Share {
                id,
                threshold,
                data,
                signature_pair,
            }
        });

        Ok(result.collect())
    }

    fn secret_share(src: &[u8], threshold: u8, shares_count: u8) -> Result<Vec<Vec<u8>>> {
        let mut result = Vec::with_capacity(shares_count as usize);
        for _ in 0..(shares_count as usize) {
            result.push(vec![0u8; src.len()]);
        }
        let mut col_in = vec![0u8, threshold];
        let mut col_out = Vec::with_capacity(shares_count as usize);
        let mut osrng = OsRng::new()?;
        for (c, &s) in src.iter().enumerate() {
            col_in[0] = s;
            osrng.fill_bytes(&mut col_in[1..]);
            col_out.clear();
            encode_secret_byte(&*col_in, shares_count, &mut col_out)?;
            for (&y, share) in col_out.iter().zip(result.iter_mut()) {
                share[c] = y;
            }
        }
        Ok(result)
    }

    /// Recovers the secret from a k-out-of-n Shamir's secret sharing.
    ///
    /// At least `k` distinct shares need to be provided to recover the share.
    pub fn recover_secret(shares: Vec<Share>, verify_signatures: bool) -> Result<Vec<u8>> {
        let (threshold, shares) = validate_signed_shares(shares, verify_signatures)?;

        let slen = shares[0].data.len();
        let mut col_in = Vec::with_capacity(threshold as usize);
        let mut secret = Vec::with_capacity(slen);
        for byteindex in 0..slen {
            col_in.clear();
            for s in shares.iter().take(threshold as usize) {
                col_in.push((s.id, s.data[byteindex]));
            }
            secret.push(interpolate_at(&*col_in));
        }

        Ok(secret)
    }
}
<commit_msg>Fix bug where threshold did not set deg of secret polynomial<commit_after>//! SSS provides Shamir's secret sharing with raw data.

use rand::{OsRng, Rng};
use merkle_sigs::sign_data_vec;

use errors::*;
use sss::{Share, HASH_ALGO};
use sss::format::format_share_for_signing;
use share::validation::{validate_share_count, validate_signed_shares};
use lagrange::interpolate_at;

use super::encode::encode_secret_byte;

/// SSS provides Shamir's secret sharing with raw data.
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
pub(crate) struct SSS;

impl SSS {
    /// Performs threshold k-out-of-n Shamir's secret sharing.
    pub fn split_secret(
        &self,
        threshold: u8,
        shares_count: u8,
        secret: &[u8],
        sign_shares: bool,
    ) -> Result<Vec<Share>> {
        let (threshold, shares_count) = validate_share_count(threshold, shares_count)?;
        let shares = Self::secret_share(secret, threshold, shares_count)?;

        let signatures = if sign_shares {
            let shares_to_sign = shares
                .iter()
                .enumerate()
                .map(|(i, x)| format_share_for_signing(threshold, (i + 1) as u8, x))
                .collect::<Vec<_>>();

            let sign = sign_data_vec(&shares_to_sign, HASH_ALGO)
                .unwrap()
                .into_iter()
                .map(Some)
                .collect::<Vec<_>>();

            Some(sign)
        } else {
            None
        };

        let sig_pairs = signatures
            .unwrap_or_else(|| vec![None; shares_count as usize])
            .into_iter()
            .map(|sig_pair| sig_pair.map(From::from));

        let shares_and_sigs = shares.into_iter().enumerate().zip(sig_pairs);

        let result = shares_and_sigs.map(|((index, data), signature_pair)| {
            // This is actually safe since we alwaays generate less than 256 shares.
            let id = (index + 1) as u8;

            Share {
                id,
                threshold,
                data,
                signature_pair,
            }
        });

        Ok(result.collect())
    }

    fn secret_share(src: &[u8], threshold: u8, shares_count: u8) -> Result<Vec<Vec<u8>>> {
        let mut result = Vec::with_capacity(shares_count as usize);
        for _ in 0..(shares_count as usize) {
            result.push(vec![0u8; src.len()]);
        }
        let mut col_in = vec![0u8; threshold as usize];
        let mut col_out = Vec::with_capacity(shares_count as usize);
        let mut osrng = OsRng::new()?;
        for (c, &s) in src.iter().enumerate() {
            col_in[0] = s;
            osrng.fill_bytes(&mut col_in[1..]);
            col_out.clear();
            encode_secret_byte(&*col_in, shares_count, &mut col_out)?;
            for (&y, share) in col_out.iter().zip(result.iter_mut()) {
                share[c] = y;
            }
        }
        Ok(result)
    }

    /// Recovers the secret from a k-out-of-n Shamir's secret sharing.
    ///
    /// At least `k` distinct shares need to be provided to recover the share.
    pub fn recover_secret(shares: Vec<Share>, verify_signatures: bool) -> Result<Vec<u8>> {
        let (threshold, shares) = validate_signed_shares(shares, verify_signatures)?;

        let slen = shares[0].data.len();
        let mut col_in = Vec::with_capacity(threshold as usize);
        let mut secret = Vec::with_capacity(slen);
        for byteindex in 0..slen {
            col_in.clear();
            for s in shares.iter().take(threshold as usize) {
                col_in.push((s.id, s.data[byteindex]));
            }
            secret.push(interpolate_at(&*col_in));
        }

        Ok(secret)
    }
}
<|endoftext|>