<commit_before>use std::io::{Error, ErrorKind, Result};
use std::ffi::CString;
use message::Message;
use std::cmp::Ordering;
use std::ptr;
use std::boxed::Box;
use std::rc::Rc;
use std::ops::Deref;
use std::slice;
use std::time::Duration;
use ffi::*;

/// An LCM instance that handles publishing and subscribing,
/// as well as encoding and decoding messages.
pub struct Lcm {
    lcm: *mut lcm_t,
    subscriptions: Vec<Rc<LcmSubscription>>,
}

pub struct LcmSubscription {
    subscription: *mut lcm_subscription_t,
    handler: Box<FnMut(*const lcm_recv_buf_t)>,
}

impl Lcm {
    /// Creates a new `Lcm` instance.
    ///
    /// ```
    /// use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// ```
    pub fn new() -> Result<Lcm> {
        trace!("Creating LCM instance");
        let lcm = unsafe { lcm_create(ptr::null()) };
        match lcm.is_null() {
            true => Err(Error::new(ErrorKind::Other, "Failed to initialize LCM.")),
            false => {
                Ok(Lcm {
                    lcm: lcm,
                    subscriptions: Vec::new(),
                })
            }
        }
    }

    pub fn get_fileno(&self) -> ::std::os::raw::c_int {
        unsafe { lcm_get_fileno(self.lcm) }
    }

    /// Subscribes a callback to a particular topic.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("GREETINGS", |name: String| println!("Hello, {}!", name) );
    /// ```
    pub fn subscribe<M, F>(&mut self, channel: &str, mut callback: F) -> Rc<LcmSubscription>
        where M: Message,
              F: FnMut(M) + 'static
    {
        trace!("Subscribing handler to channel {}", channel);

        let channel = CString::new(channel).unwrap();

        let handler = Box::new(move |rbuf: *const lcm_recv_buf_t| {
            trace!("Running handler");
            let mut buf = unsafe {
                let ref rbuf = *rbuf;
                let data = rbuf.data as *mut u8;
                let len = rbuf.data_size as usize;
                slice::from_raw_parts(data, len)
            };
            trace!("Decoding buffer: {:?}", buf);
            match M::decode_with_hash(&mut buf) {
                Ok(msg) => callback(msg),
                Err(_) => error!("Failed to decode buffer: {:?}", buf),
            }
        });

        let mut subscription = Rc::new(LcmSubscription {
            subscription: ptr::null_mut(),
            handler: handler,
        });

        let user_data = (subscription.deref() as *const _) as *mut _;

        let c_subscription = unsafe {
            lcm_subscribe(self.lcm,
                          channel.as_ptr(),
                          Some(Lcm::handler_callback::<M>),
                          user_data)
        };

        Rc::get_mut(&mut subscription).unwrap().subscription = c_subscription;
        self.subscriptions.push(subscription.clone());

        subscription
    }

