<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

//! This module implements the `Any` trait, which enables dynamic typing
//! of any type.

use cast::transmute;
use option::{Option, Some, None};
use to_str::ToStr;
use unstable::intrinsics;
use util::Void;
use unstable::intrinsics::TypeId;

///////////////////////////////////////////////////////////////////////////////
// Any trait
///////////////////////////////////////////////////////////////////////////////

/// The `Any` trait is implemented by all types, and can be used as a trait object
/// for dynamic typing
pub trait Any {
    /// Get the `TypeId` of `self`
    fn get_type_id(&self) -> TypeId;

    /// Get a void pointer to `self`
    fn as_void_ptr(&self) -> *Void;

    /// Get a mutable void pointer to `self`
    fn as_mut_void_ptr(&mut self) -> *mut Void;
}

impl<T: 'static> Any for T {
    /// Get the `TypeId` of `self`
    fn get_type_id(&self) -> TypeId {
        TypeId::of::<T>()
    }

    /// Get a void pointer to `self`
    fn as_void_ptr(&self) -> *Void {
        self as *T as *Void
    }

    /// Get a mutable void pointer to `self`
    fn as_mut_void_ptr(&mut self) -> *mut Void {
        self as *mut T as *mut Void
    }
}

///////////////////////////////////////////////////////////////////////////////
// Extension methods for Any trait objects.
// Implemented as three extension traits so that generics work.
///////////////////////////////////////////////////////////////////////////////

/// Extension methods for a referenced `Any` trait object
pub trait AnyRefExt<'self> {
    /// Returns true if the boxed type is the same as `T`
    fn is<T: 'static>(self) -> bool;

    /// Returns some reference to the boxed value if it is of type `T`, or
    /// `None` if it isn't.
    fn as_ref<T: 'static>(self) -> Option<&'self T>;
}

impl<'self> AnyRefExt<'self> for &'self Any {
    #[inline]
    fn is<T: 'static>(self) -> bool {
        // Get TypeId of the type this function is instantiated with
        let t = TypeId::of::<T>();

        // Get TypeId of the type in the trait object
        let boxed = self.get_type_id();

        // Compare both TypeIds on equality
        t == boxed
    }

    #[inline]
    fn as_ref<T: 'static>(self) -> Option<&'self T> {
        if self.is::<T>() {
            Some(unsafe { transmute(self.as_void_ptr()) })
        } else {
            None
        }
    }
}

/// Extension methods for a mutable referenced `Any` trait object
pub trait AnyMutRefExt<'self> {
    /// Returns some mutable reference to the boxed value if it is of type `T`, or
    /// `None` if it isn't.
    fn as_mut<T: 'static>(self) -> Option<&'self mut T>;
}

impl<'self> AnyMutRefExt<'self> for &'self mut Any {
    #[inline]
    fn as_mut<T: 'static>(self) -> Option<&'self mut T> {
        if self.is::<T>() {
            Some(unsafe { transmute(self.as_mut_void_ptr()) })
        } else {
            None
        }
    }
}

/// Extension methods for a owning `Any` trait object
pub trait AnyOwnExt {
    /// Returns the boxed value if it is of type `T`, or
    /// `None` if it isn't.
    fn move<T: 'static>(self) -> Option<~T>;
}

impl AnyOwnExt for ~Any {
    #[inline]
    fn move<T: 'static>(self) -> Option<~T> {
        if self.is::<T>() {
            unsafe {
                // Extract the pointer to the boxed value, temporary alias with self
                let ptr: ~T = transmute(self.as_void_ptr());

                // Prevent destructor on self being run
                intrinsics::forget(self);

                Some(ptr)
            }
        } else {
            None
        }
    }
}

///////////////////////////////////////////////////////////////////////////////
// Trait implementations
///////////////////////////////////////////////////////////////////////////////

impl ToStr for ~Any {
    fn to_str(&self) -> ~str { ~"~Any" }
}

impl<'self> ToStr for &'self Any {
    fn to_str(&self) -> ~str { ~"&Any" }
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::AnyRefExt;
    use option::{Some, None};
    use hash::Hash;

    #[deriving(Eq)]
    struct Test;

    static TEST: &'static str = "Test";

    #[test]
    fn any_as_void_ptr() {
        let (a, b, c) = (~5u as ~Any, ~TEST as ~Any, ~Test as ~Any);
        let a_r: &Any = a;
        let b_r: &Any = b;
        let c_r: &Any = c;

        assert_eq!(a.as_void_ptr(), a_r.as_void_ptr());
        assert_eq!(b.as_void_ptr(), b_r.as_void_ptr());
        assert_eq!(c.as_void_ptr(), c_r.as_void_ptr());

        let (a, b, c) = (@5u as @Any, @TEST as @Any, @Test as @Any);
        let a_r: &Any = a;
        let b_r: &Any = b;
        let c_r: &Any = c;

        assert_eq!(a.as_void_ptr(), a_r.as_void_ptr());
        assert_eq!(b.as_void_ptr(), b_r.as_void_ptr());
        assert_eq!(c.as_void_ptr(), c_r.as_void_ptr());

        let (a, b, c) = (&5u as &Any, &TEST as &Any, &Test as &Any);
        let a_r: &Any = a;
        let b_r: &Any = b;
        let c_r: &Any = c;

        assert_eq!(a.as_void_ptr(), a_r.as_void_ptr());
        assert_eq!(b.as_void_ptr(), b_r.as_void_ptr());
        assert_eq!(c.as_void_ptr(), c_r.as_void_ptr());

        let mut x = Test;
        let mut y: &'static str = "Test";
        let (a, b, c) = (&mut 5u as &mut Any,
                         &mut y as &mut Any,
                         &mut x as &mut Any);
        let a_r: &Any = a;
        let b_r: &Any = b;
        let c_r: &Any = c;

        assert_eq!(a.as_void_ptr(), a_r.as_void_ptr());
        assert_eq!(b.as_void_ptr(), b_r.as_void_ptr());
        assert_eq!(c.as_void_ptr(), c_r.as_void_ptr());

        let (a, b, c) = (5u, "hello", Test);
        let (a_r, b_r, c_r) = (&a as &Any, &b as &Any, &c as &Any);

        assert_eq!(a.as_void_ptr(), a_r.as_void_ptr());
        assert_eq!(b.as_void_ptr(), b_r.as_void_ptr());
        assert_eq!(c.as_void_ptr(), c_r.as_void_ptr());
    }

    #[test]
    fn any_as_mut_void_ptr() {
        let y: &'static str = "Test";
        let mut a = ~5u as ~Any;
        let mut b = ~y as ~Any;
        let mut c = ~Test as ~Any;

        let a_ptr = a.as_mut_void_ptr();
        let b_ptr = b.as_mut_void_ptr();
        let c_ptr = c.as_mut_void_ptr();

        let a_r: &mut Any = a;
        let b_r: &mut Any = b;
        let c_r: &mut Any = c;

        assert_eq!(a_ptr, a_r.as_mut_void_ptr());
        assert_eq!(b_ptr, b_r.as_mut_void_ptr());
        assert_eq!(c_ptr, c_r.as_mut_void_ptr());

        let mut x = Test;
        let mut y: &'static str = "Test";
        let a = &mut 5u as &mut Any;
        let b = &mut y as &mut Any;
        let c = &mut x as &mut Any;

        let a_ptr = a.as_mut_void_ptr();
        let b_ptr = b.as_mut_void_ptr();
        let c_ptr = c.as_mut_void_ptr();

        let a_r: &mut Any = a;
        let b_r: &mut Any = b;
        let c_r: &mut Any = c;

        assert_eq!(a_ptr, a_r.as_mut_void_ptr());
        assert_eq!(b_ptr, b_r.as_mut_void_ptr());
        assert_eq!(c_ptr, c_r.as_mut_void_ptr());

        let y: &'static str = "Test";
        let mut a = 5u;
        let mut b = y;
        let mut c = Test;

        let a_ptr = a.as_mut_void_ptr();
        let b_ptr = b.as_mut_void_ptr();
        let c_ptr = c.as_mut_void_ptr();

        let (a_r, b_r, c_r) = (&mut a as &mut Any, &mut b as &mut Any, &mut c as &mut Any);

        assert_eq!(a_ptr, a_r.as_mut_void_ptr());
        assert_eq!(b_ptr, b_r.as_mut_void_ptr());
        assert_eq!(c_ptr, c_r.as_mut_void_ptr());
    }

