<commit_before>pub trait GerberCode {
    fn to_code(&self) -> String;
}

#[derive(Debug)]
pub enum Command {
    FunctionCode(FunctionCode),
    ExtendedCode(ExtendedCode),
}

#[derive(Debug)]
pub enum FunctionCode {
    DCode(DCode),
    GCode(GCode),
    MCode(MCode),
}

#[derive(Debug)]
pub enum ExtendedCode {
    FS,
    MO,
    AD,
    AM,
    SR,
    LP,
    TF,
    TA,
    TD,
}

#[derive(Debug)]
pub enum DCode {
    Operation,
    SelectAperture,
}

#[derive(Debug)]
pub enum InterpolationMode {
    Linear,
    ClockwiseCircular,
    CounterclockwiseCircular,
}

#[derive(Debug)]
pub enum GCode {
    InterpolationMode(InterpolationMode),
    RegionMode(bool),
    QuadrantMode,
    Comment(String),
}

#[derive(Debug)]
pub enum MCode {
    EndOfFile,
}

impl GerberCode for GCode {
    fn to_code(&self) -> String {
        match self {
            &GCode::InterpolationMode(ref mode) => {
                match *mode {
                    InterpolationMode::Linear => "G01*".to_string(),
                    InterpolationMode::ClockwiseCircular => "G02*".to_string(),
                    InterpolationMode::CounterclockwiseCircular => "G03*".to_string(),
                }
            },
            &GCode::RegionMode(enabled) => {
                match enabled {
                    true => "G36*",
                    false => "G37*",
                }.to_string()
            },
            &GCode::QuadrantMode => format!("TODO QuadrantMode"),
            &GCode::Comment(ref comment) => format!("G04 {} *", comment),
        }
    }
}

impl<T: GerberCode> GerberCode for Vec<T> {
    fn to_code(&self) -> String {
        self.iter()
            .map(|g| g.to_code())
            .collect::<Vec<String>>()
            .join("\n")
    }
}

#[cfg(test)]
mod test {
    use super::{Command, FunctionCode, DCode, GCode, InterpolationMode};
    use super::GerberCode;

    #[test]
    fn test_gcode_debug() {
        let c = Command::FunctionCode(FunctionCode::GCode(GCode::Comment("test".to_string())));
        let debug = format!("{:?}", c);
        assert_eq!(debug, "FunctionCode(GCode(Comment(\"test\")))");
    }

    #[test]
    fn test_to_code() {
        let comment = GCode::Comment("testcomment".to_string());
        assert_eq!(comment.to_code(), "G04 testcomment *".to_string());
    }

    #[test]
    fn test_vec_to_code() {
        let mut v = Vec::new();
        v.push(GCode::Comment("comment 1".to_string()));
        v.push(GCode::Comment("another one".to_string()));
        assert_eq!(v.to_code(), "G04 comment 1 *\nG04 another one *".to_string());
    }

    #[test]
    fn test_interpolation_mode() {
        let mut commands = Vec::new();
        let c1 = GCode::InterpolationMode(InterpolationMode::Linear);
        let c2 = GCode::InterpolationMode(InterpolationMode::ClockwiseCircular);
        let c3 = GCode::InterpolationMode(InterpolationMode::CounterclockwiseCircular);
        commands.push(c1);
        commands.push(c2);
        commands.push(c3);
        assert_eq!(commands.to_code(), "G01*\nG02*\nG03*".to_string());
    }

    #[test]
    fn test_region_mode() {
        let mut commands = Vec::new();
        commands.push(GCode::RegionMode(true));
        commands.push(GCode::RegionMode(false));
        assert_eq!(commands.to_code(), "G36*\nG37*".to_string());
    }

}
<commit_msg>Implement quadrant mode<commit_after>pub trait GerberCode {
    fn to_code(&self) -> String;
}

#[derive(Debug)]
pub enum Command {
    FunctionCode(FunctionCode),
    ExtendedCode(ExtendedCode),
}

#[derive(Debug)]
pub enum FunctionCode {
    DCode(DCode),
    GCode(GCode),
    MCode(MCode),
}

#[derive(Debug)]
pub enum ExtendedCode {
    FS,
    MO,
    AD,
    AM,
    SR,
    LP,
    TF,
    TA,
    TD,
}

#[derive(Debug)]
pub enum DCode {
    Operation,
    SelectAperture,
}

#[derive(Debug)]
pub enum InterpolationMode {
    Linear,
    ClockwiseCircular,
    CounterclockwiseCircular,
}

#[derive(Debug)]
pub enum QuadrantMode {
    Single,
    Multi,
}

#[derive(Debug)]
pub enum GCode {
    InterpolationMode(InterpolationMode),
    RegionMode(bool),
    QuadrantMode(QuadrantMode),
    Comment(String),
}

#[derive(Debug)]
pub enum MCode {
    EndOfFile,
}

impl GerberCode for GCode {
    fn to_code(&self) -> String {
        match self {
            &GCode::InterpolationMode(ref mode) => {
                match *mode {
                    InterpolationMode::Linear => "G01*",
                    InterpolationMode::ClockwiseCircular => "G02*",
                    InterpolationMode::CounterclockwiseCircular => "G03*",
                }.to_string()
            },
            &GCode::RegionMode(enabled) => {
                match enabled {
                    true => "G36*",
                    false => "G37*",
                }.to_string()
            },
            &GCode::QuadrantMode(ref mode) => {
                match *mode {
                    QuadrantMode::Single => "G74*",
                    QuadrantMode::Multi => "G75*",
                }.to_string()
            },
            &GCode::Comment(ref comment) => format!("G04 {} *", comment),
        }
    }
}

impl<T: GerberCode> GerberCode for Vec<T> {
    fn to_code(&self) -> String {
        self.iter()
            .map(|g| g.to_code())
            .collect::<Vec<String>>()
            .join("\n")
    }
}

#[cfg(test)]
mod test {
    use super::{Command, FunctionCode, DCode};
    use super::{GCode, InterpolationMode, QuadrantMode};
    use super::GerberCode;

    #[test]
    fn test_gcode_debug() {
        let c = Command::FunctionCode(FunctionCode::GCode(GCode::Comment("test".to_string())));
        let debug = format!("{:?}", c);
        assert_eq!(debug, "FunctionCode(GCode(Comment(\"test\")))");
    }

    #[test]
    fn test_to_code() {
        let comment = GCode::Comment("testcomment".to_string());
        assert_eq!(comment.to_code(), "G04 testcomment *".to_string());
    }

    #[test]
    fn test_vec_to_code() {
        let mut v = Vec::new();
        v.push(GCode::Comment("comment 1".to_string()));
        v.push(GCode::Comment("another one".to_string()));
        assert_eq!(v.to_code(), "G04 comment 1 *\nG04 another one *".to_string());
    }

    #[test]
    fn test_interpolation_mode() {
        let mut commands = Vec::new();
        let c1 = GCode::InterpolationMode(InterpolationMode::Linear);
        let c2 = GCode::InterpolationMode(InterpolationMode::ClockwiseCircular);
        let c3 = GCode::InterpolationMode(InterpolationMode::CounterclockwiseCircular);
        commands.push(c1);
        commands.push(c2);
        commands.push(c3);
        assert_eq!(commands.to_code(), "G01*\nG02*\nG03*".to_string());
    }

    #[test]
    fn test_region_mode() {
        let mut commands = Vec::new();
        commands.push(GCode::RegionMode(true));
        commands.push(GCode::RegionMode(false));
        assert_eq!(commands.to_code(), "G36*\nG37*".to_string());
    }

    #[test]
    fn test_quadrant_mode() {
        let mut commands = Vec::new();
        commands.push(GCode::QuadrantMode(QuadrantMode::Single));
        commands.push(GCode::QuadrantMode(QuadrantMode::Multi));
        assert_eq!(commands.to_code(), "G74*\nG75*".to_string());
    }

}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io
#![allow(
    clippy::unused_unit,
    clippy::needless_return,
    clippy::never_loop,
    clippy::option_map_unit_fn
)]
use super::osm_utils::get_way_coord;
use super::OsmPbfReader;
use crate::admin_geofinder::AdminGeoFinder;
use crate::{labels, settings, utils, Error};
use failure::ResultExt;
use osmpbfreader::StoreObjs;
use slog_scope::info;
use std::collections::BTreeSet;
use std::ops::Deref;
use std::path::PathBuf;
use std::sync::Arc;

