<commit_before>//! References to garbage-collected JavaScript values.
//!
//! A _handle_ is a safe reference to a JavaScript value that is owned and managed
//! by the JavaScript engine's memory management system (the garbage collector).
//!
//! Neon APIs that accept and return JavaScript values never use raw pointer types
//! ([`*T`][ptr]) or reference types ([`&T`][ref]). Instead they use the special
//! Neon type [`Handle`](Handle), which encapsulates a JavaScript
//! [`Value`](crate::types::Value) and ensures that Rust only maintains access to
//! the value while it is guaranteed to be valid.
//!
//! ## Working with Handles
//!
//! The `Handle` type implements the standard [`Deref`][Deref] trait, so you can
//! use a `Handle<T>` as if it had type `&T`, including directly calling methods of
//! `T` on the handle.
//!
//! For example, a `Handle<JsNumber>` behaves like a `&JsNumber`, which means we can
//! call methods such as [`JsNumber::value()`](crate::types::JsNumber::value) on the
//! handle:
//!
//! ```ignore
//!     let n: Handle<JsNumber> = cx.argument(0)?;
//!     let v = n.value(&mut cx); // JsNumber::value()
//! ```
//!
//! ## Example
//!
//! This Neon function takes an object as its argument, extracts two properties,
//! `width` and `height`, and multiplies them together as numbers. Each JavaScript
//! value in the calculation is stored locally in a `Handle`.
//!
//! ```ignore
//! fn area(mut cx: FunctionContext) -> JsResult<JsNumber> {
//!     let rect: Handle<JsObject> = cx.argument(0)?;
//!
//!     let width: Handle<JsNumber> = rect
//!         .get(&mut cx, "width")?
//!         .downcast(&mut cx)?;
//!     let w: f64 = width.value(&mut cx);
//!
//!     let height: Handle<JsNumber> = rect
//!         .get(&mut cx, "height")?
//!         .downcast(&mut cx)?;
//!     let h: f64 = height.value(&mut cx);
//!
//!     Ok(cx.number(w * h))
//! }
//! ```
//!
//! [ptr]: https://doc.rust-lang.org/std/primitive.pointer.html
//! [ref]: https://doc.rust-lang.org/std/primitive.reference.html
//! [Deref]: https://doc.rust-lang.org/std/ops/trait.Deref.html

pub(crate) mod internal;

#[cfg(feature = "napi-1")]
pub(crate) mod root;

#[cfg(feature = "napi-1")]
pub use self::root::Root;

use self::internal::SuperType;
use context::internal::Env;
use context::Context;
use neon_runtime;
use neon_runtime::raw;
use result::{JsResult, JsResultExt};
use std::error::Error;
use std::fmt::{self, Debug, Display};
use std::marker::PhantomData;
use std::ops::{Deref, DerefMut};
use types::Value;

/// The trait of data that is managed by the JS garbage collector and can only be accessed via handles.
pub trait Managed: Copy {
    fn to_raw(self) -> raw::Local;

    fn from_raw(env: Env, h: raw::Local) -> Self;
}

/// A safely rooted _handle_ to a JS value in memory that is managed by the garbage collector.
#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct Handle<'a, T: Managed + 'a> {
    value: T,
    phantom: PhantomData<&'a T>,
}

#[cfg(feature = "legacy-runtime")]
impl<'a, T: Managed + 'a> PartialEq for Handle<'a, T> {
    fn eq(&self, other: &Self) -> bool {
        unsafe { neon_runtime::mem::same_handle(self.to_raw(), other.to_raw()) }
    }
}

#[cfg(feature = "legacy-runtime")]
impl<'a, T: Managed + 'a> Eq for Handle<'a, T> {}

impl<'a, T: Managed + 'a> Handle<'a, T> {
    pub(crate) fn new_internal(value: T) -> Handle<'a, T> {
        Handle {
            value,
            phantom: PhantomData,
        }
    }
}

/// An error representing a failed downcast.
#[derive(PartialEq, Eq, PartialOrd, Ord, Clone)]
pub struct DowncastError<F: Value, T: Value> {
    phantom_from: PhantomData<F>,
    phantom_to: PhantomData<T>,
}

impl<F: Value, T: Value> Debug for DowncastError<F, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "DowncastError")
    }
}

impl<F: Value, T: Value> DowncastError<F, T> {
    fn new() -> Self {
        DowncastError {
            phantom_from: PhantomData,
            phantom_to: PhantomData,
        }
    }
}

impl<F: Value, T: Value> Display for DowncastError<F, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "failed to downcast {} to {}", F::name(), T::name())
    }
}

impl<F: Value, T: Value> Error for DowncastError<F, T> {}

/// The result of a call to `Handle::downcast()`.
pub type DowncastResult<'a, F, T> = Result<Handle<'a, T>, DowncastError<F, T>>;

impl<'a, F: Value, T: Value> JsResultExt<'a, T> for DowncastResult<'a, F, T> {
    fn or_throw<'b, C: Context<'b>>(self, cx: &mut C) -> JsResult<'a, T> {
        match self {
            Ok(v) => Ok(v),
            Err(e) => cx.throw_type_error(&e.to_string()),
        }
    }
}

impl<'a, T: Value> Handle<'a, T> {
    /// Safely upcast a handle to a supertype.
    ///
    /// This method does not require an execution context because it only copies a handle.
    pub fn upcast<U: Value + SuperType<T>>(&self) -> Handle<'a, U> {
        Handle::new_internal(SuperType::upcast_internal(self.value))
    }

    #[cfg(feature = "legacy-runtime")]
    /// Tests whether this value is an instance of the given type.
    ///
    /// # Example:
    ///
    /// ```no_run
    /// # use neon::prelude::*;
    /// # fn my_neon_function(mut cx: FunctionContext) -> JsResult<JsUndefined> {
    /// let v: Handle<JsValue> = cx.number(17).upcast();
    /// v.is_a::<JsString>(); // false
    /// v.is_a::<JsNumber>(); // true
    /// v.is_a::<JsValue>();  // true
    /// # Ok(cx.undefined())
    /// # }
    /// ```
    pub fn is_a<U: Value>(&self) -> bool {
        U::is_typeof(Env::current(), self.value)
    }

    #[cfg(feature = "napi-1")]
    /// Tests whether this value is an instance of the given type.
    ///
    /// # Example:
    ///
    /// ```no_run
    /// # use neon::prelude::*;
    /// # fn my_neon_function(mut cx: FunctionContext) -> JsResult<JsUndefined> {
    /// let v: Handle<JsValue> = cx.number(17).upcast();
    /// v.is_a::<JsString, _>(&mut cx); // false
    /// v.is_a::<JsNumber, _>(&mut cx); // true
    /// v.is_a::<JsValue, _>(&mut cx);  // true
    /// # Ok(cx.undefined())
    /// # }
    /// ```
    pub fn is_a<'b, U: Value, C: Context<'b>>(&self, cx: &mut C) -> bool {
        U::is_typeof(cx.env(), self.value)
    }

    #[cfg(feature = "legacy-runtime")]
    /// Attempts to downcast a handle to another type, which may fail. A failure
    /// to downcast **does not** throw a JavaScript exception, so it's OK to
    /// continue interacting with the JS engine if this method produces an `Err`
    /// result.
    pub fn downcast<U: Value>(&self) -> DowncastResult<'a, T, U> {
        match U::downcast(Env::current(), self.value) {
            Some(v) => Ok(Handle::new_internal(v)),
            None => Err(DowncastError::new()),
        }
    }

    #[cfg(feature = "napi-1")]
    /// Attempts to downcast a handle to another type, which may fail. A failure
    /// to downcast **does not** throw a JavaScript exception, so it's OK to
    /// continue interacting with the JS engine if this method produces an `Err`
    /// result.
    pub fn downcast<'b, U: Value, C: Context<'b>>(&self, cx: &mut C) -> DowncastResult<'a, T, U> {
        match U::downcast(cx.env(), self.value) {
            Some(v) => Ok(Handle::new_internal(v)),
            None => Err(DowncastError::new()),
        }
    }

    #[cfg(feature = "legacy-runtime")]
    /// Attempts to downcast a handle to another type, raising a JavaScript `TypeError`
    /// exception on failure. This method is a convenient shorthand, equivalent to
    /// `self.downcast::<U>().or_throw::<C>(cx)`.
    pub fn downcast_or_throw<'b, U: Value, C: Context<'b>>(&self, cx: &mut C) -> JsResult<'a, U> {
        self.downcast().or_throw(cx)
    }

    #[cfg(feature = "napi-1")]
    /// Attempts to downcast a handle to another type, raising a JavaScript `TypeError`
    /// exception on failure. This method is a convenient shorthand, equivalent to
    /// `self.downcast::<U>().or_throw::<C>(cx)`.
    pub fn downcast_or_throw<'b, U: Value, C: Context<'b>>(&self, cx: &mut C) -> JsResult<'a, U> {
        self.downcast(cx).or_throw(cx)
    }

    #[cfg(feature = "napi-1")]
    pub fn strict_equals<'b, U: Value, C: Context<'b>>(
        &self,
        cx: &mut C,
        other: Handle<'b, U>,
    ) -> bool {
        unsafe {
            neon_runtime::mem::strict_equals(cx.env().to_raw(), self.to_raw(), other.to_raw())
        }
    }
}

impl<'a, T: Managed> Deref for Handle<'a, T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.value
    }
}

impl<'a, T: Managed> DerefMut for Handle<'a, T> {
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}
<commit_msg>Update Deref wording to be closer to the `JsBox` wording.<commit_after>//! References to garbage-collected JavaScript values.
//!
//! A _handle_ is a safe reference to a JavaScript value that is owned and managed
//! by the JavaScript engine's memory management system (the garbage collector).
//!
//! Neon APIs that accept and return JavaScript values never use raw pointer types
//! ([`*T`][ptr]) or reference types ([`&T`][ref]). Instead they use the special
//! Neon type [`Handle`](Handle), which encapsulates a JavaScript
//! [`Value`](crate::types::Value) and ensures that Rust only maintains access to
//! the value while it is guaranteed to be valid.
//!
//! ## Working with Handles
//!
//! The `Handle<T>` type automatically dereferences to `T` (via the standard
//! [`Deref`][Deref] trait), so you can call `T`'s methods on a value of type
//! `Handle<T>`. For example, we can call
//! [`JsNumber::value()`](crate::types::JsNumber::value) on a `Handle<JsNumber>`:
//!
//! ```ignore
//!     let n: Handle<JsNumber> = cx.argument(0)?;
//!     let v = n.value(&mut cx); // JsNumber::value()
//! ```
//!
//! ## Example
//!
//! This Neon function takes an object as its argument, extracts two properties,
//! `width` and `height`, and multiplies them together as numbers. Each JavaScript
//! value in the calculation is stored locally in a `Handle`.
//!
//! ```ignore
//! fn area(mut cx: FunctionContext) -> JsResult<JsNumber> {
//!     let rect: Handle<JsObject> = cx.argument(0)?;
//!
//!     let width: Handle<JsNumber> = rect
//!         .get(&mut cx, "width")?
//!         .downcast(&mut cx)?;
//!     let w: f64 = width.value(&mut cx);
//!
//!     let height: Handle<JsNumber> = rect
//!         .get(&mut cx, "height")?
//!         .downcast(&mut cx)?;
//!     let h: f64 = height.value(&mut cx);
//!
//!     Ok(cx.number(w * h))
//! }
//! ```
//!
//! [ptr]: https://doc.rust-lang.org/std/primitive.pointer.html
//! [ref]: https://doc.rust-lang.org/std/primitive.reference.html
//! [Deref]: https://doc.rust-lang.org/std/ops/trait.Deref.html

pub(crate) mod internal;

#[cfg(feature = "napi-1")]
pub(crate) mod root;

#[cfg(feature = "napi-1")]
pub use self::root::Root;

use self::internal::SuperType;
use context::internal::Env;
use context::Context;
use neon_runtime;
use neon_runtime::raw;
use result::{JsResult, JsResultExt};
use std::error::Error;
use std::fmt::{self, Debug, Display};
use std::marker::PhantomData;
use std::ops::{Deref, DerefMut};
use types::Value;

/// The trait of data that is managed by the JS garbage collector and can only be accessed via handles.
pub trait Managed: Copy {
    fn to_raw(self) -> raw::Local;

    fn from_raw(env: Env, h: raw::Local) -> Self;
}

/// A safely rooted _handle_ to a JS value in memory that is managed by the garbage collector.
#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct Handle<'a, T: Managed + 'a> {
    value: T,
    phantom: PhantomData<&'a T>,
}