    #[test]
    fn any_referenced() {
        let (a, b, c) = (&5u as &Any, &TEST as &Any, &Test as &Any);

        assert!(a.is::<uint>());
        assert!(!b.is::<uint>());
        assert!(!c.is::<uint>());

        assert!(!a.is::<&'static str>());
        assert!(b.is::<&'static str>());
        assert!(!c.is::<&'static str>());

        assert!(!a.is::<Test>());
        assert!(!b.is::<Test>());
        assert!(c.is::<Test>());
    }

    #[test]
    fn any_owning() {
        let (a, b, c) = (~5u as ~Any, ~TEST as ~Any, ~Test as ~Any);

        assert!(a.is::<uint>());
        assert!(!b.is::<uint>());
        assert!(!c.is::<uint>());

        assert!(!a.is::<&'static str>());
        assert!(b.is::<&'static str>());
        assert!(!c.is::<&'static str>());

        assert!(!a.is::<Test>());
        assert!(!b.is::<Test>());
        assert!(c.is::<Test>());
    }

    #[test]
    fn any_managed() {
        let (a, b, c) = (@5u as @Any, @TEST as @Any, @Test as @Any);

        assert!(a.is::<uint>());
        assert!(!b.is::<uint>());
        assert!(!c.is::<uint>());

        assert!(!a.is::<&'static str>());
        assert!(b.is::<&'static str>());
        assert!(!c.is::<&'static str>());

        assert!(!a.is::<Test>());
        assert!(!b.is::<Test>());
        assert!(c.is::<Test>());
    }

    #[test]
    fn any_as_ref() {
        let a = &5u as &Any;

        match a.as_ref::<uint>() {
            Some(&5) => {}
            x => fail!("Unexpected value {:?}", x)
        }

        match a.as_ref::<Test>() {
            None => {}
            x => fail!("Unexpected value {:?}", x)
        }
    }

    #[test]
    fn any_as_mut() {
        let mut a = 5u;
        let mut b = ~7u;

        let a_r = &mut a as &mut Any;
        let tmp: &mut uint = b;
        let b_r = tmp as &mut Any;

        match a_r.as_mut::<uint>() {
            Some(x) => {
                assert_eq!(*x, 5u);
                *x = 612;
            }
            x => fail!("Unexpected value {:?}", x)
        }

        match b_r.as_mut::<uint>() {
            Some(x) => {
                assert_eq!(*x, 7u);
                *x = 413;
            }
            x => fail!("Unexpected value {:?}", x)
        }

        match a_r.as_mut::<Test>() {
            None => (),
            x => fail!("Unexpected value {:?}", x)
        }

        match b_r.as_mut::<Test>() {
            None => (),
            x => fail!("Unexpected value {:?}", x)
        }

        match a_r.as_mut::<uint>() {
            Some(&612) => {}
            x => fail!("Unexpected value {:?}", x)
        }

        match b_r.as_mut::<uint>() {
            Some(&413) => {}
            x => fail!("Unexpected value {:?}", x)
        }
    }

    #[test]
    fn any_move() {
        let a = ~8u as ~Any;
        let b = ~Test as ~Any;

        assert_eq!(a.move(), Some(~8u));
        assert_eq!(b.move(), Some(~Test));

        let a = ~8u as ~Any;
        let b = ~Test as ~Any;

        assert_eq!(a.move(), None::<~Test>);
        assert_eq!(b.move(), None::<~uint>);
    }
}
<commit_msg>Add some Any docs.<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

//! This module implements the `Any` trait, which enables dynamic typing
//! of any type, through runtime reflection.
//!
//! `Any` itself can be used to get a `TypeId`, and has more features when used as a trait object.
//! As `&Any` (a borrowed trait object), it has the `is` and `as_ref` methods, to test if the
//! contained value is of a given type, and to get a reference to the inner value as a type. As
//! `&mut Any`, there is also the `as_mut` method, for getting a mutable reference to the inner
//! value. `~Any` adds the `move` method, which will unwrap a `~T` from the object.  See the
//! extension traits (`*Ext`) for the full details.

use cast::transmute;
use option::{Option, Some, None};
use to_str::ToStr;
use unstable::intrinsics;
use util::Void;
use unstable::intrinsics::TypeId;

///////////////////////////////////////////////////////////////////////////////
// Any trait
///////////////////////////////////////////////////////////////////////////////

/// The `Any` trait is implemented by all types, and can be used as a trait object
/// for dynamic typing
pub trait Any {
    /// Get the `TypeId` of `self`
    fn get_type_id(&self) -> TypeId;

    /// Get a void pointer to `self`
    fn as_void_ptr(&self) -> *Void;

    /// Get a mutable void pointer to `self`
    fn as_mut_void_ptr(&mut self) -> *mut Void;
}

impl<T: 'static> Any for T {
    /// Get the `TypeId` of `self`
    fn get_type_id(&self) -> TypeId {
        TypeId::of::<T>()
    }

    /// Get a void pointer to `self`
    fn as_void_ptr(&self) -> *Void {
        self as *T as *Void
    }

    /// Get a mutable void pointer to `self`
    fn as_mut_void_ptr(&mut self) -> *mut Void {
        self as *mut T as *mut Void
    }
}

///////////////////////////////////////////////////////////////////////////////
// Extension methods for Any trait objects.
// Implemented as three extension traits so that generics work.
///////////////////////////////////////////////////////////////////////////////

/// Extension methods for a referenced `Any` trait object
pub trait AnyRefExt<'self> {
    /// Returns true if the boxed type is the same as `T`
    fn is<T: 'static>(self) -> bool;

    /// Returns some reference to the boxed value if it is of type `T`, or
    /// `None` if it isn't.
    fn as_ref<T: 'static>(self) -> Option<&'self T>;
}

impl<'self> AnyRefExt<'self> for &'self Any {
    #[inline]
    fn is<T: 'static>(self) -> bool {
        // Get TypeId of the type this function is instantiated with
        let t = TypeId::of::<T>();

        // Get TypeId of the type in the trait object
        let boxed = self.get_type_id();

        // Compare both TypeIds on equality
        t == boxed
    }

    #[inline]
    fn as_ref<T: 'static>(self) -> Option<&'self T> {
        if self.is::<T>() {
            Some(unsafe { transmute(self.as_void_ptr()) })
        } else {
            None
        }
    }
}

/// Extension methods for a mutable referenced `Any` trait object
pub trait AnyMutRefExt<'self> {
    /// Returns some mutable reference to the boxed value if it is of type `T`, or
    /// `None` if it isn't.
    fn as_mut<T: 'static>(self) -> Option<&'self mut T>;
}

impl<'self> AnyMutRefExt<'self> for &'self mut Any {
    #[inline]
    fn as_mut<T: 'static>(self) -> Option<&'self mut T> {
        if self.is::<T>() {
            Some(unsafe { transmute(self.as_mut_void_ptr()) })
        } else {
            None
        }
    }
}

/// Extension methods for a owning `Any` trait object
pub trait AnyOwnExt {
    /// Returns the boxed value if it is of type `T`, or
    /// `None` if it isn't.
    fn move<T: 'static>(self) -> Option<~T>;
}

impl AnyOwnExt for ~Any {
    #[inline]
    fn move<T: 'static>(self) -> Option<~T> {
        if self.is::<T>() {
            unsafe {
                // Extract the pointer to the boxed value, temporary alias with self
                let ptr: ~T = transmute(self.as_void_ptr());

                // Prevent destructor on self being run
                intrinsics::forget(self);

                Some(ptr)
            }
        } else {
            None
        }
    }
}

///////////////////////////////////////////////////////////////////////////////
// Trait implementations
///////////////////////////////////////////////////////////////////////////////

impl ToStr for ~Any {
    fn to_str(&self) -> ~str { ~"~Any" }
}

impl<'self> ToStr for &'self Any {
    fn to_str(&self) -> ~str { ~"&Any" }
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::AnyRefExt;
    use option::{Some, None};
    use hash::Hash;

    #[deriving(Eq)]
    struct Test;

    static TEST: &'static str = "Test";

    #[test]
    fn any_as_void_ptr() {
        let (a, b, c) = (~5u as ~Any, ~TEST as ~Any, ~Test as ~Any);
        let a_r: &Any = a;
        let b_r: &Any = b;
        let c_r: &Any = c;

        assert_eq!(a.as_void_ptr(), a_r.as_void_ptr());
        assert_eq!(b.as_void_ptr(), b_r.as_void_ptr());
        assert_eq!(c.as_void_ptr(), c_r.as_void_ptr());

        let (a, b, c) = (@5u as @Any, @TEST as @Any, @Test as @Any);
        let a_r: &Any = a;
        let b_r: &Any = b;
        let c_r: &Any = c;

        assert_eq!(a.as_void_ptr(), a_r.as_void_ptr());
        assert_eq!(b.as_void_ptr(), b_r.as_void_ptr());
        assert_eq!(c.as_void_ptr(), c_r.as_void_ptr());

        let (a, b, c) = (&5u as &Any, &TEST as &Any, &Test as &Any);
        let a_r: &Any = a;
        let b_r: &Any = b;
        let c_r: &Any = c;

        assert_eq!(a.as_void_ptr(), a_r.as_void_ptr());
        assert_eq!(b.as_void_ptr(), b_r.as_void_ptr());
        assert_eq!(c.as_void_ptr(), c_r.as_void_ptr());

        let mut x = Test;
        let mut y: &'static str = "Test";
        let (a, b, c) = (&mut 5u as &mut Any,
                         &mut y as &mut Any,
                         &mut x as &mut Any);
        let a_r: &Any = a;
        let b_r: &Any = b;
        let c_r: &Any = c;

        assert_eq!(a.as_void_ptr(), a_r.as_void_ptr());
        assert_eq!(b.as_void_ptr(), b_r.as_void_ptr());
        assert_eq!(c.as_void_ptr(), c_r.as_void_ptr());

        let (a, b, c) = (5u, "hello", Test);
        let (a_r, b_r, c_r) = (&a as &Any, &b as &Any, &c as &Any);

        assert_eq!(a.as_void_ptr(), a_r.as_void_ptr());
        assert_eq!(b.as_void_ptr(), b_r.as_void_ptr());
        assert_eq!(c.as_void_ptr(), c_r.as_void_ptr());
    }

    #[test]
    fn any_as_mut_void_ptr() {
        let y: &'static str = "Test";
        let mut a = ~5u as ~Any;
        let mut b = ~y as ~Any;
        let mut c = ~Test as ~Any;

        let a_ptr = a.as_mut_void_ptr();
        let b_ptr = b.as_mut_void_ptr();
        let c_ptr = c.as_mut_void_ptr();

        let a_r: &mut Any = a;
        let b_r: &mut Any = b;
        let c_r: &mut Any = c;

        assert_eq!(a_ptr, a_r.as_mut_void_ptr());
        assert_eq!(b_ptr, b_r.as_mut_void_ptr());
        assert_eq!(c_ptr, c_r.as_mut_void_ptr());

        let mut x = Test;
        let mut y: &'static str = "Test";
        let a = &mut 5u as &mut Any;
        let b = &mut y as &mut Any;
        let c = &mut x as &mut Any;

        let a_ptr = a.as_mut_void_ptr();
        let b_ptr = b.as_mut_void_ptr();
        let c_ptr = c.as_mut_void_ptr();

        let a_r: &mut Any = a;
        let b_r: &mut Any = b;
        let c_r: &mut Any = c;

        assert_eq!(a_ptr, a_r.as_mut_void_ptr());
        assert_eq!(b_ptr, b_r.as_mut_void_ptr());
        assert_eq!(c_ptr, c_r.as_mut_void_ptr());

        let y: &'static str = "Test";
        let mut a = 5u;
        let mut b = y;
        let mut c = Test;

        let a_ptr = a.as_mut_void_ptr();
        let b_ptr = b.as_mut_void_ptr();
        let c_ptr = c.as_mut_void_ptr();

        let (a_r, b_r, c_r) = (&mut a as &mut Any, &mut b as &mut Any, &mut c as &mut Any);

        assert_eq!(a_ptr, a_r.as_mut_void_ptr());
        assert_eq!(b_ptr, b_r.as_mut_void_ptr());
        assert_eq!(c_ptr, c_r.as_mut_void_ptr());
    }

    #[test]
    fn any_referenced() {
        let (a, b, c) = (&5u as &Any, &TEST as &Any, &Test as &Any);

        assert!(a.is::<uint>());
        assert!(!b.is::<uint>());
        assert!(!c.is::<uint>());

        assert!(!a.is::<&'static str>());
        assert!(b.is::<&'static str>());
        assert!(!c.is::<&'static str>());

        assert!(!a.is::<Test>());
        assert!(!b.is::<Test>());
        assert!(c.is::<Test>());
    }

    #[test]
    fn any_owning() {
        let (a, b, c) = (~5u as ~Any, ~TEST as ~Any, ~Test as ~Any);

        assert!(a.is::<uint>());
        assert!(!b.is::<uint>());
        assert!(!c.is::<uint>());

        assert!(!a.is::<&'static str>());
        assert!(b.is::<&'static str>());
        assert!(!c.is::<&'static str>());

        assert!(!a.is::<Test>());
        assert!(!b.is::<Test>());
        assert!(c.is::<Test>());
    }

    #[test]
    fn any_managed() {
        let (a, b, c) = (@5u as @Any, @TEST as @Any, @Test as @Any);

        assert!(a.is::<uint>());
        assert!(!b.is::<uint>());
        assert!(!c.is::<uint>());

        assert!(!a.is::<&'static str>());
        assert!(b.is::<&'static str>());
        assert!(!c.is::<&'static str>());

        assert!(!a.is::<Test>());
        assert!(!b.is::<Test>());
        assert!(c.is::<Test>());
    }

    #[test]
    fn any_as_ref() {
        let a = &5u as &Any;

        match a.as_ref::<uint>() {
            Some(&5) => {}
            x => fail!("Unexpected value {:?}", x)
        }

        match a.as_ref::<Test>() {
            None => {}
            x => fail!("Unexpected value {:?}", x)
        }
    }

    #[test]
    fn any_as_mut() {
        let mut a = 5u;
        let mut b = ~7u;

        let a_r = &mut a as &mut Any;
        let tmp: &mut uint = b;
        let b_r = tmp as &mut Any;

        match a_r.as_mut::<uint>() {
            Some(x) => {
                assert_eq!(*x, 5u);
                *x = 612;
            }
            x => fail!("Unexpected value {:?}", x)
        }

        match b_r.as_mut::<uint>() {
            Some(x) => {
                assert_eq!(*x, 7u);
                *x = 413;
            }
            x => fail!("Unexpected value {:?}", x)
        }

        match a_r.as_mut::<Test>() {
            None => (),
            x => fail!("Unexpected value {:?}", x)
        }

        match b_r.as_mut::<Test>() {
            None => (),
            x => fail!("Unexpected value {:?}", x)
        }

        match a_r.as_mut::<uint>() {
            Some(&612) => {}
            x => fail!("Unexpected value {:?}", x)
        }

        match b_r.as_mut::<uint>() {
            Some(&413) => {}
            x => fail!("Unexpected value {:?}", x)
        }
    }

    #[test]
    fn any_move() {
        let a = ~8u as ~Any;
        let b = ~Test as ~Any;

        assert_eq!(a.move(), Some(~8u));
        assert_eq!(b.move(), Some(~Test));

        let a = ~8u as ~Any;
        let b = ~Test as ~Any;

        assert_eq!(a.move(), None::<~Test>);
        assert_eq!(b.move(), None::<~uint>);
    }
}
<|endoftext|>


<commit_before>//! Iterates over lines of a given reader by keeping an internal
//! buffer which client code is handed a reference to for consumption
//! of the line.
//!
//! Since this module implements internal buffering, there is no need
//! for clients to supply a buffering reader implementation.

use std::io::{Read, Result};
use bytes;

/// Wraps a `Read` and provides a way to iterate over lines split on a
/// newline character.
///
/// Unlike `BufReaderExt::lines` a line provided by
/// `LineReader::read_line` is represented as a slice into an internal
/// buffer which gets overwritten upon the next invocation of this
/// method.  This allows clients to read through lines without
/// enforcing memory allocation on every line.  On the other hand,
/// clients must copy the provided data if a line is to be kept for
/// later reference.
///
/// # Examples
///
/// Iterating over all the lines of a given `Read` can be implemented
/// as follows:
///
/// ```
/// let mut lr = LineReader::new(r);
/// loop {
///   match lr.read_line() {
///      Ok(b) if b.is_empty() => { /* end of file */ }
///      Ok(line) => { /* process line bytes */ }
///      Err(e) => { /* i/o error occured */ }
///   }
/// }
/// ```
///
/// For convenience, you'd usually use the `read_lines` macro provided
/// by this module.  The macro will loop over the lines and take care
/// of aborting the loop when encountering end-of-file.
///
/// ```
/// let mut reader = LineReader::new(r);
/// read_lines(line in reader, {
///   match line {
///     Ok(line) => { /* process line bytes; line is never empty */ }
///     Err(e)   => { /* i/o error occured */ }
///   }
/// }
/// ```
pub struct LineReader<R> {
    block: Vec<u8>,

    inner: R,
    buf: Vec<u8>,
    pos: usize,
    cap: usize,
}

static DEFAULT_BUF_SIZE: usize = 64 * 1024;

impl<R: Read> LineReader<R> {

    /// Constructs a new `LineReader` with an internal buffering of
    /// the specified capacity.
    pub fn with_capacity(cap: usize, inner: R) -> LineReader<R> {
        let mut buf = Vec::with_capacity(cap);
        unsafe { buf.set_len(cap); }
        LineReader {
            block: Vec::new(),

            inner: inner,
            buf: buf,
            pos: 0,
            cap: 0,
        }
    }

    /// Constructs a new `LineReader` with an internal buffering of
    /// the default capacity.
    pub fn new(inner: R) -> LineReader<R> {
        LineReader::with_capacity(DEFAULT_BUF_SIZE, inner)
    }

    // private
    fn fill_buf<'a>(&'a mut self) -> Result<usize> {
        if self.pos == self.cap {
            self.cap = try!(self.inner.read(&mut self.buf[..]));
            self.pos = 0;
        }
        Ok(self.cap-self.pos)
    }

    fn read_until<'a>(&'a mut self, byte: u8) -> Result<&'a [u8]> {
        // ~ clear our previously delivered block - if any
        unsafe { self.block.set_len(0); }

        loop {
            // ensure we have data to process
            match self.fill_buf() {
                Ok(0) => {
                    return Ok(&self.block[..]);
                },
                Ok(_) => {},
                Err(e) => return Err(e),
            };

            // note: we're dealing here with buf, pos, cap directly
            // to avoid the mutability checker from getting in our way
            let b = &self.buf[self.pos .. self.cap];
            match bytes::index(b, byte) {
                Some(mut i) => {
                    i += 1;
                    let b = if self.block.is_empty() {
                        &self.buf[self.pos .. self.pos+i]
                    } else {
                        let b: &[u8] = &b[..i];
                        self.block.extend(b.iter().cloned());
                        &self.block[..]
                    };
                    self.pos += i;
                    return Ok(b);
                }
                None => {
                    self.block.extend(b.iter().cloned());
                    self.pos += b.len();
                    // continue looping
                }
            }
        }
    }

    /// Reads the next line. Returns the empty slice if EOF is
    /// reached.  The newline character - if any - is not stripped by
    /// this method.
    #[inline]
    pub fn read_line(&mut self) -> Result<&[u8]> {
        self.read_until(b'\n')
    }
}

/// Provides a convenient way to iterate over all lines through a
/// `LineReader`.  For an example, refer to the documentation of
/// `LineReader`.  The identifier representing the read line will be
/// of type `Result<&[u8]>`.
#[macro_export]
macro_rules! read_lines {
    ($inp:ident in $expr:expr, $b:block) => {
        { let ref mut r = &mut $expr;
          loop {
              let $inp = r.read_line();
              match $inp {
                  Ok(b) if b.is_empty() => {
                      break
                  }
                  _ => { $b }
              };
          };
        }
    }
}

// XXX try_read_lines!(..)

/// Counts the lines read through the given `LineReader`.
pub fn count_lines<R: Read> (mut r: LineReader<R>)
    -> Result<usize>
{
    let mut lines = 0usize;
    read_lines!(line in r, {
        match line {
            Err(e) => return Err(e),
            Ok(_) => lines += 1,
        }
    });
    Ok(lines)
}
<commit_msg>Add convenience try_read_lines macro<commit_after>//! Iterates over lines of a given reader by keeping an internal
//! buffer which client code is handed a reference to for consumption
//! of the line.
//!
//! Since this module implements internal buffering, there is no need
//! for clients to supply a buffering reader implementation.

use std::io::{Read, Result};
use bytes;

/// Wraps a `Read` and provides a way to iterate over lines split on a
/// newline character.
///
/// Unlike `BufReaderExt::lines` a line provided by
/// `LineReader::read_line` is represented as a slice into an internal
/// buffer which gets overwritten upon the next invocation of this
/// method.  This allows clients to read through lines without
/// enforcing memory allocation on every line.  On the other hand,
/// clients must copy the provided data if a line is to be kept for
/// later reference.
///
/// # Examples
///
/// Iterating over all the lines of a given `Read` can be implemented
/// as follows:
///
/// ```
/// let mut lr = LineReader::new(r);
/// loop {
///   match lr.read_line() {
///      Ok(b) if b.is_empty() => { /* end of file */ }
///      Ok(line) => { /* process line bytes */ }
///      Err(e) => { /* i/o error occured */ }
///   }
/// }
/// ```
///
/// For convenience, you'd usually use the `read_lines` macro provided
/// by this module.  The macro will loop over the lines and take care
/// of aborting the loop when encountering end-of-file.
///
/// ```
/// let mut reader = LineReader::new(r);
/// read_lines(line in reader, {
///   match line {
///     Ok(line) => { /* process line bytes; line is never empty */ }
///     Err(e)   => { /* i/o error occured */ }
///   }
/// }
/// ```
pub struct LineReader<R> {
    block: Vec<u8>,

    inner: R,
    buf: Vec<u8>,
    pos: usize,
    cap: usize,
}

static DEFAULT_BUF_SIZE: usize = 64 * 1024;

impl<R: Read> LineReader<R> {

    /// Constructs a new `LineReader` with an internal buffering of
    /// the specified capacity.
    pub fn with_capacity(cap: usize, inner: R) -> LineReader<R> {
        let mut buf = Vec::with_capacity(cap);
        unsafe { buf.set_len(cap); }
        LineReader {
            block: Vec::new(),

            inner: inner,
            buf: buf,
            pos: 0,
            cap: 0,
        }
    }

    /// Constructs a new `LineReader` with an internal buffering of
    /// the default capacity.
    pub fn new(inner: R) -> LineReader<R> {
        LineReader::with_capacity(DEFAULT_BUF_SIZE, inner)
    }

    // private
    fn fill_buf<'a>(&'a mut self) -> Result<usize> {
        if self.pos == self.cap {
            self.cap = try!(self.inner.read(&mut self.buf[..]));
            self.pos = 0;
        }
        Ok(self.cap-self.pos)
    }

    fn read_until<'a>(&'a mut self, byte: u8) -> Result<&'a [u8]> {
        // ~ clear our previously delivered block - if any
        unsafe { self.block.set_len(0); }

        loop {
            // ensure we have data to process
            match self.fill_buf() {
                Ok(0) => {
                    return Ok(&self.block[..]);
                },
                Ok(_) => {},
                Err(e) => return Err(e),
            };

            // note: we're dealing here with buf, pos, cap directly
            // to avoid the mutability checker from getting in our way
            let b = &self.buf[self.pos .. self.cap];
            match bytes::index(b, byte) {
                Some(mut i) => {
                    i += 1;
                    let b = if self.block.is_empty() {
                        &self.buf[self.pos .. self.pos+i]
                    } else {
                        let b: &[u8] = &b[..i];
                        self.block.extend(b.iter().cloned());
                        &self.block[..]
                    };
                    self.pos += i;
                    return Ok(b);
                }
                None => {
                    self.block.extend(b.iter().cloned());
                    self.pos += b.len();
                    // continue looping
                }
            }
        }
    }

    /// Reads the next line. Returns the empty slice if EOF is
    /// reached.  The newline character - if any - is not stripped by
    /// this method.
    #[inline]
    pub fn read_line(&mut self) -> Result<&[u8]> {
        self.read_until(b'\n')
    }
}

/// Provides a convenient way to iterate over all lines through a
/// `LineReader`.  For an example, refer to the documentation of
/// `LineReader`.  The identifier representing the read line will be
/// of type `Result<&[u8]>`.
#[macro_export]
macro_rules! read_lines {
    ($inp:ident in $expr:expr, $b:block) => {
        { let ref mut r = &mut $expr;
          loop {
              let $inp = r.read_line();
              match $inp {
                  Ok(b) if b.is_empty() => {
                      break
                  }
                  _ => { $b }
              };
          };
        }
    }
}

/// Provides a convenient way to iterate over all lines through a
/// `LineReader` wrapping a `try!` on the invocation to
/// `LineReader::read_line`.  Hence, the identifer representing the
/// read line will be directly of type `&[u8]`.
#[macro_export]
macro_rules! try_read_lines {
    ($inp:ident in $expr:expr, $b:block) => {
        { let ref mut r = &mut $expr;
          loop {
              let $inp = try!(r.read_line());
              if $inp.is_empty() { break; }
              else { $b };
          };
        }
    }
}

/// Counts the lines read through the given `LineReader`.
pub fn count_lines<R: Read> (mut r: LineReader<R>)
    -> Result<usize>
{
    let mut lines = 0usize;
    try_read_lines!(line in r, {
        lines += 1;
    });
    Ok(lines)
}
<|endoftext|>


<commit_before>use mio::*;
use mio::net::*;
use mio::net::tcp::*;
use mio::buf::{ByteBuf, SliceBuf};
use mio::util::Slab;
use super::localhost;
use mio::event as evt;
use collections::DList;
use std::thread::Thread;

type TestEventLoop = EventLoop<usize, String>;

const SERVER: Token = Token(0);
const CLIENT: Token = Token(1);

struct EchoConn {
    sock: TcpSocket,
    token: Token,
    count: u32,
    buf: ByteBuf
}

impl EchoConn {
    fn new(sock: TcpSocket) -> EchoConn {
        EchoConn {
            sock: sock,
            token: Token(-1),
            buf: ByteBuf::new(10000),
            count: 0
        }
    }

    fn writable(&mut self, event_loop: &mut TestEventLoop) -> MioResult<()> {

        event_loop.reregister(&self.sock, self.token, evt::READABLE, evt::PollOpt::edge())
    }

    fn readable(&mut self, event_loop: &mut TestEventLoop) -> MioResult<()> {
        match self.sock.read(&mut self.buf) {
            Ok(NonBlock::WouldBlock) => {
                panic!("We just got readable, but were unable to read from the socket?");
            }
            Ok(NonBlock::Ready(r)) => {
                debug!("CONN : we read {} bytes!", r);
                self.buf.clear();
                self.count += 1;
                if self.count % 10000 == 0 {
                    println!("Received {} messages", self.count);
                }
            }
            Err(e) => {
                debug!("not implemented; client err={:?}", e);
            }

        };

        event_loop.reregister(&self.sock, self.token, evt::READABLE, evt::PollOpt::edge())
    }
}

struct EchoServer {
    sock: TcpAcceptor,
    conns: Slab<EchoConn>
}

impl EchoServer {
    fn accept(&mut self, event_loop: &mut TestEventLoop) -> MioResult<()> {
        debug!("server accepting socket");

        let sock = self.sock.accept().unwrap().unwrap();
        let conn = EchoConn::new(sock,);
        let tok = self.conns.insert(conn)
            .ok().expect("could not add connection to slab");

        // Register the connection
        self.conns[tok].token = tok;
        event_loop.register_opt(&self.conns[tok].sock, tok, evt::READABLE, evt::PollOpt::edge())
            .ok().expect("could not register socket with event loop");

        Ok(())
    }

    fn conn_readable(&mut self, event_loop: &mut TestEventLoop, tok: Token) -> MioResult<()> {
        debug!("server conn readable; tok={:?}", tok);
        self.conn(tok).readable(event_loop)
    }

    fn conn_writable(&mut self, event_loop: &mut TestEventLoop, tok: Token) -> MioResult<()> {
        debug!("server conn writable; tok={:?}", tok);
        self.conn(tok).writable(event_loop)
    }

    fn conn<'a>(&'a mut self, tok: Token) -> &'a mut EchoConn {
        &mut self.conns[tok]
    }
}

struct EchoClient {
    sock: TcpSocket,
    backlog: DList<String>,
    token: Token,
    count: u32
}

// Sends a message and expects to receive the same exact message, one at a time
impl EchoClient {
    fn new(sock: TcpSocket, tok: Token) -> EchoClient {

        EchoClient {
            sock: sock,
            backlog: DList::new(),
            token: tok,
            count: 0
        }
    }

    fn readable(&mut self, event_loop: &mut TestEventLoop) -> MioResult<()> {
        Ok(())
    }

    fn writable(&mut self, event_loop: &mut TestEventLoop) -> MioResult<()> {
        debug!("client socket writable");

        while self.backlog.len() > 0 {
            match self.sock.write_slice(self.backlog.front().unwrap().as_bytes()) {
                Ok(NonBlock::WouldBlock) => {
                    break;
                }
                Ok(NonBlock::Ready(r)) => {
                    self.backlog.pop_front();
                    self.count += 1;
                }
                Err(e) => { debug!("not implemented; client err={:?}", e); break; }
            }
        }
        if self.backlog.len() > 0 {
            event_loop.reregister(&self.sock, self.token, evt::WRITABLE, evt::PollOpt::edge());
        }

        Ok(())
    }
}

struct EchoHandler {
    server: EchoServer,
    client: EchoClient,
    count: u32
}

impl EchoHandler {
    fn new(srv: TcpAcceptor, client: TcpSocket) -> EchoHandler {
        EchoHandler {
            server: EchoServer {
                sock: srv,
                conns: Slab::new_starting_at(Token(2), 128),
            },
            client: EchoClient::new(client, CLIENT),
            count: 0
        }
    }
}

impl Handler<usize, String> for EchoHandler {
    fn readable(&mut self, event_loop: &mut TestEventLoop, token: Token, hint: evt::ReadHint) {
        assert_eq!(hint, evt::DATAHINT);

        match token {
            SERVER => self.server.accept(event_loop).unwrap(),
            CLIENT => self.client.readable(event_loop).unwrap(),
            i => self.server.conn_readable(event_loop, i).unwrap()
        };
    }

    fn writable(&mut self, event_loop: &mut TestEventLoop, token: Token) {
        match token {
            SERVER => panic!("received writable for token 0"),
            CLIENT => self.client.writable(event_loop).unwrap(),
            _ => self.server.conn_writable(event_loop, token).unwrap()
        };
    }

    fn notify(&mut self, event_loop: &mut TestEventLoop, msg: String) {
        match self.client.sock.write_slice(msg.as_bytes()) {
            Ok(n) => self.count += 1,
            Err(_) => {
                self.client.backlog.push_back(msg);
                event_loop.reregister(&self.client.sock, self.client.token, evt::WRITABLE, evt::PollOpt::edge());
            }
        }
    }
}

#[test]
pub fn test_echo_server() {
    debug!("Starting TEST_ECHO_SERVER");
    let config =
        EventLoopConfig {
            io_poll_timeout_ms: 1_000,
            notify_capacity: 1_048_576,
            messages_per_tick: 64,
            timer_tick_ms: 100,
            timer_wheel_size: 1_024,
            timer_capacity: 65_536,
        };
    let mut event_loop = EventLoop::configured(config).unwrap();

    let addr = SockAddr::parse(localhost().as_slice())
        .expect("could not parse InetAddr");

    let srv = TcpSocket::v4().unwrap();

    info!("setting re-use addr");
    srv.set_reuseaddr(true).unwrap();

    let srv = srv.bind(&addr).unwrap()
        .listen(256).unwrap();

    info!("listen for connections");
    event_loop.register_opt(&srv, SERVER, evt::READABLE, evt::PollOpt::edge()).unwrap();

    let sock = TcpSocket::v4().unwrap();

    // Connect to the server
    event_loop.register_opt(&sock, CLIENT, evt::WRITABLE, evt::PollOpt::edge()).unwrap();
    sock.connect(&addr).unwrap();
    let chan = event_loop.channel();

    let go = move|:| {
        let mut i = 1_000_000;

        let message = String::from_str("THIS IS A TEST MESSAGE");
        while i > 0 {
            chan.send(message.clone());
            i -= 1;
            if i % 10000 == 0 {
                println!("Sent {} messages", 1_000_000 - i);
            }
        }
    };

    let guard = Thread::spawn(go);

    // Start the event loop
    event_loop.run(EchoHandler::new(srv, sock))
        .ok().expect("failed to execute event loop");

}
<commit_msg>added termination and more telemetry to test<commit_after>use mio::*;
use mio::net::*;
use mio::net::tcp::*;
use mio::buf::{ByteBuf, SliceBuf};
use mio::util::Slab;
use super::localhost;
use mio::event as evt;
use collections::DList;
use std::thread::Thread;
use std::io::timer::Timer;
use std::time::duration::Duration;

type TestEventLoop = EventLoop<usize, String>;

const SERVER: Token = Token(0);
const CLIENT: Token = Token(1);

struct EchoConn {
    sock: TcpSocket,
    token: Token,
    count: u32,
    buf: Vec<u8>
}

impl EchoConn {
    fn new(sock: TcpSocket) -> EchoConn {
        let mut ec =
        EchoConn {
            sock: sock,
            token: Token(-1),
            buf: Vec::with_capacity(22),
            count: 0
        };
        unsafe { ec.buf.set_len(22) };
        ec
    }

    fn writable(&mut self, event_loop: &mut TestEventLoop) -> MioResult<()> {

        event_loop.reregister(&self.sock, self.token, evt::READABLE, evt::PollOpt::edge())
    }

    fn readable(&mut self, event_loop: &mut TestEventLoop) -> MioResult<()> {
        loop {
            match self.sock.read_slice(self.buf.as_mut_slice()) {
                Ok(NonBlock::WouldBlock) => {
                    break;
                }
                Ok(NonBlock::Ready(r)) => {
                    self.count += 1;
                    if self.count % 10000 == 0 {
                        println!("Received {} messages", self.count);
                    }
                    if self.count == 1_000_000 {
                        event_loop.shutdown();
                    }
                }
                Err(e) => {
                    break;
                }

            };
        }

        event_loop.reregister(&self.sock, self.token, evt::READABLE, evt::PollOpt::edge())
    }
}

struct EchoServer {
    sock: TcpAcceptor,
    conns: Slab<EchoConn>
}

impl EchoServer {
    fn accept(&mut self, event_loop: &mut TestEventLoop) -> MioResult<()> {
        debug!("server accepting socket");

        let sock = self.sock.accept().unwrap().unwrap();
        let conn = EchoConn::new(sock,);
        let tok = self.conns.insert(conn)
            .ok().expect("could not add connection to slab");

        // Register the connection
        self.conns[tok].token = tok;
        event_loop.register_opt(&self.conns[tok].sock, tok, evt::READABLE, evt::PollOpt::edge())
            .ok().expect("could not register socket with event loop");

        Ok(())
    }

    fn conn_readable(&mut self, event_loop: &mut TestEventLoop, tok: Token) -> MioResult<()> {
        debug!("server conn readable; tok={:?}", tok);
        self.conn(tok).readable(event_loop)
    }

    fn conn_writable(&mut self, event_loop: &mut TestEventLoop, tok: Token) -> MioResult<()> {
        debug!("server conn writable; tok={:?}", tok);
        self.conn(tok).writable(event_loop)
    }

    fn conn<'a>(&'a mut self, tok: Token) -> &'a mut EchoConn {
        &mut self.conns[tok]
    }
}

struct EchoClient {
    sock: TcpSocket,
    backlog: DList<String>,
    token: Token,
    count: u32
}

// Sends a message and expects to receive the same exact message, one at a time
impl EchoClient {
    fn new(sock: TcpSocket, tok: Token) -> EchoClient {

        EchoClient {
            sock: sock,
            backlog: DList::new(),
            token: tok,
            count: 0
        }
    }

