<commit_before>//! Atom-class implementation.

use std::ops::{BitXor, Sub, Range, RangeInclusive};
use std::iter::FromIterator;
use std::fmt::{self, Debug};

#[cfg(feature = "regex")] use regex_syntax;
use itertools::Itertools;

use super::{Anchor, Atom, ByteOrChar, Element, Repetition, Sequence, Union};
use super::{Distance, Step};
use util::set::{self, Contains};

// ----------------------------------------------------------------
// Class

/// Any member of a class of atoms
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Eq, Ord)]
pub enum ClassMember<T: Atom> {
    /// A single atom
    Atom(T),

    /// A range of atoms
    Range(T, T)
}

impl<T: Atom> ClassMember<T> {
    /// Map the class-member's atoms to a different type
    pub fn map_atoms<U, F>(self, f: F) -> ClassMember<U>
        where F: Fn(T) -> U,
              U: Atom
    {
        match self {
            ClassMember::Atom(a) => ClassMember::Atom(f(a)),
            ClassMember::Range(a, b) => ClassMember::Range(f(a), f(b))
        }
    }

    /// Get the number of individual atoms that this class member represents.
    pub fn len(&self) -> usize
        where T: Distance
    {
        match self {
            &ClassMember::Atom(_) => 1,
            &ClassMember::Range(first, last) => last.distance(&first) + 1,
        }
    }

    /// Fetch an iterator over each discrete value contained by this member.
    #[inline]
     pub fn iter(&self) -> impl Iterator<Item=T>
        where T: Step {
        match self {
            &ClassMember::Atom(a) => AtomIter(a...a),
            &ClassMember::Range(a, b) => AtomIter(a...b),
        }
    }
}

impl<T: Atom> set::Contains<T> for ClassMember<T> {
    /// Check if the member is or contains a particular atom.
    fn contains(&self, x: T) -> bool {
        match self {
            &ClassMember::Atom(a) => x == a,
            &ClassMember::Range(a, b) => (a...b).contains(x),
        }
    }
}

/// Iterator over the atoms specified by a class member.
struct AtomIter<T: Atom>(RangeInclusive<T>);

impl<T: Atom> Iterator for AtomIter<T>
        where T: Step
{
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        match self.0 {
            RangeInclusive::NonEmpty{start, end} => {
                let next_start = start.increment();
                if next_start > end {
                    self.0 = RangeInclusive::Empty{at: end};
                } else {
                    self.0 = RangeInclusive::NonEmpty{start: next_start, end: end};
                }
                Some(start) },
            RangeInclusive::Empty{..} => None
        }
    }
}

/*impl<T: Atom> Debug for ClassMember<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ClassMember::Atom(a) => write!(f, "{:?}", a),
            ClassMember::Range(first, last) => write!(f, "{:?}-{:?}", first, last)
        }
    }
}*/

impl<T: Atom> From<T> for ClassMember<T> {
    fn from(a: T) -> Self { ClassMember::Atom(a) }
}

impl From<regex_syntax::ClassRange> for ClassMember<ByteOrChar> {
    fn from(cr: regex_syntax::ClassRange) -> Self {
        ClassMember::Range(cr.start.into(), cr.end.into())
    }
}

impl From<regex_syntax::ByteRange> for ClassMember<ByteOrChar> {
    fn from(cr: regex_syntax::ByteRange) -> Self {
        ClassMember::Range(cr.start.into(), cr.end.into())
    }
}

impl From<regex_syntax::ClassRange> for ClassMember<char> {
    fn from(cr: regex_syntax::ClassRange) -> Self {
        ClassMember::Range(cr.start, cr.end)
    }
}

impl From<regex_syntax::ByteRange> for ClassMember<u8> {
    fn from(cr: regex_syntax::ByteRange) -> Self {
        ClassMember::Range(cr.start, cr.end)
    }
}

impl<T: Atom> From<Range<T>> for ClassMember<T>
    where T: Sub<usize,Output=T> + Sub<T,Output=usize>
{
    fn from(r: Range<T>) -> Self {
        if r.end - r.start > 1 { ClassMember::Range(r.start, r.end - 1usize) }
        else { ClassMember::Atom(r.start) }
    }
}
// ----------------------------------------------------------------

/// A set of atoms and/or ranges of atoms.
#[derive(Clone, PartialEq, PartialOrd, Eq, Ord)]
pub struct Class<T: Atom> {
    members: Vec<ClassMember<T>>
}

// FIXME: [optimize] All operations on a Class<T>, including observer and
// predicate methods, currently require at least O(N) time.  We could probably
// improve this by changing the storage representation.
impl<T: Atom> Class<T> {
    /// Create a new Class instance with members taken from the given iterator.
    pub fn new<I,U>(items: I) -> Self
        where I: IntoIterator<Item=U>,
              ClassMember<T>: From<U> {
        Class{members: Vec::from_iter(items.into_iter().map(|x| x.into()))}
    }

    /// Create a new Class instance from an iterator of ClassMember instances.
    pub fn from_members<I>(m: I) -> Self
        where I: IntoIterator<Item=ClassMember<T>>
    {
        Class{members: m.into_iter().collect()}
    }

    /// Map the atoms in the class to a different atom type.
    pub fn map_atoms<U, F>(self, f: F) -> Class<U>
        where F: Fn(T) -> U,
              U: Atom
    {
        Class::new(self.members.into_iter().map(|m| m.map_atoms(&f)))
    }

    /// Check whether the class would match a particular atom.
    pub fn matches(&self, x: T) -> bool {
        self.members.iter().any(|&range| range.contains(x))
    }

    /// Get the number of members (atoms) in the class.
    ///
    /// Ranges of atoms count as the number of atoms in each range.
    pub fn len(&self) -> usize where T: Distance {
        self.members.iter().map(|m| m.len()).sum()
    }

    /// Check if the class contains zero members.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.members.is_empty()
    }

    /// Fetch an iterator over the members specified in the class.
    pub fn iter_members<'a>(&'a self) -> impl Iterator<Item=&'a ClassMember<T>> {
        self.members.iter()
    }

    /// Fetch an iterator over the atoms specified in the class.
    pub fn iter<'a>(&'a self) -> impl Iterator<Item=T> + 'a  where T: Step {
        self.members.iter().flat_map(|m| m.iter())
    }
}

impl<T: Atom> Debug for Class<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "[{}]", self.members.iter().map(|m| format!("{:?}", m)).join(","))
    }
}

impl<'a, T: 'a + Atom> set::Contains<&'a ClassMember<T>> for Class<T> {
    fn contains(&self, cm: &'a ClassMember<T>) -> bool {
        self.iter_members().any(|m| cm == m)
    }
}
impl<T: Atom> set::Contains<T> for Class<T> {
    /// Check if any member of the class matches the given atom.
    fn contains(&self, x: T) -> bool {
        self.matches(x)
    }
}

impl<T: Atom> set::IsSubsetOf<T> for Class<T> {
    /// A class *can* be a subset of an atom -- as long as that atom is the
    /// only member of the class!
    #[inline]
    fn is_subset_of(&self, atom: &T) -> bool {
        self.contains(*atom) && self.len() == 1
    }
}

impl<T: Atom> set::IsSubsetOf<Anchor<T>> for Class<T> {
    /// A class is never a subset of an anchor.
    #[inline(always)]
    fn is_subset_of(&self, _: &Anchor<T>) -> bool {
        false
    }
}

impl<T: Atom> super::AtomicLen for Class<T> {
    fn atomic_len(&self) -> super::SizeBound {
        super::SizeBound::Exact(1)
    }
}

// FIXME [optimize] I suspect that we need some sort of specialized data
// structure or clever algorithm in order to perform this test in a time better
// than O(N·M).
impl<T: Atom> set::IsSubsetOf<Class<T>> for Class<T> {
    /// A class `A` is a subset of another class `B` if all atoms in `A` are
    /// also in `B`.
    fn is_subset_of(&self, other: &Self) -> bool {
        self.iter().all(|a| other.contains(a))
    }
}
impl<T: Atom> set::IsSubsetOf<Sequence<T>> for Class<T> {
    /// Any single-atom pattern like Class cannot match the same set of inputs
    /// as a Sequence, because a Sequence is guaranteed to contain at least
    /// two elements.
    #[inline]
    fn is_subset_of(&self, _: &Sequence<T>) -> bool {
        false
    }
}

impl<T: Atom> set::IsSubsetOf<Union<T>> for Class<T> {
    /// A class is a subset of a union if all members of the class are subsets
    /// of the union.
    fn is_subset_of(&self, union: &Union<T>) -> bool {
        self.iter().all(|a| a.is_subset_of(union))
    }
}

impl<T: Atom> set::IsSubsetOf<Repetition<T>> for Class<T> {
    fn is_subset_of(&self, rep: &Repetition<T>) -> bool {
        self.iter().all(|a| a.is_subset_of(rep.element())) && rep.count().contains(1)
    }
}

impl<T: Atom> set::IsSubsetOf<Element<T>> for Class<T> {
    /// A class is always a subset of a wildcard, and never a subset of an
    /// anchor.  See the other `IsSubsetOf` implementations for `Class` for
    /// explanations of how this trait is implements for the other variants of
    /// `Element`.
    fn is_subset_of(&self, other: &Element<T>) -> bool {
        match other {
            // Anything we can do, Wildcard can do betteeeeer...
            &Element::Wildcard => true,

            &Element::Atom(a) => self.is_subset_of(&a),
            &Element::Class(ref d) => self.is_subset_of(d),
            &Element::Tagged{ref element, ..} => self.is_subset_of(&**element),
            &Element::Sequence(ref seq) => self.is_subset_of(seq),
            &Element::Union(ref union) => self.is_subset_of(union),
            &Element::Repeat(ref rep) => self.is_subset_of(rep),
            &Element::Anchor(_) => false,
        }
    }
}

impl<T: Atom, U> FromIterator<U> for Class<T>
where ClassMember<T>: From<U> {
    fn from_iter<I>(iter: I) -> Self
        where I: IntoIterator<Item=U> {
        Class{members: Vec::from_iter(iter.into_iter().map(|x| x.into()))}
    }
}

apply_attrs! {
    cfg(feature = "regex") => {
        macro_rules! impl_class_from {
            ($T: ty, $from: ty) => {
                impl From<$from> for Class<$T> {
                    fn from(c: $from) -> Self {
                        Self::new(c.into_iter().map(|cr| if cr.end != cr.start { ClassMember::Range(cr.start.into(), cr.end.into()) }
                                                    else { ClassMember::Atom(cr.start.into()) }))
                    }
                }
            };
        }

        impl_class_from!(char, regex_syntax::CharClass);
        impl_class_from!(u8, regex_syntax::ByteClass);
        impl_class_from!(ByteOrChar, regex_syntax::CharClass);
        impl_class_from!(ByteOrChar, regex_syntax::ByteClass);
    }
}
<commit_msg>pattern/class: add `Step` bound on type param to all `IsSubsetOf` impls<commit_after>//! Atom-class implementation.