#[cfg(feature = "legacy-runtime")]
impl<'a, T: Managed + 'a> PartialEq for Handle<'a, T> {
    fn eq(&self, other: &Self) -> bool {
        unsafe { neon_runtime::mem::same_handle(self.to_raw(), other.to_raw()) }
    }
}

#[cfg(feature = "legacy-runtime")]
impl<'a, T: Managed + 'a> Eq for Handle<'a, T> {}

impl<'a, T: Managed + 'a> Handle<'a, T> {
    pub(crate) fn new_internal(value: T) -> Handle<'a, T> {
        Handle {
            value,
            phantom: PhantomData,
        }
    }
}

/// An error representing a failed downcast.
#[derive(PartialEq, Eq, PartialOrd, Ord, Clone)]
pub struct DowncastError<F: Value, T: Value> {
    phantom_from: PhantomData<F>,
    phantom_to: PhantomData<T>,
}

impl<F: Value, T: Value> Debug for DowncastError<F, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "DowncastError")
    }
}

impl<F: Value, T: Value> DowncastError<F, T> {
    fn new() -> Self {
        DowncastError {
            phantom_from: PhantomData,
            phantom_to: PhantomData,
        }
    }
}

impl<F: Value, T: Value> Display for DowncastError<F, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "failed to downcast {} to {}", F::name(), T::name())
    }
}

impl<F: Value, T: Value> Error for DowncastError<F, T> {}

/// The result of a call to `Handle::downcast()`.
pub type DowncastResult<'a, F, T> = Result<Handle<'a, T>, DowncastError<F, T>>;

impl<'a, F: Value, T: Value> JsResultExt<'a, T> for DowncastResult<'a, F, T> {
    fn or_throw<'b, C: Context<'b>>(self, cx: &mut C) -> JsResult<'a, T> {
        match self {
            Ok(v) => Ok(v),
            Err(e) => cx.throw_type_error(&e.to_string()),
        }
    }
}

impl<'a, T: Value> Handle<'a, T> {
    /// Safely upcast a handle to a supertype.
    ///
    /// This method does not require an execution context because it only copies a handle.
    pub fn upcast<U: Value + SuperType<T>>(&self) -> Handle<'a, U> {
        Handle::new_internal(SuperType::upcast_internal(self.value))
    }

    #[cfg(feature = "legacy-runtime")]
    /// Tests whether this value is an instance of the given type.
    ///
    /// # Example:
    ///
    /// ```no_run
    /// # use neon::prelude::*;
    /// # fn my_neon_function(mut cx: FunctionContext) -> JsResult<JsUndefined> {
    /// let v: Handle<JsValue> = cx.number(17).upcast();
    /// v.is_a::<JsString>(); // false
    /// v.is_a::<JsNumber>(); // true
    /// v.is_a::<JsValue>();  // true
    /// # Ok(cx.undefined())
    /// # }
    /// ```
    pub fn is_a<U: Value>(&self) -> bool {
        U::is_typeof(Env::current(), self.value)
    }

    #[cfg(feature = "napi-1")]
    /// Tests whether this value is an instance of the given type.
    ///
    /// # Example:
    ///
    /// ```no_run
    /// # use neon::prelude::*;
    /// # fn my_neon_function(mut cx: FunctionContext) -> JsResult<JsUndefined> {
    /// let v: Handle<JsValue> = cx.number(17).upcast();
    /// v.is_a::<JsString, _>(&mut cx); // false
    /// v.is_a::<JsNumber, _>(&mut cx); // true
    /// v.is_a::<JsValue, _>(&mut cx);  // true
    /// # Ok(cx.undefined())
    /// # }
    /// ```
    pub fn is_a<'b, U: Value, C: Context<'b>>(&self, cx: &mut C) -> bool {
        U::is_typeof(cx.env(), self.value)
    }

    #[cfg(feature = "legacy-runtime")]
    /// Attempts to downcast a handle to another type, which may fail. A failure
    /// to downcast **does not** throw a JavaScript exception, so it's OK to
    /// continue interacting with the JS engine if this method produces an `Err`
    /// result.
    pub fn downcast<U: Value>(&self) -> DowncastResult<'a, T, U> {
        match U::downcast(Env::current(), self.value) {
            Some(v) => Ok(Handle::new_internal(v)),
            None => Err(DowncastError::new()),
        }
    }

    #[cfg(feature = "napi-1")]
    /// Attempts to downcast a handle to another type, which may fail. A failure
    /// to downcast **does not** throw a JavaScript exception, so it's OK to
    /// continue interacting with the JS engine if this method produces an `Err`
    /// result.
    pub fn downcast<'b, U: Value, C: Context<'b>>(&self, cx: &mut C) -> DowncastResult<'a, T, U> {
        match U::downcast(cx.env(), self.value) {
            Some(v) => Ok(Handle::new_internal(v)),
            None => Err(DowncastError::new()),
        }
    }

    #[cfg(feature = "legacy-runtime")]
    /// Attempts to downcast a handle to another type, raising a JavaScript `TypeError`
    /// exception on failure. This method is a convenient shorthand, equivalent to
    /// `self.downcast::<U>().or_throw::<C>(cx)`.
    pub fn downcast_or_throw<'b, U: Value, C: Context<'b>>(&self, cx: &mut C) -> JsResult<'a, U> {
        self.downcast().or_throw(cx)
    }

    #[cfg(feature = "napi-1")]
    /// Attempts to downcast a handle to another type, raising a JavaScript `TypeError`
    /// exception on failure. This method is a convenient shorthand, equivalent to
    /// `self.downcast::<U>().or_throw::<C>(cx)`.
    pub fn downcast_or_throw<'b, U: Value, C: Context<'b>>(&self, cx: &mut C) -> JsResult<'a, U> {
        self.downcast(cx).or_throw(cx)
    }

    #[cfg(feature = "napi-1")]
    pub fn strict_equals<'b, U: Value, C: Context<'b>>(
        &self,
        cx: &mut C,
        other: Handle<'b, U>,
    ) -> bool {
        unsafe {
            neon_runtime::mem::strict_equals(cx.env().to_raw(), self.to_raw(), other.to_raw())
        }
    }
}

impl<'a, T: Managed> Deref for Handle<'a, T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.value
    }
}

impl<'a, T: Managed> DerefMut for Handle<'a, T> {
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}
<|endoftext|>


<commit_before>use super::Error;

// TODO: Implement Identifier
pub struct Identifier {
    name: String,
}

impl Identifier {
    pub fn new(name: &str) -> Result<Self, Error> {
        Ok(Identifier {
            name: name.to_string()
        })
    }

    pub fn get_name(&self) -> String {
        self.name.clone()
    }
}
<commit_msg>added todo for identifier<commit_after>use super::Error;

// TODO: Implement Identifier
pub struct Identifier {
    name: String,
}

impl Identifier {
    pub fn new(name: &str) -> Result<Self, Error> {
        Ok(Identifier {
            name: name.to_string()
        })
    }

    pub fn get_name(&self) -> String {
        self.name.clone()
    }

    fn check_string(val: String) -> bool {
        // TODO: Implement check_string for Identifier
        true
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use super::alignment::Alignment;

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Copy, Clone)]
pub struct Innovation(usize);

#[derive(Debug, PartialEq, Eq, Copy, Clone)]
pub struct InnovationRange(Innovation, Innovation);

#[derive(Debug, Clone)]
pub struct InnovationContainer<T> {
    pub map: BTreeMap<Innovation, T>,
}

impl Iterator for Innovation {
    type Item = Innovation;
    fn next(&mut self) -> Option<Innovation> {
        match self.0.checked_add(1) {
            Some(n) => {
                self.0 = n;
                Some(Innovation(n))
            }
            None => None,
        }
    }
}

impl Innovation {
    pub fn new(n: usize) -> Innovation {
        Innovation(n)
    }

    pub fn is_within(&self, range: &InnovationRange) -> bool {
        self >= &range.0 && self <= &range.1
    }

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

impl InnovationRange {
    // NOTE: This requires that the iteration is in sorted order.
    fn from_sorted_iter<I>(mut iter: I) -> Option<InnovationRange>
        where I: Iterator<Item = Innovation> + DoubleEndedIterator<Item = Innovation>
    {
        iter.next().map(|min| {
            match iter.next_back() {
                Some(max) => InnovationRange(min, max),
                None => InnovationRange(min, min),
            }
        })
    }
}

impl<T> InnovationContainer<T> {
    pub fn new() -> InnovationContainer<T> {
        InnovationContainer { map: BTreeMap::new() }
    }

    pub fn insert(&mut self, innov: Innovation, data: T) {
        // XXX
        // assert!(!self.map.contains_key(&innov));
        self.map.insert(innov, data);
    }

    pub fn get(&self, innov: &Innovation) -> Option<&T> {
        self.map.get(innov)
    }

    pub fn get_mut(&mut self, innov: &Innovation) -> Option<&mut T> {
        self.map.get_mut(innov)
    }

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

    // Returns the min and max innovation numbers (inclusive)
    pub fn innovation_range(&self) -> InnovationRange {
        InnovationRange::from_sorted_iter(self.map.keys().cloned()).unwrap()
    }

    fn align_as_container<'a>(&'a self,
                              right: &'a InnovationContainer<T>)
                              -> InnovationContainer<Alignment<'a, T>> {
        let mut c = InnovationContainer::new();
        self.align(right, &mut |innov, alignment| c.insert(innov, alignment));
        c
    }

    pub fn align<'a, F>(&'a self, right: &'a InnovationContainer<T>, f: &mut F)
        where F: FnMut(Innovation, Alignment<'a, T>)
    {
        let range_left = self.innovation_range();
        let range_right = right.innovation_range();

        for (innov_left, gene_left) in self.map.iter() {
            if innov_left.is_within(&range_right) {
                match right.get(innov_left) {
                    Some(gene_right) => f(*innov_left, Alignment::Match(gene_left, gene_right)),
                    None => f(*innov_left, Alignment::DisjointLeft(gene_left)),
                }
            } else {
                f(*innov_left, Alignment::ExcessLeft(gene_left))
            }
        }

        for (innov_right, gene_right) in right.map.iter() {
            if innov_right.is_within(&range_left) {
                if !self.map.contains_key(innov_right) {
                    f(*innov_right, Alignment::DisjointRight(gene_right))
                }
            } else {
                f(*innov_right, Alignment::ExcessRight(gene_right))
            }
        }
    }
}

#[test]
fn test_innovation() {
    let r = InnovationRange(Innovation(0), Innovation(100));
    assert_eq!(true, Innovation(0).is_within(&r));
    assert_eq!(true, Innovation(50).is_within(&r));
    assert_eq!(true, Innovation(100).is_within(&r));
    assert_eq!(false, Innovation(101).is_within(&r));
    let r = InnovationRange(Innovation(1), Innovation(100));
    assert_eq!(false, Innovation(0).is_within(&r));
    assert_eq!(true, Innovation(1).is_within(&r));
    assert_eq!(true, Innovation(50).is_within(&r));
    assert_eq!(true, Innovation(100).is_within(&r));
    assert_eq!(false, Innovation(101).is_within(&r));
}

#[test]
fn test_innovation_range() {
    let mut genome = InnovationContainer::new();
    genome.insert(Innovation(50), ());
    assert_eq!(InnovationRange(Innovation(50), Innovation(50)),
               genome.innovation_range());
    genome.insert(Innovation(20), ());
    assert_eq!(InnovationRange(Innovation(20), Innovation(50)),
               genome.innovation_range());
    genome.insert(Innovation(100), ());
    assert_eq!(InnovationRange(Innovation(20), Innovation(100)),
               genome.innovation_range());
    genome.insert(Innovation(0), ());
    assert_eq!(InnovationRange(Innovation(0), Innovation(100)),
               genome.innovation_range());
}

#[test]
fn test_innovation_alignment() {
    let mut left = InnovationContainer::new();
    let mut right = InnovationContainer::new();
    left.insert(Innovation(50), ());
    left.insert(Innovation(46), ());
    left.insert(Innovation(40), ());
    right.insert(Innovation(50), ());
    right.insert(Innovation(45), ());
    right.insert(Innovation(51), ());
    right.insert(Innovation(52), ());

    let c = left.align_as_container(&right);
    assert_eq!(6, c.len());
    assert_eq!(true, c.get(&Innovation(50)).unwrap().is_match());
    assert_eq!(true, c.get(&Innovation(40)).unwrap().is_excess_left());
    assert_eq!(true, c.get(&Innovation(45)).unwrap().is_disjoint_right());
    assert_eq!(true, c.get(&Innovation(46)).unwrap().is_disjoint_left());
    assert_eq!(true, c.get(&Innovation(51)).unwrap().is_excess_right());
    assert_eq!(true, c.get(&Innovation(52)).unwrap().is_excess_right());
}

#[test]
fn test_innovation_iter() {
    let mut innovations = Innovation(1);
    assert_eq!(Innovation(1), innovations);
    assert_eq!(Some(Innovation(2)), innovations.next());
    assert_eq!(Some(Innovation(3)), innovations.next());
    assert_eq!(Some(Innovation(4)), innovations.next());
    assert_eq!(Some(Innovation(5)), innovations.next());

    let mut innovations = Innovation(usize::max_value() - 1);
    assert_eq!(Some(Innovation(usize::max_value())), innovations.next());
    assert_eq!(None, innovations.next());
}
<commit_msg>Replace Innovation#is_within with InnovationRange#contains<commit_after>use std::collections::BTreeMap;
use super::alignment::Alignment;

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Copy, Clone)]
pub struct Innovation(usize);