use super::osm_store::{
    Getter, NameAdminMap, ObjWrapper, StreetKey, StreetWithRelationSet, StreetsVec,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
#[allow(dead_code)]
pub enum Kind {
    Node = 0,
    Way = 1,
    Relation = 2,
}

pub fn streets(
    pbf: &mut OsmPbfReader,
    admins_geofinder: &AdminGeoFinder,
    db_file: &Option<PathBuf>,
    db_buffer_size: usize,
    settings: &settings::Settings,
) -> Result<StreetsVec, Error> {
    let invalid_highways = &settings
        .street
        .clone()
        .map(|street| street.exclusion.highways.unwrap_or(Vec::new()))
        .unwrap_or(Vec::new());

    let is_valid_highway = |tag: &str| -> bool { !invalid_highways.iter().any(|k| k == tag) };

    // For the object to be a valid street, it needs to be an osm highway of a valid type,
    // or a relation of type associatedStreet.
    let is_valid_obj = |obj: &osmpbfreader::OsmObj| -> bool {
        match *obj {
            osmpbfreader::OsmObj::Way(ref way) => {
                way.tags
                    .get("highway")
                    .map_or(false, |v| !v.is_empty() && is_valid_highway(v))
                    && way.tags.get("name").map_or(false, |v| !v.is_empty())
            }
            osmpbfreader::OsmObj::Relation(ref rel) => rel
                .tags
                .get("type")
                .map_or(false, |v| v == "associatedStreet"),
            _ => false,
        }
    };

    info!("reading pbf...");
    let mut objs_map = ObjWrapper::new(db_file, db_buffer_size)?;
    pbf.get_objs_and_deps_store(is_valid_obj, &mut objs_map)
        .context("Error occurred when reading pbf")?;
    info!("reading pbf done.");
    let mut street_rel: StreetWithRelationSet = BTreeSet::new();
    let mut street_list: StreetsVec = vec![];
    // Sometimes, streets can be divided into several "way"s that still have the same street name.
    // The reason why a street is divided may be that a part of the street become
    // a bridge/tunnel/etc. In this case, a "relation" tagged with (type = associatedStreet) is used
    // to group all these "way"s. In order not to have duplicates in autocompletion, we should tag
    // the osm ways in the relation not to index them twice.

    objs_map.for_each_filter(Kind::Relation, |obj| {
        let rel = obj.relation().expect("impossible unwrap failure occured");
        let way_name = rel.tags.get("name");
        rel.refs
            .iter()
            .filter(|ref_obj| ref_obj.member.is_way() && ref_obj.role == "street")
            .filter_map(|ref_obj| {
                let obj = objs_map.get(&ref_obj.member)?;
                let way = obj.way()?;
                let way_name = way_name.or_else(|| way.tags.get("name"))?;
                let admins = get_street_admin(admins_geofinder, &objs_map, way);
                let country_codes = utils::find_country_codes(admins.iter().map(|a| a.deref()));
                let street_label = labels::format_street_label(
                    &way_name,
                    admins.iter().map(|a| a.deref()),
                    &country_codes,
                );
                let coord = get_way_coord(&objs_map, way);
                Some(mimir::Street {
                    id: format!("street:osm:relation:{}", rel.id.0.to_string()),
                    name: way_name.to_string(),
                    label: street_label,
                    weight: 0.,
                    zip_codes: utils::get_zip_codes_from_admins(&admins),
                    administrative_regions: admins,
                    coord: get_way_coord(&objs_map, way),
                    approx_coord: Some(coord.into()),
                    distance: None,
                    country_codes,
                    context: None,
                })
            })
            .next()
            .map(|street| street_list.push(street));

        // Add osmid of all the relation members in the set
        // We don't create any street for all the osmid present in street_rel
        for ref_obj in &rel.refs {
            if ref_obj.member.is_way() {
                street_rel.insert(ref_obj.member);
            }
        }
    });

    // we merge all the ways with a key = way_name + admin list of level(=city_level)
    // we use a map NameAdminMap <key, value> to manage the merging of ways
    let mut name_admin_map = NameAdminMap::default();
    objs_map.for_each_filter(Kind::Way, |obj| {
        let osmid = obj.id();
        if street_rel.contains(&osmid) {
            return;
        }
        if let Some(way) = obj.way() {
            if let Some(name) = way.tags.get("name") {
                let name = name.to_string();
                let admins = get_street_admin(admins_geofinder, &objs_map, way)
                    .into_iter()
                    .filter(|admin| admin.is_city())
                    .collect();
                name_admin_map
                    .entry(StreetKey { name, admins })
                    .or_insert_with(Vec::new)
                    .push(osmid);
            }
        }
    });

    // Create a street for each way with osmid present in objs_map
    let streets = name_admin_map.values().filter_map(|way_ids| {
        let min_id = way_ids.iter().min()?;
        let obj = objs_map.get(&min_id)?;
        let way = obj.way()?;
        let name = way.tags.get("name")?.to_string();
        let admins = get_street_admin(admins_geofinder, &objs_map, way);

        let country_codes = utils::find_country_codes(admins.iter().map(|a| a.deref()));
        let street_label =
            labels::format_street_label(&name, admins.iter().map(|a| a.deref()), &country_codes);
        let coord = get_way_coord(&objs_map, way);
        Some(mimir::Street {
            id: format!("street:osm:way:{}", way.id.0.to_string()),
            label: street_label,
            name,
            weight: 0.,
            zip_codes: utils::get_zip_codes_from_admins(&admins),
            administrative_regions: admins,
            coord: get_way_coord(&objs_map, way),
            approx_coord: Some(coord.into()),
            distance: None,
            country_codes,
            context: None,
        })
    });
    street_list.extend(streets);

    Ok(street_list)
}

fn get_street_admin<T: StoreObjs + Getter>(
    admins_geofinder: &AdminGeoFinder,
    obj_map: &T,
    way: &osmpbfreader::objects::Way,
) -> Vec<Arc<mimir::Admin>> {
    /*
        To avoid corner cases where the ends of the way are near
        administrative boundaries, the geofinder is called
        on a middle node.
    */
    let nb_nodes = way.nodes.len();
    way.nodes
        .iter()
        .skip(nb_nodes / 2)
        .filter_map(|node_id| obj_map.get(&(*node_id).into()))
        .filter_map(|node_obj| {
            node_obj.node().map(|node| geo_types::Coordinate {
                x: node.lon(),
                y: node.lat(),
            })
        })
        .next()
        .map_or(vec![], |c| admins_geofinder.get(&c))
}

pub fn compute_street_weight(streets: &mut StreetsVec) {
    for st in streets {
        for admin in &mut st.administrative_regions {
            if admin.is_city() {
                st.weight = admin.weight;
                break;
            }
        }
    }
}
<commit_msg>[FIX]: Use unwrap_or_else / unwrap_or<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io
#![allow(
    clippy::unused_unit,
    clippy::needless_return,
    clippy::never_loop,
    clippy::option_map_unit_fn
)]
use super::osm_utils::get_way_coord;
use super::OsmPbfReader;
use crate::admin_geofinder::AdminGeoFinder;
use crate::{labels, settings, utils, Error};
use failure::ResultExt;
use osmpbfreader::StoreObjs;
use slog_scope::info;
use std::collections::BTreeSet;
use std::ops::Deref;
use std::path::PathBuf;
use std::sync::Arc;

