<commit_before>// Copyright 2016 Google Inc.
//
// 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.

pub mod color;
pub mod data_provider;
pub mod errors;
pub mod iterator;
pub mod math;
pub mod octree;
pub mod read_write;
pub mod s2_cells;

use cgmath::Vector3;
use errors::{ErrorKind, Result};
use std::collections::BTreeMap;
use std::convert::{TryFrom, TryInto};

// Version 9 -> 10: Change in NodeId proto from level (u8) and index (u64) to high (u64) and low
// (u64). We are able to convert the proto on read, so the tools can still read version 9.
// Version 10 -> 11: Change in AxisAlignedCuboid proto from Vector3f min/max to Vector3d min/max.
// We are able to convert the proto on read, so the tools can still read version 9/10.
// Version 11 -> 12: Change in Meta names and structures, to allow both s2 and octree meta.
// We are able to convert the proto on read, so the tools can still read version 9/10/11.
pub const CURRENT_VERSION: i32 = 12;

/// size for batch
pub const NUM_POINTS_PER_BATCH: usize = 500_000;

pub trait NumberOfPoints {
    fn num_points(&self) -> Option<usize>;
}

#[derive(Debug, Clone)]
pub struct Point {
    pub position: cgmath::Vector3<f64>,
    // TODO(sirver): Make color optional, we might not always have it.
    pub color: color::Color<u8>,

    // The intensity of the point if it exists. This value is usually handed through directly by a
    // sensor and has therefore no defined range - or even meaning.
    pub intensity: Option<f32>,
}

pub fn attribute_extension(attribute: &str) -> &str {
    match attribute {
        "position" => "xyz",
        "color" => "rgb",
        _ => attribute,
    }
}

#[derive(Copy, Debug, Clone, PartialEq, Eq)]
pub enum AttributeDataType {
    U8,
    I64,
    U64,
    F32,
    F64,
    U8Vec3,
    F64Vec3,
}

impl AttributeDataType {
    pub fn to_proto(self) -> proto::AttributeDataType {
        match self {
            AttributeDataType::U8 => proto::AttributeDataType::U8,
            AttributeDataType::I64 => proto::AttributeDataType::I64,
            AttributeDataType::U64 => proto::AttributeDataType::U64,
            AttributeDataType::F32 => proto::AttributeDataType::F32,
            AttributeDataType::F64 => proto::AttributeDataType::F64,
            AttributeDataType::U8Vec3 => proto::AttributeDataType::U8Vec3,
            AttributeDataType::F64Vec3 => proto::AttributeDataType::F64Vec3,
        }
    }

    pub fn from_proto(attr_proto: proto::AttributeDataType) -> Result<Self> {
        let attr = match attr_proto {
            proto::AttributeDataType::U8 => AttributeDataType::U8,
            proto::AttributeDataType::I64 => AttributeDataType::I64,
            proto::AttributeDataType::U64 => AttributeDataType::U64,
            proto::AttributeDataType::F32 => AttributeDataType::F32,
            proto::AttributeDataType::F64 => AttributeDataType::F64,
            proto::AttributeDataType::U8Vec3 => AttributeDataType::U8Vec3,
            proto::AttributeDataType::F64Vec3 => AttributeDataType::F64Vec3,
            proto::AttributeDataType::U16
            | proto::AttributeDataType::U32
            | proto::AttributeDataType::I8
            | proto::AttributeDataType::I16
            | proto::AttributeDataType::I32 => {
                return Err(ErrorKind::InvalidInput(
                    "Attribute data type not supported yet".to_string(),
                )
                .into())
            }
            proto::AttributeDataType::INVALID_DATA_TYPE => {
                return Err(
                    ErrorKind::InvalidInput("Attribute data type invalid".to_string()).into(),
                )
            }
        };
        Ok(attr)
    }

    pub fn size_of(self) -> usize {
        match self {
            AttributeDataType::U8 => 1,
            AttributeDataType::F32 => 4,
            AttributeDataType::I64 | AttributeDataType::U64 | AttributeDataType::F64 => 8,
            AttributeDataType::U8Vec3 => 3,
            AttributeDataType::F64Vec3 => 3 * 8,
        }
    }
}

/// General field to describe point feature attributes such as color, intensity, ...
#[derive(Debug, Clone)]
pub enum AttributeData {
    U8(Vec<u8>),
    I64(Vec<i64>),
    U64(Vec<u64>),
    F32(Vec<f32>),
    F64(Vec<f64>),
    U8Vec3(Vec<Vector3<u8>>),
    F64Vec3(Vec<Vector3<f64>>),
}

impl AttributeData {
    pub fn len(&self) -> usize {
        match self {
            AttributeData::U8(data) => data.len(),
            AttributeData::I64(data) => data.len(),
            AttributeData::U64(data) => data.len(),
            AttributeData::F32(data) => data.len(),
            AttributeData::F64(data) => data.len(),
            AttributeData::U8Vec3(data) => data.len(),
            AttributeData::F64Vec3(data) => data.len(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn dim(&self) -> usize {
        match self {
            AttributeData::U8(_)
            | AttributeData::I64(_)
            | AttributeData::U64(_)
            | AttributeData::F32(_)
            | AttributeData::F64(_) => 1,
            AttributeData::U8Vec3(_) | AttributeData::F64Vec3(_) => 3,
        }
    }

    pub fn data_type(&self) -> AttributeDataType {
        match self {
            AttributeData::U8(_) => AttributeDataType::U8,
            AttributeData::I64(_) => AttributeDataType::I64,
            AttributeData::U64(_) => AttributeDataType::U64,
            AttributeData::F32(_) => AttributeDataType::F32,
            AttributeData::F64(_) => AttributeDataType::F64,
            AttributeData::U8Vec3(_) => AttributeDataType::U8Vec3,
            AttributeData::F64Vec3(_) => AttributeDataType::F64Vec3,
        }
    }
}

macro_rules! try_from_impl {
    ($data:ident, $attribute_data_type:ident, $vec_data_type:ty) => {
        match $data {
            AttributeData::$attribute_data_type(data) => Ok(data),
            _ => Err(format!(
                "Attribute data type '{:?}' is incompatible with requested type '{}'.",
                $data.data_type(),
                stringify!($vec_data_type)
            )),
        }
    };
}

macro_rules! try_from_attribute_data {
    ($attribute_data_type:ident, $vec_data_type:ty) => {
        impl TryFrom<AttributeData> for Vec<$vec_data_type> {
            type Error = String;

            fn try_from(data: AttributeData) -> std::result::Result<Self, Self::Error> {
                try_from_impl!(data, $attribute_data_type, $vec_data_type)
            }
        }

        impl<'a> TryFrom<&'a AttributeData> for &'a Vec<$vec_data_type> {
            type Error = String;

            fn try_from(data: &'a AttributeData) -> std::result::Result<Self, Self::Error> {
                try_from_impl!(data, $attribute_data_type, $vec_data_type)
            }
        }

        impl<'a> TryFrom<&'a mut AttributeData> for &'a mut Vec<$vec_data_type> {
            type Error = String;

            fn try_from(data: &'a mut AttributeData) -> std::result::Result<Self, Self::Error> {
                try_from_impl!(data, $attribute_data_type, $vec_data_type)
            }
        }
    };
}
try_from_attribute_data!(U8, u8);
try_from_attribute_data!(I64, i64);
try_from_attribute_data!(U64, u64);
try_from_attribute_data!(F32, f32);
try_from_attribute_data!(F64, f64);
try_from_attribute_data!(U8Vec3, Vector3<u8>);
try_from_attribute_data!(F64Vec3, Vector3<f64>);

fn append_attribute<'a, T: 'a>(
    name: &str,
    to: &'a mut PointsBatch,
    from: &'a mut PointsBatch,
) -> std::result::Result<(), String>
where
    &'a mut Vec<T>: TryFrom<&'a mut AttributeData, Error = String>,
{
    let to_vec: &'a mut Vec<T> = to.get_attribute_vec_mut(name)?;
    let from_vec: &'a mut Vec<T> = from.get_attribute_vec_mut(name)?;
    to_vec.append(from_vec);
    Ok(())
}

/// General structure that contains points and attached feature attributes.
#[derive(Debug, Clone)]
pub struct PointsBatch {
    pub position: Vec<Vector3<f64>>,
    // BTreeMap for deterministic iteration order.
    pub attributes: BTreeMap<String, AttributeData>,
}

impl PointsBatch {
    pub fn append(&mut self, other: &mut PointsBatch) -> std::result::Result<(), String> {
        if self.position.is_empty() {
            *self = other.split_off(0);
        } else {
            assert_eq!(self.attributes.len(), other.attributes.len());
            self.position.append(&mut other.position);
            let data_types: Vec<(String, AttributeDataType)> = self
                .attributes
                .iter()
                .map(|(name, attrib)| (name.clone(), attrib.data_type()))
                .collect();
            for (name, dtype) in &data_types {
                use AttributeDataType::*;
                match dtype {
                    U8 => append_attribute::<u8>(name, self, other)?,
                    U64 => append_attribute::<u64>(name, self, other)?,
                    I64 => append_attribute::<i64>(name, self, other)?,
                    F32 => append_attribute::<f32>(name, self, other)?,
                    F64 => append_attribute::<f64>(name, self, other)?,
                    U8Vec3 => append_attribute::<Vector3<u8>>(name, self, other)?,
                    F64Vec3 => append_attribute::<Vector3<f64>>(name, self, other)?,
                };
            }
        }
        Ok(())
    }

    pub fn split_off(&mut self, at: usize) -> Self {
        let mut res = Self {
            position: self.position.split_off(at),
            attributes: BTreeMap::new(),
        };
        for (name, attribute) in self.attributes.iter_mut() {
            let name = name.clone();
            use AttributeData::*;
            match attribute {
                U8(data) => res.attributes.insert(name, U8(data.split_off(at))),
                U64(data) => res.attributes.insert(name, U64(data.split_off(at))),
                I64(data) => res.attributes.insert(name, I64(data.split_off(at))),
                F32(data) => res.attributes.insert(name, F32(data.split_off(at))),
                F64(data) => res.attributes.insert(name, F64(data.split_off(at))),
                U8Vec3(data) => res.attributes.insert(name, U8Vec3(data.split_off(at))),
                F64Vec3(data) => res.attributes.insert(name, F64Vec3(data.split_off(at))),
            };
        }
        res
    }

    pub fn retain(&mut self, keep: &[bool]) {
        assert_eq!(self.position.len(), keep.len());
        let mut keep = keep.iter().copied().cycle();
        self.position.retain(|_| keep.next().unwrap());
        for a in self.attributes.values_mut() {
            match a {
                AttributeData::U8(data) => data.retain(|_| keep.next().unwrap()),
                AttributeData::U64(data) => data.retain(|_| keep.next().unwrap()),
                AttributeData::I64(data) => data.retain(|_| keep.next().unwrap()),
                AttributeData::F32(data) => data.retain(|_| keep.next().unwrap()),
                AttributeData::F64(data) => data.retain(|_| keep.next().unwrap()),
                AttributeData::U8Vec3(data) => data.retain(|_| keep.next().unwrap()),
                AttributeData::F64Vec3(data) => data.retain(|_| keep.next().unwrap()),
            }
        }
    }

    pub fn get_attribute_vec<'a, T>(
        &'a self,
        key: impl AsRef<str>,
    ) -> std::result::Result<&'a Vec<T>, String>
    where
        &'a Vec<T>: TryFrom<&'a AttributeData, Error = String>,
    {
        self.attributes
            .get(key.as_ref())
            .ok_or_else(|| format!("Attribute '{}' not found.", key.as_ref()))
            .and_then(|val| val.try_into())
    }

