<commit_before>/*
 * Copyright 2016-2017 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! Geo queries

use serde::ser::{Serialize, Serializer, SerializeMap};

use ::json::{MergeSerialize, NoOuter, serialize_map_optional_kv, ShouldSkip};
use ::units::{Distance, DistanceType, GeoBox, Location};

use super::Query;
use super::common::FieldBasedQuery;

#[derive(Debug, Serialize)]
pub enum ShapeOption {
    #[serde(rename="shape")]
    Shape(Shape),
    #[serde(rename="indexed_shape")]
    IndexedShape(IndexedShape)
}

from!(Shape, ShapeOption, Shape);
from!(IndexedShape, ShapeOption, IndexedShape);

/// GeoShape query
#[derive(Debug, Serialize)]
pub struct GeoShapeQuery(FieldBasedQuery<Option<ShapeOption>, NoOuter>);

impl Query {
    pub fn build_geo_shape<A>(field: A) -> GeoShapeQuery
        where A: Into<String> {

        GeoShapeQuery(FieldBasedQuery::new(field.into(), None, NoOuter))
    }
}

impl GeoShapeQuery {
    pub fn with_shape<A>(mut self, shape: A) -> Self
        where A: Into<Shape> {

        self.0.inner = Some(ShapeOption::Shape(shape.into()));
        self
    }

    pub fn with_indexed_shape<A>(mut self, indexed_shape: A) -> Self
        where A: Into<IndexedShape> {

        self.0.inner = Some(ShapeOption::IndexedShape(indexed_shape.into()));
        self
    }

    build!(GeoShape);
}

// Required for GeoShape
#[derive(Debug, Serialize)]
pub struct Shape {
    #[serde(rename="type")]
    shape_type: String,
    coordinates: Vec<(f64, f64)>
}

impl Shape {
    pub fn new<A: Into<String>>(shape_type: A, coordinates: Vec<(f64, f64)>) -> Shape {
        Shape {
            shape_type:  shape_type.into(),
            coordinates: coordinates
        }
    }
}

#[derive(Debug, Serialize)]
pub struct IndexedShape {
    id:       String,
    doc_type: String,
    index:    String,
    path:     String
}

impl IndexedShape {
    pub fn new<A, B, C, D>(id: A, doc_type: B, index: C, path: D) -> IndexedShape
        where A: Into<String>,
              B: Into<String>,
              C: Into<String>,
              D: Into<String> {
        IndexedShape {
            id: id.into(),
            doc_type: doc_type.into(),
            index: index.into(),
            path: path.into()
        }
    }
}

/// Geo Bounding Box Query
#[derive(Debug, Serialize)]
pub struct GeoBoundingBoxQuery(FieldBasedQuery<GeoBoundingBoxQueryInner, NoOuter>);

#[derive(Debug, Default, Serialize)]
pub struct GeoBoundingBoxQueryInner {
    geo_box: GeoBox,
    #[serde(skip_serializing_if="ShouldSkip::should_skip")]
    coerce: Option<bool>,
    #[serde(skip_serializing_if="ShouldSkip::should_skip")]
    ignore_malformed: Option<bool>,
    #[serde(skip_serializing_if="ShouldSkip::should_skip", rename="type")]
    filter_type: Option<Type>
}

impl Query {
    pub fn build_geo_bounding_box<A, B>(field: A, geo_box: B) -> GeoBoundingBoxQuery
        where A: Into<String>,
              B: Into<GeoBox> {
        GeoBoundingBoxQuery(FieldBasedQuery::new(field.into(),
                                                 GeoBoundingBoxQueryInner {
                                                     geo_box: geo_box.into(),
                                                     ..Default::default()
                                                 },
                                                 NoOuter))
    }
}

impl GeoBoundingBoxQuery {
    add_inner_field!(with_coerce, coerce, bool);
    add_inner_field!(with_ignore_malformed, ignore_malformed, bool);
    add_inner_field!(with_type, filter_type, Type);

    build!(GeoBoundingBox);
}

/// Geo Bounding Box filter type
#[derive(Debug)]
pub enum Type {
    Indexed,
    Memory
}

impl Serialize for Type {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where S: Serializer {
        use self::Type::*;
        match self {
            &Indexed => "indexed",
            &Memory => "memory"
        }.serialize(serializer)
    }
}

/// Geo Distance query
///
/// TODO: Specific full unit test for querying with a generated query from here
#[derive(Debug, Serialize)]
pub struct GeoDistanceQuery(FieldBasedQuery<Location, GeoDistanceQueryOuter>);

#[derive(Debug, Default)]
struct GeoDistanceQueryOuter {
    distance: Distance,
    distance_type: Option<DistanceType>,
    optimize_bbox: Option<OptimizeBbox>,
    coerce: Option<bool>,
    ignore_malformed: Option<bool>
}

impl MergeSerialize for GeoDistanceQueryOuter {
    fn merge_serialize<S>(&self,
                          serializer: &mut S) -> Result<(), S::Error>
        where S: SerializeMap {

        serializer.serialize_entry("distance", &self.distance)?;
        serialize_map_optional_kv(serializer, "distance_type", &self.distance_type)?;
        serialize_map_optional_kv(serializer, "optimize_bbox", &self.optimize_bbox)?;
        serialize_map_optional_kv(serializer, "coerce", &self.coerce)?;
        serialize_map_optional_kv(serializer, "ignore_malformed", &self.ignore_malformed)?;
        Ok(())
    }
}

impl Query {
    pub fn build_geo_distance<A, B, C>(field: A,
                                       location: B,
                                       distance: C) -> GeoDistanceQuery
        where A: Into<String>,
              B: Into<Location>,
              C: Into<Distance> {
        let outer = GeoDistanceQueryOuter {
            distance: distance.into(),
            ..Default::default()
        };
        GeoDistanceQuery(FieldBasedQuery::new(field.into(), location.into(), outer))
    }
}

impl GeoDistanceQuery {
    add_outer_field!(with_distance_type, distance_type, DistanceType);
    add_outer_field!(with_optimize_bbox, optimize_bbox, OptimizeBbox);
    add_outer_field!(with_coerce, coerce, bool);
    add_outer_field!(with_ignore_malformed, ignore_malformed, bool);

    build!(GeoDistance);
}

/// Options for `optimize_bbox`
#[derive(Debug)]
pub enum OptimizeBbox {
    Memory,
    Indexed,
    None
}

impl Serialize for OptimizeBbox {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where S: Serializer {
        use self::OptimizeBbox::*;
        match self {
            &Memory => "memory".serialize(serializer),
            &Indexed => "indexed".serialize(serializer),
            &None => "none".serialize(serializer)
        }
    }
}

/// Geo Polygon query
#[derive(Debug, Serialize)]
pub struct GeoPolygonQuery(FieldBasedQuery<GeoPolygonQueryInner, NoOuter>);

#[derive(Debug, Default, Serialize)]
pub struct GeoPolygonQueryInner {
    points: Vec<Location>,
    #[serde(skip_serializing_if="ShouldSkip::should_skip")]
    coerce: Option<bool>,
    #[serde(skip_serializing_if="ShouldSkip::should_skip")]
    ignore_malformed: Option<bool>
}

impl Query {
    pub fn build_geo_polygon<A, B>(field: A,
                                   points: B) -> GeoPolygonQuery
        where A: Into<String>,
              B: Into<Vec<Location>> {
        GeoPolygonQuery(FieldBasedQuery::new(field.into(),
                                             GeoPolygonQueryInner {
                                                 points: points.into(),
                                                 ..Default::default()
                                             },
                                             NoOuter))
    }
}

impl GeoPolygonQuery {
    add_inner_field!(with_coerce, coerce, bool);
    add_inner_field!(with_ignore_malformed, ignore_malformed, bool);

    build!(GeoPolygon);
}

/// Geohash cell query
#[derive(Debug, Serialize)]
pub struct GeohashCellQuery(FieldBasedQuery<Location, GeohashCellQueryOuter>);

#[derive(Debug, Default)]
pub struct GeohashCellQueryOuter {
    precision: Option<Precision>,
    neighbors: Option<bool>,
}

impl MergeSerialize for GeohashCellQueryOuter {
    fn merge_serialize<S>(&self,
                          serializer: &mut S) -> Result<(), S::Error>
        where S: SerializeMap {

        match self.precision {
            Some(ref p) => {
                serializer.serialize_entry("precision", p)?;
            },
            None => ()
        };
        match self.neighbors {
            Some(b) => {
                serializer.serialize_entry("neighbors", &b)?;
            },
            None => ()
        }
        Ok(())
    }
}

impl Query {
    pub fn build_geohash_cell<A, B>(field: A, location: B) -> GeohashCellQuery
        where A: Into<String>,
              B: Into<Location> {
        GeohashCellQuery(FieldBasedQuery::new(field.into(),
                                              location.into(),
                                              Default::default()))
    }
}

impl GeohashCellQuery {
    add_outer_field!(with_precision, precision, Precision);
    add_outer_field!(with_neighbors, neighbors, bool);

    build!(GeohashCell);
}

#[derive(Debug)]
pub enum Precision {
    Geohash(u64),
    Distance(Distance)
}

impl Default for Precision {
    fn default() -> Self {
        Precision::Distance(Default::default())
    }
}

from!(u64, Precision, Geohash);
from!(Distance, Precision, Distance);

impl Serialize for Precision {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where S: Serializer {
        use self::Precision::*;
        match self {
            &Geohash(precision) => precision.serialize(serializer),
            &Distance(ref dist) => dist.serialize(serializer)
        }
    }
}
<commit_msg>Reduce duplication<commit_after>/*
 * Copyright 2016-2017 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! Geo queries

use serde::ser::{Serialize, Serializer, SerializeMap};

use ::json::{MergeSerialize, NoOuter, serialize_map_optional_kv, ShouldSkip};
use ::units::{Distance, DistanceType, GeoBox, Location};

use super::Query;
use super::common::FieldBasedQuery;

#[derive(Debug, Serialize)]
pub enum ShapeOption {
    #[serde(rename="shape")]
    Shape(Shape),
    #[serde(rename="indexed_shape")]
    IndexedShape(IndexedShape)
}

from!(Shape, ShapeOption, Shape);
from!(IndexedShape, ShapeOption, IndexedShape);

/// GeoShape query
#[derive(Debug, Serialize)]
pub struct GeoShapeQuery(FieldBasedQuery<Option<ShapeOption>, NoOuter>);

impl Query {
    pub fn build_geo_shape<A>(field: A) -> GeoShapeQuery
        where A: Into<String> {

        GeoShapeQuery(FieldBasedQuery::new(field.into(), None, NoOuter))
    }
}

impl GeoShapeQuery {
    pub fn with_shape<A>(mut self, shape: A) -> Self
        where A: Into<Shape> {

        self.0.inner = Some(ShapeOption::Shape(shape.into()));
        self
    }

    pub fn with_indexed_shape<A>(mut self, indexed_shape: A) -> Self
        where A: Into<IndexedShape> {

        self.0.inner = Some(ShapeOption::IndexedShape(indexed_shape.into()));
        self
    }

    build!(GeoShape);
}

// Required for GeoShape
#[derive(Debug, Serialize)]
pub struct Shape {
    #[serde(rename="type")]
    shape_type: String,
    coordinates: Vec<(f64, f64)>
}

impl Shape {
    pub fn new<A: Into<String>>(shape_type: A, coordinates: Vec<(f64, f64)>) -> Shape {
        Shape {
            shape_type:  shape_type.into(),
            coordinates: coordinates
        }
    }
}

#[derive(Debug, Serialize)]
pub struct IndexedShape {
    id:       String,
    doc_type: String,
    index:    String,
    path:     String
}

impl IndexedShape {
    pub fn new<A, B, C, D>(id: A, doc_type: B, index: C, path: D) -> IndexedShape
        where A: Into<String>,
              B: Into<String>,
              C: Into<String>,
              D: Into<String> {
        IndexedShape {
            id: id.into(),
            doc_type: doc_type.into(),
            index: index.into(),
            path: path.into()
        }
    }
}

/// Geo Bounding Box Query
#[derive(Debug, Serialize)]
pub struct GeoBoundingBoxQuery(FieldBasedQuery<GeoBoundingBoxQueryInner, NoOuter>);

#[derive(Debug, Default, Serialize)]
pub struct GeoBoundingBoxQueryInner {
    geo_box: GeoBox,
    #[serde(skip_serializing_if="ShouldSkip::should_skip")]
    coerce: Option<bool>,
    #[serde(skip_serializing_if="ShouldSkip::should_skip")]
    ignore_malformed: Option<bool>,
    #[serde(skip_serializing_if="ShouldSkip::should_skip", rename="type")]
    filter_type: Option<Type>
}

impl Query {
    pub fn build_geo_bounding_box<A, B>(field: A, geo_box: B) -> GeoBoundingBoxQuery
        where A: Into<String>,
              B: Into<GeoBox> {
        GeoBoundingBoxQuery(FieldBasedQuery::new(field.into(),
                                                 GeoBoundingBoxQueryInner {
                                                     geo_box: geo_box.into(),
                                                     ..Default::default()
                                                 },
                                                 NoOuter))
    }
}

impl GeoBoundingBoxQuery {
    add_inner_field!(with_coerce, coerce, bool);
    add_inner_field!(with_ignore_malformed, ignore_malformed, bool);
    add_inner_field!(with_type, filter_type, Type);

    build!(GeoBoundingBox);
}

/// Geo Bounding Box filter type
#[derive(Debug)]
pub enum Type {
    Indexed,
    Memory
}

impl Serialize for Type {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where S: Serializer {
        use self::Type::*;
        match self {
            &Indexed => "indexed",
            &Memory => "memory"
        }.serialize(serializer)
    }
}

/// Geo Distance query
///
/// TODO: Specific full unit test for querying with a generated query from here
#[derive(Debug, Serialize)]
pub struct GeoDistanceQuery(FieldBasedQuery<Location, GeoDistanceQueryOuter>);

#[derive(Debug, Default)]
struct GeoDistanceQueryOuter {
    distance: Distance,
    distance_type: Option<DistanceType>,
    optimize_bbox: Option<OptimizeBbox>,
    coerce: Option<bool>,
    ignore_malformed: Option<bool>
}

impl MergeSerialize for GeoDistanceQueryOuter {
    fn merge_serialize<S>(&self,
                          serializer: &mut S) -> Result<(), S::Error>
        where S: SerializeMap {

        serializer.serialize_entry("distance", &self.distance)?;
        serialize_map_optional_kv(serializer, "distance_type", &self.distance_type)?;
        serialize_map_optional_kv(serializer, "optimize_bbox", &self.optimize_bbox)?;
        serialize_map_optional_kv(serializer, "coerce", &self.coerce)?;
        serialize_map_optional_kv(serializer, "ignore_malformed", &self.ignore_malformed)?;
        Ok(())
    }
}

impl Query {
    pub fn build_geo_distance<A, B, C>(field: A,
                                       location: B,
                                       distance: C) -> GeoDistanceQuery
        where A: Into<String>,
              B: Into<Location>,
              C: Into<Distance> {
        let outer = GeoDistanceQueryOuter {
            distance: distance.into(),
            ..Default::default()
        };
        GeoDistanceQuery(FieldBasedQuery::new(field.into(), location.into(), outer))
    }
}

impl GeoDistanceQuery {
    add_outer_field!(with_distance_type, distance_type, DistanceType);
    add_outer_field!(with_optimize_bbox, optimize_bbox, OptimizeBbox);
    add_outer_field!(with_coerce, coerce, bool);
    add_outer_field!(with_ignore_malformed, ignore_malformed, bool);

    build!(GeoDistance);
}

/// Options for `optimize_bbox`
#[derive(Debug)]
pub enum OptimizeBbox {
    Memory,
    Indexed,
    None
}

impl Serialize for OptimizeBbox {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where S: Serializer {
        use self::OptimizeBbox::*;
        match self {
            &Memory => "memory".serialize(serializer),
            &Indexed => "indexed".serialize(serializer),
            &None => "none".serialize(serializer)
        }
    }
}

/// Geo Polygon query
#[derive(Debug, Serialize)]
pub struct GeoPolygonQuery(FieldBasedQuery<GeoPolygonQueryInner, NoOuter>);

#[derive(Debug, Default, Serialize)]
pub struct GeoPolygonQueryInner {
    points: Vec<Location>,
    #[serde(skip_serializing_if="ShouldSkip::should_skip")]
    coerce: Option<bool>,
    #[serde(skip_serializing_if="ShouldSkip::should_skip")]
    ignore_malformed: Option<bool>
}

impl Query {
    pub fn build_geo_polygon<A, B>(field: A,
                                   points: B) -> GeoPolygonQuery
        where A: Into<String>,
              B: Into<Vec<Location>> {
        GeoPolygonQuery(FieldBasedQuery::new(field.into(),
                                             GeoPolygonQueryInner {
                                                 points: points.into(),
                                                 ..Default::default()
                                             },
                                             NoOuter))
    }
}

impl GeoPolygonQuery {
    add_inner_field!(with_coerce, coerce, bool);
    add_inner_field!(with_ignore_malformed, ignore_malformed, bool);

    build!(GeoPolygon);
}

/// Geohash cell query
#[derive(Debug, Serialize)]
pub struct GeohashCellQuery(FieldBasedQuery<Location, GeohashCellQueryOuter>);

#[derive(Debug, Default)]
pub struct GeohashCellQueryOuter {
    precision: Option<Precision>,
    neighbors: Option<bool>,
}

impl MergeSerialize for GeohashCellQueryOuter {
    fn merge_serialize<S>(&self,
                          serializer: &mut S) -> Result<(), S::Error>
        where S: SerializeMap {

        serialize_map_optional_kv(serializer, "precision", &self.precision)?;
        serialize_map_optional_kv(serializer, "neighbors", &self.neighbors)?;
        Ok(())
    }
}

impl Query {
    pub fn build_geohash_cell<A, B>(field: A, location: B) -> GeohashCellQuery
        where A: Into<String>,
              B: Into<Location> {
        GeohashCellQuery(FieldBasedQuery::new(field.into(),
                                              location.into(),
                                              Default::default()))
    }
}

impl GeohashCellQuery {
    add_outer_field!(with_precision, precision, Precision);
    add_outer_field!(with_neighbors, neighbors, bool);

    build!(GeohashCell);
}

#[derive(Debug)]
pub enum Precision {
    Geohash(u64),
    Distance(Distance)
}

impl Default for Precision {
    fn default() -> Self {
        Precision::Distance(Default::default())
    }
}

from!(u64, Precision, Geohash);
from!(Distance, Precision, Distance);

impl Serialize for Precision {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where S: Serializer {
        use self::Precision::*;
        match self {
            &Geohash(precision) => precision.serialize(serializer),
            &Distance(ref dist) => dist.serialize(serializer)
        }
    }
}
<|endoftext|>


<commit_before>use std::sync::{Arc};
use std::net::IpAddr;
use std::collections::{HashMap, HashSet};
use std::default::Default;
use std::sync::atomic::{AtomicIsize, Ordering};

use mioco::sync::Mutex;
use trust_dns::op::{Message, ResponseCode, Query};
use trust_dns::rr::{DNSClass, Name, RecordType, RData, Record};
use itertools::Itertools;

use utils::{Result, CloneExt, MessageExt, Future, AsDisplay};
use query::{query_multiple_handle_futures, query as query_one};
use cache::{Cache, RecordSource};
use nscache::NsCache;
use config::Config;
use resolver::{ErrorKind, RcResolver};

#[derive(Eq, PartialEq, Hash, Clone, Debug)]
enum QueriedItem {
    NS(IpAddr),
    CNAME(Name),
    NSDomain(Name),
}
struct RecursiveResolverState {
    queried_items: Mutex<HashSet<QueriedItem>>,
    query: Query,
    parent: RcResolver,
    query_limit: AtomicIsize,
    skip_cache: bool,
}
impl RecursiveResolverState {
    fn new(q: &Query, parent: RcResolver, skip_cache: bool) -> Self {
        RecursiveResolverState {
            queried_items: Mutex::new(Default::default()),
            query: q.clone(),
            parent: parent,
            query_limit: AtomicIsize::new(256),
            skip_cache: skip_cache,
        }
    }
    fn new_inner(&self, q: &Query) -> Self {
        RecursiveResolverState {
            queried_items: Mutex::new(Default::default()),
            query: q.clone(),
            parent: self.parent.clone(),
            query_limit: AtomicIsize::new(self.query_limit.load(Ordering::Relaxed)),
            skip_cache: self.skip_cache,
        }
    }
}
#[derive(Clone)]
pub struct RecursiveResolver {
    state: Arc<RecursiveResolverState>,
}
impl RecursiveResolver {
    fn get_cache(&self) -> &Cache {
        &self.state.parent.cache
    }
    fn get_ns_cache(&self) -> &NsCache {
        &self.state.parent.ns_cache
    }
    fn get_config(&self) -> &Config {
        &self.state.parent.config
    }
    fn query(&self) -> Result<Message> {
        let ns = self.get_ns_cache().lookup_recursive(self.state.query.get_name());
        self.query_ns_multiple(ns, None)
    }
    fn query_ns_domain(&self, auth_zone: &Name, ns: &Name) -> Result<Message> {
        debug!("Querying IP of NS {} in zone {}", ns, auth_zone);
        let mut query = Query::new();
        query.name(ns.clone())
        .query_type(RecordType::A)
        .query_class(DNSClass::IN);
        // TODO: IPv6
        let result = try!(self.resolve_next(&query));
        let ips = result.get_answers().iter().filter_map(|x| match *x.get_rdata() {
            RData::A {ref address} => Some((IpAddr::V4(*address), x.get_ttl())),
            _ => None,
        }).collect_vec();
        self.get_ns_cache().update(auth_zone, ips.iter().map(|&(ip, ttl)| (ip, ns.clone(), ttl as u64)));
        self.query_ns_multiple(ips.iter().map(|&(ip, _)| ip), None)
    }
    fn query_ns_domain_future(&self, auth_zone: &Name, ns: Name) -> Option<Future<Result<Message>>> {
        if !self.register_query(QueriedItem::NSDomain(ns.clone())) {
            return None;
        }
        let inst = self.clone();
        let auth_zone_clone = auth_zone.clone();
        Some(Future::from_fn(move || inst.query_ns_domain(&auth_zone_clone, &ns)))
    }
    fn query_ns_domain_futures<T>(&self, auth_zone: &Name, ns_iter: T) -> Vec<Future<Result<Message>>> where T: IntoIterator<Item=Name> {
        ns_iter.into_iter()
        .filter_map(|x| self.query_ns_domain_future(auth_zone, x))
        .collect()
    }
    fn query_ns_multiple<TNs, TExtra>(&self, ns: TNs, extra_futures: TExtra) -> Result<Message>
        where TNs: IntoIterator<Item=IpAddr>,
              TExtra: IntoIterator<Item=Future<Result<Message>>>,
    {
        let mut ns_iter = ns.into_iter().peekable();
        if ns_iter.peek().is_none() {
            return Err(ErrorKind::NoNameserver.into());
        }
        let mut futures: Vec<_> = ns_iter
        .filter_map(|x| self.query_ns_future(&x))
        .collect();
        futures.extend(extra_futures);
        if futures.is_empty() {
            return Err(ErrorKind::LostRace.into());
        }
        query_multiple_handle_futures(&mut futures)
    }
    fn register_query(&self, item: QueriedItem) -> bool {
        let mut guard = self.state.queried_items.lock().unwrap();
        if guard.contains(&item) {
            trace!("Already queried {:?}", item);
            return false;
        }
        if self.state.query_limit.fetch_sub(1, Ordering::Relaxed) <= 0 {
            debug_assert!(false);
            warn!("Already queried too many servers for {:?}, bug or attack?", self.state.query);
            return false;
        }
        guard.insert(item);
        true
    }
    fn query_ns_future(&self, ns: &IpAddr) -> Option<Future<Result<Message>>> {
        if !self.register_query(QueriedItem::NS(*ns)) {
            return None;
        }
        let inst = self.clone();
        let ns_clone = *ns;
        Some(Future::from_fn(move || inst.query_ns(ns_clone)))
    }
    fn query_ns(self, ns: IpAddr) -> Result<Message> {
        let result = try!(query_one(self.state.query.clone(), ns, *self.get_config().query.timeout));
        if result.get_response_code() != ResponseCode::NoError {
            return Ok(result);
        }
        if !result.get_answers().is_empty() {
            return self.handle_normal_resp(result);
        }
        if result.get_name_servers().iter().any(|x| x.get_rr_type() == RecordType::SOA) {
            // Domain exists, but no record
            return Ok(result)
        }
        self.handle_ns_referral(result)
    }
    #[allow(similar_names)]
    fn find_unresolved_cnames(&self, msg: &Message) -> Option<Vec<Record>> {
        let mut name = self.state.query.get_name();
        let query_type = self.state.query.get_query_type();
        let mut unresolved_cnames = Vec::<&Record>::new();
        loop {
            let cnames = {
                let matched_records = || msg.get_answers().iter()
                .filter(|x| x.get_name() == name);
                let have_real_record = matched_records()
                .any(|x| x.get_rr_type() == query_type);
                if have_real_record {
                    return None;
                }
                matched_records().filter(|x| x.get_rr_type() == RecordType::CNAME)
                .collect_vec()
            };
            if cnames.is_empty() {
                break;
            }
            if cnames.len() > 1 {
                debug_assert!(false);
                warn!("Server returned multiple CNAMEs for a single domain");
                // Try to resolve anyways
            }
            unresolved_cnames.push(cnames[0]);
            name = match *cnames[0].get_rdata() {
                RData::CNAME {ref cname} => cname,
                _ => panic!("Record type doesn't match RData"),
            };
        }
        if unresolved_cnames.is_empty() {
            // FIXME: This means the message has unrelated records in answer section?
            return None;
        }
        Some(unresolved_cnames.iter().map(|x| (*x).clone()).collect())
    }
    #[allow(similar_names)]
    fn handle_normal_resp(&self, msg: Message) -> Result<Message> {
        if self.state.query.get_query_type() == RecordType::CNAME {
            // No need to resolve CNAME chain(?)
            return Ok(msg);
        }
        if let Some(unresolved_cnames) = self.find_unresolved_cnames(&msg) {
            let final_record = &unresolved_cnames[unresolved_cnames.len() - 1];
            let next_name = match *final_record.get_rdata() {
                RData::CNAME {ref cname} => cname.clone(),
                _ => panic!("Record type doesn't match RData"),
            };
            if !self.register_query(QueriedItem::CNAME(next_name.clone())) {
                return Err(ErrorKind::LostRace.into());
            }
            debug!("[{}] CNAME referral: {}", self.state.query.get_name(), next_name);
            let mut next_query = self.state.query.clone();
            next_query.name(next_name);
            match self.resolve_next(&next_query) {
                Err(e) => {
                    warn!("Failed to resolve CNAME {} for {}: {:?}",
                          next_query.get_name(), self.state.query.as_disp(), e);
                    return Ok(msg);
                },
                Ok(ref next_msg) if
                    next_msg.get_response_code() != ResponseCode::NoError
                => {
                    warn!("Failed to resolve CNAME {}: Server returned error: {:?}",
                          next_query.get_name(), next_msg.get_response_code());
                    return Ok(msg);
                }
                Ok(next_msg) => {
                    // Simpler way to clear all answers
                    let mut new_msg = msg.truncate();
                    assert!(new_msg.get_queries().is_empty());
                    assert!(new_msg.get_answers().is_empty());
                    assert!(new_msg.get_name_servers().is_empty());
                    assert!(new_msg.get_additional().is_empty());
                    new_msg.truncated(false);
                    new_msg.add_query(self.state.query.clone());
                    unresolved_cnames.iter().cloned().foreach(|x| {
                        new_msg.add_answer(x);
                    });
                    new_msg.add_all_answers(next_msg.get_answers());
                    return Ok(new_msg);
                },
            };
        }
        Ok(msg)
    }
    fn handle_ns_referral(&self, msg: Message) -> Result<Message> {
        let mut ns_domains : HashMap<_, _> = msg.get_name_servers().iter()
        .filter_map(|x| match *x.get_rdata() {
            RData::NS {ref nsdname} => Some((nsdname.clone(), x.get_name().clone())),
            _ => None,
        })
        .collect();
        if ns_domains.is_empty() {
            debug!("Nameserver returned NOERROR and empty answer: {}, {}",
                   msg.get_queries()[0].as_disp(), msg.as_disp());
            return Ok(msg);
        }
        if ns_domains.values().dedup().count() > 1 {
            // FIXME: Is this legal?
            debug_assert!(false);
            warn!("Nameserver returned NS records for multiple zones");
            return Err(ErrorKind::InsaneNsReferral.into());
        }
        let referred_zone = ns_domains.values().next().unwrap().clone();
        if !referred_zone.zone_of(self.state.query.get_name()) {
            warn!("Nameserver returned NS records for incorrect zone");
            return Err(ErrorKind::InsaneNsReferral.into());
        }
        debug!("[{}] NS referral: {}", self.state.query.get_name(), referred_zone);

        // Get IPs of all NSes.
        let ns_items: Vec<_> = msg.get_additional().iter()
        .filter(|x| ns_domains.contains_key(x.get_name()))
        .filter_map(|x| match *x.get_rdata() {
            RData::A {ref address} => Some((
                IpAddr::V4(*address), x.get_name().clone(), x.get_ttl(),
            )),
            RData::AAAA {ref address} => Some((
                IpAddr::V6(*address), x.get_name().clone(), x.get_ttl(),
            )),
            _ => None,
        })
        .collect();
        self.get_ns_cache().update(
            &referred_zone,
            ns_items.iter().map(|&(ip, ref name, ttl)| {
                ns_domains.remove(name);
                (ip, name.clone(), ttl as u64)
            }),
        );
        let orphan_domain_futures = self.query_ns_domain_futures(
            &referred_zone, ns_domains.keys().cloned(),
        );
        self.query_ns_multiple(
            ns_items.iter().map(|x| x.0),
            orphan_domain_futures,
        )
    }
    fn update_cache(&self, msg: &Message) {
        self.get_cache().update_from_message(msg, RecordSource::Recursive);
    }
    pub fn resolve(q: &Query, parent: RcResolver, skip_cache: bool) -> Result<Message> {
        let resolver = RecursiveResolver {
            state: Arc::new(RecursiveResolverState::new(q, parent, skip_cache)),
        };
        let ret = try!(resolver.query());
        resolver.update_cache(&ret);
        Ok(ret)
    }
    fn resolve_next(&self, q: &Query) -> Result<Message> {
        if !self.state.skip_cache {
            if let Some(msg) = self.get_cache().lookup_with_type(
                q.get_name(), q.get_query_type(), |m| m.clone_resp_for(q),
            ) {
                return Ok(msg);
            }
        }
        let resolver = RecursiveResolver {
            state: Arc::new(self.state.new_inner(q)),
        };
        let ret = try!(resolver.query());
        self.update_cache(&ret);
        Ok(ret)
    }
}
<commit_msg>Fix incorrect check<commit_after>use std::sync::{Arc};
use std::net::IpAddr;
use std::collections::{HashMap, HashSet};
use std::default::Default;
use std::sync::atomic::{AtomicIsize, Ordering};

use mioco::sync::Mutex;
use trust_dns::op::{Message, ResponseCode, Query};
use trust_dns::rr::{DNSClass, Name, RecordType, RData, Record};
use itertools::Itertools;

use utils::{Result, CloneExt, MessageExt, Future, AsDisplay};
use query::{query_multiple_handle_futures, query as query_one};
use cache::{Cache, RecordSource};
use nscache::NsCache;
use config::Config;
use resolver::{ErrorKind, RcResolver};

#[derive(Eq, PartialEq, Hash, Clone, Debug)]
enum QueriedItem {
    NS(IpAddr),
    CNAME(Name),
    NSDomain(Name),
}
struct RecursiveResolverState {
    queried_items: Mutex<HashSet<QueriedItem>>,
    query: Query,
    parent: RcResolver,
    query_limit: AtomicIsize,
    skip_cache: bool,
}
impl RecursiveResolverState {
    fn new(q: &Query, parent: RcResolver, skip_cache: bool) -> Self {
        RecursiveResolverState {
            queried_items: Mutex::new(Default::default()),
            query: q.clone(),
            parent: parent,
            query_limit: AtomicIsize::new(256),
            skip_cache: skip_cache,
        }
    }
    fn new_inner(&self, q: &Query) -> Self {
        RecursiveResolverState {
            queried_items: Mutex::new(Default::default()),
            query: q.clone(),
            parent: self.parent.clone(),
            query_limit: AtomicIsize::new(self.query_limit.load(Ordering::Relaxed)),
            skip_cache: self.skip_cache,
        }
    }
}
#[derive(Clone)]
pub struct RecursiveResolver {
    state: Arc<RecursiveResolverState>,
}
impl RecursiveResolver {
    fn get_cache(&self) -> &Cache {
        &self.state.parent.cache
    }
    fn get_ns_cache(&self) -> &NsCache {
        &self.state.parent.ns_cache
    }
    fn get_config(&self) -> &Config {
        &self.state.parent.config
    }
    fn query(&self) -> Result<Message> {
        let ns = self.get_ns_cache().lookup_recursive(self.state.query.get_name());
        self.query_ns_multiple(ns, None)
    }
    fn query_ns_domain(&self, auth_zone: &Name, ns: &Name) -> Result<Message> {
        debug!("Querying IP of NS {} in zone {}", ns, auth_zone);
        let mut query = Query::new();
        query.name(ns.clone())
        .query_type(RecordType::A)
        .query_class(DNSClass::IN);
        // TODO: IPv6
        let result = try!(self.resolve_next(&query));
        let ips = result.get_answers().iter().filter_map(|x| match *x.get_rdata() {
            RData::A {ref address} => Some((IpAddr::V4(*address), x.get_ttl())),
            _ => None,
        }).collect_vec();
        self.get_ns_cache().update(auth_zone, ips.iter().map(|&(ip, ttl)| (ip, ns.clone(), ttl as u64)));
        self.query_ns_multiple(ips.iter().map(|&(ip, _)| ip), None)
    }
    fn query_ns_domain_future(&self, auth_zone: &Name, ns: Name) -> Option<Future<Result<Message>>> {
        if !self.register_query(QueriedItem::NSDomain(ns.clone())) {
            return None;
        }
        let inst = self.clone();
        let auth_zone_clone = auth_zone.clone();
        Some(Future::from_fn(move || inst.query_ns_domain(&auth_zone_clone, &ns)))
    }
    fn query_ns_domain_futures<T>(&self, auth_zone: &Name, ns_iter: T) -> Vec<Future<Result<Message>>> where T: IntoIterator<Item=Name> {
        ns_iter.into_iter()
        .filter_map(|x| self.query_ns_domain_future(auth_zone, x))
        .collect()
    }
    fn query_ns_multiple<TNs, TExtra>(&self, ns: TNs, extra_futures: TExtra) -> Result<Message>
        where TNs: IntoIterator<Item=IpAddr>,
              TExtra: IntoIterator<Item=Future<Result<Message>>>,
    {
        let mut ns_iter = ns.into_iter().peekable();
        let mut extra_iter = extra_futures.into_iter().peekable();
        if ns_iter.peek().is_none() && extra_iter.peek().is_none() {
            return Err(ErrorKind::NoNameserver.into());
        }
        let mut futures: Vec<_> = ns_iter
        .filter_map(|x| self.query_ns_future(&x))
        .collect();
        futures.extend(extra_iter);
        if futures.is_empty() {
            return Err(ErrorKind::LostRace.into());
        }
        query_multiple_handle_futures(&mut futures)
    }
    fn register_query(&self, item: QueriedItem) -> bool {
        let mut guard = self.state.queried_items.lock().unwrap();
        if guard.contains(&item) {
            trace!("Already queried {:?}", item);
            return false;
        }
        if self.state.query_limit.fetch_sub(1, Ordering::Relaxed) <= 0 {
            debug_assert!(false);
            warn!("Already queried too many servers for {:?}, bug or attack?", self.state.query);
            return false;
        }
        guard.insert(item);
        true
    }
    fn query_ns_future(&self, ns: &IpAddr) -> Option<Future<Result<Message>>> {
        if !self.register_query(QueriedItem::NS(*ns)) {
            return None;
        }
        let inst = self.clone();
        let ns_clone = *ns;
        Some(Future::from_fn(move || inst.query_ns(ns_clone)))
    }
    fn query_ns(self, ns: IpAddr) -> Result<Message> {
        let result = try!(query_one(self.state.query.clone(), ns, *self.get_config().query.timeout));
        if result.get_response_code() != ResponseCode::NoError {
            return Ok(result);
        }
        if !result.get_answers().is_empty() {
            return self.handle_normal_resp(result);
        }
        if result.get_name_servers().iter().any(|x| x.get_rr_type() == RecordType::SOA) {
            // Domain exists, but no record
            return Ok(result)
        }
        self.handle_ns_referral(result)
    }
    #[allow(similar_names)]
    fn find_unresolved_cnames(&self, msg: &Message) -> Option<Vec<Record>> {
        let mut name = self.state.query.get_name();
        let query_type = self.state.query.get_query_type();
        let mut unresolved_cnames = Vec::<&Record>::new();
        loop {
            let cnames = {
                let matched_records = || msg.get_answers().iter()
                .filter(|x| x.get_name() == name);
                let have_real_record = matched_records()
                .any(|x| x.get_rr_type() == query_type);
                if have_real_record {
                    return None;
                }
                matched_records().filter(|x| x.get_rr_type() == RecordType::CNAME)
                .collect_vec()
            };
            if cnames.is_empty() {
                break;
            }
            if cnames.len() > 1 {
                debug_assert!(false);
                warn!("Server returned multiple CNAMEs for a single domain");
                // Try to resolve anyways
            }
            unresolved_cnames.push(cnames[0]);
            name = match *cnames[0].get_rdata() {
                RData::CNAME {ref cname} => cname,
                _ => panic!("Record type doesn't match RData"),
            };
        }
        if unresolved_cnames.is_empty() {
            // FIXME: This means the message has unrelated records in answer section?
            return None;
        }
        Some(unresolved_cnames.iter().map(|x| (*x).clone()).collect())
    }
    #[allow(similar_names)]
    fn handle_normal_resp(&self, msg: Message) -> Result<Message> {
        if self.state.query.get_query_type() == RecordType::CNAME {
            // No need to resolve CNAME chain(?)
            return Ok(msg);
        }
        if let Some(unresolved_cnames) = self.find_unresolved_cnames(&msg) {
            let final_record = &unresolved_cnames[unresolved_cnames.len() - 1];
            let next_name = match *final_record.get_rdata() {
                RData::CNAME {ref cname} => cname.clone(),
                _ => panic!("Record type doesn't match RData"),
            };
            if !self.register_query(QueriedItem::CNAME(next_name.clone())) {
                return Err(ErrorKind::LostRace.into());
            }
            debug!("[{}] CNAME referral: {}", self.state.query.get_name(), next_name);
            let mut next_query = self.state.query.clone();
            next_query.name(next_name);
            match self.resolve_next(&next_query) {
                Err(e) => {
                    warn!("Failed to resolve CNAME {} for {}: {:?}",
                          next_query.get_name(), self.state.query.as_disp(), e);
                    return Ok(msg);
                },
                Ok(ref next_msg) if
                    next_msg.get_response_code() != ResponseCode::NoError
                => {
                    warn!("Failed to resolve CNAME {}: Server returned error: {:?}",
                          next_query.get_name(), next_msg.get_response_code());
                    return Ok(msg);
                }
                Ok(next_msg) => {
                    // Simpler way to clear all answers
                    let mut new_msg = msg.truncate();
                    assert!(new_msg.get_queries().is_empty());
                    assert!(new_msg.get_answers().is_empty());
                    assert!(new_msg.get_name_servers().is_empty());
                    assert!(new_msg.get_additional().is_empty());
                    new_msg.truncated(false);
                    new_msg.add_query(self.state.query.clone());
                    unresolved_cnames.iter().cloned().foreach(|x| {
                        new_msg.add_answer(x);
                    });
                    new_msg.add_all_answers(next_msg.get_answers());
                    return Ok(new_msg);
                },
            };
        }
        Ok(msg)
    }
    fn handle_ns_referral(&self, msg: Message) -> Result<Message> {
        let mut ns_domains : HashMap<_, _> = msg.get_name_servers().iter()
        .filter_map(|x| match *x.get_rdata() {
            RData::NS {ref nsdname} => Some((nsdname.clone(), x.get_name().clone())),
            _ => None,
        })
        .collect();
        if ns_domains.is_empty() {
            debug!("Nameserver returned NOERROR and empty answer: {}, {}",
                   msg.get_queries()[0].as_disp(), msg.as_disp());
            return Ok(msg);
        }
        if ns_domains.values().dedup().count() > 1 {
            // FIXME: Is this legal?
            debug_assert!(false);
            warn!("Nameserver returned NS records for multiple zones");
            return Err(ErrorKind::InsaneNsReferral.into());
        }
        let referred_zone = ns_domains.values().next().unwrap().clone();
        if !referred_zone.zone_of(self.state.query.get_name()) {
            warn!("Nameserver returned NS records for incorrect zone");
            return Err(ErrorKind::InsaneNsReferral.into());
        }
        debug!("[{}] NS referral: {}", self.state.query.get_name(), referred_zone);

        // Get IPs of all NSes.
        let ns_items: Vec<_> = msg.get_additional().iter()
        .filter(|x| ns_domains.contains_key(x.get_name()))
        .filter_map(|x| match *x.get_rdata() {
            RData::A {ref address} => Some((
                IpAddr::V4(*address), x.get_name().clone(), x.get_ttl(),
            )),
            RData::AAAA {ref address} => Some((
                IpAddr::V6(*address), x.get_name().clone(), x.get_ttl(),
            )),
            _ => None,
        })
        .collect();
        self.get_ns_cache().update(
            &referred_zone,
            ns_items.iter().map(|&(ip, ref name, ttl)| {
                ns_domains.remove(name);
                (ip, name.clone(), ttl as u64)
            }),
        );
        let orphan_domain_futures = self.query_ns_domain_futures(
            &referred_zone, ns_domains.keys().cloned(),
        );
        self.query_ns_multiple(
            ns_items.iter().map(|x| x.0),
            orphan_domain_futures,
        )
    }
    fn update_cache(&self, msg: &Message) {
        self.get_cache().update_from_message(msg, RecordSource::Recursive);
    }
    pub fn resolve(q: &Query, parent: RcResolver, skip_cache: bool) -> Result<Message> {
        let resolver = RecursiveResolver {
            state: Arc::new(RecursiveResolverState::new(q, parent, skip_cache)),
        };
        let ret = try!(resolver.query());
        resolver.update_cache(&ret);
        Ok(ret)
    }
    fn resolve_next(&self, q: &Query) -> Result<Message> {
        if !self.state.skip_cache {
            if let Some(msg) = self.get_cache().lookup_with_type(
                q.get_name(), q.get_query_type(), |m| m.clone_resp_for(q),
            ) {
                return Ok(msg);
            }
        }
        let resolver = RecursiveResolver {
            state: Arc::new(self.state.new_inner(q)),
        };
        let ret = try!(resolver.query());
        self.update_cache(&ret);
        Ok(ret)
    }
}
<|endoftext|>


<commit_before>use rand::{rngs::SmallRng, Rng, SeedableRng};

// Having more than 1 test does seem to make a difference
// (i.e., this calls ptr::swap which having just one test does not).
#[test]
fn simple1() {
    assert_eq!(4, 4);
}

#[test]
fn simple2() {
    assert_ne!(42, 24);
}

// A test that won't work on miri (tests disabling tests).
#[test]
#[cfg_attr(miri, ignore)]
fn does_not_work_on_miri() {
    let x = 0u8;
    assert!(&x as *const _ as usize % 4 < 4);
}

// We also use this to test some external crates, that we cannot depend on in the compiletest suite.

#[test]
fn entropy_rng() {
    // Test `getrandom` directly (in multiple different versions).
    let mut data = vec![0; 16];
    getrandom_1::getrandom(&mut data).unwrap();
    getrandom_2::getrandom(&mut data).unwrap();

    // Try seeding with "real" entropy.
    let mut rng = SmallRng::from_entropy();
    let _val = rng.gen::<i32>();
    let _val = rng.gen::<isize>();
    let _val = rng.gen::<i128>();

    // Also try per-thread RNG.
    let mut rng = rand::thread_rng();
    let _val = rng.gen::<i32>();
    let _val = rng.gen::<isize>();
    let _val = rng.gen::<i128>();
}

#[test]
fn cargo_env() {
    assert_eq!(env!("CARGO_PKG_NAME"), "cargo-miri-test");
    env!("CARGO_BIN_EXE_cargo-miri-test"); // Asserts that this exists.
}

#[test]
#[should_panic(expected = "Explicit panic")]
fn do_panic() {
    // In large, friendly letters :)
    panic!("Explicit panic from test!");
}

#[test]
#[allow(unconditional_panic)]
#[should_panic(expected = "the len is 0 but the index is 42")]
fn fail_index_check() {
    [][42]
}

#[test]
fn page_size() {
    let page_size = page_size::get();

    assert!(
        page_size.next_power_of_two() == page_size,
        "page size not a power of two: {}",
        page_size
    );
}
<commit_msg>Fix comment formatting.<commit_after>use rand::{rngs::SmallRng, Rng, SeedableRng};

// Having more than 1 test does seem to make a difference
// (i.e., this calls ptr::swap which having just one test does not).
#[test]
fn simple1() {
    assert_eq!(4, 4);
}

#[test]
fn simple2() {
    assert_ne!(42, 24);
}

// A test that won't work on miri (tests disabling tests).
#[test]
#[cfg_attr(miri, ignore)]
fn does_not_work_on_miri() {
    let x = 0u8;
    assert!(&x as *const _ as usize % 4 < 4);
}

// We also use this to test some external crates, that we cannot depend on in the compiletest suite.

#[test]
fn entropy_rng() {
    // Test `getrandom` directly (in multiple different versions).
    let mut data = vec![0; 16];
    getrandom_1::getrandom(&mut data).unwrap();
    getrandom_2::getrandom(&mut data).unwrap();

    // Try seeding with "real" entropy.
    let mut rng = SmallRng::from_entropy();
    let _val = rng.gen::<i32>();
    let _val = rng.gen::<isize>();
    let _val = rng.gen::<i128>();

    // Also try per-thread RNG.
    let mut rng = rand::thread_rng();
    let _val = rng.gen::<i32>();
    let _val = rng.gen::<isize>();
    let _val = rng.gen::<i128>();
}

#[test]
fn cargo_env() {
    assert_eq!(env!("CARGO_PKG_NAME"), "cargo-miri-test");
    env!("CARGO_BIN_EXE_cargo-miri-test"); // Asserts that this exists.
}

#[test]
#[should_panic(expected = "Explicit panic")]
fn do_panic() // In large, friendly letters :)
{
    panic!("Explicit panic from test!");
}

#[test]
#[allow(unconditional_panic)]
#[should_panic(expected = "the len is 0 but the index is 42")]
fn fail_index_check() {
    [][42]
}

#[test]
fn page_size() {
    let page_size = page_size::get();

    assert!(
        page_size.next_power_of_two() == page_size,
        "page size not a power of two: {}",
        page_size
    );
}
<|endoftext|>


<commit_before>pub trait Bits {
    fn bits(self, start: u32, end: u32) -> Self;
    fn bits_from(self, start: u32) -> Self;
    fn bit(self, offset: u32) -> Self;

    #[allow(clippy::wrong_self_convention)]
    fn is_bit_set(self, offset: u32) -> bool;
    fn replace_bits(self, start: u32, end: u32, value: Self) -> Self;
    fn replace_bit(self, offset: u32, value: bool) -> Self;
    fn sign_extend(self, bits: u32) -> Self;
}

impl Bits for u32 {
    #[inline(always)]
    fn bits(self, start: u32, end: u32) -> Self {
        (self >> start) & ((1 << (end - start + 1)) - 1)
    }

    #[inline(always)]
    fn bits_from(self, start: u32) -> Self {
        (self >> start) & 1
    }

    #[inline(always)]
    fn bit(self, offset: u32) -> Self {
        (self >> offset) & 1
    }

    #[inline(always)]
    fn is_bit_set(self, offset: u32) -> bool {
        self.bit(offset) != 0
    }

    #[inline(always)]
    fn replace_bits(self, start: u32, end: u32, value: Self) -> Self {
        (self & !(((1 << (end - start + 1)) - 1) << start))
            | ((value & ((1 << (end - start + 1)) - 1)) << start)
    }

    #[inline(always)]
    fn replace_bit(self, offset: u32, value: bool) -> Self {
        self.replace_bits(offset, offset, value as u32)
    }

    #[inline(always)]
    fn sign_extend(self, bits: u32) -> Self {
        (((self << (32 - bits)) as i32) >> (32 - bits)) as u32
    }
}

impl Bits for u16 {
    #[inline(always)]
    fn bits(self, start: u32, end: u32) -> Self {
        (self >> start) & ((1 << (end - start + 1)) - 1)
    }

    #[inline(always)]
    fn bits_from(self, start: u32) -> Self {
        (self >> start) & 1
    }

    #[inline(always)]
    fn bit(self, offset: u32) -> Self {
        (self >> offset) & 1
    }

    #[inline(always)]
    fn is_bit_set(self, offset: u32) -> bool {
        self.bit(offset) != 0
    }

    #[inline(always)]
    fn replace_bits(self, start: u32, end: u32, value: Self) -> Self {
        (self & !(((1 << (end - start + 1)) - 1) << start))
            | ((value & ((1 << (end - start + 1)) - 1)) << start)
    }

    #[inline(always)]
    fn replace_bit(self, offset: u32, value: bool) -> Self {
        self.replace_bits(offset, offset, value as u16)
    }

    #[inline(always)]
    fn sign_extend(self, bits: u32) -> Self {
        (((self << (16 - bits as u16)) as i16) >> (16 - bits as u16)) as u16
    }
}
<commit_msg>Refactor bits and add u64.<commit_after>pub trait Bits {
    const BITS: u32 = 0;
    type Signed;
    type Unsigned;

    fn bits(self, start: u32, end: u32) -> Self;
    fn bits_from(self, start: u32) -> Self;
    fn bit(self, offset: u32) -> Self;

    #[allow(clippy::wrong_self_convention)]
    fn is_bit_set(self, offset: u32) -> bool;
    fn replace_bits(self, start: u32, end: u32, value: Self) -> Self;
    fn replace_bit(self, offset: u32, value: bool) -> Self;
    fn sign_extend(self, bits: u32) -> Self;
}

impl Bits for u32 {
    const BITS: u32 = 32;
    type Signed = i32;
    type Unsigned = Self;

    #[inline(always)]
    fn bits(self, start: u32, end: u32) -> Self {
        (self >> start) & ((1 << (end - start + 1)) - 1)
    }

    #[inline(always)]
    fn bits_from(self, start: u32) -> Self {
        (self >> start) & 1
    }

    #[inline(always)]
    fn bit(self, offset: u32) -> Self {
        (self >> offset) & 1
    }

    #[inline(always)]
    fn is_bit_set(self, offset: u32) -> bool {
        self.bit(offset) != 0
    }

    #[inline(always)]
    fn replace_bits(self, start: u32, end: u32, value: Self) -> Self {
        (self & !(((1 << (end - start + 1)) - 1) << start))
            | ((value & ((1 << (end - start + 1)) - 1)) << start)
    }

    #[inline(always)]
    fn replace_bit(self, offset: u32, value: bool) -> Self {
        self.replace_bits(offset, offset, value as Self)
    }

    #[inline(always)]
    fn sign_extend(self, bits: u32) -> Self {
        (((self << (Self::BITS as Self - bits as Self)) as Self::Signed) >> (32 - bits)) as Self
    }
}

impl Bits for u16 {
    const BITS: u32 = 16;
    type Signed = i32;
    type Unsigned = Self;

    #[inline(always)]
    fn bits(self, start: u32, end: u32) -> Self {
        (self >> start) & ((1 << (end - start + 1)) - 1)
    }

    #[inline(always)]
    fn bits_from(self, start: u32) -> Self {
        (self >> start) & 1
    }

    #[inline(always)]
    fn bit(self, offset: u32) -> Self {
        (self >> offset) & 1
    }

    #[inline(always)]
    fn is_bit_set(self, offset: u32) -> bool {
        self.bit(offset) != 0
    }

    #[inline(always)]
    fn replace_bits(self, start: u32, end: u32, value: Self) -> Self {
        (self & !(((1 << (end - start + 1)) - 1) << start))
            | ((value & ((1 << (end - start + 1)) - 1)) << start)
    }

    #[inline(always)]
    fn replace_bit(self, offset: u32, value: bool) -> Self {
        self.replace_bits(offset, offset, value as Self)
    }

    #[inline(always)]
    fn sign_extend(self, bits: u32) -> Self {
        (((self << (Self::BITS as Self - bits as Self)) as Self::Signed) >> (32 - bits)) as Self
    }
}

impl Bits for u64 {
    const BITS: u32 = 64;
    type Signed = i64;
    type Unsigned = Self;

    #[inline(always)]
    fn bits(self, start: u32, end: u32) -> Self {
        (self >> start) & ((1 << (end - start + 1)) - 1)
    }

    #[inline(always)]
    fn bits_from(self, start: u32) -> Self {
        (self >> start) & 1
    }

    #[inline(always)]
    fn bit(self, offset: u32) -> Self {
        (self >> offset) & 1
    }

    #[inline(always)]
    fn is_bit_set(self, offset: u32) -> bool {
        self.bit(offset) != 0
    }

    #[inline(always)]
    fn replace_bits(self, start: u32, end: u32, value: Self) -> Self {
        (self & !(((1 << (end - start + 1)) - 1) << start))
            | ((value & ((1 << (end - start + 1)) - 1)) << start)
    }

    #[inline(always)]
    fn replace_bit(self, offset: u32, value: bool) -> Self {
        self.replace_bits(offset, offset, value as Self)
    }

    #[inline(always)]
    fn sign_extend(self, bits: u32) -> Self {
        (((self << (Self::BITS as Self - bits as Self)) as Self::Signed) >> (32 - bits)) as Self
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Alex Regueiro
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

use {DB, Error};
use ffi;

use libc::{c_int, uint32_t};
use std::ffi::CString;
use std::path::Path;

pub struct BackupEngine {
    inner: *mut ffi::rocksdb_backup_engine_t,
}

pub struct BackupEngineOptions {
    inner: *mut ffi::rocksdb_options_t,
}

pub struct RestoreOptions {
    inner: *mut ffi::rocksdb_restore_options_t,
}

impl BackupEngine {
    /// Open a backup engine with the specified options.
    pub fn open<P: AsRef<Path>>(opts: &BackupEngineOptions,
                                path: P)
                                -> Result<BackupEngine, Error> {
        let path = path.as_ref();
        let cpath = match CString::new(path.to_string_lossy().as_bytes()) {
            Ok(c) => c,
            Err(_) => {
                return Err(Error::new("Failed to convert path to CString \
                                       when opening backup engine"
                    .to_owned()))
            }
        };

        let be: *mut ffi::rocksdb_backup_engine_t;
        unsafe { be = ffi_try!(ffi::rocksdb_backup_engine_open(opts.inner, cpath.as_ptr())) }

        if be.is_null() {
            return Err(Error::new("Could not initialize backup engine.".to_owned()));
        }

        Ok(BackupEngine { inner: be })
    }

    fn create_new_backup(&mut self, db: &DB) -> Result<(), Error> {
        unsafe {
            ffi_try!(ffi::rocksdb_backup_engine_create_new_backup(self.inner, db.inner));
            Ok(())
        }
    }

    fn purge_old_backups(&mut self, num_backups_to_keep: usize) -> Result<(), Error> {
        unsafe {
            ffi_try!(ffi::rocksdb_backup_engine_purge_old_backups(self.inner,
                                                                  num_backups_to_keep as uint32_t));
            Ok(())
        }
    }
}

impl BackupEngineOptions {
    //
}

impl RestoreOptions {
    pub fn set_keep_log_files(&mut self, keep_log_files: bool) {
        unsafe {
            ffi::rocksdb_restore_options_set_keep_log_files(self.inner, keep_log_files as c_int);
        }
    }
}

impl Default for BackupEngineOptions {
    fn default() -> BackupEngineOptions {
        unsafe {
            let opts = ffi::rocksdb_options_create();
            if opts.is_null() {
                panic!("Could not create RocksDB backup options".to_owned());
            }
            BackupEngineOptions { inner: opts }
        }
    }
}

impl Default for RestoreOptions {
    fn default() -> RestoreOptions {
        unsafe {
            let opts = ffi::rocksdb_restore_options_create();
            if opts.is_null() {
                panic!("Could not create RocksDB restore options".to_owned());
            }
            RestoreOptions { inner: opts }
        }
    }
}

impl Drop for BackupEngine {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_backup_engine_close(self.inner);
        }
    }
}

impl Drop for BackupEngineOptions {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_options_destroy(self.inner);
        }
    }
}

impl Drop for RestoreOptions {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_restore_options_destroy(self.inner);
        }
    }
}
<commit_msg>Make BackupEngine methods public<commit_after>// Copyright 2016 Alex Regueiro
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

use {DB, Error};
use ffi;

use libc::{c_int, uint32_t};
use std::ffi::CString;
use std::path::Path;

pub struct BackupEngine {
    inner: *mut ffi::rocksdb_backup_engine_t,
}

pub struct BackupEngineOptions {
    inner: *mut ffi::rocksdb_options_t,
}

pub struct RestoreOptions {
    inner: *mut ffi::rocksdb_restore_options_t,
}

impl BackupEngine {
    /// Open a backup engine with the specified options.
    pub fn open<P: AsRef<Path>>(opts: &BackupEngineOptions,
                                path: P)
                                -> Result<BackupEngine, Error> {
        let path = path.as_ref();
        let cpath = match CString::new(path.to_string_lossy().as_bytes()) {
            Ok(c) => c,
            Err(_) => {
                return Err(Error::new("Failed to convert path to CString \
                                       when opening backup engine"
                    .to_owned()))
            }
        };

        let be: *mut ffi::rocksdb_backup_engine_t;
        unsafe { be = ffi_try!(ffi::rocksdb_backup_engine_open(opts.inner, cpath.as_ptr())) }

        if be.is_null() {
            return Err(Error::new("Could not initialize backup engine.".to_owned()));
        }

        Ok(BackupEngine { inner: be })
    }

    pub fn create_new_backup(&mut self, db: &DB) -> Result<(), Error> {
        unsafe {
            ffi_try!(ffi::rocksdb_backup_engine_create_new_backup(self.inner, db.inner));
            Ok(())
        }
    }

    pub fn purge_old_backups(&mut self, num_backups_to_keep: usize) -> Result<(), Error> {
        unsafe {
            ffi_try!(ffi::rocksdb_backup_engine_purge_old_backups(self.inner,
                                                                  num_backups_to_keep as uint32_t));
            Ok(())
        }
    }
}

impl BackupEngineOptions {
    //
}

impl RestoreOptions {
    pub fn set_keep_log_files(&mut self, keep_log_files: bool) {
        unsafe {
            ffi::rocksdb_restore_options_set_keep_log_files(self.inner, keep_log_files as c_int);
        }
    }
}

impl Default for BackupEngineOptions {
    fn default() -> BackupEngineOptions {
        unsafe {
            let opts = ffi::rocksdb_options_create();
            if opts.is_null() {
                panic!("Could not create RocksDB backup options".to_owned());
            }
            BackupEngineOptions { inner: opts }
        }
    }
}

impl Default for RestoreOptions {
    fn default() -> RestoreOptions {
        unsafe {
            let opts = ffi::rocksdb_restore_options_create();
            if opts.is_null() {
                panic!("Could not create RocksDB restore options".to_owned());
            }
            RestoreOptions { inner: opts }
        }
    }
}

impl Drop for BackupEngine {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_backup_engine_close(self.inner);
        }
    }
}

impl Drop for BackupEngineOptions {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_options_destroy(self.inner);
        }
    }
}

impl Drop for RestoreOptions {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_restore_options_destroy(self.inner);
        }
    }
}
<|endoftext|>


<commit_before>use std::{error, fmt, result};
use std::net::SocketAddr;

use store::Store;

use LogIndex;
use Term;

#[derive(Clone, Debug)]
pub struct MemStore {
    current_term: Term,
    voted_for: Option<SocketAddr>,
    entries: Vec<(Term, Vec<u8>)>,
}

/// Non-instantiable error type for MemStore
pub enum Error { }

impl fmt::Display for Error {
    fn fmt(&self, _fmt: &mut fmt::Formatter) -> fmt::Result {
        unreachable!()
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, _fmt: &mut fmt::Formatter) -> fmt::Result {
        unreachable!()
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        unreachable!()
    }
}

impl MemStore {

    pub fn new() -> MemStore {
        MemStore {
            current_term: Term(0),
            voted_for: None,
            entries: Vec::new(),
        }
    }
}

impl Store for MemStore {

    type Error = Error;

    fn current_term(&self) -> result::Result<Term, Error> {
        Ok(self.current_term)
    }

    fn set_current_term(&mut self, term: Term) -> result::Result<(), Error> {
        self.voted_for = None;
        Ok(self.current_term = term)
    }

    fn inc_current_term(&mut self) -> result::Result<Term, Error> {
        self.voted_for = None;
        self.current_term = self.current_term + 1;
        self.current_term()
    }

    fn voted_for(&self) -> result::Result<Option<SocketAddr>, Error> {
        Ok(self.voted_for)
    }

    fn set_voted_for(&mut self, address: SocketAddr) -> result::Result<(), Error> {
        Ok(self.voted_for = Some(address))
    }

    fn latest_log_index(&self) -> result::Result<LogIndex, Error> {
        Ok(LogIndex(self.entries.len() as u64))
    }

    fn latest_log_term(&self) -> result::Result<Term, Error> {
        let len = self.entries.len();
        if len == 0 {
            Ok(Term::from(0))
        } else {
            Ok(self.entries[len - 1].0)
        }
    }

    fn entry(&self, index: LogIndex) -> result::Result<(Term, &[u8]), Error> {
        let (term, ref bytes) = self.entries[Into::<u64>::into(index) as usize - 1];
        Ok((term, &bytes))
    }

    fn append_entries(&mut self,
                      from: LogIndex,
                      entries: &[(Term, &[u8])])
                      -> result::Result<(), Error> {
        assert!(self.latest_log_index().unwrap() + 1 >= from);
        self.entries.truncate(Into::<u64>::into(from) as usize - 1);
        Ok(self.entries.extend(entries.iter().map(|&(term, command)| (term, command.to_vec()))))
    }

    fn truncate_entries(&mut self, index: LogIndex) -> result::Result<(), Error> {
        assert!(self.latest_log_index().unwrap() >= index);
        Ok(self.entries.truncate(Into::<u64>::into(index) as usize - 1))
    }
}

#[cfg(test)]
mod test {

    use std::str::FromStr;
    use std::net::SocketAddr;

    use super::*;
    use LogIndex;
    use Term;
    use store::Store;

    #[test]
    fn test_current_term() {
        let mut store = MemStore::new();
        assert_eq!(Term(0), store.current_term().unwrap());
        store.set_voted_for(SocketAddr::from_str("127.0.0.1:0").unwrap()).unwrap();
        store.set_current_term(Term(42)).unwrap();
        assert_eq!(None, store.voted_for().unwrap());
        assert_eq!(Term(42), store.current_term().unwrap());
        store.inc_current_term().unwrap();
        assert_eq!(Term(43), store.current_term().unwrap());
    }

    #[test]
    fn test_voted_for() {
        let mut store = MemStore::new();
        assert_eq!(None, store.voted_for().unwrap());
        let addr = SocketAddr::from_str("127.0.0.1:0").unwrap();
        store.set_voted_for(addr.clone()).unwrap();
        assert_eq!(Some(addr), store.voted_for().unwrap());
    }

    #[test]
    fn test_append_entries() {
        let mut store = MemStore::new();
        assert_eq!(LogIndex::from(0), store.latest_log_index().unwrap());
        assert_eq!(Term::from(0), store.latest_log_term().unwrap());

        store.append_entries(LogIndex(1), &[(Term::from(0), &[1]),
                                            (Term::from(0), &[2]),
                                            (Term::from(0), &[3]),
                                            (Term::from(1), &[4])]).unwrap();
        assert_eq!(LogIndex::from(4), store.latest_log_index().unwrap());
        assert_eq!(Term::from(1), store.latest_log_term().unwrap());
        assert_eq!((Term::from(0), &*vec![1u8]), store.entry(LogIndex::from(1)).unwrap());
        assert_eq!((Term::from(0), &*vec![2u8]), store.entry(LogIndex::from(2)).unwrap());
        assert_eq!((Term::from(0), &*vec![3u8]), store.entry(LogIndex::from(3)).unwrap());
        assert_eq!((Term::from(1), &*vec![4u8]), store.entry(LogIndex::from(4)).unwrap());

        store.append_entries(LogIndex::from(4), &[]).unwrap();
        assert_eq!(LogIndex(3), store.latest_log_index().unwrap());
        assert_eq!(Term::from(0), store.latest_log_term().unwrap());
        assert_eq!((Term::from(0), &*vec![1u8]), store.entry(LogIndex::from(1)).unwrap());
        assert_eq!((Term::from(0), &*vec![2u8]), store.entry(LogIndex::from(2)).unwrap());
        assert_eq!((Term::from(0), &*vec![3u8]), store.entry(LogIndex::from(3)).unwrap());

        store.append_entries(LogIndex::from(3), &[(Term(2), &[3]), (Term(3), &[4])]).unwrap();
        assert_eq!(LogIndex(4), store.latest_log_index().unwrap());
        assert_eq!(Term::from(3), store.latest_log_term().unwrap());
        assert_eq!((Term::from(0), &*vec![1u8]), store.entry(LogIndex::from(1)).unwrap());
        assert_eq!((Term::from(0), &*vec![2u8]), store.entry(LogIndex::from(2)).unwrap());
        assert_eq!((Term::from(2), &*vec![3u8]), store.entry(LogIndex::from(3)).unwrap());
        assert_eq!((Term::from(3), &*vec![4u8]), store.entry(LogIndex::from(4)).unwrap());
    }
}
<commit_msg>Comment on memstore<commit_after>use std::{error, fmt, result};
use std::net::SocketAddr;

use store::Store;

use LogIndex;
use Term;

/// This is a `Store` implementation that stores entries in a simple in-memory vector. Other data
/// is stored in a struct. It is chiefly intended for testing.
#[derive(Clone, Debug)]
pub struct MemStore {
    current_term: Term,
    voted_for: Option<SocketAddr>,
    entries: Vec<(Term, Vec<u8>)>,
}

/// Non-instantiable error type for MemStore
pub enum Error { }

impl fmt::Display for Error {
    fn fmt(&self, _fmt: &mut fmt::Formatter) -> fmt::Result {
        unreachable!()
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, _fmt: &mut fmt::Formatter) -> fmt::Result {
        unreachable!()
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        unreachable!()
    }
}

impl MemStore {

    pub fn new() -> MemStore {
        MemStore {
            current_term: Term(0),
            voted_for: None,
            entries: Vec::new(),
        }
    }
}

impl Store for MemStore {

    type Error = Error;

    fn current_term(&self) -> result::Result<Term, Error> {
        Ok(self.current_term)
    }

    fn set_current_term(&mut self, term: Term) -> result::Result<(), Error> {
        self.voted_for = None;
        Ok(self.current_term = term)
    }

    fn inc_current_term(&mut self) -> result::Result<Term, Error> {
        self.voted_for = None;
        self.current_term = self.current_term + 1;
        self.current_term()
    }

    fn voted_for(&self) -> result::Result<Option<SocketAddr>, Error> {
        Ok(self.voted_for)
    }

    fn set_voted_for(&mut self, address: SocketAddr) -> result::Result<(), Error> {
        Ok(self.voted_for = Some(address))
    }

    fn latest_log_index(&self) -> result::Result<LogIndex, Error> {
        Ok(LogIndex(self.entries.len() as u64))
    }

    fn latest_log_term(&self) -> result::Result<Term, Error> {
        let len = self.entries.len();
        if len == 0 {
            Ok(Term::from(0))
        } else {
            Ok(self.entries[len - 1].0)
        }
    }

    fn entry(&self, index: LogIndex) -> result::Result<(Term, &[u8]), Error> {
        let (term, ref bytes) = self.entries[Into::<u64>::into(index) as usize - 1];
        Ok((term, &bytes))
    }

    fn append_entries(&mut self,
                      from: LogIndex,
                      entries: &[(Term, &[u8])])
                      -> result::Result<(), Error> {
        assert!(self.latest_log_index().unwrap() + 1 >= from);
        self.entries.truncate(Into::<u64>::into(from) as usize - 1);
        Ok(self.entries.extend(entries.iter().map(|&(term, command)| (term, command.to_vec()))))
    }

    fn truncate_entries(&mut self, index: LogIndex) -> result::Result<(), Error> {
        assert!(self.latest_log_index().unwrap() >= index);
        Ok(self.entries.truncate(Into::<u64>::into(index) as usize - 1))
    }
}

#[cfg(test)]
mod test {

    use std::str::FromStr;
    use std::net::SocketAddr;

    use super::*;
    use LogIndex;
    use Term;
    use store::Store;

    #[test]
    fn test_current_term() {
        let mut store = MemStore::new();
        assert_eq!(Term(0), store.current_term().unwrap());
        store.set_voted_for(SocketAddr::from_str("127.0.0.1:0").unwrap()).unwrap();
        store.set_current_term(Term(42)).unwrap();
        assert_eq!(None, store.voted_for().unwrap());
        assert_eq!(Term(42), store.current_term().unwrap());
        store.inc_current_term().unwrap();
        assert_eq!(Term(43), store.current_term().unwrap());
    }

    #[test]
    fn test_voted_for() {
        let mut store = MemStore::new();
        assert_eq!(None, store.voted_for().unwrap());
        let addr = SocketAddr::from_str("127.0.0.1:0").unwrap();
        store.set_voted_for(addr.clone()).unwrap();
        assert_eq!(Some(addr), store.voted_for().unwrap());
    }

    #[test]
    fn test_append_entries() {
        let mut store = MemStore::new();
        assert_eq!(LogIndex::from(0), store.latest_log_index().unwrap());
        assert_eq!(Term::from(0), store.latest_log_term().unwrap());

        store.append_entries(LogIndex(1), &[(Term::from(0), &[1]),
                                            (Term::from(0), &[2]),
                                            (Term::from(0), &[3]),
                                            (Term::from(1), &[4])]).unwrap();
        assert_eq!(LogIndex::from(4), store.latest_log_index().unwrap());
        assert_eq!(Term::from(1), store.latest_log_term().unwrap());
        assert_eq!((Term::from(0), &*vec![1u8]), store.entry(LogIndex::from(1)).unwrap());
        assert_eq!((Term::from(0), &*vec![2u8]), store.entry(LogIndex::from(2)).unwrap());
        assert_eq!((Term::from(0), &*vec![3u8]), store.entry(LogIndex::from(3)).unwrap());
        assert_eq!((Term::from(1), &*vec![4u8]), store.entry(LogIndex::from(4)).unwrap());

        store.append_entries(LogIndex::from(4), &[]).unwrap();
        assert_eq!(LogIndex(3), store.latest_log_index().unwrap());
        assert_eq!(Term::from(0), store.latest_log_term().unwrap());
        assert_eq!((Term::from(0), &*vec![1u8]), store.entry(LogIndex::from(1)).unwrap());
        assert_eq!((Term::from(0), &*vec![2u8]), store.entry(LogIndex::from(2)).unwrap());
        assert_eq!((Term::from(0), &*vec![3u8]), store.entry(LogIndex::from(3)).unwrap());

        store.append_entries(LogIndex::from(3), &[(Term(2), &[3]), (Term(3), &[4])]).unwrap();
        assert_eq!(LogIndex(4), store.latest_log_index().unwrap());
        assert_eq!(Term::from(3), store.latest_log_term().unwrap());
        assert_eq!((Term::from(0), &*vec![1u8]), store.entry(LogIndex::from(1)).unwrap());
        assert_eq!((Term::from(0), &*vec![2u8]), store.entry(LogIndex::from(2)).unwrap());
        assert_eq!((Term::from(2), &*vec![3u8]), store.entry(LogIndex::from(3)).unwrap());
        assert_eq!((Term::from(3), &*vec![4u8]), store.entry(LogIndex::from(4)).unwrap());
    }
}
<|endoftext|>


<commit_before>//! An unbounded channel that only stores last value sent

use std::sync::{Arc, Weak, Mutex};

use task::{self, Task};
use {Sink, Stream, AsyncSink, Async, Poll, StartSend};

/// Slot is very similar to unbounded channel but only stores last value sent
///
/// I.e. if you want to send some value between from producer to a consumer
/// and if consumer is slow it should skip old values, the slot is
/// a structure for the task.

/// The transmission end of a channel which is used to send values
///
/// If the receiver is not fast enough only the last value is preserved and
/// other ones are discarded.
#[derive(Debug)]
pub struct Sender<T> {
    inner: Option<Weak<Mutex<Inner<T>>>>,
}

/// The receiving end of a channel which preserves only the last value
#[derive(Debug)]
pub struct Receiver<T> {
    inner: Arc<Mutex<Inner<T>>>,
}

/// Error type for sending, used when the receiving end of a channel is
/// dropped
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct SendError<T>(T);

#[derive(Debug)]
struct Inner<T> {
    value: Option<T>,
    task: Option<Task>,
}

trait AssertKindsSender: Send + Sync + Clone {}
impl AssertKindsSender for Sender<u32> {}

trait AssertKindsReceiver: Send + Sync {}
impl AssertKindsReceiver for Receiver<u32> {}

impl<T> Sender<T> {
    /// Sets the new new value of the stream and notifies the consumer if any
    pub fn swap(&self, value: T) -> Result<Option<T>, SendError<T>> {
        let result;
        // Do this step first so that the lock is dropped when
        // `unpark` is called
        let task = {
            let strong = self.inner.as_ref()
                .expect("sending to a closed slot");
            if let Some(ref lock) = strong.upgrade() {
                let mut inner = lock.lock().unwrap();
                result = inner.value.take();
                inner.value = Some(value);
                inner.task.take()
            } else {
                return Err(SendError(value));
            }
        };
        if let Some(task) = task {
            task.notify();
        }
        return Ok(result);
    }
}

impl<T> Sink for Sender<T> {
    type SinkItem = T;
    type SinkError = SendError<T>;
    fn start_send(&mut self, item: T) -> StartSend<T, SendError<T>> {
        self.swap(item)?;
        Ok(AsyncSink::Ready)
    }
    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        Ok(Async::Ready(()))
    }
    fn close(&mut self) -> Poll<(), Self::SinkError> {
        // Do this step first so that the lock is dropped *and*
        // weakref is dropped when `unpark` is called
        let task = {
            if let Some(weak) = self.inner.take() {
                if let Some(ref lock) = weak.upgrade() {
                    drop(weak);
                    let mut inner = lock.lock().unwrap();
                    inner.task.take()
                } else {
                    None
                }
            } else {
                None
            }
        };
        // notify on any drop of a sender, so eventually receiver wakes up
        // when there are no senders and closes the stream
        if let Some(task) = task {
            task.notify();
        }
        Ok(Async::Ready(()))
    }
}

impl<T> Drop for Sender<T> {
    fn drop(&mut self) {
        self.close().ok();
    }
}

impl<T> Stream for Receiver<T> {
    type Item = T;
    type Error = ();  // actually void
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        let result = {
            let mut inner = self.inner.lock().unwrap();
            if inner.value.is_none() {
                if Arc::weak_count(&self.inner) == 0 {
                    // no senders, terminate the stream
                    return Ok(Async::Ready(None));
                } else {
                    inner.task = Some(task::current());
                }
            }
            inner.value.take()
        };
        match result {
            Some(value) => Ok(Async::Ready(Some(value))),
            None => Ok(Async::NotReady),
        }
    }
}

/// Creates an in-memory Stream which only preserves last value
///
/// This method is somewhat similar to `channel(1)` but instead of preserving
/// first value sent (and erroring on sender side) it replaces value if
/// consumer is not fast enough and preserves last values sent on any
/// poll of a stream.
///
/// # Example
///
/// ```
/// use std::thread;
/// use futures::prelude::*;
/// use futures::stream::iter_ok;
/// use futures::sync::slot;
///
/// let (tx, rx) = slot::channel::<i32>();
///
/// tx.send_all(iter_ok(vec![1, 2, 3])).wait();
///
/// let received = rx.collect().wait().unwrap();
/// assert_eq!(received, vec![3]);
/// ```
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
    let inner = Arc::new(Mutex::new(Inner {
        value: None,
        task: None,
    }));
    return (Sender { inner: Some(Arc::downgrade(&inner)) },
            Receiver { inner: inner });
}

impl<T> Clone for Sender<T> {
    fn clone(&self) -> Sender<T> {
        Sender { inner: self.inner.clone() }
    }
}
<commit_msg>Use `Weak::new` instead of `Option<Weak>`<commit_after>//! An unbounded channel that only stores last value sent

use std::mem;
use std::sync::{Arc, Weak, Mutex};

use task::{self, Task};
use {Sink, Stream, AsyncSink, Async, Poll, StartSend};

/// Slot is very similar to unbounded channel but only stores last value sent
///
/// I.e. if you want to send some value between from producer to a consumer
/// and if consumer is slow it should skip old values, the slot is
/// a structure for the task.

/// The transmission end of a channel which is used to send values
///
/// If the receiver is not fast enough only the last value is preserved and
/// other ones are discarded.
#[derive(Debug)]
pub struct Sender<T> {
    inner: Weak<Mutex<Inner<T>>>,
}

/// The receiving end of a channel which preserves only the last value
#[derive(Debug)]
pub struct Receiver<T> {
    inner: Arc<Mutex<Inner<T>>>,
}

/// Error type for sending, used when the receiving end of a channel is
/// dropped
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct SendError<T>(T);

#[derive(Debug)]
struct Inner<T> {
    value: Option<T>,
    task: Option<Task>,
}

trait AssertKindsSender: Send + Sync + Clone {}
impl AssertKindsSender for Sender<u32> {}

trait AssertKindsReceiver: Send + Sync {}
impl AssertKindsReceiver for Receiver<u32> {}

impl<T> Sender<T> {
    /// Sets the new new value of the stream and notifies the consumer if any
    pub fn swap(&self, value: T) -> Result<Option<T>, SendError<T>> {
        let result;
        // Do this step first so that the lock is dropped when
        // `unpark` is called
        let task = {
            if let Some(ref lock) = self.inner.upgrade() {
                let mut inner = lock.lock().unwrap();
                result = inner.value.take();
                inner.value = Some(value);
                inner.task.take()
            } else {
                return Err(SendError(value));
            }
        };
        if let Some(task) = task {
            task.notify();
        }
        return Ok(result);
    }
}

impl<T> Sink for Sender<T> {
    type SinkItem = T;
    type SinkError = SendError<T>;
    fn start_send(&mut self, item: T) -> StartSend<T, SendError<T>> {
        self.swap(item)?;
        Ok(AsyncSink::Ready)
    }
    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        Ok(Async::Ready(()))
    }
    fn close(&mut self) -> Poll<(), Self::SinkError> {
        // Do this step first so that the lock is dropped *and*
        // weakref is dropped when `unpark` is called
        let task = {
            let weak = mem::replace(&mut self.inner, Weak::new());
            if let Some(ref lock) = weak.upgrade() {
                drop(weak);
                let mut inner = lock.lock().unwrap();
                inner.task.take()
            } else {
                None
            }
        };
        // notify on any drop of a sender, so eventually receiver wakes up
        // when there are no senders and closes the stream
        if let Some(task) = task {
            task.notify();
        }
        Ok(Async::Ready(()))
    }
}

impl<T> Drop for Sender<T> {
    fn drop(&mut self) {
        self.close().ok();
    }
}

impl<T> Stream for Receiver<T> {
    type Item = T;
    type Error = ();  // actually void
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        let result = {
            let mut inner = self.inner.lock().unwrap();
            if inner.value.is_none() {
                if Arc::weak_count(&self.inner) == 0 {
                    // no senders, terminate the stream
                    return Ok(Async::Ready(None));
                } else {
                    inner.task = Some(task::current());
                }
            }
            inner.value.take()
        };
        match result {
            Some(value) => Ok(Async::Ready(Some(value))),
            None => Ok(Async::NotReady),
        }
    }
}

/// Creates an in-memory Stream which only preserves last value
///
/// This method is somewhat similar to `channel(1)` but instead of preserving
/// first value sent (and erroring on sender side) it replaces value if
/// consumer is not fast enough and preserves last values sent on any
/// poll of a stream.
///
/// # Example
///
/// ```
/// use std::thread;
/// use futures::prelude::*;
/// use futures::stream::iter_ok;
/// use futures::sync::slot;
///
/// let (tx, rx) = slot::channel::<i32>();
///
/// tx.send_all(iter_ok(vec![1, 2, 3])).wait();
///
/// let received = rx.collect().wait().unwrap();
/// assert_eq!(received, vec![3]);
/// ```
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
    let inner = Arc::new(Mutex::new(Inner {
        value: None,
        task: None,
    }));
    return (Sender { inner: Arc::downgrade(&inner) },
            Receiver { inner: inner });
}

impl<T> Clone for Sender<T> {
    fn clone(&self) -> Sender<T> {
        Sender { inner: self.inner.clone() }
    }
}
<|endoftext|>


<commit_before>
use vertex::{Vertex, VertexProxy};
use edge::{Edge, EdgeProxy};
// define a traversal as a series of stages

struct VertexVisit {
    vertex: VertexProxy
}

struct EdgeVisit {
    edge: EdgeProxy
}

enum TraversalStage {

}

pub enum GraphIterable {
    Vertex(VertexProxy),
    Edge(EdgeProxy)
}

// used to track a single graph traversal
// will be used in conjunction with a GraphQuery

struct Path {
    path: Vec<GraphIterable>
}

impl Path {
    fn new(vertex: VertexProxy) -> Path {
		Path{path:vec![GraphIterable::Vertex(vertex)]}
    }
}

pub struct GraphQuery {
	// many paths
    paths: Vec<Path>
}

impl GraphQuery {
    pub fn new(vertices: Vec<VertexProxy>) -> GraphQuery {
        // create a new path for each proxy, add the proxy
        let mut paths = Vec::new();
        for p in vertices {
			let mut path = Path::new(p);
            paths.push(path)
        }
        GraphQuery{paths:paths}
    }

    pub fn empty() -> GraphQuery {
        // return an empty graph query
        GraphQuery{paths:Vec::new()}
    }

    fn map<F: Fn(Path) -> Vec<Path>>(&self, closure: F) -> GraphQuery  {
        GraphQuery::empty()
    }

    pub fn outV(&self) -> GraphQuery {
//        let g = GraphQuery::empty();

		let f = |path: Path| {
			let mut result : Vec<Path>  = Vec::new();
			result
		};
		self.map(f)
    }

	pub fn inv(&self) -> GraphQuery {
		let g = GraphQuery::empty();
		g
	}
}

<commit_msg>iteration over path elements<commit_after>
use vertex::{Vertex, VertexProxy};
use edge::{Edge, EdgeProxy};
// define a traversal as a series of stages

struct VertexVisit {
    vertex: VertexProxy
}

struct EdgeVisit {
    edge: EdgeProxy
}

enum TraversalStage {

}

pub enum GraphIterable {
    Vertex(VertexProxy),
    Edge(EdgeProxy)
}

// used to track a single graph traversal
// will be used in conjunction with a GraphQuery

struct Path {
    path: Vec<GraphIterable>
}

impl Path {
    fn new(vertex: VertexProxy) -> Path {
		Path{path:vec![GraphIterable::Vertex(vertex)]}
    }
}

pub struct GraphQuery {
	// many paths
    paths: Vec<Path>
}

impl GraphQuery {
    pub fn new(vertices: Vec<VertexProxy>) -> GraphQuery {
        // create a new path for each proxy, add the proxy
        let mut paths = Vec::new();
        for p in vertices {
			let mut path = Path::new(p);
            paths.push(path)
        }
        GraphQuery{paths:paths}
    }

    pub fn empty() -> GraphQuery {
        // return an empty graph query
        GraphQuery{paths:Vec::new()}
    }

    fn map<F: Fn(Path) -> Vec<Path>>(&self, closure: F) -> GraphQuery  {
        let mut result = GraphQuery::empty(); // result
		for x in self.paths.iter() {

		}
		result
    }

    pub fn outV(&self) -> GraphQuery {
		let f = |path: Path| {
			let mut result : Vec<Path>  = Vec::new();
			result
		};
		self.map(f)
    }

	pub fn inv(&self) -> GraphQuery {
		let g = GraphQuery::empty();
		g
	}
}

<|endoftext|>