use super::osm_store::{
    Getter, NameAdminMap, ObjWrapper, StreetKey, StreetWithRelationSet, StreetsVec,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
#[allow(dead_code)]
pub enum Kind {
    Node = 0,
    Way = 1,
    Relation = 2,
}

pub fn streets(
    pbf: &mut OsmPbfReader,
    admins_geofinder: &AdminGeoFinder,
    db_file: &Option<PathBuf>,
    db_buffer_size: usize,
    settings: &settings::Settings,
) -> Result<StreetsVec, Error> {
    let invalid_highways = &settings
        .street
        .clone()
        .map(|street| street.exclusion.highways.unwrap_or_else(Vec::new))
        .unwrap_or(Vec::new());

    let is_valid_highway = |tag: &str| -> bool { !invalid_highways.iter().any(|k| k == tag) };

    // For the object to be a valid street, it needs to be an osm highway of a valid type,
    // or a relation of type associatedStreet.
    let is_valid_obj = |obj: &osmpbfreader::OsmObj| -> bool {
        match *obj {
            osmpbfreader::OsmObj::Way(ref way) => {
                way.tags
                    .get("highway")
                    .map_or(false, |v| !v.is_empty() && is_valid_highway(v))
                    && way.tags.get("name").map_or(false, |v| !v.is_empty())
            }
            osmpbfreader::OsmObj::Relation(ref rel) => rel
                .tags
                .get("type")
                .map_or(false, |v| v == "associatedStreet"),
            _ => false,
        }
    };

    info!("reading pbf...");
    let mut objs_map = ObjWrapper::new(db_file, db_buffer_size)?;
    pbf.get_objs_and_deps_store(is_valid_obj, &mut objs_map)
        .context("Error occurred when reading pbf")?;
    info!("reading pbf done.");
    let mut street_rel: StreetWithRelationSet = BTreeSet::new();
    let mut street_list: StreetsVec = vec![];
    // Sometimes, streets can be divided into several "way"s that still have the same street name.
    // The reason why a street is divided may be that a part of the street become
    // a bridge/tunnel/etc. In this case, a "relation" tagged with (type = associatedStreet) is used
    // to group all these "way"s. In order not to have duplicates in autocompletion, we should tag
    // the osm ways in the relation not to index them twice.

    objs_map.for_each_filter(Kind::Relation, |obj| {
        let rel = obj.relation().expect("impossible unwrap failure occured");
        let way_name = rel.tags.get("name");
        rel.refs
            .iter()
            .filter(|ref_obj| ref_obj.member.is_way() && ref_obj.role == "street")
            .filter_map(|ref_obj| {
                let obj = objs_map.get(&ref_obj.member)?;
                let way = obj.way()?;
                let way_name = way_name.or_else(|| way.tags.get("name"))?;
                let admins = get_street_admin(admins_geofinder, &objs_map, way);
                let country_codes = utils::find_country_codes(admins.iter().map(|a| a.deref()));
                let street_label = labels::format_street_label(
                    &way_name,
                    admins.iter().map(|a| a.deref()),
                    &country_codes,
                );
                let coord = get_way_coord(&objs_map, way);
                Some(mimir::Street {
                    id: format!("street:osm:relation:{}", rel.id.0.to_string()),
                    name: way_name.to_string(),
                    label: street_label,
                    weight: 0.,
                    zip_codes: utils::get_zip_codes_from_admins(&admins),
                    administrative_regions: admins,
                    coord: get_way_coord(&objs_map, way),
                    approx_coord: Some(coord.into()),
                    distance: None,
                    country_codes,
                    context: None,
                })
            })
            .next()
            .map(|street| street_list.push(street));

        // Add osmid of all the relation members in the set
        // We don't create any street for all the osmid present in street_rel
        for ref_obj in &rel.refs {
            if ref_obj.member.is_way() {
                street_rel.insert(ref_obj.member);
            }
        }
    });

    // we merge all the ways with a key = way_name + admin list of level(=city_level)
    // we use a map NameAdminMap <key, value> to manage the merging of ways
    let mut name_admin_map = NameAdminMap::default();
    objs_map.for_each_filter(Kind::Way, |obj| {
        let osmid = obj.id();
        if street_rel.contains(&osmid) {
            return;
        }
        if let Some(way) = obj.way() {
            if let Some(name) = way.tags.get("name") {
                let name = name.to_string();
                let admins = get_street_admin(admins_geofinder, &objs_map, way)
                    .into_iter()
                    .filter(|admin| admin.is_city())
                    .collect();
                name_admin_map
                    .entry(StreetKey { name, admins })
                    .or_insert_with(Vec::new)
                    .push(osmid);
            }
        }
    });

    // Create a street for each way with osmid present in objs_map
    let streets = name_admin_map.values().filter_map(|way_ids| {
        let min_id = way_ids.iter().min()?;
        let obj = objs_map.get(&min_id)?;
        let way = obj.way()?;
        let name = way.tags.get("name")?.to_string();
        let admins = get_street_admin(admins_geofinder, &objs_map, way);

        let country_codes = utils::find_country_codes(admins.iter().map(|a| a.deref()));
        let street_label =
            labels::format_street_label(&name, admins.iter().map(|a| a.deref()), &country_codes);
        let coord = get_way_coord(&objs_map, way);
        Some(mimir::Street {
            id: format!("street:osm:way:{}", way.id.0.to_string()),
            label: street_label,
            name,
            weight: 0.,
            zip_codes: utils::get_zip_codes_from_admins(&admins),
            administrative_regions: admins,
            coord: get_way_coord(&objs_map, way),
            approx_coord: Some(coord.into()),
            distance: None,
            country_codes,
            context: None,
        })
    });
    street_list.extend(streets);

    Ok(street_list)
}

fn get_street_admin<T: StoreObjs + Getter>(
    admins_geofinder: &AdminGeoFinder,
    obj_map: &T,
    way: &osmpbfreader::objects::Way,
) -> Vec<Arc<mimir::Admin>> {
    /*
        To avoid corner cases where the ends of the way are near
        administrative boundaries, the geofinder is called
        on a middle node.
    */
    let nb_nodes = way.nodes.len();
    way.nodes
        .iter()
        .skip(nb_nodes / 2)
        .filter_map(|node_id| obj_map.get(&(*node_id).into()))
        .filter_map(|node_obj| {
            node_obj.node().map(|node| geo_types::Coordinate {
                x: node.lon(),
                y: node.lat(),
            })
        })
        .next()
        .map_or(vec![], |c| admins_geofinder.get(&c))
}

pub fn compute_street_weight(streets: &mut StreetsVec) {
    for st in streets {
        for admin in &mut st.administrative_regions {
            if admin.is_city() {
                st.weight = admin.weight;
                break;
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::old_io::IoErrorKind::*;
use std::collections::VecMap;
use unicode::str::*;

use std::str;
use std::old_io;
use std::fmt;

use constants::*;
use util::*;
use types::*;
use ioctl::*;

type Stdr = old_io::stdio::StdReader;
type Stdw = old_io::stdio::StdWriter;

pub struct Controls {
    stdin: Stdr,
    stdout: Stdw,
    stderr: Stdw,
    cursor: Position,
    tsize: WinSize,
    rows: VecMap<usize>,
    roff: isize
}

impl Controls {
    pub fn new() -> Controls {
        Controls {
            stdin: old_io::stdio::stdin_raw(),
            stdout: old_io::stdio::stdout_raw(),
            stderr: old_io::stdio::stderr_raw(),
            cursor: Position::new(),
            tsize: WinSize::new(),
            rows: VecMap::new(),
            roff: 0
        }
    }

    pub fn update_cursor(&mut self, pos:Position) {
        self.cursor = pos;
    }

    pub fn update_size(&mut self, size:WinSize) {
        self.tsize = size;
    }

    pub fn clear_rows(&mut self) {
        self.rows.clear();
        self.roff = 0;
    }

    pub fn get_row(&mut self) -> usize {
        self.cursor.row
    }

    pub fn width(&mut self) -> usize {
        self.tsize.col as usize
    }

    pub fn get_pos(&mut self) -> Position {
        self.cursor
    }

    fn row_length(&mut self) -> usize {
        if !self.rows.contains_key(&((self.cursor.row as isize + self.roff) as usize)) {
            return 1;
        } else {
            return self.rows[(self.cursor.row as isize + self.roff) as usize];
        }
    }

    fn save_row(&mut self, len:usize) {
        if !self.rows.contains_key(&((self.cursor.row as isize + self.roff) as usize)) {
            self.rows.insert((self.cursor.row as isize + self.roff) as usize, len);
        } else {
            self.rows[(self.cursor.row as isize + self.roff) as usize] = len;
        }
    }

    fn move_right(&mut self, by:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        self.cursor.col += by;
        let mut row_size = self.row_length();
        if self.cursor.col > row_size {
            loop {
                if self.cursor.col > row_size {
                    self.cursor.col -= row_size;
                    self.cursor.row += 1;
                    row_size = self.row_length();
                } else {
                    break;
                }
            }
            if self.cursor.row > self.tsize.row as usize {
                self.roff += self.cursor.row as isize - self.tsize.row as isize;
                self.cursor.row = self.tsize.row as usize;
            }
        }
    }

    fn move_left(&mut self, mut by:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        loop {
            if by >= self.cursor.col {
                by -= self.cursor.col;
                if self.cursor.row == 0 {
                    self.roff -= 1;
                } else {
                    self.cursor.row -= 1;
                }
                self.cursor.col = self.row_length();
            } else {
                self.cursor.col -= by;
                break;
            }
        }
    }

    pub fn grow_check(&mut self, len:usize) -> bool {
        self.cursor.col + len > self.tsize.col as usize
    }

    pub fn grow(&mut self, by:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        self.cursor.col += by;
        if self.cursor.col > self.tsize.col as usize {
            let col = self.tsize.col as usize;
            self.save_row(col);
            self.cursor.row += self.cursor.col / self.tsize.col as usize;
            self.cursor.col = self.cursor.col % self.tsize.col as usize;
            if self.cursor.row > self.tsize.row as usize {
                self.roff += self.cursor.row as isize - self.tsize.row as isize;
                self.cursor.row = self.tsize.row as usize;
            }
        }
        let col = self.cursor.col;
        self.save_row(col);
    }

    pub fn shrink(&mut self, by:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if by >= self.cursor.col {
            self.rows.remove(&((self.cursor.row as isize + self.roff) as usize));
            let diff = by - self.cursor.col;
            self.cursor.col = self.tsize.col as usize -
                (diff % self.tsize.col as usize);
            let rdiff = (diff / self.tsize.col as usize) + 1;
            if rdiff > self.cursor.row {
                self.cursor.row = 0;
            } else {
                self.cursor.row -= rdiff;
            }
        } else {
            self.cursor.col -= by;
        }
        let col = self.cursor.col;
        self.save_row(col);
    }

    fn new_row(&mut self) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if self.cursor.row == self.tsize.row as usize {
            self.roff += 1;
        } else {
            self.cursor.row += 1;
        }
        self.cursor.col = 1;
    }

    /* out* and err* functions are assumed to be output
       functions */

    pub fn outc(&mut self, ch:char) {
        self.stdout.write_char(ch).unwrap();
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if ch == NL {
            self.new_row();
        } else {
            if self.grow_check(1) {
                self.stdout.write_char(NL).unwrap();
            }
            self.grow(1);
        }
    }

    pub fn outs(&mut self, s:&str) {
        if self.tsize.col == 0 || self.tsize.row == 0 {
            self.stdout.write_str(s).unwrap();
            return;
        }
        let mut splits = NL_REGEX.split(s);
        match splits.next() {
            Some(part) => {
                self.stdout.write_str(part).unwrap();
                if self.cursor.col + part.len() - 1 == self.tsize.col as usize {
                    self.stdout.write_char(NL).unwrap();
                }
                self.grow(part.len());
                for part in splits {
                    self.new_row();
                    self.stdout.write_char(NL).unwrap();
                    self.stdout.write_str(part).unwrap();
                    if self.cursor.col + part.len() - 1 == self.tsize.col as usize {
                        self.stdout.write_char(NL).unwrap();
                    }
                    self.grow(part.len());
                }
            },
            _ => {}
        }
    }

    pub fn outf(&mut self, args:fmt::Arguments) {
        self.outs(fmt::format(args).as_slice());
    }

    pub fn err(&mut self, s:&str) {
        if self.tsize.col == 0 || self.tsize.row == 0 {
            self.stderr.write_str(s).unwrap();
            return;
        }
        self.stderr.write_str(s).unwrap();
        let mut splits = NL_REGEX.split(s);
        match splits.next() {
            Some(part) => {
                self.grow(part.len());
                for part in splits {
                    self.new_row();
                    if self.cursor.col + part.len() == self.tsize.col as usize {
                        self.stdout.write_char(NL).unwrap();
                    }
                    self.grow(part.len());
                }
            },
            _ => {}
        }
    }

    pub fn errf(&mut self, args:fmt::Arguments) {
        self.err(fmt::format(args).as_slice());
    }

    pub fn del(&mut self) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if self.cursor.col > 1 {
            self.stdout.write_char(DEL).unwrap();
            self.shrink(1);
        } else {
            self.shrink(1);
            self.move_to_pointer();
        }
    }

    pub fn read(&mut self) -> old_io::IoResult<char> {
        // Below lifted almost verbatim from rust's read_char.
        // In compliance with MIT, nothing to worry about
        let first_byte = try!(self.stdin.read_byte());
        let width = utf8_char_width(first_byte);
        if width == 1 { return Ok(first_byte as char) }
        if width == 0 { return Err(old_io::standard_error(InvalidInput)) } // not utf8
        let mut buf = [first_byte, 0, 0, 0];
        {
            let mut start = 1;
            while start < width {
                match try!(self.stdin.read(&mut buf[start..width])) {
                    n if n == width - start => break,
                    n if n < width - start => { start += n; }
                    _ => return Err(old_io::standard_error(InvalidInput)),
                }
            }
        }
        match str::from_utf8(&buf[..width]).ok() {
            Some(s) => Ok(s.char_at(0)),
            None => Err(old_io::standard_error(InvalidInput))
        }
    }

    pub fn cursor_left(&mut self) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if self.cursor.col > 1 {
            self.stdout.write_char(DEL).unwrap();
            self.move_left(1);
        } else {
            self.move_left(1);
            self.move_to_pointer();
        }
    }

    pub fn cursor_right(&mut self) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if self.cursor.col < self.row_length() {
            self.stdout.write_str(CRSR_RIGHT).unwrap();
            self.move_right(1);
        } else {
            self.move_right(1);
            self.move_to_pointer();
        }
    }

    pub fn cursors_left(&mut self, by:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if by == 0 {
            return;
        } else if by <= 3 && self.cursor.col > by {
            self.stdout.write_str(build_string(DEL, by).as_slice()).unwrap();
            self.move_left(by);
        } else if self.cursor.col > by {
            self.stdout.write_fmt(format_args!("{}{}D", ANSI_BEGIN, by)).unwrap();
            self.move_left(by);
        } else {
            self.move_left(by);
            self.move_to_pointer();
        }
    }

    pub fn cursors_right(&mut self, by:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if by == 0 {
            return;
        } else if by + self.cursor.col <= self.row_length() {
            self.stdout.write_fmt(format_args!("{}{}C", ANSI_BEGIN, by)).unwrap();
            self.move_right(by);
        } else {
            self.move_right(by);
            self.move_to_pointer();
        }
    }

    pub fn clear_line(&mut self) {
        self.stdout.write_str(ANSI_BEGIN).unwrap();
        self.stdout.write_char('K').unwrap();
        let col = self.cursor.col;
        self.save_row(col);
    }

    pub fn clear_line_to(&mut self, len:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        let old = self.cursor;
        let mut total = self.cursor.col + len;
        let mut rlen = self.row_length();
        loop {
            self.clear_line();
            if total > rlen {
                total -= rlen;
                self.next_start();
                rlen = self.row_length();
            } else {
                break;
            }
        }
        self.move_to(old);
    }

    pub fn next_start(&mut self) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        self.new_row();
        let row = self.cursor.row;
        self.move_to(Position {
            col: 1,
            row: row
        });
    }

    pub fn flush(&mut self) {
        self.stdout.flush().unwrap();
        self.stderr.flush().unwrap();
    }

    pub fn query_cursor(&mut self) {
        self.stdout.write_str(CRSR_POS).unwrap();
    }

    fn move_to_pointer(&mut self) {
        let (row, col) = (self.cursor.row, self.cursor.col);
        self.stdout.write_fmt(format_args!("{}{};{}f", ANSI_BEGIN, row, col)).unwrap();
    }

    pub fn move_to(&mut self, pos:Position) {
        self.cursor = pos;
        self.move_to_pointer();
    }

    pub fn bell(&mut self) {
        self.stdout.write_char(BEL).unwrap();
    }
}