    /// Unsubscribes a message handler.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// # let mut lcm = Lcm::new().unwrap();
    /// let handler = lcm.subscribe("GREETINGS", handler_function);
    /// // ...
    /// lcm.unsubscribe(handler);
    /// ```
    pub fn unsubscribe(&mut self, handler: Rc<LcmSubscription>) -> Result<()> {
        trace!("Unsubscribing handler {:?}", handler.subscription);
        let result = unsafe { lcm_unsubscribe(self.lcm, handler.subscription) };
        // TODO: Remove subscription from Lcm object
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM: Failed to unsubscribe")),
        }
    }

    /// Publishes a message on the specified channel.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.publish("GREETINGS", &"Charles".to_string()).unwrap();
    /// ```
    pub fn publish<M>(&mut self, channel: &str, message: &M) -> Result<()>
        where M: Message
    {
        let channel = CString::new(channel).unwrap();
        let buffer = message.encode_with_hash()?;
        let result = unsafe {
            lcm_publish(self.lcm,
                        channel.as_ptr(),
                        buffer.as_ptr() as *mut _,
                        buffer.len() as _)
        };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// loop {
    /// # break;
    ///     lcm.handle().unwrap();
    /// }
    /// ```
    pub fn handle(&mut self) -> Result<()> {
        let result = unsafe { lcm_handle(self.lcm) };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message, up to a time limit.
    ///
    /// ```
    /// # use std::time::Duration;
    /// # use lcm::Lcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// let wait_dur = Duration::from_millis(100);
    /// loop {
    /// # break;
    ///     lcm.handle_timeout(Duration::from_millis(1000)).unwrap();
    /// }
    /// ```
    pub fn handle_timeout(&mut self, timeout: Duration) -> Result<()> {
        let result = unsafe { lcm_handle_timeout(self.lcm, (timeout.as_secs() * 1000) as i32 + (timeout.subsec_nanos() / 1000_000) as i32) };
        match result.cmp(&0) {
            Ordering::Less => Err(Error::new(ErrorKind::Other, "LCM Error")),
            Ordering::Equal => Err(Error::new(ErrorKind::Other, "LCM Timeout")),
            Ordering::Greater => Ok(()),
        }
    }

    /// Adjusts the maximum number of received messages that can be queued up for a subscription.
    /// The default is `30`.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// # let mut lcm = Lcm::new().unwrap();
    /// let handler = lcm.subscribe("POSITION", handler_function);
    /// lcm.subscription_set_queue_capacity(handler, 30);
    /// ```
    pub fn subscription_set_queue_capacity(&self, handler: Rc<LcmSubscription>, num_messages: usize) {
        let handler = handler.subscription;
        let num_messages = num_messages as _;
        unsafe { lcm_subscription_set_queue_capacity(handler, num_messages) };
    }

    extern "C" fn handler_callback<M>(rbuf: *const lcm_recv_buf_t,
                                      _: *const ::std::os::raw::c_char,
                                      user_data: *mut ::std::os::raw::c_void)
        where M: Message
    {
        trace!("Received data");
        let sub = user_data as *mut LcmSubscription;
        let sub = unsafe { &mut *sub };
        (sub.handler)(rbuf);
    }
}

impl Drop for Lcm {
    fn drop(&mut self) {
        trace!("Destroying Lcm instance");
        unsafe { lcm_destroy(self.lcm) };
    }
}
<commit_msg>Remove unsubscribed handlers from vec.<commit_after>use std::io::{Error, ErrorKind, Result};
use std::ffi::CString;
use message::Message;
use std::cmp::Ordering;
use std::ptr;
use std::boxed::Box;
use std::rc::Rc;
use std::ops::Deref;
use std::slice;
use std::time::Duration;
use ffi::*;

/// An LCM instance that handles publishing and subscribing,
/// as well as encoding and decoding messages.
pub struct Lcm {
    lcm: *mut lcm_t,
    subscriptions: Vec<Rc<LcmSubscription>>,
}

pub struct LcmSubscription {
    subscription: *mut lcm_subscription_t,
    handler: Box<FnMut(*const lcm_recv_buf_t)>,
}

impl Lcm {
    /// Creates a new `Lcm` instance.
    ///
    /// ```
    /// use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// ```
    pub fn new() -> Result<Lcm> {
        trace!("Creating LCM instance");
        let lcm = unsafe { lcm_create(ptr::null()) };
        match lcm.is_null() {
            true => Err(Error::new(ErrorKind::Other, "Failed to initialize LCM.")),
            false => {
                Ok(Lcm {
                    lcm: lcm,
                    subscriptions: Vec::new(),
                })
            }
        }
    }

    pub fn get_fileno(&self) -> ::std::os::raw::c_int {
        unsafe { lcm_get_fileno(self.lcm) }
    }

    /// Subscribes a callback to a particular topic.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("GREETINGS", |name: String| println!("Hello, {}!", name) );
    /// ```
    pub fn subscribe<M, F>(&mut self, channel: &str, mut callback: F) -> Rc<LcmSubscription>
        where M: Message,
              F: FnMut(M) + 'static
    {
        trace!("Subscribing handler to channel {}", channel);

        let channel = CString::new(channel).unwrap();

        let handler = Box::new(move |rbuf: *const lcm_recv_buf_t| {
            trace!("Running handler");
            let mut buf = unsafe {
                let ref rbuf = *rbuf;
                let data = rbuf.data as *mut u8;
                let len = rbuf.data_size as usize;
                slice::from_raw_parts(data, len)
            };
            trace!("Decoding buffer: {:?}", buf);
            match M::decode_with_hash(&mut buf) {
                Ok(msg) => callback(msg),
                Err(_) => error!("Failed to decode buffer: {:?}", buf),
            }
        });

        let mut subscription = Rc::new(LcmSubscription {
            subscription: ptr::null_mut(),
            handler: handler,
        });

        let user_data = (subscription.deref() as *const _) as *mut _;

        let c_subscription = unsafe {
            lcm_subscribe(self.lcm,
                          channel.as_ptr(),
                          Some(Lcm::handler_callback::<M>),
                          user_data)
        };

        Rc::get_mut(&mut subscription).unwrap().subscription = c_subscription;
        self.subscriptions.push(subscription.clone());

        subscription
    }

    /// Unsubscribes a message handler.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// # let mut lcm = Lcm::new().unwrap();
    /// let handler = lcm.subscribe("GREETINGS", handler_function);
    /// // ...
    /// lcm.unsubscribe(handler);
    /// ```
    pub fn unsubscribe(&mut self, handler: Rc<LcmSubscription>) -> Result<()> {
        trace!("Unsubscribing handler {:?}", handler.subscription);
        let result = unsafe { lcm_unsubscribe(self.lcm, handler.subscription) };

        self.subscriptions.retain(|sub| { sub.subscription != handler.subscription });

        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM: Failed to unsubscribe")),
        }
    }

    /// Publishes a message on the specified channel.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.publish("GREETINGS", &"Charles".to_string()).unwrap();
    /// ```
    pub fn publish<M>(&mut self, channel: &str, message: &M) -> Result<()>
        where M: Message
    {
        let channel = CString::new(channel).unwrap();
        let buffer = message.encode_with_hash()?;
        let result = unsafe {
            lcm_publish(self.lcm,
                        channel.as_ptr(),
                        buffer.as_ptr() as *mut _,
                        buffer.len() as _)
        };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// loop {
    /// # break;
    ///     lcm.handle().unwrap();
    /// }
    /// ```
    pub fn handle(&mut self) -> Result<()> {
        let result = unsafe { lcm_handle(self.lcm) };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message, up to a time limit.
    ///
    /// ```
    /// # use std::time::Duration;
    /// # use lcm::Lcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// let wait_dur = Duration::from_millis(100);
    /// loop {
    /// # break;
    ///     lcm.handle_timeout(Duration::from_millis(1000)).unwrap();
    /// }
    /// ```
    pub fn handle_timeout(&mut self, timeout: Duration) -> Result<()> {
        let result = unsafe { lcm_handle_timeout(self.lcm, (timeout.as_secs() * 1000) as i32 + (timeout.subsec_nanos() / 1000_000) as i32) };
        match result.cmp(&0) {
            Ordering::Less => Err(Error::new(ErrorKind::Other, "LCM Error")),
            Ordering::Equal => Err(Error::new(ErrorKind::Other, "LCM Timeout")),
            Ordering::Greater => Ok(()),
        }
    }

    /// Adjusts the maximum number of received messages that can be queued up for a subscription.
    /// The default is `30`.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// # let handler_function = |name: String| println!("Hello, {}!", name);
    /// # let mut lcm = Lcm::new().unwrap();
    /// let handler = lcm.subscribe("POSITION", handler_function);
    /// lcm.subscription_set_queue_capacity(handler, 30);
    /// ```
    pub fn subscription_set_queue_capacity(&self, handler: Rc<LcmSubscription>, num_messages: usize) {
        let handler = handler.subscription;
        let num_messages = num_messages as _;
        unsafe { lcm_subscription_set_queue_capacity(handler, num_messages) };
    }

    extern "C" fn handler_callback<M>(rbuf: *const lcm_recv_buf_t,
                                      _: *const ::std::os::raw::c_char,
                                      user_data: *mut ::std::os::raw::c_void)
        where M: Message
    {
        trace!("Received data");
        let sub = user_data as *mut LcmSubscription;
        let sub = unsafe { &mut *sub };
        (sub.handler)(rbuf);
    }
}

impl Drop for Lcm {
    fn drop(&mut self) {
        trace!("Destroying Lcm instance");
        unsafe { lcm_destroy(self.lcm) };
    }
}

///
/// Tests
///

#[test]
fn initialized() {
    let _lcm = Lcm::new().unwrap();
}

#[test]
fn test_subscribe() {
    let mut lcm = Lcm::new().unwrap();
    lcm.subscribe("channel", |_: String| {});
    assert_eq!(lcm.subscriptions.len(), 1);
}

#[test]
fn test_unsubscribe() {
    let mut lcm = Lcm::new().unwrap();
    let sub = lcm.subscribe("channel", |_: String| {});
    lcm.unsubscribe(sub).unwrap();
    assert_eq!(lcm.subscriptions.len(), 0);
}
<|endoftext|>


<commit_before>use rustc::ty::{self, Ty, TypeAndMut};
use rustc::ty::layout::{self, TyLayout, Size};
use rustc::ty::adjustment::{PointerCast};
use syntax::ast::FloatTy;
use syntax::symbol::sym;

use rustc_apfloat::ieee::{Single, Double};
use rustc_apfloat::{Float, FloatConvert};
use rustc::mir::interpret::{
    Scalar, InterpResult, Pointer, PointerArithmetic,
};
use rustc::mir::CastKind;

use super::{InterpCx, Machine, PlaceTy, OpTy, ImmTy, Immediate, FnVal};

impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
    fn type_is_fat_ptr(&self, ty: Ty<'tcx>) -> bool {
        match ty.sty {
            ty::RawPtr(ty::TypeAndMut { ty, .. }) |
            ty::Ref(_, ty, _) => !self.type_is_sized(ty),
            ty::Adt(def, _) if def.is_box() => !self.type_is_sized(ty.boxed_ty()),
            _ => false,
        }
    }

    pub fn cast(
        &mut self,
        src: OpTy<'tcx, M::PointerTag>,
        kind: CastKind,
        dest: PlaceTy<'tcx, M::PointerTag>,
    ) -> InterpResult<'tcx> {
        use rustc::mir::CastKind::*;
        match kind {
            Pointer(PointerCast::Unsize) => {
                self.unsize_into(src, dest)?;
            }

            Misc | Pointer(PointerCast::MutToConstPointer) => {
                let src = self.read_immediate(src)?;
                let res = self.cast_immediate(src, dest.layout)?;
                self.write_immediate(res, dest)?;
            }

            Pointer(PointerCast::ReifyFnPointer) => {
                // The src operand does not matter, just its type
                match src.layout.ty.sty {
                    ty::FnDef(def_id, substs) => {
                        if self.tcx.has_attr(def_id, sym::rustc_args_required_const) {
                            bug!("reifying a fn ptr that requires const arguments");
                        }
                        let instance = ty::Instance::resolve(
                            *self.tcx,
                            self.param_env,
                            def_id,
                            substs,
                        ).ok_or_else(|| err_inval!(TooGeneric))?;
                        let fn_ptr = self.memory.create_fn_alloc(FnVal::Instance(instance));
                        self.write_scalar(Scalar::Ptr(fn_ptr.into()), dest)?;
                    }
                    _ => bug!("reify fn pointer on {:?}", src.layout.ty),
                }
            }

            Pointer(PointerCast::UnsafeFnPointer) => {
                let src = self.read_immediate(src)?;
                match dest.layout.ty.sty {
                    ty::FnPtr(_) => {
                        // No change to value
                        self.write_immediate(*src, dest)?;
                    }
                    _ => bug!("fn to unsafe fn cast on {:?}", dest.layout.ty),
                }
            }

            Pointer(PointerCast::ClosureFnPointer(_)) => {
                // The src operand does not matter, just its type
                match src.layout.ty.sty {
                    ty::Closure(def_id, substs) => {
                        let substs = self.subst_and_normalize_erasing_regions(substs)?;
                        let instance = ty::Instance::resolve_closure(
                            *self.tcx,
                            def_id,
                            substs,
                            ty::ClosureKind::FnOnce,
                        );
                        let fn_ptr = self.memory.create_fn_alloc(FnVal::Instance(instance));
                        let val = Immediate::Scalar(Scalar::Ptr(fn_ptr.into()).into());
                        self.write_immediate(val, dest)?;
                    }
                    _ => bug!("closure fn pointer on {:?}", src.layout.ty),
                }
            }
        }
        Ok(())
    }

    fn cast_immediate(
        &self,
        src: ImmTy<'tcx, M::PointerTag>,
        dest_layout: TyLayout<'tcx>,
    ) -> InterpResult<'tcx, Immediate<M::PointerTag>> {
        if self.type_is_fat_ptr(src.layout.ty) {
            return match (*src, self.type_is_fat_ptr(dest_layout.ty)) {
                // pointers to extern types
                (Immediate::Scalar(_),_) |
                // slices and trait objects to other slices/trait objects
                (Immediate::ScalarPair(..), true) => {
                    // No change to immediate
                    Ok(*src)
                }
                // slices and trait objects to thin pointers (dropping the metadata)
                (Immediate::ScalarPair(data, _), false) => {
                    Ok(data.into())
                }
            };
        } else {
            match src.layout.variants {
                layout::Variants::Single { index } => {
                    if let Some(discr) =
                        src.layout.ty.discriminant_for_variant(*self.tcx, index)
                    {
                        // Cast from a univariant enum
                        assert!(src.layout.is_zst());
                        return Ok(Scalar::from_uint(discr.val, dest_layout.size).into());
                    }
                }
                layout::Variants::Multiple { .. } => {},
            }

            return Ok(self.cast_scalar(src.to_scalar()?, src.layout, dest_layout)?.into());
        }
    }

    fn cast_scalar(
        &self,
        val: Scalar<M::PointerTag>,
        src_layout: TyLayout<'tcx>,
        dest_layout: TyLayout<'tcx>,
    ) -> InterpResult<'tcx, Scalar<M::PointerTag>> {
        use rustc::ty::TyKind::*;
        trace!("Casting {:?}: {:?} to {:?}", val, src_layout.ty, dest_layout.ty);

        match src_layout.ty.sty {
            // Floating point
            Float(FloatTy::F32) => self.cast_from_float(val.to_f32()?, dest_layout.ty),
            Float(FloatTy::F64) => self.cast_from_float(val.to_f64()?, dest_layout.ty),
            // Integer(-like), including fn ptr casts and casts from enums that
            // are represented as integers (this excludes univariant enums, which
            // are handled in `cast` directly).
            _ => {
                assert!(
                    src_layout.ty.is_bool()       || src_layout.ty.is_char()     ||
                    src_layout.ty.is_enum()       || src_layout.ty.is_integral() ||
                    src_layout.ty.is_unsafe_ptr() || src_layout.ty.is_fn_ptr()   ||
                    src_layout.ty.is_region_ptr(),
                    "Unexpected cast from type {:?}", src_layout.ty
                );
                match val.to_bits_or_ptr(src_layout.size, self) {
                    Err(ptr) => self.cast_from_ptr(ptr, src_layout, dest_layout),
                    Ok(data) => self.cast_from_int(data, src_layout, dest_layout),
                }
            }
        }
    }

    fn cast_from_int(
        &self,
        v: u128, // raw bits
        src_layout: TyLayout<'tcx>,
        dest_layout: TyLayout<'tcx>,
    ) -> InterpResult<'tcx, Scalar<M::PointerTag>> {
        // Let's make sure v is sign-extended *if* it has a signed type.
        let signed = src_layout.abi.is_signed();
        let v = if signed {
            self.sign_extend(v, src_layout)
        } else {
            v
        };
        trace!("cast_from_int: {}, {}, {}", v, src_layout.ty, dest_layout.ty);
        use rustc::ty::TyKind::*;
        match dest_layout.ty.sty {
            Int(_) | Uint(_) | RawPtr(_) => {
                let v = self.truncate(v, dest_layout);
                Ok(Scalar::from_uint(v, dest_layout.size))
            }

            Float(FloatTy::F32) if signed => Ok(Scalar::from_f32(
                Single::from_i128(v as i128).value
            )),
            Float(FloatTy::F64) if signed => Ok(Scalar::from_f64(
                Double::from_i128(v as i128).value
            )),
            Float(FloatTy::F32) => Ok(Scalar::from_f32(
                Single::from_u128(v).value
            )),
            Float(FloatTy::F64) => Ok(Scalar::from_f64(
                Double::from_u128(v).value
            )),

            Char => {
                // `u8` to `char` cast
                debug_assert_eq!(v as u8 as u128, v);
                Ok(Scalar::from_uint(v, Size::from_bytes(4)))
            },

            // Casts to bool are not permitted by rustc, no need to handle them here.
            _ => throw_unsup!(Unimplemented(format!("int to {:?} cast", dest_layout.ty))),
        }
    }

    fn cast_from_float<F>(
        &self,
        f: F,
        dest_ty: Ty<'tcx>
    ) -> InterpResult<'tcx, Scalar<M::PointerTag>>
    where F: Float + Into<Scalar<M::PointerTag>> + FloatConvert<Single> + FloatConvert<Double>
    {
        use rustc::ty::TyKind::*;
        match dest_ty.sty {
            // float -> uint
            Uint(t) => {
                let width = t.bit_width().unwrap_or_else(|| self.pointer_size().bits() as usize);
                let v = f.to_u128(width).value;
                // This should already fit the bit width
                Ok(Scalar::from_uint(v, Size::from_bits(width as u64)))
            },
            // float -> int
            Int(t) => {
                let width = t.bit_width().unwrap_or_else(|| self.pointer_size().bits() as usize);
                let v = f.to_i128(width).value;
                Ok(Scalar::from_int(v, Size::from_bits(width as u64)))
            },
            // float -> f32
            Float(FloatTy::F32) =>
                Ok(Scalar::from_f32(f.convert(&mut false).value)),
            // float -> f64
            Float(FloatTy::F64) =>
                Ok(Scalar::from_f64(f.convert(&mut false).value)),
            // That's it.
            _ => bug!("invalid float to {:?} cast", dest_ty),
        }
    }

    fn cast_from_ptr(
        &self,
        ptr: Pointer<M::PointerTag>,
        src_layout: TyLayout<'tcx>,
        dest_layout: TyLayout<'tcx>,
    ) -> InterpResult<'tcx, Scalar<M::PointerTag>> {
        use rustc::ty::TyKind::*;

        match dest_layout.ty.sty {
            // Casting to a reference or fn pointer is not permitted by rustc,
            // no need to support it here.
            RawPtr(_) => Ok(ptr.into()),
            Int(_) | Uint(_) => {
                let size = self.memory.pointer_size();
                let bits = self.force_bits(Scalar::Ptr(ptr), size)?;
                self.cast_from_int(bits, src_layout, dest_layout)
            }
            _ => bug!("invalid MIR: ptr to {:?} cast", dest_layout.ty)
        }
    }

    fn unsize_into_ptr(
        &mut self,
        src: OpTy<'tcx, M::PointerTag>,
        dest: PlaceTy<'tcx, M::PointerTag>,
        // The pointee types
        sty: Ty<'tcx>,
        dty: Ty<'tcx>,
    ) -> InterpResult<'tcx> {
        // A<Struct> -> A<Trait> conversion
        let (src_pointee_ty, dest_pointee_ty) =
            self.tcx.struct_lockstep_tails_erasing_lifetimes(sty, dty, self.param_env);

        match (&src_pointee_ty.sty, &dest_pointee_ty.sty) {
            (&ty::Array(_, length), &ty::Slice(_)) => {
                let ptr = self.read_immediate(src)?.to_scalar_ptr()?;
                // u64 cast is from usize to u64, which is always good
                let val = Immediate::new_slice(
                    ptr,
                    length.unwrap_usize(self.tcx.tcx),
                    self,
                );
                self.write_immediate(val, dest)
            }
            (&ty::Dynamic(..), &ty::Dynamic(..)) => {
                // For now, upcasts are limited to changes in marker
                // traits, and hence never actually require an actual
                // change to the vtable.
                let val = self.read_immediate(src)?;
                self.write_immediate(*val, dest)
            }
            (_, &ty::Dynamic(ref data, _)) => {
                // Initial cast from sized to dyn trait
                let vtable = self.get_vtable(src_pointee_ty, data.principal())?;
                let ptr = self.read_immediate(src)?.to_scalar_ptr()?;
                let val = Immediate::new_dyn_trait(ptr, vtable);
                self.write_immediate(val, dest)
            }

            _ => bug!("invalid unsizing {:?} -> {:?}", src.layout.ty, dest.layout.ty),
        }
    }

    fn unsize_into(
        &mut self,
        src: OpTy<'tcx, M::PointerTag>,
        dest: PlaceTy<'tcx, M::PointerTag>,
    ) -> InterpResult<'tcx> {
        trace!("Unsizing {:?} into {:?}", src, dest);
        match (&src.layout.ty.sty, &dest.layout.ty.sty) {
            (&ty::Ref(_, s, _), &ty::Ref(_, d, _)) |
            (&ty::Ref(_, s, _), &ty::RawPtr(TypeAndMut { ty: d, .. })) |
            (&ty::RawPtr(TypeAndMut { ty: s, .. }),
             &ty::RawPtr(TypeAndMut { ty: d, .. })) => {
                self.unsize_into_ptr(src, dest, s, d)
            }
            (&ty::Adt(def_a, _), &ty::Adt(def_b, _)) => {
                assert_eq!(def_a, def_b);
                if def_a.is_box() || def_b.is_box() {
                    if !def_a.is_box() || !def_b.is_box() {
                        bug!("invalid unsizing between {:?} -> {:?}", src.layout, dest.layout);
                    }
                    return self.unsize_into_ptr(
                        src,
                        dest,
                        src.layout.ty.boxed_ty(),
                        dest.layout.ty.boxed_ty(),
                    );
                }

                // unsizing of generic struct with pointer fields
                // Example: `Arc<T>` -> `Arc<Trait>`
                // here we need to increase the size of every &T thin ptr field to a fat ptr
                for i in 0..src.layout.fields.count() {
                    let dst_field = self.place_field(dest, i as u64)?;
                    if dst_field.layout.is_zst() {
                        continue;
                    }
                    let src_field = self.operand_field(src, i as u64)?;
                    if src_field.layout.ty == dst_field.layout.ty {
                        self.copy_op(src_field, dst_field)?;
                    } else {
                        self.unsize_into(src_field, dst_field)?;
                    }
                }
                Ok(())
            }
            _ => {
                bug!(
                    "unsize_into: invalid conversion: {:?} -> {:?}",
                    src.layout,
                    dest.layout
                )
            }
        }
    }
}
<commit_msg>refactor cast_immediate to dispatch on the type first, and on the value second<commit_after>use rustc::ty::{self, Ty, TypeAndMut};
use rustc::ty::layout::{self, TyLayout, Size};
use rustc::ty::adjustment::{PointerCast};
use syntax::ast::FloatTy;
use syntax::symbol::sym;