    fn readable(&mut self, event_loop: &mut TestEventLoop) -> MioResult<()> {
        Ok(())
    }

    fn writable(&mut self, event_loop: &mut TestEventLoop) -> MioResult<()> {
        debug!("client socket writable");

        while self.backlog.len() > 0 {
            match self.sock.write_slice(self.backlog.front().unwrap().as_bytes()) {
                Ok(NonBlock::WouldBlock) => {
                    break;
                }
                Ok(NonBlock::Ready(r)) => {
                    self.backlog.pop_front();
                    self.count += 1;
                    if self.count % 10000 == 0 {
                        println!("Sent {} messages", self.count);
                    }
                }
                Err(e) => { debug!("not implemented; client err={:?}", e); break; }
            }
        }
        if self.backlog.len() > 0 {
            event_loop.reregister(&self.sock, self.token, evt::WRITABLE, evt::PollOpt::edge());
        }

        Ok(())
    }
}

struct EchoHandler {
    server: EchoServer,
    client: EchoClient,
    count: u32
}

impl EchoHandler {
    fn new(srv: TcpAcceptor, client: TcpSocket) -> EchoHandler {
        EchoHandler {
            server: EchoServer {
                sock: srv,
                conns: Slab::new_starting_at(Token(2), 128),
            },
            client: EchoClient::new(client, CLIENT),
            count: 0
        }
    }
}

impl Handler<usize, String> for EchoHandler {
    fn readable(&mut self, event_loop: &mut TestEventLoop, token: Token, hint: evt::ReadHint) {
        assert_eq!(hint, evt::DATAHINT);

        match token {
            SERVER => self.server.accept(event_loop).unwrap(),
            CLIENT => self.client.readable(event_loop).unwrap(),
            i => self.server.conn_readable(event_loop, i).unwrap()
        };
    }