#[derive(Debug, PartialEq, Eq, Copy, Clone)]
pub struct InnovationRange(Innovation, Innovation);

#[derive(Debug, Clone)]
pub struct InnovationContainer<T> {
    pub map: BTreeMap<Innovation, T>,
}

impl Iterator for Innovation {
    type Item = Innovation;
    fn next(&mut self) -> Option<Innovation> {
        match self.0.checked_add(1) {
            Some(n) => {
                self.0 = n;
                Some(Innovation(n))
            }
            None => None,
        }
    }
}

impl Innovation {
    pub fn new(n: usize) -> Innovation {
        Innovation(n)
    }

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

impl InnovationRange {
    pub fn contains(&self, innov: &Innovation) -> bool {
        innov >= &self.0 && innov <= &self.1
    }

    // NOTE: This requires that the iteration is in sorted order.
    fn from_sorted_iter<I>(mut iter: I) -> Option<InnovationRange>
        where I: Iterator<Item = Innovation> + DoubleEndedIterator<Item = Innovation>
    {
        iter.next().map(|min| {
            match iter.next_back() {
                Some(max) => InnovationRange(min, max),
                None => InnovationRange(min, min),
            }
        })
    }
}

impl<T> InnovationContainer<T> {
    pub fn new() -> InnovationContainer<T> {
        InnovationContainer { map: BTreeMap::new() }
    }

    pub fn insert(&mut self, innov: Innovation, data: T) {
        // XXX
        // assert!(!self.map.contains_key(&innov));
        self.map.insert(innov, data);
    }

    pub fn get(&self, innov: &Innovation) -> Option<&T> {
        self.map.get(innov)
    }

    pub fn get_mut(&mut self, innov: &Innovation) -> Option<&mut T> {
        self.map.get_mut(innov)
    }

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

    // Returns the min and max innovation numbers (inclusive)
    pub fn innovation_range(&self) -> InnovationRange {
        InnovationRange::from_sorted_iter(self.map.keys().cloned()).unwrap()
    }

    fn align_as_container<'a>(&'a self,
                              right: &'a InnovationContainer<T>)
                              -> InnovationContainer<Alignment<'a, T>> {
        let mut c = InnovationContainer::new();
        self.align(right, &mut |innov, alignment| c.insert(innov, alignment));
        c
    }

    pub fn align<'a, F>(&'a self, right: &'a InnovationContainer<T>, f: &mut F)
        where F: FnMut(Innovation, Alignment<'a, T>)
    {
        let range_left = self.innovation_range();
        let range_right = right.innovation_range();

        for (innov_left, gene_left) in self.map.iter() {
            if range_right.contains(innov_left) {
                match right.get(innov_left) {
                    Some(gene_right) => f(*innov_left, Alignment::Match(gene_left, gene_right)),
                    None => f(*innov_left, Alignment::DisjointLeft(gene_left)),
                }
            } else {
                f(*innov_left, Alignment::ExcessLeft(gene_left))
            }
        }

        for (innov_right, gene_right) in right.map.iter() {
            if range_left.contains(innov_right) {
                if !self.map.contains_key(innov_right) {
                    f(*innov_right, Alignment::DisjointRight(gene_right))
                }
            } else {
                f(*innov_right, Alignment::ExcessRight(gene_right))
            }
        }
    }
}

#[test]
fn test_innovation() {
    let r = InnovationRange(Innovation(0), Innovation(100));
    assert_eq!(true, r.contains(&Innovation(0)));
    assert_eq!(true, r.contains(&Innovation(50)));
    assert_eq!(true, r.contains(&Innovation(100)));
    assert_eq!(false, r.contains(&Innovation(101)));
    let r = InnovationRange(Innovation(1), Innovation(100));
    assert_eq!(false, r.contains(&Innovation(0)));
    assert_eq!(true, r.contains(&Innovation(1)));
    assert_eq!(true, r.contains(&Innovation(50)));
    assert_eq!(true, r.contains(&Innovation(100)));
    assert_eq!(false, r.contains(&Innovation(101)));
}

#[test]
fn test_innovation_range() {
    let mut genome = InnovationContainer::new();
    genome.insert(Innovation(50), ());
    assert_eq!(InnovationRange(Innovation(50), Innovation(50)),
               genome.innovation_range());
    genome.insert(Innovation(20), ());
    assert_eq!(InnovationRange(Innovation(20), Innovation(50)),
               genome.innovation_range());
    genome.insert(Innovation(100), ());
    assert_eq!(InnovationRange(Innovation(20), Innovation(100)),
               genome.innovation_range());
    genome.insert(Innovation(0), ());
    assert_eq!(InnovationRange(Innovation(0), Innovation(100)),
               genome.innovation_range());
}

#[test]
fn test_innovation_alignment() {
    let mut left = InnovationContainer::new();
    let mut right = InnovationContainer::new();
    left.insert(Innovation(50), ());
    left.insert(Innovation(46), ());
    left.insert(Innovation(40), ());
    right.insert(Innovation(50), ());
    right.insert(Innovation(45), ());
    right.insert(Innovation(51), ());
    right.insert(Innovation(52), ());

    let c = left.align_as_container(&right);
    assert_eq!(6, c.len());
    assert_eq!(true, c.get(&Innovation(50)).unwrap().is_match());
    assert_eq!(true, c.get(&Innovation(40)).unwrap().is_excess_left());
    assert_eq!(true, c.get(&Innovation(45)).unwrap().is_disjoint_right());
    assert_eq!(true, c.get(&Innovation(46)).unwrap().is_disjoint_left());
    assert_eq!(true, c.get(&Innovation(51)).unwrap().is_excess_right());
    assert_eq!(true, c.get(&Innovation(52)).unwrap().is_excess_right());
}

#[test]
fn test_innovation_iter() {
    let mut innovations = Innovation(1);
    assert_eq!(Innovation(1), innovations);
    assert_eq!(Some(Innovation(2)), innovations.next());
    assert_eq!(Some(Innovation(3)), innovations.next());
    assert_eq!(Some(Innovation(4)), innovations.next());
    assert_eq!(Some(Innovation(5)), innovations.next());

    let mut innovations = Innovation(usize::max_value() - 1);
    assert_eq!(Some(Innovation(usize::max_value())), innovations.next());
    assert_eq!(None, innovations.next());
}
<|endoftext|>


<commit_before>pub mod files;

use std::cmp;
use std::iter;
use std::io::Write;

const MAX_PATCH_LEN: usize = 0xFFFF;

#[derive(Debug)]
pub struct Ips {
    buffer: Vec<u8>
}

impl Ips {
    pub fn new(buffer: &Vec<u8>) -> Self {
        Ips {
            buffer: buffer.clone()
        }
    }

    pub fn diff(&self, change: &Vec<u8>) -> Vec<Patch> {
        let (change_within, additional) = change.split_at(self.buffer.len());
        let mut pairs = (0..).zip(self.buffer.iter().zip(change_within));

        let mut patches: Vec<Patch> = vec![];
        let mut index = 0;
        let mut data = vec![];

        while let Some((offset, (&before, &after))) = pairs.next() {
            if before != after {
                if data.is_empty() {
                    index = offset;
                }
                data.push(after);
            } else if ! data.is_empty() {
                let patch = Patch::new(index, data);
                patches.push(patch);
                data = vec![];
            }
        };

        for chunk in additional.chunks(MAX_PATCH_LEN) {
            let patch = Patch::new(self.buffer.len() as u32, chunk.to_vec());
            patches.push(patch);
        }

        patches
    }

    pub fn patch(&self, change: &Vec<Patch>) -> Vec<u8> {
        let total_size = cmp::max(self.buffer.len(), patch_max_len(&change).unwrap());

        let mut output = Vec::with_capacity(total_size);
        output.write(&self.buffer);

        for patch in change {
            {
                let (_, mut skipped) = output.split_at_mut(patch.addr as usize);
                skipped.write(&patch.data())
            }.map(|written_len| {
                let data = patch.data();
                let (_, unwritten_data) = data.split_at(written_len);
                for &b in unwritten_data {
                    output.push(b);
                }
            });
        }
        output
    }

}

pub fn patch_max_len(patch: &[Patch]) -> Option<usize> {
    patch.iter().map(|x| x.addr as usize + x.data().len()).max()
}

pub fn serialize_patches(patches: Vec<Patch>) -> Vec<u8> {
    let patch_contents: Vec<u8> = patches.iter().flat_map(|p| p.bytes()).collect();
    let mut binary: Vec<u8> = b"PATCH".to_vec();
    binary.extend(patch_contents);
    binary.extend(b"EOF".to_vec());
    binary
}

pub fn unserialize_patches(binary: Vec<u8>) -> Option<Vec<Patch>> {
    let mut patches: Vec<Patch> = Vec::new();

    let slice = &binary;
    let (header, mut slice) = slice.split_at(5);

    if header != b"PATCH" {
        return None;
    }

    loop {
        // TODO support RLE'd patches
        let (addr_slice, rest) = slice.split_at(3);
        let (len_slice, rest) = rest.split_at(2);

        if (! addr_slice.is_empty()) && len_slice.len() < 2 {
            // Malformed data
            return None;
        }

        let mut addr_array = [0; 3];
        addr_array.clone_from_slice(addr_slice);
        let addr = Patch::unserialize_addr_array(addr_array);

        let mut len_array = [0; 2];
        len_array.clone_from_slice(len_slice);
        let len = Patch::unserialize_len(len_array);

        if rest.len() < len {
            // Malformed data
            return None;
        }

        let (patch, rest) = if len == 0 {
            // should decode RLE
            let (rle_len_slice, rest) = rest.split_at(2);

            let mut rle_len_array = [0; 2];
            rle_len_array.clone_from_slice(rle_len_slice);
            let rle_len = Patch::unserialize_len(rle_len_array);

            let (rle_val_slice, rest) = rest.split_at(1);
            let rle_val = rle_val_slice[0];

            let patch = Patch::new(addr, from_rle(rle_len, rle_val));
            (patch, rest)
        } else {
            let (data, rest) = rest.split_at(len);

            (Patch::new(addr, data.to_vec()), rest)
        };

        patches.push(patch);

        if rest.len() < 3 {
            // Malformed data
            return None;
        }

        let (possible_eof, _) = rest.split_at(3);
        if possible_eof == b"EOF" {
            // We are done here
            break;
        }

        slice = rest;
    }

    Some(patches)
}

pub fn into_rle(data: Vec<u8>) -> (usize, u8) {
    (data.len(), data[0])
}

pub fn from_rle(len: usize, val: u8) -> Vec<u8> {
    iter::repeat(val).take(len).collect()
}

#[derive(Debug, PartialEq, Clone)]
pub struct Patch {
    addr: u32,
    data: Vec<u8>
}

impl Patch {
    pub fn new(addr: u32, data: Vec<u8>) -> Self {
        assert!(addr <= 0xFFFFFF);
        Patch {
            addr: addr,
            data: data
        }
    }

    pub fn data(&self) -> &Vec<u8> {
        &self.data
    }

    pub fn bytes(&self) -> Vec<u8> {
        let mut bytes: Vec<u8> = self.serialize_addr();
        let len = self.serialize_len();
        bytes.extend(len);
        bytes.extend(self.data());
        bytes
    }

    pub fn unserialize_addr_array(addr: [u8; 3]) -> u32 {
        ((addr[0] as u32) << 16) | ((addr[1] as u32) << 8) | addr[2] as u32
    }