use rustc_apfloat::ieee::{Single, Double};
use rustc_apfloat::{Float, FloatConvert};
use rustc::mir::interpret::{
    Scalar, InterpResult, Pointer, PointerArithmetic,
};
use rustc::mir::CastKind;

use super::{InterpCx, Machine, PlaceTy, OpTy, ImmTy, Immediate, FnVal};

impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
    pub fn cast(
        &mut self,
        src: OpTy<'tcx, M::PointerTag>,
        kind: CastKind,
        dest: PlaceTy<'tcx, M::PointerTag>,
    ) -> InterpResult<'tcx> {
        use rustc::mir::CastKind::*;
        match kind {
            Pointer(PointerCast::Unsize) => {
                self.unsize_into(src, dest)?;
            }

            Misc | Pointer(PointerCast::MutToConstPointer) => {
                let src = self.read_immediate(src)?;
                let res = self.cast_immediate(src, dest.layout)?;
                self.write_immediate(res, dest)?;
            }

            Pointer(PointerCast::ReifyFnPointer) => {
                // The src operand does not matter, just its type
                match src.layout.ty.sty {
                    ty::FnDef(def_id, substs) => {
                        if self.tcx.has_attr(def_id, sym::rustc_args_required_const) {
                            bug!("reifying a fn ptr that requires const arguments");
                        }
                        let instance = ty::Instance::resolve(
                            *self.tcx,
                            self.param_env,
                            def_id,
                            substs,
                        ).ok_or_else(|| err_inval!(TooGeneric))?;
                        let fn_ptr = self.memory.create_fn_alloc(FnVal::Instance(instance));
                        self.write_scalar(Scalar::Ptr(fn_ptr.into()), dest)?;
                    }
                    _ => bug!("reify fn pointer on {:?}", src.layout.ty),
                }
            }

            Pointer(PointerCast::UnsafeFnPointer) => {
                let src = self.read_immediate(src)?;
                match dest.layout.ty.sty {
                    ty::FnPtr(_) => {
                        // No change to value
                        self.write_immediate(*src, dest)?;
                    }
                    _ => bug!("fn to unsafe fn cast on {:?}", dest.layout.ty),
                }
            }

            Pointer(PointerCast::ClosureFnPointer(_)) => {
                // The src operand does not matter, just its type
                match src.layout.ty.sty {
                    ty::Closure(def_id, substs) => {
                        let substs = self.subst_and_normalize_erasing_regions(substs)?;
                        let instance = ty::Instance::resolve_closure(
                            *self.tcx,
                            def_id,
                            substs,
                            ty::ClosureKind::FnOnce,
                        );
                        let fn_ptr = self.memory.create_fn_alloc(FnVal::Instance(instance));
                        let val = Immediate::Scalar(Scalar::Ptr(fn_ptr.into()).into());
                        self.write_immediate(val, dest)?;
                    }
                    _ => bug!("closure fn pointer on {:?}", src.layout.ty),
                }
            }
        }
        Ok(())
    }

    fn cast_immediate(
        &self,
        src: ImmTy<'tcx, M::PointerTag>,
        dest_layout: TyLayout<'tcx>,
    ) -> InterpResult<'tcx, Immediate<M::PointerTag>> {
        use rustc::ty::TyKind::*;
        trace!("Casting {:?}: {:?} to {:?}", *src, src.layout.ty, dest_layout.ty);

        match src.layout.ty.sty {
            // Floating point
            Float(FloatTy::F32) =>
                return Ok(self.cast_from_float(src.to_scalar()?.to_f32()?, dest_layout.ty)?.into()),
            Float(FloatTy::F64) =>
                return Ok(self.cast_from_float(src.to_scalar()?.to_f64()?, dest_layout.ty)?.into()),
            // The rest is integer/pointer-"like", including fn ptr casts and casts from enums that
            // are represented as integers.
            _ =>
                assert!(
                    src.layout.ty.is_bool()       || src.layout.ty.is_char()     ||
                    src.layout.ty.is_enum()       || src.layout.ty.is_integral() ||
                    src.layout.ty.is_unsafe_ptr() || src.layout.ty.is_fn_ptr()   ||
                    src.layout.ty.is_region_ptr(),
                    "Unexpected cast from type {:?}", src.layout.ty
                )
        }

        // Handle cast the metadata away from a fat pointer.
        if dest_layout.size != src.layout.size {
            assert_eq!(dest_layout.size, self.memory.pointer_size());
            return match *src {
                Immediate::ScalarPair(data, _) => Ok(data.into()),
                Immediate::Scalar(..) =>
                    bug!(
                        "{:?} input to a fat-to-thin cast ({:?} -> {:?})",
                        *src, src.layout.ty, dest_layout.ty
                    ),
            };
        }

        // Handle cast from a univariant (ZST) enum
        match src.layout.variants {
            layout::Variants::Single { index } => {
                if let Some(discr) =
                    src.layout.ty.discriminant_for_variant(*self.tcx, index)
                {
                    assert!(src.layout.is_zst());
                    return Ok(Scalar::from_uint(discr.val, dest_layout.size).into());
                }
            }
            layout::Variants::Multiple { .. } => {},
        }

        // Handle all the rest.
        let val = src.to_scalar()?;
        Ok(match val.to_bits_or_ptr(src.layout.size, self) {
            Err(ptr) => self.cast_from_ptr(ptr, src.layout, dest_layout)?,
            Ok(data) => self.cast_from_int(data, src.layout, dest_layout)?,
        }.into())
    }

    fn cast_from_int(
        &self,
        v: u128, // raw bits
        src_layout: TyLayout<'tcx>,
        dest_layout: TyLayout<'tcx>,
    ) -> InterpResult<'tcx, Scalar<M::PointerTag>> {
        // Let's make sure v is sign-extended *if* it has a signed type.
        let signed = src_layout.abi.is_signed();
        let v = if signed {
            self.sign_extend(v, src_layout)
        } else {
            v
        };
        trace!("cast_from_int: {}, {}, {}", v, src_layout.ty, dest_layout.ty);
        use rustc::ty::TyKind::*;
        match dest_layout.ty.sty {
            Int(_) | Uint(_) | RawPtr(_) => {
                let v = self.truncate(v, dest_layout);
                Ok(Scalar::from_uint(v, dest_layout.size))
            }

            Float(FloatTy::F32) if signed => Ok(Scalar::from_f32(
                Single::from_i128(v as i128).value
            )),
            Float(FloatTy::F64) if signed => Ok(Scalar::from_f64(
                Double::from_i128(v as i128).value
            )),
            Float(FloatTy::F32) => Ok(Scalar::from_f32(
                Single::from_u128(v).value
            )),
            Float(FloatTy::F64) => Ok(Scalar::from_f64(
                Double::from_u128(v).value
            )),

            Char => {
                // `u8` to `char` cast
                debug_assert_eq!(v as u8 as u128, v);
                Ok(Scalar::from_uint(v, Size::from_bytes(4)))
            },

            // Casts to bool are not permitted by rustc, no need to handle them here.
            _ => throw_unsup!(Unimplemented(format!("int to {:?} cast", dest_layout.ty))),
        }
    }

    fn cast_from_float<F>(
        &self,
        f: F,
        dest_ty: Ty<'tcx>
    ) -> InterpResult<'tcx, Scalar<M::PointerTag>>
    where F: Float + Into<Scalar<M::PointerTag>> + FloatConvert<Single> + FloatConvert<Double>
    {
        use rustc::ty::TyKind::*;
        match dest_ty.sty {
            // float -> uint
            Uint(t) => {
                let width = t.bit_width().unwrap_or_else(|| self.pointer_size().bits() as usize);
                let v = f.to_u128(width).value;
                // This should already fit the bit width
                Ok(Scalar::from_uint(v, Size::from_bits(width as u64)))
            },
            // float -> int
            Int(t) => {
                let width = t.bit_width().unwrap_or_else(|| self.pointer_size().bits() as usize);
                let v = f.to_i128(width).value;
                Ok(Scalar::from_int(v, Size::from_bits(width as u64)))
            },
            // float -> f32
            Float(FloatTy::F32) =>
                Ok(Scalar::from_f32(f.convert(&mut false).value)),
            // float -> f64
            Float(FloatTy::F64) =>
                Ok(Scalar::from_f64(f.convert(&mut false).value)),
            // That's it.
            _ => bug!("invalid float to {:?} cast", dest_ty),
        }
    }

    fn cast_from_ptr(
        &self,
        ptr: Pointer<M::PointerTag>,
        src_layout: TyLayout<'tcx>,
        dest_layout: TyLayout<'tcx>,
    ) -> InterpResult<'tcx, Scalar<M::PointerTag>> {
        use rustc::ty::TyKind::*;

        match dest_layout.ty.sty {
            // Casting to a reference or fn pointer is not permitted by rustc,
            // no need to support it here.
            RawPtr(_) => Ok(ptr.into()),
            Int(_) | Uint(_) => {
                let size = self.memory.pointer_size();
                let bits = self.force_bits(Scalar::Ptr(ptr), size)?;
                self.cast_from_int(bits, src_layout, dest_layout)
            }
            _ => bug!("invalid MIR: ptr to {:?} cast", dest_layout.ty)
        }
    }

    fn unsize_into_ptr(
        &mut self,
        src: OpTy<'tcx, M::PointerTag>,
        dest: PlaceTy<'tcx, M::PointerTag>,
        // The pointee types
        sty: Ty<'tcx>,
        dty: Ty<'tcx>,
    ) -> InterpResult<'tcx> {
        // A<Struct> -> A<Trait> conversion
        let (src_pointee_ty, dest_pointee_ty) =
            self.tcx.struct_lockstep_tails_erasing_lifetimes(sty, dty, self.param_env);

        match (&src_pointee_ty.sty, &dest_pointee_ty.sty) {
            (&ty::Array(_, length), &ty::Slice(_)) => {
                let ptr = self.read_immediate(src)?.to_scalar_ptr()?;
                // u64 cast is from usize to u64, which is always good
                let val = Immediate::new_slice(
                    ptr,
                    length.unwrap_usize(self.tcx.tcx),
                    self,
                );
                self.write_immediate(val, dest)
            }
            (&ty::Dynamic(..), &ty::Dynamic(..)) => {
                // For now, upcasts are limited to changes in marker
                // traits, and hence never actually require an actual
                // change to the vtable.
                let val = self.read_immediate(src)?;
                self.write_immediate(*val, dest)
            }
            (_, &ty::Dynamic(ref data, _)) => {
                // Initial cast from sized to dyn trait
                let vtable = self.get_vtable(src_pointee_ty, data.principal())?;
                let ptr = self.read_immediate(src)?.to_scalar_ptr()?;
                let val = Immediate::new_dyn_trait(ptr, vtable);
                self.write_immediate(val, dest)
            }

            _ => bug!("invalid unsizing {:?} -> {:?}", src.layout.ty, dest.layout.ty),
        }
    }

    fn unsize_into(
        &mut self,
        src: OpTy<'tcx, M::PointerTag>,
        dest: PlaceTy<'tcx, M::PointerTag>,
    ) -> InterpResult<'tcx> {
        trace!("Unsizing {:?} into {:?}", src, dest);
        match (&src.layout.ty.sty, &dest.layout.ty.sty) {
            (&ty::Ref(_, s, _), &ty::Ref(_, d, _)) |
            (&ty::Ref(_, s, _), &ty::RawPtr(TypeAndMut { ty: d, .. })) |
            (&ty::RawPtr(TypeAndMut { ty: s, .. }),
             &ty::RawPtr(TypeAndMut { ty: d, .. })) => {
                self.unsize_into_ptr(src, dest, s, d)
            }
            (&ty::Adt(def_a, _), &ty::Adt(def_b, _)) => {
                assert_eq!(def_a, def_b);
                if def_a.is_box() || def_b.is_box() {
                    if !def_a.is_box() || !def_b.is_box() {
                        bug!("invalid unsizing between {:?} -> {:?}", src.layout, dest.layout);
                    }
                    return self.unsize_into_ptr(
                        src,
                        dest,
                        src.layout.ty.boxed_ty(),
                        dest.layout.ty.boxed_ty(),
                    );
                }

                // unsizing of generic struct with pointer fields
                // Example: `Arc<T>` -> `Arc<Trait>`
                // here we need to increase the size of every &T thin ptr field to a fat ptr
                for i in 0..src.layout.fields.count() {
                    let dst_field = self.place_field(dest, i as u64)?;
                    if dst_field.layout.is_zst() {
                        continue;
                    }
                    let src_field = self.operand_field(src, i as u64)?;
                    if src_field.layout.ty == dst_field.layout.ty {
                        self.copy_op(src_field, dst_field)?;
                    } else {
                        self.unsize_into(src_field, dst_field)?;
                    }
                }
                Ok(())
            }
            _ => {
                bug!(
                    "unsize_into: invalid conversion: {:?} -> {:?}",
                    src.layout,
                    dest.layout
                )
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 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.

use prelude::v1::*;
use io::prelude::*;

use io;

#[cfg(target_pointer_width = "64")]
pub const HEX_WIDTH: usize = 18;

#[cfg(target_pointer_width = "32")]
pub const HEX_WIDTH: usize = 10;

// All rust symbols are in theory lists of "::"-separated identifiers. Some
// assemblers, however, can't handle these characters in symbol names. To get
// around this, we use C++-style mangling. The mangling method is:
//
// 1. Prefix the symbol with "_ZN"
// 2. For each element of the path, emit the length plus the element
// 3. End the path with "E"
//
// For example, "_ZN4testE" => "test" and "_ZN3foo3bar" => "foo::bar".
//
// We're the ones printing our backtraces, so we can't rely on anything else to
// demangle our symbols. It's *much* nicer to look at demangled symbols, so
// this function is implemented to give us nice pretty output.
//
// Note that this demangler isn't quite as fancy as it could be. We have lots
// of other information in our symbols like hashes, version, type information,
// etc. Additionally, this doesn't handle glue symbols at all.
pub fn demangle(writer: &mut Write, s: &str) -> io::Result<()> {
    // First validate the symbol. If it doesn't look like anything we're
    // expecting, we just print it literally. Note that we must handle non-rust
    // symbols because we could have any function in the backtrace.
    let mut valid = true;
    let mut inner = s;
    if s.len() > 4 && s.starts_with("_ZN") && s.ends_with("E") {
        inner = &s[3 .. s.len() - 1];
    // On Windows, dbghelp strips leading underscores, so we accept "ZN...E" form too.
    } else if s.len() > 3 && s.starts_with("ZN") && s.ends_with("E") {
        inner = &s[2 .. s.len() - 1];
    } else {
        valid = false;
    }

    if valid {
        let mut chars = inner.chars();
        while valid {
            let mut i = 0;
            for c in chars.by_ref() {
                if c.is_numeric() {
                    i = i * 10 + c as usize - '0' as usize;
                } else {
                    break
                }
            }
            if i == 0 {
                valid = chars.next().is_none();
                break
            } else if chars.by_ref().take(i - 1).count() != i - 1 {
                valid = false;
            }
        }
    }

    // Alright, let's do this.
    if !valid {
        try!(writer.write_all(s.as_bytes()));
    } else {
        let mut first = true;
        while !inner.is_empty() {
            if !first {
                try!(writer.write_all(b"::"));
            } else {
                first = false;
            }
            let mut rest = inner;
            while rest.char_at(0).is_numeric() {
                rest = &rest[1..];
            }
            let i: usize = inner[.. (inner.len() - rest.len())].parse().unwrap();
            inner = &rest[i..];
            rest = &rest[..i];
            while !rest.is_empty() {
                if rest.starts_with("$") {
                    macro_rules! demangle {
                        ($($pat:expr, => $demangled:expr),*) => ({
                            $(if rest.starts_with($pat) {
                                try!(writer.write_all($demangled));
                                rest = &rest[$pat.len()..];
                              } else)*
                            {
                                try!(writer.write_all(rest.as_bytes()));
                                break;
                            }

                        })
                    }

                    // see src/librustc/back/link.rs for these mappings
                    demangle! (
                        "$SP$", => b"@",
                        "$BP$", => b"*",
                        "$RF$", => b"&",
                        "$LT$", => b"<",
                        "$GT$", => b">",
                        "$LP$", => b"(",
                        "$RP$", => b")",
                        "$C$", => b",",

                        // in theory we can demangle any Unicode code point, but
                        // for simplicity we just catch the common ones.
                        "$u7e$", => b"~",
                        "$u20$", => b" ",
                        "$u27$", => b"'",
                        "$u5b$", => b"[",
                        "$u5d$", => b"]"
                    )
                } else {
                    let idx = match rest.find('$') {
                        None => rest.len(),
                        Some(i) => i,
                    };
                    try!(writer.write_all(rest[..idx].as_bytes()));
                    rest = &rest[idx..];
                }
            }
        }
    }

    Ok(())
}
<commit_msg>fix example in comments about demangling<commit_after>// Copyright 2014 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.

use prelude::v1::*;
use io::prelude::*;

use io;

#[cfg(target_pointer_width = "64")]
pub const HEX_WIDTH: usize = 18;

#[cfg(target_pointer_width = "32")]
pub const HEX_WIDTH: usize = 10;

// All rust symbols are in theory lists of "::"-separated identifiers. Some
// assemblers, however, can't handle these characters in symbol names. To get
// around this, we use C++-style mangling. The mangling method is:
//
// 1. Prefix the symbol with "_ZN"
// 2. For each element of the path, emit the length plus the element
// 3. End the path with "E"
//
// For example, "_ZN4testE" => "test" and "_ZN3foo3barE" => "foo::bar".
//
// We're the ones printing our backtraces, so we can't rely on anything else to
// demangle our symbols. It's *much* nicer to look at demangled symbols, so
// this function is implemented to give us nice pretty output.
//
// Note that this demangler isn't quite as fancy as it could be. We have lots
// of other information in our symbols like hashes, version, type information,
// etc. Additionally, this doesn't handle glue symbols at all.
pub fn demangle(writer: &mut Write, s: &str) -> io::Result<()> {
    // First validate the symbol. If it doesn't look like anything we're
    // expecting, we just print it literally. Note that we must handle non-rust
    // symbols because we could have any function in the backtrace.
    let mut valid = true;
    let mut inner = s;
    if s.len() > 4 && s.starts_with("_ZN") && s.ends_with("E") {
        inner = &s[3 .. s.len() - 1];
    // On Windows, dbghelp strips leading underscores, so we accept "ZN...E" form too.
    } else if s.len() > 3 && s.starts_with("ZN") && s.ends_with("E") {
        inner = &s[2 .. s.len() - 1];
    } else {
        valid = false;
    }

    if valid {
        let mut chars = inner.chars();
        while valid {
            let mut i = 0;
            for c in chars.by_ref() {
                if c.is_numeric() {
                    i = i * 10 + c as usize - '0' as usize;
                } else {
                    break
                }
            }
            if i == 0 {
                valid = chars.next().is_none();
                break
            } else if chars.by_ref().take(i - 1).count() != i - 1 {
                valid = false;
            }
        }
    }

    // Alright, let's do this.
    if !valid {
        try!(writer.write_all(s.as_bytes()));
    } else {
        let mut first = true;
        while !inner.is_empty() {
            if !first {
                try!(writer.write_all(b"::"));
            } else {
                first = false;
            }
            let mut rest = inner;
            while rest.char_at(0).is_numeric() {
                rest = &rest[1..];
            }
            let i: usize = inner[.. (inner.len() - rest.len())].parse().unwrap();
            inner = &rest[i..];
            rest = &rest[..i];
            while !rest.is_empty() {
                if rest.starts_with("$") {
                    macro_rules! demangle {
                        ($($pat:expr, => $demangled:expr),*) => ({
                            $(if rest.starts_with($pat) {
                                try!(writer.write_all($demangled));
                                rest = &rest[$pat.len()..];
                              } else)*
                            {
                                try!(writer.write_all(rest.as_bytes()));
                                break;
                            }

                        })
                    }

                    // see src/librustc/back/link.rs for these mappings
                    demangle! (
                        "$SP$", => b"@",
                        "$BP$", => b"*",
                        "$RF$", => b"&",
                        "$LT$", => b"<",
                        "$GT$", => b">",
                        "$LP$", => b"(",
                        "$RP$", => b")",
                        "$C$", => b",",

                        // in theory we can demangle any Unicode code point, but
                        // for simplicity we just catch the common ones.
                        "$u7e$", => b"~",
                        "$u20$", => b" ",
                        "$u27$", => b"'",
                        "$u5b$", => b"[",
                        "$u5d$", => b"]"
                    )
                } else {
                    let idx = match rest.find('$') {
                        None => rest.len(),
                        Some(i) => i,
                    };
                    try!(writer.write_all(rest[..idx].as_bytes()));
                    rest = &rest[idx..];
                }
            }
        }
    }

    Ok(())
}
<|endoftext|>


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

use std::ops::Deref;
use std::ops::DerefMut;
use std::io::{self, Write};
use std::{slice, thread};
use std::net::{ToSocketAddrs, TcpStream, SocketAddr};
use std::time::{Duration, Instant};
use std::collections::hash_map::Entry;
use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};

use time;
use rand::{self, ThreadRng};
use protobuf::Message;
use prometheus::{self, Encoder, TextEncoder};
use log::{self, Log, LogMetadata, LogRecord, SetLoggerError};

#[macro_use]
pub mod macros;
pub mod logger;
pub mod panic_hook;
pub mod worker;
pub mod codec;
pub mod xeval;
#[macro_use]
pub mod event;
pub mod rocksdb;
pub mod config;
pub mod fs;
pub mod buf;
pub mod sockopt;
pub mod transport;
pub mod time_monitor;
pub mod file_log;

pub use self::fs::{DiskStat, get_disk_stat};
pub use log::LogLevelFilter;

pub fn init_log(level: LogLevelFilter) -> Result<(), SetLoggerError> {
    log::set_logger(|filter| {
        filter.set(level);
        Box::new(DefaultLogger { level: level })
    })
}

struct DefaultLogger {
    level: LogLevelFilter,
}

impl Log for DefaultLogger {
    fn enabled(&self, meta: &LogMetadata) -> bool {
        meta.level() <= self.level
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            let t = time::now();
            // TODO allow formatter to be configurable.
            let _ = write!(io::stderr(),
                           "{},{:03} {}:{} - {:5} - {}\n",
                           time::strftime("%Y-%m-%d %H:%M:%S", &t).unwrap(),
                           t.tm_nsec / 1000_000,
                           record.location().file().rsplit('/').nth(0).unwrap(),
                           record.location().line(),
                           record.level(),
                           record.args());
        }
    }
}