    fn writable(&mut self, event_loop: &mut TestEventLoop, token: Token) {
        match token {
            SERVER => panic!("received writable for token 0"),
            CLIENT => self.client.writable(event_loop).unwrap(),
            _ => self.server.conn_writable(event_loop, token).unwrap()
        };
    }

    fn notify(&mut self, event_loop: &mut TestEventLoop, msg: String) {
        match self.client.sock.write_slice(msg.as_bytes()) {
            Ok(NonBlock::Ready(n)) => {
                self.client.count += 1;
                if self.client.count % 10000 == 0 {
                    println!("Sent {} bytes:   count {}", n, self.client.count);
                }
            },

            _ => {
                self.client.backlog.push_back(msg);
                event_loop.reregister(&self.client.sock, self.client.token, evt::WRITABLE, evt::PollOpt::edge());
            }
        }
    }
}

#[test]
pub fn test_echo_server() {
    debug!("Starting TEST_ECHO_SERVER");
    let config =
        EventLoopConfig {
            io_poll_timeout_ms: 1_000,
            notify_capacity: 1_048_576,
            messages_per_tick: 64,
            timer_tick_ms: 100,
            timer_wheel_size: 1_024,
            timer_capacity: 65_536,
        };
    let mut event_loop = EventLoop::configured(config).unwrap();

    let addr = SockAddr::parse(localhost().as_slice())
        .expect("could not parse InetAddr");

    let srv = TcpSocket::v4().unwrap();

    info!("setting re-use addr");
    srv.set_reuseaddr(true).unwrap();

    let srv = srv.bind(&addr).unwrap()
        .listen(256).unwrap();

    info!("listen for connections");
    event_loop.register_opt(&srv, SERVER, evt::READABLE, evt::PollOpt::edge()).unwrap();

    let sock = TcpSocket::v4().unwrap();

    // Connect to the server
    event_loop.register_opt(&sock, CLIENT, evt::WRITABLE, evt::PollOpt::edge()).unwrap();
    sock.connect(&addr).unwrap();
    let chan = event_loop.channel();

    let go = move|:| {
        let mut i = 1_000_000;

        let mut timer = Timer::new().unwrap();
        timer.sleep(Duration::seconds(1));

        let message = String::from_str("THIS IS A TEST MESSAGE");
        while i > 0 {
            chan.send(message.clone());
            i -= 1;
            if i % 10000 == 0 {
                println!("Enqueued {} messages", 1_000_000 - i);
            }
        }
    };

    let guard = Thread::spawn(go);

    // Start the event loop
    event_loop.run(EchoHandler::new(srv, sock))
        .ok().expect("failed to execute event loop");

}
<|endoftext|>


<commit_before>//! Validates the MIR to ensure that invariants are upheld.

use super::{MirPass, MirSource};
use rustc_middle::mir::visit::Visitor;
use rustc_middle::{
    mir::{
        BasicBlock, Body, Location, Operand, Rvalue, Statement, StatementKind, Terminator,
        TerminatorKind,
    },
    ty::{self, ParamEnv, TyCtxt},
};
use rustc_span::def_id::DefId;

enum EdgeKind {
    Unwind,
    Normal,
}

pub struct Validator {
    /// Describes at which point in the pipeline this validation is happening.
    pub when: String,
}

impl<'tcx> MirPass<'tcx> for Validator {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut Body<'tcx>) {
        let def_id = source.def_id();
        let param_env = tcx.param_env(def_id);
        TypeChecker { when: &self.when, def_id, body, tcx, param_env }.visit_body(body);
    }
}