<commit_msg>Fix clear_entire_line on new lines<commit_after>use std::old_io::IoErrorKind::*;
use std::collections::VecMap;
use unicode::str::*;

use std::str;
use std::old_io;
use std::fmt;

use constants::*;
use util::*;
use types::*;
use ioctl::*;

type Stdr = old_io::stdio::StdReader;
type Stdw = old_io::stdio::StdWriter;

pub struct Controls {
    stdin: Stdr,
    stdout: Stdw,
    stderr: Stdw,
    cursor: Position,
    tsize: WinSize,
    rows: VecMap<usize>,
    roff: isize
}

impl Controls {
    pub fn new() -> Controls {
        Controls {
            stdin: old_io::stdio::stdin_raw(),
            stdout: old_io::stdio::stdout_raw(),
            stderr: old_io::stdio::stderr_raw(),
            cursor: Position::new(),
            tsize: WinSize::new(),
            rows: VecMap::new(),
            roff: 0
        }
    }

    pub fn update_cursor(&mut self, pos:Position) {
        self.cursor = pos;
        self.save_row(pos.col);
    }

    pub fn update_size(&mut self, size:WinSize) {
        self.tsize = size;
    }

    pub fn clear_rows(&mut self) {
        self.rows.clear();
        self.roff = 0;
    }

    pub fn get_row(&mut self) -> usize {
        self.cursor.row
    }

    pub fn width(&mut self) -> usize {
        self.tsize.col as usize
    }

    pub fn get_pos(&mut self) -> Position {
        self.cursor
    }

    fn row_length(&mut self) -> usize {
        if !self.rows.contains_key(&((self.cursor.row as isize + self.roff) as usize)) {
            return 1;
        } else {
            return self.rows[(self.cursor.row as isize + self.roff) as usize];
        }
    }

    fn save_row(&mut self, len:usize) {
        if !self.rows.contains_key(&((self.cursor.row as isize + self.roff) as usize)) {
            self.rows.insert((self.cursor.row as isize + self.roff) as usize, len);
        } else {
            self.rows[(self.cursor.row as isize + self.roff) as usize] = len;
        }
    }

    fn move_right(&mut self, by:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        self.cursor.col += by;
        let mut row_size = self.row_length();
        if self.cursor.col > row_size {
            loop {
                if self.cursor.col > row_size {
                    self.cursor.col -= row_size;
                    self.cursor.row += 1;
                    row_size = self.row_length();
                } else {
                    break;
                }
            }
            if self.cursor.row > self.tsize.row as usize {
                self.roff += self.cursor.row as isize - self.tsize.row as isize;
                self.cursor.row = self.tsize.row as usize;
            }
        }
    }

    fn move_left(&mut self, mut by:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        loop {
            if by >= self.cursor.col {
                by -= self.cursor.col;
                if self.cursor.row == 0 {
                    self.roff -= 1;
                } else {
                    self.cursor.row -= 1;
                }
                self.cursor.col = self.row_length();
            } else {
                self.cursor.col -= by;
                break;
            }
        }
    }

    pub fn grow_check(&mut self, len:usize) -> bool {
        self.cursor.col + len > self.tsize.col as usize
    }

    pub fn grow(&mut self, by:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        self.cursor.col += by;
        if self.cursor.col > self.tsize.col as usize {
            let col = self.tsize.col as usize;
            self.save_row(col);
            self.cursor.row += self.cursor.col / self.tsize.col as usize;
            self.cursor.col = self.cursor.col % self.tsize.col as usize;
            if self.cursor.row > self.tsize.row as usize {
                self.roff += self.cursor.row as isize - self.tsize.row as isize;
                self.cursor.row = self.tsize.row as usize;
            }
        }
        let col = self.cursor.col;
        self.save_row(col);
    }

    pub fn shrink(&mut self, by:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if by >= self.cursor.col {
            self.rows.remove(&((self.cursor.row as isize + self.roff) as usize));
            let diff = by - self.cursor.col;
            self.cursor.col = self.tsize.col as usize -
                (diff % self.tsize.col as usize);
            let rdiff = (diff / self.tsize.col as usize) + 1;
            if rdiff > self.cursor.row {
                self.cursor.row = 0;
            } else {
                self.cursor.row -= rdiff;
            }
        } else {
            self.cursor.col -= by;
        }
        let col = self.cursor.col;
        self.save_row(col);
    }