pub fn limit_size<T: Message + Clone>(entries: &[T], max: u64) -> Vec<T> {
    if entries.len() == 0 {
        return entries.to_vec();
    }

    let mut size = Message::compute_size(&entries[0]) as u64;
    let mut limit = 1usize;
    while limit < entries.len() {
        size += Message::compute_size(&entries[limit]) as u64;
        if size > max {
            break;
        }
        limit += 1;
    }
    entries[..limit].to_vec()
}

pub struct DefaultRng {
    rng: ThreadRng,
}

impl DefaultRng {
    fn new() -> DefaultRng {
        DefaultRng { rng: rand::thread_rng() }
    }
}

impl Default for DefaultRng {
    fn default() -> DefaultRng {
        DefaultRng::new()
    }
}

impl Deref for DefaultRng {
    type Target = ThreadRng;

    fn deref(&self) -> &ThreadRng {
        &self.rng
    }
}

impl DerefMut for DefaultRng {
    fn deref_mut(&mut self) -> &mut ThreadRng {
        &mut self.rng
    }
}

/// A handy shortcut to replace `RwLock` write/read().unwrap() pattern to
/// shortcut wl and rl.
pub trait HandyRwLock<T> {
    fn wl(&self) -> RwLockWriteGuard<T>;
    fn rl(&self) -> RwLockReadGuard<T>;
}