    pub fn get_attribute_vec_mut<'a, T>(
        &'a mut self,
        key: impl AsRef<str>,
    ) -> std::result::Result<&'a mut Vec<T>, String>
    where
        &'a mut Vec<T>: TryFrom<&'a mut AttributeData, Error = String>,
    {
        self.attributes
            .get_mut(key.as_ref())
            .ok_or_else(|| format!("Attribute '{}' not found.", key.as_ref()))
            .and_then(|val| val.try_into())
    }

    pub fn remove_attribute_vec<T>(
        &mut self,
        key: impl AsRef<str>,
    ) -> std::result::Result<Vec<T>, String>
    where
        Vec<T>: TryFrom<AttributeData, Error = String>,
    {
        self.attributes
            .remove(key.as_ref())
            .ok_or_else(|| format!("Attribute '{}' not found.", key.as_ref()))
            .and_then(|val| val.try_into())
    }
}

pub use point_viewer_proto_rust::proto;
<commit_msg>Simplified append and split_off. (#358)<commit_after>// Copyright 2016 Google Inc.
//
// 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.

pub mod color;
pub mod data_provider;
pub mod errors;
pub mod iterator;
pub mod math;
pub mod octree;
pub mod read_write;
pub mod s2_cells;

use cgmath::Vector3;
use errors::{ErrorKind, Result};
use std::collections::BTreeMap;
use std::convert::{TryFrom, TryInto};

// Version 9 -> 10: Change in NodeId proto from level (u8) and index (u64) to high (u64) and low
// (u64). We are able to convert the proto on read, so the tools can still read version 9.
// Version 10 -> 11: Change in AxisAlignedCuboid proto from Vector3f min/max to Vector3d min/max.
// We are able to convert the proto on read, so the tools can still read version 9/10.
// Version 11 -> 12: Change in Meta names and structures, to allow both s2 and octree meta.
// We are able to convert the proto on read, so the tools can still read version 9/10/11.
pub const CURRENT_VERSION: i32 = 12;

/// size for batch
pub const NUM_POINTS_PER_BATCH: usize = 500_000;

pub trait NumberOfPoints {
    fn num_points(&self) -> Option<usize>;
}

#[derive(Debug, Clone)]
pub struct Point {
    pub position: cgmath::Vector3<f64>,
    // TODO(sirver): Make color optional, we might not always have it.
    pub color: color::Color<u8>,

    // The intensity of the point if it exists. This value is usually handed through directly by a
    // sensor and has therefore no defined range - or even meaning.
    pub intensity: Option<f32>,
}

pub fn attribute_extension(attribute: &str) -> &str {
    match attribute {
        "position" => "xyz",
        "color" => "rgb",
        _ => attribute,
    }
}

#[derive(Copy, Debug, Clone, PartialEq, Eq)]
pub enum AttributeDataType {
    U8,
    I64,
    U64,
    F32,
    F64,
    U8Vec3,
    F64Vec3,
}

impl AttributeDataType {
    pub fn to_proto(self) -> proto::AttributeDataType {
        match self {
            AttributeDataType::U8 => proto::AttributeDataType::U8,
            AttributeDataType::I64 => proto::AttributeDataType::I64,
            AttributeDataType::U64 => proto::AttributeDataType::U64,
            AttributeDataType::F32 => proto::AttributeDataType::F32,
            AttributeDataType::F64 => proto::AttributeDataType::F64,
            AttributeDataType::U8Vec3 => proto::AttributeDataType::U8Vec3,
            AttributeDataType::F64Vec3 => proto::AttributeDataType::F64Vec3,
        }
    }

    pub fn from_proto(attr_proto: proto::AttributeDataType) -> Result<Self> {
        let attr = match attr_proto {
            proto::AttributeDataType::U8 => AttributeDataType::U8,
            proto::AttributeDataType::I64 => AttributeDataType::I64,
            proto::AttributeDataType::U64 => AttributeDataType::U64,
            proto::AttributeDataType::F32 => AttributeDataType::F32,
            proto::AttributeDataType::F64 => AttributeDataType::F64,
            proto::AttributeDataType::U8Vec3 => AttributeDataType::U8Vec3,
            proto::AttributeDataType::F64Vec3 => AttributeDataType::F64Vec3,
            proto::AttributeDataType::U16
            | proto::AttributeDataType::U32
            | proto::AttributeDataType::I8
            | proto::AttributeDataType::I16
            | proto::AttributeDataType::I32 => {
                return Err(ErrorKind::InvalidInput(
                    "Attribute data type not supported yet".to_string(),
                )
                .into())
            }
            proto::AttributeDataType::INVALID_DATA_TYPE => {
                return Err(
                    ErrorKind::InvalidInput("Attribute data type invalid".to_string()).into(),
                )
            }
        };
        Ok(attr)
    }

    pub fn size_of(self) -> usize {
        match self {
            AttributeDataType::U8 => 1,
            AttributeDataType::F32 => 4,
            AttributeDataType::I64 | AttributeDataType::U64 | AttributeDataType::F64 => 8,
            AttributeDataType::U8Vec3 => 3,
            AttributeDataType::F64Vec3 => 3 * 8,
        }
    }
}

/// General field to describe point feature attributes such as color, intensity, ...
#[derive(Debug, Clone)]
pub enum AttributeData {
    U8(Vec<u8>),
    I64(Vec<i64>),
    U64(Vec<u64>),
    F32(Vec<f32>),
    F64(Vec<f64>),
    U8Vec3(Vec<Vector3<u8>>),
    F64Vec3(Vec<Vector3<f64>>),
}

impl AttributeData {
    pub fn len(&self) -> usize {
        match self {
            AttributeData::U8(data) => data.len(),
            AttributeData::I64(data) => data.len(),
            AttributeData::U64(data) => data.len(),
            AttributeData::F32(data) => data.len(),
            AttributeData::F64(data) => data.len(),
            AttributeData::U8Vec3(data) => data.len(),
            AttributeData::F64Vec3(data) => data.len(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn dim(&self) -> usize {
        match self {
            AttributeData::U8(_)
            | AttributeData::I64(_)
            | AttributeData::U64(_)
            | AttributeData::F32(_)
            | AttributeData::F64(_) => 1,
            AttributeData::U8Vec3(_) | AttributeData::F64Vec3(_) => 3,
        }
    }

    pub fn data_type(&self) -> AttributeDataType {
        match self {
            AttributeData::U8(_) => AttributeDataType::U8,
            AttributeData::I64(_) => AttributeDataType::I64,
            AttributeData::U64(_) => AttributeDataType::U64,
            AttributeData::F32(_) => AttributeDataType::F32,
            AttributeData::F64(_) => AttributeDataType::F64,
            AttributeData::U8Vec3(_) => AttributeDataType::U8Vec3,
            AttributeData::F64Vec3(_) => AttributeDataType::F64Vec3,
        }
    }

    pub fn append(&mut self, other: &mut Self) -> std::result::Result<(), String> {
        match (self, other) {
            (AttributeData::U8(s), AttributeData::U8(o)) => s.append(o),
            (AttributeData::U64(s), AttributeData::U64(o)) => s.append(o),
            (AttributeData::I64(s), AttributeData::I64(o)) => s.append(o),
            (AttributeData::F32(s), AttributeData::F32(o)) => s.append(o),
            (AttributeData::F64(s), AttributeData::F64(o)) => s.append(o),
            (AttributeData::U8Vec3(s), AttributeData::U8Vec3(o)) => s.append(o),
            (AttributeData::F64Vec3(s), AttributeData::F64Vec3(o)) => s.append(o),
            (s, o) => {
                return Err(format!(
                    "Own data type '{:?}' is incompatible with other type '{:?}'.",
                    s.data_type(),
                    o.data_type(),
                ))
            }
        };
        Ok(())
    }

    pub fn split_off(&mut self, at: usize) -> Self {
        match self {
            AttributeData::U8(data) => AttributeData::U8(data.split_off(at)),
            AttributeData::U64(data) => AttributeData::U64(data.split_off(at)),
            AttributeData::I64(data) => AttributeData::I64(data.split_off(at)),
            AttributeData::F32(data) => AttributeData::F32(data.split_off(at)),
            AttributeData::F64(data) => AttributeData::F64(data.split_off(at)),
            AttributeData::U8Vec3(data) => AttributeData::U8Vec3(data.split_off(at)),
            AttributeData::F64Vec3(data) => AttributeData::F64Vec3(data.split_off(at)),
        }
    }
}

macro_rules! try_from_impl {
    ($data:ident, $attribute_data_type:ident, $vec_data_type:ty) => {
        match $data {
            AttributeData::$attribute_data_type(data) => Ok(data),
            _ => Err(format!(
                "Attribute data type '{:?}' is incompatible with requested type '{}'.",
                $data.data_type(),
                stringify!($vec_data_type)
            )),
        }
    };
}

macro_rules! try_from_attribute_data {
    ($attribute_data_type:ident, $vec_data_type:ty) => {
        impl TryFrom<AttributeData> for Vec<$vec_data_type> {
            type Error = String;

            fn try_from(data: AttributeData) -> std::result::Result<Self, Self::Error> {
                try_from_impl!(data, $attribute_data_type, $vec_data_type)
            }
        }

        impl<'a> TryFrom<&'a AttributeData> for &'a Vec<$vec_data_type> {
            type Error = String;

            fn try_from(data: &'a AttributeData) -> std::result::Result<Self, Self::Error> {
                try_from_impl!(data, $attribute_data_type, $vec_data_type)
            }
        }

        impl<'a> TryFrom<&'a mut AttributeData> for &'a mut Vec<$vec_data_type> {
            type Error = String;

            fn try_from(data: &'a mut AttributeData) -> std::result::Result<Self, Self::Error> {
                try_from_impl!(data, $attribute_data_type, $vec_data_type)
            }
        }
    };
}
try_from_attribute_data!(U8, u8);
try_from_attribute_data!(I64, i64);
try_from_attribute_data!(U64, u64);
try_from_attribute_data!(F32, f32);
try_from_attribute_data!(F64, f64);
try_from_attribute_data!(U8Vec3, Vector3<u8>);
try_from_attribute_data!(F64Vec3, Vector3<f64>);

/// General structure that contains points and attached feature attributes.
#[derive(Debug, Clone)]
pub struct PointsBatch {
    pub position: Vec<Vector3<f64>>,
    // BTreeMap for deterministic iteration order.
    pub attributes: BTreeMap<String, AttributeData>,
}

impl PointsBatch {
    pub fn append(&mut self, other: &mut PointsBatch) -> std::result::Result<(), String> {
        if self.position.is_empty() {
            *self = other.split_off(0);
        } else {
            assert_eq!(self.attributes.len(), other.attributes.len());
            self.position.append(&mut other.position);
            for (s, o) in self
                .attributes
                .values_mut()
                .zip(other.attributes.values_mut())
            {
                s.append(o)?;
            }
        }
        Ok(())
    }

    pub fn split_off(&mut self, at: usize) -> Self {
        let position = self.position.split_off(at);
        let attributes = self
            .attributes
            .iter_mut()
            .map(|(n, a)| (n.clone(), a.split_off(at)))
            .collect();
        Self {
            position,
            attributes,
        }
    }

    pub fn retain(&mut self, keep: &[bool]) {
        assert_eq!(self.position.len(), keep.len());
        let mut keep = keep.iter().copied().cycle();
        self.position.retain(|_| keep.next().unwrap());
        for a in self.attributes.values_mut() {
            match a {
                AttributeData::U8(data) => data.retain(|_| keep.next().unwrap()),
                AttributeData::U64(data) => data.retain(|_| keep.next().unwrap()),
                AttributeData::I64(data) => data.retain(|_| keep.next().unwrap()),
                AttributeData::F32(data) => data.retain(|_| keep.next().unwrap()),
                AttributeData::F64(data) => data.retain(|_| keep.next().unwrap()),
                AttributeData::U8Vec3(data) => data.retain(|_| keep.next().unwrap()),
                AttributeData::F64Vec3(data) => data.retain(|_| keep.next().unwrap()),
            }
        }
    }

    pub fn get_attribute_vec<'a, T>(
        &'a self,
        key: impl AsRef<str>,
    ) -> std::result::Result<&'a Vec<T>, String>
    where
        &'a Vec<T>: TryFrom<&'a AttributeData, Error = String>,
    {
        self.attributes
            .get(key.as_ref())
            .ok_or_else(|| format!("Attribute '{}' not found.", key.as_ref()))
            .and_then(|val| val.try_into())
    }

    pub fn get_attribute_vec_mut<'a, T>(
        &'a mut self,
        key: impl AsRef<str>,
    ) -> std::result::Result<&'a mut Vec<T>, String>
    where
        &'a mut Vec<T>: TryFrom<&'a mut AttributeData, Error = String>,
    {
        self.attributes
            .get_mut(key.as_ref())
            .ok_or_else(|| format!("Attribute '{}' not found.", key.as_ref()))
            .and_then(|val| val.try_into())
    }

    pub fn remove_attribute_vec<T>(
        &mut self,
        key: impl AsRef<str>,
    ) -> std::result::Result<Vec<T>, String>
    where
        Vec<T>: TryFrom<AttributeData, Error = String>,
    {
        self.attributes
            .remove(key.as_ref())
            .ok_or_else(|| format!("Attribute '{}' not found.", key.as_ref()))
            .and_then(|val| val.try_into())
    }
}

