<commit_before>use std::io::Read;
use std::io::Write;
use std::path::Path;
use std::fs::File;

use tomson::Toml;
use handlebars::Handlebars;
use rustc_serialize::json::Json;

fn generate_code(mut toml: String) -> String {
    // turn the toml string into json for compatibility with handlebars
    let mut json = Toml::as_json(&mut toml).unwrap();

    for (id, field) in json.as_object_mut().unwrap().get_mut("field").unwrap().as_object_mut().unwrap().iter_mut() {
        let field_obj = field.as_object_mut().unwrap();
        let field_type = field_obj.get("type").unwrap().as_string().unwrap().to_string();
        let field_name = id.clone();
        let component_name = if let Some(component_json) = field_obj.get("component") {
            component_json.as_string().unwrap().to_string()
        } else {
            id.clone()
        };

        let component_has_type = if let Some(component_has_type_json) = field_obj.get("component_has_type") {
            component_has_type_json.as_boolean().unwrap()
        } else {
            false
        };

        field_obj.insert("struct_field_name".to_string(), Json::String(field_name.clone()));
        field_obj.insert("getter_name".to_string(), Json::String(id.clone()));
        field_obj.insert("component_name".to_string(), Json::String(component_name.clone()));
        field_obj.insert(field_type.clone(), Json::Boolean(true));

        if component_has_type {
            field_obj.insert("component_has_type".to_string(), Json::Boolean(true));
        }

        match field_type.as_ref() {
            "sum_f64" => {
                field_obj.insert("struct_field_type".to_string(), Json::String("f64".to_string()));
                field_obj.insert("struct_field_cons".to_string(), Json::String("0.0".to_string()));
                field_obj.insert("getter_type".to_string(), Json::String("f64".to_string()));
                field_obj.insert("getter_expr".to_string(), Json::String(format!("self.{}", field_name)));
                field_obj.insert("remove_entity".to_string(), Json::String(format!("if let Some(v) = entity.{}() {{ cell.{} -= v; }}",
                                                                                component_name, field_name)));
                field_obj.insert("insert_entity".to_string(), Json::String(format!("if let Some(v) = entity.{}() {{ cell.{} += v; }}",
                                                                                component_name, field_name)));
            }
            "count_bool" => {
                field_obj.insert("struct_field_type".to_string(), Json::String("usize".to_string()));
                field_obj.insert("struct_field_cons".to_string(), Json::String("0".to_string()));
                field_obj.insert("getter_type".to_string(), Json::String("bool".to_string()));
                field_obj.insert("getter_expr".to_string(), Json::String(format!("self.{} != 0", field_name)));
                field_obj.insert("remove_entity".to_string(), Json::String(format!("if entity.contains_{}() {{ cell.{} -= 1; }}",
                                                                                component_name, field_name)));
                field_obj.insert("insert_entity".to_string(), Json::String(format!("if entity.contains_{}() {{ cell.{} += 1; }}",
                                                                                component_name, field_name)));
            }
            "any" => {
                field_obj.insert("struct_field_type".to_string(), Json::String("AnySet<EntityId>".to_string()));
                field_obj.insert("struct_field_cons".to_string(), Json::String("AnySet::new()".to_string()));
                field_obj.insert("getter_type".to_string(), Json::String("bool".to_string()));
                field_obj.insert("getter_expr".to_string(), Json::String(format!("!self.{}.is_empty()", field_name)));
                field_obj.insert("any_name".to_string(), Json::String(format!("any_{}", field_name)));
                field_obj.insert("any_type".to_string(), Json::String("Option<EntityId>".to_string()));
                field_obj.insert("any_expr".to_string(), Json::String(format!("self.{}.any()", field_name)));
                field_obj.insert("remove_entity".to_string(), Json::String(format!("if entity.contains_{}() {{ cell.{}.remove(entity.id()); }}",
                                                                                component_name, field_name)));
                field_obj.insert("insert_entity".to_string(), Json::String(format!("if entity.contains_{}() {{ cell.{}.insert(entity.id()); }}",
                                                                                component_name, field_name)));
                field_obj.insert("is_set_type".to_string(), Json::Boolean(true));
            }
            "set" => {
                field_obj.insert("struct_field_type".to_string(), Json::String("EntitySet".to_string()));
                field_obj.insert("struct_field_cons".to_string(), Json::String("EntitySet::new()".to_string()));
                field_obj.insert("getter_type".to_string(), Json::String("bool".to_string()));
                field_obj.insert("getter_expr".to_string(), Json::String(format!("!self.{}.is_empty()", field_name)));
                field_obj.insert("iter_name".to_string(), Json::String(format!("{}_iter", field_name)));
                field_obj.insert("iter_type".to_string(), Json::String("EntitySetIter".to_string()));
                field_obj.insert("iter_expr".to_string(), Json::String(format!("self.{}.iter()", field_name)));
                field_obj.insert("remove_entity".to_string(), Json::String(format!("if entity.contains_{}() {{ cell.{}.remove(entity.id()); }}",
                                                                                component_name, field_name)));
                field_obj.insert("insert_entity".to_string(), Json::String(format!("if entity.contains_{}() {{ cell.{}.insert(entity.id()); }}",
                                                                                component_name, field_name)));
                field_obj.insert("is_set_type".to_string(), Json::Boolean(true));
            }
            other => panic!("unknown field type {}", other),
        }
    }

    let mut handlebars = Handlebars::new();

    // prevent xml escaping
    handlebars.register_escape_fn(|input| input.to_string());
    handlebars.template_render(TEMPLATE, &json).unwrap()

}

fn read_file_to_string<P: AsRef<Path>>(path: P) -> String {
    let mut file = File::open(path).unwrap();
    let mut string = String::new();
    file.read_to_string(&mut string).unwrap();

    string
}

pub fn generate_spatial_hash<P: AsRef<Path>, Q: AsRef<Path>>(in_path: P, out_path: Q) {
    let string = read_file_to_string(in_path);

    let output_string = generate_code(string);

    let mut outfile = File::create(out_path).unwrap();
    write!(outfile, "{}", output_string).unwrap();
}

const TEMPLATE: &'static str = r#"// Automatically generated. Do not edit.
#![allow(unused_imports)]
use ecs::*;
use coord::Coord;
use grid::{Grid, StaticGrid, DefaultGrid, IterGrid, CoordIterGrid};
use util::AnySet;

pub type SpatialHashCellIter<'a> = <StaticGrid<SpatialHashCell> as IterGrid<'a>>::Iter;
pub type SpatialHashCoordIter = <StaticGrid<SpatialHashCell> as CoordIterGrid>::CoordIter;