impl<T> HandyRwLock<T> for RwLock<T> {
    fn wl(&self) -> RwLockWriteGuard<T> {
        self.write().unwrap()
    }

    fn rl(&self) -> RwLockReadGuard<T> {
        self.read().unwrap()
    }
}

pub fn make_std_tcp_conn<A: ToSocketAddrs>(addr: A) -> io::Result<TcpStream> {
    let stream = try!(TcpStream::connect(addr));
    try!(stream.set_nodelay(true));
    Ok(stream)
}

// A helper function to parse SocketAddr for mio.
// In mio example, it uses "127.0.0.1:80".parse() to get the SocketAddr,
// but it is just ok for "ip:port", not "host:port".
pub fn to_socket_addr<A: ToSocketAddrs>(addr: A) -> io::Result<SocketAddr> {
    let addrs = try!(addr.to_socket_addrs());
    Ok(addrs.collect::<Vec<SocketAddr>>()[0])
}

/// A function to escape a byte array to a readable ascii string.
///
/// # Examples
///
/// ```
/// use tikv::util::escape;
///
/// assert_eq!(r"ab", escape(b"ab"));
/// assert_eq!(r"a\\023", escape(b"a\\023"));
/// assert_eq!("a\\r\\n\\t '\\\"\\\\", escape(b"a\r\n\t '\"\\"));
/// assert_eq!(r"\342\235\244\360\237\220\267", escape("❤🐷".as_bytes()));
/// ```
pub fn escape(data: &[u8]) -> String {
    let mut escaped = Vec::with_capacity(data.len() * 4);
    for &c in data {
        match c {
            b'\n' => escaped.extend_from_slice(br"\n"),
            b'\r' => escaped.extend_from_slice(br"\r"),
            b'\t' => escaped.extend_from_slice(br"\t"),
            b'"' => escaped.extend_from_slice(b"\\\""),
            b'\\' => escaped.extend_from_slice(br"\\"),
            b'\x20'...b'\x7e' => escaped.push(c),
            _ => {
                escaped.push(b'\\');
                escaped.push(b'0' + (c >> 6));
                escaped.push(b'0' + ((c >> 3) & 7));
                escaped.push(b'0' + (c & 7));
            }
        }
    }
    escaped.shrink_to_fit();
    unsafe { String::from_utf8_unchecked(escaped) }
}