pub use point_viewer_proto_rust::proto;
<|endoftext|>


<commit_before>#![feature(collections, unboxed_closures, core)]
extern crate itertools;

use std::collections::{BTreeMap, HashMap, HashSet};
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use std::iter;
use std::ops;

use itertools::{GroupBy, Itertools};

/// A Document contains an id and content.
/// Hashing and equality are based only on the id field.
#[derive(Clone, Eq, Debug)]
pub struct Document {
    id: String,
    content: String,
}

impl Drop for Document {
    fn drop(&mut self) {
        println!("Dropped doc {}", self.id);
    }
}

impl Document {
    /// Construct a new Document from an id and content.
    /// Both two arguments can be anything that can be turned into a String.
    pub fn new<S, T>(id: S, content: T) -> Document
        where S: Into<String>,
              T: Into<String>
    {
        Document { id: id.into(), content: content.into(), }
    }

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

    pub fn content(&self) -> &str {
        &self.content
    }
}

impl Hash for Document {
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        state.write(self.id.as_bytes());
    }
}

impl PartialEq<Document> for Document {
    fn eq(&self, other: &Document) -> bool {
        self.id == other.id
    }
}

#[derive(Clone, Eq, PartialEq, Hash, Debug)]
struct IndexedDocument {
    pub doc: Arc<Document>,
    pub highlighted: (usize, usize),
}