use std::ops::{BitXor, Sub, Range, RangeInclusive};
use std::iter::FromIterator;
use std::fmt::{self, Debug};

#[cfg(feature = "regex")] use regex_syntax;
use itertools::Itertools;

use super::{Anchor, Atom, ByteOrChar, Element, Repetition, Sequence, Union};
use super::{Distance, Step};
use util::set::{self, Contains};

// ----------------------------------------------------------------
// Class

/// Any member of a class of atoms
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Eq, Ord)]
pub enum ClassMember<T: Atom> {
    /// A single atom
    Atom(T),

    /// A range of atoms
    Range(T, T)
}

impl<T: Atom> ClassMember<T> {
    /// Map the class-member's atoms to a different type
    pub fn map_atoms<U, F>(self, f: F) -> ClassMember<U>
        where F: Fn(T) -> U,
              U: Atom
    {
        match self {
            ClassMember::Atom(a) => ClassMember::Atom(f(a)),
            ClassMember::Range(a, b) => ClassMember::Range(f(a), f(b))
        }
    }

    /// Get the number of individual atoms that this class member represents.
    pub fn len(&self) -> usize
        where T: Distance
    {
        match self {
            &ClassMember::Atom(_) => 1,
            &ClassMember::Range(first, last) => last.distance(&first) + 1,
        }
    }

    /// Fetch an iterator over each discrete value contained by this member.
    #[inline]
     pub fn iter(&self) -> impl Iterator<Item=T>
        where T: Step {
        match self {
            &ClassMember::Atom(a) => AtomIter(a...a),
            &ClassMember::Range(a, b) => AtomIter(a...b),
        }
    }
}

impl<T: Atom> set::Contains<T> for ClassMember<T> {
    /// Check if the member is or contains a particular atom.
    fn contains(&self, x: T) -> bool {
        match self {
            &ClassMember::Atom(a) => x == a,
            &ClassMember::Range(a, b) => (a...b).contains(x),
        }
    }
}

/// Iterator over the atoms specified by a class member.
struct AtomIter<T: Atom>(RangeInclusive<T>);

impl<T: Atom> Iterator for AtomIter<T>
        where T: Step
{
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        match self.0 {
            RangeInclusive::NonEmpty{start, end} => {
                let next_start = start.increment();
                if next_start > end {
                    self.0 = RangeInclusive::Empty{at: end};
                } else {
                    self.0 = RangeInclusive::NonEmpty{start: next_start, end: end};
                }
                Some(start) },
            RangeInclusive::Empty{..} => None
        }
    }
}

/*impl<T: Atom> Debug for ClassMember<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ClassMember::Atom(a) => write!(f, "{:?}", a),
            ClassMember::Range(first, last) => write!(f, "{:?}-{:?}", first, last)
        }
    }
}*/

impl<T: Atom> From<T> for ClassMember<T> {
    fn from(a: T) -> Self { ClassMember::Atom(a) }
}

impl From<regex_syntax::ClassRange> for ClassMember<ByteOrChar> {
    fn from(cr: regex_syntax::ClassRange) -> Self {
        ClassMember::Range(cr.start.into(), cr.end.into())
    }
}

impl From<regex_syntax::ByteRange> for ClassMember<ByteOrChar> {
    fn from(cr: regex_syntax::ByteRange) -> Self {
        ClassMember::Range(cr.start.into(), cr.end.into())
    }
}

impl From<regex_syntax::ClassRange> for ClassMember<char> {
    fn from(cr: regex_syntax::ClassRange) -> Self {
        ClassMember::Range(cr.start, cr.end)
    }
}

impl From<regex_syntax::ByteRange> for ClassMember<u8> {
    fn from(cr: regex_syntax::ByteRange) -> Self {
        ClassMember::Range(cr.start, cr.end)
    }
}

impl<T: Atom> From<Range<T>> for ClassMember<T>
    where T: Sub<usize,Output=T> + Sub<T,Output=usize>
{
    fn from(r: Range<T>) -> Self {
        if r.end - r.start > 1 { ClassMember::Range(r.start, r.end - 1usize) }
        else { ClassMember::Atom(r.start) }
    }
}
// ----------------------------------------------------------------

/// A set of atoms and/or ranges of atoms.
#[derive(Clone, PartialEq, PartialOrd, Eq, Ord)]
pub struct Class<T: Atom> {
    members: Vec<ClassMember<T>>
}

// FIXME: [optimize] All operations on a Class<T>, including observer and
// predicate methods, currently require at least O(N) time.  We could probably
// improve this by changing the storage representation.
impl<T: Atom> Class<T> {
    /// Create a new Class instance with members taken from the given iterator.
    pub fn new<I,U>(items: I) -> Self
        where I: IntoIterator<Item=U>,
              ClassMember<T>: From<U> {
        Class{members: Vec::from_iter(items.into_iter().map(|x| x.into()))}
    }

    /// Create a new Class instance from an iterator of ClassMember instances.
    pub fn from_members<I>(m: I) -> Self
        where I: IntoIterator<Item=ClassMember<T>>
    {
        Class{members: m.into_iter().collect()}
    }

    /// Map the atoms in the class to a different atom type.
    pub fn map_atoms<U, F>(self, f: F) -> Class<U>
        where F: Fn(T) -> U,
              U: Atom
    {
        Class::new(self.members.into_iter().map(|m| m.map_atoms(&f)))
    }

    /// Check whether the class would match a particular atom.
    pub fn matches(&self, x: T) -> bool {
        self.members.iter().any(|&range| range.contains(x))
    }

    /// Get the number of members (atoms) in the class.
    ///
    /// Ranges of atoms count as the number of atoms in each range.
    pub fn len(&self) -> usize where T: Distance {
        self.members.iter().map(|m| m.len()).sum()
    }

    /// Check if the class contains zero members.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.members.is_empty()
    }

    /// Fetch an iterator over the members specified in the class.
    pub fn iter_members<'a>(&'a self) -> impl Iterator<Item=&'a ClassMember<T>> {
        self.members.iter()
    }

    /// Fetch an iterator over the atoms specified in the class.
    pub fn iter<'a>(&'a self) -> impl Iterator<Item=T> + 'a  where T: Step {
        self.members.iter().flat_map(|m| m.iter())
    }
}

impl<T: Atom> Debug for Class<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "[{}]", self.members.iter().map(|m| format!("{:?}", m)).join(","))
    }
}

impl<'a, T: 'a + Atom> set::Contains<&'a ClassMember<T>> for Class<T> {
    fn contains(&self, cm: &'a ClassMember<T>) -> bool {
        self.iter_members().any(|m| cm == m)
    }
}
impl<T: Atom> set::Contains<T> for Class<T> {
    /// Check if any member of the class matches the given atom.
    fn contains(&self, x: T) -> bool {
        self.matches(x)
    }
}

impl<T: Atom + Step> set::IsSubsetOf<T> for Class<T> {
    /// A class *can* be a subset of an atom -- as long as that atom is the
    /// only member of the class!
    #[inline]
    fn is_subset_of(&self, atom: &T) -> bool {
        self.contains(*atom) && self.len() == 1
    }
}

impl<T: Atom> set::IsSubsetOf<Anchor<T>> for Class<T> {
    /// A class is never a subset of an anchor.
    #[inline(always)]
    fn is_subset_of(&self, _: &Anchor<T>) -> bool {
        false
    }
}

impl<T: Atom> super::AtomicLen for Class<T> {
    fn atomic_len(&self) -> super::SizeBound {
        super::SizeBound::Exact(1)
    }
}

// FIXME [optimize] I suspect that we need some sort of specialized data
// structure or clever algorithm in order to perform this test in a time better
// than O(N·M).
impl<T: Atom + Step> set::IsSubsetOf<Class<T>> for Class<T> {
    /// A class `A` is a subset of another class `B` if all atoms in `A` are
    /// also in `B`.
    fn is_subset_of(&self, other: &Self) -> bool {
        self.iter().all(|m| other.contains(m))
    }
}
impl<T: Atom> set::IsSubsetOf<Sequence<T>> for Class<T> {
    /// Any single-atom pattern like Class cannot match the same set of inputs
    /// as a Sequence, because a Sequence is guaranteed to contain at least
    /// two elements.
    #[inline]
    fn is_subset_of(&self, _: &Sequence<T>) -> bool {
        false
    }
}

impl<T: Atom + Step> set::IsSubsetOf<Union<T>> for Class<T> {
    /// A class is a subset of a union if all members of the class are subsets
    /// of the union.
    fn is_subset_of(&self, union: &Union<T>) -> bool {
        self.iter().all(|a| a.is_subset_of(union))
    }
}

impl<T: Atom + Step> set::IsSubsetOf<Repetition<T>> for Class<T> {
    fn is_subset_of(&self, rep: &Repetition<T>) -> bool {
        self.iter().all(|a| a.is_subset_of(rep.element())) && rep.count().contains(1)
    }
}

impl<T: Atom> set::IsSubsetOf<Element<T>> for Class<T> {
    /// A class is always a subset of a wildcard, and never a subset of an
    /// anchor.  See the other `IsSubsetOf` implementations for `Class` for
    /// explanations of how this trait is implements for the other variants of
    /// `Element`.
    fn is_subset_of(&self, other: &Element<T>) -> bool {
        match other {
            // Anything we can do, Wildcard can do betteeeeer...
            &Element::Wildcard => true,

            &Element::Atom(a) => self.is_subset_of(&a),
            &Element::Class(ref d) => self.is_subset_of(d),
            &Element::Tagged{ref element, ..} => self.is_subset_of(&**element),
            &Element::Sequence(ref seq) => self.is_subset_of(seq),
            &Element::Union(ref union) => self.is_subset_of(union),
            &Element::Repeat(ref rep) => self.is_subset_of(rep),
            &Element::Anchor(_) => false,
        }
    }
}

impl<T: Atom, U> FromIterator<U> for Class<T>
where ClassMember<T>: From<U> {
    fn from_iter<I>(iter: I) -> Self
        where I: IntoIterator<Item=U> {
        Class{members: Vec::from_iter(iter.into_iter().map(|x| x.into()))}
    }
}

apply_attrs! {
    cfg(feature = "regex") => {
        macro_rules! impl_class_from {
            ($T: ty, $from: ty) => {
                impl From<$from> for Class<$T> {
                    fn from(c: $from) -> Self {
                        Self::new(c.into_iter().map(|cr| if cr.end != cr.start { ClassMember::Range(cr.start.into(), cr.end.into()) }
                                                    else { ClassMember::Atom(cr.start.into()) }))
                    }
                }
            };
        }

        impl_class_from!(char, regex_syntax::CharClass);
        impl_class_from!(u8, regex_syntax::ByteClass);
        impl_class_from!(ByteOrChar, regex_syntax::CharClass);
        impl_class_from!(ByteOrChar, regex_syntax::ByteClass);
    }
}
<|endoftext|>