/// A function to unescape an escaped string to a byte array.
///
/// # Panic
///
/// If s is not a properly encoded string.
///
/// # Examples
///
/// ```
/// use tikv::util::unescape;
///
/// assert_eq!(unescape(r"ab"), b"ab");
/// assert_eq!(unescape(r"a\\023"), b"a\\023");
/// assert_eq!(unescape("a\\r\\n\\t '\\\"\\\\"), b"a\r\n\t '\"\\");
/// assert_eq!(unescape(r"\342\235\244\360\237\220\267"), "❤🐷".as_bytes());
/// ```
///
pub fn unescape(s: &str) -> Vec<u8> {
    let mut buf = Vec::with_capacity(s.len());
    let mut bytes = s.bytes();
    loop {
        let b = match bytes.next() {
            None => break,
            Some(t) => t,
        };
        if b != b'\\' {
            buf.push(b);
            continue;
        }
        match bytes.next().unwrap() {
            b'"' => buf.push(b'"'),
            b'\'' => buf.push(b'\''),
            b'\\' => buf.push(b'\\'),
            b'n' => buf.push(b'\n'),
            b't' => buf.push(b'\t'),
            b'r' => buf.push(b'\r'),
            b => {
                let b1 = b - b'0';
                let b2 = bytes.next().unwrap() - b'0';
                let b3 = bytes.next().unwrap() - b'0';
                buf.push((b1 << 6) + (b2 << 3) + b3);
            }
        }
    }
    buf.shrink_to_fit();
    buf
}

/// Convert a borrow to a slice.
pub fn as_slice<T>(t: &T) -> &[T] {
    unsafe {
        let ptr = t as *const T;
        slice::from_raw_parts(ptr, 1)
    }
}

pub struct SlowTimer {
    slow_time: Duration,
    t: Instant,
}

impl SlowTimer {
    pub fn new() -> SlowTimer {
        SlowTimer::default()
    }

    pub fn from(slow_time: Duration) -> SlowTimer {
        SlowTimer {
            slow_time: slow_time,
            t: Instant::now(),
        }
    }

    pub fn from_secs(secs: u64) -> SlowTimer {
        SlowTimer::from(Duration::from_secs(secs))
    }

    pub fn from_millis(millis: u64) -> SlowTimer {
        SlowTimer::from(Duration::from_millis(millis))
    }