    fn new_row(&mut self) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if self.cursor.row == self.tsize.row as usize {
            self.roff += 1;
        } else {
            self.cursor.row += 1;
        }
        self.cursor.col = 1;
    }

    /* out* and err* functions are assumed to be output
       functions */

    pub fn outc(&mut self, ch:char) {
        self.stdout.write_char(ch).unwrap();
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if ch == NL {
            self.new_row();
        } else {
            if self.grow_check(1) {
                self.stdout.write_char(NL).unwrap();
            }
            self.grow(1);
        }
    }

    pub fn outs(&mut self, s:&str) {
        if self.tsize.col == 0 || self.tsize.row == 0 {
            self.stdout.write_str(s).unwrap();
            return;
        }
        let mut splits = NL_REGEX.split(s);
        match splits.next() {
            Some(part) => {
                self.stdout.write_str(part).unwrap();
                if self.cursor.col + part.len() - 1 == self.tsize.col as usize {
                    self.stdout.write_char(NL).unwrap();
                }
                self.grow(part.len());
                for part in splits {
                    self.new_row();
                    self.stdout.write_char(NL).unwrap();
                    self.stdout.write_str(part).unwrap();
                    if self.cursor.col + part.len() - 1 == self.tsize.col as usize {
                        self.stdout.write_char(NL).unwrap();
                    }
                    self.grow(part.len());
                }
            },
            _ => {}
        }
    }

    pub fn outf(&mut self, args:fmt::Arguments) {
        self.outs(fmt::format(args).as_slice());
    }

    pub fn err(&mut self, s:&str) {
        if self.tsize.col == 0 || self.tsize.row == 0 {
            self.stderr.write_str(s).unwrap();
            return;
        }
        self.stderr.write_str(s).unwrap();
        let mut splits = NL_REGEX.split(s);
        match splits.next() {
            Some(part) => {
                self.grow(part.len());
                for part in splits {
                    self.new_row();
                    if self.cursor.col + part.len() == self.tsize.col as usize {
                        self.stdout.write_char(NL).unwrap();
                    }
                    self.grow(part.len());
                }
            },
            _ => {}
        }
    }

    pub fn errf(&mut self, args:fmt::Arguments) {
        self.err(fmt::format(args).as_slice());
    }

    pub fn del(&mut self) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if self.cursor.col > 1 {
            self.stdout.write_char(DEL).unwrap();
            self.shrink(1);
        } else {
            self.shrink(1);
            self.move_to_pointer();
        }
    }

    pub fn read(&mut self) -> old_io::IoResult<char> {
        // Below lifted almost verbatim from rust's read_char.
        // In compliance with MIT, nothing to worry about
        let first_byte = try!(self.stdin.read_byte());
        let width = utf8_char_width(first_byte);
        if width == 1 { return Ok(first_byte as char) }
        if width == 0 { return Err(old_io::standard_error(InvalidInput)) } // not utf8
        let mut buf = [first_byte, 0, 0, 0];
        {
            let mut start = 1;
            while start < width {
                match try!(self.stdin.read(&mut buf[start..width])) {
                    n if n == width - start => break,
                    n if n < width - start => { start += n; }
                    _ => return Err(old_io::standard_error(InvalidInput)),
                }
            }
        }
        match str::from_utf8(&buf[..width]).ok() {
            Some(s) => Ok(s.char_at(0)),
            None => Err(old_io::standard_error(InvalidInput))
        }
    }

    pub fn cursor_left(&mut self) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if self.cursor.col > 1 {
            self.stdout.write_char(DEL).unwrap();
            self.move_left(1);
        } else {
            self.move_left(1);
            self.move_to_pointer();
        }
    }

    pub fn cursor_right(&mut self) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if self.cursor.col < self.row_length() {
            self.stdout.write_str(CRSR_RIGHT).unwrap();
            self.move_right(1);
        } else {
            self.move_right(1);
            self.move_to_pointer();
        }
    }

    pub fn cursors_left(&mut self, by:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if by == 0 {
            return;
        } else if by <= 3 && self.cursor.col > by {
            self.stdout.write_str(build_string(DEL, by).as_slice()).unwrap();
            self.move_left(by);
        } else if self.cursor.col > by {
            self.stdout.write_fmt(format_args!("{}{}D", ANSI_BEGIN, by)).unwrap();
            self.move_left(by);
        } else {
            self.move_left(by);
            self.move_to_pointer();
        }
    }

    pub fn cursors_right(&mut self, by:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        if by == 0 {
            return;
        } else if by + self.cursor.col <= self.row_length() {
            self.stdout.write_fmt(format_args!("{}{}C", ANSI_BEGIN, by)).unwrap();
            self.move_right(by);
        } else {
            self.move_right(by);
            self.move_to_pointer();
        }
    }

    pub fn clear_line(&mut self) {
        self.stdout.write_str(ANSI_BEGIN).unwrap();
        self.stdout.write_char('K').unwrap();
        let col = self.cursor.col;
        self.save_row(col);
    }

    pub fn clear_line_to(&mut self, len:usize) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        let old = self.cursor;
        let mut total = self.cursor.col + len;
        let mut rlen = self.row_length();
        loop {
            self.clear_line();
            if total > rlen {
                total -= rlen;
                self.next_start();
                rlen = self.row_length();
            } else {
                break;
            }
        }
        self.move_to(old);
    }

    pub fn next_start(&mut self) {
        if self.tsize.col == 0 || self.tsize.row == 0 {return}
        self.new_row();
        let row = self.cursor.row;
        self.move_to(Position {
            col: 1,
            row: row
        });
    }

    pub fn flush(&mut self) {
        self.stdout.flush().unwrap();
        self.stderr.flush().unwrap();
    }

    pub fn query_cursor(&mut self) {
        self.stdout.write_str(CRSR_POS).unwrap();
    }

    fn move_to_pointer(&mut self) {
        let (row, col) = (self.cursor.row, self.cursor.col);
        self.stdout.write_fmt(format_args!("{}{};{}f", ANSI_BEGIN, row, col)).unwrap();
    }

    pub fn move_to(&mut self, pos:Position) {
        self.cursor = pos;
        self.move_to_pointer();
    }

    pub fn bell(&mut self) {
        self.stdout.write_char(BEL).unwrap();
    }
}

<|endoftext|>


<commit_before>/* Copyright (c) 2016
 * Jeff Nettleton
 *
 * Licensed under the MIT license (http://opensource.org/licenses/MIT). This
 * file may not be copied, modified, or distributed except according to those
 * terms
 */

extern crate chrono;

use std::collections::HashMap;

pub trait ToOutput {
    fn to_output(&self) -> &[u8];
}

impl ToOutput for str {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for &'static str {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for String {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for Vec<u8> {
    fn to_output(&self) -> &[u8] {
        self.as_slice()
    }
}

#[derive(Debug)]
pub struct Response {
    code:       i32,
    cmsg:       String,
    ctype:      String,
    headers:    HashMap<String, String>,
    payload:    Vec<u8>,
}

impl Response {
    pub fn new() -> Response {
        let mut res = Response {
            code:       200,
            cmsg:       String::from("OK"),
            ctype:      String::from("text/plain"),
            headers:    HashMap::new(),
            payload:    Vec::with_capacity(2048),
        };

        let now = chrono::UTC::now().format("%a, %d %b %Y, %H:%M:%S GMT").to_string();

        res.add_header("Connection", "close");
        res.add_header("Server", "canteen/0.0.1");
        res.add_header("Date", &now);

        res
    }

    pub fn get_http_message(code: i32) -> String {
        let msg = match code {
            100 => "Continue",
            101 => "Switching Protocols",
            200 => "OK",
            201 => "Created",
            202 => "Accepted",
            203 => "Non-Authoritative Information",
            204 => "No Content",
            205 => "Reset Content",
            206 => "Partial Content",
            300 => "Multiple Choices",
            301 => "Moved Permanently",
            302 => "Found",
            303 => "See Other",
            304 => "Not Modified",
            305 => "Use Proxy",
            307 => "Temporary Redirect",
            400 => "Bad Request",
            401 => "Unauthorized",
            402 => "Payment Required",
            403 => "Forbidden",
            404 => "Not Found",
            405 => "Method Not Allowed",
            406 => "Not Acceptable",
            407 => "Proxy Authentication Required",
            408 => "Request Time Out",
            409 => "Conflict",
            410 => "Gone",
            411 => "Length Required",
            412 => "Precondition Failed",
            413 => "Request Entity Too Large",
            414 => "Request-URI Too Large",
            415 => "Unsupported Media Type",
            416 => "Requested Range Not Satisfiable",
            417 => "Expectation Failed",
            500 => "Internal Server Error",
            501 => "Not Implemented",
            502 => "Bad Gateway",
            503 => "Service Unavailable",
            504 => "Gateway Time-out",
            505 => "HTTP Version Not Supported",
            _     => "OK",
        };

        String::from(msg)
    }