    pub fn serialize_addr(&self) -> Vec<u8> {
        vec![(self.addr >> 16) as u8, (self.addr >> 8) as u8, self.addr as u8]
    }

    pub fn unserialize_len(addr: [u8; 2]) -> usize {
        (((addr[0] as u32) << 8) | addr[1] as u32) as usize
    }

    pub fn serialize_len(&self) -> Vec<u8> {
        vec![(self.data().len() >> 8) as u8, self.data().len() as u8]
    }
}
<commit_msg>Rename rle functions (hopefully more clear)<commit_after>pub mod files;

use std::cmp;
use std::iter;
use std::io::Write;

const MAX_PATCH_LEN: usize = 0xFFFF;

#[derive(Debug)]
pub struct Ips {
    buffer: Vec<u8>
}

impl Ips {
    pub fn new(buffer: &Vec<u8>) -> Self {
        Ips {
            buffer: buffer.clone()
        }
    }

    pub fn diff(&self, change: &Vec<u8>) -> Vec<Patch> {
        let (change_within, additional) = change.split_at(self.buffer.len());
        let mut pairs = (0..).zip(self.buffer.iter().zip(change_within));

        let mut patches: Vec<Patch> = vec![];
        let mut index = 0;
        let mut data = vec![];

        while let Some((offset, (&before, &after))) = pairs.next() {
            if before != after {
                if data.is_empty() {
                    index = offset;
                }
                data.push(after);
            } else if ! data.is_empty() {
                let patch = Patch::new(index, data);
                patches.push(patch);
                data = vec![];
            }
        };

        for chunk in additional.chunks(MAX_PATCH_LEN) {
            let patch = Patch::new(self.buffer.len() as u32, chunk.to_vec());
            patches.push(patch);
        }

        patches
    }

    pub fn patch(&self, change: &Vec<Patch>) -> Vec<u8> {
        let total_size = cmp::max(self.buffer.len(), patch_max_len(&change).unwrap());

        let mut output = Vec::with_capacity(total_size);
        output.write(&self.buffer);

        for patch in change {
            {
                let (_, mut skipped) = output.split_at_mut(patch.addr as usize);
                skipped.write(&patch.data())
            }.map(|written_len| {
                let data = patch.data();
                let (_, unwritten_data) = data.split_at(written_len);
                for &b in unwritten_data {
                    output.push(b);
                }
            });
        }
        output
    }

}

pub fn patch_max_len(patch: &[Patch]) -> Option<usize> {
    patch.iter().map(|x| x.addr as usize + x.data().len()).max()
}

pub fn serialize_patches(patches: Vec<Patch>) -> Vec<u8> {
    let patch_contents: Vec<u8> = patches.iter().flat_map(|p| p.bytes()).collect();
    let mut binary: Vec<u8> = b"PATCH".to_vec();
    binary.extend(patch_contents);
    binary.extend(b"EOF".to_vec());
    binary
}

pub fn unserialize_patches(binary: Vec<u8>) -> Option<Vec<Patch>> {
    let mut patches: Vec<Patch> = Vec::new();

    let slice = &binary;
    let (header, mut slice) = slice.split_at(5);

    if header != b"PATCH" {
        return None;
    }

    loop {
        // TODO support RLE'd patches
        let (addr_slice, rest) = slice.split_at(3);
        let (len_slice, rest) = rest.split_at(2);

        if (! addr_slice.is_empty()) && len_slice.len() < 2 {
            // Malformed data
            return None;
        }

        let mut addr_array = [0; 3];
        addr_array.clone_from_slice(addr_slice);
        let addr = Patch::unserialize_addr_array(addr_array);

        let mut len_array = [0; 2];
        len_array.clone_from_slice(len_slice);
        let len = Patch::unserialize_len(len_array);

        if rest.len() < len {
            // Malformed data
            return None;
        }

        let (patch, rest) = if len == 0 {
            // should decode RLE
            let (rle_len_slice, rest) = rest.split_at(2);

            let mut rle_len_array = [0; 2];
            rle_len_array.clone_from_slice(rle_len_slice);
            let rle_len = Patch::unserialize_len(rle_len_array);

            let (rle_val_slice, rest) = rest.split_at(1);
            let rle_val = rle_val_slice[0];

            let patch = Patch::new(addr, rle_decompress(rle_len, rle_val));
            (patch, rest)
        } else {
            let (data, rest) = rest.split_at(len);

            (Patch::new(addr, data.to_vec()), rest)
        };

        patches.push(patch);

        if rest.len() < 3 {
            // Malformed data
            return None;
        }

        let (possible_eof, _) = rest.split_at(3);
        if possible_eof == b"EOF" {
            // We are done here
            break;
        }

        slice = rest;
    }

    Some(patches)
}

pub fn rle_compress(data: Vec<u8>) -> (usize, u8) {
    (data.len(), data[0])
}

pub fn rle_decompress(len: usize, val: u8) -> Vec<u8> {
    iter::repeat(val).take(len).collect()
}

#[derive(Debug, PartialEq, Clone)]
pub struct Patch {
    addr: u32,
    data: Vec<u8>
}

impl Patch {
    pub fn new(addr: u32, data: Vec<u8>) -> Self {
        assert!(addr <= 0xFFFFFF);
        Patch {
            addr: addr,
            data: data
        }
    }

    pub fn data(&self) -> &Vec<u8> {
        &self.data
    }

    pub fn bytes(&self) -> Vec<u8> {
        let mut bytes: Vec<u8> = self.serialize_addr();
        let len = self.serialize_len();
        bytes.extend(len);
        bytes.extend(self.data());
        bytes
    }

    pub fn unserialize_addr_array(addr: [u8; 3]) -> u32 {
        ((addr[0] as u32) << 16) | ((addr[1] as u32) << 8) | addr[2] as u32
    }

    pub fn serialize_addr(&self) -> Vec<u8> {
        vec![(self.addr >> 16) as u8, (self.addr >> 8) as u8, self.addr as u8]
    }

    pub fn unserialize_len(addr: [u8; 2]) -> usize {
        (((addr[0] as u32) << 8) | addr[1] as u32) as usize
    }

    pub fn serialize_len(&self) -> Vec<u8> {
        vec![(self.data().len() >> 8) as u8, self.data().len() as u8]
    }
}
<|endoftext|>


<commit_before>extern crate rocksdb;
extern crate rustc_serialize;

use std::collections::HashMap;

use self::rustc_serialize::json::{JsonEvent, Parser, StackElement};
// Needed for a trait in order to `put()` into a `rocksdb::WriteBatch`
use self::rocksdb::Writable;

use key_builder::{KeyBuilder, SegmentType};
use records_capnp::payload;
use stems::Stems;

// Good example of using rustc_serialize: https://github.com/ajroetker/beautician/blob/master/src/lib.rs
// Callback based JSON streaming parser: https://github.com/gyscos/json-streamer.rs
// Another parser pased on rustc_serializ: https://github.com/isagalaev/ijson-rust/blob/master/src/test.rs#L11

#[derive(Debug, PartialEq)]
struct WordInfo {
    //offset in the text field where the stemmed text starts
    stemmed_offset: u64,

    // the suffix of the stemmed text. When applied over stemmed, the original
    // text is returned.
    suffix_text: String,

    // the start of the suffixText
    suffix_offset: u64,
}

type ArrayOffsets = Vec<u64>;
type ArrayOffsetsToWordInfo = HashMap<ArrayOffsets, Vec<WordInfo>>;
type WordPathInfoMap = HashMap<String, ArrayOffsetsToWordInfo>;

#[derive(Debug)]
pub struct Shredder {
    keybuilder: KeyBuilder,
    map: WordPathInfoMap,
    path_array_offsets: ArrayOffsets,
}

impl Shredder {
    pub fn new() -> Shredder {
        Shredder{
            keybuilder: KeyBuilder::new(),
            map: WordPathInfoMap::new(),
            path_array_offsets: Vec::new(),
        }
    }
    fn add_entries(&mut self, text: String, docseq: u64) {
        let stems = Stems::new(text.as_str());
        for stem in stems {
            self.keybuilder.push_word(&stem.stemmed);
            self.keybuilder.push_doc_seq(docseq);
            let map_path_array_offsets = self.map.entry(self.keybuilder.key())
                                                        .or_insert(ArrayOffsetsToWordInfo::new());
            let map_word_infos = map_path_array_offsets.entry(self.path_array_offsets.clone())
                .or_insert(Vec::new());
            map_word_infos.push(WordInfo{
                stemmed_offset: stem.stemmed_offset as u64,
                suffix_text: stem.suffix.to_string(),
                suffix_offset: stem.suffix_offset as u64,
            });
            self.keybuilder.pop_doc_seq();
            self.keybuilder.pop_word();
        }
        println!("add_entries: map: {:?}", self.map);
    }

    fn inc_top_array_offset(&mut self) {
        // we encounter a new element. if we are a child element of an array
        // increment the offset. If we aren't (we are the root value or a map
        // value) we don't increment
        if let Some(SegmentType::Array) = self.keybuilder.last_pushed_segment_type() {
            if let Some(last) = self.path_array_offsets.last_mut() {
                *last += 1;
            }
        }
    }

    pub fn shred(&mut self, json: &str, docseq: u64) -> Result<&str, String> {
        println!("{}", json);
        let mut parser = Parser::new(json.chars());
        let mut token = parser.next();

        loop {
            // Get the next token, so that in case of an `ObjectStart` the key is already
            // on the stack.
            let nexttoken = parser.next();

            match token.take() {
                Some(JsonEvent::ObjectStart) => {
                    match parser.stack().top() {
                        Some(StackElement::Key(key)) => {
                            println!("object start: {:?}", key);
                            self.keybuilder.push_object_key(key.to_string());
                            self.inc_top_array_offset();
                        },
                        _ => {
                            panic!("XXX This is probably an object end");
                        }
                    }
                },
                Some(JsonEvent::ObjectEnd) => {
                    self.keybuilder.pop_object_key();
                },
                Some(JsonEvent::ArrayStart) => {
                    println!("array start");
                    self.keybuilder.push_array();
                    //self.inc_top_array_offset();
                    self.path_array_offsets.push(0);
                },
                Some(JsonEvent::ArrayEnd) => {
                    self.path_array_offsets.pop();
                    self.keybuilder.pop_array();
                },
                Some(JsonEvent::StringValue(value)) => {
                    self.add_entries(value, docseq);
                    self.inc_top_array_offset();
                    //self.keybuilder.pop_object_key();
                },
                not_implemented => {
                    panic!("Not yet implemented other JSON types! {:?}", not_implemented);
                }
            };

            token = nexttoken;
            if token == None {
                break;
            }
        }
        println!("keybuilder: {}", self.keybuilder.key());
        println!("shredder: keys:");
        for key in self.map.keys() {
            println!("  {}", key);
        }
        Ok(&"thedocid")
    }