    pub fn elapsed(&self) -> Duration {
        self.t.elapsed()
    }

    pub fn is_slow(&self) -> bool {
        self.elapsed() >= self.slow_time
    }
}

const DEFAULT_SLOW_SECS: u64 = 1;

impl Default for SlowTimer {
    fn default() -> SlowTimer {
        SlowTimer::from_secs(DEFAULT_SLOW_SECS)
    }
}

/// `TryInsertWith` is a helper trait for `Entry` to accept a failable closure.
pub trait TryInsertWith<'a, V, E> {
    fn or_try_insert_with<F: FnOnce() -> Result<V, E>>(self, default: F) -> Result<&'a mut V, E>;
}

impl<'a, T: 'a, V: 'a, E> TryInsertWith<'a, V, E> for Entry<'a, T, V> {
    fn or_try_insert_with<F: FnOnce() -> Result<V, E>>(self, default: F) -> Result<&'a mut V, E> {
        match self {
            Entry::Occupied(e) => Ok(e.into_mut()),
            Entry::Vacant(e) => {
                let v = try!(default());
                Ok(e.insert(v))
            }
        }
    }
}

/// Convert Duration to milliseconds.
pub fn duration_to_ms(d: Duration) -> u64 {
    let nanos = d.subsec_nanos() as u64;
    // Most of case, we can't have so large Duration, so here just panic if overflow now.
    d.as_secs() * 1_000 + (nanos / 1_000_000)
}

pub fn get_tag_from_thread_name() -> Option<String> {
    thread::current().name().and_then(|name| name.split("::").skip(1).last()).map(From::from)
}

/// `DeferContext` will invoke the wrapped closure when dropped.
pub struct DeferContext<T: FnOnce()> {
    t: Option<T>,
}

impl<T: FnOnce()> DeferContext<T> {
    pub fn new(t: T) -> DeferContext<T> {
        DeferContext { t: Some(t) }
    }
}

impl<T: FnOnce()> Drop for DeferContext<T> {
    fn drop(&mut self) {
        self.t.take().unwrap()()
    }
}

/// Represents a value of one of two possible types (a more generic Result.)
#[derive(Debug)]
pub enum Either<L, R> {
    Left(L),
    Right(R),
}

impl<L, R> Either<L, R> {
    #[inline]
    pub fn as_ref(&self) -> Either<&L, &R> {
        match *self {
            Either::Left(ref l) => Either::Left(l),
            Either::Right(ref r) => Either::Right(r),
        }
    }

    #[inline]
    pub fn left(self) -> Option<L> {
        match self {
            Either::Left(l) => Some(l),
            _ => None,
        }
    }

    #[inline]
    pub fn right(self) -> Option<R> {
        match self {
            Either::Right(r) => Some(r),
            _ => None,
        }
    }
}

/// `build_info` returns a tuple of Strings that contains build utc time and commit hash.
fn build_info() -> (String, String) {
    let raw = include_str!(concat!(env!("OUT_DIR"), "/build-info.txt"));
    let mut parts = raw.split('\n');

    (parts.next().unwrap_or("None").to_owned(), parts.next().unwrap_or("None").to_owned())
}

/// `print_tikv_info` prints the tikv version information to the standard output.
pub fn print_tikv_info() {
    let (hash, date) = build_info();
    info!("Welcome to the TiKV.");
    info!("Version:");
    info!("Git Commit Hash: {}", hash);
    info!("UTC Build Time:  {}", date);
}

/// `run_prometheus` runs a background prometheus client.
pub fn run_prometheus(interval: Duration) -> Option<thread::JoinHandle<()>> {
    if interval == Duration::from_secs(0) {
        return None;
    }

    Some(thread::spawn(move || {
        let encoder = TextEncoder::new();
        let mut buffer = Vec::<u8>::new();
        loop {
            let metric_familys = prometheus::gather();
            encoder.encode(&metric_familys, &mut buffer).unwrap();

            // Output to the standard output.
            info!("{}", String::from_utf8(buffer.clone()).unwrap());

            buffer.clear();
            thread::sleep(interval);
        }
    }))
}

#[cfg(test)]
mod tests {
    use std::net::{SocketAddr, AddrParseError};
    use std::time::Duration;
    use std::rc::Rc;
    use std::sync::atomic::{AtomicBool, Ordering};
    use super::*;

    #[test]
    fn test_to_socket_addr() {
        let tbls = vec![
            ("", false),
            ("127.0.0.1", false),
            ("localhost", false),
            ("127.0.0.1:80", true),
            ("localhost:80", true),
        ];

        for (addr, ok) in tbls {
            assert_eq!(to_socket_addr(addr).is_ok(), ok);
        }

        let tbls = vec![
            ("localhost:80", false),
            ("127.0.0.1:80", true),
        ];

        for (addr, ok) in tbls {
            let ret: Result<SocketAddr, AddrParseError> = addr.parse();
            assert_eq!(ret.is_ok(), ok);
        }
    }

    #[test]
    fn test_duration_to_ms() {
        let tbl = vec![0, 100, 1_000, 5_000, 9999, 1_000_000, 1_000_000_000];
        for ms in tbl {
            let d = Duration::from_millis(ms);
            assert_eq!(ms, duration_to_ms(d));
        }
    }

    #[test]
    fn test_defer() {
        let should_panic = Rc::new(AtomicBool::new(true));
        let sp = should_panic.clone();
        defer!(assert!(!sp.load(Ordering::SeqCst)));
        should_panic.store(false, Ordering::SeqCst);
    }
}
<commit_msg>util: make clippy88 happy (#1031)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::ascii;
use std::ops::Deref;
use std::ops::DerefMut;
use std::io::{self, Write};
use std::{slice, thread};
use std::net::{ToSocketAddrs, TcpStream, SocketAddr};
use std::time::{Duration, Instant};
use std::collections::hash_map::Entry;
use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};

use time;
use rand::{self, ThreadRng};
use protobuf::Message;
use prometheus::{self, Encoder, TextEncoder};
use log::{self, Log, LogMetadata, LogRecord, SetLoggerError};

#[macro_use]
pub mod macros;
pub mod logger;
pub mod panic_hook;
pub mod worker;
pub mod codec;
pub mod xeval;
#[macro_use]
pub mod event;
pub mod rocksdb;
pub mod config;
pub mod fs;
pub mod buf;
pub mod sockopt;
pub mod transport;
pub mod time_monitor;
pub mod file_log;

pub use self::fs::{DiskStat, get_disk_stat};
pub use log::LogLevelFilter;

pub fn init_log(level: LogLevelFilter) -> Result<(), SetLoggerError> {
    log::set_logger(|filter| {
        filter.set(level);
        Box::new(DefaultLogger { level: level })
    })
}

struct DefaultLogger {
    level: LogLevelFilter,
}

impl Log for DefaultLogger {
    fn enabled(&self, meta: &LogMetadata) -> bool {
        meta.level() <= self.level
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            let t = time::now();
            // TODO allow formatter to be configurable.
            let _ = write!(io::stderr(),
                           "{},{:03} {}:{} - {:5} - {}\n",
                           time::strftime("%Y-%m-%d %H:%M:%S", &t).unwrap(),
                           t.tm_nsec / 1000_000,
                           record.location().file().rsplit('/').nth(0).unwrap(),
                           record.location().line(),
                           record.level(),
                           record.args());
        }
    }
}

pub fn limit_size<T: Message + Clone>(entries: &[T], max: u64) -> Vec<T> {
    if entries.len() == 0 {
        return entries.to_vec();
    }

    let mut size = Message::compute_size(&entries[0]) as u64;
    let mut limit = 1usize;
    while limit < entries.len() {
        size += Message::compute_size(&entries[limit]) as u64;
        if size > max {
            break;
        }
        limit += 1;
    }
    entries[..limit].to_vec()
}

pub struct DefaultRng {
    rng: ThreadRng,
}

impl DefaultRng {
    fn new() -> DefaultRng {
        DefaultRng { rng: rand::thread_rng() }
    }
}

impl Default for DefaultRng {
    fn default() -> DefaultRng {
        DefaultRng::new()
    }
}

impl Deref for DefaultRng {
    type Target = ThreadRng;

    fn deref(&self) -> &ThreadRng {
        &self.rng
    }
}

impl DerefMut for DefaultRng {
    fn deref_mut(&mut self) -> &mut ThreadRng {
        &mut self.rng
    }
}

/// A handy shortcut to replace `RwLock` write/read().unwrap() pattern to
/// shortcut wl and rl.
pub trait HandyRwLock<T> {
    fn wl(&self) -> RwLockWriteGuard<T>;
    fn rl(&self) -> RwLockReadGuard<T>;
}

impl<T> HandyRwLock<T> for RwLock<T> {
    fn wl(&self) -> RwLockWriteGuard<T> {
        self.write().unwrap()
    }

    fn rl(&self) -> RwLockReadGuard<T> {
        self.read().unwrap()
    }
}

pub fn make_std_tcp_conn<A: ToSocketAddrs>(addr: A) -> io::Result<TcpStream> {
    let stream = try!(TcpStream::connect(addr));
    try!(stream.set_nodelay(true));
    Ok(stream)
}