    fn err_body(message: &str, path: &str) -> String {
        format!("\
                 \

Your request failed

{}: {}
", message, path)
    }

    pub fn err_403(path: &str) -> Response {
        let mut res = Response::new();

        res.set_code(403);
        res.set_content_type("text/html");
        res.append(Response::err_body("forbidden", path));

        res
    }

    pub fn err_404(path: &str) -> Response {
        let mut res = Response::new();

        res.set_code(403);
        res.set_content_type("text/html");
        res.append(Response::err_body("not found", path));

        res
    }

    pub fn err_500(path: &str) -> Response {
        let mut res = Response::new();

        res.set_code(500);
        res.set_content_type("text/html");
        res.append(Response::err_body("internal server error", path));

        res
    }

    /* set the response code
     * ex: res.set_code(200, "OK");
     */
    pub fn set_code(&mut self, code: i32) {
        self.code = code;
        self.cmsg = Response::get_http_message(code);
    }

    /* set the content type
     * ex: res.set_content_type("text/html");
     */
    pub fn set_content_type(&mut self, ctype: &str) {
        self.ctype = String::from(ctype);
    }

    /* add an HTTP header
     * ex: res.add_header("Connection", "close");
     */
    pub fn add_header(&mut self, key: &str, value: &str) {
        if !self.headers.contains_key(key) {
            self.headers.insert(String::from(key), String::from(value));
        }
    }

    /* add data to the payload -- can take any type that has
     * implemented the canteen::response::ToOutput trait
     */
    pub fn append<T: ToOutput>(&mut self, payload: T) {
        self.payload.extend(payload.to_output().into_iter());
    }

    pub fn gen_output(&self) -> Vec<u8> {
        let mut output: Vec<u8> = Vec::with_capacity(self.payload.len() + 500);
        let mut inter = String::new();

        inter.push_str(&format!("HTTP/1.1 {} {}\r\n", self.code, self.cmsg));

        for (key, value) in &self.headers {
            inter.push_str(&format!("{}: {}\r\n", key, value));
        }

        inter.push_str(&format!("Content-Type: {}\r\n", self.ctype));
        inter.push_str(&format!("Content-Length: {}\r\n", self.payload.len()));
        inter.push_str("\r\n");

        output.extend(inter.as_bytes());
        output.extend(self.payload.iter());

        output
    }
}
<commit_msg>implement utils::make_response for brevity's sake<commit_after>/* Copyright (c) 2016
 * Jeff Nettleton
 *
 * Licensed under the MIT license (http://opensource.org/licenses/MIT). This
 * file may not be copied, modified, or distributed except according to those
 * terms
 */

extern crate chrono;

use std::collections::HashMap;
use utils::make_response;

pub trait ToOutput {
    fn to_output(&self) -> &[u8];
}

impl ToOutput for str {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for &'static str {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for String {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for Vec<u8> {
    fn to_output(&self) -> &[u8] {
        self.as_slice()
    }
}

#[derive(Debug)]
pub struct Response {
    code:       i32,
    cmsg:       String,
    ctype:      String,
    headers:    HashMap<String, String>,
    payload:    Vec<u8>,
}

impl Response {
    pub fn new() -> Response {
        let mut res = Response {
            code:       200,
            cmsg:       String::from("OK"),
            ctype:      String::from("text/plain"),
            headers:    HashMap::new(),
            payload:    Vec::with_capacity(2048),
        };

        let now = chrono::UTC::now().format("%a, %d %b %Y, %H:%M:%S GMT").to_string();

        res.add_header("Connection", "close");
        res.add_header("Server", "canteen/0.0.1");
        res.add_header("Date", &now);

        res
    }

    pub fn get_http_message(code: i32) -> String {
        let msg = match code {
            100 => "Continue",
            101 => "Switching Protocols",
            200 => "OK",
            201 => "Created",
            202 => "Accepted",
            203 => "Non-Authoritative Information",
            204 => "No Content",
            205 => "Reset Content",
            206 => "Partial Content",
            300 => "Multiple Choices",
            301 => "Moved Permanently",
            302 => "Found",
            303 => "See Other",
            304 => "Not Modified",
            305 => "Use Proxy",
            307 => "Temporary Redirect",
            400 => "Bad Request",
            401 => "Unauthorized",
            402 => "Payment Required",
            403 => "Forbidden",
            404 => "Not Found",
            405 => "Method Not Allowed",
            406 => "Not Acceptable",
            407 => "Proxy Authentication Required",
            408 => "Request Time Out",
            409 => "Conflict",
            410 => "Gone",
            411 => "Length Required",
            412 => "Precondition Failed",
            413 => "Request Entity Too Large",
            414 => "Request-URI Too Large",
            415 => "Unsupported Media Type",
            416 => "Requested Range Not Satisfiable",
            417 => "Expectation Failed",
            500 => "Internal Server Error",
            501 => "Not Implemented",
            502 => "Bad Gateway",
            503 => "Service Unavailable",
            504 => "Gateway Time-out",
            505 => "HTTP Version Not Supported",
            _     => "OK",
        };

        String::from(msg)
    }

    fn err_body(message: &str, path: &str) -> String {
        format!("\
                 \

Your request failed

{}: {}
", message, path)
    }

    pub fn err_403(path: &str) -> Response {
        make_response(Response::err_body("forbidden", path), "text/html", 403)
    }

    pub fn err_404(path: &str) -> Response {
        make_response(Response::err_body("not found", path), "text/html", 403)
    }

    pub fn err_500(path: &str) -> Response {
        make_response(Response::err_body("internal server error", path), "text/html", 500)
    }

    /* set the response code
     * ex: res.set_code(200, "OK");
     */
    pub fn set_code(&mut self, code: i32) {
        self.code = code;
        self.cmsg = Response::get_http_message(code);
    }

    /* set the content type
     * ex: res.set_content_type("text/html");
     */
    pub fn set_content_type(&mut self, ctype: &str) {
        self.ctype = String::from(ctype);
    }

    /* add an HTTP header
     * ex: res.add_header("Connection", "close");
     */
    pub fn add_header(&mut self, key: &str, value: &str) {
        if !self.headers.contains_key(key) {
            self.headers.insert(String::from(key), String::from(value));
        }
    }

    /* add data to the payload -- can take any type that has
     * implemented the canteen::response::ToOutput trait
     */
    pub fn append<T: ToOutput>(&mut self, payload: T) {
        self.payload.extend(payload.to_output().into_iter());
    }

    pub fn gen_output(&self) -> Vec<u8> {
        let mut output: Vec<u8> = Vec::with_capacity(self.payload.len() + 500);
        let mut inter = String::new();

        inter.push_str(&format!("HTTP/1.1 {} {}\r\n", self.code, self.cmsg));

        for (key, value) in &self.headers {
            inter.push_str(&format!("{}: {}\r\n", key, value));
        }

        inter.push_str(&format!("Content-Type: {}\r\n", self.ctype));
        inter.push_str(&format!("Content-Length: {}\r\n", self.payload.len()));
        inter.push_str("\r\n");

        output.extend(inter.as_bytes());
        output.extend(self.payload.iter());

        output
    }
}
<|endoftext|>


<commit_before>pub use Response = http::server::response::ResponseWriter;
<commit_msg>(doc) Documented Response.<commit_after>//! An alias of the rust-http Response struct.

pub use Response = http::server::response::ResponseWriter;
<|endoftext|>


<commit_before>//! Generic functions for working with (primarily nucleic acid) sequences
use std::borrow::Cow;

use memchr::memchr2;

use crate::bitkmer::BitNuclKmer;
use crate::kmer::{CanonicalKmers, Kmers};

/// Transform a nucleic acid sequence into its "normalized" form.
///
/// The normalized form is:
///  - only AGCTN and possibly - (for gaps)
///  - strip out any whitespace or line endings
///  - lowercase versions of these are uppercased
///  - U is converted to T (make everything a DNA sequence)
///  - some other punctuation is converted to gaps
///  - IUPAC bases may be converted to N's depending on the parameter passed in
///  - everything else is considered a N
pub fn normalize(seq: &[u8], allow_iupac: bool) -> Option<Vec<u8>> {
    let mut buf: Vec<u8> = Vec::with_capacity(seq.len());
    let mut changed: bool = false;

    for n in seq.iter() {
        let (new_char, char_changed) = match (*n, allow_iupac) {
            c @ (b'A', _)
            | c @ (b'C', _)
            | c @ (b'G', _)
            | c @ (b'T', _)
            | c @ (b'N', _)
            | c @ (b'-', _) => (c.0, false),
            (b'a', _) => (b'A', true),
            (b'c', _) => (b'C', true),
            (b'g', _) => (b'G', true),
            // normalize uridine to thymine
            (b't', _) | (b'u', _) | (b'U', _) => (b'T', true),
            // normalize gaps
            (b'.', _) | (b'~', _) => (b'-', true),
            // logic for IUPAC bases (a little messy)
            c @ (b'B', true)
            | c @ (b'D', true)
            | c @ (b'H', true)
            | c @ (b'V', true)
            | c @ (b'R', true)
            | c @ (b'Y', true)
            | c @ (b'S', true)
            | c @ (b'W', true)
            | c @ (b'K', true)
            | c @ (b'M', true) => (c.0, false),
            (b'b', true) => (b'B', true),
            (b'd', true) => (b'D', true),
            (b'h', true) => (b'H', true),
            (b'v', true) => (b'V', true),
            (b'r', true) => (b'R', true),
            (b'y', true) => (b'Y', true),
            (b's', true) => (b'S', true),
            (b'w', true) => (b'W', true),
            (b'k', true) => (b'K', true),
            (b'm', true) => (b'M', true),
            // remove all whitespace and line endings
            (b' ', _) | (b'\t', _) | (b'\r', _) | (b'\n', _) => (b' ', true),
            // everything else is an N
            _ => (b'N', true),
        };
        changed = changed || char_changed;
        if new_char != b' ' {
            buf.push(new_char);
        }
    }
    if changed {
        Some(buf)
    } else {
        None
    }
}

#[test]
fn test_normalize() {
    assert_eq!(normalize(b"ACGTU", false), Some(b"ACGTT".to_vec()));
    assert_eq!(normalize(b"acgtu", false), Some(b"ACGTT".to_vec()));

    assert_eq!(normalize(b"N.N-N~N N", false), Some(b"N-N-N-NN".to_vec()));

    assert_eq!(normalize(b"BDHVRYSWKM", true), None);
    assert_eq!(normalize(b"bdhvryswkm", true), Some(b"BDHVRYSWKM".to_vec()));
    assert_eq!(
        normalize(b"BDHVRYSWKM", false),
        Some(b"NNNNNNNNNN".to_vec())
    );
    assert_eq!(
        normalize(b"bdhvryswkm", false),
        Some(b"NNNNNNNNNN".to_vec())
    );
}

/// Returns the complementary base for a given IUPAC base code.
///
/// Does not work for RNA sequences (maybe we should raise an error or something?)
#[inline]
pub fn complement(n: u8) -> u8 {
    match n {
        b'a' => b't',
        b'A' => b'T',
        b'c' => b'g',
        b'C' => b'G',
        b'g' => b'c',
        b'G' => b'C',
        b't' => b'a',
        b'T' => b'A',

        // IUPAC codes
        b'r' => b'y',
        b'y' => b'r',
        b'k' => b'm',
        b'm' => b'k',
        b'b' => b'v',
        b'v' => b'b',
        b'd' => b'h',
        b'h' => b'd',
        b's' => b's',
        b'w' => b'w',
        b'R' => b'Y',
        b'Y' => b'R',
        b'K' => b'M',
        b'M' => b'K',
        b'B' => b'V',
        b'V' => b'B',
        b'D' => b'H',
        b'H' => b'D',
        b'S' => b'S',
        b'W' => b'W',

        // anything else just pass through
        // 'u' | 'U' => panic!("Does not support complements of U"),
        x => x,
    }
}

/// Taking in a sequence string, return the canonical form of the sequence
/// (e.g. the lexigraphically lowest of either the original sequence or its
/// reverse complement)
pub fn canonical(seq: &[u8]) -> Cow<[u8]> {
    let mut buf: Vec<u8> = Vec::with_capacity(seq.len());
    // enough just keeps our comparisons from happening after they need to
    let mut enough = false;
    let mut original_was_canonical = false;

    // loop through the kmer and its reverse complement simultaneously
    for (rn, n) in seq.iter().rev().map(|n| complement(*n)).zip(seq.iter()) {
        buf.push(rn);
        if !enough && (*n < rn) {
            original_was_canonical = true;
            break;
        } else if !enough && (rn < *n) {
            enough = true;
        }
        // unstated if branch: if rn == n, keep comparing
    }
    match (original_was_canonical, enough) {
        (true, true) => panic!("Bug: should never set original_was_canonical if enough == true"),
        (true, false) => seq.into(),
        (false, true) => buf.into(),
        // the sequences were completely equal, return the ref
        (false, false) => seq.into(),
    }
}

/// Find the lexigraphically smallest substring of `seq` of length `length`
///
/// There's probably a faster algorithm for this somewhere...
pub fn minimizer(seq: &[u8], length: usize) -> Cow<[u8]> {
    let reverse_complement: Vec<u8> = seq.iter().rev().map(|n| complement(*n)).collect();
    let mut minmer = Cow::Borrowed(&seq[..length]);

    for (kmer, rc_kmer) in seq.windows(length).zip(reverse_complement.windows(length)) {
        if *kmer < minmer[..] {
            minmer = kmer.into();
        }
        if *rc_kmer < minmer[..] {
            minmer = rc_kmer.to_vec().into();
        }
    }
    minmer
}

/// A generic FASTX record that also abstracts over several logical operations
/// that can be performed on nucleic acid sequences.
pub trait Sequence<'a> {
    fn sequence(&'a self) -> &'a [u8];

    /// Remove newlines from the sequence; this handles `\r`, `\n`, and `\r\n`
    /// and removes internal newlines in addition to ones at the end.
    /// Primarily used for FASTA multiline records, but can also help process
    /// (the much rarer) multiline FASTQs. Always use before iteration methods
    /// below to ensure no newlines are being returned with e.g. `.kmers`.
    fn strip_returns(&'a self) -> Cow<'a, [u8]> {
        let seq = self.sequence();

        // first part is a fast check to see if we need to do any allocations
        let mut i;
        match memchr2(b'\r', b'\n', &seq) {
            Some(break_loc) => i = break_loc,
            None => return seq.into(),
        }
        // we found a newline; create a new buffer and stripping out newlines
        // and writing into it
        let mut new_buf = Vec::with_capacity(seq.len() - 1);
        new_buf.extend_from_slice(&seq[..i]);
        while i < seq.len() {
            match memchr2(b'\r', b'\n', &seq[i..]) {
                None => {
                    new_buf.extend_from_slice(&seq[i..]);
                    break;
                }
                Some(match_pos) => {
                    new_buf.extend_from_slice(&seq[i..i + match_pos]);
                    i += match_pos + 1;
                }
            }
        }
        new_buf.into()
    }

    /// Returns the reverse complement of a sequence. Biologically this is
    /// equivalent to the sequence of the strand opposite the one you pass
    /// in.
    ///
    /// ```
    /// use needletail::Sequence;
    ///
    /// assert_eq!(b"AACC".reverse_complement(), b"GGTT");
    /// ```
    fn reverse_complement(&'a self) -> Vec<u8> {
        self.sequence()
            .iter()
            .rev()
            .map(|n| complement(*n))
            .collect()
    }

    /// [Nucleic Acids] Normalizes the sequence. See documentation for
    /// `needletail::sequence::normalize`. Do not use on amino acid
    /// sequences. Note that this returns a Cow so you may have to coerce
    /// to a Vec<u8> or &[u8] as necessary.
    ///
    /// ```
    /// use needletail::Sequence;
    ///
    /// // IUPAC bases are coerced to N's if `false`
    /// assert_eq!(b"ADGH".normalize(false).as_ref(), b"ANGN");
    /// // otherwise they're preserved
    /// assert_eq!(b"ADGH".normalize(true).as_ref(), b"ADGH");
    ///
    /// // Uridine residues are converted to thymidine
    /// assert_eq!(b"ACGU".normalize(true).as_ref(), b"ACGT");
    /// ```
    fn normalize(&'a self, iupac: bool) -> Cow<'a, [u8]> {
        if let Some(s) = normalize(&self.sequence(), iupac) {
            s.into()
        } else {
            self.sequence().into()
        }
    }

    /// [Nucleic Acids] Returns an iterator over the sequence that skips
    /// non-ACGT bases and returns a tuple containing (position, the
    /// canonicalized kmer, if the sequence is the complement of the original).
    fn canonical_kmers(&'a self, k: u8, reverse_complement: &'a [u8]) -> CanonicalKmers<'a> {
        CanonicalKmers::new(self.sequence().as_ref(), reverse_complement, k)
    }

    /// Returns an iterator that returns a sliding window of k-sized
    /// sequences (k-mers). Does not skip whitespace or correct bases in the
    /// original sequence so `.normalize` or `.strip_returns` may be
    /// appropriate to use first.
    fn kmers(&'a self, k: u8) -> Kmers<'a> {
        Kmers::new(self.sequence().as_ref(), k)
    }

    /// Return an iterator that returns valid kmers in 4-bit form
    fn bit_kmers(&'a self, k: u8, canonical: bool) -> BitNuclKmer<'a> {
        BitNuclKmer::new(self.sequence(), k, canonical)
    }
}

impl<'a> Sequence<'a> for &'a [u8] {
    fn sequence(&'a self) -> &'a [u8] {
        &self
    }
}

impl<'a> Sequence<'a> for [u8] {
    fn sequence(&'a self) -> &'a [u8] {
        &self
    }
}

impl<'a> Sequence<'a> for Cow<'a, [u8]> {
    fn sequence(&'a self) -> &'a [u8] {
        &self
    }
}

/// [⚠Unstable] A trait to wrap over sequence data that has associated
/// quality information.
///
/// Will be stabilized once we figure out a good way to handle sequences that
/// have _optional_ quality information (like SequenceRecord) because the
/// return trait requires a slice from an immutable reference and
/// SequenceRecords can't return that without modifying themselves.
pub trait QualitySequence<'a>: Sequence<'a> {
    fn quality(&'a self) -> &'a [u8];

    /// Given a SeqRecord and a quality cutoff, mask out low-quality bases with
    /// `N` characters.
    fn quality_mask(&'a self, score: u8) -> Cow<'a, [u8]> {
        let qual = self.quality();
        // could maybe speed this up by doing a copy of base and then
        // iterating though qual and masking?
        let seq: Vec<u8> = self
            .sequence()
            .iter()
            .zip(qual.iter())
            .map(|(base, qual)| if *qual < score { b'N' } else { *base })
            .collect();
        seq.into()
    }
}

impl<'a> Sequence<'a> for (&'a [u8], &'a [u8]) {
    fn sequence(&'a self) -> &'a [u8] {
        &self.0
    }
}

impl<'a> QualitySequence<'a> for (&'a [u8], &'a [u8]) {
    fn quality(&'a self) -> &'a [u8] {
        &self.1
    }
}

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

    #[test]
    fn test_complement() {
        assert_eq!(complement(b'a'), b't');
        assert_eq!(complement(b'c'), b'g');
        assert_eq!(complement(b'g'), b'c');
        assert_eq!(complement(b'n'), b'n');
    }

    #[test]
    fn can_canonicalize() {
        assert!(canonical(b"A") == Cow::Borrowed(b"A"));
        assert!(canonical(b"T") == Cow::Owned::<[u8]>(b"A".to_vec()));
        assert!(canonical(b"AAGT") == Cow::Borrowed(b"AAGT"));
        assert!(canonical(b"ACTT") == Cow::Owned::<[u8]>(b"AAGT".to_vec()));
        assert!(canonical(b"GC") == Cow::Borrowed(b"GC"));
    }

    #[test]
    fn can_minimize() {
        let minmer = minimizer(&b"ATTTCG"[..], 3);
        assert_eq!(&minmer[..], b"AAA");
    }

    #[test]
    fn test_quality_mask() {
        let seq_rec = (&b"AGCT"[..], &b"AAA0"[..]);
        let filtered_rec = seq_rec.quality_mask(b'5');
        assert_eq!(&filtered_rec[..], &b"AGCN"[..]);
    }
}
<commit_msg>Move stray test<commit_after>//! Generic functions for working with (primarily nucleic acid) sequences
use std::borrow::Cow;

use memchr::memchr2;

use crate::bitkmer::BitNuclKmer;
use crate::kmer::{CanonicalKmers, Kmers};

/// Transform a nucleic acid sequence into its "normalized" form.
///
/// The normalized form is:
///  - only AGCTN and possibly - (for gaps)
///  - strip out any whitespace or line endings
///  - lowercase versions of these are uppercased
///  - U is converted to T (make everything a DNA sequence)
///  - some other punctuation is converted to gaps
///  - IUPAC bases may be converted to N's depending on the parameter passed in
///  - everything else is considered a N
pub fn normalize(seq: &[u8], allow_iupac: bool) -> Option<Vec<u8>> {
    let mut buf: Vec<u8> = Vec::with_capacity(seq.len());
    let mut changed: bool = false;

    for n in seq.iter() {
        let (new_char, char_changed) = match (*n, allow_iupac) {
            c @ (b'A', _)
            | c @ (b'C', _)
            | c @ (b'G', _)
            | c @ (b'T', _)
            | c @ (b'N', _)
            | c @ (b'-', _) => (c.0, false),
            (b'a', _) => (b'A', true),
            (b'c', _) => (b'C', true),
            (b'g', _) => (b'G', true),
            // normalize uridine to thymine
            (b't', _) | (b'u', _) | (b'U', _) => (b'T', true),
            // normalize gaps
            (b'.', _) | (b'~', _) => (b'-', true),
            // logic for IUPAC bases (a little messy)
            c @ (b'B', true)
            | c @ (b'D', true)
            | c @ (b'H', true)
            | c @ (b'V', true)
            | c @ (b'R', true)
            | c @ (b'Y', true)
            | c @ (b'S', true)
            | c @ (b'W', true)
            | c @ (b'K', true)
            | c @ (b'M', true) => (c.0, false),
            (b'b', true) => (b'B', true),
            (b'd', true) => (b'D', true),
            (b'h', true) => (b'H', true),
            (b'v', true) => (b'V', true),
            (b'r', true) => (b'R', true),
            (b'y', true) => (b'Y', true),
            (b's', true) => (b'S', true),
            (b'w', true) => (b'W', true),
            (b'k', true) => (b'K', true),
            (b'm', true) => (b'M', true),
            // remove all whitespace and line endings
            (b' ', _) | (b'\t', _) | (b'\r', _) | (b'\n', _) => (b' ', true),
            // everything else is an N
            _ => (b'N', true),
        };
        changed = changed || char_changed;
        if new_char != b' ' {
            buf.push(new_char);
        }
    }
    if changed {
        Some(buf)
    } else {
        None
    }
}

/// Returns the complementary base for a given IUPAC base code.
///
/// Does not work for RNA sequences (maybe we should raise an error or something?)
#[inline]
pub fn complement(n: u8) -> u8 {
    match n {
        b'a' => b't',
        b'A' => b'T',
        b'c' => b'g',
        b'C' => b'G',
        b'g' => b'c',
        b'G' => b'C',
        b't' => b'a',
        b'T' => b'A',

        // IUPAC codes
        b'r' => b'y',
        b'y' => b'r',
        b'k' => b'm',
        b'm' => b'k',
        b'b' => b'v',
        b'v' => b'b',
        b'd' => b'h',
        b'h' => b'd',
        b's' => b's',
        b'w' => b'w',
        b'R' => b'Y',
        b'Y' => b'R',
        b'K' => b'M',
        b'M' => b'K',
        b'B' => b'V',
        b'V' => b'B',
        b'D' => b'H',
        b'H' => b'D',
        b'S' => b'S',
        b'W' => b'W',

        // anything else just pass through
        // 'u' | 'U' => panic!("Does not support complements of U"),
        x => x,
    }
}

/// Taking in a sequence string, return the canonical form of the sequence
/// (e.g. the lexigraphically lowest of either the original sequence or its
/// reverse complement)
pub fn canonical(seq: &[u8]) -> Cow<[u8]> {
    let mut buf: Vec<u8> = Vec::with_capacity(seq.len());
    // enough just keeps our comparisons from happening after they need to
    let mut enough = false;
    let mut original_was_canonical = false;

    // loop through the kmer and its reverse complement simultaneously
    for (rn, n) in seq.iter().rev().map(|n| complement(*n)).zip(seq.iter()) {
        buf.push(rn);
        if !enough && (*n < rn) {
            original_was_canonical = true;
            break;
        } else if !enough && (rn < *n) {
            enough = true;
        }
        // unstated if branch: if rn == n, keep comparing
    }
    match (original_was_canonical, enough) {
        (true, true) => panic!("Bug: should never set original_was_canonical if enough == true"),
        (true, false) => seq.into(),
        (false, true) => buf.into(),
        // the sequences were completely equal, return the ref
        (false, false) => seq.into(),
    }
}

/// Find the lexigraphically smallest substring of `seq` of length `length`
///
/// There's probably a faster algorithm for this somewhere...
pub fn minimizer(seq: &[u8], length: usize) -> Cow<[u8]> {
    let reverse_complement: Vec<u8> = seq.iter().rev().map(|n| complement(*n)).collect();
    let mut minmer = Cow::Borrowed(&seq[..length]);

    for (kmer, rc_kmer) in seq.windows(length).zip(reverse_complement.windows(length)) {
        if *kmer < minmer[..] {
            minmer = kmer.into();
        }
        if *rc_kmer < minmer[..] {
            minmer = rc_kmer.to_vec().into();
        }
    }
    minmer
}

/// A generic FASTX record that also abstracts over several logical operations
/// that can be performed on nucleic acid sequences.
pub trait Sequence<'a> {
    fn sequence(&'a self) -> &'a [u8];

    /// Remove newlines from the sequence; this handles `\r`, `\n`, and `\r\n`
    /// and removes internal newlines in addition to ones at the end.
    /// Primarily used for FASTA multiline records, but can also help process
    /// (the much rarer) multiline FASTQs. Always use before iteration methods
    /// below to ensure no newlines are being returned with e.g. `.kmers`.
    fn strip_returns(&'a self) -> Cow<'a, [u8]> {
        let seq = self.sequence();

        // first part is a fast check to see if we need to do any allocations
        let mut i;
        match memchr2(b'\r', b'\n', &seq) {
            Some(break_loc) => i = break_loc,
            None => return seq.into(),
        }
        // we found a newline; create a new buffer and stripping out newlines
        // and writing into it
        let mut new_buf = Vec::with_capacity(seq.len() - 1);
        new_buf.extend_from_slice(&seq[..i]);
        while i < seq.len() {
            match memchr2(b'\r', b'\n', &seq[i..]) {
                None => {
                    new_buf.extend_from_slice(&seq[i..]);
                    break;
                }
                Some(match_pos) => {
                    new_buf.extend_from_slice(&seq[i..i + match_pos]);
                    i += match_pos + 1;
                }
            }
        }
        new_buf.into()
    }

    /// Returns the reverse complement of a sequence. Biologically this is
    /// equivalent to the sequence of the strand opposite the one you pass
    /// in.
    ///
    /// ```
    /// use needletail::Sequence;
    ///
    /// assert_eq!(b"AACC".reverse_complement(), b"GGTT");
    /// ```
    fn reverse_complement(&'a self) -> Vec<u8> {
        self.sequence()
            .iter()
            .rev()
            .map(|n| complement(*n))
            .collect()
    }

    /// [Nucleic Acids] Normalizes the sequence. See documentation for
    /// `needletail::sequence::normalize`. Do not use on amino acid
    /// sequences. Note that this returns a Cow so you may have to coerce
    /// to a Vec<u8> or &[u8] as necessary.
    ///
    /// ```
    /// use needletail::Sequence;
    ///
    /// // IUPAC bases are coerced to N's if `false`
    /// assert_eq!(b"ADGH".normalize(false).as_ref(), b"ANGN");
    /// // otherwise they're preserved
    /// assert_eq!(b"ADGH".normalize(true).as_ref(), b"ADGH");
    ///
    /// // Uridine residues are converted to thymidine
    /// assert_eq!(b"ACGU".normalize(true).as_ref(), b"ACGT");
    /// ```
    fn normalize(&'a self, iupac: bool) -> Cow<'a, [u8]> {
        if let Some(s) = normalize(&self.sequence(), iupac) {
            s.into()
        } else {
            self.sequence().into()
        }
    }

    /// [Nucleic Acids] Returns an iterator over the sequence that skips
    /// non-ACGT bases and returns a tuple containing (position, the
    /// canonicalized kmer, if the sequence is the complement of the original).
    fn canonical_kmers(&'a self, k: u8, reverse_complement: &'a [u8]) -> CanonicalKmers<'a> {
        CanonicalKmers::new(self.sequence().as_ref(), reverse_complement, k)
    }

    /// Returns an iterator that returns a sliding window of k-sized
    /// sequences (k-mers). Does not skip whitespace or correct bases in the
    /// original sequence so `.normalize` or `.strip_returns` may be
    /// appropriate to use first.
    fn kmers(&'a self, k: u8) -> Kmers<'a> {
        Kmers::new(self.sequence().as_ref(), k)
    }

    /// Return an iterator that returns valid kmers in 4-bit form
    fn bit_kmers(&'a self, k: u8, canonical: bool) -> BitNuclKmer<'a> {
        BitNuclKmer::new(self.sequence(), k, canonical)
    }
}

impl<'a> Sequence<'a> for &'a [u8] {
    fn sequence(&'a self) -> &'a [u8] {
        &self
    }
}

impl<'a> Sequence<'a> for [u8] {
    fn sequence(&'a self) -> &'a [u8] {
        &self
    }
}

impl<'a> Sequence<'a> for Cow<'a, [u8]> {
    fn sequence(&'a self) -> &'a [u8] {
        &self
    }
}

/// [⚠Unstable] A trait to wrap over sequence data that has associated
/// quality information.
///
/// Will be stabilized once we figure out a good way to handle sequences that
/// have _optional_ quality information (like SequenceRecord) because the
/// return trait requires a slice from an immutable reference and
/// SequenceRecords can't return that without modifying themselves.
pub trait QualitySequence<'a>: Sequence<'a> {
    fn quality(&'a self) -> &'a [u8];

    /// Given a SeqRecord and a quality cutoff, mask out low-quality bases with
    /// `N` characters.
    fn quality_mask(&'a self, score: u8) -> Cow<'a, [u8]> {
        let qual = self.quality();
        // could maybe speed this up by doing a copy of base and then
        // iterating though qual and masking?
        let seq: Vec<u8> = self
            .sequence()
            .iter()
            .zip(qual.iter())
            .map(|(base, qual)| if *qual < score { b'N' } else { *base })
            .collect();
        seq.into()
    }
}

impl<'a> Sequence<'a> for (&'a [u8], &'a [u8]) {
    fn sequence(&'a self) -> &'a [u8] {
        &self.0
    }
}

impl<'a> QualitySequence<'a> for (&'a [u8], &'a [u8]) {
    fn quality(&'a self) -> &'a [u8] {
        &self.1
    }
}

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

    #[test]
    fn test_normalize() {
        assert_eq!(normalize(b"ACGTU", false), Some(b"ACGTT".to_vec()));
        assert_eq!(normalize(b"acgtu", false), Some(b"ACGTT".to_vec()));

        assert_eq!(normalize(b"N.N-N~N N", false), Some(b"N-N-N-NN".to_vec()));

        assert_eq!(normalize(b"BDHVRYSWKM", true), None);
        assert_eq!(normalize(b"bdhvryswkm", true), Some(b"BDHVRYSWKM".to_vec()));
        assert_eq!(
            normalize(b"BDHVRYSWKM", false),
            Some(b"NNNNNNNNNN".to_vec())
        );
        assert_eq!(
            normalize(b"bdhvryswkm", false),
            Some(b"NNNNNNNNNN".to_vec())
        );
    }

    #[test]
    fn test_complement() {
        assert_eq!(complement(b'a'), b't');
        assert_eq!(complement(b'c'), b'g');
        assert_eq!(complement(b'g'), b'c');
        assert_eq!(complement(b'n'), b'n');
    }

    #[test]
    fn can_canonicalize() {
        assert!(canonical(b"A") == Cow::Borrowed(b"A"));
        assert!(canonical(b"T") == Cow::Owned::<[u8]>(b"A".to_vec()));
        assert!(canonical(b"AAGT") == Cow::Borrowed(b"AAGT"));
        assert!(canonical(b"ACTT") == Cow::Owned::<[u8]>(b"AAGT".to_vec()));
        assert!(canonical(b"GC") == Cow::Borrowed(b"GC"));
    }

    #[test]
    fn can_minimize() {
        let minmer = minimizer(&b"ATTTCG"[..], 3);
        assert_eq!(&minmer[..], b"AAA");
    }

    #[test]
    fn test_quality_mask() {
        let seq_rec = (&b"AGCT"[..], &b"AAA0"[..]);
        let filtered_rec = seq_rec.quality_mask(b'5');
        assert_eq!(&filtered_rec[..], &b"AGCN"[..]);
    }
}
<|endoftext|>