struct TypeChecker<'a, 'tcx> {
    when: &'a str,
    def_id: DefId,
    body: &'a Body<'tcx>,
    tcx: TyCtxt<'tcx>,
    param_env: ParamEnv<'tcx>,
}

impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
    fn fail(&self, location: Location, msg: impl AsRef<str>) {
        let span = self.body.source_info(location).span;
        // We use `delay_span_bug` as we might see broken MIR when other errors have already
        // occurred.
        self.tcx.sess.diagnostic().delay_span_bug(
            span,
            &format!(
                "broken MIR in {:?} ({}) at {:?}:\n{}",
                self.def_id,
                self.when,
                location,
                msg.as_ref()
            ),
        );
    }

    fn check_bb(&self, location: Location, bb: BasicBlock, edge_kind: EdgeKind) {
        if let Some(bb) = self.body.basic_blocks().get(bb) {
            let src = self.body.basic_blocks().get(location.block).unwrap();
            match (src.is_cleanup, bb.is_cleanup, edge_kind) {
                // Non-cleanup blocks can jump to non-cleanup blocks along non-unwind edges
                (false, false, EdgeKind::Normal)
                // Non-cleanup blocks can jump to cleanup blocks along unwind edges
                | (false, true, EdgeKind::Unwind)
                // Cleanup blocks can jump to cleanup blocks along non-unwind edges
                | (true, true, EdgeKind::Normal) => {}
                // All other jumps are invalid
                _ => {
                    self.fail(
                        location,
                        format!("encountered jump that does not respect unwind invariants {:?}", bb)
                    )
                }
            }
        } else {
            self.fail(location, format!("encountered jump to invalid basic block {:?}", bb))
        }
    }
}

impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
    fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) {
        // `Operand::Copy` is only supposed to be used with `Copy` types.
        if let Operand::Copy(place) = operand {
            let ty = place.ty(&self.body.local_decls, self.tcx).ty;
            let span = self.body.source_info(location).span;

            if !ty.is_copy_modulo_regions(self.tcx, self.param_env, span) {
                self.fail(location, format!("`Operand::Copy` with non-`Copy` type {}", ty));
            }
        }

        self.super_operand(operand, location);
    }

    fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
        // The sides of an assignment must not alias. Currently this just checks whether the places
        // are identical.
        if let StatementKind::Assign(box (dest, rvalue)) = &statement.kind {
            match rvalue {
                Rvalue::Use(Operand::Copy(src) | Operand::Move(src)) => {
                    if dest == src {
                        self.fail(
                            location,
                            "encountered `Assign` statement with overlapping memory",
                        );
                    }
                }
                _ => {}
            }
        }
    }

    fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
        match &terminator.kind {
            TerminatorKind::Goto { target } => {
                self.check_bb(location, *target, EdgeKind::Normal);
            }
            TerminatorKind::SwitchInt { targets, values, .. } => {
                if targets.len() != values.len() + 1 {
                    self.fail(
                        location,
                        format!(
                            "encountered `SwitchInt` terminator with {} values, but {} targets (should be values+1)",
                            values.len(),
                            targets.len(),
                        ),
                    );
                }
                for target in targets {
                    self.check_bb(location, *target, EdgeKind::Normal);
                }
            }
            TerminatorKind::Drop { target, unwind, .. } => {
                self.check_bb(location, *target, EdgeKind::Normal);
                if let Some(unwind) = unwind {
                    self.check_bb(location, *unwind, EdgeKind::Unwind);
                }
            }
            TerminatorKind::DropAndReplace { target, unwind, .. } => {
                self.check_bb(location, *target, EdgeKind::Normal);
                if let Some(unwind) = unwind {
                    self.check_bb(location, *unwind, EdgeKind::Unwind);
                }
            }
            TerminatorKind::Call { func, destination, cleanup, .. } => {
                let func_ty = func.ty(&self.body.local_decls, self.tcx);
                match func_ty.kind {
                    ty::FnPtr(..) | ty::FnDef(..) => {}
                    _ => self.fail(
                        location,
                        format!("encountered non-callable type {} in `Call` terminator", func_ty),
                    ),
                }
                if let Some((_, target)) = destination {
                    self.check_bb(location, *target, EdgeKind::Normal);
                }
                if let Some(cleanup) = cleanup {
                    self.check_bb(location, *cleanup, EdgeKind::Unwind);
                }
            }
            TerminatorKind::Assert { cond, target, cleanup, .. } => {
                let cond_ty = cond.ty(&self.body.local_decls, self.tcx);
                if cond_ty != self.tcx.types.bool {
                    self.fail(
                        location,
                        format!(
                            "encountered non-boolean condition of type {} in `Assert` terminator",
                            cond_ty
                        ),
                    );
                }
                self.check_bb(location, *target, EdgeKind::Normal);
                if let Some(cleanup) = cleanup {
                    self.check_bb(location, *cleanup, EdgeKind::Unwind);
                }
            }
            TerminatorKind::Yield { resume, drop, .. } => {
                self.check_bb(location, *resume, EdgeKind::Normal);
                if let Some(drop) = drop {
                    self.check_bb(location, *drop, EdgeKind::Normal);
                }
            }
            TerminatorKind::FalseEdge { real_target, imaginary_target } => {
                self.check_bb(location, *real_target, EdgeKind::Normal);
                self.check_bb(location, *imaginary_target, EdgeKind::Normal);
            }
            TerminatorKind::FalseUnwind { real_target, unwind } => {
                self.check_bb(location, *real_target, EdgeKind::Normal);
                if let Some(unwind) = unwind {
                    self.check_bb(location, *unwind, EdgeKind::Unwind);
                }
            }
            TerminatorKind::InlineAsm { destination, .. } => {
                if let Some(destination) = destination {
                    self.check_bb(location, *destination, EdgeKind::Normal);
                }
            }
            // Nothing to validate for these.
            TerminatorKind::Resume
            | TerminatorKind::Abort
            | TerminatorKind::Return
            | TerminatorKind::Unreachable
            | TerminatorKind::GeneratorDrop => {}
        }
    }
}
<commit_msg>Address comments<commit_after>//! Validates the MIR to ensure that invariants are upheld.

use super::{MirPass, MirSource};
use rustc_middle::mir::visit::Visitor;
use rustc_middle::{
    mir::{
        BasicBlock, Body, Location, Operand, Rvalue, Statement, StatementKind, Terminator,
        TerminatorKind,
    },
    ty::{self, ParamEnv, TyCtxt},
};
use rustc_span::def_id::DefId;

#[derive(Copy, Clone, Debug)]
enum EdgeKind {
    Unwind,
    Normal,
}

pub struct Validator {
    /// Describes at which point in the pipeline this validation is happening.
    pub when: String,
}

impl<'tcx> MirPass<'tcx> for Validator {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut Body<'tcx>) {
        let def_id = source.def_id();
        let param_env = tcx.param_env(def_id);
        TypeChecker { when: &self.when, def_id, body, tcx, param_env }.visit_body(body);
    }
}

struct TypeChecker<'a, 'tcx> {
    when: &'a str,
    def_id: DefId,
    body: &'a Body<'tcx>,
    tcx: TyCtxt<'tcx>,
    param_env: ParamEnv<'tcx>,
}

impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
    fn fail(&self, location: Location, msg: impl AsRef<str>) {
        let span = self.body.source_info(location).span;
        // We use `delay_span_bug` as we might see broken MIR when other errors have already
        // occurred.
        self.tcx.sess.diagnostic().delay_span_bug(
            span,
            &format!(
                "broken MIR in {:?} ({}) at {:?}:\n{}",
                self.def_id,
                self.when,
                location,
                msg.as_ref()
            ),
        );
    }

    fn check_edge(&self, location: Location, bb: BasicBlock, edge_kind: EdgeKind) {
        if let Some(bb) = self.body.basic_blocks().get(bb) {
            let src = self.body.basic_blocks().get(location.block).unwrap();
            match (src.is_cleanup, bb.is_cleanup, edge_kind) {
                // Non-cleanup blocks can jump to non-cleanup blocks along non-unwind edges
                (false, false, EdgeKind::Normal)
                // Non-cleanup blocks can jump to cleanup blocks along unwind edges
                | (false, true, EdgeKind::Unwind)
                // Cleanup blocks can jump to cleanup blocks along non-unwind edges
                | (true, true, EdgeKind::Normal) => {}
                // All other jumps are invalid
                _ => {
                    self.fail(
                        location,
                        format!(
                            "{:?} edge to {:?} violates unwind invariants (cleanup {:?} -> {:?})",
                            edge_kind,
                            bb,
                            src.is_cleanup,
                            bb.is_cleanup,
                        )
                    )
                }
            }
        } else {
            self.fail(location, format!("encountered jump to invalid basic block {:?}", bb))
        }
    }
}

impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
    fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) {
        // `Operand::Copy` is only supposed to be used with `Copy` types.
        if let Operand::Copy(place) = operand {
            let ty = place.ty(&self.body.local_decls, self.tcx).ty;
            let span = self.body.source_info(location).span;

            if !ty.is_copy_modulo_regions(self.tcx, self.param_env, span) {
                self.fail(location, format!("`Operand::Copy` with non-`Copy` type {}", ty));
            }
        }

        self.super_operand(operand, location);
    }

    fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
        // The sides of an assignment must not alias. Currently this just checks whether the places
        // are identical.
        if let StatementKind::Assign(box (dest, rvalue)) = &statement.kind {
            match rvalue {
                Rvalue::Use(Operand::Copy(src) | Operand::Move(src)) => {
                    if dest == src {
                        self.fail(
                            location,
                            "encountered `Assign` statement with overlapping memory",
                        );
                    }
                }
                _ => {}
            }
        }
    }

    fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
        match &terminator.kind {
            TerminatorKind::Goto { target } => {
                self.check_edge(location, *target, EdgeKind::Normal);
            }
            TerminatorKind::SwitchInt { targets, values, .. } => {
                if targets.len() != values.len() + 1 {
                    self.fail(
                        location,
                        format!(
                            "encountered `SwitchInt` terminator with {} values, but {} targets (should be values+1)",
                            values.len(),
                            targets.len(),
                        ),
                    );
                }
                for target in targets {
                    self.check_edge(location, *target, EdgeKind::Normal);
                }
            }
            TerminatorKind::Drop { target, unwind, .. } => {
                self.check_edge(location, *target, EdgeKind::Normal);
                if let Some(unwind) = unwind {
                    self.check_edge(location, *unwind, EdgeKind::Unwind);
                }
            }
            TerminatorKind::DropAndReplace { target, unwind, .. } => {
                self.check_edge(location, *target, EdgeKind::Normal);
                if let Some(unwind) = unwind {
                    self.check_edge(location, *unwind, EdgeKind::Unwind);
                }
            }
            TerminatorKind::Call { func, destination, cleanup, .. } => {
                let func_ty = func.ty(&self.body.local_decls, self.tcx);
                match func_ty.kind {
                    ty::FnPtr(..) | ty::FnDef(..) => {}
                    _ => self.fail(
                        location,
                        format!("encountered non-callable type {} in `Call` terminator", func_ty),
                    ),
                }
                if let Some((_, target)) = destination {
                    self.check_edge(location, *target, EdgeKind::Normal);
                }
                if let Some(cleanup) = cleanup {
                    self.check_edge(location, *cleanup, EdgeKind::Unwind);
                }
            }
            TerminatorKind::Assert { cond, target, cleanup, .. } => {
                let cond_ty = cond.ty(&self.body.local_decls, self.tcx);
                if cond_ty != self.tcx.types.bool {
                    self.fail(
                        location,
                        format!(
                            "encountered non-boolean condition of type {} in `Assert` terminator",
                            cond_ty
                        ),
                    );
                }
                self.check_edge(location, *target, EdgeKind::Normal);
                if let Some(cleanup) = cleanup {
                    self.check_edge(location, *cleanup, EdgeKind::Unwind);
                }
            }
            TerminatorKind::Yield { resume, drop, .. } => {
                self.check_edge(location, *resume, EdgeKind::Normal);
                if let Some(drop) = drop {
                    self.check_edge(location, *drop, EdgeKind::Normal);
                }
            }
            TerminatorKind::FalseEdge { real_target, imaginary_target } => {
                self.check_edge(location, *real_target, EdgeKind::Normal);
                self.check_edge(location, *imaginary_target, EdgeKind::Normal);
            }
            TerminatorKind::FalseUnwind { real_target, unwind } => {
                self.check_edge(location, *real_target, EdgeKind::Normal);
                if let Some(unwind) = unwind {
                    self.check_edge(location, *unwind, EdgeKind::Unwind);
                }
            }
            TerminatorKind::InlineAsm { destination, .. } => {
                if let Some(destination) = destination {
                    self.check_edge(location, *destination, EdgeKind::Normal);
                }
            }
            // Nothing to validate for these.
            TerminatorKind::Resume
            | TerminatorKind::Abort
            | TerminatorKind::Return
            | TerminatorKind::Unreachable
            | TerminatorKind::GeneratorDrop => {}
        }
    }
}
<|endoftext|>


<commit_before>extern crate bazel_protos;
extern crate boxfuture;
extern crate clap;
extern crate fs;
extern crate futures;
extern crate protobuf;

use boxfuture::{Boxable, BoxFuture};
use clap::{App, Arg, SubCommand};
use fs::hash::Fingerprint;
use fs::store::Store;
use futures::future::{Future, join_all};
use std::error::Error;
use std::fs::File;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::process::exit;
use std::sync::Arc;

#[derive(Debug)]
enum ExitCode {
  UnknownError = 1,
  NotFound = 2,
}

#[derive(Debug)]
struct ExitError(pub String, pub ExitCode);

impl From<String> for ExitError {
  fn from(s: String) -> Self {
    ExitError(s, ExitCode::UnknownError)
  }
}

fn main() {
  match execute(
    App::new("fs_util")
      .subcommand(
        SubCommand::with_name("file")
          .subcommand(
            SubCommand::with_name("cat")
              .about("Output the contents of a file by fingerprint.")
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              )),
          )
          .subcommand(
            SubCommand::with_name("save")
              .about(
                "Ingest a file by path, which allows it to be used in Directories/Snapshots. \
Outputs a fingerprint of its contents and its size in bytes, separated by a space.",
              )
              .arg(Arg::with_name("path").required(true).takes_value(true)),
          ),
      )
      .subcommand(
        SubCommand::with_name("directory")
          .subcommand(
            SubCommand::with_name("materialize")
              .about(
                "Materialize a directory by fingerprint to the filesystem. \
Destination must not exist before this command is run.",
              )
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              ))
              .arg(Arg::with_name("destination").required(true).takes_value(
                true,
              )),
          )
          .subcommand(
            SubCommand::with_name("save")
              .about(
                "Ingest a directory recursively. Saves all files found therein and saves Directory \
protos for each directory found. Outputs a fingerprint of the canonical top-level Directory proto \
and the size of the serialized proto in bytes, separated by a space.",
              )
              .arg(Arg::with_name("source").required(true).takes_value(true)),
          )
          .subcommand(
            SubCommand::with_name("cat-proto")
              .about(
                "Output the bytes of a serialized Directory proto addressed by fingerprint.",
              )
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              )),
          ),
      )
      .arg(
        Arg::with_name("local_store_path")
          .takes_value(true)
          .long("local_store_path")
          .required(true),
      )
      .get_matches(),
  ) {
    Ok(_) => {}
    Err(err) => {
      eprintln!("{}", err.0);
      exit(err.1 as i32)
    }
  };
}

fn execute(top_match: clap::ArgMatches) -> Result<(), ExitError> {
  let store_dir = top_match.value_of("local_store_path").unwrap();
  let store = Arc::new(Store::new(store_dir).map_err(|e| {
    format!(
      "Failed to open/create store for directory {}: {}",
      store_dir,
      e
    )
  })?);

  match top_match.subcommand() {
    ("file", Some(sub_match)) => {
      match sub_match.subcommand() {
        ("cat", Some(args)) => {
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          match store.load_file_bytes(&fingerprint)? {
            Some(bytes) => {
              io::stdout().write(&bytes).unwrap();
              Ok(())
            }
            None => Err(ExitError(
              format!("File with fingerprint {} not found", fingerprint),
              ExitCode::NotFound,
            )),
          }
        }
        ("save", Some(args)) => {
          let path = PathBuf::from(args.value_of("path").unwrap());
          // Canonicalize path to guarantee that a relative path has a parent.
          let posix_fs = make_posix_fs(path
            .canonicalize()
            .map_err(|e| {
              format!("Error canonicalizing path {:?}: {}", path, e.description())
            })?
            .parent()
            .ok_or_else(|| {
              format!("File being saved must have parent but {:?} did not", path)
            })?);
          let file = posix_fs
            .stat(PathBuf::from(path.file_name().unwrap()))
            .unwrap();
          match file {
            fs::Stat::File(f) => {
              let (fingerprint, size_bytes) = save_file(store, &posix_fs, f).wait().unwrap();
              Ok(println!("{} {}", fingerprint, size_bytes))
            }
            o => Err(ExitError(
              format!(
                "Tried to save file {:?} but it was not a file, was a {:?}",
                path,
                o
              ),
              ExitCode::UnknownError,
            )),
          }

        }
        (_, _) => unimplemented!(),
      }
    }
    ("directory", Some(sub_match)) => {
      match sub_match.subcommand() {
        ("materialize", Some(args)) => {
          let destination = Path::new(args.value_of("destination").unwrap());
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          Ok(materialize_directory(store, destination, &fingerprint)?)
        }
        ("save", Some(args)) => {
          let posix_fs = Arc::new(make_posix_fs(args.value_of("source").unwrap()));
          let (fingerprint, size_bytes) =
            save_directory(store, posix_fs, Arc::new(fs::Dir(PathBuf::from("."))))
              .wait()?;
          Ok(println!("{} {}", fingerprint, size_bytes))
        }
        ("cat-proto", Some(args)) => {
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          match store.load_directory_proto_bytes(&fingerprint)? {
            Some(bytes) => {
              io::stdout().write(&bytes).unwrap();
              Ok(())
            }
            None => Err(ExitError(
              format!(
                "Directory with fingerprint {} not found",
                fingerprint
              ),
              ExitCode::NotFound,
            )),
          }
        }
        (_, _) => unimplemented!(),
      }
    }
    (_, _) => unimplemented!(),
  }
}

fn make_posix_fs<P: AsRef<Path>>(root: P) -> fs::PosixFS {
  fs::PosixFS::new(&root, vec![]).unwrap()
}

fn save_file(
  store: Arc<Store>,
  posix_fs: &fs::PosixFS,
  file: fs::File,
) -> BoxFuture<(Fingerprint, usize), String> {
  posix_fs
    .read_file(&file)
    .map_err(move |err| {
      format!("Error reading file {:?}: {}", file, err.description())
    })
    .and_then(move |content| {
      Ok((
        store.store_file_bytes(&content.content)?,
        content.content.len(),
      ))
    })
    .to_boxed()
}