pub struct SpatialHashCell {

{{#each field}}
    {{ struct_field_name }}: {{ struct_field_type }},
{{/each}}

    // set of entities currently in this cell
    entities: EntitySet,

    // action on which this cell was last updated
    last_updated: u64,
}

impl SpatialHashCell {
    pub fn new() -> Self {
        SpatialHashCell {
{{#each field}}
            {{ struct_field_name }}: {{ struct_field_cons }},
{{/each}}

            entities: EntitySet::new(),
            last_updated: 0,
        }
    }

    pub fn last_updated(&self) -> u64 {
        self.last_updated
    }

    pub fn entity_ids(&self) -> &EntitySet {
        &self.entities
    }

    pub fn entity_id_iter(&self) -> EntitySetIter {
        self.entities.iter()
    }

{{#each field}}
    pub fn {{ getter_name }}(&self) -> {{ getter_type }} {
        {{ getter_expr }}
    }
    {{#if iter_name}}
    pub fn {{ iter_name }}(&self) -> {{ iter_type }} {
        {{ iter_expr }}
    }
    {{/if}}
    {{#if any_name}}
    pub fn {{ any_name }}(&self) -> {{ any_type }} {
        {{ any_expr }}
    }
    {{/if}}
{{/each}}
}

impl Default for SpatialHashCell {
    fn default() -> Self {
        Self::new()
    }
}

pub struct SpatialHashTable {
    grid: StaticGrid<SpatialHashCell>,
    empty: SpatialHashCell,
}

impl SpatialHashTable {
    pub fn new(width: usize, height: usize) -> Self {
        SpatialHashTable {
            grid: StaticGrid::new_default(width, height),
            empty: SpatialHashCell::new(),
        }
    }

    pub fn cell_iter(&self) -> SpatialHashCellIter {
        self.grid.iter()
    }

    pub fn coord_iter(&self) -> SpatialHashCoordIter {
        self.grid.coord_iter()
    }

    pub fn get(&self, coord: Coord) -> &SpatialHashCell {
        self.grid.get(coord).unwrap_or(&self.empty)
    }

    pub fn limits_min(&self) -> Coord {
        self.grid.limits_min()
    }

    pub fn limits_max(&self) -> Coord {
        self.grid.limits_max()
    }

    pub fn width(&self) -> usize {
        self.grid.width()
    }

    pub fn height(&self) -> usize {
        self.grid.height()
    }

    fn get_mut(&mut self, coord: Coord) -> &mut SpatialHashCell {
        unsafe { self.grid.get_unchecked_mut(coord) }
    }

    fn change_entity_position(&mut self, entity: EntityRef, current_position: Coord, new_position: Coord, action_id: u64) {
        self.remove_entity_position(entity, current_position, action_id);
        self.add_entity_position(entity, new_position, action_id);
    }

    fn remove_entity_position(&mut self, entity: EntityRef, position: Coord, action_id: u64) {
        let mut cell = self.get_mut(position);

{{#each field}}
        {{ remove_entity }}
{{/each}}

        cell.entities.remove(entity.id());
        cell.last_updated = action_id;
    }

    fn add_entity_position(&mut self, entity: EntityRef, position: Coord, action_id: u64) {
        let mut cell = self.get_mut(position);

{{#each field}}
        {{ insert_entity }}
{{/each}}

        cell.entities.insert(entity.id());
        cell.last_updated = action_id;
    }

    pub fn update(&mut self, ecs: &EcsCtx, action: &EcsAction, action_id: u64) {

        for (entity_id, new_position) in action.position_positive_iter(ecs) {
            let entity = ecs.entity(entity_id);
            // Add and remove tracked components based on the current data stored about the
            // entity, ignoring any component changes in the current action. These will be
            // applied later.
            if let Some(current_position) = entity.position() {
                // the entity is changing position
                self.change_entity_position(entity, current_position, *new_position, action_id);
            } else {
                // the entity is gaining a position
                self.add_entity_position(entity, *new_position, action_id);
            }
        }

        for entity_id in action.position_negative_iter(ecs) {
            let entity = ecs.entity(entity_id);
            if let Some(position) = entity.position() {
                self.remove_entity_position(entity, position, action_id);
            }
        }

{{#each field}}
        self.update_{{ struct_field_name }}(ecs, action, action_id);
{{/each}}
    }

{{#each field}}
    fn update_{{ struct_field_name }}(&mut self, ecs: &EcsCtx, action: &EcsAction, action_id: u64) {
    {{#if count_bool}}
        for entity_id in action.{{ component_name }}_positive_iter(ecs) {
            let entity = ecs.post_action_entity(entity_id, action);
            if let Some(position) = entity.position() {
                if !entity.current_contains_{{ component_name }}() {
                    let cell = self.get_mut(position);
                    cell.{{ struct_field_name }} += 1;
                    cell.last_updated = action_id;
                }
            }
        }

        for entity_id in action.{{ component_name }}_negative_iter(ecs) {
            let entity = ecs.post_action_entity(entity_id, action);
            if let Some(position) = entity.position() {
                if entity.current_contains_{{ component_name }}() {
                    let cell = self.get_mut(position);
                    cell.{{ struct_field_name }} -= 1;
                    cell.last_updated = action_id;
                }
            }
        }
    {{/if}}
    {{#if sum_f64}}
        for (entity_id, new) in action.{{ component_name }}_positive_iter(ecs) {
            let entity = ecs.post_action_entity(entity_id, action);
            if let Some(position) = entity.position() {
                let current = entity.current_{{ component_name }}().unwrap_or(0.0);
                let increase = new - current;
                let cell = self.get_mut(position);
                cell.{{ struct_field_name }} += increase;
                cell.last_updated = action_id;
            }
        }
        for entity_id in action.{{ component_name }}_negative_iter(ecs) {
            let entity = ecs.post_action_entity(entity_id, action);
            if let Some(position) = entity.position() {
                if let Some(value) = entity.current_{{ component_name }}() {
                    let cell = self.get_mut(position);
                    cell.{{ struct_field_name }} -= value;
                    cell.last_updated = action_id;
                }
            }
        }
    {{/if}}
    {{#if is_set_type}}
        {{#if component_has_type}}
        for (entity_id, _) in action.{{ component_name }}_positive_iter(ecs) {
        {{else}}
        for entity_id in action.{{ component_name }}_positive_iter(ecs) {
        {{/if}}
            let entity = ecs.post_action_entity(entity_id, action);
            if let Some(position) = entity.position() {
        {{#if component_has_type}}
                if entity.current_{{ component_name }}().is_none() {
        {{else}}
                if !entity.current_contains_{{ component_name }}() {
        {{/if}}
                    let cell = self.get_mut(position);
                    cell.{{ struct_field_name }}.insert(entity_id);
                    cell.last_updated = action_id;
                }
            }
        }
        for entity_id in action.{{ component_name }}_negative_iter(ecs) {
            let entity = ecs.post_action_entity(entity_id, action);
            if let Some(position) = entity.position() {
        {{#if component_has_type}}
                if entity.current_{{ component_name }}().is_some() {
        {{else}}
                if entity.current_contains_{{ component_name }}() {
        {{/if}}
                    let cell = self.get_mut(position);
                    cell.{{ struct_field_name }}.remove(entity_id);
                    cell.last_updated = action_id;
                }
            }
        }
    {{/if}}
    }
{{/each}}
}
"#;

<commit_msg>Spatial hash updates timestamp on changes<commit_after>use std::io::Read;
use std::io::Write;
use std::path::Path;
use std::fs::File;

use tomson::Toml;
use handlebars::Handlebars;
use rustc_serialize::json::Json;

fn generate_code(mut toml: String) -> String {
    // turn the toml string into json for compatibility with handlebars
    let mut json = Toml::as_json(&mut toml).unwrap();

    for (id, field) in json.as_object_mut().unwrap().get_mut("field").unwrap().as_object_mut().unwrap().iter_mut() {
        let field_obj = field.as_object_mut().unwrap();
        let field_type = field_obj.get("type").unwrap().as_string().unwrap().to_string();
        let field_name = id.clone();
        let component_name = if let Some(component_json) = field_obj.get("component") {
            component_json.as_string().unwrap().to_string()
        } else {
            id.clone()
        };

        let component_has_type = if let Some(component_has_type_json) = field_obj.get("component_has_type") {
            component_has_type_json.as_boolean().unwrap()
        } else {
            false
        };

        field_obj.insert("struct_field_name".to_string(), Json::String(field_name.clone()));
        field_obj.insert("getter_name".to_string(), Json::String(id.clone()));
        field_obj.insert("component_name".to_string(), Json::String(component_name.clone()));
        field_obj.insert(field_type.clone(), Json::Boolean(true));

        if component_has_type {
            field_obj.insert("component_has_type".to_string(), Json::Boolean(true));
        }

        match field_type.as_ref() {
            "sum_f64" => {
                field_obj.insert("struct_field_type".to_string(), Json::String("f64".to_string()));
                field_obj.insert("struct_field_cons".to_string(), Json::String("0.0".to_string()));
                field_obj.insert("getter_type".to_string(), Json::String("f64".to_string()));
                field_obj.insert("getter_expr".to_string(), Json::String(format!("self.{}", field_name)));
                field_obj.insert("remove_entity".to_string(), Json::String(format!("if let Some(v) = entity.{}() {{ cell.{} -= v; }}",
                                                                                component_name, field_name)));
                field_obj.insert("insert_entity".to_string(), Json::String(format!("if let Some(v) = entity.{}() {{ cell.{} += v; }}",
                                                                                component_name, field_name)));
            }
            "count_bool" => {
                field_obj.insert("struct_field_type".to_string(), Json::String("usize".to_string()));
                field_obj.insert("struct_field_cons".to_string(), Json::String("0".to_string()));
                field_obj.insert("getter_type".to_string(), Json::String("bool".to_string()));
                field_obj.insert("getter_expr".to_string(), Json::String(format!("self.{} != 0", field_name)));
                field_obj.insert("remove_entity".to_string(), Json::String(format!("if entity.contains_{}() {{ cell.{} -= 1; }}",
                                                                                component_name, field_name)));
                field_obj.insert("insert_entity".to_string(), Json::String(format!("if entity.contains_{}() {{ cell.{} += 1; }}",
                                                                                component_name, field_name)));
            }
            "any" => {
                field_obj.insert("struct_field_type".to_string(), Json::String("AnySet<EntityId>".to_string()));
                field_obj.insert("struct_field_cons".to_string(), Json::String("AnySet::new()".to_string()));
                field_obj.insert("getter_type".to_string(), Json::String("bool".to_string()));
                field_obj.insert("getter_expr".to_string(), Json::String(format!("!self.{}.is_empty()", field_name)));
                field_obj.insert("any_name".to_string(), Json::String(format!("any_{}", field_name)));
                field_obj.insert("any_type".to_string(), Json::String("Option<EntityId>".to_string()));
                field_obj.insert("any_expr".to_string(), Json::String(format!("self.{}.any()", field_name)));
                field_obj.insert("remove_entity".to_string(), Json::String(format!("if entity.contains_{}() {{ cell.{}.remove(entity.id()); }}",
                                                                                component_name, field_name)));
                field_obj.insert("insert_entity".to_string(), Json::String(format!("if entity.contains_{}() {{ cell.{}.insert(entity.id()); }}",
                                                                                component_name, field_name)));
                field_obj.insert("is_set_type".to_string(), Json::Boolean(true));
            }
            "set" => {
                field_obj.insert("struct_field_type".to_string(), Json::String("EntitySet".to_string()));
                field_obj.insert("struct_field_cons".to_string(), Json::String("EntitySet::new()".to_string()));
                field_obj.insert("getter_type".to_string(), Json::String("bool".to_string()));
                field_obj.insert("getter_expr".to_string(), Json::String(format!("!self.{}.is_empty()", field_name)));
                field_obj.insert("iter_name".to_string(), Json::String(format!("{}_iter", field_name)));
                field_obj.insert("iter_type".to_string(), Json::String("EntitySetIter".to_string()));
                field_obj.insert("iter_expr".to_string(), Json::String(format!("self.{}.iter()", field_name)));
                field_obj.insert("remove_entity".to_string(), Json::String(format!("if entity.contains_{}() {{ cell.{}.remove(entity.id()); }}",
                                                                                component_name, field_name)));
                field_obj.insert("insert_entity".to_string(), Json::String(format!("if entity.contains_{}() {{ cell.{}.insert(entity.id()); }}",
                                                                                component_name, field_name)));
                field_obj.insert("is_set_type".to_string(), Json::Boolean(true));
            }
            other => panic!("unknown field type {}", other),
        }
    }

    let mut handlebars = Handlebars::new();

    // prevent xml escaping
    handlebars.register_escape_fn(|input| input.to_string());
    handlebars.template_render(TEMPLATE, &json).unwrap()

}

fn read_file_to_string<P: AsRef<Path>>(path: P) -> String {
    let mut file = File::open(path).unwrap();
    let mut string = String::new();
    file.read_to_string(&mut string).unwrap();

    string
}

pub fn generate_spatial_hash<P: AsRef<Path>, Q: AsRef<Path>>(in_path: P, out_path: Q) {
    let string = read_file_to_string(in_path);

    let output_string = generate_code(string);

    let mut outfile = File::create(out_path).unwrap();
    write!(outfile, "{}", output_string).unwrap();
}

const TEMPLATE: &'static str = r#"// Automatically generated. Do not edit.
#![allow(unused_imports)]
use ecs::*;
use coord::Coord;
use grid::{Grid, StaticGrid, DefaultGrid, IterGrid, CoordIterGrid};
use util::AnySet;

pub type SpatialHashCellIter<'a> = <StaticGrid<SpatialHashCell> as IterGrid<'a>>::Iter;
pub type SpatialHashCoordIter = <StaticGrid<SpatialHashCell> as CoordIterGrid>::CoordIter;

pub struct SpatialHashCell {

{{#each field}}
    {{ struct_field_name }}: {{ struct_field_type }},
{{/each}}

    // set of entities currently in this cell
    entities: EntitySet,

    // action on which this cell was last updated
    last_updated: u64,
}

impl SpatialHashCell {
    pub fn new() -> Self {
        SpatialHashCell {
{{#each field}}
            {{ struct_field_name }}: {{ struct_field_cons }},
{{/each}}

            entities: EntitySet::new(),
            last_updated: 0,
        }
    }

    pub fn last_updated(&self) -> u64 {
        self.last_updated
    }

    pub fn entity_ids(&self) -> &EntitySet {
        &self.entities
    }

    pub fn entity_id_iter(&self) -> EntitySetIter {
        self.entities.iter()
    }

{{#each field}}
    pub fn {{ getter_name }}(&self) -> {{ getter_type }} {
        {{ getter_expr }}
    }
    {{#if iter_name}}
    pub fn {{ iter_name }}(&self) -> {{ iter_type }} {
        {{ iter_expr }}
    }
    {{/if}}
    {{#if any_name}}
    pub fn {{ any_name }}(&self) -> {{ any_type }} {
        {{ any_expr }}
    }
    {{/if}}
{{/each}}
}

impl Default for SpatialHashCell {
    fn default() -> Self {
        Self::new()
    }
}

pub struct SpatialHashTable {
    grid: StaticGrid<SpatialHashCell>,
    empty: SpatialHashCell,
}

impl SpatialHashTable {
    pub fn new(width: usize, height: usize) -> Self {
        SpatialHashTable {
            grid: StaticGrid::new_default(width, height),
            empty: SpatialHashCell::new(),
        }
    }

    pub fn cell_iter(&self) -> SpatialHashCellIter {
        self.grid.iter()
    }

    pub fn coord_iter(&self) -> SpatialHashCoordIter {
        self.grid.coord_iter()
    }

    pub fn get(&self, coord: Coord) -> &SpatialHashCell {
        self.grid.get(coord).unwrap_or(&self.empty)
    }

    pub fn limits_min(&self) -> Coord {
        self.grid.limits_min()
    }

    pub fn limits_max(&self) -> Coord {
        self.grid.limits_max()
    }

    pub fn width(&self) -> usize {
        self.grid.width()
    }

    pub fn height(&self) -> usize {
        self.grid.height()
    }

    fn get_mut(&mut self, coord: Coord) -> &mut SpatialHashCell {
        unsafe { self.grid.get_unchecked_mut(coord) }
    }

    fn change_entity_position(&mut self, entity: EntityRef, current_position: Coord, new_position: Coord, action_id: u64) {
        self.remove_entity_position(entity, current_position, action_id);
        self.add_entity_position(entity, new_position, action_id);
    }

    fn remove_entity_position(&mut self, entity: EntityRef, position: Coord, action_id: u64) {
        let mut cell = self.get_mut(position);

{{#each field}}
        {{ remove_entity }}
{{/each}}

        cell.entities.remove(entity.id());
        cell.last_updated = action_id;
    }

    fn add_entity_position(&mut self, entity: EntityRef, position: Coord, action_id: u64) {
        let mut cell = self.get_mut(position);

{{#each field}}
        {{ insert_entity }}
{{/each}}

        cell.entities.insert(entity.id());
        cell.last_updated = action_id;
    }

    pub fn update(&mut self, ecs: &EcsCtx, action: &EcsAction, action_id: u64) {

        for (entity_id, new_position) in action.position_positive_iter(ecs) {
            let entity = ecs.entity(entity_id);
            // Add and remove tracked components based on the current data stored about the
            // entity, ignoring any component changes in the current action. These will be
            // applied later.
            if let Some(current_position) = entity.position() {
                // the entity is changing position
                self.change_entity_position(entity, current_position, *new_position, action_id);
            } else {
                // the entity is gaining a position
                self.add_entity_position(entity, *new_position, action_id);
            }
        }

        for entity_id in action.position_negative_iter(ecs) {
            let entity = ecs.entity(entity_id);
            if let Some(position) = entity.position() {
                self.remove_entity_position(entity, position, action_id);
            }
        }

{{#each field}}
        self.update_{{ struct_field_name }}(ecs, action, action_id);
{{/each}}
    }

{{#each field}}
    fn update_{{ struct_field_name }}(&mut self, ecs: &EcsCtx, action: &EcsAction, action_id: u64) {
    {{#if count_bool}}
        for entity_id in action.{{ component_name }}_positive_iter(ecs) {
            let entity = ecs.post_action_entity(entity_id, action);
            if let Some(position) = entity.position() {
                if !entity.current_contains_{{ component_name }}() {
                    let cell = self.get_mut(position);
                    cell.{{ struct_field_name }} += 1;
                    cell.last_updated = action_id;
                }
            }
        }

        for entity_id in action.{{ component_name }}_negative_iter(ecs) {
            let entity = ecs.post_action_entity(entity_id, action);
            if let Some(position) = entity.position() {
                if entity.current_contains_{{ component_name }}() {
                    let cell = self.get_mut(position);
                    cell.{{ struct_field_name }} -= 1;
                    cell.last_updated = action_id;
                }
            }
        }
    {{/if}}
    {{#if sum_f64}}
        for (entity_id, new) in action.{{ component_name }}_positive_iter(ecs) {
            let entity = ecs.post_action_entity(entity_id, action);
            if let Some(position) = entity.position() {
                let current = entity.current_{{ component_name }}().unwrap_or(0.0);
                let increase = new - current;
                let cell = self.get_mut(position);
                cell.{{ struct_field_name }} += increase;
                cell.last_updated = action_id;
            }
        }
        for entity_id in action.{{ component_name }}_negative_iter(ecs) {
            let entity = ecs.post_action_entity(entity_id, action);
            if let Some(position) = entity.position() {
                if let Some(value) = entity.current_{{ component_name }}() {
                    let cell = self.get_mut(position);
                    cell.{{ struct_field_name }} -= value;
                    cell.last_updated = action_id;
                }
            }
        }
    {{/if}}
    {{#if is_set_type}}
        {{#if component_has_type}}
        for (entity_id, _) in action.{{ component_name }}_positive_iter(ecs) {
        {{else}}
        for entity_id in action.{{ component_name }}_positive_iter(ecs) {
        {{/if}}
            let entity = ecs.post_action_entity(entity_id, action);
            if let Some(position) = entity.position() {
                let cell = self.get_mut(position);
        {{#if component_has_type}}
                if entity.current_{{ component_name }}().is_none() {
        {{else}}
                if !entity.current_contains_{{ component_name }}() {
        {{/if}}
                    cell.{{ struct_field_name }}.insert(entity_id);
                }
                cell.last_updated = action_id;
            }
        }
        for entity_id in action.{{ component_name }}_negative_iter(ecs) {
            let entity = ecs.post_action_entity(entity_id, action);
            if let Some(position) = entity.position() {
        {{#if component_has_type}}
                if entity.current_{{ component_name }}().is_some() {
        {{else}}
                if entity.current_contains_{{ component_name }}() {
        {{/if}}
                    let cell = self.get_mut(position);
                    cell.{{ struct_field_name }}.remove(entity_id);
                    cell.last_updated = action_id;
                }
            }
        }
    {{/if}}
    }
{{/each}}
}
"#;

<|endoftext|>


<commit_before>use std::{cmp, fmt, ptr};
#[cfg(not(target_os = "netbsd"))]
use std::os::raw::{c_int, c_short};
use std::os::unix::io::AsRawFd;
use std::os::unix::io::RawFd;
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
use std::time::Duration;

use libc::{self, time_t};

use {io, Ready, PollOpt, Token};
use event_imp::{self as event, Event};
use sys::unix::{cvt, UnixReady};
use sys::unix::io::set_cloexec;

/// Each Selector has a globally unique(ish) ID associated with it. This ID
/// gets tracked by `TcpStream`, `TcpListener`, etc... when they are first
/// registered with the `Selector`. If a type that is previously associated with
/// a `Selector` attempts to register itself with a different `Selector`, the
/// operation will return with an error. This matches windows behavior.
static NEXT_ID: AtomicUsize = ATOMIC_USIZE_INIT;

#[cfg(not(target_os = "netbsd"))]
type Filter = c_short;
#[cfg(not(target_os = "netbsd"))]
type UData = *mut ::libc::c_void;
#[cfg(not(target_os = "netbsd"))]
type Count = c_int;

#[cfg(target_os = "netbsd")]
type Filter = u32;
#[cfg(target_os = "netbsd")]
type UData = ::libc::intptr_t;
#[cfg(target_os = "netbsd")]
type Count = usize;

macro_rules! kevent {
    ($id: expr, $filter: expr, $flags: expr, $data: expr) => {
        libc::kevent {
            ident: $id as ::libc::uintptr_t,
            filter: $filter as Filter,
            flags: $flags,
            fflags: 0,
            data: 0,
            udata: $data as UData,
        }
    }
}

pub struct Selector {
    id: usize,
    kq: RawFd,
}

impl Selector {
    pub fn new() -> io::Result<Selector> {
        // offset by 1 to avoid choosing 0 as the id of a selector
        let id = NEXT_ID.fetch_add(1, Ordering::Relaxed) + 1;
        let kq = unsafe { cvt(libc::kqueue())? };
        drop(set_cloexec(kq));

        Ok(Selector {
            id,
            kq,
        })
    }

    pub fn id(&self) -> usize {
        self.id
    }

    pub fn select(&self, evts: &mut Events, awakener: Token, timeout: Option<Duration>) -> io::Result<bool> {
        let timeout = timeout.map(|to| {
            libc::timespec {
                tv_sec: cmp::min(to.as_secs(), time_t::max_value() as u64) as time_t,
                tv_nsec: libc::c_long::from(to.subsec_nanos()),
            }
        });
        let timeout = timeout.as_ref().map(|s| s as *const _).unwrap_or(ptr::null_mut());

        evts.clear();
        unsafe {
            let cnt = cvt(libc::kevent(self.kq,
                                            ptr::null(),
                                            0,
                                            evts.sys_events.0.as_mut_ptr(),
                                            evts.sys_events.0.capacity() as Count,
                                            timeout))?;
            evts.sys_events.0.set_len(cnt as usize);
            Ok(evts.coalesce(awakener))
        }
    }

    pub fn register(&self, fd: RawFd, token: Token, interests: Ready, opts: PollOpt) -> io::Result<()> {
        trace!("registering; token={:?}; interests={:?}", token, interests);

        let flags = if opts.contains(PollOpt::edge()) { libc::EV_CLEAR } else { 0 } |
                    if opts.contains(PollOpt::oneshot()) { libc::EV_ONESHOT } else { 0 } |
                    libc::EV_RECEIPT;

        unsafe {
            let r = if interests.contains(Ready::readable()) { libc::EV_ADD } else { libc::EV_DELETE };
            let w = if interests.contains(Ready::writable()) { libc::EV_ADD } else { libc::EV_DELETE };
            let mut changes = [
                kevent!(fd, libc::EVFILT_READ, flags | r, usize::from(token)),
                kevent!(fd, libc::EVFILT_WRITE, flags | w, usize::from(token)),
            ];

            cvt(libc::kevent(self.kq,
                             changes.as_ptr(),
                             changes.len() as Count,
                             changes.as_mut_ptr(),
                             changes.len() as Count,
                             ::std::ptr::null()))?;

            for change in changes.iter() {
                debug_assert_eq!(change.flags & libc::EV_ERROR, libc::EV_ERROR);

                // Test to see if an error happened
                if change.data == 0 {
                    continue
                }

                // Older versions of OSX (10.11 and 10.10 have been witnessed)
                // can return EPIPE when registering a pipe file descriptor
                // where the other end has already disappeared. For example code
                // that creates a pipe, closes a file descriptor, and then
                // registers the other end will see an EPIPE returned from
                // `register`.
                //
                // It also turns out that kevent will still report events on the
                // file descriptor, telling us that it's readable/hup at least
                // after we've done this registration. As a result we just
                // ignore `EPIPE` here instead of propagating it.
                //
                // More info can be found at carllerche/mio#582
                if change.data as i32 == libc::EPIPE &&
                   change.filter == libc::EVFILT_WRITE as Filter {
                    continue
                }

                // ignore ENOENT error for EV_DELETE
                let orig_flags = if change.filter == libc::EVFILT_READ as Filter { r } else { w };
                if change.data as i32 == libc::ENOENT && orig_flags & libc::EV_DELETE != 0 {
                    continue
                }

                return Err(::std::io::Error::from_raw_os_error(change.data as i32));
            }
            Ok(())
        }
    }

    pub fn reregister(&self, fd: RawFd, token: Token, interests: Ready, opts: PollOpt) -> io::Result<()> {
        // Just need to call register here since EV_ADD is a mod if already
        // registered
        self.register(fd, token, interests, opts)
    }

    pub fn deregister(&self, fd: RawFd) -> io::Result<()> {
        unsafe {
            // EV_RECEIPT is a nice way to apply changes and get back per-event results while not
            // draining the actual changes.
            let filter = libc::EV_DELETE | libc::EV_RECEIPT;
#[cfg(not(target_os = "netbsd"))]
            let mut changes = [
                kevent!(fd, libc::EVFILT_READ, filter, ptr::null_mut()),
                kevent!(fd, libc::EVFILT_WRITE, filter, ptr::null_mut()),
            ];

#[cfg(target_os = "netbsd")]
            let mut changes = [
                kevent!(fd, libc::EVFILT_READ, filter, 0),
                kevent!(fd, libc::EVFILT_WRITE, filter, 0),
            ];

            cvt(libc::kevent(self.kq,
                             changes.as_ptr(),
                             changes.len() as Count,
                             changes.as_mut_ptr(),
                             changes.len() as Count,
                             ::std::ptr::null())).map(|_| ())?;

            if changes[0].data as i32 == libc::ENOENT && changes[1].data as i32 == libc::ENOENT {
                return Err(::std::io::Error::from_raw_os_error(changes[0].data as i32));
            }
            for change in changes.iter() {
                debug_assert_eq!(libc::EV_ERROR & change.flags, libc::EV_ERROR);
                if change.data != 0 && change.data as i32 != libc::ENOENT {
                    return Err(::std::io::Error::from_raw_os_error(changes[0].data as i32));
                }
            }
            Ok(())
        }
    }
}

impl fmt::Debug for Selector {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Selector")
            .field("id", &self.id)
            .field("kq", &self.kq)
            .finish()
    }
}

impl AsRawFd for Selector {
    fn as_raw_fd(&self) -> RawFd {
        self.kq
    }
}

impl Drop for Selector {
    fn drop(&mut self) {
        unsafe {
            let _ = libc::close(self.kq);
        }
    }
}

pub struct Events {
    sys_events: KeventList,
    events: Vec<Event>,
    event_map: HashMap<Token, usize>,
}

struct KeventList(Vec<libc::kevent>);

unsafe impl Send for KeventList {}
unsafe impl Sync for KeventList {}

impl Events {
    pub fn with_capacity(cap: usize) -> Events {
        Events {
            sys_events: KeventList(Vec::with_capacity(cap)),
            events: Vec::with_capacity(cap),
            event_map: HashMap::with_capacity(cap)
        }
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.events.len()
    }

    #[inline]
    pub fn capacity(&self) -> usize {
        self.events.capacity()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.events.is_empty()
    }

    pub fn get(&self, idx: usize) -> Option<Event> {
        self.events.get(idx).cloned()
    }

    fn coalesce(&mut self, awakener: Token) -> bool {
        let mut ret = false;
        self.events.clear();
        self.event_map.clear();

        for e in self.sys_events.0.iter() {
            let token = Token(e.udata as usize);
            let len = self.events.len();

            if token == awakener {
                // TODO: Should this return an error if event is an error. It
                // is not critical as spurious wakeups are permitted.
                ret = true;
                continue;
            }

            let idx = *self.event_map.entry(token)
                .or_insert(len);

            if idx == len {
                // New entry, insert the default
                self.events.push(Event::new(Ready::empty(), token));

            }

            if e.flags & libc::EV_ERROR != 0 {
                event::kind_mut(&mut self.events[idx]).insert(*UnixReady::error());
            }

            if e.filter == libc::EVFILT_READ as Filter {
                event::kind_mut(&mut self.events[idx]).insert(Ready::readable());
            } else if e.filter == libc::EVFILT_WRITE as Filter {
                event::kind_mut(&mut self.events[idx]).insert(Ready::writable());

                // `EV_EOF` set with `EVFILT_WRITE` indicates the connection has been fully
                // disconnected (read and write), but only sockets...
                if e.flags & libc::EV_EOF != 0 {
                    event::kind_mut(&mut self.events[idx]).insert(UnixReady::hup());

                    // When the read end of the socket is closed, EV_EOF is set on
                    // flags, and fflags contains the error if there is one.
                    if e.fflags != 0 {
                        event::kind_mut(&mut self.events[idx]).insert(UnixReady::error());
                    }
                }
            }
#[cfg(any(target_os = "dragonfly",
    target_os = "freebsd", target_os = "ios", target_os = "macos"))]
            {
                if e.filter == libc::EVFILT_AIO {
                    event::kind_mut(&mut self.events[idx]).insert(UnixReady::aio());
                }
            }
#[cfg(any(target_os = "freebsd"))]
            {
                if e.filter == libc::EVFILT_LIO {
                    event::kind_mut(&mut self.events[idx]).insert(UnixReady::lio());
                }
            }
        }

        ret
    }

    pub fn push_event(&mut self, event: Event) {
        self.events.push(event);
    }

    pub fn clear(&mut self) {
        self.sys_events.0.truncate(0);
        self.events.truncate(0);
        self.event_map.clear();
    }
}

impl fmt::Debug for Events {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Events")
            .field("len", &self.sys_events.0.len())
            .finish()
    }
}

#[test]
fn does_not_register_rw() {
    use {Poll, Ready, PollOpt, Token};
    use unix::EventedFd;

    let kq = unsafe { libc::kqueue() };
    let kqf = EventedFd(&kq);
    let poll = Poll::new().unwrap();

    // registering kqueue fd will fail if write is requested (On anything but some versions of OS
    // X)
    poll.register(&kqf, Token(1234), Ready::readable(),
                  PollOpt::edge() | PollOpt::oneshot()).unwrap();
}

#[cfg(any(target_os = "dragonfly",
    target_os = "freebsd", target_os = "ios", target_os = "macos"))]
#[test]
fn test_coalesce_aio() {
    let mut events = Events::with_capacity(1);
    events.sys_events.0.push(kevent!(0x1234, libc::EVFILT_AIO, 0, 42));
    events.coalesce(Token(0));
    assert!(events.events[0].readiness() == UnixReady::aio().into());
    assert!(events.events[0].token() == Token(42));
}
<commit_msg>Fix kqueue on platforms where C's long is i32<commit_after>use std::{cmp, fmt, ptr};
#[cfg(not(target_os = "netbsd"))]
use std::os::raw::{c_int, c_short};
use std::os::unix::io::AsRawFd;
use std::os::unix::io::RawFd;
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
use std::time::Duration;

use libc::{self, time_t};

use {io, Ready, PollOpt, Token};
use event_imp::{self as event, Event};
use sys::unix::{cvt, UnixReady};
use sys::unix::io::set_cloexec;

/// Each Selector has a globally unique(ish) ID associated with it. This ID
/// gets tracked by `TcpStream`, `TcpListener`, etc... when they are first
/// registered with the `Selector`. If a type that is previously associated with
/// a `Selector` attempts to register itself with a different `Selector`, the
/// operation will return with an error. This matches windows behavior.
static NEXT_ID: AtomicUsize = ATOMIC_USIZE_INIT;

#[cfg(not(target_os = "netbsd"))]
type Filter = c_short;
#[cfg(not(target_os = "netbsd"))]
type UData = *mut ::libc::c_void;
#[cfg(not(target_os = "netbsd"))]
type Count = c_int;

#[cfg(target_os = "netbsd")]
type Filter = u32;
#[cfg(target_os = "netbsd")]
type UData = ::libc::intptr_t;
#[cfg(target_os = "netbsd")]
type Count = usize;

macro_rules! kevent {
    ($id: expr, $filter: expr, $flags: expr, $data: expr) => {
        libc::kevent {
            ident: $id as ::libc::uintptr_t,
            filter: $filter as Filter,
            flags: $flags,
            fflags: 0,
            data: 0,
            udata: $data as UData,
        }
    }
}

pub struct Selector {
    id: usize,
    kq: RawFd,
}

impl Selector {
    pub fn new() -> io::Result<Selector> {
        // offset by 1 to avoid choosing 0 as the id of a selector
        let id = NEXT_ID.fetch_add(1, Ordering::Relaxed) + 1;
        let kq = unsafe { cvt(libc::kqueue())? };
        drop(set_cloexec(kq));

        Ok(Selector {
            id,
            kq,
        })
    }

    pub fn id(&self) -> usize {
        self.id
    }

    pub fn select(&self, evts: &mut Events, awakener: Token, timeout: Option<Duration>) -> io::Result<bool> {
        let timeout = timeout.map(|to| {
            libc::timespec {
                tv_sec: cmp::min(to.as_secs(), time_t::max_value() as u64) as time_t,
                // `Duration::subsec_nanos` is guaranteed to be less than one
                // billion (the number of nanoseconds in a second), making the
                // cast to i32 safe. The cast itself is needed for platforms
                // where C's long is only 32 bits.
                tv_nsec: libc::c_long::from(to.subsec_nanos() as i32),
            }
        });
        let timeout = timeout.as_ref().map(|s| s as *const _).unwrap_or(ptr::null_mut());

        evts.clear();
        unsafe {
            let cnt = cvt(libc::kevent(self.kq,
                                            ptr::null(),
                                            0,
                                            evts.sys_events.0.as_mut_ptr(),
                                            evts.sys_events.0.capacity() as Count,
                                            timeout))?;
            evts.sys_events.0.set_len(cnt as usize);
            Ok(evts.coalesce(awakener))
        }
    }

    pub fn register(&self, fd: RawFd, token: Token, interests: Ready, opts: PollOpt) -> io::Result<()> {
        trace!("registering; token={:?}; interests={:?}", token, interests);

        let flags = if opts.contains(PollOpt::edge()) { libc::EV_CLEAR } else { 0 } |
                    if opts.contains(PollOpt::oneshot()) { libc::EV_ONESHOT } else { 0 } |
                    libc::EV_RECEIPT;

        unsafe {
            let r = if interests.contains(Ready::readable()) { libc::EV_ADD } else { libc::EV_DELETE };
            let w = if interests.contains(Ready::writable()) { libc::EV_ADD } else { libc::EV_DELETE };
            let mut changes = [
                kevent!(fd, libc::EVFILT_READ, flags | r, usize::from(token)),
                kevent!(fd, libc::EVFILT_WRITE, flags | w, usize::from(token)),
            ];

            cvt(libc::kevent(self.kq,
                             changes.as_ptr(),
                             changes.len() as Count,
                             changes.as_mut_ptr(),
                             changes.len() as Count,
                             ::std::ptr::null()))?;

            for change in changes.iter() {
                debug_assert_eq!(change.flags & libc::EV_ERROR, libc::EV_ERROR);

                // Test to see if an error happened
                if change.data == 0 {
                    continue
                }

                // Older versions of OSX (10.11 and 10.10 have been witnessed)
                // can return EPIPE when registering a pipe file descriptor
                // where the other end has already disappeared. For example code
                // that creates a pipe, closes a file descriptor, and then
                // registers the other end will see an EPIPE returned from
                // `register`.
                //
                // It also turns out that kevent will still report events on the
                // file descriptor, telling us that it's readable/hup at least
                // after we've done this registration. As a result we just
                // ignore `EPIPE` here instead of propagating it.
                //
                // More info can be found at carllerche/mio#582
                if change.data as i32 == libc::EPIPE &&
                   change.filter == libc::EVFILT_WRITE as Filter {
                    continue
                }

                // ignore ENOENT error for EV_DELETE
                let orig_flags = if change.filter == libc::EVFILT_READ as Filter { r } else { w };
                if change.data as i32 == libc::ENOENT && orig_flags & libc::EV_DELETE != 0 {
                    continue
                }

                return Err(::std::io::Error::from_raw_os_error(change.data as i32));
            }
            Ok(())
        }
    }

    pub fn reregister(&self, fd: RawFd, token: Token, interests: Ready, opts: PollOpt) -> io::Result<()> {
        // Just need to call register here since EV_ADD is a mod if already
        // registered
        self.register(fd, token, interests, opts)
    }

    pub fn deregister(&self, fd: RawFd) -> io::Result<()> {
        unsafe {
            // EV_RECEIPT is a nice way to apply changes and get back per-event results while not
            // draining the actual changes.
            let filter = libc::EV_DELETE | libc::EV_RECEIPT;
#[cfg(not(target_os = "netbsd"))]
            let mut changes = [
                kevent!(fd, libc::EVFILT_READ, filter, ptr::null_mut()),
                kevent!(fd, libc::EVFILT_WRITE, filter, ptr::null_mut()),
            ];

#[cfg(target_os = "netbsd")]
            let mut changes = [
                kevent!(fd, libc::EVFILT_READ, filter, 0),
                kevent!(fd, libc::EVFILT_WRITE, filter, 0),
            ];

            cvt(libc::kevent(self.kq,
                             changes.as_ptr(),
                             changes.len() as Count,
                             changes.as_mut_ptr(),
                             changes.len() as Count,
                             ::std::ptr::null())).map(|_| ())?;

            if changes[0].data as i32 == libc::ENOENT && changes[1].data as i32 == libc::ENOENT {
                return Err(::std::io::Error::from_raw_os_error(changes[0].data as i32));
            }
            for change in changes.iter() {
                debug_assert_eq!(libc::EV_ERROR & change.flags, libc::EV_ERROR);
                if change.data != 0 && change.data as i32 != libc::ENOENT {
                    return Err(::std::io::Error::from_raw_os_error(changes[0].data as i32));
                }
            }
            Ok(())
        }
    }
}

impl fmt::Debug for Selector {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Selector")
            .field("id", &self.id)
            .field("kq", &self.kq)
            .finish()
    }
}

impl AsRawFd for Selector {
    fn as_raw_fd(&self) -> RawFd {
        self.kq
    }
}

impl Drop for Selector {
    fn drop(&mut self) {
        unsafe {
            let _ = libc::close(self.kq);
        }
    }
}

pub struct Events {
    sys_events: KeventList,
    events: Vec<Event>,
    event_map: HashMap<Token, usize>,
}

struct KeventList(Vec<libc::kevent>);

unsafe impl Send for KeventList {}
unsafe impl Sync for KeventList {}

impl Events {
    pub fn with_capacity(cap: usize) -> Events {
        Events {
            sys_events: KeventList(Vec::with_capacity(cap)),
            events: Vec::with_capacity(cap),
            event_map: HashMap::with_capacity(cap)
        }
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.events.len()
    }

    #[inline]
    pub fn capacity(&self) -> usize {
        self.events.capacity()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.events.is_empty()
    }

    pub fn get(&self, idx: usize) -> Option<Event> {
        self.events.get(idx).cloned()
    }

    fn coalesce(&mut self, awakener: Token) -> bool {
        let mut ret = false;
        self.events.clear();
        self.event_map.clear();

        for e in self.sys_events.0.iter() {
            let token = Token(e.udata as usize);
            let len = self.events.len();

            if token == awakener {
                // TODO: Should this return an error if event is an error. It
                // is not critical as spurious wakeups are permitted.
                ret = true;
                continue;
            }

            let idx = *self.event_map.entry(token)
                .or_insert(len);

            if idx == len {
                // New entry, insert the default
                self.events.push(Event::new(Ready::empty(), token));

            }

            if e.flags & libc::EV_ERROR != 0 {
                event::kind_mut(&mut self.events[idx]).insert(*UnixReady::error());
            }

            if e.filter == libc::EVFILT_READ as Filter {
                event::kind_mut(&mut self.events[idx]).insert(Ready::readable());
            } else if e.filter == libc::EVFILT_WRITE as Filter {
                event::kind_mut(&mut self.events[idx]).insert(Ready::writable());

                // `EV_EOF` set with `EVFILT_WRITE` indicates the connection has been fully
                // disconnected (read and write), but only sockets...
                if e.flags & libc::EV_EOF != 0 {
                    event::kind_mut(&mut self.events[idx]).insert(UnixReady::hup());

                    // When the read end of the socket is closed, EV_EOF is set on
                    // flags, and fflags contains the error if there is one.
                    if e.fflags != 0 {
                        event::kind_mut(&mut self.events[idx]).insert(UnixReady::error());
                    }
                }
            }
#[cfg(any(target_os = "dragonfly",
    target_os = "freebsd", target_os = "ios", target_os = "macos"))]
            {
                if e.filter == libc::EVFILT_AIO {
                    event::kind_mut(&mut self.events[idx]).insert(UnixReady::aio());
                }
            }
#[cfg(any(target_os = "freebsd"))]
            {
                if e.filter == libc::EVFILT_LIO {
                    event::kind_mut(&mut self.events[idx]).insert(UnixReady::lio());
                }
            }
        }

        ret
    }

    pub fn push_event(&mut self, event: Event) {
        self.events.push(event);
    }

    pub fn clear(&mut self) {
        self.sys_events.0.truncate(0);
        self.events.truncate(0);
        self.event_map.clear();
    }
}

impl fmt::Debug for Events {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Events")
            .field("len", &self.sys_events.0.len())
            .finish()
    }
}

#[test]
fn does_not_register_rw() {
    use {Poll, Ready, PollOpt, Token};
    use unix::EventedFd;

    let kq = unsafe { libc::kqueue() };
    let kqf = EventedFd(&kq);
    let poll = Poll::new().unwrap();

    // registering kqueue fd will fail if write is requested (On anything but some versions of OS
    // X)
    poll.register(&kqf, Token(1234), Ready::readable(),
                  PollOpt::edge() | PollOpt::oneshot()).unwrap();
}

#[cfg(any(target_os = "dragonfly",
    target_os = "freebsd", target_os = "ios", target_os = "macos"))]
#[test]
fn test_coalesce_aio() {
    let mut events = Events::with_capacity(1);
    events.sys_events.0.push(kevent!(0x1234, libc::EVFILT_AIO, 0, 42));
    events.coalesce(Token(0));
    assert!(events.events[0].readiness() == UnixReady::aio().into());
    assert!(events.events[0].token() == Token(42));
}
<|endoftext|>


<commit_before>#![crate_type = "dylib"]
#![feature(plugin_registrar, quote)]
#![feature(collections, rustc_private)]

extern crate syntax;
extern crate rustc;
extern crate maud;

use syntax::ast::TokenTree;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, MacExpr, MacResult};
use rustc::plugin::Registry;

mod parse;
mod render;

fn expand_html<'cx>(cx: &'cx mut ExtCtxt, sp: Span, args: &[TokenTree]) -> Box<MacResult + 'cx> {
    MacExpr::new(parse::parse(cx, args, sp))
}

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("html", expand_html);
}
<commit_msg>Update Rust<commit_after>#![crate_type = "dylib"]
#![feature(plugin_registrar, quote)]
#![feature(collections, rustc_private)]

extern crate syntax;
extern crate rustc;
extern crate maud;

use syntax::ast::TokenTree;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, MacEager, MacResult};
use rustc::plugin::Registry;

mod parse;
mod render;

fn expand_html<'cx>(cx: &'cx mut ExtCtxt, sp: Span, args: &[TokenTree]) -> Box<MacResult + 'cx> {
    MacEager::expr(parse::parse(cx, args, sp))
}

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("html", expand_html);
}
<|endoftext|>


<commit_before>extern crate lumol_tutorial_potential;
use lumol_tutorial_potential::Mie;

extern crate lumol;
extern crate lumol_input as input;
use input::Input;
use lumol::energy::PairInteraction;
use lumol::units;

// This function runs a small Monte-Carlo simulation for Argon using a Mie potential.
fn main() {
    // read configuration from input file located at the data/ folder
    let config = Input::new("data/argon.toml").unwrap().read().unwrap();
    let mut system = config.system;
    let mut simulation = config.simulation;

    // build a potential
    let mie = Mie::new(units::from(3.405, "A").unwrap(),
                       units::from(1.0, "kJ/mol").unwrap(),
                       12.0,
                       6.0);
    // use the potential with a cut off radius of 3 * sigma
    let mut pairs = PairInteraction::new(Box::new(mie), units::from(3.0 * 3.405, "A").unwrap());
    // use tail corrections to account for our truncation
    pairs.enable_tail_corrections();
    // finally add this interaction to the system for Argon atoms
    system.add_pair_potential("Ar", "Ar", pairs);

    // report the initial system energy
    let ener_init = units::to(system.total_energy(), "kJ/mol").unwrap();
    println!("initial energy     : {}", ener_init);

    // run the simulation
    simulation.run(&mut system, config.nsteps);

    // print some final information
    println!("final energy       : {}",
             units::to(system.total_energy(), "kJ/mol").unwrap());
    println!("It worked. Hooray!")
}
<commit_msg>Fix tutorial<commit_after>extern crate lumol_tutorial_potential;
use lumol_tutorial_potential::Mie;

extern crate lumol;
extern crate lumol_input as input;
use input::Input;
use lumol::energy::PairInteraction;
use lumol::units;

// This function runs a small Monte-Carlo simulation for Argon using a Mie potential.
fn main() {
    // read configuration from input file located at the data/ folder
    let config = Input::new("data/argon.toml").unwrap().read().unwrap();
    let mut system = config.system;
    let mut simulation = config.simulation;

    // build a potential
    let mie = Mie::new(units::from(3.405, "A").unwrap(),
                       units::from(1.0, "kJ/mol").unwrap(),
                       12.0,
                       6.0);
    // use the potential with a cut off radius of 3 * sigma
    let mut pairs = PairInteraction::new(Box::new(mie), units::from(3.0 * 3.405, "A").unwrap());
    // use tail corrections to account for our truncation
    pairs.enable_tail_corrections();
    // finally add this interaction to the system for Argon atoms
    system.add_pair_potential(("Ar", "Ar"), pairs);

    // report the initial system energy
    let ener_init = units::to(system.total_energy(), "kJ/mol").unwrap();
    println!("initial energy     : {}", ener_init);

    // run the simulation
    simulation.run(&mut system, config.nsteps);

    // print some final information
    println!("final energy       : {}",
             units::to(system.total_energy(), "kJ/mol").unwrap());
    println!("It worked. Hooray!")
}
<|endoftext|>


<commit_before>use super::start_server;
use graph;
use graph::*;
use graph::edge::*;
use graph::vertex::*;
use server::schema::{MorpheusSchema, SchemaError, EMPTY_FIELDS};
use neb::ram::schema::Field;
use neb::ram::types::{TypeId, Value, Map};
use neb::ram::cell::Cell;

#[test]
pub fn schemas() {
    let server = start_server(4001, "schemas");
    let graph = &server.graph;
    let mut edge_schema = MorpheusSchema::new(
        "test_edge_schema",
        None,
        &vec![Field::new(&"test_field", TypeId::U32 as u32, false, false, None)], false
    );
    assert_eq!(edge_schema.id, 0);
    assert!(
    graph.new_edge_group(
        &mut edge_schema,
        graph::edge::EdgeAttributes::new
            (
                graph::edge::EdgeType::Directed,
                false
            )).is_err());
    assert_eq!(edge_schema.id, 0);
    graph.new_edge_group(
        &mut edge_schema,
        graph::edge::EdgeAttributes::new(graph::edge::EdgeType::Directed, true)).unwrap();
    let mut vertex_schema = edge_schema.clone();
    vertex_schema.name = "test_vertex_schema".to_string();
    graph.new_vertex_group(&mut vertex_schema).unwrap();
    assert_eq!(edge_schema.id, 1);
    assert_eq!(vertex_schema.id, 2);
    let mut test_data = Map::new();
    test_data.insert("test_field", Value::U32(1));
    graph.new_vertex(&vertex_schema, test_data.clone()).unwrap();
    graph.new_vertex("test_edge_schema", test_data.clone()).is_err();
}

#[test]
pub fn relationship() {
    let server = start_server(4002, "relationship");
    let graph = &server.graph;
    let mut people_schema = MorpheusSchema::new("people", Some(&vec!["name".to_string()]), &vec! [
        Field::new("name", TypeId::String as u32, false, false, None)
    ], true);
    let mut movie_schema = MorpheusSchema::new("movie", Some(&vec!["name".to_string()]), &vec! [
        Field::new("name", TypeId::String as u32, false, false, None),
        Field::new("year", TypeId::U32 as u32, false, false, None)
    ], true);
    let mut acted_in_schema = MorpheusSchema::new("acted-in", None, &vec! [
        Field::new("role", TypeId::String as u32, false, false, None)
    ], true);
    let mut spouse_schema = MorpheusSchema::new("spouse", None, &EMPTY_FIELDS, false);
    graph.new_vertex_group(&mut people_schema).unwrap();
    graph.new_vertex_group(&mut movie_schema).unwrap();
    graph.new_edge_group(
        &mut acted_in_schema,
        EdgeAttributes::new(
            EdgeType::Directed,
            true
        )
    ).unwrap();
    graph.new_edge_group(
        &mut spouse_schema,
        EdgeAttributes::new(
            EdgeType::Undirected,
            false
        )
    ).unwrap();
    assert_eq!(people_schema.id, 1);
    assert_eq!(movie_schema.id, 2);
    assert_eq!(acted_in_schema.id, 3);
    assert_eq!(spouse_schema.id, 4);
    let morgan_freeman_name = "Morgan Freeman";
    let batman_begins_name = "Batman Begins";
    let the_dark_knight_name = "The Dark Knight";
    let the_dark_knight_rises_name = "The Dark Knight Rises";
    let oblivion_name = "Oblivion";
    let jeanette_name = "Jeanette Adair Bradshaw";
    graph.new_vertex("people", data_map!{
        name: morgan_freeman_name, age: 80 as u8
    }).unwrap();
    graph.new_vertex("movie", data_map!{
        name: batman_begins_name, year: 2005 as u32
    }).unwrap();
    graph.new_vertex("movie", data_map!{
        name: the_dark_knight_name, year: 2008 as u32
    }).unwrap();
    graph.new_vertex("movie", data_map!{
        name: the_dark_knight_rises_name, year: 2012 as u32
    }).unwrap();
    graph.new_vertex("movie", data_map!{
        name: oblivion_name, year: 2010 as u32
    }).unwrap();
    graph.new_vertex("people", data_map!{
        name: jeanette_name
    }).unwrap();

    assert_eq!(
        graph.vertex_by_key("people", morgan_freeman_name)
            .unwrap().unwrap()["name"].String().unwrap(),
        morgan_freeman_name
    );
    assert_eq!(
        graph.vertex_by_key("movie", batman_begins_name)
            .unwrap().unwrap()["name"].String().unwrap(),
        batman_begins_name
    );
    assert_eq!(
        graph.vertex_by_key("movie", the_dark_knight_name)
            .unwrap().unwrap()["name"].String().unwrap(),
        the_dark_knight_name
    );
    assert_eq!(
        graph.vertex_by_key("movie", the_dark_knight_rises_name)
            .unwrap().unwrap()["name"].String().unwrap(),
        the_dark_knight_rises_name
    );
    assert_eq!(
        graph.vertex_by_key("movie", oblivion_name)
            .unwrap().unwrap()["name"].String().unwrap(),
        oblivion_name
    );
    assert_eq!(
        graph.vertex_by_key("people", jeanette_name)
            .unwrap().unwrap()["name"].String().unwrap(),
        jeanette_name
    );
    assert_eq!(
        graph.vertex_by_key("people", morgan_freeman_name)
            .unwrap().unwrap()["age"].U8().unwrap(),
        80u8
    );

    let morgan_freeman =
        graph.vertex_by_key("people", morgan_freeman_name)
        .unwrap().unwrap();

    let batman_begins =
        graph.vertex_by_key("movie", batman_begins_name)
        .unwrap().unwrap();

    let the_dark_knight =
        graph.vertex_by_key("movie", the_dark_knight_name)
            .unwrap().unwrap();

    let the_dark_knight_rises =
        graph.vertex_by_key("movie", the_dark_knight_rises_name)
        .unwrap().unwrap();

    let oblivion =
        graph.vertex_by_key("movie", oblivion_name)
            .unwrap().unwrap();

    graph.link(&morgan_freeman, "acted-in", &batman_begins, Some(&data_map!{
        role: "Lucius Fox", works_for: "Bruce Wayne"
    })).unwrap().unwrap();
    graph.link(&morgan_freeman, "acted-in", &the_dark_knight, Some(&data_map!{
        role: "Lucius Fox"
    })).unwrap().unwrap();
    graph.link(&morgan_freeman, "acted-in", &the_dark_knight_rises, Some(&data_map!{
        role: "Lucius Fox"
    })).unwrap().unwrap();
}
<commit_msg>detected false postive<commit_after>use super::start_server;
use graph;
use graph::*;
use graph::edge::*;
use graph::vertex::*;
use server::schema::{MorpheusSchema, SchemaError, EMPTY_FIELDS};
use neb::ram::schema::Field;
use neb::ram::types::{TypeId, Value, Map};
use neb::ram::cell::Cell;

#[test]
pub fn schemas() {
    let server = start_server(4001, "schemas");
    let graph = &server.graph;
    let mut edge_schema = MorpheusSchema::new(
        "test_edge_schema",
        None,
        &vec![Field::new(&"test_field", TypeId::U32 as u32, false, false, None)], false
    );
    assert_eq!(edge_schema.id, 0);
    assert!(
    graph.new_edge_group(
        &mut edge_schema,
        graph::edge::EdgeAttributes::new
            (
                graph::edge::EdgeType::Directed,
                false
            )).is_err());
    assert_eq!(edge_schema.id, 0);
    graph.new_edge_group(
        &mut edge_schema,
        graph::edge::EdgeAttributes::new(graph::edge::EdgeType::Directed, true)).unwrap();
    let mut vertex_schema = edge_schema.clone();
    vertex_schema.name = "test_vertex_schema".to_string();
    graph.new_vertex_group(&mut vertex_schema).unwrap();
    assert_eq!(edge_schema.id, 1);
    assert_eq!(vertex_schema.id, 2);
    let mut test_data = Map::new();
    test_data.insert("test_field", Value::U32(1));
    graph.new_vertex(&vertex_schema, test_data.clone()).unwrap();
    graph.new_vertex("test_edge_schema", test_data.clone()).is_err();
}

#[test]
pub fn relationship() {
    let server = start_server(4002, "relationship");
    let graph = &server.graph;
    let mut people_schema = MorpheusSchema::new("people", Some(&vec!["name".to_string()]), &vec! [
        Field::new("name", TypeId::String as u32, false, false, None)
    ], true);
    let mut movie_schema = MorpheusSchema::new("movie", Some(&vec!["name".to_string()]), &vec! [
        Field::new("name", TypeId::String as u32, false, false, None),
        Field::new("year", TypeId::U32 as u32, false, false, None)
    ], true);
    let mut acted_in_schema = MorpheusSchema::new("acted-in", None, &vec! [
        Field::new("role", TypeId::String as u32, false, false, None)
    ], true);
    let mut spouse_schema = MorpheusSchema::new("spouse", None, &EMPTY_FIELDS, false);
    graph.new_vertex_group(&mut people_schema).unwrap();
    graph.new_vertex_group(&mut movie_schema).unwrap();
    graph.new_edge_group(
        &mut acted_in_schema,
        EdgeAttributes::new(
            EdgeType::Directed,
            true
        )
    ).unwrap();
    graph.new_edge_group(
        &mut spouse_schema,
        EdgeAttributes::new(
            EdgeType::Undirected,
            false
        )
    ).unwrap();
    assert_eq!(people_schema.id, 1);
    assert_eq!(movie_schema.id, 2);
    assert_eq!(acted_in_schema.id, 3);
    assert_eq!(spouse_schema.id, 4);
    let morgan_freeman_name = "Morgan Freeman";
    let batman_begins_name = "Batman Begins";
    let the_dark_knight_name = "The Dark Knight";
    let the_dark_knight_rises_name = "The Dark Knight Rises";
    let oblivion_name = "Oblivion";
    let jeanette_name = "Jeanette Adair Bradshaw";
    graph.new_vertex("people", data_map!{
        name: morgan_freeman_name, age: 80 as u8
    }).unwrap();
    graph.new_vertex("movie", data_map!{
        name: batman_begins_name, year: 2005 as u32
    }).unwrap();
    graph.new_vertex("movie", data_map!{
        name: the_dark_knight_name, year: 2008 as u32
    }).unwrap();
    graph.new_vertex("movie", data_map!{
        name: the_dark_knight_rises_name, year: 2012 as u32
    }).unwrap();
    graph.new_vertex("movie", data_map!{
        name: oblivion_name, year: 2010 as u32
    }).unwrap();
    graph.new_vertex("people", data_map!{
        name: jeanette_name
    }).unwrap();

    assert_eq!(
        graph.vertex_by_key("people", morgan_freeman_name)
            .unwrap().unwrap()["name"].String().unwrap(),
        morgan_freeman_name
    );
    assert_eq!(
        graph.vertex_by_key("movie", batman_begins_name)
            .unwrap().unwrap()["name"].String().unwrap(),
        batman_begins_name
    );
    assert_eq!(
        graph.vertex_by_key("movie", the_dark_knight_name)
            .unwrap().unwrap()["name"].String().unwrap(),
        the_dark_knight_name
    );
    assert_eq!(
        graph.vertex_by_key("movie", the_dark_knight_rises_name)
            .unwrap().unwrap()["name"].String().unwrap(),
        the_dark_knight_rises_name
    );
    assert_eq!(
        graph.vertex_by_key("movie", oblivion_name)
            .unwrap().unwrap()["name"].String().unwrap(),
        oblivion_name
    );
    assert_eq!(
        graph.vertex_by_key("people", jeanette_name)
            .unwrap().unwrap()["name"].String().unwrap(),
        jeanette_name
    );
    assert_eq!(
        graph.vertex_by_key("people", morgan_freeman_name)
            .unwrap().unwrap()["age"].U8().unwrap(),
        80u8
    );

    let morgan_freeman =
        graph.vertex_by_key("people", morgan_freeman_name)
        .unwrap().unwrap();

    let batman_begins =
        graph.vertex_by_key("movie", batman_begins_name)
        .unwrap().unwrap();

    let the_dark_knight =
        graph.vertex_by_key("movie", the_dark_knight_name)
            .unwrap().unwrap();

    let the_dark_knight_rises =
        graph.vertex_by_key("movie", the_dark_knight_rises_name)
        .unwrap().unwrap();

    let oblivion =
        graph.vertex_by_key("movie", oblivion_name)
            .unwrap().unwrap();

    graph.link(&morgan_freeman, "acted-in", &batman_begins, Some(&data_map!{
        role: "Lucius Fox", works_for: "Bruce Wayne"
    })).unwrap().unwrap();
    graph.link(&morgan_freeman, "acted-in", &the_dark_knight, Some(&data_map!{
        role: "Lucius Fox"
    })).unwrap().unwrap();
    graph.link(&morgan_freeman, "acted-in", &the_dark_knight_rises, Some(&data_map!{
        role: "Lucius Fox"
    })).unwrap().unwrap();

    graph.link(&morgan_freeman, "acted-in", &oblivion, Some(&data_map!{
        // missing required field should fail
    })).unwrap().err().unwrap();

}
<|endoftext|>


<commit_before>//! SOCKSv5 protocol support

use std::comm::channel;
use std::io::{IoError, IoResult, Reader, TcpStream};
use std::io::net::ip::{Ipv4Addr, Ipv6Addr, SocketAddr};
use std::io::net::addrinfo;
use std::str;
use std::time::duration::Duration;

#[allow(dead_code)]
mod consts {
    pub static VERSION_NUMBER: u8 = 0x05;

    pub mod auth {
        pub static NONE: u8 = 0x00;
    }

    pub mod atype {
        pub static IPV4:   u8 = 0x1;
        pub static IPV6:   u8 = 0x4;
        pub static DOMAIN: u8 = 0x3;
    }

    pub mod command {
        pub static CONNECT:       u8 = 0x01;
        pub static BIND:          u8 = 0x02;
        pub static UDP_ASSOCIATE: u8 = 0x04;
    }

    pub mod reply {
        pub static SUCCESS:          u8 = 0x00;
        pub static GENERAL_FAILURE:  u8 = 0x01;
        pub static CONN_DENIED:      u8 = 0x02;
        pub static NET_UNREACHABLE:  u8 = 0x03;
        pub static HOST_UNREACHABLE: u8 = 0x04;
        pub static CONN_REFUSED:     u8 = 0x05;
        pub static TTL_EXPIRED:      u8 = 0x06;
        pub static CMD_UNSUPPORTED:  u8 = 0x07;
        pub static ADDR_UNSUPPORTED: u8 = 0x08;
    }
}

struct SocksConnection<'a> {
    stream: &'a mut TcpStream
}

impl <'a> SocksConnection<'a> {
    /// Make sure that the client sends the right version number at
    /// the start of the message, or fail
    fn read_client_version(&mut self) -> IoResult<()> {
        match try!(self.stream.read_byte()) {
            consts::VERSION_NUMBER => Ok(()),
            _ => Err(IoError::last_error())
        }
    }

    /// Client -> Server:
    ///
    /// |VER | NMETHODS | METHODS  |
    /// |----+----------+----------|
    /// | 1  |    1     | 1 to 255 |
    ///
    /// Client <- Server:
    /// +----+--------+
    /// |VER | METHOD |
    /// +----+--------+
    /// | 1  |   1    |
    fn read_client_hello(&mut self) -> IoResult<()> {
        try!(self.read_client_version());

        // Number of authentication methods
        let nmethods = try!(self.stream.read_byte());

        // Authentication methods
        let methods = try!(self.stream.read_exact(nmethods as uint));

        // If the client doesn't respond to NULL-auth, we can't proceed.
        if !methods.contains(&consts::auth::NONE) {
            // Send the error back
            try!(self.stream.write(&[consts::VERSION_NUMBER, 0xFF]));

            return Err(IoError::last_error());
        }

        // Now respond that we're using NULL auth.
        try!(self.stream.write(&[consts::VERSION_NUMBER, consts::auth::NONE]));

        Ok(())
    }

    /// Client -> Server:
    /// +----+-----+-------+------+----------+----------+
    /// |VER | CMD |  RSV  | ATYP | DST.ADDR | DST.PORT |
    /// +----+-----+-------+------+----------+----------+
    /// | 1  |  1  | X'00' |  1   | Variable |    2     |
    ///
    /// o  VER    protocol version: X'05'
    /// o  CMD
    /// o  RSV    RESERVED
    /// o  ATYP   address type of following address
    /// o  DST.ADDR desired destination address
    /// o  DST.PORT desired destination port in network octet
    ///    order
    fn read_client_request(&mut self) -> IoResult<TcpStream> {
        try!(self.read_client_version());

        // Commands (for now only connect is supported)
        match try!(self.stream.read_byte()) {
            consts::command::CONNECT => (),
            _ => {
                // Tell client we don't support that command
                try!(self.send_error_reply(consts::reply::CMD_UNSUPPORTED));
                return Err(IoError::last_error());
            }
        };

        // Reserved byte
        match try!(self.stream.read_byte()) {
            0u8 => (),
            _ => { return Err(IoError::last_error()); }
        };

        // Address type
        let addr = match try!(self.stream.read_byte()) {
            // 4 bytes
            consts::atype::IPV4 => {
                let a = try!(self.stream.read_byte());
                let b = try!(self.stream.read_byte());
                let c = try!(self.stream.read_byte());
                let d = try!(self.stream.read_byte());

                Ipv4Addr(a, b, c, d)
            },

            // 16 bytes
            consts::atype::IPV6 => {
                let a = try!(self.stream.read_be_u16());
                let b = try!(self.stream.read_be_u16());
                let c = try!(self.stream.read_be_u16());
                let d = try!(self.stream.read_be_u16());
                let e = try!(self.stream.read_be_u16());
                let f = try!(self.stream.read_be_u16());
                let g = try!(self.stream.read_be_u16());
                let h = try!(self.stream.read_be_u16());

                Ipv6Addr(a, b, c, d, e, f, g, h)
            },

            // len | hostname
            consts::atype::DOMAIN => {
                let len = try!(self.stream.read_byte());
                let hostname = try!(self.stream.read_exact(len as uint));

                let mut hosts = match str::from_utf8(hostname.as_slice()) {
                    Some(host) => try!(addrinfo::get_host_addresses(host)),
                    None => { return Err(IoError::last_error()); }
                };

                match hosts.remove(0) {
                    Some(host) => host,
                    None => { return Err(IoError::last_error()); }
                }
            },

            // Unsupported address type
            _ => {
                try!(self.send_error_reply(consts::reply::ADDR_UNSUPPORTED));
                return Err(IoError::last_error());
            }
        };

        // Network port
        let port = try!(self.stream.read_be_u16());

        Ok(try!(TcpStream::connect_timeout(
            SocketAddr { ip: addr, port: port },
            Duration::seconds(1000))))
    }

    /// Client <- Server:
    /// +----+-----+-------+------+----------+----------+
    /// |VER | REP |  RSV  | ATYP | BND.ADDR | BND.PORT |
    /// +----+-----+-------+------+----------+----------+
    /// | 1  |  1  | X'00' |  1   | Variable |    2     |
    ///
    ///   o  VER    protocol version: X'05'
    ///   o  REP    Reply field
    ///   o  RSV    RESERVED
    ///   o  ATYP   address type of following address
    ///   o  BND.ADDR       server bound address
    ///   o  BND.PORT       server bound port in network octet order
    fn send_reply(&mut self, rep: u8, atype: u8, addr: &[u8], port: u16) -> IoResult<()> {
        try!(self.stream.write(&[consts::VERSION_NUMBER, rep, 0x00, atype]));
        try!(self.stream.write(addr));
        try!(self.stream.write_be_u16(port));

        Ok(())
    }

    fn send_error_reply(&mut self, rep: u8) -> IoResult<()> {
        self.send_reply(rep, 0, &[0u8, 0, 0, 0], 0)
    }
}

pub fn handle_stream<'a>(client_stream: &'a mut TcpStream) -> () {
    let mut conn = SocksConnection { stream: client_stream };

    if conn.read_client_hello().is_err() {
        return;
    }

    // FIXME: A ton of probably redundant clone()s here.
    match conn.read_client_request() {
        Ok(server_stream) => {
            let (client_send, client_recv) = channel();
            let (server_send, server_recv) = channel();

            let client = conn.stream.clone();
            let server = server_stream.clone();

            // client -> server
            spawn(proc() {
                let mut client = client.clone();
                loop {
                    let mut buf = [0u8, ..4096];
                    match client.read(buf) {
                        Ok(_) => client_send.send(buf),
                        Err(_) => break
                    }
                }
            });

            // client <- server
            spawn(proc() {
                let mut server = server.clone();
                loop {
                    let mut buf = [0u8, ..4096];
                    match server.read(buf) {
                        Ok(_) => server_send.send(buf),
                        Err(_) => break
                    }
                }
            });

            loop {
                let mut client = conn.stream.clone();
                let mut server = server_stream.clone();

                let result = select! {
                    buf = client_recv.recv() => server.write(buf),
                    buf = server_recv.recv() => client.write(buf)
                };

                match result {
                    Err(e) => println!("Something broke: {}", e),
                    _ => ()
                }
            }
        },
        Err(_) => ()
}
    }
<commit_msg>Comment out ICE for now<commit_after>//! SOCKSv5 protocol support

use std::comm::channel;
use std::io::{IoError, IoResult, Reader, TcpStream};
use std::io::net::ip::{Ipv4Addr, Ipv6Addr, SocketAddr};
use std::io::net::addrinfo;
use std::str;
use std::time::duration::Duration;

#[allow(dead_code)]
mod consts {
    pub static VERSION_NUMBER: u8 = 0x05;

    pub mod auth {
        pub static NONE: u8 = 0x00;
    }

    pub mod atype {
        pub static IPV4:   u8 = 0x1;
        pub static IPV6:   u8 = 0x4;
        pub static DOMAIN: u8 = 0x3;
    }

    pub mod command {
        pub static CONNECT:       u8 = 0x01;
        pub static BIND:          u8 = 0x02;
        pub static UDP_ASSOCIATE: u8 = 0x04;
    }

    pub mod reply {
        pub static SUCCESS:          u8 = 0x00;
        pub static GENERAL_FAILURE:  u8 = 0x01;
        pub static CONN_DENIED:      u8 = 0x02;
        pub static NET_UNREACHABLE:  u8 = 0x03;
        pub static HOST_UNREACHABLE: u8 = 0x04;
        pub static CONN_REFUSED:     u8 = 0x05;
        pub static TTL_EXPIRED:      u8 = 0x06;
        pub static CMD_UNSUPPORTED:  u8 = 0x07;
        pub static ADDR_UNSUPPORTED: u8 = 0x08;
    }
}

struct SocksConnection<'a> {
    stream: &'a mut TcpStream
}

impl <'a> SocksConnection<'a> {
    /// Make sure that the client sends the right version number at
    /// the start of the message, or fail
    fn read_client_version(&mut self) -> IoResult<()> {
        match try!(self.stream.read_byte()) {
            consts::VERSION_NUMBER => Ok(()),
            _ => Err(IoError::last_error())
        }
    }

    /// Client -> Server:
    ///
    /// |VER | NMETHODS | METHODS  |
    /// |----+----------+----------|
    /// | 1  |    1     | 1 to 255 |
    ///
    /// Client <- Server:
    /// +----+--------+
    /// |VER | METHOD |
    /// +----+--------+
    /// | 1  |   1    |
    fn read_client_hello(&mut self) -> IoResult<()> {
        try!(self.read_client_version());

        // Number of authentication methods
        let nmethods = try!(self.stream.read_byte());

        // Authentication methods
        let methods = try!(self.stream.read_exact(nmethods as uint));

        // If the client doesn't respond to NULL-auth, we can't proceed.
        if !methods.contains(&consts::auth::NONE) {
            // Send the error back
            try!(self.stream.write(&[consts::VERSION_NUMBER, 0xFF]));

            return Err(IoError::last_error());
        }

        // Now respond that we're using NULL auth.
        try!(self.stream.write(&[consts::VERSION_NUMBER, consts::auth::NONE]));

        Ok(())
    }

    /// Client -> Server:
    /// +----+-----+-------+------+----------+----------+
    /// |VER | CMD |  RSV  | ATYP | DST.ADDR | DST.PORT |
    /// +----+-----+-------+------+----------+----------+
    /// | 1  |  1  | X'00' |  1   | Variable |    2     |
    ///
    /// o  VER    protocol version: X'05'
    /// o  CMD
    /// o  RSV    RESERVED
    /// o  ATYP   address type of following address
    /// o  DST.ADDR desired destination address
    /// o  DST.PORT desired destination port in network octet
    ///    order
    fn read_client_request(&mut self) -> IoResult<TcpStream> {
        try!(self.read_client_version());

        // Commands (for now only connect is supported)
        match try!(self.stream.read_byte()) {
            consts::command::CONNECT => (),
            _ => {
                // Tell client we don't support that command
                try!(self.send_error_reply(consts::reply::CMD_UNSUPPORTED));
                return Err(IoError::last_error());
            }
        };

        // Reserved byte
        match try!(self.stream.read_byte()) {
            0u8 => (),
            _ => { return Err(IoError::last_error()); }
        };

        // Address type
        let addr = match try!(self.stream.read_byte()) {
            // 4 bytes
            consts::atype::IPV4 => {
                let a = try!(self.stream.read_byte());
                let b = try!(self.stream.read_byte());
                let c = try!(self.stream.read_byte());
                let d = try!(self.stream.read_byte());

                Ipv4Addr(a, b, c, d)
            },

            // 16 bytes
            consts::atype::IPV6 => {
                let a = try!(self.stream.read_be_u16());
                let b = try!(self.stream.read_be_u16());
                let c = try!(self.stream.read_be_u16());
                let d = try!(self.stream.read_be_u16());
                let e = try!(self.stream.read_be_u16());
                let f = try!(self.stream.read_be_u16());
                let g = try!(self.stream.read_be_u16());
                let h = try!(self.stream.read_be_u16());

                Ipv6Addr(a, b, c, d, e, f, g, h)
            },

            // len | hostname
            consts::atype::DOMAIN => {
                let len = try!(self.stream.read_byte());
                let hostname = try!(self.stream.read_exact(len as uint));

                let mut hosts = match str::from_utf8(hostname.as_slice()) {
                    Some(host) => try!(addrinfo::get_host_addresses(host)),
                    None => { return Err(IoError::last_error()); }
                };

                match hosts.remove(0) {
                    Some(host) => host,
                    None => { return Err(IoError::last_error()); }
                }
            },

            // Unsupported address type
            _ => {
                try!(self.send_error_reply(consts::reply::ADDR_UNSUPPORTED));
                return Err(IoError::last_error());
            }
        };

        // Network port
        let port = try!(self.stream.read_be_u16());

        Ok(try!(TcpStream::connect_timeout(
            SocketAddr { ip: addr, port: port },
            Duration::seconds(1000))))
    }

    /// Client <- Server:
    /// +----+-----+-------+------+----------+----------+
    /// |VER | REP |  RSV  | ATYP | BND.ADDR | BND.PORT |
    /// +----+-----+-------+------+----------+----------+
    /// | 1  |  1  | X'00' |  1   | Variable |    2     |
    ///
    ///   o  VER    protocol version: X'05'
    ///   o  REP    Reply field
    ///   o  RSV    RESERVED
    ///   o  ATYP   address type of following address
    ///   o  BND.ADDR       server bound address
    ///   o  BND.PORT       server bound port in network octet order
    fn send_reply(&mut self, rep: u8, atype: u8, addr: &[u8], port: u16) -> IoResult<()> {
        try!(self.stream.write(&[consts::VERSION_NUMBER, rep, 0x00, atype]));
        try!(self.stream.write(addr));
        try!(self.stream.write_be_u16(port));

        Ok(())
    }

    fn send_error_reply(&mut self, rep: u8) -> IoResult<()> {
        self.send_reply(rep, 0, &[0u8, 0, 0, 0], 0)
    }
}

pub fn handle_stream<'a>(client_stream: &'a mut TcpStream) -> () {
    let mut conn = SocksConnection { stream: client_stream };

    if conn.read_client_hello().is_err() {
        return;
    }

    // FIXME: A ton of probably redundant clone()s here.
    // FIXME: This causes an ICE:
    // match conn.read_client_request() {
    //     Ok(server_stream) => {
    //         let (client_send, client_recv) = channel();
    //         let (server_send, server_recv) = channel();

    //         let client = conn.stream.clone();
    //         let server = server_stream.clone();

    //         // client -> server
    //         spawn(proc() {
    //             let mut client = client.clone();
    //             loop {
    //                 let mut buf = [0u8, ..4096];
    //                 match client.read(buf) {
    //                     Ok(_) => client_send.send(buf),
    //                     Err(_) => break
    //                 }
    //             }
    //         });

    //         // client <- server
    //         spawn(proc() {
    //             let mut server = server.clone();
    //             loop {
    //                 let mut buf = [0u8, ..4096];
    //                 match server.read(buf) {
    //                     Ok(_) => server_send.send(buf),
    //                     Err(_) => break
    //                 }
    //             }
    //         });

    //         loop {
    //             let mut client = conn.stream.clone();
    //             let mut server = server_stream.clone();

    //             let result = select! {
    //                 buf = client_recv.recv() => server.write(buf),
    //                 buf = server_recv.recv() => client.write(buf)
    //             };

    //             match result {
    //                 Err(e) => println!("Something broke: {}", e),
    //                 _ => ()
    //             }
    //         }
    //     },
    //     Err(_) => ()
    // }
}
<|endoftext|>