    pub fn add_to_batch(&self, batch: &rocksdb::WriteBatch) -> Result<(), String> {
        for (key_path, word_path_infos) in &self.map {
            let mut message = ::capnp::message::Builder::new_default();
            {
                let capn_payload = message.init_root::<payload::Builder>();
                let mut capn_arrayoffsets_to_wordinfo = capn_payload.init_arrayoffsets_to_wordinfos(
                    word_path_infos.len() as u32);
                for (infos_pos, (arrayoffsets, wordinfos)) in word_path_infos.iter().enumerate() {

                    let mut capn_a2w = capn_arrayoffsets_to_wordinfo.borrow().get(infos_pos as u32);
                    {
                        let mut capn_arrayoffsets =
                            capn_a2w.borrow().init_arrayoffsets(arrayoffsets.len() as u32);
                        for (pos, arrayoffset) in arrayoffsets.iter().enumerate() {
                            capn_arrayoffsets.set(pos as u32, arrayoffset.clone());
                        }
                    }
                    {
                        let mut capn_wordinfos = capn_a2w.init_wordinfos(wordinfos.len() as u32);
                        for (pos, wordinfo) in wordinfos.iter().enumerate() {
                            let mut capn_wordinfo = capn_wordinfos.borrow().get(pos as u32);
                            capn_wordinfo.set_stemmed_offset(wordinfo.stemmed_offset);
                            capn_wordinfo.set_suffix_text(&wordinfo.suffix_text);
                            capn_wordinfo.set_suffix_offset(wordinfo.suffix_offset);
                        }
                    }
                }
            }
            let mut bytes = Vec::new();
            ::capnp::serialize_packed::write_message(&mut bytes, &message).unwrap();
            try!(batch.put(&key_path.clone().into_bytes(), &bytes));
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::{ArrayOffsetsToWordInfo, WordInfo, WordPathInfoMap};

    #[test]
    fn test_shred_nested() {
        let mut shredder = super::Shredder::new();
        //let json = r#"{"hello": {"my": "world!"}, "anumber": 2}"#;
        //let json = r#"{"A":[{"B":"B2VMX two three","C":"C2"},{"B": "b1","C":"C2"}]}"#;
        //let json = r#"{"A":[[[{"B": "string within deeply nested array should be stemmed"}]]]}"#;
        //let json = r#"[{"A": 1, "B": 2, "C": 3}]"#;
        //let json = r#"{"foo": {"bar": 1}}"#;
        let json = r#"{"some": ["array", "data", ["also", "nested"]]}"#;
        let docseq = 123;
        shredder.shred(json, docseq).unwrap();
        let expected = vec![
            ("W.some$!array#123", vec![
                (vec![0], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 5 }])]),
            ("W.some$!data#123", vec![
                (vec![1], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 4 }])]),
            ("W.some$$!also#123", vec![
                (vec![2, 0], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 4 }])]),
            ("W.some$$!nest#123", vec![
                (vec![2, 1], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "ed".to_string(), suffix_offset: 4 }])]),
            ];
        compare_shredded(&shredder.map, &expected);
    }

    #[test]
    fn test_shred_objects() {
        let mut shredder = super::Shredder::new();
        let json = r#"{"A":[{"B":"B2VMX two three","C":"..C2"},{"B": "b1","C":"..C2"}]}"#;
        let docseq = 1234;
        shredder.shred(json, docseq).unwrap();
        let expected = vec![
            ("W.A$.B!b1#1234", vec![
                (vec![0], vec![
                    WordInfo {
                        stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 2 }])]),
            ("W.A$.B!b2vmx#1234", vec![
                (vec![0], vec![
                    WordInfo {
                        stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 5 }])]),
            ("W.A$.B!c2#1234", vec![
                (vec![0], vec![
                    WordInfo {
                        stemmed_offset: 2, suffix_text: "".to_string(), suffix_offset: 4 },
                    WordInfo {
                        stemmed_offset: 2, suffix_text: "".to_string(), suffix_offset: 4 }])]),
            ("W.A$.B!three#1234", vec![
                (vec![0], vec![WordInfo {
                    stemmed_offset: 10, suffix_text: "".to_string(), suffix_offset: 15 }])]),
            ("W.A$.B!two#1234", vec![
                (vec![0], vec![WordInfo {
                    stemmed_offset: 6, suffix_text: "".to_string(), suffix_offset: 9 }])]),
            ];
        compare_shredded(&shredder.map, &expected);
    }

    fn compare_shredded(result_map: &WordPathInfoMap,
                        expected: &Vec<(&str, Vec<(Vec<u64>, Vec<WordInfo>)>)>) {
        // HashMap have an arbitraty order of the elements
        let mut result: Vec<(&String, &ArrayOffsetsToWordInfo)> = result_map.into_iter().collect();
        result.sort_by(|a, b| Ord::cmp(&a.0, &b.0));
        for (ii, &(key, values)) in result.iter().enumerate() {
            assert_eq!(key, expected[ii].0);
            let mut wordinfos: Vec<(&Vec<u64>, &Vec<WordInfo>)> = values.iter().collect();
            wordinfos.sort_by_key(|item| item.0);
            for (jj, wordinfo) in wordinfos.iter().enumerate() {
                assert_eq!(wordinfo.0, &expected[ii].1[jj].0);
                assert_eq!(wordinfo.1, &expected[ii].1[jj].1);
            }
        }
    }
}
<commit_msg>Deal with empty objects properly<commit_after>extern crate rocksdb;
extern crate rustc_serialize;

use std::collections::HashMap;

use self::rustc_serialize::json::{JsonEvent, Parser, StackElement};
// Needed for a trait in order to `put()` into a `rocksdb::WriteBatch`
use self::rocksdb::Writable;

use key_builder::{KeyBuilder, SegmentType};
use records_capnp::payload;
use stems::Stems;

// Good example of using rustc_serialize: https://github.com/ajroetker/beautician/blob/master/src/lib.rs
// Callback based JSON streaming parser: https://github.com/gyscos/json-streamer.rs
// Another parser pased on rustc_serializ: https://github.com/isagalaev/ijson-rust/blob/master/src/test.rs#L11

#[derive(Debug, PartialEq)]
struct WordInfo {
    //offset in the text field where the stemmed text starts
    stemmed_offset: u64,

    // the suffix of the stemmed text. When applied over stemmed, the original
    // text is returned.
    suffix_text: String,

    // the start of the suffixText
    suffix_offset: u64,
}

type ArrayOffsets = Vec<u64>;
type ArrayOffsetsToWordInfo = HashMap<ArrayOffsets, Vec<WordInfo>>;
type WordPathInfoMap = HashMap<String, ArrayOffsetsToWordInfo>;

#[derive(Debug)]
pub struct Shredder {
    keybuilder: KeyBuilder,
    map: WordPathInfoMap,
    path_array_offsets: ArrayOffsets,
}

impl Shredder {
    pub fn new() -> Shredder {
        Shredder{
            keybuilder: KeyBuilder::new(),
            map: WordPathInfoMap::new(),
            path_array_offsets: Vec::new(),
        }
    }
    fn add_entries(&mut self, text: String, docseq: u64) {
        let stems = Stems::new(text.as_str());
        for stem in stems {
            self.keybuilder.push_word(&stem.stemmed);
            self.keybuilder.push_doc_seq(docseq);
            let map_path_array_offsets = self.map.entry(self.keybuilder.key())
                                                        .or_insert(ArrayOffsetsToWordInfo::new());
            let map_word_infos = map_path_array_offsets.entry(self.path_array_offsets.clone())
                .or_insert(Vec::new());
            map_word_infos.push(WordInfo{
                stemmed_offset: stem.stemmed_offset as u64,
                suffix_text: stem.suffix.to_string(),
                suffix_offset: stem.suffix_offset as u64,
            });
            self.keybuilder.pop_doc_seq();
            self.keybuilder.pop_word();
        }
        println!("add_entries: map: {:?}", self.map);
    }

    fn inc_top_array_offset(&mut self) {
        // we encounter a new element. if we are a child element of an array
        // increment the offset. If we aren't (we are the root value or a map
        // value) we don't increment
        if let Some(SegmentType::Array) = self.keybuilder.last_pushed_segment_type() {
            if let Some(last) = self.path_array_offsets.last_mut() {
                *last += 1;
            }
        }
    }

    pub fn shred(&mut self, json: &str, docseq: u64) -> Result<&str, String> {
        let mut parser = Parser::new(json.chars());
        let mut token = parser.next();

        loop {
            // Get the next token, so that in case of an `ObjectStart` the key is already
            // on the stack.
            let nexttoken = parser.next();

            match token.take() {
                Some(JsonEvent::ObjectStart) => {
                    match parser.stack().top() {
                        Some(StackElement::Key(key)) => {
                            self.keybuilder.push_object_key(key.to_string());
                            self.inc_top_array_offset();
                        },
                        // We won't hit this case as we are within an object and not within
                        // an array
                        Some(StackElement::Index(_)) => {}
                        // It's an empty object
                        None => {
                            // Just push something to make `ObjectEnd` happy
                            self.keybuilder.push_object_key("".to_string());
                        },
                    }
                },
                Some(JsonEvent::ObjectEnd) => {
                    self.keybuilder.pop_object_key();
                },
                Some(JsonEvent::ArrayStart) => {
                    self.keybuilder.push_array();
                    self.path_array_offsets.push(0);
                },
                Some(JsonEvent::ArrayEnd) => {
                    self.path_array_offsets.pop();
                    self.keybuilder.pop_array();
                },
                Some(JsonEvent::StringValue(value)) => {
                    self.add_entries(value, docseq);
                    self.inc_top_array_offset();
                },
                not_implemented => {
                    panic!("Not yet implemented other JSON types! {:?}", not_implemented);
                }
            };

            token = nexttoken;
            if token == None {
                break;
            }
        }
        println!("keybuilder: {}", self.keybuilder.key());
        println!("shredder: keys:");
        for key in self.map.keys() {
            println!("  {}", key);
        }
        Ok(&"thedocid")
    }

    pub fn add_to_batch(&self, batch: &rocksdb::WriteBatch) -> Result<(), String> {
        for (key_path, word_path_infos) in &self.map {
            let mut message = ::capnp::message::Builder::new_default();
            {
                let capn_payload = message.init_root::<payload::Builder>();
                let mut capn_arrayoffsets_to_wordinfo = capn_payload.init_arrayoffsets_to_wordinfos(
                    word_path_infos.len() as u32);
                for (infos_pos, (arrayoffsets, wordinfos)) in word_path_infos.iter().enumerate() {

                    let mut capn_a2w = capn_arrayoffsets_to_wordinfo.borrow().get(infos_pos as u32);
                    {
                        let mut capn_arrayoffsets =
                            capn_a2w.borrow().init_arrayoffsets(arrayoffsets.len() as u32);
                        for (pos, arrayoffset) in arrayoffsets.iter().enumerate() {
                            capn_arrayoffsets.set(pos as u32, arrayoffset.clone());
                        }
                    }
                    {
                        let mut capn_wordinfos = capn_a2w.init_wordinfos(wordinfos.len() as u32);
                        for (pos, wordinfo) in wordinfos.iter().enumerate() {
                            let mut capn_wordinfo = capn_wordinfos.borrow().get(pos as u32);
                            capn_wordinfo.set_stemmed_offset(wordinfo.stemmed_offset);
                            capn_wordinfo.set_suffix_text(&wordinfo.suffix_text);
                            capn_wordinfo.set_suffix_offset(wordinfo.suffix_offset);
                        }
                    }
                }
            }
            let mut bytes = Vec::new();
            ::capnp::serialize_packed::write_message(&mut bytes, &message).unwrap();
            try!(batch.put(&key_path.clone().into_bytes(), &bytes));
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::{ArrayOffsetsToWordInfo, WordInfo, WordPathInfoMap};

    #[test]
    fn test_shred_nested() {
        let mut shredder = super::Shredder::new();
        //let json = r#"{"hello": {"my": "world!"}, "anumber": 2}"#;
        //let json = r#"{"A":[{"B":"B2VMX two three","C":"C2"},{"B": "b1","C":"C2"}]}"#;
        //let json = r#"{"A":[[[{"B": "string within deeply nested array should be stemmed"}]]]}"#;
        //let json = r#"[{"A": 1, "B": 2, "C": 3}]"#;
        //let json = r#"{"foo": {"bar": 1}}"#;
        let json = r#"{"some": ["array", "data", ["also", "nested"]]}"#;
        let docseq = 123;
        shredder.shred(json, docseq).unwrap();
        let expected = vec![
            ("W.some$!array#123", vec![
                (vec![0], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 5 }])]),
            ("W.some$!data#123", vec![
                (vec![1], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 4 }])]),
            ("W.some$$!also#123", vec![
                (vec![2, 0], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 4 }])]),
            ("W.some$$!nest#123", vec![
                (vec![2, 1], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "ed".to_string(), suffix_offset: 4 }])]),
            ];
        compare_shredded(&shredder.map, &expected);
    }

    #[test]
    fn test_shred_objects() {
        let mut shredder = super::Shredder::new();
        let json = r#"{"A":[{"B":"B2VMX two three","C":"..C2"},{"B": "b1","C":"..C2"}]}"#;
        let docseq = 1234;
        shredder.shred(json, docseq).unwrap();
        let expected = vec![
            ("W.A$.B!b1#1234", vec![
                (vec![0], vec![
                    WordInfo {
                        stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 2 }])]),
            ("W.A$.B!b2vmx#1234", vec![
                (vec![0], vec![
                    WordInfo {
                        stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 5 }])]),
            ("W.A$.B!c2#1234", vec![
                (vec![0], vec![
                    WordInfo {
                        stemmed_offset: 2, suffix_text: "".to_string(), suffix_offset: 4 },
                    WordInfo {
                        stemmed_offset: 2, suffix_text: "".to_string(), suffix_offset: 4 }])]),
            ("W.A$.B!three#1234", vec![
                (vec![0], vec![WordInfo {
                    stemmed_offset: 10, suffix_text: "".to_string(), suffix_offset: 15 }])]),
            ("W.A$.B!two#1234", vec![
                (vec![0], vec![WordInfo {
                    stemmed_offset: 6, suffix_text: "".to_string(), suffix_offset: 9 }])]),
            ];
        compare_shredded(&shredder.map, &expected);
    }

    fn compare_shredded(result_map: &WordPathInfoMap,
                        expected: &Vec<(&str, Vec<(Vec<u64>, Vec<WordInfo>)>)>) {
        // HashMap have an arbitraty order of the elements
        let mut result: Vec<(&String, &ArrayOffsetsToWordInfo)> = result_map.into_iter().collect();
        result.sort_by(|a, b| Ord::cmp(&a.0, &b.0));
        for (ii, &(key, values)) in result.iter().enumerate() {
            assert_eq!(key, expected[ii].0);
            let mut wordinfos: Vec<(&Vec<u64>, &Vec<WordInfo>)> = values.iter().collect();
            wordinfos.sort_by_key(|item| item.0);
            for (jj, wordinfo) in wordinfos.iter().enumerate() {
                assert_eq!(wordinfo.0, &expected[ii].1[jj].0);
                assert_eq!(wordinfo.1, &expected[ii].1[jj].1);
            }
        }
    }

    #[test]
    fn test_shred_empty_object() {
        let mut shredder = super::Shredder::new();
        let json = r#"{}"#;
        let docseq = 123;
        shredder.shred(json, docseq).unwrap();
        // TODO vmx 2016-11-08: Finish this test
    }
}
<|endoftext|>