fn ngrams((_, chars): (bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams> {
    (1..chars.len() + 1).map(Ngrams::new(chars))
}

fn not_whitespace(&(is_whitespace, _): &(bool, Vec<(usize, char)>)) -> bool {
    !is_whitespace
}

fn is_whitespace(&(_, c): &(usize, char)) -> bool {
    c.is_whitespace()
}

struct Ngrams {
    chars: Vec<(usize, char)>
}

impl Ngrams {
    fn new(chars: Vec<(usize, char)>) -> Ngrams {
        Ngrams { chars: chars }
    }
}

impl Fn<(usize,)> for Ngrams {
    extern "rust-call" fn call(&self, (to,): (usize,)) -> (String, (usize, usize)) {
        let word = self.chars[..to].iter().flat_map(|&(_, c)| c.to_lowercase()).collect();
        let start = self.chars[0].0;
        let (last_idx, last_char) = self.chars[to - 1];
        let finish = last_idx + last_char.len_utf8();
        (word, (start, finish))
    }
}

impl FnMut<(usize,)> for Ngrams {
    extern "rust-call" fn call_mut(&mut self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

impl FnOnce<(usize,)> for Ngrams {
    type Output = (String, (usize, usize));
    extern "rust-call" fn call_once(self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

fn analyze_doc(doc: &str)
-> iter::FlatMap<
    iter::Filter<
        GroupBy<bool, std::str::CharIndices, fn(&(usize, char)) -> bool>,
        fn(&(bool, Vec<(usize, char)>)) -> bool>,
    iter::Map<ops::Range<usize>, Ngrams>,
    fn((bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams>>
{
    doc.char_indices()
        .group_by(is_whitespace as fn(&(usize, char)) -> bool)
        .filter(not_whitespace as fn(&(bool, Vec<(usize, char)>)) -> bool)
        .flat_map(ngrams)
}

/// A basic implementation of an `Index`, the inverted index is a data structure that maps
/// from words to sets of Documents.
pub struct InvertedIndex {
    index: BTreeMap<String, HashSet<IndexedDocument>>,
    docs: BTreeMap<String, Arc<Document>>,
}

impl InvertedIndex {
    pub fn new() -> InvertedIndex {
        InvertedIndex {
            index: BTreeMap::new(),
            docs: BTreeMap::new(),
        }
    }

    /// A basic implementation of index, splits the document's content into whitespace-separated
    /// words, and inserts each word-document pair into the map.
    pub fn index(&mut self, doc: Document) {
        let doc = Arc::new(doc);
        let analyzed  = analyze_doc(doc.content());
        let previous_version = self.docs.insert(doc.id.clone(), doc.clone());
        if let Some(previous_version) = previous_version {
            let previous_analyzed = analyze_doc(previous_version.content());
            for (ngram, highlights) in previous_analyzed {
                let is_empty = {
                    let docs_for_ngram = self.index.get_mut(&ngram).unwrap();
                    docs_for_ngram.remove(&IndexedDocument { doc: doc.clone(), highlighted: highlights });
                    docs_for_ngram.is_empty()
                };
                if is_empty {
                    self.index.remove(&ngram);
                }
            }
        }

        for (ngram, highlighted) in analyzed {
            self.index.entry(ngram).or_insert_with(|| HashSet::new())
                .insert(IndexedDocument { doc: doc.clone(), highlighted: highlighted });
        }
    }

    /// A basic search implementation that splits the query's content into whitespace-separated
    /// words, looks up the set of Documents for each word, and then concatenates the sets.
    pub fn search(&self, query: &str) -> HashSet<SearchResult> {
        let map = query.split_whitespace()
            .flat_map(|word| self.index.get(&word.to_lowercase()))
            .flat_map(|docs| docs)
            .cloned()
            .fold(HashMap::new(), |mut map, search_result| {
                map.entry(search_result.doc)
                    .or_insert(Vec::new())
                    .push(search_result.highlighted);
                map
            });
        map.into_iter()
            .map(|(doc, mut highlights)| {
                highlights.sort();
                SearchResult { doc: doc, highlights: highlights }
            })
        .collect()
    }
}

#[derive(Clone, Eq, PartialEq, Hash, Debug)]
pub struct SearchResult {
    doc: Arc<Document>,
    highlights: Vec<(usize, usize)>,
}

impl SearchResult {
    #[cfg(test)]
    fn new(doc: Arc<Document>, highlights: Vec<(usize, usize)>) -> SearchResult {
        SearchResult {
            doc: doc,
            highlights: highlights,
        }
    }

    pub fn doc(&self) -> &Arc<Document> {
        &self.doc
    }

    pub fn highlights(&self) -> &Vec<(usize, usize)> {
        &self.highlights
    }

    pub fn highlighted_content(&self) -> String {
        let &SearchResult {ref doc, ref highlights} = self;
        let content = doc.content();
        let mut begin_idx = 0;
        let mut parts = vec![];
        for &(begin, end) in highlights {
            parts.push(&content[begin_idx..begin]);
            parts.push("");
            parts.push(&content[begin..end]);
            parts.push("");
            begin_idx = end;
        }
        parts.push(&content[begin_idx..]);
        parts.into_iter().join("")
    }
}

#[test]
fn test_search1() {
    let mut index = InvertedIndex::new();
    let doc1 = Document::new("1", "learn to program in rust today");
    index.index(doc1.clone());
    let doc2 = Document::new("2", "what did you today do");
    index.index(doc2.clone());
    let search_results = index.search("today");
    let expected = [
        SearchResult::new(Arc::new(doc1), vec![(25, 30)]),
        SearchResult::new(Arc::new(doc2), vec![(13, 18)])
    ];
    assert_eq!(search_results, expected.iter().cloned().collect());
    assert_eq!("learn to program in rust today", expected[0].highlighted_content());
}

#[test]
fn test_ngrams() {
    let mut index = InvertedIndex::new();
    let doc1 = Document::new("1", "learn to program in rust today");
    let doc2 = Document::new("2", "what did you today do");
    index.index(doc1.clone());
    index.index(doc2.clone());
    let search_results = index.search("to");
    let expected = [
        SearchResult::new(Arc::new(doc1), vec![(6, 8), (25, 27)]),
        SearchResult::new(Arc::new(doc2), vec![(13, 15)]),
    ];
    assert_eq!(search_results, expected.iter().cloned().collect());
    assert_eq!("learn to program in rust today", expected[0].highlighted_content());

}

#[test]
fn test_search2() {
    let mut index = InvertedIndex::new();
    let doc1 = Document::new("2", "what to do today");
    let doc2 = Document::new("3", "hey today");
    index.index(doc1.clone());
    index.index(doc2.clone());
    let search_results = index.search("to");
    let expected = [
        SearchResult::new(Arc::new(doc1), vec![(5, 7), (11, 13)]),
        SearchResult::new(Arc::new(doc2), vec![(4, 6)]),
    ];
    assert_eq!(search_results, expected.iter().cloned().collect());
    assert_eq!(expected[0].highlighted_content(), "what to do today");
}

#[test]
fn test_unicode() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "abc åäö");
    index.index(doc.clone());
    let to_search = "å";
    let search_results = index.search(to_search);
    let &SearchResult { ref doc, ref highlights } = search_results.iter().next().unwrap();
    let (begin, end) = highlights[0];
    assert_eq!(&doc.content()[begin..end], to_search);
}
<commit_msg>derive Debug for InvertedIndex, and add a test for document updates<commit_after>#![feature(collections, unboxed_closures, core)]
extern crate itertools;

use std::collections::{BTreeMap, HashMap, HashSet};
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use std::iter;
use std::ops;

use itertools::{GroupBy, Itertools};

/// A Document contains an id and content.
/// Hashing and equality are based only on the id field.
#[derive(Clone, Eq, Debug)]
pub struct Document {
    id: String,
    content: String,
}

impl Drop for Document {
    fn drop(&mut self) {
        println!("Dropped doc {}", self.id);
    }
}

impl Document {
    /// Construct a new Document from an id and content.
    /// Both two arguments can be anything that can be turned into a String.
    pub fn new<S, T>(id: S, content: T) -> Document
        where S: Into<String>,
              T: Into<String>
    {
        Document { id: id.into(), content: content.into(), }
    }

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

    pub fn content(&self) -> &str {
        &self.content
    }
}

impl Hash for Document {
    fn hash<H>(&self, state: &mut H) where H: Hasher {
        state.write(self.id.as_bytes());
    }
}

impl PartialEq<Document> for Document {
    fn eq(&self, other: &Document) -> bool {
        self.id == other.id
    }
}

#[derive(Clone, Eq, PartialEq, Hash, Debug)]
struct IndexedDocument {
    pub doc: Arc<Document>,
    pub highlighted: (usize, usize),
}

fn ngrams((_, chars): (bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams> {
    (1..chars.len() + 1).map(Ngrams::new(chars))
}

fn not_whitespace(&(is_whitespace, _): &(bool, Vec<(usize, char)>)) -> bool {
    !is_whitespace
}

fn is_whitespace(&(_, c): &(usize, char)) -> bool {
    c.is_whitespace()
}

struct Ngrams {
    chars: Vec<(usize, char)>
}

impl Ngrams {
    fn new(chars: Vec<(usize, char)>) -> Ngrams {
        Ngrams { chars: chars }
    }
}

impl Fn<(usize,)> for Ngrams {
    extern "rust-call" fn call(&self, (to,): (usize,)) -> (String, (usize, usize)) {
        let word = self.chars[..to].iter().flat_map(|&(_, c)| c.to_lowercase()).collect();
        let start = self.chars[0].0;
        let (last_idx, last_char) = self.chars[to - 1];
        let finish = last_idx + last_char.len_utf8();
        (word, (start, finish))
    }
}

impl FnMut<(usize,)> for Ngrams {
    extern "rust-call" fn call_mut(&mut self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

impl FnOnce<(usize,)> for Ngrams {
    type Output = (String, (usize, usize));
    extern "rust-call" fn call_once(self, to: (usize,)) -> (String, (usize, usize)) {
        self.call(to)
    }
}

fn analyze_doc(doc: &str)
-> iter::FlatMap<
    iter::Filter<
        GroupBy<bool, std::str::CharIndices, fn(&(usize, char)) -> bool>,
        fn(&(bool, Vec<(usize, char)>)) -> bool>,
    iter::Map<ops::Range<usize>, Ngrams>,
    fn((bool, Vec<(usize, char)>)) -> iter::Map<ops::Range<usize>, Ngrams>>
{
    doc.char_indices()
        .group_by(is_whitespace as fn(&(usize, char)) -> bool)
        .filter(not_whitespace as fn(&(bool, Vec<(usize, char)>)) -> bool)
        .flat_map(ngrams)
}

/// A basic implementation of an `Index`, the inverted index is a data structure that maps
/// from words to sets of Documents.
#[derive(Debug)]
pub struct InvertedIndex {
    index: BTreeMap<String, HashSet<IndexedDocument>>,
    docs: BTreeMap<String, Arc<Document>>,
}

impl InvertedIndex {
    pub fn new() -> InvertedIndex {
        InvertedIndex {
            index: BTreeMap::new(),
            docs: BTreeMap::new(),
        }
    }

    /// A basic implementation of index, splits the document's content into whitespace-separated
    /// words, and inserts each word-document pair into the map.
    pub fn index(&mut self, doc: Document) {
        let doc = Arc::new(doc);
        let analyzed  = analyze_doc(doc.content());
        let previous_version = self.docs.insert(doc.id.clone(), doc.clone());
        if let Some(previous_version) = previous_version {
            let previous_analyzed = analyze_doc(previous_version.content());
            for (ngram, highlights) in previous_analyzed {
                let is_empty = {
                    let docs_for_ngram = self.index.get_mut(&ngram).unwrap();
                    docs_for_ngram.remove(&IndexedDocument { doc: doc.clone(), highlighted: highlights });
                    docs_for_ngram.is_empty()
                };
                if is_empty {
                    self.index.remove(&ngram);
                }
            }
        }

        for (ngram, highlighted) in analyzed {
            self.index.entry(ngram).or_insert_with(|| HashSet::new())
                .insert(IndexedDocument { doc: doc.clone(), highlighted: highlighted });
        }
    }

    /// A basic search implementation that splits the query's content into whitespace-separated
    /// words, looks up the set of Documents for each word, and then concatenates the sets.
    pub fn search(&self, query: &str) -> HashSet<SearchResult> {
        let map = query.split_whitespace()
            .flat_map(|word| self.index.get(&word.to_lowercase()))
            .flat_map(|docs| docs)
            .cloned()
            .fold(HashMap::new(), |mut map, search_result| {
                map.entry(search_result.doc)
                    .or_insert(Vec::new())
                    .push(search_result.highlighted);
                map
            });
        map.into_iter()
            .map(|(doc, mut highlights)| {
                highlights.sort();
                SearchResult { doc: doc, highlights: highlights }
            })
        .collect()
    }
}

#[derive(Clone, Eq, PartialEq, Hash, Debug)]
pub struct SearchResult {
    doc: Arc<Document>,
    highlights: Vec<(usize, usize)>,
}

impl SearchResult {
    #[cfg(test)]
    fn new(doc: Arc<Document>, highlights: Vec<(usize, usize)>) -> SearchResult {
        SearchResult {
            doc: doc,
            highlights: highlights,
        }
    }

    pub fn doc(&self) -> &Arc<Document> {
        &self.doc
    }

    pub fn highlights(&self) -> &Vec<(usize, usize)> {
        &self.highlights
    }

    pub fn highlighted_content(&self) -> String {
        let &SearchResult {ref doc, ref highlights} = self;
        let content = doc.content();
        let mut begin_idx = 0;
        let mut parts = vec![];
        for &(begin, end) in highlights {
            parts.push(&content[begin_idx..begin]);
            parts.push("");
            parts.push(&content[begin..end]);
            parts.push("");
            begin_idx = end;
        }
        parts.push(&content[begin_idx..]);
        parts.into_iter().join("")
    }
}

#[test]
fn test_search1() {
    let mut index = InvertedIndex::new();
    let doc1 = Document::new("1", "learn to program in rust today");
    index.index(doc1.clone());
    let doc2 = Document::new("2", "what did you today do");
    index.index(doc2.clone());
    let search_results = index.search("today");
    let expected = [
        SearchResult::new(Arc::new(doc1), vec![(25, 30)]),
        SearchResult::new(Arc::new(doc2), vec![(13, 18)])
    ];
    assert_eq!(search_results, expected.iter().cloned().collect());
    assert_eq!("learn to program in rust today", expected[0].highlighted_content());
}

#[test]
fn test_ngrams() {
    let mut index = InvertedIndex::new();
    let doc1 = Document::new("1", "learn to program in rust today");
    let doc2 = Document::new("2", "what did you today do");
    index.index(doc1.clone());
    index.index(doc2.clone());
    let search_results = index.search("to");
    let expected = [
        SearchResult::new(Arc::new(doc1), vec![(6, 8), (25, 27)]),
        SearchResult::new(Arc::new(doc2), vec![(13, 15)]),
    ];
    assert_eq!(search_results, expected.iter().cloned().collect());
    assert_eq!("learn to program in rust today", expected[0].highlighted_content());

}

#[test]
fn test_search2() {
    let mut index = InvertedIndex::new();
    let doc1 = Document::new("2", "what to do today");
    let doc2 = Document::new("3", "hey today");
    index.index(doc1.clone());
    index.index(doc2.clone());
    let search_results = index.search("to");
    let expected = [
        SearchResult::new(Arc::new(doc1), vec![(5, 7), (11, 13)]),
        SearchResult::new(Arc::new(doc2), vec![(4, 6)]),
    ];
    assert_eq!(search_results, expected.iter().cloned().collect());
    assert_eq!(expected[0].highlighted_content(), "what to do today");
}

#[test]
fn test_unicode() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "abc åäö");
    index.index(doc.clone());
    let to_search = "å";
    let search_results = index.search(to_search);
    let &SearchResult { ref doc, ref highlights } = search_results.iter().next().unwrap();
    let (begin, end) = highlights[0];
    assert_eq!(&doc.content()[begin..end], to_search);
}

#[test]
fn test_update_doc() {
    let mut index = InvertedIndex::new();
    let doc = Document::new("0", "abc åäö");
    index.index(doc);
    let doc = Document::new("0", "different");
    index.index(doc);
    let search_results = index.search("å");
    assert!(search_results.is_empty());
    assert_eq!(index.docs.len(), 1);
}
<|endoftext|>


<commit_before>#![no_std]

#[doc(hidden)]
pub extern crate core as _core;

/// Asserts at compile-time that the types have equal sizes.
///
/// # Example
///
/// ```
/// #[macro_use]
/// extern crate static_assertions;
///
/// struct Byte(u8);
///
/// fn main() {
///     assert_eq_size!(Byte, u8);
///     // assert_eq_size!(Byte, u16); // Fails to compile
///
///     // Supports unlimited arguments:
///     assert_eq_size!([Byte; 4], [u16; 2], u32);
/// }
/// ```
#[macro_export]
macro_rules! assert_eq_size {
    ($x:ty, $y:ty) => {
        #[allow(unused_unsafe)]
        unsafe {
            use $crate::_core::mem::{forget, transmute, uninitialized};
            forget::<$y>(transmute(uninitialized::<$x>()));
        }
    };
    ($x:ty, $y:ty, $($rest:ty),+) => {
        assert_eq_size!($x, $y);
        assert_eq_size!($y, $($rest),+);
    };
}

/// Asserts at compile-time that the values have equal sizes.
#[macro_export]
macro_rules! assert_eq_size_val {
    ($x:expr, $($xs:expr),+) => {
        #[allow(unused_unsafe)]
        unsafe {
            use $crate::_core::{mem, ptr};
            let mut copy = ptr::read(&$x);
            $(ptr::write(&mut copy, mem::transmute(ptr::read(&$xs)));)+
            mem::forget(copy);
        }
    }
}
<commit_msg>Make assert_eq_size iterative<commit_after>#![no_std]

#[doc(hidden)]
pub extern crate core as _core;

/// Asserts at compile-time that the types have equal sizes.
///
/// # Example
///
/// ```
/// #[macro_use]
/// extern crate static_assertions;
///
/// struct Byte(u8);
///
/// fn main() {
///     assert_eq_size!(Byte, u8);
///     // assert_eq_size!(Byte, u16); // Fails to compile
///
///     // Supports unlimited arguments:
///     assert_eq_size!([Byte; 4], [u16; 2], u32);
/// }
/// ```
#[macro_export]
macro_rules! assert_eq_size {
    ($x:ty, $($xs:ty),+) => {
        #[allow(unused_unsafe)]
        unsafe {
            use $crate::_core::mem::{forget, transmute, uninitialized};
            $(forget::<$xs>(transmute(uninitialized::<$x>()));)+
        }
    }
}

/// Asserts at compile-time that the values have equal sizes.
#[macro_export]
macro_rules! assert_eq_size_val {
    ($x:expr, $($xs:expr),+) => {
        #[allow(unused_unsafe)]
        unsafe {
            use $crate::_core::{mem, ptr};
            let mut copy = ptr::read(&$x);
            $(ptr::write(&mut copy, mem::transmute(ptr::read(&$xs)));)+
            mem::forget(copy);
        }
    }
}
<|endoftext|>


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

extern crate libc;
extern crate rustc;
extern crate syntax;

use rustc::plugin::Registry;
use std::c_str::CString;
use std::mem;
use std::gc::Gc;
use syntax::ast::{TokenTree, ExprLit, LitStr, Expr, Ident};
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, MacResult, MacExpr, DummyResult};
use syntax::ext::build::AstBuilder;
use syntax::parse::token;
use syntax::parse::token::{InternedString, COMMA, EOF};
use syntax::parse;
use syntax::parse::parser::Parser;

mod ffi {
    use libc::{c_char, c_int};

    pub struct ParseResult {
        pub success: c_int,
        pub error_message: *c_char,
        pub index: c_int,
        pub num_params: c_int,
    }

    #[link(name="parser", kind="static")]
    extern {
        pub fn init_parser();
        pub fn parse_query(query: *c_char, result: *mut ParseResult);
    }
}

struct ParseInfo {
    num_params: uint,
}

struct ParseError {
    message: CString,
    index: uint,
}

#[plugin_registrar]
#[doc(hidden)]
pub fn registrar(reg: &mut Registry) {
    reg.register_macro("sql", expand_sql);
    reg.register_macro("execute", expand_execute);
    unsafe { ffi::init_parser() }
}

fn expand_sql(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree])
              -> Box<MacResult> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(),
                                                Vec::from_slice(tts));

    let query_expr = cx.expand_expr(parser.parse_expr());
    let query = match parse_str_lit(cx, query_expr) {
        Some(query) => query,
        None => return DummyResult::expr(sp)
    };

    match parse(query.get()) {
        Ok(_) => {}
        Err(err) => parse_error(cx, query_expr.span, err),
    }

    MacExpr::new(query_expr)
}

fn expand_execute(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree])
                  -> Box<MacResult> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(),
                                                Vec::from_slice(tts));

    let conn = parser.parse_expr();

    if !parser.eat(&COMMA) {
        cx.span_err(parser.span, "expected `,`");
        return DummyResult::expr(sp);
    }

    let query_expr = cx.expand_expr(parser.parse_expr());
    let query = match parse_str_lit(cx, query_expr) {
        Some(query) => query,
        None => return DummyResult::expr(sp),
    };

    if !parser.eat(&COMMA) {
        cx.span_err(parser.span, "expected `,`");
        return DummyResult::expr(sp);
    }

    let args = match parse_args(cx, &mut parser) {
        Some(args) => args,
        None => return DummyResult::expr(sp),
    };

    match parse(query.get()) {
        Ok(ref info) if info.num_params != args.len() =>
            cx.span_err(sp, format!("Expected {} query parameters but got {}",
                                    info.num_params, args.len()).as_slice()),
        Ok(_) => {}
        Err(err) => parse_error(cx, query_expr.span, err),
    }

    let ident = Ident::new(token::intern("execute"));
    let args = cx.expr_vec(sp, args);
    MacExpr::new(cx.expr_method_call(sp, conn, ident, vec![query_expr, args]))
}

fn parse_error(cx: &mut ExtCtxt, sp: Span, err: ParseError) {
    cx.span_err(sp, format!("Invalid syntax at position {}: {}",
                            err.index,
                            err.message.as_str().unwrap()).as_slice());
}

fn parse_str_lit(cx: &mut ExtCtxt, e: &Expr) -> Option<InternedString> {
    match e.node {
        ExprLit(lit) => {
            match lit.node {
                LitStr(ref s, _) => Some(s.clone()),
                _ => {
                    cx.span_err(e.span, "expected string literal");
                    None
                }
            }
        }
        _ => {
            cx.span_err(e.span, "expected string literal");
            None
        }
    }
}

fn parse_args(cx: &mut ExtCtxt, parser: &mut Parser) -> Option<Vec<Gc<Expr>>> {
    let mut args = Vec::new();

    while parser.token != EOF {
        args.push(parser.parse_expr());

        if !parser.eat(&COMMA) && parser.token != EOF {
            cx.span_err(parser.span, "expected `,`");
            return None;
        }
    }

    Some(args)
}

fn parse(query: &str) -> Result<ParseInfo, ParseError> {
    unsafe {
        ffi::init_parser();
        let mut result = mem::uninitialized();
        query.with_c_str(|query| {
            ffi::parse_query(query, &mut result);
        });
        match result.success != 0 {
            true => Ok(ParseInfo {
                num_params: result.num_params as uint,
            }),
            false => Err(ParseError {
                message: CString::new(result.error_message, true),
                index: result.index as uint,
            }),
        }
    }
}
<commit_msg>Remove unneeded feature gate<commit_after>#![crate_id="postgres_macros"]
#![crate_type="dylib"]
#![feature(plugin_registrar)]

extern crate libc;
extern crate rustc;
extern crate syntax;

use rustc::plugin::Registry;
use std::c_str::CString;
use std::mem;
use std::gc::Gc;
use syntax::ast::{TokenTree, ExprLit, LitStr, Expr, Ident};
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, MacResult, MacExpr, DummyResult};
use syntax::ext::build::AstBuilder;
use syntax::parse::token;
use syntax::parse::token::{InternedString, COMMA, EOF};
use syntax::parse;
use syntax::parse::parser::Parser;