<commit_before>use crate::math::{EARTH_RADIUS_MAX_M, EARTH_RADIUS_MIN_M};
use crate::read_write::{Encoding, NodeWriter, OpenMode};
use crate::s2_cells::{S2CellMeta, S2Meta};
use crate::{AttributeData, AttributeDataType, PointsBatch};
use cgmath::{EuclideanSpace, InnerSpace, Point3};
use collision::{Aabb, Aabb3};
use fnv::FnvHashMap;
use lru::LruCache;
use s2::cellid::CellID;
use s2::point::Point;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::io::{Error, ErrorKind, Result};
use std::iter::Iterator;
use std::path::PathBuf;

/// The actual number of underlying writers is MAX_NUM_NODE_WRITERS * num_attributes.
const MAX_NUM_NODE_WRITERS: usize = 25;
/// Corresponds to cells of up to about 160m x 160m.
const S2_SPLIT_LEVEL: u64 = 16;

pub struct S2Splitter<W> {
    writers: LruCache<CellID, W>,
    already_opened_writers: HashSet<CellID>,
    cell_stats: FnvHashMap<CellID, S2CellMeta>,
    bounding_box: Option<Aabb3<f64>>,
    attributes_seen: BTreeMap<String, AttributeDataType>,
    encoding: Encoding,
    open_mode: OpenMode,
    stem: PathBuf,
}

impl<W> NodeWriter<PointsBatch> for S2Splitter<W>
where
    W: NodeWriter<PointsBatch>,
{
    fn new(path: impl Into<PathBuf>, encoding: Encoding, open_mode: OpenMode) -> Self {
        let writers = LruCache::new(MAX_NUM_NODE_WRITERS);
        let already_opened_writers = HashSet::new();
        let stem = path.into();
        let cell_stats = FnvHashMap::default();
        let bounding_box = None;
        let attributes_seen = BTreeMap::new();
        S2Splitter {
            writers,
            already_opened_writers,
            cell_stats,
            bounding_box,
            attributes_seen,
            encoding,
            open_mode,
            stem,
        }
    }

    fn write(&mut self, points_batch: &PointsBatch) -> Result<()> {
        self.check_attributes(points_batch)?;
        let mut batches_by_s2_cell = HashMap::new();
        for (i, pos) in points_batch.position.iter().enumerate() {
            let radius = pos.magnitude();
            if radius > EARTH_RADIUS_MAX_M || radius < EARTH_RADIUS_MIN_M {
                let msg = format!(
                    "Point ({}, {}, {}) is not a valid ECEF point",
                    pos.x, pos.y, pos.z
                );
                return Err(Error::new(ErrorKind::InvalidInput, msg));
            }
            let p3 = Point3::from_vec(*pos);
            let b = self.bounding_box.get_or_insert(Aabb3::new(p3, p3));
            *b = b.grow(p3);
            let s2_point = Point::from_coords(pos.x, pos.y, pos.z);
            let s2_cell_id = CellID::from(s2_point).parent(S2_SPLIT_LEVEL);
            self.cell_stats
                .entry(s2_cell_id)
                .or_insert(S2CellMeta { num_points: 0 })
                .num_points += 1;
            let s2_cell_batch = batches_by_s2_cell.entry(s2_cell_id).or_insert(PointsBatch {
                position: Vec::new(),
                attributes: BTreeMap::new(),
            });
            s2_cell_batch.position.push(*pos);
            for (in_key, in_data) in &points_batch.attributes {
                use AttributeData::*;
                let key = in_key.to_string();
                s2_cell_batch
                    .attributes
                    .entry(key)
                    .and_modify(|out_data| match (in_data, out_data) {
                        (U8(in_vec), U8(out_vec)) => out_vec.push(in_vec[i]),
                        (I64(in_vec), I64(out_vec)) => out_vec.push(in_vec[i]),
                        (U64(in_vec), U64(out_vec)) => out_vec.push(in_vec[i]),
                        (F32(in_vec), F32(out_vec)) => out_vec.push(in_vec[i]),
                        (F64(in_vec), F64(out_vec)) => out_vec.push(in_vec[i]),
                        (U8Vec3(in_vec), U8Vec3(out_vec)) => out_vec.push(in_vec[i]),
                        (F64Vec3(in_vec), F64Vec3(out_vec)) => out_vec.push(in_vec[i]),
                        _ => panic!("Input data type unequal output data type."),
                    })
                    .or_insert(match in_data {
                        U8(in_vec) => U8(vec![in_vec[i]]),
                        I64(in_vec) => I64(vec![in_vec[i]]),
                        U64(in_vec) => U64(vec![in_vec[i]]),
                        F32(in_vec) => F32(vec![in_vec[i]]),
                        F64(in_vec) => F64(vec![in_vec[i]]),
                        U8Vec3(in_vec) => U8Vec3(vec![in_vec[i]]),
                        F64Vec3(in_vec) => F64Vec3(vec![in_vec[i]]),
                    });
            }
        }

        for (cell_id, batch) in &batches_by_s2_cell {
            self.writer(cell_id).write(batch)?;
        }
        Ok(())
    }
}