<commit_before>use error::Error;
use num_traits::NumCast;
use serde::de::{self, Visitor};
use serde::{Serialize, Serializer, Deserialize, Deserializer};
use std::{fmt, i64};

/// Represents a JSON number, whether integer or floating point.
#[derive(Clone, PartialEq)]
pub struct Number(N);

// "N" is a prefix of "NegInt"... this is a false positive.
// https://github.com/Manishearth/rust-clippy/issues/1241
#[cfg_attr(feature = "clippy", allow(pub_enum_variant_names))]
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum N {
    PosInt(u64),
    /// Always less than zero.
    NegInt(i64),
    /// Always finite.
    Float(f64),
}

impl Number {
    /// Returns true if the number can be represented by i64.
    #[inline]
    pub fn is_i64(&self) -> bool {
        match self.0 {
            N::PosInt(v) => v <= i64::MAX as u64,
            N::NegInt(_) => true,
            N::Float(_) => false,
        }
    }

    /// Returns true if the number can be represented as u64.
    #[inline]
    pub fn is_u64(&self) -> bool {
        match self.0 {
            N::PosInt(_) => true,
            N::NegInt(_) | N::Float(_) => false,
        }
    }

    /// Returns true if the number can be represented as f64.
    #[inline]
    pub fn is_f64(&self) -> bool {
        match self.0 {
            N::Float(_) => true,
            N::PosInt(_) | N::NegInt(_) => false,
        }
    }

    /// Returns the number represented as i64 if possible, or else None.
    #[inline]
    pub fn as_i64(&self) -> Option<i64> {
        match self.0 {
            N::PosInt(n) => NumCast::from(n),
            N::NegInt(n) => Some(n),
            N::Float(_) => None,
        }
    }

    /// Returns the number represented as u64 if possible, or else None.
    #[inline]
    pub fn as_u64(&self) -> Option<u64> {
        match self.0 {
            N::PosInt(n) => Some(n),
            N::NegInt(n) => NumCast::from(n),
            N::Float(_) => None,
        }
    }

    /// Returns the number represented as f64 if possible, or else None.
    #[inline]
    pub fn as_f64(&self) -> Option<f64> {
        match self.0 {
            N::PosInt(n) => NumCast::from(n),
            N::NegInt(n) => NumCast::from(n),
            N::Float(n) => Some(n),
        }
    }

    /// Converts a finite f64 to a Number. Infinite or NaN values are not JSON
    /// numbers.
    #[inline]
    pub fn from_f64(f: f64) -> Option<Number> {
        if f.is_finite() {
            Some(Number(N::Float(f)))
        } else {
            None
        }
    }
}

impl fmt::Display for Number {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        match self.0 {
            N::PosInt(i) => i.fmt(formatter),
            N::NegInt(i) => i.fmt(formatter),
            N::Float(f) => f.fmt(formatter),
        }
    }
}

impl Serialize for Number {
    #[inline]
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where S: Serializer
    {
        match self.0 {
            N::PosInt(i) => serializer.serialize_u64(i),
            N::NegInt(i) => serializer.serialize_i64(i),
            N::Float(f) => serializer.serialize_f64(f),
        }
    }
}

impl Deserialize for Number {
    #[inline]
    fn deserialize<D>(deserializer: D) -> Result<Number, D::Error>
        where D: Deserializer
    {
        struct NumberVisitor;

        impl Visitor for NumberVisitor {
            type Value = Number;

            #[inline]
            fn visit_i64<E>(self, value: i64) -> Result<Number, E> {
                Ok(value.into())
            }

            #[inline]
            fn visit_u64<E>(self, value: u64) -> Result<Number, E> {
                Ok(value.into())
            }

            #[inline]
            fn visit_f64<E>(self, value: f64) -> Result<Number, E>
                where E: de::Error
            {
                Number::from_f64(value).ok_or_else(|| de::Error::custom("not a JSON number"))
            }
        }

        deserializer.deserialize(NumberVisitor)
    }
}

impl Deserializer for Number {
    type Error = Error;

    #[inline]
    fn deserialize<V>(self, visitor: V) -> Result<V::Value, Error>
        where V: Visitor
    {
        match self.0 {
            N::PosInt(i) => visitor.visit_u64(i),
            N::NegInt(i) => visitor.visit_i64(i),
            N::Float(f) => visitor.visit_f64(f),
        }
    }

    forward_to_deserialize! {
        bool usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 char str string
        unit option seq seq_fixed_size bytes map unit_struct newtype_struct
        tuple_struct struct struct_field tuple enum ignored_any
    }
}

macro_rules! from_signed {
    ($($signed_ty:ident)*) => {
        $(
            impl From<$signed_ty> for Number {
                #[inline]
                fn from(i: $signed_ty) -> Self {
                    if i < 0 {
                        Number(N::NegInt(i as i64))
                    } else {
                        Number(N::PosInt(i as u64))
                    }
                }
            }
        )*
    };
}

macro_rules! from_unsigned {
    ($($unsigned_ty:ident)*) => {
        $(
            impl From<$unsigned_ty> for Number {
                #[inline]
                fn from(u: $unsigned_ty) -> Self {
                    Number(N::PosInt(u as u64))
                }
            }
        )*
    };
}

from_signed!(i8 i16 i32 i64 isize);
from_unsigned!(u8 u16 u32 u64 usize);
<commit_msg>Private N<commit_after>use error::Error;
use num_traits::NumCast;
use serde::de::{self, Visitor};
use serde::{Serialize, Serializer, Deserialize, Deserializer};
use std::{fmt, i64};

/// Represents a JSON number, whether integer or floating point.
#[derive(Clone, PartialEq)]
pub struct Number(N);

// "N" is a prefix of "NegInt"... this is a false positive.
// https://github.com/Manishearth/rust-clippy/issues/1241
#[cfg_attr(feature = "clippy", allow(pub_enum_variant_names))]
#[derive(Copy, Clone, Debug, PartialEq)]
enum N {
    PosInt(u64),
    /// Always less than zero.
    NegInt(i64),
    /// Always finite.
    Float(f64),
}

impl Number {
    /// Returns true if the number can be represented by i64.
    #[inline]
    pub fn is_i64(&self) -> bool {
        match self.0 {
            N::PosInt(v) => v <= i64::MAX as u64,
            N::NegInt(_) => true,
            N::Float(_) => false,
        }
    }

    /// Returns true if the number can be represented as u64.
    #[inline]
    pub fn is_u64(&self) -> bool {
        match self.0 {
            N::PosInt(_) => true,
            N::NegInt(_) | N::Float(_) => false,
        }
    }

    /// Returns true if the number can be represented as f64.
    #[inline]
    pub fn is_f64(&self) -> bool {
        match self.0 {
            N::Float(_) => true,
            N::PosInt(_) | N::NegInt(_) => false,
        }
    }

    /// Returns the number represented as i64 if possible, or else None.
    #[inline]
    pub fn as_i64(&self) -> Option<i64> {
        match self.0 {
            N::PosInt(n) => NumCast::from(n),
            N::NegInt(n) => Some(n),
            N::Float(_) => None,
        }
    }

    /// Returns the number represented as u64 if possible, or else None.
    #[inline]
    pub fn as_u64(&self) -> Option<u64> {
        match self.0 {
            N::PosInt(n) => Some(n),
            N::NegInt(n) => NumCast::from(n),
            N::Float(_) => None,
        }
    }

    /// Returns the number represented as f64 if possible, or else None.
    #[inline]
    pub fn as_f64(&self) -> Option<f64> {
        match self.0 {
            N::PosInt(n) => NumCast::from(n),
            N::NegInt(n) => NumCast::from(n),
            N::Float(n) => Some(n),
        }
    }

    /// Converts a finite f64 to a Number. Infinite or NaN values are not JSON
    /// numbers.
    #[inline]
    pub fn from_f64(f: f64) -> Option<Number> {
        if f.is_finite() {
            Some(Number(N::Float(f)))
        } else {
            None
        }
    }
}

impl fmt::Display for Number {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        match self.0 {
            N::PosInt(i) => i.fmt(formatter),
            N::NegInt(i) => i.fmt(formatter),
            N::Float(f) => f.fmt(formatter),
        }
    }
}

impl Serialize for Number {
    #[inline]
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where S: Serializer
    {
        match self.0 {
            N::PosInt(i) => serializer.serialize_u64(i),
            N::NegInt(i) => serializer.serialize_i64(i),
            N::Float(f) => serializer.serialize_f64(f),
        }
    }
}

impl Deserialize for Number {
    #[inline]
    fn deserialize<D>(deserializer: D) -> Result<Number, D::Error>
        where D: Deserializer
    {
        struct NumberVisitor;

        impl Visitor for NumberVisitor {
            type Value = Number;

            #[inline]
            fn visit_i64<E>(self, value: i64) -> Result<Number, E> {
                Ok(value.into())
            }

            #[inline]
            fn visit_u64<E>(self, value: u64) -> Result<Number, E> {
                Ok(value.into())
            }

            #[inline]
            fn visit_f64<E>(self, value: f64) -> Result<Number, E>
                where E: de::Error
            {
                Number::from_f64(value).ok_or_else(|| de::Error::custom("not a JSON number"))
            }
        }

        deserializer.deserialize(NumberVisitor)
    }
}

impl Deserializer for Number {
    type Error = Error;

    #[inline]
    fn deserialize<V>(self, visitor: V) -> Result<V::Value, Error>
        where V: Visitor
    {
        match self.0 {
            N::PosInt(i) => visitor.visit_u64(i),
            N::NegInt(i) => visitor.visit_i64(i),
            N::Float(f) => visitor.visit_f64(f),
        }
    }

    forward_to_deserialize! {
        bool usize u8 u16 u32 u64 isize i8 i16 i32 i64 f32 f64 char str string
        unit option seq seq_fixed_size bytes map unit_struct newtype_struct
        tuple_struct struct struct_field tuple enum ignored_any
    }
}

macro_rules! from_signed {
    ($($signed_ty:ident)*) => {
        $(
            impl From<$signed_ty> for Number {
                #[inline]
                fn from(i: $signed_ty) -> Self {
                    if i < 0 {
                        Number(N::NegInt(i as i64))
                    } else {
                        Number(N::PosInt(i as u64))
                    }
                }
            }
        )*
    };
}

macro_rules! from_unsigned {
    ($($unsigned_ty:ident)*) => {
        $(
            impl From<$unsigned_ty> for Number {
                #[inline]
                fn from(u: $unsigned_ty) -> Self {
                    Number(N::PosInt(u as u64))
                }
            }
        )*
    };
}

from_signed!(i8 i16 i32 i64 isize);
from_unsigned!(u8 u16 u32 u64 usize);
<|endoftext|>


<commit_before>// Copyright 2018 Doug Goldstein
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use large_example::bar;
use log::*;
use structopt::StructOpt;

/// A basic example
#[derive(StructOpt, Debug)]
#[structopt()]
struct Opt {
    /// Silence all output
    #[structopt(short = "q", long = "quiet")]
    quiet: bool,
    /// Verbose mode (-v, -vv, -vvv, etc.)
    #[structopt(short = "v", long = "verbose", parse(from_occurrences))]
    verbose: usize,
    /// Allow module to log
    #[structopt(short = "m", long = "module")]
    modules: Vec<String>,
    /// Timestamp (sec, ms, us, ns, none)
    #[structopt(short = "t", long = "timestamp")]
    ts: Option<stderrlog::Timestamp>,
}