mod ffi {
    use libc::{c_char, c_int};

    pub struct ParseResult {
        pub success: c_int,
        pub error_message: *c_char,
        pub index: c_int,
        pub num_params: c_int,
    }

    #[link(name="parser", kind="static")]
    extern {
        pub fn init_parser();
        pub fn parse_query(query: *c_char, result: *mut ParseResult);
    }
}

struct ParseInfo {
    num_params: uint,
}

struct ParseError {
    message: CString,
    index: uint,
}

#[plugin_registrar]
#[doc(hidden)]
pub fn registrar(reg: &mut Registry) {
    reg.register_macro("sql", expand_sql);
    reg.register_macro("execute", expand_execute);
    unsafe { ffi::init_parser() }
}

fn expand_sql(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree])
              -> Box<MacResult> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(),
                                                Vec::from_slice(tts));

    let query_expr = cx.expand_expr(parser.parse_expr());
    let query = match parse_str_lit(cx, query_expr) {
        Some(query) => query,
        None => return DummyResult::expr(sp)
    };

    match parse(query.get()) {
        Ok(_) => {}
        Err(err) => parse_error(cx, query_expr.span, err),
    }

    MacExpr::new(query_expr)
}

fn expand_execute(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree])
                  -> Box<MacResult> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(),
                                                Vec::from_slice(tts));

    let conn = parser.parse_expr();

    if !parser.eat(&COMMA) {
        cx.span_err(parser.span, "expected `,`");
        return DummyResult::expr(sp);
    }

    let query_expr = cx.expand_expr(parser.parse_expr());
    let query = match parse_str_lit(cx, query_expr) {
        Some(query) => query,
        None => return DummyResult::expr(sp),
    };

    if !parser.eat(&COMMA) {
        cx.span_err(parser.span, "expected `,`");
        return DummyResult::expr(sp);
    }

    let args = match parse_args(cx, &mut parser) {
        Some(args) => args,
        None => return DummyResult::expr(sp),
    };

    match parse(query.get()) {
        Ok(ref info) if info.num_params != args.len() =>
            cx.span_err(sp, format!("Expected {} query parameters but got {}",
                                    info.num_params, args.len()).as_slice()),
        Ok(_) => {}
        Err(err) => parse_error(cx, query_expr.span, err),
    }

    let ident = Ident::new(token::intern("execute"));
    let args = cx.expr_vec(sp, args);
    MacExpr::new(cx.expr_method_call(sp, conn, ident, vec![query_expr, args]))
}

fn parse_error(cx: &mut ExtCtxt, sp: Span, err: ParseError) {
    cx.span_err(sp, format!("Invalid syntax at position {}: {}",
                            err.index,
                            err.message.as_str().unwrap()).as_slice());
}

fn parse_str_lit(cx: &mut ExtCtxt, e: &Expr) -> Option<InternedString> {
    match e.node {
        ExprLit(lit) => {
            match lit.node {
                LitStr(ref s, _) => Some(s.clone()),
                _ => {
                    cx.span_err(e.span, "expected string literal");
                    None
                }
            }
        }
        _ => {
            cx.span_err(e.span, "expected string literal");
            None
        }
    }
}

fn parse_args(cx: &mut ExtCtxt, parser: &mut Parser) -> Option<Vec<Gc<Expr>>> {
    let mut args = Vec::new();

    while parser.token != EOF {
        args.push(parser.parse_expr());

        if !parser.eat(&COMMA) && parser.token != EOF {
            cx.span_err(parser.span, "expected `,`");
            return None;
        }
    }

    Some(args)
}

fn parse(query: &str) -> Result<ParseInfo, ParseError> {
    unsafe {
        ffi::init_parser();
        let mut result = mem::uninitialized();
        query.with_c_str(|query| {
            ffi::parse_query(query, &mut result);
        });
        match result.success != 0 {
            true => Ok(ParseInfo {
                num_params: result.num_params as uint,
            }),
            false => Err(ParseError {
                message: CString::new(result.error_message, true),
                index: result.index as uint,
            }),
        }
    }
}
<|endoftext|>


<commit_before>use std::ffi::{OsStr, OsString};
use std::fmt::Debug;
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::os::unix::io::{IntoRawFd, FromRawFd};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio, Output, ExitStatus};
use std::thread::JoinHandle;

mod pipe;

pub trait Expression: Clone + Debug {
    fn exec(&self, context: IoContext) -> io::Result<ExitStatus>;

    fn run(&self) -> Result<Output, Error> {
        let context = IoContext {
            stdin: CloneableStdio::Inherit,
            stdout: CloneableStdio::Inherit,
            stderr: CloneableStdio::Inherit,
        };
        let status = try!(self.exec(context));
        Ok(Output{
            status: status,
            stdout: Vec::new(),
            stderr: Vec::new(),
        })
    }

    fn read(&self) -> Result<String, Error> {
        let (stdout, stdout_reader) = pipe_with_reader_thread();
        let context = IoContext {
            stdin: CloneableStdio::Inherit,
            stdout: CloneableStdio::Handle(stdout),
            stderr: CloneableStdio::Inherit,
        };
        let status = try!(self.exec(context));
        let output = Output{
            status: status,
            stdout: try!(stdout_reader.join().unwrap()),
            stderr: Vec::new(),
        };
        if output.status.success() {
            // TODO: should only trim newlines
            Ok(try!(String::from_utf8(output.stdout)).trim_right().to_string())
        } else {
            Err(Error::Status(output))
        }
    }
}

#[derive(Debug, Clone)]
pub struct ArgvCommand {
    argv: Vec<OsString>,
    stdout: Option<PathBuf>,
}

impl ArgvCommand {
    pub fn new<T: AsRef<OsStr>>(prog: T) -> ArgvCommand {
        ArgvCommand{
            argv: vec![prog.as_ref().to_owned()],
            stdout: None,
        }
    }

    pub fn arg<T: AsRef<OsStr>>(&mut self, arg: T) -> &mut Self {
        self.argv.push(arg.as_ref().to_owned());
        self
    }

    pub fn stdout<T: AsRef<Path>>(&mut self, path: T) -> &mut Self {
        self.stdout = Some(path.as_ref().to_owned());
        self
    }
}

impl Expression for ArgvCommand {
    fn exec(&self, context: IoContext) -> io::Result<ExitStatus> {
        let IoContext{stdin, stdout, stderr} = context;
        // Create a Command and add the args.
        let mut command = Command::new(&self.argv[0]);
        command.args(&self.argv[1..]);
        command.stdin(stdin.to_stdio());
        command.stdout(stdout.to_stdio());
        command.stderr(stderr.to_stdio());
        if let Some(ref path) = self.stdout {
            let file = try!(File::create(path));
            command.stdout(unsafe {
                Stdio::from_raw_fd(file.into_raw_fd())
            });
        }
        Ok(try!(command.status()))
    }
}

#[derive(Debug, Clone)]
pub struct Pipe {
    // TODO: Make this hold any Expression.
    left: ArgvCommand,
    right: ArgvCommand,
}

impl Pipe {
    pub fn new(left: &ArgvCommand, right: &ArgvCommand) -> Pipe {
        Pipe{left: left.clone(), right: right.clone()}
    }
}

impl Expression for Pipe {
    fn exec(&self, context: IoContext) -> io::Result<ExitStatus> {
        let IoContext{stdin, stdout, stderr} = context;
        let (read_pipe, write_pipe) = pipe::open_pipe();
        let left_context = IoContext{
            stdin: stdin,
            stdout: CloneableStdio::Handle(write_pipe),
            stderr: stderr.clone(),
        };
        let right_context = IoContext{
            stdin: CloneableStdio::Handle(read_pipe),
            stdout: stdout,
            stderr: stderr,
        };
        let left_clone = self.left.clone();
        let left_thread = std::thread::spawn(move || {
            left_clone.exec(left_context)
        });
        let right_status = self.right.exec(right_context);
        let left_status = left_thread.join().unwrap();  // TODO: handle errors here?
        match right_status {
            Err(_) => right_status,
            _ => left_status,
        }
    }
}