// A helper function to parse SocketAddr for mio.
// In mio example, it uses "127.0.0.1:80".parse() to get the SocketAddr,
// but it is just ok for "ip:port", not "host:port".
pub fn to_socket_addr<A: ToSocketAddrs>(addr: A) -> io::Result<SocketAddr> {
    let addrs = try!(addr.to_socket_addrs());
    Ok(addrs.collect::<Vec<SocketAddr>>()[0])
}

/// A function to escape a byte array to a readable ascii string.
/// escape rules follow golang/protobuf.
/// https://github.com/golang/protobuf/blob/master/proto/text.go#L578
///
/// # Examples
///
/// ```
/// use tikv::util::escape;
///
/// assert_eq!(r"ab", escape(b"ab"));
/// assert_eq!(r"a\\023", escape(b"a\\023"));
/// assert_eq!("a\\r\\n\\t '\\\"\\\\", escape(b"a\r\n\t '\"\\"));
/// assert_eq!(r"\342\235\244\360\237\220\267", escape("❤🐷".as_bytes()));
/// ```
pub fn escape(data: &[u8]) -> String {
    let mut escaped = Vec::with_capacity(data.len() * 4);
    for &c in data {
        match c {
            b'"' => escaped.extend_from_slice(b"\\\""),
            b'\\' => escaped.extend_from_slice(br"\\"),
            b'\'' => escaped.push(b'\''),
            _ if c > b'\x7e' => {
                escaped.push(b'\\');
                escaped.push(b'0' + (c >> 6));
                escaped.push(b'0' + ((c >> 3) & 7));
                escaped.push(b'0' + (c & 7));
            }
            _ => escaped.extend(ascii::escape_default(c)),
        }
    }
    escaped.shrink_to_fit();
    unsafe { String::from_utf8_unchecked(escaped) }
}

/// A function to unescape an escaped string to a byte array.
///
/// # Panic
///
/// If s is not a properly encoded string.
///
/// # Examples
///
/// ```
/// use tikv::util::unescape;
///
/// assert_eq!(unescape(r"ab"), b"ab");
/// assert_eq!(unescape(r"a\\023"), b"a\\023");
/// assert_eq!(unescape("a\\r\\n\\t '\\\"\\\\"), b"a\r\n\t '\"\\");
/// assert_eq!(unescape(r"\342\235\244\360\237\220\267"), "❤🐷".as_bytes());
/// ```
///
pub fn unescape(s: &str) -> Vec<u8> {
    let mut buf = Vec::with_capacity(s.len());
    let mut bytes = s.bytes();
    loop {
        let b = match bytes.next() {
            None => break,
            Some(t) => t,
        };
        if b != b'\\' {
            buf.push(b);
            continue;
        }
        match bytes.next().unwrap() {
            b'"' => buf.push(b'"'),
            b'\'' => buf.push(b'\''),
            b'\\' => buf.push(b'\\'),
            b'n' => buf.push(b'\n'),
            b't' => buf.push(b'\t'),
            b'r' => buf.push(b'\r'),
            b => {
                let b1 = b - b'0';
                let b2 = bytes.next().unwrap() - b'0';
                let b3 = bytes.next().unwrap() - b'0';
                buf.push((b1 << 6) + (b2 << 3) + b3);
            }
        }
    }
    buf.shrink_to_fit();
    buf
}

/// Convert a borrow to a slice.
pub fn as_slice<T>(t: &T) -> &[T] {
    unsafe {
        let ptr = t as *const T;
        slice::from_raw_parts(ptr, 1)
    }
}

pub struct SlowTimer {
    slow_time: Duration,
    t: Instant,
}

impl SlowTimer {
    pub fn new() -> SlowTimer {
        SlowTimer::default()
    }

    pub fn from(slow_time: Duration) -> SlowTimer {
        SlowTimer {
            slow_time: slow_time,
            t: Instant::now(),
        }
    }

    pub fn from_secs(secs: u64) -> SlowTimer {
        SlowTimer::from(Duration::from_secs(secs))
    }

    pub fn from_millis(millis: u64) -> SlowTimer {
        SlowTimer::from(Duration::from_millis(millis))
    }

    pub fn elapsed(&self) -> Duration {
        self.t.elapsed()
    }

    pub fn is_slow(&self) -> bool {
        self.elapsed() >= self.slow_time
    }
}

const DEFAULT_SLOW_SECS: u64 = 1;

impl Default for SlowTimer {
    fn default() -> SlowTimer {
        SlowTimer::from_secs(DEFAULT_SLOW_SECS)
    }
}

/// `TryInsertWith` is a helper trait for `Entry` to accept a failable closure.
pub trait TryInsertWith<'a, V, E> {
    fn or_try_insert_with<F: FnOnce() -> Result<V, E>>(self, default: F) -> Result<&'a mut V, E>;
}

impl<'a, T: 'a, V: 'a, E> TryInsertWith<'a, V, E> for Entry<'a, T, V> {
    fn or_try_insert_with<F: FnOnce() -> Result<V, E>>(self, default: F) -> Result<&'a mut V, E> {
        match self {
            Entry::Occupied(e) => Ok(e.into_mut()),
            Entry::Vacant(e) => {
                let v = try!(default());
                Ok(e.insert(v))
            }
        }
    }
}

/// Convert Duration to milliseconds.
pub fn duration_to_ms(d: Duration) -> u64 {
    let nanos = d.subsec_nanos() as u64;
    // Most of case, we can't have so large Duration, so here just panic if overflow now.
    d.as_secs() * 1_000 + (nanos / 1_000_000)
}

pub fn get_tag_from_thread_name() -> Option<String> {
    thread::current().name().and_then(|name| name.split("::").skip(1).last()).map(From::from)
}

/// `DeferContext` will invoke the wrapped closure when dropped.
pub struct DeferContext<T: FnOnce()> {
    t: Option<T>,
}

impl<T: FnOnce()> DeferContext<T> {
    pub fn new(t: T) -> DeferContext<T> {
        DeferContext { t: Some(t) }
    }
}

impl<T: FnOnce()> Drop for DeferContext<T> {
    fn drop(&mut self) {
        self.t.take().unwrap()()
    }
}

/// Represents a value of one of two possible types (a more generic Result.)
#[derive(Debug)]
pub enum Either<L, R> {
    Left(L),
    Right(R),
}

impl<L, R> Either<L, R> {
    #[inline]
    pub fn as_ref(&self) -> Either<&L, &R> {
        match *self {
            Either::Left(ref l) => Either::Left(l),
            Either::Right(ref r) => Either::Right(r),
        }
    }

    #[inline]
    pub fn left(self) -> Option<L> {
        match self {
            Either::Left(l) => Some(l),
            _ => None,
        }
    }

    #[inline]
    pub fn right(self) -> Option<R> {
        match self {
            Either::Right(r) => Some(r),
            _ => None,
        }
    }
}

/// `build_info` returns a tuple of Strings that contains build utc time and commit hash.
fn build_info() -> (String, String) {
    let raw = include_str!(concat!(env!("OUT_DIR"), "/build-info.txt"));
    let mut parts = raw.split('\n');

    (parts.next().unwrap_or("None").to_owned(), parts.next().unwrap_or("None").to_owned())
}

/// `print_tikv_info` prints the tikv version information to the standard output.
pub fn print_tikv_info() {
    let (hash, date) = build_info();
    info!("Welcome to the TiKV.");
    info!("Version:");
    info!("Git Commit Hash: {}", hash);
    info!("UTC Build Time:  {}", date);
}

/// `run_prometheus` runs a background prometheus client.
pub fn run_prometheus(interval: Duration) -> Option<thread::JoinHandle<()>> {
    if interval == Duration::from_secs(0) {
        return None;
    }

    Some(thread::spawn(move || {
        let encoder = TextEncoder::new();
        let mut buffer = Vec::<u8>::new();
        loop {
            let metric_familys = prometheus::gather();
            encoder.encode(&metric_familys, &mut buffer).unwrap();

            // Output to the standard output.
            info!("{}", String::from_utf8(buffer.clone()).unwrap());

            buffer.clear();
            thread::sleep(interval);
        }
    }))
}

#[cfg(test)]
mod tests {
    use std::net::{SocketAddr, AddrParseError};
    use std::time::Duration;
    use std::rc::Rc;
    use std::sync::atomic::{AtomicBool, Ordering};
    use super::*;

    #[test]
    fn test_to_socket_addr() {
        let tbls = vec![
            ("", false),
            ("127.0.0.1", false),
            ("localhost", false),
            ("127.0.0.1:80", true),
            ("localhost:80", true),
        ];

        for (addr, ok) in tbls {
            assert_eq!(to_socket_addr(addr).is_ok(), ok);
        }

        let tbls = vec![
            ("localhost:80", false),
            ("127.0.0.1:80", true),
        ];

        for (addr, ok) in tbls {
            let ret: Result<SocketAddr, AddrParseError> = addr.parse();
            assert_eq!(ret.is_ok(), ok);
        }
    }

    #[test]
    fn test_duration_to_ms() {
        let tbl = vec![0, 100, 1_000, 5_000, 9999, 1_000_000, 1_000_000_000];
        for ms in tbl {
            let d = Duration::from_millis(ms);
            assert_eq!(ms, duration_to_ms(d));
        }
    }

    #[test]
    fn test_defer() {
        let should_panic = Rc::new(AtomicBool::new(true));
        let sp = should_panic.clone();
        defer!(assert!(!sp.load(Ordering::SeqCst)));
        should_panic.store(false, Ordering::SeqCst);
    }
}
<|endoftext|>