fn main() {
    let opt = Opt::from_args();

    stderrlog::new()
        .module(module_path!())
        .quiet(opt.quiet)
        .verbosity(opt.verbose)
        .timestamp(opt.ts.unwrap_or(stderrlog::Timestamp::Off))
        .modules(opt.modules)
        .init()
        .unwrap();

    large_example::libfn();
    bar::barfn();

    println!("back to main()");
    trace!("{} - trace message", module_path!());
    debug!("{} - debug message", module_path!());
    info!("{} - info message", module_path!());
    warn!("{} - warn message", module_path!());
    error!("{} - error message", module_path!());

    println!("Some modules to try:");
    println!("  -m large_example");
    println!("  -m large_example::foo");
    println!("  -m large_example::bar");
}
<commit_msg>fix example to allow for filtering<commit_after>// Copyright 2018 Doug Goldstein
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use large_example::bar;
use log::*;
use structopt::StructOpt;

/// A basic example
#[derive(StructOpt, Debug)]
#[structopt()]
struct Opt {
    /// Silence all output
    #[structopt(short = "q", long = "quiet")]
    quiet: bool,
    /// Verbose mode (-v, -vv, -vvv, etc.)
    #[structopt(short = "v", long = "verbose", parse(from_occurrences))]
    verbose: usize,
    /// Allow module to log
    #[structopt(short = "m", long = "module")]
    modules: Vec<String>,
    /// Timestamp (sec, ms, us, ns, none)
    #[structopt(short = "t", long = "timestamp")]
    ts: Option<stderrlog::Timestamp>,
}

fn main() {
    let opt = Opt::from_args();

    stderrlog::new()
        .quiet(opt.quiet)
        .verbosity(opt.verbose)
        .timestamp(opt.ts.unwrap_or(stderrlog::Timestamp::Off))
        .modules(opt.modules)
        .init()
        .unwrap();

    large_example::libfn();
    bar::barfn();

    println!("back to main()");
    trace!("{} - trace message", module_path!());
    debug!("{} - debug message", module_path!());
    info!("{} - info message", module_path!());
    warn!("{} - warn message", module_path!());
    error!("{} - error message", module_path!());

    println!("Some modules to try:");
    println!("  -m large_example");
    println!("  -m large_example::foo");
    println!("  -m large_example::bar");
}
<|endoftext|>


<commit_before>//
// Sysinfo
//
// Copyright (c) 2017 Guillaume Gomez
//

use crate::sys::utils::get_all_data;
use crate::{utils, DiskExt, DiskType};

use libc::statvfs;
use std::ffi::{OsStr, OsString};
use std::fs;
use std::mem;
use std::os::unix::ffi::OsStrExt;
use std::path::{Path, PathBuf};

macro_rules! cast {
    ($x:expr) => {
        u64::from($x)
    };
}

/// Struct containing a disk information.
#[derive(PartialEq)]
pub struct Disk {
    type_: DiskType,
    device_name: OsString,
    file_system: Vec<u8>,
    mount_point: PathBuf,
    total_space: u64,
    available_space: u64,
    is_removable: bool,
}

impl DiskExt for Disk {
    fn type_(&self) -> DiskType {
        self.type_
    }

    fn name(&self) -> &OsStr {
        &self.device_name
    }

    fn file_system(&self) -> &[u8] {
        &self.file_system
    }

    fn mount_point(&self) -> &Path {
        &self.mount_point
    }

    fn total_space(&self) -> u64 {
        self.total_space
    }

    fn available_space(&self) -> u64 {
        self.available_space
    }

    fn is_removable(&self) -> bool {
        self.is_removable
    }

    fn refresh(&mut self) -> bool {
        unsafe {
            let mut stat: statvfs = mem::zeroed();
            let mount_point_cpath = utils::to_cpath(&self.mount_point);
            if statvfs(mount_point_cpath.as_ptr() as *const _, &mut stat) == 0 {
                let tmp = cast!(stat.f_bsize) * cast!(stat.f_bavail);
                self.available_space = cast!(tmp);
                true
            } else {
                false
            }
        }
    }
}

fn new_disk(
    device_name: &OsStr,
    mount_point: &Path,
    file_system: &[u8],
    removable_entries: &[PathBuf],
) -> Option<Disk> {
    let mount_point_cpath = utils::to_cpath(mount_point);
    let type_ = find_type_for_device_name(device_name);
    let mut total = 0;
    let mut available = 0;
    unsafe {
        let mut stat: statvfs = mem::zeroed();
        if statvfs(mount_point_cpath.as_ptr() as *const _, &mut stat) == 0 {
            total = cast!(stat.f_bsize) * cast!(stat.f_blocks);
            available = cast!(stat.f_bsize) * cast!(stat.f_bavail);
        }
    }
    if total == 0 {
        return None;
    }
    let mount_point = mount_point.to_owned();
    let is_removable = removable_entries
        .iter()
        .any(|e| e.as_os_str() == device_name);
    Some(Disk {
        type_,
        device_name: device_name.to_owned(),
        file_system: file_system.to_owned(),
        mount_point,
        total_space: cast!(total),
        available_space: cast!(available),
        is_removable,
    })
}

#[allow(clippy::manual_range_contains)]
fn find_type_for_device_name(device_name: &OsStr) -> DiskType {
    // The format of devices are as follows:
    //  - device_name is symbolic link in the case of /dev/mapper/
    //     and /dev/root, and the target is corresponding device under
    //     /sys/block/
    //  - In the case of /dev/sd, the format is /dev/sd[a-z][1-9],
    //     corresponding to /sys/block/sd[a-z]
    //  - In the case of /dev/nvme, the format is /dev/nvme[0-9]n[0-9]p[0-9],
    //     corresponding to /sys/block/nvme[0-9]n[0-9]
    //  - In the case of /dev/mmcblk, the format is /dev/mmcblk[0-9]p[0-9],
    //     corresponding to /sys/block/mmcblk[0-9]
    let device_name_path = device_name.to_str().unwrap_or_default();
    let real_path = fs::canonicalize(device_name).unwrap_or_else(|_| PathBuf::from(device_name));
    let mut real_path = real_path.to_str().unwrap_or_default();
    if device_name_path.starts_with("/dev/mapper/") {
        // Recursively solve, for example /dev/dm-0
        if real_path != device_name_path {
            return find_type_for_device_name(OsStr::new(&real_path));
        }
    } else if device_name_path.starts_with("/dev/sd") || device_name_path.starts_with("/dev/vd") {
        // Turn "sda1" into "sda" or "vda1" into "vda"
        real_path = real_path.trim_start_matches("/dev/");
        real_path = real_path.trim_end_matches(|c| c >= '0' && c <= '9');
    } else if device_name_path.starts_with("/dev/nvme") {
        // Turn "nvme0n1p1" into "nvme0n1"
        real_path = real_path.trim_start_matches("/dev/");
        real_path = real_path.trim_end_matches(|c| c >= '0' && c <= '9');
        real_path = real_path.trim_end_matches(|c| c == 'p');
    } else if device_name_path.starts_with("/dev/root") {
        // Recursively solve, for example /dev/mmcblk0p1
        if real_path != device_name_path {
            return find_type_for_device_name(OsStr::new(&real_path));
        }
    } else if device_name_path.starts_with("/dev/mmcblk") {
        // Turn "mmcblk0p1" into "mmcblk0"
        real_path = real_path.trim_start_matches("/dev/");
        real_path = real_path.trim_end_matches(|c| c >= '0' && c <= '9');
        real_path = real_path.trim_end_matches(|c| c == 'p');
    } else {
        // Default case: remove /dev/ and expects the name presents under /sys/block/
        // For example, /dev/dm-0 to dm-0
        real_path = real_path.trim_start_matches("/dev/");
    }

    let trimmed: &OsStr = OsStrExt::from_bytes(real_path.as_bytes());

    let path = Path::new("/sys/block/")
        .to_owned()
        .join(trimmed)
        .join("queue/rotational");
    // Normally, this file only contains '0' or '1' but just in case, we get 8 bytes...
    match get_all_data(path, 8)
        .unwrap_or_default()
        .trim()
        .parse()
        .ok()
    {
        // The disk is marked as rotational so it's a HDD.
        Some(1) => DiskType::HDD,
        // The disk is marked as non-rotational so it's very likely a SSD.
        Some(0) => DiskType::SSD,
        // Normally it shouldn't happen but welcome to the wonderful world of IT! :D
        Some(x) => DiskType::Unknown(x),
        // The information isn't available...
        None => DiskType::Unknown(-1),
    }
}

fn get_all_disks_inner(content: &str) -> Vec<Disk> {
    // The goal of this array is to list all removable devices (the ones whose name starts with
    // "usb-"). Then we check if
    let removable_entries = match fs::read_dir("/dev/disk/by-id/") {
        Ok(r) => r
            .filter_map(|res| Some(res.ok()?.path()))
            .filter_map(|e| {
                if e.file_name()
                    .and_then(|x| Some(x.to_str()?.starts_with("usb-")))
                    .unwrap_or_default()
                {
                    e.canonicalize().ok()
                } else {
                    None
                }
            })
            .collect::<Vec<PathBuf>>(),
        _ => Vec::new(),
    };

    content
        .lines()
        .map(|line| {
            let line = line.trim_start();
            // mounts format
            // http://man7.org/linux/man-pages/man5/fstab.5.html
            // fs_spec<tab>fs_file<tab>fs_vfstype<tab>other fields
            let mut fields = line.split_whitespace();
            let fs_spec = fields.next().unwrap_or("");
            let fs_file = fields
                .next()
                .unwrap_or("")
                .replace("\\134", "\\")
                .replace("\\040", " ")
                .replace("\\011", "\t")
                .replace("\\012", "\n");
            let fs_vfstype = fields.next().unwrap_or("");
            (fs_spec, fs_file, fs_vfstype)
        })
        .filter(|(fs_spec, fs_file, fs_vfstype)| {
            // Check if fs_vfstype is one of our 'ignored' file systems.
            let filtered = matches!(
                *fs_vfstype,
                "sysfs" | // pseudo file system for kernel objects
                "proc" |  // another pseudo file system
                "tmpfs" |
                "devtmpfs" |
                "cgroup" |
                "cgroup2" |
                "pstore" | // https://www.kernel.org/doc/Documentation/ABI/testing/pstore
                "squashfs" | // squashfs is a compressed read-only file system (for snaps)
                "rpc_pipefs" | // The pipefs pseudo file system service
                "iso9660" // optical media
            );

            !(filtered ||
               fs_file.starts_with("/sys") || // check if fs_file is an 'ignored' mount point
               fs_file.starts_with("/proc") ||
               (fs_file.starts_with("/run") && !fs_file.starts_with("/run/media")) ||
               fs_spec.starts_with("sunrpc"))
        })
        .filter_map(|(fs_spec, fs_file, fs_vfstype)| {
            new_disk(
                fs_spec.as_ref(),
                Path::new(&fs_file),
                fs_vfstype.as_bytes(),
                &removable_entries,
            )
        })
        .collect()
}

pub fn get_all_disks() -> Vec<Disk> {
    get_all_disks_inner(&get_all_data("/proc/mounts", 16_385).unwrap_or_default())
}