#[derive(Debug)]
pub enum Error {
    Io(io::Error),
    Utf8(std::string::FromUtf8Error),
    Status(Output),
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::Io(err)
    }
}

impl From<std::string::FromUtf8Error> for Error {
    fn from(err: std::string::FromUtf8Error) -> Error {
        Error::Utf8(err)
    }
}

#[derive(Clone, Debug)]
pub struct IoContext {
    stdin: CloneableStdio,
    stdout: CloneableStdio,
    stderr: CloneableStdio,
}

#[derive(Clone, Debug)]
enum CloneableStdio {
    Inherit,
    Handle(pipe::Handle),
}

impl CloneableStdio {
    fn to_stdio(self) -> Stdio {
        match self {
            CloneableStdio::Inherit => Stdio::inherit(),
            CloneableStdio::Handle(handle) => handle.to_stdio(),
        }
    }
}

fn pipe_with_reader_thread() -> (pipe::Handle, JoinHandle<io::Result<Vec<u8>>>) {
    let (read_pipe, write_pipe) = pipe::open_pipe();
    let thread = std::thread::spawn(move || {
        let mut read_file = read_pipe.to_file();
        let mut output = Vec::new();
        try!(read_file.read_to_end(&mut output));
        Ok(output)
    });
    (write_pipe, thread)
}

#[cfg(test)]
mod test {
    use super::{ArgvCommand, Pipe, Expression};
    use std::fs::File;
    use std::io::prelude::*;
    use std::path::PathBuf;
    use std::process::Command;
    use std::str::from_utf8;

    fn mktemp() -> PathBuf {
        // TODO: use duct for this :)
        let output = Command::new("mktemp").output().unwrap();
        let path: PathBuf = from_utf8(&output.stdout).unwrap().trim().into();
        println!("here's the path we're using: {:?}", path);
        path
    }

    #[test]
    fn test_run() {
        let result = ArgvCommand::new("true").arg("foo").arg("bar").run();
        assert!(result.unwrap().status.success());
    }

    #[test]
    fn test_read() {
        let output = ArgvCommand::new("echo").arg("hi").read().unwrap();
        assert_eq!(output, "hi");
    }

    #[test]
    fn test_stdout() {
        let path = mktemp();
        let result = ArgvCommand::new("echo").arg("hi").stdout(&path).run();
        assert!(result.unwrap().status.success());
        let mut contents = String::new();
        File::open(&path).unwrap().read_to_string(&mut contents).unwrap();
        assert_eq!(contents, "hi\n");
    }

    #[test]
    fn test_pipe() {
        let mut left = ArgvCommand::new("echo");
        left.arg("hi");
        let mut right = ArgvCommand::new("sed");
        right.arg("s/i/o/");
        let pipe = Pipe::new(&left, &right);
        let output = pipe.read().unwrap();
        assert_eq!(output, "ho");
    }
}
<commit_msg>use ArgvCommand in mktemp()<commit_after>use std::ffi::{OsStr, OsString};
use std::fmt::Debug;
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::os::unix::io::{IntoRawFd, FromRawFd};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio, Output, ExitStatus};
use std::thread::JoinHandle;

mod pipe;

pub trait Expression: Clone + Debug {
    fn exec(&self, context: IoContext) -> io::Result<ExitStatus>;

    fn run(&self) -> Result<Output, Error> {
        let context = IoContext {
            stdin: CloneableStdio::Inherit,
            stdout: CloneableStdio::Inherit,
            stderr: CloneableStdio::Inherit,
        };
        let status = try!(self.exec(context));
        Ok(Output{
            status: status,
            stdout: Vec::new(),
            stderr: Vec::new(),
        })
    }

    fn read(&self) -> Result<String, Error> {
        let (stdout, stdout_reader) = pipe_with_reader_thread();
        let context = IoContext {
            stdin: CloneableStdio::Inherit,
            stdout: CloneableStdio::Handle(stdout),
            stderr: CloneableStdio::Inherit,
        };
        let status = try!(self.exec(context));
        let output = Output{
            status: status,
            stdout: try!(stdout_reader.join().unwrap()),
            stderr: Vec::new(),
        };
        if output.status.success() {
            // TODO: should only trim newlines
            Ok(try!(String::from_utf8(output.stdout)).trim_right().to_string())
        } else {
            Err(Error::Status(output))
        }
    }
}

#[derive(Debug, Clone)]
pub struct ArgvCommand {
    argv: Vec<OsString>,
    stdout: Option<PathBuf>,
}

impl ArgvCommand {
    pub fn new<T: AsRef<OsStr>>(prog: T) -> ArgvCommand {
        ArgvCommand{
            argv: vec![prog.as_ref().to_owned()],
            stdout: None,
        }
    }

    pub fn arg<T: AsRef<OsStr>>(&mut self, arg: T) -> &mut Self {
        self.argv.push(arg.as_ref().to_owned());
        self
    }

    pub fn stdout<T: AsRef<Path>>(&mut self, path: T) -> &mut Self {
        self.stdout = Some(path.as_ref().to_owned());
        self
    }
}

impl Expression for ArgvCommand {
    fn exec(&self, context: IoContext) -> io::Result<ExitStatus> {
        let IoContext{stdin, stdout, stderr} = context;
        // Create a Command and add the args.
        let mut command = Command::new(&self.argv[0]);
        command.args(&self.argv[1..]);
        command.stdin(stdin.to_stdio());
        command.stdout(stdout.to_stdio());
        command.stderr(stderr.to_stdio());
        if let Some(ref path) = self.stdout {
            let file = try!(File::create(path));
            command.stdout(unsafe {
                Stdio::from_raw_fd(file.into_raw_fd())
            });
        }
        Ok(try!(command.status()))
    }
}

#[derive(Debug, Clone)]
pub struct Pipe {
    // TODO: Make this hold any Expression.
    left: ArgvCommand,
    right: ArgvCommand,
}

impl Pipe {
    pub fn new(left: &ArgvCommand, right: &ArgvCommand) -> Pipe {
        Pipe{left: left.clone(), right: right.clone()}
    }
}

impl Expression for Pipe {
    fn exec(&self, context: IoContext) -> io::Result<ExitStatus> {
        let IoContext{stdin, stdout, stderr} = context;
        let (read_pipe, write_pipe) = pipe::open_pipe();
        let left_context = IoContext{
            stdin: stdin,
            stdout: CloneableStdio::Handle(write_pipe),
            stderr: stderr.clone(),
        };
        let right_context = IoContext{
            stdin: CloneableStdio::Handle(read_pipe),
            stdout: stdout,
            stderr: stderr,
        };
        let left_clone = self.left.clone();
        let left_thread = std::thread::spawn(move || {
            left_clone.exec(left_context)
        });
        let right_status = self.right.exec(right_context);
        let left_status = left_thread.join().unwrap();  // TODO: handle errors here?
        match right_status {
            Err(_) => right_status,
            _ => left_status,
        }
    }
}

#[derive(Debug)]
pub enum Error {
    Io(io::Error),
    Utf8(std::string::FromUtf8Error),
    Status(Output),
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::Io(err)
    }
}

impl From<std::string::FromUtf8Error> for Error {
    fn from(err: std::string::FromUtf8Error) -> Error {
        Error::Utf8(err)
    }
}

#[derive(Clone, Debug)]
pub struct IoContext {
    stdin: CloneableStdio,
    stdout: CloneableStdio,
    stderr: CloneableStdio,
}

#[derive(Clone, Debug)]
enum CloneableStdio {
    Inherit,
    Handle(pipe::Handle),
}

impl CloneableStdio {
    fn to_stdio(self) -> Stdio {
        match self {
            CloneableStdio::Inherit => Stdio::inherit(),
            CloneableStdio::Handle(handle) => handle.to_stdio(),
        }
    }
}

fn pipe_with_reader_thread() -> (pipe::Handle, JoinHandle<io::Result<Vec<u8>>>) {
    let (read_pipe, write_pipe) = pipe::open_pipe();
    let thread = std::thread::spawn(move || {
        let mut read_file = read_pipe.to_file();
        let mut output = Vec::new();
        try!(read_file.read_to_end(&mut output));
        Ok(output)
    });
    (write_pipe, thread)
}

#[cfg(test)]
mod test {
    use super::{ArgvCommand, Pipe, Expression};
    use std::fs::File;
    use std::io::prelude::*;
    use std::path::PathBuf;

    fn mktemp() -> PathBuf {
        let output = ArgvCommand::new("mktemp").read().unwrap();
        let path: PathBuf = output.trim().into();
        println!("here's the path we're using: {:?}", path);
        path
    }

    #[test]
    fn test_run() {
        let result = ArgvCommand::new("true").arg("foo").arg("bar").run();
        assert!(result.unwrap().status.success());
    }

    #[test]
    fn test_read() {
        let output = ArgvCommand::new("echo").arg("hi").read().unwrap();
        assert_eq!(output, "hi");
    }

    #[test]
    fn test_stdout() {
        let path = mktemp();
        let result = ArgvCommand::new("echo").arg("hi").stdout(&path).run();
        assert!(result.unwrap().status.success());
        let mut contents = String::new();
        File::open(&path).unwrap().read_to_string(&mut contents).unwrap();
        assert_eq!(contents, "hi\n");
    }

    #[test]
    fn test_pipe() {
        let mut left = ArgvCommand::new("echo");
        left.arg("hi");
        let mut right = ArgvCommand::new("sed");
        right.arg("s/i/o/");
        let pipe = Pipe::new(&left, &right);
        let output = pipe.read().unwrap();
        assert_eq!(output, "ho");
    }
}
<|endoftext|>


<commit_before>//! A generic [Markov chain](https://en.wikipedia.org/wiki/Markov_chain) for almost any type. This
//! uses HashMaps internally, and so Eq and Hash are both required.
//!
//! # Examples
//!
//! ```
//! use markov::Chain;
//!
//! let mut chain = Chain::new();
//! chain.feed_str("I like cats and I like dogs.");
//! println!("{}", chain.generate_str());
//! ```
//!
//! ```
//! use markov::Chain;
//!
//! let mut chain = Chain::new();
//! chain.feed(vec![1u8, 2, 3, 5]).feed(vec![3u8, 9, 2]);
//! println!("{:?}", chain.generate());
//! ```
#![warn(missing_docs)]

extern crate rand;

use std::borrow::ToOwned;
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::fs::File;
use std::hash::Hash;
use std::io::BufReader;
use std::io::prelude::*;
use std::iter::Map;
use std::path::Path;
use std::rc::Rc;
use rand::{Rng, thread_rng};

/// The definition of all types that can be used in a Chain.
pub trait Chainable: Eq + Hash {}
impl<T> Chainable for T where T: Eq + Hash {}

/// A generic [Markov chain](https://en.wikipedia.org/wiki/Markov_chain) for almost any type. This
/// uses HashMaps internally, and so Eq and Hash are both required.
#[derive(PartialEq, Debug)]
pub struct Chain<T> where T: Chainable {
    map: HashMap<Vec<Option<Rc<T>>>, HashMap<Option<Rc<T>>, usize>>,
    order: usize,
}

impl<T> Chain<T> where T: Chainable {
    /// Constructs a new Markov chain.
    pub fn new() -> Chain<T> {
        Chain {
            map: {
                let mut map = HashMap::new();
                map.insert(vec!(None; 1), HashMap::new());
                map
            },
            order: 1,
        }
    }

    /// Choose a specific Markov chain order. The order is the number of previous tokens to use
    /// as the index into the map.
    pub fn order(&mut self, order: usize) -> &mut Chain<T> {
        assert!(order > 0);
        self.order = order;
        self.map.insert(vec!(None; self.order), HashMap::new());
        self
    }