fn save_directory(
  store: Arc<Store>,
  posix_fs: Arc<fs::PosixFS>,
  base_dir: Arc<fs::Dir>,
) -> BoxFuture<(Fingerprint, usize), String> {
  let store_copy = store.clone();
  Box::new(
    posix_fs
      .scandir(&base_dir)
      .map_err(move |err| {
        format!("Failed to scandir {:?}: {}", &base_dir, err.description())
      })
      .and_then(move |stats| {
        let mut dirs = Vec::new();
        let mut files = Vec::new();
        for stat in stats.into_iter() {
          match stat {
            fs::Stat::Dir(dir) => dirs.push(dir),
            fs::Stat::File(file) => files.push(file),
            fs::Stat::Link(fs::Link(path)) => {
              let err: BoxFuture<_, _> =
                futures::future::err(format!("Don't yet know how to handle symlinks: {:?}", path))
                  .to_boxed();
              return err;
            }
          }
        }

        let store_inner_copy = store_copy.clone();
        let posix_fs_copy = posix_fs.clone();
        let dir_futures: Vec<_> = dirs
          .into_iter()
          .map(move |dir| {
            let dir = Arc::new(dir);
            save_directory(store_inner_copy.clone(), posix_fs_copy.clone(), dir.clone())
              .and_then(move |(fingerprint, proto_size_bytes)| {
                let mut dir_node = bazel_protos::remote_execution::DirectoryNode::new();
                dir_node.set_name(file_name_as_utf8(&dir.0)?);
                dir_node.set_digest(fingerprint_to_digest(&fingerprint, proto_size_bytes as i64));
                Ok(dir_node)
              })
          })
          .collect();

        let file_futures: Vec<_> = files
          .into_iter()
          .map(move |file| {
            file_to_file_node(store_copy.clone(), posix_fs.clone(), file)
          })
          .collect();

        join_all(dir_futures)
          .join(join_all(file_futures))
          .to_boxed()
      })
      .and_then(move |(dirs, files)| {
        let mut directory = bazel_protos::remote_execution::Directory::new();
        directory.set_directories(protobuf::RepeatedField::from_vec(dirs));
        directory.set_files(protobuf::RepeatedField::from_vec(files));
        store.record_directory(&directory)
      }),
  )
}

fn materialize_directory(
  store: Arc<Store>,
  destination: &Path,
  fingerprint: &Fingerprint,
) -> Result<(), ExitError> {
  let directory = store.load_directory_proto(&fingerprint)?.ok_or_else(|| {
    ExitError(
      format!("Directory with fingerprint {} not found", fingerprint),
      ExitCode::NotFound,
    )
  })?;
  make_clean_dir(&destination).map_err(|e| {
    format!(
      "Error making directory {:?}: {}",
      destination,
      e.description()
    )
  })?;
  for file_node in directory.get_files() {
    let fingerprint = &Fingerprint::from_hex_string(&file_node.get_digest().get_hash())?;
    match store.load_file_bytes(fingerprint)? {
      Some(bytes) => {
        let path = destination.join(file_node.get_name());
        File::create(&path)
          .and_then(|mut f| f.write_all(&bytes))
          .map_err(|e| {
            format!("Error writing file {:?}: {}", path, e.description())
          })?;
      }
      None => {
        return Err(ExitError(
          format!(
            "File with fingerprint {} not found",
            file_node.get_digest().get_hash()
          ),
          ExitCode::NotFound,
        ))
      }
    }
  }
  for directory_node in directory.get_directories() {
    materialize_directory(
      store.clone(),
      &destination.join(directory_node.get_name()),
      &Fingerprint::from_hex_string(directory_node.get_digest().get_hash())?,
    )?;
  }
  Ok(())
}

fn file_to_file_node(
  store: Arc<Store>,
  posix_fs: Arc<fs::PosixFS>,
  file: fs::File,
) -> BoxFuture<bazel_protos::remote_execution::FileNode, String> {
  let file_copy = file.clone();
  save_file(store, &posix_fs, file)
    .and_then(move |(fingerprint, size_bytes)| {
      let mut file_node = bazel_protos::remote_execution::FileNode::new();
      file_node.set_name(file_name_as_utf8(&file_copy.path)?);
      file_node.set_digest(fingerprint_to_digest(&fingerprint, size_bytes as i64));
      file_node.set_is_executable(file_copy.is_executable);
      Ok(file_node)
    })
    .to_boxed()
}

pub fn fingerprint_to_digest(
  fingerprint: &Fingerprint,
  size_bytes: i64,
) -> bazel_protos::remote_execution::Digest {
  let mut digest = bazel_protos::remote_execution::Digest::new();
  digest.set_hash(fingerprint.to_hex());
  digest.set_size_bytes(size_bytes);
  digest
}

fn file_name_as_utf8(path: &Path) -> Result<String, String> {
  match path.file_name() {
    Some(name) => {
      match name.to_str() {
        Some(name_utf8) => Ok(name_utf8.to_string()),
        None => Err(format!(
          "{:?}'s file_name is not representable in UTF8",
          path
        )),
      }
    }
    None => Err(format!("{:?} did not have a file_name", path)),
  }
}

fn make_clean_dir(path: &Path) -> Result<(), io::Error> {
  let parent = path.parent().ok_or_else(|| {
    io::Error::new(io::ErrorKind::NotFound, format!("{:?} had no parent", path))
  })?;
  std::fs::create_dir_all(parent)?;
  std::fs::create_dir(path)
}
<commit_msg>`fs_util directory cat-proto` supports text format output (#5083)<commit_after>extern crate bazel_protos;
extern crate boxfuture;
extern crate clap;
extern crate fs;
extern crate futures;
extern crate protobuf;

use boxfuture::{Boxable, BoxFuture};
use clap::{App, Arg, SubCommand};
use fs::hash::Fingerprint;
use fs::store::Store;
use futures::future::{Future, join_all};
use std::error::Error;
use std::fs::File;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::process::exit;
use std::sync::Arc;

#[derive(Debug)]
enum ExitCode {
  UnknownError = 1,
  NotFound = 2,
}

#[derive(Debug)]
struct ExitError(pub String, pub ExitCode);

impl From<String> for ExitError {
  fn from(s: String) -> Self {
    ExitError(s, ExitCode::UnknownError)
  }
}

fn main() {
  match execute(
    App::new("fs_util")
      .subcommand(
        SubCommand::with_name("file")
          .subcommand(
            SubCommand::with_name("cat")
              .about("Output the contents of a file by fingerprint.")
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              )),
          )
          .subcommand(
            SubCommand::with_name("save")
              .about(
                "Ingest a file by path, which allows it to be used in Directories/Snapshots. \
Outputs a fingerprint of its contents and its size in bytes, separated by a space.",
              )
              .arg(Arg::with_name("path").required(true).takes_value(true)),
          ),
      )
      .subcommand(
        SubCommand::with_name("directory")
          .subcommand(
            SubCommand::with_name("materialize")
              .about(
                "Materialize a directory by fingerprint to the filesystem. \
Destination must not exist before this command is run.",
              )
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              ))
              .arg(Arg::with_name("destination").required(true).takes_value(
                true,
              )),
          )
          .subcommand(
            SubCommand::with_name("save")
              .about(
                "Ingest a directory recursively. Saves all files found therein and saves Directory \
protos for each directory found. Outputs a fingerprint of the canonical top-level Directory proto \
and the size of the serialized proto in bytes, separated by a space.",
              )
              .arg(Arg::with_name("source").required(true).takes_value(true)),
          )
          .subcommand(
            SubCommand::with_name("cat-proto")
              .about(
                "Output the bytes of a serialized Directory proto addressed by fingerprint.",
              )
              .arg(
                Arg::with_name("output-format")
                  .long("output-format")
                  .takes_value(true)
                  .default_value("binary")
                  .possible_values(&["binary", "text"]),
              )
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              )),
          ),
      )
      .arg(
        Arg::with_name("local_store_path")
          .takes_value(true)
          .long("local_store_path")
          .required(true),
      )
      .get_matches(),
  ) {
    Ok(_) => {}
    Err(err) => {
      eprintln!("{}", err.0);
      exit(err.1 as i32)
    }
  };
}

fn execute(top_match: clap::ArgMatches) -> Result<(), ExitError> {
  let store_dir = top_match.value_of("local_store_path").unwrap();
  let store = Arc::new(Store::new(store_dir).map_err(|e| {
    format!(
      "Failed to open/create store for directory {}: {}",
      store_dir,
      e
    )
  })?);

  match top_match.subcommand() {
    ("file", Some(sub_match)) => {
      match sub_match.subcommand() {
        ("cat", Some(args)) => {
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          match store.load_file_bytes(&fingerprint)? {
            Some(bytes) => {
              io::stdout().write(&bytes).unwrap();
              Ok(())
            }
            None => Err(ExitError(
              format!("File with fingerprint {} not found", fingerprint),
              ExitCode::NotFound,
            )),
          }
        }
        ("save", Some(args)) => {
          let path = PathBuf::from(args.value_of("path").unwrap());
          // Canonicalize path to guarantee that a relative path has a parent.
          let posix_fs = make_posix_fs(path
            .canonicalize()
            .map_err(|e| {
              format!("Error canonicalizing path {:?}: {}", path, e.description())
            })?
            .parent()
            .ok_or_else(|| {
              format!("File being saved must have parent but {:?} did not", path)
            })?);
          let file = posix_fs
            .stat(PathBuf::from(path.file_name().unwrap()))
            .unwrap();
          match file {
            fs::Stat::File(f) => {
              let (fingerprint, size_bytes) = save_file(store, &posix_fs, f).wait().unwrap();
              Ok(println!("{} {}", fingerprint, size_bytes))
            }
            o => Err(ExitError(
              format!(
                "Tried to save file {:?} but it was not a file, was a {:?}",
                path,
                o
              ),
              ExitCode::UnknownError,
            )),
          }

        }
        (_, _) => unimplemented!(),
      }
    }
    ("directory", Some(sub_match)) => {
      match sub_match.subcommand() {
        ("materialize", Some(args)) => {
          let destination = Path::new(args.value_of("destination").unwrap());
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          Ok(materialize_directory(store, destination, &fingerprint)?)
        }
        ("save", Some(args)) => {
          let posix_fs = Arc::new(make_posix_fs(args.value_of("source").unwrap()));
          let (fingerprint, size_bytes) =
            save_directory(store, posix_fs, Arc::new(fs::Dir(PathBuf::from("."))))
              .wait()?;
          Ok(println!("{} {}", fingerprint, size_bytes))
        }
        ("cat-proto", Some(args)) => {
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          let proto_bytes = match args.value_of("output-format").unwrap() {
            "binary" => store.load_directory_proto_bytes(&fingerprint),
            "text" => {
              store.load_directory_proto(&fingerprint).map(|maybe_p| {
                maybe_p.map(|p| format!("{:?}\n", p).as_bytes().to_vec())
              })
            }
            format => Err(format!(
              "Unexpected value of --output-format arg: {}",
              format
            )),
          }?;
          match proto_bytes {
            Some(bytes) => {
              io::stdout().write(&bytes).unwrap();
              Ok(())
            }
            None => Err(ExitError(
              format!(
                "Directory with fingerprint {} not found",
                fingerprint
              ),
              ExitCode::NotFound,
            )),
          }
        }
        (_, _) => unimplemented!(),
      }
    }
    (_, _) => unimplemented!(),
  }
}

fn make_posix_fs<P: AsRef<Path>>(root: P) -> fs::PosixFS {
  fs::PosixFS::new(&root, vec![]).unwrap()
}

fn save_file(
  store: Arc<Store>,
  posix_fs: &fs::PosixFS,
  file: fs::File,
) -> BoxFuture<(Fingerprint, usize), String> {
  posix_fs
    .read_file(&file)
    .map_err(move |err| {
      format!("Error reading file {:?}: {}", file, err.description())
    })
    .and_then(move |content| {
      Ok((
        store.store_file_bytes(&content.content)?,
        content.content.len(),
      ))
    })
    .to_boxed()
}