// #[test]
// fn check_all_disks() {
//     let disks = get_all_disks_inner(
//         r#"tmpfs /proc tmpfs rw,seclabel,relatime 0 0
// proc /proc proc rw,nosuid,nodev,noexec,relatime 0 0
// systemd-1 /proc/sys/fs/binfmt_misc autofs rw,relatime,fd=29,pgrp=1,timeout=0,minproto=5,maxproto=5,direct,pipe_ino=17771 0 0
// tmpfs /sys tmpfs rw,seclabel,relatime 0 0
// sysfs /sys sysfs rw,seclabel,nosuid,nodev,noexec,relatime 0 0
// securityfs /sys/kernel/security securityfs rw,nosuid,nodev,noexec,relatime 0 0
// cgroup2 /sys/fs/cgroup cgroup2 rw,seclabel,nosuid,nodev,noexec,relatime,nsdelegate 0 0
// pstore /sys/fs/pstore pstore rw,seclabel,nosuid,nodev,noexec,relatime 0 0
// none /sys/fs/bpf bpf rw,nosuid,nodev,noexec,relatime,mode=700 0 0
// configfs /sys/kernel/config configfs rw,nosuid,nodev,noexec,relatime 0 0
// selinuxfs /sys/fs/selinux selinuxfs rw,relatime 0 0
// debugfs /sys/kernel/debug debugfs rw,seclabel,nosuid,nodev,noexec,relatime 0 0
// tmpfs /dev/shm tmpfs rw,seclabel,relatime 0 0
// devpts /dev/pts devpts rw,seclabel,relatime,gid=5,mode=620,ptmxmode=666 0 0
// tmpfs /sys/fs/selinux tmpfs rw,seclabel,relatime 0 0
// /dev/vda2 /proc/filesystems xfs rw,seclabel,relatime,attr2,inode64,logbufs=8,logbsize=32k,noquota 0 0
// "#,
//     );
//     assert_eq!(disks.len(), 1);
//     assert_eq!(
//         disks[0],
//         Disk {
//             type_: DiskType::Unknown(-1),
//             name: OsString::from("devpts"),
//             file_system: vec![100, 101, 118, 112, 116, 115],
//             mount_point: PathBuf::from("/dev/pts"),
//             total_space: 0,
//             available_space: 0,
//         }
//     );
// }
<commit_msg>Use Wrapping on the stat numbers due to a potential overflow in the following multiplications<commit_after>//
// Sysinfo
//
// Copyright (c) 2017 Guillaume Gomez
//

use crate::sys::utils::get_all_data;
use crate::{utils, DiskExt, DiskType};

use libc::statvfs;
use std::ffi::{OsStr, OsString};
use std::fs;
use std::mem;
use std::num::Wrapping;
use std::os::unix::ffi::OsStrExt;
use std::path::{Path, PathBuf};

macro_rules! cast {
    ($x:expr) => {
        u64::from($x)
    };
}

/// Struct containing a disk information.
#[derive(PartialEq)]
pub struct Disk {
    type_: DiskType,
    device_name: OsString,
    file_system: Vec<u8>,
    mount_point: PathBuf,
    total_space: u64,
    available_space: u64,
    is_removable: bool,
}

impl DiskExt for Disk {
    fn type_(&self) -> DiskType {
        self.type_
    }

    fn name(&self) -> &OsStr {
        &self.device_name
    }

    fn file_system(&self) -> &[u8] {
        &self.file_system
    }

    fn mount_point(&self) -> &Path {
        &self.mount_point
    }

    fn total_space(&self) -> u64 {
        self.total_space
    }

    fn available_space(&self) -> u64 {
        self.available_space
    }

    fn is_removable(&self) -> bool {
        self.is_removable
    }

    fn refresh(&mut self) -> bool {
        unsafe {
            let mut stat: statvfs = mem::zeroed();
            let mount_point_cpath = utils::to_cpath(&self.mount_point);
            if statvfs(mount_point_cpath.as_ptr() as *const _, &mut stat) == 0 {
                let tmp = cast!(stat.f_bsize) * cast!(stat.f_bavail);
                self.available_space = cast!(tmp);
                true
            } else {
                false
            }
        }
    }
}

fn new_disk(
    device_name: &OsStr,
    mount_point: &Path,
    file_system: &[u8],
    removable_entries: &[PathBuf],
) -> Option<Disk> {
    let mount_point_cpath = utils::to_cpath(mount_point);
    let type_ = find_type_for_device_name(device_name);
    let mut total = 0;
    let mut available = 0;
    unsafe {
        let mut stat: statvfs = mem::zeroed();
        if statvfs(mount_point_cpath.as_ptr() as *const _, &mut stat) == 0 {
            let bsize = Wrapping(cast!(stat.f_bsize));
            let blocks = Wrapping(cast!(stat.f_blocks));
            let bavail = Wrapping(cast!(stat.f_bavail));
            total = (bsize * blocks).0;
            available = (bsize * bavail).0;
        }
    }
    if total == 0 {
        return None;
    }
    let mount_point = mount_point.to_owned();
    let is_removable = removable_entries
        .iter()
        .any(|e| e.as_os_str() == device_name);
    Some(Disk {
        type_,
        device_name: device_name.to_owned(),
        file_system: file_system.to_owned(),
        mount_point,
        total_space: cast!(total),
        available_space: cast!(available),
        is_removable,
    })
}

#[allow(clippy::manual_range_contains)]
fn find_type_for_device_name(device_name: &OsStr) -> DiskType {
    // The format of devices are as follows:
    //  - device_name is symbolic link in the case of /dev/mapper/
    //     and /dev/root, and the target is corresponding device under
    //     /sys/block/
    //  - In the case of /dev/sd, the format is /dev/sd[a-z][1-9],
    //     corresponding to /sys/block/sd[a-z]
    //  - In the case of /dev/nvme, the format is /dev/nvme[0-9]n[0-9]p[0-9],
    //     corresponding to /sys/block/nvme[0-9]n[0-9]
    //  - In the case of /dev/mmcblk, the format is /dev/mmcblk[0-9]p[0-9],
    //     corresponding to /sys/block/mmcblk[0-9]
    let device_name_path = device_name.to_str().unwrap_or_default();
    let real_path = fs::canonicalize(device_name).unwrap_or_else(|_| PathBuf::from(device_name));
    let mut real_path = real_path.to_str().unwrap_or_default();
    if device_name_path.starts_with("/dev/mapper/") {
        // Recursively solve, for example /dev/dm-0
        if real_path != device_name_path {
            return find_type_for_device_name(OsStr::new(&real_path));
        }
    } else if device_name_path.starts_with("/dev/sd") || device_name_path.starts_with("/dev/vd") {
        // Turn "sda1" into "sda" or "vda1" into "vda"
        real_path = real_path.trim_start_matches("/dev/");
        real_path = real_path.trim_end_matches(|c| c >= '0' && c <= '9');
    } else if device_name_path.starts_with("/dev/nvme") {
        // Turn "nvme0n1p1" into "nvme0n1"
        real_path = real_path.trim_start_matches("/dev/");
        real_path = real_path.trim_end_matches(|c| c >= '0' && c <= '9');
        real_path = real_path.trim_end_matches(|c| c == 'p');
    } else if device_name_path.starts_with("/dev/root") {
        // Recursively solve, for example /dev/mmcblk0p1
        if real_path != device_name_path {
            return find_type_for_device_name(OsStr::new(&real_path));
        }
    } else if device_name_path.starts_with("/dev/mmcblk") {
        // Turn "mmcblk0p1" into "mmcblk0"
        real_path = real_path.trim_start_matches("/dev/");
        real_path = real_path.trim_end_matches(|c| c >= '0' && c <= '9');
        real_path = real_path.trim_end_matches(|c| c == 'p');
    } else {
        // Default case: remove /dev/ and expects the name presents under /sys/block/
        // For example, /dev/dm-0 to dm-0
        real_path = real_path.trim_start_matches("/dev/");
    }

    let trimmed: &OsStr = OsStrExt::from_bytes(real_path.as_bytes());

    let path = Path::new("/sys/block/")
        .to_owned()
        .join(trimmed)
        .join("queue/rotational");
    // Normally, this file only contains '0' or '1' but just in case, we get 8 bytes...
    match get_all_data(path, 8)
        .unwrap_or_default()
        .trim()
        .parse()
        .ok()
    {
        // The disk is marked as rotational so it's a HDD.
        Some(1) => DiskType::HDD,
        // The disk is marked as non-rotational so it's very likely a SSD.
        Some(0) => DiskType::SSD,
        // Normally it shouldn't happen but welcome to the wonderful world of IT! :D
        Some(x) => DiskType::Unknown(x),
        // The information isn't available...
        None => DiskType::Unknown(-1),
    }
}

fn get_all_disks_inner(content: &str) -> Vec<Disk> {
    // The goal of this array is to list all removable devices (the ones whose name starts with
    // "usb-"). Then we check if
    let removable_entries = match fs::read_dir("/dev/disk/by-id/") {
        Ok(r) => r
            .filter_map(|res| Some(res.ok()?.path()))
            .filter_map(|e| {
                if e.file_name()
                    .and_then(|x| Some(x.to_str()?.starts_with("usb-")))
                    .unwrap_or_default()
                {
                    e.canonicalize().ok()
                } else {
                    None
                }
            })
            .collect::<Vec<PathBuf>>(),
        _ => Vec::new(),
    };

    content
        .lines()
        .map(|line| {
            let line = line.trim_start();
            // mounts format
            // http://man7.org/linux/man-pages/man5/fstab.5.html
            // fs_spec<tab>fs_file<tab>fs_vfstype<tab>other fields
            let mut fields = line.split_whitespace();
            let fs_spec = fields.next().unwrap_or("");
            let fs_file = fields
                .next()
                .unwrap_or("")
                .replace("\\134", "\\")
                .replace("\\040", " ")
                .replace("\\011", "\t")
                .replace("\\012", "\n");
            let fs_vfstype = fields.next().unwrap_or("");
            (fs_spec, fs_file, fs_vfstype)
        })
        .filter(|(fs_spec, fs_file, fs_vfstype)| {
            // Check if fs_vfstype is one of our 'ignored' file systems.
            let filtered = matches!(
                *fs_vfstype,
                "sysfs" | // pseudo file system for kernel objects
                "proc" |  // another pseudo file system
                "tmpfs" |
                "devtmpfs" |
                "cgroup" |
                "cgroup2" |
                "pstore" | // https://www.kernel.org/doc/Documentation/ABI/testing/pstore
                "squashfs" | // squashfs is a compressed read-only file system (for snaps)
                "rpc_pipefs" | // The pipefs pseudo file system service
                "iso9660" // optical media
            );

            !(filtered ||
               fs_file.starts_with("/sys") || // check if fs_file is an 'ignored' mount point
               fs_file.starts_with("/proc") ||
               (fs_file.starts_with("/run") && !fs_file.starts_with("/run/media")) ||
               fs_spec.starts_with("sunrpc"))
        })
        .filter_map(|(fs_spec, fs_file, fs_vfstype)| {
            new_disk(
                fs_spec.as_ref(),
                Path::new(&fs_file),
                fs_vfstype.as_bytes(),
                &removable_entries,
            )
        })
        .collect()
}

pub fn get_all_disks() -> Vec<Disk> {
    get_all_disks_inner(&get_all_data("/proc/mounts", 16_385).unwrap_or_default())
}

// #[test]
// fn check_all_disks() {
//     let disks = get_all_disks_inner(
//         r#"tmpfs /proc tmpfs rw,seclabel,relatime 0 0
// proc /proc proc rw,nosuid,nodev,noexec,relatime 0 0
// systemd-1 /proc/sys/fs/binfmt_misc autofs rw,relatime,fd=29,pgrp=1,timeout=0,minproto=5,maxproto=5,direct,pipe_ino=17771 0 0
// tmpfs /sys tmpfs rw,seclabel,relatime 0 0
// sysfs /sys sysfs rw,seclabel,nosuid,nodev,noexec,relatime 0 0
// securityfs /sys/kernel/security securityfs rw,nosuid,nodev,noexec,relatime 0 0
// cgroup2 /sys/fs/cgroup cgroup2 rw,seclabel,nosuid,nodev,noexec,relatime,nsdelegate 0 0
// pstore /sys/fs/pstore pstore rw,seclabel,nosuid,nodev,noexec,relatime 0 0
// none /sys/fs/bpf bpf rw,nosuid,nodev,noexec,relatime,mode=700 0 0
// configfs /sys/kernel/config configfs rw,nosuid,nodev,noexec,relatime 0 0
// selinuxfs /sys/fs/selinux selinuxfs rw,relatime 0 0
// debugfs /sys/kernel/debug debugfs rw,seclabel,nosuid,nodev,noexec,relatime 0 0
// tmpfs /dev/shm tmpfs rw,seclabel,relatime 0 0
// devpts /dev/pts devpts rw,seclabel,relatime,gid=5,mode=620,ptmxmode=666 0 0
// tmpfs /sys/fs/selinux tmpfs rw,seclabel,relatime 0 0
// /dev/vda2 /proc/filesystems xfs rw,seclabel,relatime,attr2,inode64,logbufs=8,logbsize=32k,noquota 0 0
// "#,
//     );
//     assert_eq!(disks.len(), 1);
//     assert_eq!(
//         disks[0],
//         Disk {
//             type_: DiskType::Unknown(-1),
//             name: OsString::from("devpts"),
//             file_system: vec![100, 101, 118, 112, 116, 115],
//             mount_point: PathBuf::from("/dev/pts"),
//             total_space: 0,
//             available_space: 0,
//         }
//     );
// }
<|endoftext|>