    /// Determines whether or not the chain is empty. A chain is considered empty if nothing has
    /// been fed into it.
    pub fn is_empty(&self) -> bool {
        self.map[&vec!(None; self.order)].is_empty()
    }

    /// Feeds the chain a collection of tokens. This operation is O(n) where n is the number of
    /// tokens to be fed into the chain.
    pub fn feed(&mut self, tokens: Vec<T>) -> &mut Chain<T> {
        if tokens.is_empty() { return self }
        let mut toks = vec!(None; self.order);
        toks.extend(tokens.into_iter().map(|token| {
            Some(Rc::new(token))
        }));
        toks.push(None);
        for p in toks.windows(self.order + 1) {
            if !self.map.contains_key(&p[0..self.order].to_vec()) {
                self.map.insert(p[0..self.order].to_vec(), HashMap::new());
            }
            self.map.get_mut(&p[0..self.order].to_vec()).unwrap().add(p[self.order].clone());
        }
        self
    }

    /// Generates a collection of tokens from the chain. This operation is O(mn) where m is the
    /// length of the generated collection, and n is the number of possible states from a given
    /// state.
    pub fn generate(&self) -> Vec<Rc<T>> {
        let mut ret = Vec::new();
        let mut curs = vec!(None; self.order);
        loop {
            let next = self.map[&curs].next();
            curs = curs[1..self.order].to_vec();
            curs.push(next.clone());
            if let Some(next) = next { ret.push(next) };
            if curs[self.order - 1].is_none() { break }
        }
        ret
    }

    /// Generates a collection of tokens from the chain, starting with the given token. This
    /// operation is O(mn) where m is the length of the generated collection, and n is the number
    /// of possible states from a given state. This returns an empty vector if the token is not
    /// found.
    pub fn generate_from_token(&self, token: T) -> Vec<Rc<T>> {
        let token = Rc::new(token);
        if !self.map.contains_key(&vec!(Some(token.clone()); self.order)) { return Vec::new() }
        let mut ret = vec![token.clone()];
        let mut curs = vec!(Some(token.clone()); self.order);
        loop {
            let next = self.map[&curs].next();
            curs = curs[1..self.order].to_vec();
            curs.push(next.clone());
            if let Some(next) = next { ret.push(next) };
            if curs[self.order - 1].is_none() { break }
        }
        ret
    }

    /// Produces an infinite iterator of generated token collections.
    pub fn iter(&self) -> InfiniteChainIterator<T> {
        InfiniteChainIterator { chain: self }
    }

    /// Produces an iterator for the specified number of generated token collections.
    pub fn iter_for(&self, size: usize) -> SizedChainIterator<T> {
        SizedChainIterator { chain: self, size: size }
    }
}

impl Chain<String> {
    /// Feeds a string of text into the chain.
    pub fn feed_str(&mut self, string: &str) -> &mut Chain<String> {
        self.feed(string.split(" ").map(|s| s.to_owned()).collect())
    }

    /// Feeds a properly formatted file into the chain. This file should be formatted such that
    /// each line is a new sentence. Punctuation may be included if it is desired.
    pub fn feed_file(&mut self, path: &Path) -> &mut Chain<String> {
        let reader = BufReader::new(File::open(path).unwrap());
        for line in reader.lines() {
            let line = line.unwrap();
            let words: Vec<_> = line.split(&[' ', '\t', '\n', '\r'][..])
                                    .filter(|word| !word.is_empty())
                                    .collect();
            self.feed(words.iter().map(|&s| s.to_owned()).collect());
        }
        self
    }

    /// Converts the output of generate(...) on a String chain to a single String.
    fn vec_to_string(vec: Vec<Rc<String>>) -> String {
        let mut ret = String::new();
        for s in vec.iter() {
            ret.push_str(&s);
            ret.push_str(" ");
        }
        let len = ret.len();
        if len > 0 {
            ret.truncate(len - 1);
        }
        ret
    }

    /// Generates a random string of text.
    pub fn generate_str(&self) -> String {
        Chain::vec_to_string(self.generate())
    }

    /// Generates a random string of text starting with the desired token. This returns an empty
    /// string if the token is not found.
    pub fn generate_str_from_token(&self, string: &str) -> String {
        Chain::vec_to_string(self.generate_from_token(string.to_owned()))
    }

    /// Produces an infinite iterator of generated strings.
    pub fn str_iter(&self) -> InfiniteChainStringIterator {
        let vec_to_string: fn(Vec<Rc<String>>) -> String = Chain::vec_to_string;
        self.iter().map(vec_to_string)
    }

    /// Produces a sized iterator of generated strings.
    pub fn str_iter_for(&self, size: usize) -> SizedChainStringIterator {
        let vec_to_string: fn(Vec<Rc<String>>) -> String = Chain::vec_to_string;
        self.iter_for(size).map(vec_to_string)
    }
}

/// A sized iterator over a Markov chain of strings.
pub type SizedChainStringIterator<'a> =
Map<SizedChainIterator<'a, String>, fn(Vec<Rc<String>>) -> String>;

/// A sized iterator over a Markov chain.
pub struct SizedChainIterator<'a, T: Chainable + 'a> {
    chain: &'a Chain<T>,
    size: usize,
}

impl<'a, T> Iterator for SizedChainIterator<'a, T> where T: Chainable + 'a {
    type Item = Vec<Rc<T>>;
    fn next(&mut self) -> Option<Vec<Rc<T>>> {
        if self.size > 0 {
            self.size -= 1;
            Some(self.chain.generate())
        } else {
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.size, Some(self.size))
    }
}

/// An infinite iterator over a Markov chain of strings.
pub type InfiniteChainStringIterator<'a> =
Map<InfiniteChainIterator<'a, String>, fn(Vec<Rc<String>>) -> String>;

/// An infinite iterator over a Markov chain.
pub struct InfiniteChainIterator<'a, T: Chainable + 'a> {
    chain: &'a Chain<T>
}

impl<'a, T> Iterator for InfiniteChainIterator<'a, T> where T: Chainable + 'a {
    type Item = Vec<Rc<T>>;
    fn next(&mut self) -> Option<Vec<Rc<T>>> {
        Some(self.chain.generate())
    }
}

/// A collection of states for the Markov chain.
trait States<T: PartialEq> {
    /// Adds a state to this states collection.
    fn add(&mut self, token: Option<Rc<T>>);
    /// Gets the next state from this collection of states.
    fn next(&self) -> Option<Rc<T>>;
}

impl<T> States<T> for HashMap<Option<Rc<T>>, usize> where T: Chainable {
    fn add(&mut self, token: Option<Rc<T>>) {
        match self.entry(token) {
            Occupied(mut e) => *e.get_mut() += 1,
            Vacant(e) => { e.insert(1); },
        }
    }

    fn next(&self) -> Option<Rc<T>> {
        let mut sum = 0;
        for &value in self.values() {
            sum += value;
        }
        let mut rng = thread_rng();
        let cap = rng.gen_range(0, sum);
        sum = 0;
        for (key, &value) in self.iter() {
            sum += value;
            if sum > cap {
                return key.clone()
            }
        }
        unreachable!("The random number generator failed.")
    }
}

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

    #[test]
    fn new() {
        Chain::<u8>::new();
        Chain::<String>::new();
    }

    #[test]
    fn is_empty() {
        let mut chain = Chain::new();
        assert!(chain.is_empty());
        chain.feed(vec![1u8, 2, 3]);
        assert!(!chain.is_empty());
    }

    #[test]
    fn feed() {
        let mut chain = Chain::new();
        chain.feed(vec![3, 5, 10]).feed(vec![5, 12]);
    }

    #[test]
    fn generate() {
        let mut chain = Chain::new();
        chain.feed(vec![3u8, 5, 10]).feed(vec![5, 12]);
        let v = chain.generate().into_iter().map(|v| *v).collect();
        assert!([vec![3, 5, 10], vec![3, 5, 12], vec![5, 10], vec![5, 12]].contains(&v));
    }

    #[test]
    fn generate_for_higher_order() {
        let mut chain = Chain::new();
        chain.order(2);
        chain.feed(vec![3u8, 5, 10]).feed(vec![2, 3, 5, 12]);
        let v = chain.generate().into_iter().map(|v| *v).collect();
        assert!([vec![3, 5, 10], vec![3, 5, 12], vec![2, 3, 5, 10], vec![2, 3, 5, 12]].contains(&v));
    }

    #[test]
    fn generate_from_token() {
        let mut chain = Chain::new();
        chain.feed(vec![3u8, 5, 10]).feed(vec![5, 12]);
        let v = chain.generate_from_token(5).into_iter().map(|v| *v).collect();
        assert!([vec![5, 10], vec![5, 12]].contains(&v));
    }

    #[test]
    fn generate_from_unfound_token() {
        let mut chain = Chain::new();
        chain.feed(vec![3u8, 5, 10]).feed(vec![5, 12]);
        let v: Vec<_> = chain.generate_from_token(9).into_iter().map(|v| *v).collect();
        assert_eq!(v, vec![]);
    }

    #[test]
    fn iter() {
        let mut chain = Chain::new();
        chain.feed(vec![3u8, 5, 10]).feed(vec![5, 12]);
        assert_eq!(chain.iter().size_hint().1, None);
    }

    #[test]
    fn iter_for() {
        let mut chain = Chain::new();
        chain.feed(vec![3u8, 5, 10]).feed(vec![5, 12]);
        assert_eq!(chain.iter_for(5).collect::<Vec<_>>().len(), 5);
    }

    #[test]
    fn feed_str() {
        let mut chain = Chain::new();
        chain.feed_str("I like cats and dogs");
    }

    #[test]
    fn generate_str() {
        let mut chain = Chain::new();
        chain.feed_str("I like cats").feed_str("I hate cats");
        assert!(["I like cats", "I hate cats"].contains(&&chain.generate_str()[..]));
    }

    #[test]
    fn generate_str_from_token() {
        let mut chain = Chain::new();
        chain.feed_str("I like cats").feed_str("cats are cute");
        assert!(["cats", "cats are cute"].contains(&&chain.generate_str_from_token("cats")[..]));
    }

    #[test]
    fn generate_str_from_unfound_token() {
        let mut chain = Chain::new();
        chain.feed_str("I like cats").feed_str("cats are cute");
        assert_eq!(chain.generate_str_from_token("test"), "");
    }

    #[test]
    fn str_iter() {
        let mut chain = Chain::new();
        chain.feed_str("I like cats and I like dogs");
        assert_eq!(chain.str_iter().size_hint().1, None);
    }

    #[test]
    fn str_iter_for() {
        let mut chain = Chain::new();
        chain.feed_str("I like cats and I like dogs");
        assert_eq!(chain.str_iter_for(5).collect::<Vec<_>>().len(), 5);
    }
}
<commit_msg>Replace explicit iter with implicit on loop.<commit_after>//! A generic [Markov chain](https://en.wikipedia.org/wiki/Markov_chain) for almost any type. This
//! uses HashMaps internally, and so Eq and Hash are both required.
//!
//! # Examples
//!
//! ```
//! use markov::Chain;
//!
//! let mut chain = Chain::new();
//! chain.feed_str("I like cats and I like dogs.");
//! println!("{}", chain.generate_str());
//! ```
//!
//! ```
//! use markov::Chain;
//!
//! let mut chain = Chain::new();
//! chain.feed(vec![1u8, 2, 3, 5]).feed(vec![3u8, 9, 2]);
//! println!("{:?}", chain.generate());
//! ```
#![warn(missing_docs)]

extern crate rand;

use std::borrow::ToOwned;
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::fs::File;
use std::hash::Hash;
use std::io::BufReader;
use std::io::prelude::*;
use std::iter::Map;
use std::path::Path;
use std::rc::Rc;
use rand::{Rng, thread_rng};

/// The definition of all types that can be used in a Chain.
pub trait Chainable: Eq + Hash {}
impl<T> Chainable for T where T: Eq + Hash {}