impl<W> S2Splitter<W>
where
    W: NodeWriter<PointsBatch>,
{
    fn writer(&mut self, cell_id: &CellID) -> &mut W {
        let path = self.stem.join(cell_id.to_token());
        if !self.writers.contains(cell_id) {
            let open_mode = if self.open_mode == OpenMode::Append
                || self.already_opened_writers.contains(cell_id)
            {
                OpenMode::Append
            } else {
                self.already_opened_writers.insert(*cell_id);
                OpenMode::Truncate
            };
            self.writers
                .put(*cell_id, W::new(path, self.encoding.clone(), open_mode));
        }
        self.writers.get_mut(cell_id).unwrap()
    }

    /// Records the list of attributes seen in the first batch, and checks
    /// that the following batches contain the same attributes.
    fn check_attributes(&mut self, batch: &PointsBatch) -> Result<()> {
        let mut attr_iter = batch
            .attributes
            .iter()
            .map(|(name, data)| (name, data.data_type()));
        if self.attributes_seen.is_empty() {
            self.attributes_seen
                .extend(attr_iter.map(|(key, val)| (key.to_owned(), val)));
            Ok(())
        } else {
            attr_iter.try_for_each(|(name, dtype)| {
                self.attributes_seen
                    .get(name)
                    .filter(|seen_dtype| **seen_dtype == dtype)
                    .ok_or_else(|| {
                        let msg = format!(
                            "S2Splitter received incompatible data types for attribute {}",
                            name
                        );
                        Error::new(ErrorKind::InvalidInput, msg)
                    })?;
                Ok(())
            })
        }
    }

    pub fn get_meta(self) -> S2Meta {
        S2Meta::new(
            self.cell_stats,
            self.attributes_seen.into_iter().collect(),
            self.bounding_box.unwrap_or_else(Aabb3::zero),
        )
    }
}
<commit_msg>Make split level configurable, set default to 20. (#384)<commit_after>use crate::math::{EARTH_RADIUS_MAX_M, EARTH_RADIUS_MIN_M};
use crate::read_write::{Encoding, NodeWriter, OpenMode};
use crate::s2_cells::{S2CellMeta, S2Meta};
use crate::{AttributeData, AttributeDataType, PointsBatch};
use cgmath::{EuclideanSpace, InnerSpace, Point3};
use collision::{Aabb, Aabb3};
use fnv::FnvHashMap;
use lru::LruCache;
use s2::cellid::CellID;
use s2::point::Point;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::io::{Error, ErrorKind, Result};
use std::iter::Iterator;
use std::path::PathBuf;

/// The actual number of underlying writers is MAX_NUM_NODE_WRITERS * num_attributes.
const MAX_NUM_NODE_WRITERS: usize = 25;
/// Corresponds to cells of up to about 10m x 10m.
const DEFAULT_S2_SPLIT_LEVEL: u64 = 20;

pub struct S2Splitter<W> {
    split_level: u64,
    writers: LruCache<CellID, W>,
    already_opened_writers: HashSet<CellID>,
    cell_stats: FnvHashMap<CellID, S2CellMeta>,
    bounding_box: Option<Aabb3<f64>>,
    attributes_seen: BTreeMap<String, AttributeDataType>,
    encoding: Encoding,
    open_mode: OpenMode,
    stem: PathBuf,
}

impl<W> S2Splitter<W> {
    pub fn with_split_level(
        split_level: u64,
        path: impl Into<PathBuf>,
        encoding: Encoding,
        open_mode: OpenMode,
    ) -> Self {
        S2Splitter {
            split_level,
            writers: LruCache::new(MAX_NUM_NODE_WRITERS),
            already_opened_writers: HashSet::new(),
            cell_stats: FnvHashMap::default(),
            bounding_box: None,
            attributes_seen: BTreeMap::new(),
            encoding,
            open_mode,
            stem: path.into(),
        }
    }
}

impl<W> NodeWriter<PointsBatch> for S2Splitter<W>
where
    W: NodeWriter<PointsBatch>,
{
    fn new(path: impl Into<PathBuf>, encoding: Encoding, open_mode: OpenMode) -> Self {
        Self::with_split_level(DEFAULT_S2_SPLIT_LEVEL, path, encoding, open_mode)
    }

    fn write(&mut self, points_batch: &PointsBatch) -> Result<()> {
        self.check_attributes(points_batch)?;
        let mut batches_by_s2_cell = HashMap::new();
        for (i, pos) in points_batch.position.iter().enumerate() {
            let radius = pos.magnitude();
            if radius > EARTH_RADIUS_MAX_M || radius < EARTH_RADIUS_MIN_M {
                let msg = format!(
                    "Point ({}, {}, {}) is not a valid ECEF point",
                    pos.x, pos.y, pos.z
                );
                return Err(Error::new(ErrorKind::InvalidInput, msg));
            }
            let p3 = Point3::from_vec(*pos);
            let b = self.bounding_box.get_or_insert(Aabb3::new(p3, p3));
            *b = b.grow(p3);
            let s2_point = Point::from_coords(pos.x, pos.y, pos.z);
            let s2_cell_id = CellID::from(s2_point).parent(self.split_level);
            self.cell_stats
                .entry(s2_cell_id)
                .or_insert(S2CellMeta { num_points: 0 })
                .num_points += 1;
            let s2_cell_batch = batches_by_s2_cell.entry(s2_cell_id).or_insert(PointsBatch {
                position: Vec::new(),
                attributes: BTreeMap::new(),
            });
            s2_cell_batch.position.push(*pos);
            for (in_key, in_data) in &points_batch.attributes {
                use AttributeData::*;
                let key = in_key.to_string();
                s2_cell_batch
                    .attributes
                    .entry(key)
                    .and_modify(|out_data| match (in_data, out_data) {
                        (U8(in_vec), U8(out_vec)) => out_vec.push(in_vec[i]),
                        (I64(in_vec), I64(out_vec)) => out_vec.push(in_vec[i]),
                        (U64(in_vec), U64(out_vec)) => out_vec.push(in_vec[i]),
                        (F32(in_vec), F32(out_vec)) => out_vec.push(in_vec[i]),
                        (F64(in_vec), F64(out_vec)) => out_vec.push(in_vec[i]),
                        (U8Vec3(in_vec), U8Vec3(out_vec)) => out_vec.push(in_vec[i]),
                        (F64Vec3(in_vec), F64Vec3(out_vec)) => out_vec.push(in_vec[i]),
                        _ => panic!("Input data type unequal output data type."),
                    })
                    .or_insert(match in_data {
                        U8(in_vec) => U8(vec![in_vec[i]]),
                        I64(in_vec) => I64(vec![in_vec[i]]),
                        U64(in_vec) => U64(vec![in_vec[i]]),
                        F32(in_vec) => F32(vec![in_vec[i]]),
                        F64(in_vec) => F64(vec![in_vec[i]]),
                        U8Vec3(in_vec) => U8Vec3(vec![in_vec[i]]),
                        F64Vec3(in_vec) => F64Vec3(vec![in_vec[i]]),
                    });
            }
        }

        for (cell_id, batch) in &batches_by_s2_cell {
            self.writer(cell_id).write(batch)?;
        }
        Ok(())
    }
}

impl<W> S2Splitter<W>
where
    W: NodeWriter<PointsBatch>,
{
    fn writer(&mut self, cell_id: &CellID) -> &mut W {
        let path = self.stem.join(cell_id.to_token());
        if !self.writers.contains(cell_id) {
            let open_mode = if self.open_mode == OpenMode::Append
                || self.already_opened_writers.contains(cell_id)
            {
                OpenMode::Append
            } else {
                self.already_opened_writers.insert(*cell_id);
                OpenMode::Truncate
            };
            self.writers
                .put(*cell_id, W::new(path, self.encoding.clone(), open_mode));
        }
        self.writers.get_mut(cell_id).unwrap()
    }

    /// Records the list of attributes seen in the first batch, and checks
    /// that the following batches contain the same attributes.
    fn check_attributes(&mut self, batch: &PointsBatch) -> Result<()> {
        let mut attr_iter = batch
            .attributes
            .iter()
            .map(|(name, data)| (name, data.data_type()));
        if self.attributes_seen.is_empty() {
            self.attributes_seen
                .extend(attr_iter.map(|(key, val)| (key.to_owned(), val)));
            Ok(())
        } else {
            attr_iter.try_for_each(|(name, dtype)| {
                self.attributes_seen
                    .get(name)
                    .filter(|seen_dtype| **seen_dtype == dtype)
                    .ok_or_else(|| {
                        let msg = format!(
                            "S2Splitter received incompatible data types for attribute {}",
                            name
                        );
                        Error::new(ErrorKind::InvalidInput, msg)
                    })?;
                Ok(())
            })
        }
    }

    pub fn get_meta(self) -> S2Meta {
        S2Meta::new(
            self.cell_stats,
            self.attributes_seen.into_iter().collect(),
            self.bounding_box.unwrap_or_else(Aabb3::zero),
        )
    }
}
<|endoftext|>


<commit_before>use syn;
use quote;

use model::{infer_association_name, Model};
use util::{str_value_of_meta_item, wrap_item_in_const};

pub fn derive_associations(input: syn::DeriveInput) -> quote::Tokens {
    let model = t!(Model::from_item(&input, "Associations"));

    let derived_associations = input.attrs.as_slice().iter().filter_map(|attr| {
        if attr.name() == "belongs_to" {
            Some(
                build_association_options(attr)
                    .map(|options| expand_belongs_to(&model, options))
                    .unwrap_or_else(|e| e),
            )
        } else {
            None
        }
    });

    wrap_item_in_const(
        model.dummy_const_name("ASSOCIATIONS"),
        quote!(#(#derived_associations)*),
    )
}

fn expand_belongs_to(model: &Model, options: AssociationOptions) -> quote::Tokens {
    let parent_struct = options.name;
    let struct_name = &model.name;
    let fields = model.attrs.as_slice();

    let foreign_key_name = options
        .foreign_key_name
        .unwrap_or_else(|| to_foreign_key(parent_struct.as_ref()));
    let foreign_key_name = &foreign_key_name;
    let child_table_name = model.table_name();
    let child_table_name = &child_table_name;
    let foreign_key_attr = fields
        .iter()
        .find(|attr| attr.column_name().as_ref() == foreign_key_name.as_ref())
        .unwrap();
    let foreign_key_ty = &foreign_key_attr.ty;

    // we need to special case foreign keys on with an Option type
    // to allow self referencing joins
    let (foreign_key, foreign_key_ty) = match *foreign_key_ty {
        syn::Ty::Path(None, ref p) if p.segments[0].ident == "Option" => {
            let segment = &p.segments[0];
            let t = match segment.parameters {
                syn::PathParameters::AngleBracketed(ref p) => p.types[0].clone(),
                syn::PathParameters::Parenthesized(_) => unreachable!(),
            };
            (quote!(self.#foreign_key_name.as_ref()), t)
        }
        ref t => (quote!(Some(&self.#foreign_key_name)), t.clone()),
    };

    quote!(
        impl diesel::associations::BelongsTo<#parent_struct> for #struct_name {
            type ForeignKey = #foreign_key_ty;
            type ForeignKeyColumn = #child_table_name::#foreign_key_name;

            fn foreign_key(&self) -> Option<&Self::ForeignKey> {
                #foreign_key
            }

            fn foreign_key_column() -> Self::ForeignKeyColumn {
                #child_table_name::#foreign_key_name
            }
        }
    )
}

struct AssociationOptions {
    name: syn::Ident,
    foreign_key_name: Option<syn::Ident>,
}

fn build_association_options(attr: &syn::Attribute) -> Result<AssociationOptions, quote::Tokens> {
    let usage_error = Err(quote!(
        compile_error!("`#[belongs_to]` must be in the form `#[belongs_to(table_name, foreign_key=\"column_name\")]`")
    ));
    match attr.value {
        syn::MetaItem::List(_, ref options) if options.len() >= 1 => {
            let association_name = match options[0] {
                syn::NestedMetaItem::MetaItem(syn::MetaItem::Word(ref name)) => name.clone(),
                _ => return usage_error,
            };
            let foreign_key_name = options
                .iter()
                .filter_map(|o| match *o {
                    syn::NestedMetaItem::MetaItem(ref mi) => Some(mi),
                    _ => None,
                })
                .find(|a| a.name() == "foreign_key")
                .map(|a| str_value_of_meta_item(a, "foreign_key"))
                .map(syn::Ident::new);
            if options.len() == 1 || (options.len() == 2 && foreign_key_name.is_some()) {
                Ok(AssociationOptions {
                    name: association_name,
                    foreign_key_name: foreign_key_name,
                })
            } else {
                usage_error
            }
        }
        _ => usage_error,
    }
}

fn to_foreign_key(model_name: &str) -> syn::Ident {
    let lower_cased = infer_association_name(model_name);
    syn::Ident::new(format!("{}_id", &lower_cased))
}
<commit_msg>Address review comments<commit_after>use syn;
use quote;

use model::{infer_association_name, Model};
use util::{str_value_of_meta_item, wrap_item_in_const};

pub fn derive_associations(input: syn::DeriveInput) -> quote::Tokens {
    let model = t!(Model::from_item(&input, "Associations"));

    let derived_associations = input
        .attrs
        .as_slice()
        .iter()
        .filter_map(|attr| {
            if attr.name() == "belongs_to" {
                Some(
                    build_association_options(attr)
                        .map(|options| expand_belongs_to(&model, options))
                        .unwrap_or_else(|e| e),
                )
            } else {
                None
            }
        })
        .collect::<Vec<_>>();
    if derived_associations.is_empty() {
        quote!(compile_error!("Could not derive `Associations` without any `#[belongs_to(table_name)]` attribute");)
    } else {
        wrap_item_in_const(
            model.dummy_const_name("ASSOCIATIONS"),
            quote!(#(#derived_associations)*),
        )
    }
}

fn expand_belongs_to(model: &Model, options: AssociationOptions) -> quote::Tokens {
    let parent_struct = options.name;
    let struct_name = &model.name;
    let fields = model.attrs.as_slice();

    let foreign_key_name = options
        .foreign_key_name
        .unwrap_or_else(|| to_foreign_key(parent_struct.as_ref()));
    let child_table_name = model.table_name();
    let child_table_name = &child_table_name;
    let foreign_key_attr = fields
        .iter()
        .find(|attr| *attr.column_name() == foreign_key_name);
    let foreign_key_ty = if let Some(foreign_key_attr) = foreign_key_attr {
        &foreign_key_attr.ty
    } else {
        let msg = format!("The foreign key column {} is not found", foreign_key_name);
        return quote!{
            compile_error!(#msg);
        };
    };

    // we need to special case foreign keys on with an Option type
    // to allow self referencing joins
    let (foreign_key, foreign_key_ty) = ::util::inner_of_option_ty(foreign_key_ty)
        .map(|foreign_key_ty| (quote!(self.#foreign_key_name.as_ref()), foreign_key_ty))
        .unwrap_or_else(|| (quote!(Some(&self.#foreign_key_name)), foreign_key_ty));

    quote!(
        impl diesel::associations::BelongsTo<#parent_struct> for #struct_name {
            type ForeignKey = #foreign_key_ty;
            type ForeignKeyColumn = #child_table_name::#foreign_key_name;

            fn foreign_key(&self) -> Option<&Self::ForeignKey> {
                #foreign_key
            }

            fn foreign_key_column() -> Self::ForeignKeyColumn {
                #child_table_name::#foreign_key_name
            }
        }
    )
}

struct AssociationOptions {
    name: syn::Ident,
    foreign_key_name: Option<syn::Ident>,
}

fn build_association_options(attr: &syn::Attribute) -> Result<AssociationOptions, quote::Tokens> {
    let usage_error = Err(quote!(
        compile_error!("`#[belongs_to]` must be in the form `#[belongs_to(table_name, foreign_key=\"column_name\")]`")
    ));
    match attr.value {
        syn::MetaItem::List(_, ref options) if options.len() >= 1 => {
            let association_name = match options[0] {
                syn::NestedMetaItem::MetaItem(syn::MetaItem::Word(ref name)) => name.clone(),
                _ => return usage_error,
            };
            let foreign_key_name = options
                .iter()
                .filter_map(|o| match *o {
                    syn::NestedMetaItem::MetaItem(ref mi) => Some(mi),
                    _ => None,
                })
                .find(|a| a.name() == "foreign_key")
                .map(|a| str_value_of_meta_item(a, "foreign_key"))
                .map(syn::Ident::new);
            if options.len() == 1 || (options.len() == 2 && foreign_key_name.is_some()) {
                Ok(AssociationOptions {
                    name: association_name,
                    foreign_key_name: foreign_key_name,
                })
            } else {
                usage_error
            }
        }
        _ => usage_error,
    }
}

fn to_foreign_key(model_name: &str) -> syn::Ident {
    let lower_cased = infer_association_name(model_name);
    syn::Ident::new(format!("{}_id", &lower_cased))
}
<|endoftext|>


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

use futures::{Async, Future, Poll, Stream};

use tokio_core;
use tokio_core::reactor::Core;

use trust_dns::udp::UdpStream;
use trust_dns::tcp::TcpStream;

#[cfg(feature = "tls")]
use trust_dns_openssl::{TlsStream, tls_server};

#[cfg(feature = "tls")]
use trust_dns_openssl::tls_server::*;

use server::{Request, RequestHandler, RequestStream, ResponseHandle, TimeoutStream};

// TODO, would be nice to have a Slab for buffers here...

/// A Futures based implementation of a DNS server
pub struct ServerFuture<T: RequestHandler + 'static> {
    io_loop: Core,
    handler: Arc<T>,
}

impl<T: RequestHandler> ServerFuture<T> {
    /// Creates a new ServerFuture with the specified Handler.
    pub fn new(handler: T) -> io::Result<ServerFuture<T>> {
        Ok(ServerFuture {
            io_loop: try!(Core::new()),
            handler: Arc::new(handler),
        })
    }

    /// Register a UDP socket. Should be bound before calling this function.
    pub fn register_socket(&self, socket: std::net::UdpSocket) {
        debug!("registered udp: {:?}", socket);

        // create the new UdpStream
        let (buf_stream, stream_handle) = UdpStream::with_bound(socket, &self.io_loop.handle());
        let request_stream = RequestStream::new(buf_stream, stream_handle);
        let handler = self.handler.clone();

        // this spawns a ForEach future which handles all the requests into a Handler.
        self.io_loop.handle().spawn(
            // TODO dedup with below into generic func
            request_stream
                .for_each(move |(request, response_handle)| {
                    Self::handle_request(request, response_handle, handler.clone())
                })
                .map_err(|e| debug!("error in UDP request_stream handler: {}", e)),
        );
    }

    /// Register a TcpListener to the Server. This should already be bound to either an IPv6 or an
    ///  IPv4 address.
    ///
    /// To make the server more resilient to DOS issues, there is a timeout. Care should be taken
    ///  to not make this too low depending on use cases.
    ///
    /// # Arguments
    /// * `listener` - a bound TCP socket
    /// * `timeout` - timeout duration of incoming requests, any connection that does not send
    ///               requests within this time period will be closed. In the future it should be
    ///               possible to create long-lived queries, but these should be from trusted sources
    ///               only, this would require some type of whitelisting.
    pub fn register_listener(
        &self,
        listener: std::net::TcpListener,
        timeout: Duration,
    ) -> io::Result<()> {
        let handle = self.io_loop.handle();
        let handler = self.handler.clone();
        // TODO: this is an awkward interface with socketaddr...
        let addr = try!(listener.local_addr());
        let listener = tokio_core::net::TcpListener::from_listener(listener, &addr, &handle)
            .expect("could not register listener");
        debug!("registered tcp: {:?}", listener);

        // for each incoming request...
        self.io_loop.handle().spawn(
            listener
                .incoming()
                .for_each(move |(tcp_stream, src_addr)| {
                    debug!("accepted request from: {}", src_addr);
                    // take the created stream...
                    let (buf_stream, stream_handle) = TcpStream::from_stream(tcp_stream, src_addr);
                    let timeout_stream = try!(TimeoutStream::new(buf_stream, timeout, &handle));
                    let request_stream = RequestStream::new(timeout_stream, stream_handle);
                    let handler = handler.clone();

                    // and spawn to the io_loop
                    handle.spawn(
                        request_stream
                            .for_each(move |(request, response_handle)| {
                                Self::handle_request(request, response_handle, handler.clone())
                            })
                            .map_err(move |e| {
                                debug!(
                                    "error in TCP request_stream src: {:?} error: {}",
                                    src_addr,
                                    e
                                )
                            }),
                    );

                    Ok(())
                })
                .map_err(|e| debug!("error in inbound tcp_stream: {}", e)),
        );

        Ok(())
    }

    /// Register a TlsListener to the Server. The TlsListener should already be bound to either an
    /// IPv6 or an IPv4 address.
    ///
    /// To make the server more resilient to DOS issues, there is a timeout. Care should be taken
    ///  to not make this too low depending on use cases.
    ///
    /// # Arguments
    /// * `listener` - a bound TCP (needs to be on a different port from standard TCP connections) socket
    /// * `timeout` - timeout duration of incoming requests, any connection that does not send
    ///               requests within this time period will be closed. In the future it should be
    ///               possible to create long-lived queries, but these should be from trusted sources
    ///               only, this would require some type of whitelisting.
    /// * `pkcs12` - certificate used to announce to clients
    #[cfg(feature = "tls")]
    pub fn register_tls_listener(
        &self,
        listener: std::net::TcpListener,
        timeout: Duration,
        pkcs12: ParsedPkcs12,
    ) -> io::Result<()> {
        let handle = self.io_loop.handle();
        let handler = self.handler.clone();
        // TODO: this is an awkward interface with socketaddr...
        let addr = listener.local_addr().expect("listener is not bound?");
        let listener = tokio_core::net::TcpListener::from_listener(listener, &addr, &handle)
            .expect("could not register listener");
        debug!("registered tcp: {:?}", listener);

        let tls_acceptor = tls_server::new_acceptor(&pkcs12)?;

        // for each incoming request...
        self.io_loop.handle().spawn(
        listener.incoming()
                .for_each(move |(tcp_stream, src_addr)| {
                  debug!("accepted request from: {}", src_addr);
                  let timeout = timeout.clone();
                  let handle = handle.clone();
                  let handler = handler.clone();

                  // take the created stream...
                  tls_acceptor.accept_async(tcp_stream)
                              .map_err(|e| io::Error::new(io::ErrorKind::ConnectionRefused, format!("tls error: {}", e)))
                              .and_then(move |tls_stream| {
                                  let (buf_stream, stream_handle) =
                                      TlsStream::from_stream(tls_stream, src_addr.clone());
                                  let timeout_stream = try!(TimeoutStream::new(buf_stream, timeout, &handle));
                                  let request_stream = RequestStream::new(timeout_stream, stream_handle);
                                  let handler = handler.clone();

                                  // and spawn to the io_loop
                                  handle.spawn(
                                  request_stream.for_each(move |(request, response_handle)| {
                                      Self::handle_request(request, response_handle, handler.clone())
                                  })
                              .map_err(move |e| debug!("error in TCP request_stream src: {:?} error: {}", src_addr, e))
                              );

                              Ok(())
                            })
                            //.map_err(move |e| debug!("error TLS handshake: {:?} error: {:?}", src_addr, e))
              })
              .map_err(|e| debug!("error in inbound tcp_stream: {}", e))
    );

        Ok(())
    }

    /// TODO how to do threads? should we do a bunch of listener threads and then query threads?
    /// Ideally the processing would be n-threads for recieving, which hand off to m-threads for
    ///  request handling. It would generally be the case that n <= m.
    pub fn listen(&mut self) -> io::Result<()> {
        info!("Server starting up");
        try!(self.io_loop.run(Forever));

        Err(io::Error::new(
            io::ErrorKind::Interrupted,
            "Server stopping due to interruption",
        ))
    }

    /// Returns a reference to the tokio core loop driving this Server instance
    pub fn tokio_core(&mut self) -> &mut Core {
        &mut self.io_loop
    }

    fn handle_request(
        request: Request,
        mut response_handle: ResponseHandle,
        handler: Arc<T>,
    ) -> io::Result<()> {
        info!(
            "request: {} src: {} type: {:?} op_code: {:?} {}",
            request.message.id(),
            request.src,
            request.message.message_type(),
            request.message.op_code(),
            request.message.queries().first().map(|q| q.to_string()).unwrap_or_else(|| "empty_queries".to_string()),
        );
        let response = handler.handle_request(&request);

        info!(
            "request: {} response_code: {} answers: {} name_servers: {} additionals: {}",
            request.message.id(),
            response.response_code(),
            response.answers().len(),
            response.name_servers().len(),
            response.additionals().len(),
        );
        response_handle.send(response)
    }
}

struct Forever;

impl Future for Forever {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        // run forever...
        Ok(Async::NotReady)
    }
}
<commit_msg>add dnssec to request log<commit_after>// Copyright 2015-2016 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.
use std;
use std::io;
use std::sync::Arc;
use std::time::Duration;

use futures::{Async, Future, Poll, Stream};

use tokio_core;
use tokio_core::reactor::Core;

use trust_dns::udp::UdpStream;
use trust_dns::tcp::TcpStream;

#[cfg(feature = "tls")]
use trust_dns_openssl::{TlsStream, tls_server};

#[cfg(feature = "tls")]
use trust_dns_openssl::tls_server::*;

use server::{Request, RequestHandler, RequestStream, ResponseHandle, TimeoutStream};

// TODO, would be nice to have a Slab for buffers here...

/// A Futures based implementation of a DNS server
pub struct ServerFuture<T: RequestHandler + 'static> {
    io_loop: Core,
    handler: Arc<T>,
}

impl<T: RequestHandler> ServerFuture<T> {
    /// Creates a new ServerFuture with the specified Handler.
    pub fn new(handler: T) -> io::Result<ServerFuture<T>> {
        Ok(ServerFuture {
            io_loop: try!(Core::new()),
            handler: Arc::new(handler),
        })
    }

    /// Register a UDP socket. Should be bound before calling this function.
    pub fn register_socket(&self, socket: std::net::UdpSocket) {
        debug!("registered udp: {:?}", socket);

        // create the new UdpStream
        let (buf_stream, stream_handle) = UdpStream::with_bound(socket, &self.io_loop.handle());
        let request_stream = RequestStream::new(buf_stream, stream_handle);
        let handler = self.handler.clone();

        // this spawns a ForEach future which handles all the requests into a Handler.
        self.io_loop.handle().spawn(
            // TODO dedup with below into generic func
            request_stream
                .for_each(move |(request, response_handle)| {
                    Self::handle_request(request, response_handle, handler.clone())
                })
                .map_err(|e| debug!("error in UDP request_stream handler: {}", e)),
        );
    }

    /// Register a TcpListener to the Server. This should already be bound to either an IPv6 or an
    ///  IPv4 address.
    ///
    /// To make the server more resilient to DOS issues, there is a timeout. Care should be taken
    ///  to not make this too low depending on use cases.
    ///
    /// # Arguments
    /// * `listener` - a bound TCP socket
    /// * `timeout` - timeout duration of incoming requests, any connection that does not send
    ///               requests within this time period will be closed. In the future it should be
    ///               possible to create long-lived queries, but these should be from trusted sources
    ///               only, this would require some type of whitelisting.
    pub fn register_listener(
        &self,
        listener: std::net::TcpListener,
        timeout: Duration,
    ) -> io::Result<()> {
        let handle = self.io_loop.handle();
        let handler = self.handler.clone();
        // TODO: this is an awkward interface with socketaddr...
        let addr = try!(listener.local_addr());
        let listener = tokio_core::net::TcpListener::from_listener(listener, &addr, &handle)
            .expect("could not register listener");
        debug!("registered tcp: {:?}", listener);

        // for each incoming request...
        self.io_loop.handle().spawn(
            listener
                .incoming()
                .for_each(move |(tcp_stream, src_addr)| {
                    debug!("accepted request from: {}", src_addr);
                    // take the created stream...
                    let (buf_stream, stream_handle) = TcpStream::from_stream(tcp_stream, src_addr);
                    let timeout_stream = try!(TimeoutStream::new(buf_stream, timeout, &handle));
                    let request_stream = RequestStream::new(timeout_stream, stream_handle);
                    let handler = handler.clone();

                    // and spawn to the io_loop
                    handle.spawn(
                        request_stream
                            .for_each(move |(request, response_handle)| {
                                Self::handle_request(request, response_handle, handler.clone())
                            })
                            .map_err(move |e| {
                                debug!(
                                    "error in TCP request_stream src: {:?} error: {}",
                                    src_addr,
                                    e
                                )
                            }),
                    );

                    Ok(())
                })
                .map_err(|e| debug!("error in inbound tcp_stream: {}", e)),
        );

        Ok(())
    }

    /// Register a TlsListener to the Server. The TlsListener should already be bound to either an
    /// IPv6 or an IPv4 address.
    ///
    /// To make the server more resilient to DOS issues, there is a timeout. Care should be taken
    ///  to not make this too low depending on use cases.
    ///
    /// # Arguments
    /// * `listener` - a bound TCP (needs to be on a different port from standard TCP connections) socket
    /// * `timeout` - timeout duration of incoming requests, any connection that does not send
    ///               requests within this time period will be closed. In the future it should be
    ///               possible to create long-lived queries, but these should be from trusted sources
    ///               only, this would require some type of whitelisting.
    /// * `pkcs12` - certificate used to announce to clients
    #[cfg(feature = "tls")]
    pub fn register_tls_listener(
        &self,
        listener: std::net::TcpListener,
        timeout: Duration,
        pkcs12: ParsedPkcs12,
    ) -> io::Result<()> {
        let handle = self.io_loop.handle();
        let handler = self.handler.clone();
        // TODO: this is an awkward interface with socketaddr...
        let addr = listener.local_addr().expect("listener is not bound?");
        let listener = tokio_core::net::TcpListener::from_listener(listener, &addr, &handle)
            .expect("could not register listener");
        debug!("registered tcp: {:?}", listener);

        let tls_acceptor = tls_server::new_acceptor(&pkcs12)?;

        // for each incoming request...
        self.io_loop.handle().spawn(
        listener.incoming()
                .for_each(move |(tcp_stream, src_addr)| {
                  debug!("accepted request from: {}", src_addr);
                  let timeout = timeout.clone();
                  let handle = handle.clone();
                  let handler = handler.clone();

                  // take the created stream...
                  tls_acceptor.accept_async(tcp_stream)
                              .map_err(|e| io::Error::new(io::ErrorKind::ConnectionRefused, format!("tls error: {}", e)))
                              .and_then(move |tls_stream| {
                                  let (buf_stream, stream_handle) =
                                      TlsStream::from_stream(tls_stream, src_addr.clone());
                                  let timeout_stream = try!(TimeoutStream::new(buf_stream, timeout, &handle));
                                  let request_stream = RequestStream::new(timeout_stream, stream_handle);
                                  let handler = handler.clone();

                                  // and spawn to the io_loop
                                  handle.spawn(
                                  request_stream.for_each(move |(request, response_handle)| {
                                      Self::handle_request(request, response_handle, handler.clone())
                                  })
                              .map_err(move |e| debug!("error in TCP request_stream src: {:?} error: {}", src_addr, e))
                              );

                              Ok(())
                            })
                            //.map_err(move |e| debug!("error TLS handshake: {:?} error: {:?}", src_addr, e))
              })
              .map_err(|e| debug!("error in inbound tcp_stream: {}", e))
    );

        Ok(())
    }

    /// TODO how to do threads? should we do a bunch of listener threads and then query threads?
    /// Ideally the processing would be n-threads for recieving, which hand off to m-threads for
    ///  request handling. It would generally be the case that n <= m.
    pub fn listen(&mut self) -> io::Result<()> {
        info!("Server starting up");
        try!(self.io_loop.run(Forever));

        Err(io::Error::new(
            io::ErrorKind::Interrupted,
            "Server stopping due to interruption",
        ))
    }

    /// Returns a reference to the tokio core loop driving this Server instance
    pub fn tokio_core(&mut self) -> &mut Core {
        &mut self.io_loop
    }

    fn handle_request(
        request: Request,
        mut response_handle: ResponseHandle,
        handler: Arc<T>,
    ) -> io::Result<()> {
        info!(
            "request: {} src: {} type: {:?} op_code: {:?} dnssec: {} {}",
            request.message.id(),
            request.src,
            request.message.message_type(),
            request.message.op_code(),
            request.message.edns().map_or(false, |edns| edns.dnssec_ok()),
            request.message.queries().first().map(|q| q.to_string()).unwrap_or_else(|| "empty_queries".to_string()),
        );
        let response = handler.handle_request(&request);

        info!(
            "request: {} response_code: {} answers: {} name_servers: {} additionals: {}",
            request.message.id(),
            response.response_code(),
            response.answers().len(),
            response.name_servers().len(),
            response.additionals().len(),
        );
        response_handle.send(response)
    }
}

struct Forever;

impl Future for Forever {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        // run forever...
        Ok(Async::NotReady)
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 Jeffery Olson
//
// Licensed under the 3-Clause BSD License, see LICENSE.txt
// at the top-level of this repository.
// This file may not be copied, modified, or distributed
// except according to those terms.
#![crate_id="http://olsonjeffery.github.io#schemers:0.0.1"]
#![desc = "Simple Scheme Interpreter"]
#![license = "3-Clause BSD"]

extern crate collections = "collections";

use std::from_str::FromStr;
use collections::hashmap::HashMap;

// this is starting out as a straight port of Peter Norvig's
// lis.py (the first iteration) to Rust

#[cfg(not(test))]
fn main() {
    use std::os::args;
    let args: Vec<~str> = args().move_iter().skip(1).collect();
    if args.len() == 0 {
        fail!("no program provided");
    }
    let program: ~str = args.move_iter()
        .fold(~"", |memo, arg| memo + arg + " ").trim().to_owned();
    let parsed_items = parse_str(program);
    println!("{}", parsed_items.print());
}

pub struct Env {
    entries: HashMap<~str, Expr>,
    outer: Option<~Env>
}

#[deriving(Eq, Show, Clone)]
pub enum Expr {
    Atom(AtomVal),
    List(Vec<~Expr>)
}

#[deriving(Eq, Show, Clone)]
pub enum AtomVal {
    Symbol(~str),
    Integer(i64),
    Float(f64)
}

fn eval(expr: Expr, env: Env) -> (Expr, Env) {
    match expr {
        // Atom values and values in the Env
        Atom(Symbol(ref var_name)) => (env.find(var_name), env),
        val @ Atom(_) => (val, env),
        _ => fail!("un-implemented case")
    }
}

fn parse_str(input: ~str) -> Expr {
    parse(&mut tokenize(input))
}

fn pad_input(input: ~str) -> ~str {
    input.replace("(", " ( ").replace(")", " ) ")
}

fn tokenize(input: ~str) -> Vec<~str> {
    pad_input(input).split_str(" ")
        .filter(|i| *i != "").map(|i| i.to_owned()).collect()
}

fn parse(tokens: &mut Vec<~str>) -> Expr {
    let current_token =
        tokens.shift().expect("calling parse() w/ empty token list; shouldn't happen");
    match current_token {
        ref x if *x == ~"(" => {
            let mut list = Vec::new();
            while *tokens.get(0) != ~")" {
                list.push(box parse(tokens));
            }
            tokens.shift();
            List(list)
        },
        ref x if *x == ~")" => fail!("hit ) token; shouldn't happen"),
        x => Expr::new_atom(x)
    }
}

impl Env {
    pub fn new(params: Option<Vec<~str>>, args: Option<Vec<Expr>>, outer: Option<~Env>) -> Env {
        let mut entries = HashMap::new();
        if params.is_some() && args.is_some() {
            let params = params.expect("params should be a value");
            let args = args.expect("args should be a value");
            if params.len() == args.len() {
                for ctr in range(0,params.len()) {
                    entries.insert(params.get(ctr).to_owned(),
                                   args.get(ctr).clone());
                }
            } else {
                fail!("params and args length doesn't match")
            }
        }
        else if params.is_none() && args.is_none() {
            // do nothing!
        } else {
            fail!("cannot have params & args unset")
        }
        Env { entries: entries, outer: outer }
    }

    pub fn find(&self, symbol: &~str) -> Expr {
        match self.entries.find(symbol) {
            Some(expr) => expr.clone(),
            None => {
                match &self.outer {
                    &Some(ref outer_env) => outer_env.find(symbol),
                    &None => fail!("No variable named {} defined in the environment."
                                  , *symbol)
                }
            }
        }
    }
}

impl Expr {
    pub fn new_atom(input: ~str) -> Expr {
        let first_char = input.char_at(0);
        match first_char {
            '0' | '1' | '2' | '3' | '4' | '5' |
            '6' | '7' | '8' | '9' => {
                match input.contains(".") {
                    true => {
                        match FromStr::from_str(input) {
                            Some(val) => Atom(Float(val)),
                            None => fail!("Cannot parse number-like input of: {}", input)
                        }
                    },
                    false => {
                        match FromStr::from_str(input) {
                            Some(val) => Atom(Integer(val)),
                            None => fail!("Cannot parse number-like input of: {}", input)
                        }
                    }
                }
            },
            _ => Atom(Symbol(input))
        }
    }

    pub fn is_atom(&self) -> bool {
        match *self {
            Atom(_) => true,
            _ => false
        }
    }

    pub fn print(&self) -> ~str {
        match self {
            &List(ref items) => {
                let out = items.iter().map(|i| i.print())
                    .fold(~"(", |m, v| m + v + " ");
                out.trim() + ")"
            },
            &Atom(ref v) => v.print()
        }
    }
    pub fn un_cons(self) -> (Expr, Option<Expr>) {
        match self {
            Atom(_) => fail!("Cannot un_cons an Atom value"),
            List(mut items) => {
                if items.len() == 0 {
                    fail!("cannot build car/cdr of empty list");
                }
                let car = items.shift()
                    .expect("at least one item in the least; shouldn't happen");
                (*car, Some(List(items)))
            }
        }
    }
}

impl AtomVal {
    pub fn print(&self) -> ~str {
        match self {
            &Symbol(ref v) => v.to_owned(),
            &Integer(ref v) => v.to_str(),
            &Float(ref v) => v.to_str()
        }
    }
}

#[cfg(test)]
mod parser_test {
    use super::{pad_input, tokenize, parse, Atom, List, Symbol, Integer, Float};

    #[test]
    fn pad_input_should_insert_spaces_before_and_after_parens() {
        assert_eq!(pad_input(~"(x 1 2 3)"), ~" ( x 1 2 3 ) ");
    }

    #[test]
    fn tokenize_string_into_tokens() {
        let tokens = tokenize(~"(foo 12 3)");
        assert_eq!(tokens.len(), 5);
        assert_eq!(tokens.get(0).to_owned(), ~"(");
        assert_eq!(tokens.get(1).to_owned(), ~"foo");
        assert_eq!(tokens.get(2).to_owned(), ~"12");
        assert_eq!(tokens.get(3).to_owned(), ~"3");
        assert_eq!(tokens.get(4).to_owned(), ~")");
    }

    #[test]
    fn should_parse_tokens_consisting_of_a_str_atom_and_convert_it_to_a_parse_item() {
        let tokens = &mut tokenize(~"bar");
        let parsed_item = parse(tokens);
        println!("{:?}", parsed_item);
        match parsed_item {
            Atom(Symbol(val)) => {
                assert_eq!(val, ~"bar")
            },
            _ => assert!(false)
        }
    }
    #[test]
    fn should_parse_tokens_consisting_of_an_int_atom_and_convert_it_to_a_parse_item() {
        let tokens = &mut tokenize(~"42");
        let parsed_item = parse(tokens);
        println!("{:?}", parsed_item);
        match parsed_item {
            Atom(Integer(val)) => {
                assert_eq!(val, 42)
            },
            _ => assert!(false)
        }
    }
    #[test]
    fn should_parse_tokens_consisting_of_a_float_atom_and_convert_it_to_a_parse_item() {
        let tokens = &mut tokenize(~"1.1");
        let parsed_item = parse(tokens);
        println!("{:?}", parsed_item);
        match parsed_item {
            Atom(Float(val)) => {
                assert_eq!(val, 1.1)
            },
            _ => assert!(false)
        }
    }

    #[test]
    fn should_parse_tokens_consisting_of_a_list_of_atoms_and_parse_it() {
        let tokens = &mut tokenize(~"(bar 12 45)");
        let parsed_item = parse(tokens);
        assert_eq!(parsed_item.is_atom(), false);
        match parsed_item {
            List(items) => {
                let items_len = items.len();
                assert!(3 == items_len);
                assert_eq!(items.get(0).is_atom(), true);
                assert_eq!(*items.get(0), ~Atom(Symbol(~"bar")));
                assert_eq!(items.get(1).is_atom(), true);
                assert_eq!(*items.get(1), ~Atom(Integer(12)));
                assert_eq!(items.get(2).is_atom(), true);
                assert_eq!(*items.get(2), ~Atom(Integer(45)));
            },
            _ => fail!("got back an atom, it seems")
        }
    }
    #[test]
    fn should_parse_tokens_consisting_of_a_list_of_atoms_and_a_nested_list() {
        let tokens = &mut tokenize(~"((2 3) bar 12 (hee (hah)))");
        let parsed_item = parse(tokens);
        assert_eq!(parsed_item.is_atom(), false);
        match parsed_item {
            List(items) => {
                let items_len = items.len();
                assert!(4 == items_len);
                match *items.get(0) {
                    ~List(ref items) => {
                        assert_eq!(items.len(), 2);
                        assert_eq!(*items.get(0), ~Atom(Integer(2)));
                        assert_eq!(*items.get(1), ~Atom(Integer(3)));
                    },
                    _ => fail!("shoulda got a list")
                }
                assert_eq!(*items.get(1), ~Atom(Symbol(~"bar")));
                assert_eq!(*items.get(2), ~Atom(Integer(12)));
                match *items.get(3) {
                    ~List(ref items) => {
                        assert_eq!(items.len(), 2);
                        assert_eq!(*items.get(0), ~Atom(Symbol(~"hee")));
                        match *items.get(1) {
                            ~List(ref items) => {
                                assert_eq!(items.len(), 1);
                                assert_eq!(*items.get(0), ~Atom(Symbol(~"hah")));
                            },
                            _ => fail!("shoulda got a list")
                        }
                    },
                    _ => fail!("shoulda got a list")
                }
            },
            _ => fail!("got back an atom, it seems")
        }
    }
    #[test]
    fn parsed_tokens_of_a_list_of_atoms_and_a_nested_list_should_to_str_correctly() {
        let input = ~"((2 3) bar 12 (hee (hah)))";
        let tokens = &mut tokenize(input.to_owned());
        let parsed_item = parse(tokens);
        let output = parsed_item.print();
        assert_eq!(input, output);
    }
}

#[cfg(test)]
mod eval_test {
    mod un_cons {
        use super::super::{parse_str, Atom, List, Symbol, Integer, Float,
                           Env, eval};
        #[should_fail]
        #[test]
        fn an_atom_expr_returns_the_atom_in_the_car_with_none_in_the_cdr() {
            let expr = parse_str(~"x");
            expr.un_cons();
        }

        #[test]
        fn a_list_with_one_atom_element_returns_an_atom_car_and_an_empty_list_in_the_cdr() {
            let expr = parse_str(~"(x)");
            let (car, cdr) = expr.un_cons();
            assert_eq!(car, Atom(Symbol(~"x")));
            match cdr.expect("should be a value here") {
                List(items) => assert_eq!(items.len(), 0),
                _ => fail!("expected a list")
            }
        }

        #[test]
        fn a_list_with_multiple_elems_puts_the_first_the_car_and_the_rest_in_the_cdr() {
            let expr = parse_str(~"(y 2 3)");
            let (car, cdr) = expr.un_cons();
            assert_eq!(car, Atom(Symbol(~"y")));
            match cdr.expect("should be a value here") {
                List(items) => {
                    assert_eq!(items.len(), 2)
                },
                _ => fail!("expected a list")
            }
        }

        #[test]
        #[should_fail]
        fn calling_un_cons_on_an_empty_list_should_fail() {
            let expr = parse_str(~"()");
            expr.un_cons();
        }

        #[test]
        fn given_a_symbol_in_the_env_then_calling_eval_should_resolve_its_value() {
            let env = Env::new(
                Some(vec!(~"x")),
                Some(vec!(Atom(Integer(42)))),
                None);
            let in_expr = parse_str(~"x");
            let (out_expr, _) = eval(in_expr, env);
            assert_eq!(out_expr, Atom(Integer(42)));
        }

        #[test]
        #[should_fail]
        fn given_a_symbol_NOT_in_the_env_then_calling_eval_should_fail() {
            let env = Env::new(None, None, None);
            let in_expr = parse_str(~"x");
            eval(in_expr, env);
        }

        #[test]
        fn given_a_float_literal_then_calling_eval_should_return_it_back() {
            let env = Env::new(None, None, None);
            let in_expr = parse_str(~"34.3");
            let (out_expr, _) = eval(in_expr, env);
            assert_eq!(out_expr, Atom(Float(34.3)));
        }

        #[test]
        fn given_a_integer_literal_then_calling_eval_should_return_it_back() {
            let env = Env::new(None, None, None);
            let in_expr = parse_str(~"34");
            let (out_expr, _) = eval(in_expr, env);
            assert_eq!(out_expr, Atom(Integer(34)));
        }
    }

    mod env_tests {
        use super::super::{Env};

        #[test]
        fn a_new_env_with_no_parent_params_or_vars_should_be_set_up_empty() {
            let _env = Env::new(None, None, None);
        }
    }
}
<commit_msg>more work on eval, impl'd quote .. reached 20 tests!<commit_after>// Copyright 2014 Jeffery Olson
//
// Licensed under the 3-Clause BSD License, see LICENSE.txt
// at the top-level of this repository.
// This file may not be copied, modified, or distributed
// except according to those terms.
#![crate_id="http://olsonjeffery.github.io#schemers:0.0.1"]
#![desc = "Simple Scheme Interpreter"]
#![license = "3-Clause BSD"]

extern crate collections = "collections";

use std::from_str::FromStr;
use collections::hashmap::HashMap;

// this is starting out as a straight port of Peter Norvig's
// lis.py (the first iteration) to Rust

#[cfg(not(test))]
fn main() {
    use std::os::args;
    let args: Vec<~str> = args().move_iter().skip(1).collect();
    if args.len() == 0 {
        fail!("no program provided");
    }
    let program: ~str = args.move_iter()
        .fold(~"", |memo, arg| memo + arg + " ").trim().to_owned();
    let parsed_items = parse_str(program);
    println!("{}", parsed_items.print());
}

pub struct Env {
    entries: HashMap<~str, Expr>,
    outer: Option<~Env>
}

#[deriving(Eq, Show, Clone)]
pub enum Expr {
    Atom(AtomVal),
    List(Vec<~Expr>)
}

#[deriving(Eq, Show, Clone)]
pub enum AtomVal {
    Symbol(~str),
    Integer(i64),
    Float(f64)
}

fn eval(expr: Expr, env: Env) -> (Expr, Env) {
    match expr {
        // Atom values and values in the Env
        Atom(Symbol(ref var_name)) => (env.find(var_name), env),
        val @ Atom(_) => (val, env),
        list @ List(_) => {
            let (car, cdr) = list.un_cons();
            match car {
                Atom(ref val) if *val == Symbol(~"quote")  => {
                    match cdr {
                        List(mut items) => {
                            (*items.shift().expect("eval: quote list shouldnt be empty"), env)
                        },
                        _ => fail!("eval: expected List in cdr position")
                    }
                },
                _ => fail!("un-implemented case of eval")
            }
        }
    }
}

fn parse_str(input: ~str) -> Expr {
    parse(&mut tokenize(input))
}

fn pad_input(input: ~str) -> ~str {
    input.replace("(", " ( ").replace(")", " ) ")
}

fn tokenize(input: ~str) -> Vec<~str> {
    pad_input(input).split_str(" ")
        .filter(|i| *i != "").map(|i| i.to_owned()).collect()
}

fn parse(tokens: &mut Vec<~str>) -> Expr {
    let current_token =
        tokens.shift().expect("calling parse() w/ empty token list; shouldn't happen");
    match current_token {
        ref x if *x == ~"(" => {
            let mut list = Vec::new();
            while *tokens.get(0) != ~")" {
                list.push(box parse(tokens));
            }
            tokens.shift();
            List(list)
        },
        ref x if *x == ~")" => fail!("hit ) token; shouldn't happen"),
        x => Expr::new_atom(x)
    }
}

impl Env {
    pub fn new(params: Option<Vec<~str>>, args: Option<Vec<Expr>>, outer: Option<~Env>) -> Env {
        let mut entries = HashMap::new();
        if params.is_some() && args.is_some() {
            let params = params.expect("params should be a value");
            let args = args.expect("args should be a value");
            if params.len() == args.len() {
                for ctr in range(0,params.len()) {
                    entries.insert(params.get(ctr).to_owned(),
                                   args.get(ctr).clone());
                }
            } else {
                fail!("params and args length doesn't match")
            }
        }
        else if params.is_none() && args.is_none() {
            // do nothing!
        } else {
            fail!("cannot have params & args unset")
        }
        Env { entries: entries, outer: outer }
    }

    pub fn find(&self, symbol: &~str) -> Expr {
        match self.entries.find(symbol) {
            Some(expr) => expr.clone(),
            None => {
                match &self.outer {
                    &Some(ref outer_env) => outer_env.find(symbol),
                    &None => fail!("No variable named {} defined in the environment."
                                   , *symbol)
                }
            }
        }
    }
}

impl Expr {
    pub fn new_atom(input: ~str) -> Expr {
        let first_char = input.char_at(0);
        match first_char {
            '0' | '1' | '2' | '3' | '4' | '5' |
            '6' | '7' | '8' | '9' => {
                match input.contains(".") {
                    true => {
                        match FromStr::from_str(input) {
                            Some(val) => Atom(Float(val)),
                            None => fail!("Cannot parse number-like input of: {}", input)
                        }
                    },
                    false => {
                        match FromStr::from_str(input) {
                            Some(val) => Atom(Integer(val)),
                            None => fail!("Cannot parse number-like input of: {}", input)
                        }
                    }
                }
            },
            _ => Atom(Symbol(input))
        }
    }

    pub fn is_atom(&self) -> bool {
        match *self {
            Atom(_) => true,
            _ => false
        }
    }

    pub fn print(&self) -> ~str {
        match self {
            &List(ref items) => {
                let out = items.iter().map(|i| i.print())
                    .fold(~"(", |m, v| m + v + " ");
                out.trim() + ")"
            },
            &Atom(ref v) => v.print()
        }
    }
    pub fn un_cons(self) -> (Expr, Expr) {
        match self {
            Atom(_) => fail!("Cannot un_cons an Atom value"),
            List(mut items) => {
                if items.len() == 0 {
                    fail!("cannot build car/cdr of empty list");
                }
                let car = items.shift()
                    .expect("at least one item in the least; shouldn't happen");
                (*car, List(items))
            }
        }
    }
}

impl AtomVal {
    pub fn print(&self) -> ~str {
        match self {
            &Symbol(ref v) => v.to_owned(),
            &Integer(ref v) => v.to_str(),
            &Float(ref v) => v.to_str()
        }
    }
}

#[cfg(test)]
mod parser_test {
    use super::{pad_input, tokenize, parse, Atom, List, Symbol, Integer, Float};

    #[test]
    fn pad_input_should_insert_spaces_before_and_after_parens() {
        assert_eq!(pad_input(~"(x 1 2 3)"), ~" ( x 1 2 3 ) ");
    }

    #[test]
    fn tokenize_string_into_tokens() {
        let tokens = tokenize(~"(foo 12 3)");
        assert_eq!(tokens.len(), 5);
        assert_eq!(tokens.get(0).to_owned(), ~"(");
        assert_eq!(tokens.get(1).to_owned(), ~"foo");
        assert_eq!(tokens.get(2).to_owned(), ~"12");
        assert_eq!(tokens.get(3).to_owned(), ~"3");
        assert_eq!(tokens.get(4).to_owned(), ~")");
    }

    #[test]
    fn should_parse_tokens_consisting_of_a_str_atom_and_convert_it_to_a_parse_item() {
        let tokens = &mut tokenize(~"bar");
        let parsed_item = parse(tokens);
        println!("{:?}", parsed_item);
        match parsed_item {
            Atom(Symbol(val)) => {
                assert_eq!(val, ~"bar")
            },
            _ => assert!(false)
        }
    }
    #[test]
    fn should_parse_tokens_consisting_of_an_int_atom_and_convert_it_to_a_parse_item() {
        let tokens = &mut tokenize(~"42");
        let parsed_item = parse(tokens);
        println!("{:?}", parsed_item);
        match parsed_item {
            Atom(Integer(val)) => {
                assert_eq!(val, 42)
            },
            _ => assert!(false)
        }
    }
    #[test]
    fn should_parse_tokens_consisting_of_a_float_atom_and_convert_it_to_a_parse_item() {
        let tokens = &mut tokenize(~"1.1");
        let parsed_item = parse(tokens);
        println!("{:?}", parsed_item);
        match parsed_item {
            Atom(Float(val)) => {
                assert_eq!(val, 1.1)
            },
            _ => assert!(false)
        }
    }

    #[test]
    fn should_parse_tokens_consisting_of_a_list_of_atoms_and_parse_it() {
        let tokens = &mut tokenize(~"(bar 12 45)");
        let parsed_item = parse(tokens);
        assert_eq!(parsed_item.is_atom(), false);
        match parsed_item {
            List(items) => {
                let items_len = items.len();
                assert!(3 == items_len);
                assert_eq!(items.get(0).is_atom(), true);
                assert_eq!(*items.get(0), ~Atom(Symbol(~"bar")));
                assert_eq!(items.get(1).is_atom(), true);
                assert_eq!(*items.get(1), ~Atom(Integer(12)));
                assert_eq!(items.get(2).is_atom(), true);
                assert_eq!(*items.get(2), ~Atom(Integer(45)));
            },
            _ => fail!("got back an atom, it seems")
        }
    }
    #[test]
    fn should_parse_tokens_consisting_of_a_list_of_atoms_and_a_nested_list() {
        let tokens = &mut tokenize(~"((2 3) bar 12 (hee (hah)))");
        let parsed_item = parse(tokens);
        assert_eq!(parsed_item.is_atom(), false);
        match parsed_item {
            List(items) => {
                let items_len = items.len();
                assert!(4 == items_len);
                match *items.get(0) {
                    ~List(ref items) => {
                        assert_eq!(items.len(), 2);
                        assert_eq!(*items.get(0), ~Atom(Integer(2)));
                        assert_eq!(*items.get(1), ~Atom(Integer(3)));
                    },
                    _ => fail!("shoulda got a list")
                }
                assert_eq!(*items.get(1), ~Atom(Symbol(~"bar")));
                assert_eq!(*items.get(2), ~Atom(Integer(12)));
                match *items.get(3) {
                    ~List(ref items) => {
                        assert_eq!(items.len(), 2);
                        assert_eq!(*items.get(0), ~Atom(Symbol(~"hee")));
                        match *items.get(1) {
                            ~List(ref items) => {
                                assert_eq!(items.len(), 1);
                                assert_eq!(*items.get(0), ~Atom(Symbol(~"hah")));
                            },
                            _ => fail!("shoulda got a list")
                        }
                    },
                    _ => fail!("shoulda got a list")
                }
            },
            _ => fail!("got back an atom, it seems")
        }
    }
    #[test]
    fn parsed_tokens_of_a_list_of_atoms_and_a_nested_list_should_to_str_correctly() {
        let input = ~"((2 3) bar 12 (hee (hah)))";
        let tokens = &mut tokenize(input.to_owned());
        let parsed_item = parse(tokens);
        let output = parsed_item.print();
        assert_eq!(input, output);
    }
}

#[cfg(test)]
mod eval_test {
    use super::{parse_str, Atom, List, Symbol, Integer, Float,
                       Env, eval};
    mod un_cons {
        use super::super::{parse_str, Atom, List, Symbol};
        #[should_fail]
        #[test]
        fn an_atom_expr_returns_the_atom_in_the_car_with_none_in_the_cdr() {
            let expr = parse_str(~"x");
            expr.un_cons();
        }

        #[test]
        fn a_list_with_one_atom_element_returns_an_atom_car_and_an_empty_list_in_the_cdr() {
            let expr = parse_str(~"(x)");
            let (car, cdr) = expr.un_cons();
            assert_eq!(car, Atom(Symbol(~"x")));
            match cdr {
                List(items) => assert_eq!(items.len(), 0),
                _ => fail!("expected a list")
            }
        }

        #[test]
        fn a_list_with_multiple_elems_puts_the_first_the_car_and_the_rest_in_the_cdr() {
            let expr = parse_str(~"(y 2 3)");
            let (car, cdr) = expr.un_cons();
            assert_eq!(car, Atom(Symbol(~"y")));
            match cdr {
                List(items) => {
                    assert_eq!(items.len(), 2)
                },
                _ => fail!("expected a list")
            }
        }

        #[test]
        #[should_fail]
        fn calling_un_cons_on_an_empty_list_should_fail() {
            let expr = parse_str(~"()");
            expr.un_cons();
        }
    }

    #[test]
    fn given_a_symbol_in_the_env_then_calling_eval_should_resolve_its_value() {
        let env = Env::new(
            Some(vec!(~"x")),
            Some(vec!(Atom(Integer(42)))),
            None);
        let in_expr = parse_str(~"x");
        let (out_expr, _) = eval(in_expr, env);
        assert_eq!(out_expr, Atom(Integer(42)));
    }

    #[test]
    #[should_fail]
    fn given_a_symbol_NOT_in_the_env_then_calling_eval_should_fail() {
        let env = Env::new(None, None, None);
        let in_expr = parse_str(~"x");
        eval(in_expr, env);
    }

    #[test]
    fn given_a_float_literal_then_calling_eval_should_return_it_back() {
        let env = Env::new(None, None, None);
        let in_expr = parse_str(~"34.3");
        let (out_expr, _) = eval(in_expr, env);
        assert_eq!(out_expr, Atom(Float(34.3)));
    }

    #[test]
    fn given_a_integer_literal_then_calling_eval_should_return_it_back() {
        let env = Env::new(None, None, None);
        let in_expr = parse_str(~"34");
        let (out_expr, _) = eval(in_expr, env);
        assert_eq!(out_expr, Atom(Integer(34)));
    }

    #[test]
    fn given_a_quote_of_an_atom_expr_should_resolve_to_just_the_atom() {
        let env = Env::new(None, None, None);
        let in_expr = parse_str(~"(quote 42)");
        let (out_expr, _) = eval(in_expr, env);
        assert_eq!(out_expr, Atom(Integer(42)));
    }

    #[test]
    #[should_fail]
    fn a_quote_with_no_params_should_fail() {
        let env = Env::new(None, None, None);
        let in_expr = parse_str(~"(quote)");
        eval(in_expr, env);
    }

    #[test]
    fn given_a_quote_of_a_list_expr_it_should_resolve_to_the_list_without_resolving_symbols() {
        let env = Env::new(None, None, None);
        let in_expr = parse_str(~"(quote (x 37))");
        let (out_expr, _) = eval(in_expr, env);
        match out_expr {
            List(items) => {
                assert_eq!(items.len(), 2);
                assert_eq!(*items.get(0), ~Atom(Symbol(~"x")));
                assert_eq!(*items.get(1), ~Atom(Integer(37)));
            },
            _ => fail!("expected a list")
        }
    }

    mod env_tests {
        use super::super::{Env};

        #[test]
        fn a_new_env_with_no_parent_params_or_vars_should_be_set_up_empty() {
            let _env = Env::new(None, None, None);
        }
    }
}
<|endoftext|>