fn save_directory(
  store: Arc<Store>,
  posix_fs: Arc<fs::PosixFS>,
  base_dir: Arc<fs::Dir>,
) -> BoxFuture<(Fingerprint, usize), String> {
  let store_copy = store.clone();
  Box::new(
    posix_fs
      .scandir(&base_dir)
      .map_err(move |err| {
        format!("Failed to scandir {:?}: {}", &base_dir, err.description())
      })
      .and_then(move |stats| {
        let mut dirs = Vec::new();
        let mut files = Vec::new();
        for stat in stats.into_iter() {
          match stat {
            fs::Stat::Dir(dir) => dirs.push(dir),
            fs::Stat::File(file) => files.push(file),
            fs::Stat::Link(fs::Link(path)) => {
              let err: BoxFuture<_, _> =
                futures::future::err(format!("Don't yet know how to handle symlinks: {:?}", path))
                  .to_boxed();
              return err;
            }
          }
        }

        let store_inner_copy = store_copy.clone();
        let posix_fs_copy = posix_fs.clone();
        let dir_futures: Vec<_> = dirs
          .into_iter()
          .map(move |dir| {
            let dir = Arc::new(dir);
            save_directory(store_inner_copy.clone(), posix_fs_copy.clone(), dir.clone())
              .and_then(move |(fingerprint, proto_size_bytes)| {
                let mut dir_node = bazel_protos::remote_execution::DirectoryNode::new();
                dir_node.set_name(file_name_as_utf8(&dir.0)?);
                dir_node.set_digest(fingerprint_to_digest(&fingerprint, proto_size_bytes as i64));
                Ok(dir_node)
              })
          })
          .collect();

        let file_futures: Vec<_> = files
          .into_iter()
          .map(move |file| {
            file_to_file_node(store_copy.clone(), posix_fs.clone(), file)
          })
          .collect();

        join_all(dir_futures)
          .join(join_all(file_futures))
          .to_boxed()
      })
      .and_then(move |(dirs, files)| {
        let mut directory = bazel_protos::remote_execution::Directory::new();
        directory.set_directories(protobuf::RepeatedField::from_vec(dirs));
        directory.set_files(protobuf::RepeatedField::from_vec(files));
        store.record_directory(&directory)
      }),
  )
}

fn materialize_directory(
  store: Arc<Store>,
  destination: &Path,
  fingerprint: &Fingerprint,
) -> Result<(), ExitError> {
  let directory = store.load_directory_proto(&fingerprint)?.ok_or_else(|| {
    ExitError(
      format!("Directory with fingerprint {} not found", fingerprint),
      ExitCode::NotFound,
    )
  })?;
  make_clean_dir(&destination).map_err(|e| {
    format!(
      "Error making directory {:?}: {}",
      destination,
      e.description()
    )
  })?;
  for file_node in directory.get_files() {
    let fingerprint = &Fingerprint::from_hex_string(&file_node.get_digest().get_hash())?;
    match store.load_file_bytes(fingerprint)? {
      Some(bytes) => {
        let path = destination.join(file_node.get_name());
        File::create(&path)
          .and_then(|mut f| f.write_all(&bytes))
          .map_err(|e| {
            format!("Error writing file {:?}: {}", path, e.description())
          })?;
      }
      None => {
        return Err(ExitError(
          format!(
            "File with fingerprint {} not found",
            file_node.get_digest().get_hash()
          ),
          ExitCode::NotFound,
        ))
      }
    }
  }
  for directory_node in directory.get_directories() {
    materialize_directory(
      store.clone(),
      &destination.join(directory_node.get_name()),
      &Fingerprint::from_hex_string(directory_node.get_digest().get_hash())?,
    )?;
  }
  Ok(())
}

fn file_to_file_node(
  store: Arc<Store>,
  posix_fs: Arc<fs::PosixFS>,
  file: fs::File,
) -> BoxFuture<bazel_protos::remote_execution::FileNode, String> {
  let file_copy = file.clone();
  save_file(store, &posix_fs, file)
    .and_then(move |(fingerprint, size_bytes)| {
      let mut file_node = bazel_protos::remote_execution::FileNode::new();
      file_node.set_name(file_name_as_utf8(&file_copy.path)?);
      file_node.set_digest(fingerprint_to_digest(&fingerprint, size_bytes as i64));
      file_node.set_is_executable(file_copy.is_executable);
      Ok(file_node)
    })
    .to_boxed()
}

pub fn fingerprint_to_digest(
  fingerprint: &Fingerprint,
  size_bytes: i64,
) -> bazel_protos::remote_execution::Digest {
  let mut digest = bazel_protos::remote_execution::Digest::new();
  digest.set_hash(fingerprint.to_hex());
  digest.set_size_bytes(size_bytes);
  digest
}

fn file_name_as_utf8(path: &Path) -> Result<String, String> {
  match path.file_name() {
    Some(name) => {
      match name.to_str() {
        Some(name_utf8) => Ok(name_utf8.to_string()),
        None => Err(format!(
          "{:?}'s file_name is not representable in UTF8",
          path
        )),
      }
    }
    None => Err(format!("{:?} did not have a file_name", path)),
  }
}

fn make_clean_dir(path: &Path) -> Result<(), io::Error> {
  let parent = path.parent().ok_or_else(|| {
    io::Error::new(io::ErrorKind::NotFound, format!("{:?} had no parent", path))
  })?;
  std::fs::create_dir_all(parent)?;
  std::fs::create_dir(path)
}
<|endoftext|>


<commit_before>extern crate time;

use std::fs::{File, read_dir};
use std::io::{BufReader, BufRead, Read, Write};
use std::io;
use std::path::Path;
use std::collections::HashMap;
use models::LanguageModelBuilder;
use error::Error;

pub fn find_most_common_words(corpus: &Path, num: usize) -> Vec<(String, u64)> {
    let words = files(corpus).flat_map(|file| read_words(file));
    let mut word_counts = HashMap::new();
    for word in words {
        *word_counts.entry(word).or_insert(0) += 1;
    }

    let mut counts: Vec<_> = word_counts.into_iter().collect();
    counts.sort_by(|a, b| b.1.cmp(&a.1));
    counts.truncate(num);

    counts
}

pub fn save_words(path: &Path, words: &Vec<(String, u64)>) -> io::Result<()> {
    let mut out = try!(File::create(path));
    for &(ref word, count) in words {
        try!(out.write_all(format!("{}, {}\n", word, count).as_bytes()));
    }

    Ok(())
}

pub fn load_most_common_words<R: Read>(file: R, num: usize) -> Result<Vec<(String, u64)>, Error> {
    let reader = BufReader::new(file);
    let mut words = vec![];
    for line in reader.lines().take(num) {
        if let Ok(line) = line {
            let mut columns = line.split(',');
            let word = columns.next().map(|w| w.trim().to_string());
            let num = columns.next().map(|n| n.trim().parse());

            match (word, num) {
                (Some(word), Some(num)) => words.push((word, try!(num))),
                _ => return Err(Error::FormatError),
            }
        } else {
            return Err(Error::FormatError);
        }
    }

    Ok(words)
}

fn sentences<T: Read + 'static>(reader: BufReader<T>) -> Box<Iterator<Item = String>> {
    Box::new(reader.split('.' as u8)
                   .filter_map(|v| {
                       let lowercase = v.unwrap();
                       String::from_utf8(lowercase).ok().map(|s| s.to_lowercase())
                   }))
}

fn read_words<R: BufRead + 'static>(reader: R) -> Box<Iterator<Item = String>> {
    Box::new(reader.lines()
                   .filter_map(|line| line.ok())
                   .flat_map(|line| {
                       line.split(|c| {
                               match c {
                                   'a'...'z' | 'A'...'Z' => false,
                                   _ => true,
                               }
                           })
                           .filter(|word| !word.is_empty())
                           .map(|word| word.to_lowercase())
                           .collect::<Vec<_>>()
                           .into_iter()
                   }))
}

fn files(path: &Path) -> Box<Iterator<Item = BufReader<File>>> {
    Box::new(read_dir(path)
                 .unwrap()
                 .map(|path| path.unwrap().path())
                 .filter(|path| path.to_str().unwrap().ends_with(".txt"))
                 .map(|path| {
                     let file = File::open(path).unwrap();
                     BufReader::new(file)
                 }))
}

pub fn create_model(corpus: &Path, words: Vec<String>) -> LanguageModelBuilder {
    let mut builder = LanguageModelBuilder::new(10, words);

    for sentence in files(corpus).flat_map(sentences) {
        let mut acc = builder.new_sentence();

        let words = sentence.split(|c| {
                                match c {
                                    'a'...'z' => false,
                                    _ => true,
                                }
                            })
                            .filter(|w| !w.is_empty());
        for word in words {
            acc.add_word(word);
        }
    }

    builder
}
<commit_msg>Refactor to take advantage of collect magic.<commit_after>extern crate time;

use std::fs::{File, read_dir};
use std::io::{BufReader, BufRead, Read, Write};
use std::io;
use std::path::Path;
use std::collections::HashMap;
use models::LanguageModelBuilder;
use error::Error;

pub fn find_most_common_words(corpus: &Path, num: usize) -> Vec<(String, u64)> {
    let words = files(corpus).flat_map(|file| read_words(file));
    let mut word_counts = HashMap::new();
    for word in words {
        *word_counts.entry(word).or_insert(0) += 1;
    }

    let mut counts: Vec<_> = word_counts.into_iter().collect();
    counts.sort_by(|a, b| b.1.cmp(&a.1));
    counts.truncate(num);

    counts
}

pub fn save_words(path: &Path, words: &Vec<(String, u64)>) -> io::Result<()> {
    let mut out = try!(File::create(path));
    for &(ref word, count) in words {
        try!(out.write_all(format!("{}, {}\n", word, count).as_bytes()));
    }

    Ok(())
}

pub fn load_most_common_words<R: Read>(file: R, num: usize) -> Result<Vec<(String, u64)>, Error> {
    let reader = BufReader::new(file);
    reader.lines()
          .take(num)
          .map(|r| {
              let line = try!(r);
              let mut columns = line.split(',');
              let word = columns.next().map(|w| w.trim().to_string());
              let num = columns.next().map(|n| n.trim().parse());
              match (word, num) {
                  (Some(word), Some(num)) => Ok((word, try!(num))),
                  _ => Err(Error::FormatError),
              }
          })
          .collect()
}

fn sentences<T: Read + 'static>(reader: BufReader<T>) -> Box<Iterator<Item = String>> {
    Box::new(reader.split('.' as u8)
                   .filter_map(|v| {
                       let lowercase = v.unwrap();
                       String::from_utf8(lowercase).ok().map(|s| s.to_lowercase())
                   }))
}

fn read_words<R: BufRead + 'static>(reader: R) -> Box<Iterator<Item = String>> {
    Box::new(reader.lines()
                   .filter_map(|line| line.ok())
                   .flat_map(|line| {
                       line.split(|c| {
                               match c {
                                   'a'...'z' | 'A'...'Z' => false,
                                   _ => true,
                               }
                           })
                           .filter(|word| !word.is_empty())
                           .map(|word| word.to_lowercase())
                           .collect::<Vec<_>>()
                           .into_iter()
                   }))
}

fn files(path: &Path) -> Box<Iterator<Item = BufReader<File>>> {
    Box::new(read_dir(path)
                 .unwrap()
                 .map(|path| path.unwrap().path())
                 .filter(|path| path.to_str().unwrap().ends_with(".txt"))
                 .map(|path| {
                     let file = File::open(path).unwrap();
                     BufReader::new(file)
                 }))
}

pub fn create_model(corpus: &Path, words: Vec<String>) -> LanguageModelBuilder {
    let mut builder = LanguageModelBuilder::new(10, words);

    for sentence in files(corpus).flat_map(sentences) {
        let mut acc = builder.new_sentence();

        let words = sentence.split(|c| {
                                match c {
                                    'a'...'z' => false,
                                    _ => true,
                                }
                            })
                            .filter(|w| !w.is_empty());
        for word in words {
            acc.add_word(word);
        }
    }

    builder
}
<|endoftext|>