/// A generic [Markov chain](https://en.wikipedia.org/wiki/Markov_chain) for almost any type. This
/// uses HashMaps internally, and so Eq and Hash are both required.
#[derive(PartialEq, Debug)]
pub struct Chain<T> where T: Chainable {
    map: HashMap<Vec<Option<Rc<T>>>, HashMap<Option<Rc<T>>, usize>>,
    order: usize,
}

impl<T> Chain<T> where T: Chainable {
    /// Constructs a new Markov chain.
    pub fn new() -> Chain<T> {
        Chain {
            map: {
                let mut map = HashMap::new();
                map.insert(vec!(None; 1), HashMap::new());
                map
            },
            order: 1,
        }
    }

    /// Choose a specific Markov chain order. The order is the number of previous tokens to use
    /// as the index into the map.
    pub fn order(&mut self, order: usize) -> &mut Chain<T> {
        assert!(order > 0);
        self.order = order;
        self.map.insert(vec!(None; self.order), HashMap::new());
        self
    }

    /// Determines whether or not the chain is empty. A chain is considered empty if nothing has
    /// been fed into it.
    pub fn is_empty(&self) -> bool {
        self.map[&vec!(None; self.order)].is_empty()
    }

    /// Feeds the chain a collection of tokens. This operation is O(n) where n is the number of
    /// tokens to be fed into the chain.
    pub fn feed(&mut self, tokens: Vec<T>) -> &mut Chain<T> {
        if tokens.is_empty() { return self }
        let mut toks = vec!(None; self.order);
        toks.extend(tokens.into_iter().map(|token| {
            Some(Rc::new(token))
        }));
        toks.push(None);
        for p in toks.windows(self.order + 1) {
            if !self.map.contains_key(&p[0..self.order].to_vec()) {
                self.map.insert(p[0..self.order].to_vec(), HashMap::new());
            }
            self.map.get_mut(&p[0..self.order].to_vec()).unwrap().add(p[self.order].clone());
        }
        self
    }

    /// Generates a collection of tokens from the chain. This operation is O(mn) where m is the
    /// length of the generated collection, and n is the number of possible states from a given
    /// state.
    pub fn generate(&self) -> Vec<Rc<T>> {
        let mut ret = Vec::new();
        let mut curs = vec!(None; self.order);
        loop {
            let next = self.map[&curs].next();
            curs = curs[1..self.order].to_vec();
            curs.push(next.clone());
            if let Some(next) = next { ret.push(next) };
            if curs[self.order - 1].is_none() { break }
        }
        ret
    }

    /// Generates a collection of tokens from the chain, starting with the given token. This
    /// operation is O(mn) where m is the length of the generated collection, and n is the number
    /// of possible states from a given state. This returns an empty vector if the token is not
    /// found.
    pub fn generate_from_token(&self, token: T) -> Vec<Rc<T>> {
        let token = Rc::new(token);
        if !self.map.contains_key(&vec!(Some(token.clone()); self.order)) { return Vec::new() }
        let mut ret = vec![token.clone()];
        let mut curs = vec!(Some(token.clone()); self.order);
        loop {
            let next = self.map[&curs].next();
            curs = curs[1..self.order].to_vec();
            curs.push(next.clone());
            if let Some(next) = next { ret.push(next) };
            if curs[self.order - 1].is_none() { break }
        }
        ret
    }

    /// Produces an infinite iterator of generated token collections.
    pub fn iter(&self) -> InfiniteChainIterator<T> {
        InfiniteChainIterator { chain: self }
    }

    /// Produces an iterator for the specified number of generated token collections.
    pub fn iter_for(&self, size: usize) -> SizedChainIterator<T> {
        SizedChainIterator { chain: self, size: size }
    }
}

impl Chain<String> {
    /// Feeds a string of text into the chain.
    pub fn feed_str(&mut self, string: &str) -> &mut Chain<String> {
        self.feed(string.split(" ").map(|s| s.to_owned()).collect())
    }

    /// Feeds a properly formatted file into the chain. This file should be formatted such that
    /// each line is a new sentence. Punctuation may be included if it is desired.
    pub fn feed_file(&mut self, path: &Path) -> &mut Chain<String> {
        let reader = BufReader::new(File::open(path).unwrap());
        for line in reader.lines() {
            let line = line.unwrap();
            let words: Vec<_> = line.split(&[' ', '\t', '\n', '\r'][..])
                                    .filter(|word| !word.is_empty())
                                    .collect();
            self.feed(words.iter().map(|&s| s.to_owned()).collect());
        }
        self
    }

    /// Converts the output of generate(...) on a String chain to a single String.
    fn vec_to_string(vec: Vec<Rc<String>>) -> String {
        let mut ret = String::new();
        for s in &vec {
            ret.push_str(&s);
            ret.push_str(" ");
        }
        let len = ret.len();
        if len > 0 {
            ret.truncate(len - 1);
        }
        ret
    }

    /// Generates a random string of text.
    pub fn generate_str(&self) -> String {
        Chain::vec_to_string(self.generate())
    }

    /// Generates a random string of text starting with the desired token. This returns an empty
    /// string if the token is not found.
    pub fn generate_str_from_token(&self, string: &str) -> String {
        Chain::vec_to_string(self.generate_from_token(string.to_owned()))
    }

    /// Produces an infinite iterator of generated strings.
    pub fn str_iter(&self) -> InfiniteChainStringIterator {
        let vec_to_string: fn(Vec<Rc<String>>) -> String = Chain::vec_to_string;
        self.iter().map(vec_to_string)
    }

    /// Produces a sized iterator of generated strings.
    pub fn str_iter_for(&self, size: usize) -> SizedChainStringIterator {
        let vec_to_string: fn(Vec<Rc<String>>) -> String = Chain::vec_to_string;
        self.iter_for(size).map(vec_to_string)
    }
}

/// A sized iterator over a Markov chain of strings.
pub type SizedChainStringIterator<'a> =
Map<SizedChainIterator<'a, String>, fn(Vec<Rc<String>>) -> String>;

/// A sized iterator over a Markov chain.
pub struct SizedChainIterator<'a, T: Chainable + 'a> {
    chain: &'a Chain<T>,
    size: usize,
}

impl<'a, T> Iterator for SizedChainIterator<'a, T> where T: Chainable + 'a {
    type Item = Vec<Rc<T>>;
    fn next(&mut self) -> Option<Vec<Rc<T>>> {
        if self.size > 0 {
            self.size -= 1;
            Some(self.chain.generate())
        } else {
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.size, Some(self.size))
    }
}

/// An infinite iterator over a Markov chain of strings.
pub type InfiniteChainStringIterator<'a> =
Map<InfiniteChainIterator<'a, String>, fn(Vec<Rc<String>>) -> String>;

/// An infinite iterator over a Markov chain.
pub struct InfiniteChainIterator<'a, T: Chainable + 'a> {
    chain: &'a Chain<T>
}

impl<'a, T> Iterator for InfiniteChainIterator<'a, T> where T: Chainable + 'a {
    type Item = Vec<Rc<T>>;
    fn next(&mut self) -> Option<Vec<Rc<T>>> {
        Some(self.chain.generate())
    }
}

/// A collection of states for the Markov chain.
trait States<T: PartialEq> {
    /// Adds a state to this states collection.
    fn add(&mut self, token: Option<Rc<T>>);
    /// Gets the next state from this collection of states.
    fn next(&self) -> Option<Rc<T>>;
}

impl<T> States<T> for HashMap<Option<Rc<T>>, usize> where T: Chainable {
    fn add(&mut self, token: Option<Rc<T>>) {
        match self.entry(token) {
            Occupied(mut e) => *e.get_mut() += 1,
            Vacant(e) => { e.insert(1); },
        }
    }

    fn next(&self) -> Option<Rc<T>> {
        let mut sum = 0;
        for &value in self.values() {
            sum += value;
        }
        let mut rng = thread_rng();
        let cap = rng.gen_range(0, sum);
        sum = 0;
        for (key, &value) in self.iter() {
            sum += value;
            if sum > cap {
                return key.clone()
            }
        }
        unreachable!("The random number generator failed.")
    }
}

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

    #[test]
    fn new() {
        Chain::<u8>::new();
        Chain::<String>::new();
    }

    #[test]
    fn is_empty() {
        let mut chain = Chain::new();
        assert!(chain.is_empty());
        chain.feed(vec![1u8, 2, 3]);
        assert!(!chain.is_empty());
    }

    #[test]
    fn feed() {
        let mut chain = Chain::new();
        chain.feed(vec![3, 5, 10]).feed(vec![5, 12]);
    }

    #[test]
    fn generate() {
        let mut chain = Chain::new();
        chain.feed(vec![3u8, 5, 10]).feed(vec![5, 12]);
        let v = chain.generate().into_iter().map(|v| *v).collect();
        assert!([vec![3, 5, 10], vec![3, 5, 12], vec![5, 10], vec![5, 12]].contains(&v));
    }

    #[test]
    fn generate_for_higher_order() {
        let mut chain = Chain::new();
        chain.order(2);
        chain.feed(vec![3u8, 5, 10]).feed(vec![2, 3, 5, 12]);
        let v = chain.generate().into_iter().map(|v| *v).collect();
        assert!([vec![3, 5, 10], vec![3, 5, 12], vec![2, 3, 5, 10], vec![2, 3, 5, 12]].contains(&v));
    }

    #[test]
    fn generate_from_token() {
        let mut chain = Chain::new();
        chain.feed(vec![3u8, 5, 10]).feed(vec![5, 12]);
        let v = chain.generate_from_token(5).into_iter().map(|v| *v).collect();
        assert!([vec![5, 10], vec![5, 12]].contains(&v));
    }

    #[test]
    fn generate_from_unfound_token() {
        let mut chain = Chain::new();
        chain.feed(vec![3u8, 5, 10]).feed(vec![5, 12]);
        let v: Vec<_> = chain.generate_from_token(9).into_iter().map(|v| *v).collect();
        assert_eq!(v, vec![]);
    }

    #[test]
    fn iter() {
        let mut chain = Chain::new();
        chain.feed(vec![3u8, 5, 10]).feed(vec![5, 12]);
        assert_eq!(chain.iter().size_hint().1, None);
    }

    #[test]
    fn iter_for() {
        let mut chain = Chain::new();
        chain.feed(vec![3u8, 5, 10]).feed(vec![5, 12]);
        assert_eq!(chain.iter_for(5).collect::<Vec<_>>().len(), 5);
    }

    #[test]
    fn feed_str() {
        let mut chain = Chain::new();
        chain.feed_str("I like cats and dogs");
    }

    #[test]
    fn generate_str() {
        let mut chain = Chain::new();
        chain.feed_str("I like cats").feed_str("I hate cats");
        assert!(["I like cats", "I hate cats"].contains(&&chain.generate_str()[..]));
    }

    #[test]
    fn generate_str_from_token() {
        let mut chain = Chain::new();
        chain.feed_str("I like cats").feed_str("cats are cute");
        assert!(["cats", "cats are cute"].contains(&&chain.generate_str_from_token("cats")[..]));
    }

    #[test]
    fn generate_str_from_unfound_token() {
        let mut chain = Chain::new();
        chain.feed_str("I like cats").feed_str("cats are cute");
        assert_eq!(chain.generate_str_from_token("test"), "");
    }

    #[test]
    fn str_iter() {
        let mut chain = Chain::new();
        chain.feed_str("I like cats and I like dogs");
        assert_eq!(chain.str_iter().size_hint().1, None);
    }

    #[test]
    fn str_iter_for() {
        let mut chain = Chain::new();
        chain.feed_str("I like cats and I like dogs");
        assert_eq!(chain.str_iter_for(5).collect::<Vec<_>>().len(), 5);
    }
}
<|endoftext|>