<commit_before>use std::fmt;
use std::slice;
use std::mem;
use num::{Float, NumCast};
use channel::{ColorChannel, BoundedChannel, cast};
use approx;

use color;

pub struct RgbTag;

#[derive(Clone, Debug, PartialEq, PartialOrd, Eq, Ord, Hash)]
pub struct Rgb<T> {
    r: BoundedChannel<T>,
    g: BoundedChannel<T>,
    b: BoundedChannel<T>,
}

impl<T: ColorChannel> Rgb<T> {
    #[inline]
    pub fn new() -> Self {
        Rgb{r: BoundedChannel::min(),
            g: BoundedChannel::min(),
            b: BoundedChannel::min()}
    }

    #[inline]
    pub fn from_channels(red: T, green: T, blue: T) -> Self {
        Rgb{r: BoundedChannel::new(red),
            g: BoundedChannel::new(green),
            b: BoundedChannel::new(blue)}
    }

    #[inline]
    pub fn from_array(values: &[T; 3]) -> Self {
        Rgb{r: BoundedChannel::new(values[0]),
            g: BoundedChannel::new(values[1]),
            b: BoundedChannel::new(values[2])}
    }

    #[inline]
    pub fn red(&self) -> T {
        self.r.0.clone()
    }
    #[inline]
    pub fn green(&self) -> T {
        self.g.0.clone()
    }
    #[inline]
    pub fn blue(&self) -> T {
        self.b.0.clone()
    }
    #[inline]
    pub fn red_mut(&mut self) -> &mut T {
        &mut self.r.0
    }
    #[inline]
    pub fn green_mut(&mut self) -> &mut T {
        &mut self.g.0
    }
    #[inline]
    pub fn blue_mut(&mut self) -> &mut T {
        &mut self.b.0
    }
    #[inline]
    pub fn red_chan(&self) -> &BoundedChannel<T> {
        &self.r
    }
    #[inline]
    pub fn green_chan(&self) -> &BoundedChannel<T> {
        &self.g
    }
    #[inline]
    pub fn blue_chan(&self) -> &BoundedChannel<T> {
        &self.b
    }

    #[inline]
    pub fn lerp<P>(&self, right: &Self, pos: P) -> Self
            where P: Float + NumCast {
        assert!(pos <= cast(1.0).unwrap() && pos >= cast(0.0).unwrap());
        Rgb{
            r: self.r.lerp(right.r, pos),
            g: self.g.lerp(right.g, pos),
            b: self.b.lerp(right.b, pos)
        }
    }

    pub fn channel_labels(&self) -> &'static str {
        "rgb"
    }
}

impl<T: ColorChannel> color::Color for Rgb<T> {
    type Component = T;
    type Tag = RgbTag;

    fn num_channels() -> u32 {
        3
    }

    fn from_slice(values: &[Self::Component]) -> Self {
        Rgb{r: BoundedChannel::new(values[0]),
            g: BoundedChannel::new(values[1]),
            b: BoundedChannel::new(values[2])}
    }

    fn as_slice(&self) -> &[Self::Component] {
        unsafe {
            let ptr: *const Self::Component = mem::transmute(self);
            slice::from_raw_parts(ptr, Self::num_channels() as usize)
        }
    }

    fn broadcast(value: T) -> Self {
        Rgb::from_channels(value.clone(), value.clone(), value.clone())
    }
}

impl<T: ColorChannel> color::Color3 for Rgb<T> {
    fn as_tuple(&self) -> (Self::Component, Self::Component, Self::Component) {
        (self.r.0, self.g.0, self.b.0)
    }

    fn as_array(&self) -> [Self::Component; 3] {
        [self.r.0, self.g.0, self.b.0]
    }

    #[inline]
    fn from_tuple(values: &(Self::Component, Self::Component, Self::Component)) -> Self {
        Rgb{r: BoundedChannel::new(values.0),
            g: BoundedChannel::new(values.1),
            b: BoundedChannel::new(values.2)}
    }

}

impl<T: ColorChannel> color::ComponentMap for Rgb<T> {
    fn component_map<F> (&self, mut f: F) -> Self
            where F: FnMut(Self::Component) -> Self::Component {
        Rgb{r: BoundedChannel::new(f(self.red().clone())),
            g: BoundedChannel::new(f(self.green().clone())),
            b: BoundedChannel::new(f(self.blue().clone()))}
    }

    fn component_map_binary<F>(&self, other: &Self, mut f: F) -> Self
            where F: FnMut(Self::Component, Self::Component) -> Self::Component {
        Rgb{r: BoundedChannel::new(
                f(self.red(), other.red())),
            g: BoundedChannel::new(
                f(self.green(), other.green())),
            b: BoundedChannel::new(
                f(self.blue(), other.blue()))}
    }
}

impl<T: ColorChannel> color::Invert for Rgb<T> {
    fn invert(&self) -> Self {
        Rgb{
            r: self.r.clone().invert(),
            g: self.g.clone().invert(),
            b: self.b.clone().invert(),
        }
    }
}

impl<T: ColorChannel> color::Bounded for Rgb<T> {
    fn clamp(&self, min: Self::Component, max: Self::Component) -> Self {
        Rgb{r: self.r.clamp(min, max),
            g: self.g.clamp(min, max),
            b: self.b.clamp(min, max)}
    }

    fn normalize(&self) -> Self {
        Rgb{r: self.r.normalize(),
            g: self.g.normalize(),
            b: self.b.normalize()}
    }

    fn is_normalized(&self) -> bool {
        self.r.is_normalized() && self.g.is_normalized() && self.b.is_normalized()
    }
}

impl<T: ColorChannel> Default for Rgb<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T: ColorChannel + Float + approx::ApproxEq> approx::ApproxEq for Rgb<T>
        where T::Epsilon: Clone {
    type Epsilon = T::Epsilon;

    fn default_epsilon() -> Self::Epsilon {
        T::default_epsilon()
    }

    fn default_max_relative() -> Self::Epsilon {
        T::default_max_relative()
    }

    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    fn relative_eq(&self, other: &Self, epsilon: Self::Epsilon,
           max_relative: Self::Epsilon) -> bool {
        self.red().relative_eq(&other.red(), epsilon.clone(), max_relative.clone())
        && self.green().relative_eq(&other.green(), epsilon.clone(), max_relative.clone())
        && self.blue().relative_eq(&other.blue(), epsilon, max_relative)
    }

    fn ulps_eq(&self, other: &Self, epsilon: Self::Epsilon, max_ulps: u32) -> bool {
        self.red().ulps_eq(&other.red(), epsilon.clone(), max_ulps)
        && self.green().ulps_eq(&other.green(), epsilon.clone(), max_ulps)
        && self.blue().ulps_eq(&other.blue(), epsilon.clone(), max_ulps)
    }

}

impl<T: ColorChannel + fmt::Display> fmt::Display for Rgb<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Rgb({}, {}, {})", self.r, self.g, self.b)
    }
}

#[cfg(test)]
mod test {
    use ::rgb::*;
    use ::color;
    use ::color::{Color, Invert};

    #[test]
    fn test_construct() {
        {
            let color = Rgb::from_channels(0u8, 0, 0);
            assert_eq!(color.red(), 0u8);
            assert_eq!(color.green(), 0u8);
            assert_eq!(color.blue(), 0u8);

            let c2 = color.clone();
            assert_eq!(color, c2);
        }
        {
            let color: Rgb<u8> = Rgb::new();
            assert_eq!(color.red(), 0u8);
            assert_eq!(color.green(), 0u8);
            assert_eq!(color.blue(), 0u8);
        }
        {
            let color = Rgb::broadcast(0.5_f32);
            assert_ulps_eq!(color, Rgb::from_channels(0.5_f32, 0.5, 0.5));
        }
    }

    #[test]
    fn test_lerp_int() {
        let c1 = Rgb::from_channels(100u8, 200u8, 0u8);
        let c2 = Rgb::from_channels(200u8, 0u8, 255u8);

        assert_eq!(c1.lerp(&c2, 0.5_f64), Rgb::from_channels(150u8, 100, 127));
        assert_eq!(c1.lerp(&c2, 0.0_f64), c1);
        assert_eq!(c1.lerp(&c2, 1.0_f64), c2);
    }

    #[test]
    fn test_lerp_float() {
        let c1 = Rgb::from_channels(0.2_f32, 0.5, 1.0);
        let c2 = Rgb::from_channels(0.8_f32, 0.5, 0.1);

        assert_ulps_eq!(c1.lerp(&c2, 0.5_f32), Rgb::from_channels(0.5_f32, 0.5, 0.55));
        assert_ulps_eq!(c1.lerp(&c2, 0.0_f64), Rgb::from_channels(0.2_f32, 0.5, 1.0));
        assert_ulps_eq!(c1.lerp(&c2, 1.0_f32), Rgb::from_channels(0.8_f32, 0.5, 0.1));
    }

    #[test]
    fn test_invert() {
        let c = Rgb::from_channels(200u8, 0, 255);
        let c2 = Rgb::from_channels(0.8_f32, 0.0, 0.25);

        assert_eq!(c.invert(), Rgb::from_channels(55u8, 255, 0));
        assert_ulps_eq!(c2.invert(), Rgb::from_channels(0.2_f32, 1.0, 0.75));
    }

    #[test]
    fn as_slice() {
        let c = Rgb::from_channels(100u8, 0, 125);
        let c2 = Rgb::from_channels(1.0, 0.25, 0.125);

        assert_eq!(c.as_slice()[0], 100u8);
        assert_eq!(c.as_slice()[1], 0u8);
        assert_eq!(c.as_slice()[2], 125u8);

        assert_ulps_eq!(Rgb::from_slice(c2.as_slice()), c2);
    }

    #[test]
    fn color_cast() {
        let c = Rgb::from_channels(127, 0, 255);
        let c2 = color::color_cast::<Rgb<f32>, _>(&c);
        let c3 = color::color_cast::<Rgb<u8>, _>(&c2);

        assert_ulps_eq!(c2.red(), 127.0 / 255.0);
        assert_ulps_eq!(c2.green(), 0.0);
        assert_ulps_eq!(c2.blue(), 1.0);

        assert_eq!(c3.red(), 127);
        assert_eq!(c3.green(), 0);
        assert_eq!(c3.blue(), 255);

        println!("{}", c2);
    }
}
<commit_msg>Added convenience type `Rgba`.<commit_after>use std::fmt;
use std::slice;
use std::mem;
use num::{Float, NumCast};
use channel::{ColorChannel, BoundedChannel, cast};
use approx;
use alpha::Alpha;

use color;

pub struct RgbTag;

#[derive(Clone, Debug, PartialEq, PartialOrd, Eq, Ord, Hash)]
pub struct Rgb<T> {
    r: BoundedChannel<T>,
    g: BoundedChannel<T>,
    b: BoundedChannel<T>,
}

pub type Rgba<T> = Alpha<T, Rgb<T>>;

impl<T: ColorChannel> Rgb<T> {
    #[inline]
    pub fn new() -> Self {
        Rgb{r: BoundedChannel::min(),
            g: BoundedChannel::min(),
            b: BoundedChannel::min()}
    }

    #[inline]
    pub fn from_channels(red: T, green: T, blue: T) -> Self {
        Rgb{r: BoundedChannel::new(red),
            g: BoundedChannel::new(green),
            b: BoundedChannel::new(blue)}
    }

    #[inline]
    pub fn from_array(values: &[T; 3]) -> Self {
        Rgb{r: BoundedChannel::new(values[0]),
            g: BoundedChannel::new(values[1]),
            b: BoundedChannel::new(values[2])}
    }

    #[inline]
    pub fn red(&self) -> T {
        self.r.0.clone()
    }
    #[inline]
    pub fn green(&self) -> T {
        self.g.0.clone()
    }
    #[inline]
    pub fn blue(&self) -> T {
        self.b.0.clone()
    }
    #[inline]
    pub fn red_mut(&mut self) -> &mut T {
        &mut self.r.0
    }
    #[inline]
    pub fn green_mut(&mut self) -> &mut T {
        &mut self.g.0
    }
    #[inline]
    pub fn blue_mut(&mut self) -> &mut T {
        &mut self.b.0
    }
    #[inline]
    pub fn red_chan(&self) -> &BoundedChannel<T> {
        &self.r
    }
    #[inline]
    pub fn green_chan(&self) -> &BoundedChannel<T> {
        &self.g
    }
    #[inline]
    pub fn blue_chan(&self) -> &BoundedChannel<T> {
        &self.b
    }

    #[inline]
    pub fn lerp<P>(&self, right: &Self, pos: P) -> Self
            where P: Float + NumCast {
        assert!(pos <= cast(1.0).unwrap() && pos >= cast(0.0).unwrap());
        Rgb{
            r: self.r.lerp(right.r, pos),
            g: self.g.lerp(right.g, pos),
            b: self.b.lerp(right.b, pos)
        }
    }

    pub fn channel_labels(&self) -> &'static str {
        "rgb"
    }
}

impl<T: ColorChannel> color::Color for Rgb<T> {
    type Component = T;
    type Tag = RgbTag;

    fn num_channels() -> u32 {
        3
    }

    fn from_slice(values: &[Self::Component]) -> Self {
        Rgb{r: BoundedChannel::new(values[0]),
            g: BoundedChannel::new(values[1]),
            b: BoundedChannel::new(values[2])}
    }

    fn as_slice(&self) -> &[Self::Component] {
        unsafe {
            let ptr: *const Self::Component = mem::transmute(self);
            slice::from_raw_parts(ptr, Self::num_channels() as usize)
        }
    }

    fn broadcast(value: T) -> Self {
        Rgb::from_channels(value.clone(), value.clone(), value.clone())
    }
}

impl<T: ColorChannel> color::Color3 for Rgb<T> {
    fn as_tuple(&self) -> (Self::Component, Self::Component, Self::Component) {
        (self.r.0, self.g.0, self.b.0)
    }

    fn as_array(&self) -> [Self::Component; 3] {
        [self.r.0, self.g.0, self.b.0]
    }

    #[inline]
    fn from_tuple(values: &(Self::Component, Self::Component, Self::Component)) -> Self {
        Rgb{r: BoundedChannel::new(values.0),
            g: BoundedChannel::new(values.1),
            b: BoundedChannel::new(values.2)}
    }

}

impl<T: ColorChannel> color::ComponentMap for Rgb<T> {
    fn component_map<F> (&self, mut f: F) -> Self
            where F: FnMut(Self::Component) -> Self::Component {
        Rgb{r: BoundedChannel::new(f(self.red().clone())),
            g: BoundedChannel::new(f(self.green().clone())),
            b: BoundedChannel::new(f(self.blue().clone()))}
    }

    fn component_map_binary<F>(&self, other: &Self, mut f: F) -> Self
            where F: FnMut(Self::Component, Self::Component) -> Self::Component {
        Rgb{r: BoundedChannel::new(
                f(self.red(), other.red())),
            g: BoundedChannel::new(
                f(self.green(), other.green())),
            b: BoundedChannel::new(
                f(self.blue(), other.blue()))}
    }
}

impl<T: ColorChannel> color::Invert for Rgb<T> {
    fn invert(&self) -> Self {
        Rgb{
            r: self.r.clone().invert(),
            g: self.g.clone().invert(),
            b: self.b.clone().invert(),
        }
    }
}

impl<T: ColorChannel> color::Bounded for Rgb<T> {
    fn clamp(&self, min: Self::Component, max: Self::Component) -> Self {
        Rgb{r: self.r.clamp(min, max),
            g: self.g.clamp(min, max),
            b: self.b.clamp(min, max)}
    }

    fn normalize(&self) -> Self {
        Rgb{r: self.r.normalize(),
            g: self.g.normalize(),
            b: self.b.normalize()}
    }

    fn is_normalized(&self) -> bool {
        self.r.is_normalized() && self.g.is_normalized() && self.b.is_normalized()
    }
}

impl<T: ColorChannel> Default for Rgb<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T: ColorChannel + Float + approx::ApproxEq> approx::ApproxEq for Rgb<T>
        where T::Epsilon: Clone {
    type Epsilon = T::Epsilon;

    fn default_epsilon() -> Self::Epsilon {
        T::default_epsilon()
    }

    fn default_max_relative() -> Self::Epsilon {
        T::default_max_relative()
    }

    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    fn relative_eq(&self, other: &Self, epsilon: Self::Epsilon,
           max_relative: Self::Epsilon) -> bool {
        self.red().relative_eq(&other.red(), epsilon.clone(), max_relative.clone())
        && self.green().relative_eq(&other.green(), epsilon.clone(), max_relative.clone())
        && self.blue().relative_eq(&other.blue(), epsilon, max_relative)
    }

    fn ulps_eq(&self, other: &Self, epsilon: Self::Epsilon, max_ulps: u32) -> bool {
        self.red().ulps_eq(&other.red(), epsilon.clone(), max_ulps)
        && self.green().ulps_eq(&other.green(), epsilon.clone(), max_ulps)
        && self.blue().ulps_eq(&other.blue(), epsilon.clone(), max_ulps)
    }

}

impl<T: ColorChannel + fmt::Display> fmt::Display for Rgb<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Rgb({}, {}, {})", self.r, self.g, self.b)
    }
}

#[cfg(test)]
mod test {
    use ::rgb::*;
    use ::color;
    use ::color::{Color, Invert};

    #[test]
    fn test_construct() {
        {
            let color = Rgb::from_channels(0u8, 0, 0);
            assert_eq!(color.red(), 0u8);
            assert_eq!(color.green(), 0u8);
            assert_eq!(color.blue(), 0u8);

            let c2 = color.clone();
            assert_eq!(color, c2);
        }
        {
            let color: Rgb<u8> = Rgb::new();
            assert_eq!(color.red(), 0u8);
            assert_eq!(color.green(), 0u8);
            assert_eq!(color.blue(), 0u8);
        }
        {
            let color = Rgb::broadcast(0.5_f32);
            assert_ulps_eq!(color, Rgb::from_channels(0.5_f32, 0.5, 0.5));
        }
    }

    #[test]
    fn test_lerp_int() {
        let c1 = Rgb::from_channels(100u8, 200u8, 0u8);
        let c2 = Rgb::from_channels(200u8, 0u8, 255u8);

        assert_eq!(c1.lerp(&c2, 0.5_f64), Rgb::from_channels(150u8, 100, 127));
        assert_eq!(c1.lerp(&c2, 0.0_f64), c1);
        assert_eq!(c1.lerp(&c2, 1.0_f64), c2);
    }

    #[test]
    fn test_lerp_float() {
        let c1 = Rgb::from_channels(0.2_f32, 0.5, 1.0);
        let c2 = Rgb::from_channels(0.8_f32, 0.5, 0.1);

        assert_ulps_eq!(c1.lerp(&c2, 0.5_f32), Rgb::from_channels(0.5_f32, 0.5, 0.55));
        assert_ulps_eq!(c1.lerp(&c2, 0.0_f64), Rgb::from_channels(0.2_f32, 0.5, 1.0));
        assert_ulps_eq!(c1.lerp(&c2, 1.0_f32), Rgb::from_channels(0.8_f32, 0.5, 0.1));
    }

    #[test]
    fn test_invert() {
        let c = Rgb::from_channels(200u8, 0, 255);
        let c2 = Rgb::from_channels(0.8_f32, 0.0, 0.25);

        assert_eq!(c.invert(), Rgb::from_channels(55u8, 255, 0));
        assert_ulps_eq!(c2.invert(), Rgb::from_channels(0.2_f32, 1.0, 0.75));
    }

    #[test]
    fn as_slice() {
        let c = Rgb::from_channels(100u8, 0, 125);
        let c2 = Rgb::from_channels(1.0, 0.25, 0.125);

        assert_eq!(c.as_slice()[0], 100u8);
        assert_eq!(c.as_slice()[1], 0u8);
        assert_eq!(c.as_slice()[2], 125u8);

        assert_ulps_eq!(Rgb::from_slice(c2.as_slice()), c2);
    }

    #[test]
    fn color_cast() {
        let c = Rgb::from_channels(127, 0, 255);
        let c2 = color::color_cast::<Rgb<f32>, _>(&c);
        let c3 = color::color_cast::<Rgb<u8>, _>(&c2);

        assert_ulps_eq!(c2.red(), 127.0 / 255.0);
        assert_ulps_eq!(c2.green(), 0.0);
        assert_ulps_eq!(c2.blue(), 1.0);

        assert_eq!(c3.red(), 127);
        assert_eq!(c3.green(), 0);
        assert_eq!(c3.blue(), 255);

        println!("{}", c2);
    }
}
<|endoftext|>


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

extern crate libc;
extern crate rocksdb;

mod util;

use libc::size_t;

use rocksdb::{DBVector, Error, IteratorMode, Options, WriteBatch, DB};
use util::DBPath;

#[test]
fn test_db_vector() {
    use std::mem;
    let len: size_t = 4;
    let data: *mut u8 = unsafe { libc::calloc(len, mem::size_of::<u8>()) as *mut u8 };
    let v = unsafe { DBVector::from_c(data, len) };
    let ctrl = [0u8, 0, 0, 0];
    assert_eq!(&*v, &ctrl[..]);
}

#[test]
fn external() {
    let path = DBPath::new("_rust_rocksdb_externaltest");

    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1111").is_ok());

        let r: Result<Option<DBVector>, Error> = db.get(b"k1");

        assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
        assert!(db.delete(b"k1").is_ok());
        assert!(db.get(b"k1").unwrap().is_none());
    }
}

#[test]
fn errors_do_stuff() {
    let path = DBPath::new("_rust_rocksdb_error");
    let _db = DB::open_default(&path).unwrap();
    let opts = Options::default();
    // The DB will still be open when we try to destroy it and the lock should fail.
    match DB::destroy(&opts, &path) {
        Err(s) => {
            let message = s.to_string();
            assert!(message.find("IO error:").is_some());
            assert!(message.find("_rust_rocksdb_error").is_some());
            assert!(message.find("/LOCK:").is_some());
        }
        Ok(_) => panic!("should fail"),
    }
}

#[test]
fn writebatch_works() {
    let path = DBPath::new("_rust_rocksdb_writebacktest");
    {
        let db = DB::open_default(&path).unwrap();
        {
            // test put
            let mut batch = WriteBatch::default();
            assert!(db.get(b"k1").unwrap().is_none());
            assert_eq!(batch.len(), 0);
            assert!(batch.is_empty());
            let _ = batch.put(b"k1", b"v1111");
            assert_eq!(batch.len(), 1);
            assert!(!batch.is_empty());
            assert!(db.get(b"k1").unwrap().is_none());
            assert!(db.write(batch).is_ok());
            let r: Result<Option<DBVector>, Error> = db.get(b"k1");
            assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
        }
        {
            // test delete
            let mut batch = WriteBatch::default();
            let _ = batch.delete(b"k1");
            assert_eq!(batch.len(), 1);
            assert!(!batch.is_empty());
            assert!(db.write(batch).is_ok());
            assert!(db.get(b"k1").unwrap().is_none());
        }
        {
            // test size_in_bytes
            let mut batch = WriteBatch::default();
            let before = batch.size_in_bytes();
            let _ = batch.put(b"k1", b"v1234567890");
            let after = batch.size_in_bytes();
            assert!(before + 10 <= after);
        }
    }
}

#[test]
fn iterator_test() {
    let path = DBPath::new("_rust_rocksdb_iteratortest");
    {
        let data = [(b"k1", b"v1111"), (b"k2", b"v2222"), (b"k3", b"v3333")];
        let db = DB::open_default(&path).unwrap();

        for (key, value) in &data {
            assert!(db.put(key, value).is_ok());
        }

        let iter = db.iterator(IteratorMode::Start);

        for (idx, (db_key, db_value)) in iter.enumerate() {
            let (key, value) = data[idx];
            assert_eq!((&key[..], &value[..]), (db_key.as_ref(), db_value.as_ref()));
        }
    }
}

#[test]
fn snapshot_test() {
    let path = DBPath::new("_rust_rocksdb_snapshottest");
    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1111").is_ok());

        let snap = db.snapshot();
        assert!(snap.get(b"k1").unwrap().unwrap().to_utf8().unwrap() == "v1111");

        assert!(db.put(b"k2", b"v2222").is_ok());

        assert!(db.get(b"k2").unwrap().is_some());
        assert!(snap.get(b"k2").unwrap().is_none());
    }
}

#[test]
fn set_option_test() {
    let path = DBPath::new("_rust_rocksdb_set_optionstest");
    {
        let db = DB::open_default(&path).unwrap();
        // set an option to valid values
        assert!(db
            .set_options(&[("disable_auto_compactions", "true")])
            .is_ok());
        assert!(db
            .set_options(&[("disable_auto_compactions", "false")])
            .is_ok());
        // invalid names/values should result in an error
        assert!(db
            .set_options(&[("disable_auto_compactions", "INVALID_VALUE")])
            .is_err());
        assert!(db
            .set_options(&[("INVALID_NAME", "INVALID_VALUE")])
            .is_err());
        // option names/values must not contain NULLs
        assert!(db
            .set_options(&[("disable_auto_compactions", "true\0")])
            .is_err());
        assert!(db
            .set_options(&[("disable_auto_compactions\0", "true")])
            .is_err());
        // empty options are not allowed
        assert!(db.set_options(&[]).is_err());
        // multiple options can be set in a single API call
        let multiple_options = [
            ("paranoid_file_checks", "true"),
            ("report_bg_io_stats", "true"),
        ];
        db.set_options(&multiple_options).unwrap();
    }
}
<commit_msg>Adding test for AsRef<[u8]> impl on DBVector...<commit_after>// Copyright 2019 Tyler Neely
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate libc;
extern crate rocksdb;

mod util;

use libc::size_t;

use rocksdb::{DBVector, Error, IteratorMode, Options, WriteBatch, DB};
use util::DBPath;

#[test]
fn test_db_vector() {
    use std::mem;
    let len: size_t = 4;
    let data: *mut u8 = unsafe { libc::calloc(len, mem::size_of::<u8>()) as *mut u8 };
    let v = unsafe { DBVector::from_c(data, len) };
    let ctrl = [0u8, 0, 0, 0];
    assert_eq!(&*v, &ctrl[..]);
}

#[test]
fn external() {
    let path = DBPath::new("_rust_rocksdb_externaltest");

    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1111").is_ok());

        let r: Result<Option<DBVector>, Error> = db.get(b"k1");

        assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
        assert!(db.delete(b"k1").is_ok());
        assert!(db.get(b"k1").unwrap().is_none());
    }
}

#[test]
fn db_vector_as_ref_byte_slice() {
    let path = DBPath::new("_rust_rocksdb_db_vector_as_ref_byte_slice");

    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1111").is_ok());

        let r: Result<Option<DBVector>, Error> = db.get(b"k1");
        let vector = r.unwrap().unwrap();

        assert!(get_byte_slice(&vector) == b"v1111");
    }
}

fn get_byte_slice<T: AsRef<[u8]>>(source: &'_ T) -> &'_ [u8] {
    source.as_ref()
}

#[test]
fn errors_do_stuff() {
    let path = DBPath::new("_rust_rocksdb_error");
    let _db = DB::open_default(&path).unwrap();
    let opts = Options::default();
    // The DB will still be open when we try to destroy it and the lock should fail.
    match DB::destroy(&opts, &path) {
        Err(s) => {
            let message = s.to_string();
            assert!(message.find("IO error:").is_some());
            assert!(message.find("_rust_rocksdb_error").is_some());
            assert!(message.find("/LOCK:").is_some());
        }
        Ok(_) => panic!("should fail"),
    }
}

#[test]
fn writebatch_works() {
    let path = DBPath::new("_rust_rocksdb_writebacktest");
    {
        let db = DB::open_default(&path).unwrap();
        {
            // test put
            let mut batch = WriteBatch::default();
            assert!(db.get(b"k1").unwrap().is_none());
            assert_eq!(batch.len(), 0);
            assert!(batch.is_empty());
            let _ = batch.put(b"k1", b"v1111");
            assert_eq!(batch.len(), 1);
            assert!(!batch.is_empty());
            assert!(db.get(b"k1").unwrap().is_none());
            assert!(db.write(batch).is_ok());
            let r: Result<Option<DBVector>, Error> = db.get(b"k1");
            assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
        }
        {
            // test delete
            let mut batch = WriteBatch::default();
            let _ = batch.delete(b"k1");
            assert_eq!(batch.len(), 1);
            assert!(!batch.is_empty());
            assert!(db.write(batch).is_ok());
            assert!(db.get(b"k1").unwrap().is_none());
        }
        {
            // test size_in_bytes
            let mut batch = WriteBatch::default();
            let before = batch.size_in_bytes();
            let _ = batch.put(b"k1", b"v1234567890");
            let after = batch.size_in_bytes();
            assert!(before + 10 <= after);
        }
    }
}

#[test]
fn iterator_test() {
    let path = DBPath::new("_rust_rocksdb_iteratortest");
    {
        let data = [(b"k1", b"v1111"), (b"k2", b"v2222"), (b"k3", b"v3333")];
        let db = DB::open_default(&path).unwrap();

        for (key, value) in &data {
            assert!(db.put(key, value).is_ok());
        }

        let iter = db.iterator(IteratorMode::Start);

        for (idx, (db_key, db_value)) in iter.enumerate() {
            let (key, value) = data[idx];
            assert_eq!((&key[..], &value[..]), (db_key.as_ref(), db_value.as_ref()));
        }
    }
}

#[test]
fn snapshot_test() {
    let path = DBPath::new("_rust_rocksdb_snapshottest");
    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1111").is_ok());

        let snap = db.snapshot();
        assert!(snap.get(b"k1").unwrap().unwrap().to_utf8().unwrap() == "v1111");

        assert!(db.put(b"k2", b"v2222").is_ok());

        assert!(db.get(b"k2").unwrap().is_some());
        assert!(snap.get(b"k2").unwrap().is_none());
    }
}

#[test]
fn set_option_test() {
    let path = DBPath::new("_rust_rocksdb_set_optionstest");
    {
        let db = DB::open_default(&path).unwrap();
        // set an option to valid values
        assert!(db
            .set_options(&[("disable_auto_compactions", "true")])
            .is_ok());
        assert!(db
            .set_options(&[("disable_auto_compactions", "false")])
            .is_ok());
        // invalid names/values should result in an error
        assert!(db
            .set_options(&[("disable_auto_compactions", "INVALID_VALUE")])
            .is_err());
        assert!(db
            .set_options(&[("INVALID_NAME", "INVALID_VALUE")])
            .is_err());
        // option names/values must not contain NULLs
        assert!(db
            .set_options(&[("disable_auto_compactions", "true\0")])
            .is_err());
        assert!(db
            .set_options(&[("disable_auto_compactions\0", "true")])
            .is_err());
        // empty options are not allowed
        assert!(db.set_options(&[]).is_err());
        // multiple options can be set in a single API call
        let multiple_options = [
            ("paranoid_file_checks", "true"),
            ("report_bg_io_stats", "true"),
        ];
        db.set_options(&multiple_options).unwrap();
    }
}
<|endoftext|>


<commit_before>extern crate actix;
extern crate actix_web;
extern crate bytes;
extern crate futures;
extern crate http;
extern crate rand;

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::{thread, time};

use bytes::Bytes;
use futures::Stream;
use rand::distributions::Alphanumeric;
use rand::Rng;

#[cfg(feature = "alpn")]
extern crate openssl;
#[cfg(feature = "rust-tls")]
extern crate rustls;

use actix::prelude::*;
use actix_web::*;

struct Ws;

impl Actor for Ws {
    type Context = ws::WebsocketContext<Self>;
}

impl StreamHandler<ws::Message, ws::ProtocolError> for Ws {
    fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
        match msg {
            ws::Message::Ping(msg) => ctx.pong(&msg),
            ws::Message::Text(text) => ctx.text(text),
            ws::Message::Binary(bin) => ctx.binary(bin),
            ws::Message::Close(reason) => ctx.close(reason),
            _ => (),
        }
    }
}

#[test]
fn test_simple() {
    let mut srv = test::TestServer::new(|app| app.handler(|req| ws::start(req, Ws)));
    let (reader, mut writer) = srv.ws().unwrap();

    writer.text("text");
    let (item, reader) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(item, Some(ws::Message::Text("text".to_owned())));

    writer.binary(b"text".as_ref());
    let (item, reader) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(
        item,
        Some(ws::Message::Binary(Bytes::from_static(b"text").into()))
    );

    writer.ping("ping");
    let (item, reader) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(item, Some(ws::Message::Pong("ping".to_owned())));

    writer.close(Some(ws::CloseCode::Normal.into()));
    let (item, _) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(
        item,
        Some(ws::Message::Close(Some(ws::CloseCode::Normal.into())))
    );
}

// websocket resource helper function
fn start_ws_resource(req: &HttpRequest) -> Result<HttpResponse, Error> {
    ws::start(req, Ws)
}

#[test]
fn test_simple_path() {
    const PATH: &str = "/v1/ws/";

    // Create a websocket at a specific path.
    let mut srv = test::TestServer::new(|app| {
        app.resource(PATH, |r| r.route().f(start_ws_resource));
    });
    // fetch the sockets for the resource at a given path.
    let (reader, mut writer) = srv.ws_at(PATH).unwrap();

    writer.text("text");
    let (item, reader) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(item, Some(ws::Message::Text("text".to_owned())));

    writer.binary(b"text".as_ref());
    let (item, reader) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(
        item,
        Some(ws::Message::Binary(Bytes::from_static(b"text").into()))
    );

    writer.ping("ping");
    let (item, reader) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(item, Some(ws::Message::Pong("ping".to_owned())));

    writer.close(Some(ws::CloseCode::Normal.into()));
    let (item, _) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(
        item,
        Some(ws::Message::Close(Some(ws::CloseCode::Normal.into())))
    );
}

#[test]
fn test_empty_close_code() {
    let mut srv = test::TestServer::new(|app| app.handler(|req| ws::start(req, Ws)));
    let (reader, mut writer) = srv.ws().unwrap();

    writer.close(None);
    let (item, _) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(item, Some(ws::Message::Close(None)));
}

#[test]
fn test_close_description() {
    let mut srv = test::TestServer::new(|app| app.handler(|req| ws::start(req, Ws)));
    let (reader, mut writer) = srv.ws().unwrap();

    let close_reason: ws::CloseReason =
        (ws::CloseCode::Normal, "close description").into();
    writer.close(Some(close_reason.clone()));
    let (item, _) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(item, Some(ws::Message::Close(Some(close_reason))));
}

#[test]
fn test_large_text() {
    let data = rand::thread_rng()
        .sample_iter(&Alphanumeric)
        .take(65_536)
        .collect::<String>();

    let mut srv = test::TestServer::new(|app| app.handler(|req| ws::start(req, Ws)));
    let (mut reader, mut writer) = srv.ws().unwrap();

    for _ in 0..100 {
        writer.text(data.clone());
        let (item, r) = srv.execute(reader.into_future()).unwrap();
        reader = r;
        assert_eq!(item, Some(ws::Message::Text(data.clone())));
    }
}

#[test]
fn test_large_bin() {
    let data = rand::thread_rng()
        .sample_iter(&Alphanumeric)
        .take(65_536)
        .collect::<String>();

    let mut srv = test::TestServer::new(|app| app.handler(|req| ws::start(req, Ws)));
    let (mut reader, mut writer) = srv.ws().unwrap();

    for _ in 0..100 {
        writer.binary(data.clone());
        let (item, r) = srv.execute(reader.into_future()).unwrap();
        reader = r;
        assert_eq!(item, Some(ws::Message::Binary(Binary::from(data.clone()))));
    }
}

#[test]
fn test_client_frame_size() {
    let data = rand::thread_rng()
        .sample_iter(&Alphanumeric)
        .take(131_072)
        .collect::<String>();

    let mut srv = test::TestServer::new(|app| {
        app.handler(|req| -> Result<HttpResponse> {
            let mut resp = ws::handshake(req)?;
            let stream = ws::WsStream::new(req.payload()).max_size(131_072);

            let body = ws::WebsocketContext::create(req.clone(), Ws, stream);
            Ok(resp.body(body))
        })
    });
    let (reader, mut writer) = srv.ws().unwrap();

    writer.binary(data.clone());
    match srv.execute(reader.into_future()).err().unwrap().0 {
        ws::ProtocolError::Overflow => (),
        _ => panic!(),
    }
}

struct Ws2 {
    count: usize,
    bin: bool,
}

impl Actor for Ws2 {
    type Context = ws::WebsocketContext<Self>;

    fn started(&mut self, ctx: &mut Self::Context) {
        self.send(ctx);
    }
}

impl Ws2 {
    fn send(&mut self, ctx: &mut ws::WebsocketContext<Self>) {
        if self.bin {
            ctx.binary(Vec::from("0".repeat(65_536)));
        } else {
            ctx.text("0".repeat(65_536));
        }
        ctx.drain()
            .and_then(|_, act, ctx| {
                act.count += 1;
                if act.count != 10_000 {
                    act.send(ctx);
                }
                actix::fut::ok(())
            }).wait(ctx);
    }
}

impl StreamHandler<ws::Message, ws::ProtocolError> for Ws2 {
    fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
        match msg {
            ws::Message::Ping(msg) => ctx.pong(&msg),
            ws::Message::Text(text) => ctx.text(text),
            ws::Message::Binary(bin) => ctx.binary(bin),
            ws::Message::Close(reason) => ctx.close(reason),
            _ => (),
        }
    }
}

#[test]
fn test_server_send_text() {
    let data = Some(ws::Message::Text("0".repeat(65_536)));

    let mut srv = test::TestServer::new(|app| {
        app.handler(|req| {
            ws::start(
                req,
                Ws2 {
                    count: 0,
                    bin: false,
                },
            )
        })
    });
    let (mut reader, _writer) = srv.ws().unwrap();

    for _ in 0..10_000 {
        let (item, r) = srv.execute(reader.into_future()).unwrap();
        reader = r;
        assert_eq!(item, data);
    }
}

#[test]
fn test_server_send_bin() {
    let data = Some(ws::Message::Binary(Binary::from("0".repeat(65_536))));

    let mut srv = test::TestServer::new(|app| {
        app.handler(|req| {
            ws::start(
                req,
                Ws2 {
                    count: 0,
                    bin: true,
                },
            )
        })
    });
    let (mut reader, _writer) = srv.ws().unwrap();

    for _ in 0..10_000 {
        let (item, r) = srv.execute(reader.into_future()).unwrap();
        reader = r;
        assert_eq!(item, data);
    }
}

#[test]
#[cfg(feature = "alpn")]
fn test_ws_server_ssl() {
    extern crate openssl;
    use openssl::ssl::{SslAcceptor, SslFiletype, SslMethod};

    // load ssl keys
    let mut builder = SslAcceptor::mozilla_intermediate(SslMethod::tls()).unwrap();
    builder
        .set_private_key_file("tests/key.pem", SslFiletype::PEM)
        .unwrap();
    builder
        .set_certificate_chain_file("tests/cert.pem")
        .unwrap();

    let mut srv = test::TestServer::build().ssl(builder).start(|app| {
        app.handler(|req| {
            ws::start(
                req,
                Ws2 {
                    count: 0,
                    bin: false,
                },
            )
        })
    });
    let (mut reader, _writer) = srv.ws().unwrap();

    let data = Some(ws::Message::Text("0".repeat(65_536)));
    for _ in 0..10_000 {
        let (item, r) = srv.execute(reader.into_future()).unwrap();
        reader = r;
        assert_eq!(item, data);
    }
}

#[test]
#[cfg(feature = "rust-tls")]
fn test_ws_server_rust_tls() {
    extern crate rustls;
    use rustls::internal::pemfile::{certs, rsa_private_keys};
    use rustls::{NoClientAuth, ServerConfig};
    use std::fs::File;
    use std::io::BufReader;

    // load ssl keys
    let mut config = ServerConfig::new(NoClientAuth::new());
    let cert_file = &mut BufReader::new(File::open("tests/cert.pem").unwrap());
    let key_file = &mut BufReader::new(File::open("tests/key.pem").unwrap());
    let cert_chain = certs(cert_file).unwrap();
    let mut keys = rsa_private_keys(key_file).unwrap();
    config.set_single_cert(cert_chain, keys.remove(0)).unwrap();

    let mut srv = test::TestServer::build().rustls(config).start(|app| {
        app.handler(|req| {
            ws::start(
                req,
                Ws2 {
                    count: 0,
                    bin: false,
                },
            )
        })
    });

    let (mut reader, _writer) = srv.ws().unwrap();

    let data = Some(ws::Message::Text("0".repeat(65_536)));
    for _ in 0..10_000 {
        let (item, r) = srv.execute(reader.into_future()).unwrap();
        reader = r;
        assert_eq!(item, data);
    }
}

struct WsStopped(Arc<AtomicUsize>);

impl Actor for WsStopped {
    type Context = ws::WebsocketContext<Self>;

    fn stopped(&mut self, ctx: &mut Self::Context) {
        self.0.fetch_add(1, Ordering::Relaxed);
    }
}

impl StreamHandler<ws::Message, ws::ProtocolError> for WsStopped {
    fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
        match msg {
            ws::Message::Text(text) => ctx.text(text),
            _ => (),
        }
    }
}

#[test]
fn test_ws_stopped() {
    let num = Arc::new(AtomicUsize::new(0));
    let num2 = num.clone();

    let _ = thread::spawn(move || {
        let num3 = num2.clone();
        let mut srv = test::TestServer::new(move |app| {
            let num4 = num3.clone();
            app.handler(move |req| ws::start(req, WsStopped(num4.clone())))
        });
        let (reader, mut writer) = srv.ws().unwrap();

        writer.text("text");
        let (item, reader) = srv.execute(reader.into_future()).unwrap();
        assert_eq!(item, Some(ws::Message::Text("text".to_owned())));
    });

    thread::sleep(time::Duration::from_secs(1));
    assert_eq!(num.load(Ordering::Relaxed), 1);
}
<commit_msg>increase sleep time<commit_after>extern crate actix;
extern crate actix_web;
extern crate bytes;
extern crate futures;
extern crate http;
extern crate rand;

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::{thread, time};

use bytes::Bytes;
use futures::Stream;
use rand::distributions::Alphanumeric;
use rand::Rng;

#[cfg(feature = "alpn")]
extern crate openssl;
#[cfg(feature = "rust-tls")]
extern crate rustls;

use actix::prelude::*;
use actix_web::*;

struct Ws;

impl Actor for Ws {
    type Context = ws::WebsocketContext<Self>;
}

impl StreamHandler<ws::Message, ws::ProtocolError> for Ws {
    fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
        match msg {
            ws::Message::Ping(msg) => ctx.pong(&msg),
            ws::Message::Text(text) => ctx.text(text),
            ws::Message::Binary(bin) => ctx.binary(bin),
            ws::Message::Close(reason) => ctx.close(reason),
            _ => (),
        }
    }
}

#[test]
fn test_simple() {
    let mut srv = test::TestServer::new(|app| app.handler(|req| ws::start(req, Ws)));
    let (reader, mut writer) = srv.ws().unwrap();

    writer.text("text");
    let (item, reader) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(item, Some(ws::Message::Text("text".to_owned())));

    writer.binary(b"text".as_ref());
    let (item, reader) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(
        item,
        Some(ws::Message::Binary(Bytes::from_static(b"text").into()))
    );

    writer.ping("ping");
    let (item, reader) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(item, Some(ws::Message::Pong("ping".to_owned())));

    writer.close(Some(ws::CloseCode::Normal.into()));
    let (item, _) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(
        item,
        Some(ws::Message::Close(Some(ws::CloseCode::Normal.into())))
    );
}

// websocket resource helper function
fn start_ws_resource(req: &HttpRequest) -> Result<HttpResponse, Error> {
    ws::start(req, Ws)
}

#[test]
fn test_simple_path() {
    const PATH: &str = "/v1/ws/";

    // Create a websocket at a specific path.
    let mut srv = test::TestServer::new(|app| {
        app.resource(PATH, |r| r.route().f(start_ws_resource));
    });
    // fetch the sockets for the resource at a given path.
    let (reader, mut writer) = srv.ws_at(PATH).unwrap();

    writer.text("text");
    let (item, reader) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(item, Some(ws::Message::Text("text".to_owned())));

    writer.binary(b"text".as_ref());
    let (item, reader) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(
        item,
        Some(ws::Message::Binary(Bytes::from_static(b"text").into()))
    );

    writer.ping("ping");
    let (item, reader) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(item, Some(ws::Message::Pong("ping".to_owned())));

    writer.close(Some(ws::CloseCode::Normal.into()));
    let (item, _) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(
        item,
        Some(ws::Message::Close(Some(ws::CloseCode::Normal.into())))
    );
}

#[test]
fn test_empty_close_code() {
    let mut srv = test::TestServer::new(|app| app.handler(|req| ws::start(req, Ws)));
    let (reader, mut writer) = srv.ws().unwrap();

    writer.close(None);
    let (item, _) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(item, Some(ws::Message::Close(None)));
}

#[test]
fn test_close_description() {
    let mut srv = test::TestServer::new(|app| app.handler(|req| ws::start(req, Ws)));
    let (reader, mut writer) = srv.ws().unwrap();

    let close_reason: ws::CloseReason =
        (ws::CloseCode::Normal, "close description").into();
    writer.close(Some(close_reason.clone()));
    let (item, _) = srv.execute(reader.into_future()).unwrap();
    assert_eq!(item, Some(ws::Message::Close(Some(close_reason))));
}

#[test]
fn test_large_text() {
    let data = rand::thread_rng()
        .sample_iter(&Alphanumeric)
        .take(65_536)
        .collect::<String>();

    let mut srv = test::TestServer::new(|app| app.handler(|req| ws::start(req, Ws)));
    let (mut reader, mut writer) = srv.ws().unwrap();

    for _ in 0..100 {
        writer.text(data.clone());
        let (item, r) = srv.execute(reader.into_future()).unwrap();
        reader = r;
        assert_eq!(item, Some(ws::Message::Text(data.clone())));
    }
}

#[test]
fn test_large_bin() {
    let data = rand::thread_rng()
        .sample_iter(&Alphanumeric)
        .take(65_536)
        .collect::<String>();

    let mut srv = test::TestServer::new(|app| app.handler(|req| ws::start(req, Ws)));
    let (mut reader, mut writer) = srv.ws().unwrap();

    for _ in 0..100 {
        writer.binary(data.clone());
        let (item, r) = srv.execute(reader.into_future()).unwrap();
        reader = r;
        assert_eq!(item, Some(ws::Message::Binary(Binary::from(data.clone()))));
    }
}

#[test]
fn test_client_frame_size() {
    let data = rand::thread_rng()
        .sample_iter(&Alphanumeric)
        .take(131_072)
        .collect::<String>();

    let mut srv = test::TestServer::new(|app| {
        app.handler(|req| -> Result<HttpResponse> {
            let mut resp = ws::handshake(req)?;
            let stream = ws::WsStream::new(req.payload()).max_size(131_072);

            let body = ws::WebsocketContext::create(req.clone(), Ws, stream);
            Ok(resp.body(body))
        })
    });
    let (reader, mut writer) = srv.ws().unwrap();

    writer.binary(data.clone());
    match srv.execute(reader.into_future()).err().unwrap().0 {
        ws::ProtocolError::Overflow => (),
        _ => panic!(),
    }
}

struct Ws2 {
    count: usize,
    bin: bool,
}

impl Actor for Ws2 {
    type Context = ws::WebsocketContext<Self>;

    fn started(&mut self, ctx: &mut Self::Context) {
        self.send(ctx);
    }
}

impl Ws2 {
    fn send(&mut self, ctx: &mut ws::WebsocketContext<Self>) {
        if self.bin {
            ctx.binary(Vec::from("0".repeat(65_536)));
        } else {
            ctx.text("0".repeat(65_536));
        }
        ctx.drain()
            .and_then(|_, act, ctx| {
                act.count += 1;
                if act.count != 10_000 {
                    act.send(ctx);
                }
                actix::fut::ok(())
            }).wait(ctx);
    }
}

impl StreamHandler<ws::Message, ws::ProtocolError> for Ws2 {
    fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
        match msg {
            ws::Message::Ping(msg) => ctx.pong(&msg),
            ws::Message::Text(text) => ctx.text(text),
            ws::Message::Binary(bin) => ctx.binary(bin),
            ws::Message::Close(reason) => ctx.close(reason),
            _ => (),
        }
    }
}

#[test]
fn test_server_send_text() {
    let data = Some(ws::Message::Text("0".repeat(65_536)));

    let mut srv = test::TestServer::new(|app| {
        app.handler(|req| {
            ws::start(
                req,
                Ws2 {
                    count: 0,
                    bin: false,
                },
            )
        })
    });
    let (mut reader, _writer) = srv.ws().unwrap();

    for _ in 0..10_000 {
        let (item, r) = srv.execute(reader.into_future()).unwrap();
        reader = r;
        assert_eq!(item, data);
    }
}

#[test]
fn test_server_send_bin() {
    let data = Some(ws::Message::Binary(Binary::from("0".repeat(65_536))));

    let mut srv = test::TestServer::new(|app| {
        app.handler(|req| {
            ws::start(
                req,
                Ws2 {
                    count: 0,
                    bin: true,
                },
            )
        })
    });
    let (mut reader, _writer) = srv.ws().unwrap();

    for _ in 0..10_000 {
        let (item, r) = srv.execute(reader.into_future()).unwrap();
        reader = r;
        assert_eq!(item, data);
    }
}

#[test]
#[cfg(feature = "alpn")]
fn test_ws_server_ssl() {
    extern crate openssl;
    use openssl::ssl::{SslAcceptor, SslFiletype, SslMethod};

    // load ssl keys
    let mut builder = SslAcceptor::mozilla_intermediate(SslMethod::tls()).unwrap();
    builder
        .set_private_key_file("tests/key.pem", SslFiletype::PEM)
        .unwrap();
    builder
        .set_certificate_chain_file("tests/cert.pem")
        .unwrap();

    let mut srv = test::TestServer::build().ssl(builder).start(|app| {
        app.handler(|req| {
            ws::start(
                req,
                Ws2 {
                    count: 0,
                    bin: false,
                },
            )
        })
    });
    let (mut reader, _writer) = srv.ws().unwrap();

    let data = Some(ws::Message::Text("0".repeat(65_536)));
    for _ in 0..10_000 {
        let (item, r) = srv.execute(reader.into_future()).unwrap();
        reader = r;
        assert_eq!(item, data);
    }
}

#[test]
#[cfg(feature = "rust-tls")]
fn test_ws_server_rust_tls() {
    extern crate rustls;
    use rustls::internal::pemfile::{certs, rsa_private_keys};
    use rustls::{NoClientAuth, ServerConfig};
    use std::fs::File;
    use std::io::BufReader;

    // load ssl keys
    let mut config = ServerConfig::new(NoClientAuth::new());
    let cert_file = &mut BufReader::new(File::open("tests/cert.pem").unwrap());
    let key_file = &mut BufReader::new(File::open("tests/key.pem").unwrap());
    let cert_chain = certs(cert_file).unwrap();
    let mut keys = rsa_private_keys(key_file).unwrap();
    config.set_single_cert(cert_chain, keys.remove(0)).unwrap();

    let mut srv = test::TestServer::build().rustls(config).start(|app| {
        app.handler(|req| {
            ws::start(
                req,
                Ws2 {
                    count: 0,
                    bin: false,
                },
            )
        })
    });

    let (mut reader, _writer) = srv.ws().unwrap();

    let data = Some(ws::Message::Text("0".repeat(65_536)));
    for _ in 0..10_000 {
        let (item, r) = srv.execute(reader.into_future()).unwrap();
        reader = r;
        assert_eq!(item, data);
    }
}

struct WsStopped(Arc<AtomicUsize>);

impl Actor for WsStopped {
    type Context = ws::WebsocketContext<Self>;

    fn stopped(&mut self, ctx: &mut Self::Context) {
        self.0.fetch_add(1, Ordering::Relaxed);
    }
}

impl StreamHandler<ws::Message, ws::ProtocolError> for WsStopped {
    fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
        match msg {
            ws::Message::Text(text) => ctx.text(text),
            _ => (),
        }
    }
}

#[test]
fn test_ws_stopped() {
    let num = Arc::new(AtomicUsize::new(0));
    let num2 = num.clone();

    let _ = thread::spawn(move || {
        let num3 = num2.clone();
        let mut srv = test::TestServer::new(move |app| {
            let num4 = num3.clone();
            app.handler(move |req| ws::start(req, WsStopped(num4.clone())))
        });
        let (reader, mut writer) = srv.ws().unwrap();

        writer.text("text");
        let (item, reader) = srv.execute(reader.into_future()).unwrap();
        assert_eq!(item, Some(ws::Message::Text("text".to_owned())));
    });

    thread::sleep(time::Duration::from_secs(3));
    assert_eq!(num.load(Ordering::Relaxed), 1);
}
<|endoftext|>


<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.

// Coroutines represent nothing more than a context and a stack
// segment.

use std::default::Default;
use std::rt::util::min_stack;
use std::thunk::Thunk;
use std::mem::transmute;
use std::rt::unwind::try;
use std::boxed::BoxAny;
use std::ops::Deref;
use std::cell::RefCell;
use std::rc::Rc;

use context::Context;
use stack::{StackPool, Stack};

#[derive(Debug, Copy)]
pub enum State {
    Suspended,
    Running,
    Finished,
}

/// A coroutine is nothing more than a (register context, stack) pair.
#[allow(raw_pointer_derive)]
#[derive(Debug)]
pub struct Coroutine {
    /// The segment of stack on which the task is currently running or
    /// if the task is blocked, on which the task will resume
    /// execution.
    current_stack_segment: Option<Stack>,

    /// Always valid if the task is alive and not running.
    saved_context: Context,

    /// Parent coroutine, may always be valid.
    parent: Option<Rc<RefCell<Coroutine>>>,

    /// State
    state: State
}

/// Destroy coroutine and try to reuse std::stack segment.
impl Drop for Coroutine {
    fn drop(&mut self) {
        match self.current_stack_segment.take() {
            Some(stack) => {
                COROUTINE_ENVIRONMENT.with(|env| {
                    env.borrow_mut().stack_pool.give_stack(stack);
                });
            },
            None => {}
        }
    }
}

/// Coroutine spawn options
#[derive(Debug)]
pub struct Options {
    stack_size: usize,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            stack_size: min_stack(),
        }
    }
}

/// Initialization function for make context
extern "C" fn coroutine_initialize(_: usize, f: *mut ()) -> ! {
    let func: Box<Thunk> = unsafe { transmute(f) };

    if let Err(cause) = unsafe { try(move|| func.invoke(())) } {
        error!("Panicked inside: {:?}", cause.downcast::<&str>());
    }

    COROUTINE_ENVIRONMENT.with(|env| {
        let mut env = env.borrow_mut();
        env.current_running.borrow_mut().state = State::Finished;
    });

    loop {
        Coroutine::sched();
    }
}

impl Coroutine {
    pub fn empty() -> Rc<RefCell<Coroutine>> {
        Rc::new(RefCell::new(Coroutine {
            current_stack_segment: None,
            saved_context: Context::empty(),
            parent: None,
            state: State::Running,
        }))
    }

    pub fn spawn_opts<F>(f: F, opts: Options) -> Rc<RefCell<Coroutine>>
            where F: FnOnce() + Send + 'static {
        let stack = RefCell::new(unsafe { Stack::dummy_stack() });
        COROUTINE_ENVIRONMENT.with(|env| {
            *stack.borrow_mut() = env.borrow_mut().stack_pool.take_stack(opts.stack_size);
        });

        let mut stack = stack.into_inner();
        let coro = Coroutine::empty();

        let ctx = Context::new(coroutine_initialize,
                               0,
                               f,
                               &mut stack);
        coro.borrow_mut().saved_context = ctx;
        coro.borrow_mut().current_stack_segment = Some(stack);

        Coroutine::resume(&coro);
        coro
    }

    /// Spawn a coroutine with default options
    pub fn spawn<F>(f: F) -> Rc<RefCell<Coroutine>>
            where F: FnOnce() + Send + 'static {
        Coroutine::spawn_opts(f, Default::default())
    }

    pub fn resume(coro: &Rc<RefCell<Coroutine>>) {
        match coro.borrow().state {
            State::Finished => return,
            _ => {}
        }

        COROUTINE_ENVIRONMENT.with(|env| {
            let current_running = env.borrow().current_running.clone();

            let to_coro: &mut Coroutine = unsafe { transmute(coro.borrow().deref()) };
            let from_coro: &mut Coroutine = unsafe { transmute(current_running.borrow().deref()) };

            to_coro.parent = Some(current_running.clone());
            env.borrow_mut().current_running = coro.clone();

            from_coro.state = State::Suspended;
            to_coro.state = State::Running;

            Context::swap(&mut from_coro.saved_context, &to_coro.saved_context);
        })
    }

    pub fn sched() {
        COROUTINE_ENVIRONMENT.with(|env| {
            let from_coro: &mut Coroutine = unsafe { transmute(env.borrow().current_running.deref()) };
            let to_coro: &mut Coroutine = match from_coro.parent {
                Some(ref parent) => unsafe {
                    env.borrow_mut().current_running = parent.clone();
                    transmute(parent.borrow().deref())
                },
                None => return,
            };

            match from_coro.state {
                State::Finished => {},
                _ => from_coro.state = State::Suspended,
            }
            to_coro.state = State::Running;

            Context::swap(&mut from_coro.saved_context, &to_coro.saved_context);
        });
    }

    pub fn current() -> Rc<RefCell<Coroutine>> {
        let opt = RefCell::new(None);
        COROUTINE_ENVIRONMENT.with(|env| {
            *opt.borrow_mut() = Some(env.borrow().current_running.clone());
        });
        opt.into_inner().unwrap()
    }

    pub fn state(&self) -> State {
        self.state
    }
}

thread_local!(static COROUTINE_ENVIRONMENT: RefCell<Environment> = RefCell::new(Environment::new()));

/// Coroutine managing environment
#[allow(raw_pointer_derive)]
#[derive(Debug)]
struct Environment {
    stack_pool: StackPool,
    current_running: Rc<RefCell<Coroutine>>,
    main_coroutine: Rc<RefCell<Coroutine>>,
}

impl Environment {
    /// Initialize a new environment
    fn new() -> Environment {
        let main_coro = Coroutine::empty();
        Environment {
            stack_pool: StackPool::new(),
            current_running: main_coro.clone(),
            main_coroutine: main_coro,
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::mpsc::channel;

    use test::Bencher;

    use coroutine::Coroutine;

    #[test]
    fn test_coroutine_basic() {
        let (tx, rx) = channel();
        Coroutine::spawn(move|| {
            tx.send(1).unwrap();
        });

        assert_eq!(rx.recv().unwrap(), 1);
    }

    #[test]
    fn test_coroutine_yield() {
        let (tx, rx) = channel();
        let coro = Coroutine::spawn(move|| {
            tx.send(1).unwrap();

            Coroutine::sched();

            tx.send(2).unwrap();
        });

        assert_eq!(rx.recv().unwrap(), 1);
        assert!(rx.try_recv().is_err());

        Coroutine::resume(&coro);

        assert_eq!(rx.recv().unwrap(), 2);
    }

    #[test]
    fn test_coroutine_spawn_inside() {
        let (tx, rx) = channel();
        Coroutine::spawn(move|| {
            tx.send(1).unwrap();

            Coroutine::spawn(move|| {
                tx.send(2).unwrap();
            });
        });

        assert_eq!(rx.recv().unwrap(), 1);
        assert_eq!(rx.recv().unwrap(), 2);
    }

    #[test]
    #[should_fail]
    fn test_coroutine_panic() {
        Coroutine::spawn(move|| {
            panic!("Panic inside a coroutine!!");
        });

        unreachable!();
    }

    #[test]
    #[should_fail]
    fn test_coroutine_child_panic() {
        Coroutine::spawn(move|| {

            Coroutine::spawn(move|| {
                panic!("Panic inside a coroutine's child!!");
            });

        });

        unreachable!();
    }

    #[test]
    fn test_coroutine_resume_after_finished() {
        let coro = Coroutine::spawn(move|| {});

        // It is already finished, but we try to resume it
        // Idealy it would come back immediately
        Coroutine::resume(&coro);

        // Again?
        Coroutine::resume(&coro);
    }

    #[bench]
    fn bench_coroutine_spawning_with_recycle(b: &mut Bencher) {
        b.iter(|| {
            Coroutine::spawn(move|| {});
        });
    }

    #[bench]
    fn bench_coroutine_counting(b: &mut Bencher) {
        b.iter(|| {
            const MAX_NUMBER: usize = 100;
            let (tx, rx) = channel();

            let coro = Coroutine::spawn(move|| {
                for _ in 0..MAX_NUMBER {
                    tx.send(1).unwrap();
                    Coroutine::sched();
                }
            });

            let result = rx.iter().fold(0, |a, b| {
                Coroutine::resume(&coro);
                a + b
            });
            assert_eq!(result, MAX_NUMBER);
        });
    }
}
<commit_msg>do not need to maintain a copy of main coroutine in the environment<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.

// Coroutines represent nothing more than a context and a stack
// segment.

use std::default::Default;
use std::rt::util::min_stack;
use std::thunk::Thunk;
use std::mem::transmute;
use std::rt::unwind::try;
use std::boxed::BoxAny;
use std::ops::Deref;
use std::cell::RefCell;
use std::rc::Rc;

use context::Context;
use stack::{StackPool, Stack};

#[derive(Debug, Copy)]
pub enum State {
    Suspended,
    Running,
    Finished,
}

/// A coroutine is nothing more than a (register context, stack) pair.
#[allow(raw_pointer_derive)]
#[derive(Debug)]
pub struct Coroutine {
    /// The segment of stack on which the task is currently running or
    /// if the task is blocked, on which the task will resume
    /// execution.
    current_stack_segment: Option<Stack>,

    /// Always valid if the task is alive and not running.
    saved_context: Context,

    /// Parent coroutine, may always be valid.
    parent: Option<Rc<RefCell<Coroutine>>>,

    /// State
    state: State
}

/// Destroy coroutine and try to reuse std::stack segment.
impl Drop for Coroutine {
    fn drop(&mut self) {
        match self.current_stack_segment.take() {
            Some(stack) => {
                COROUTINE_ENVIRONMENT.with(|env| {
                    env.borrow_mut().stack_pool.give_stack(stack);
                });
            },
            None => {}
        }
    }
}

/// Coroutine spawn options
#[derive(Debug)]
pub struct Options {
    stack_size: usize,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            stack_size: min_stack(),
        }
    }
}

/// Initialization function for make context
extern "C" fn coroutine_initialize(_: usize, f: *mut ()) -> ! {
    let func: Box<Thunk> = unsafe { transmute(f) };

    if let Err(cause) = unsafe { try(move|| func.invoke(())) } {
        error!("Panicked inside: {:?}", cause.downcast::<&str>());
    }

    COROUTINE_ENVIRONMENT.with(|env| {
        let mut env = env.borrow_mut();
        env.current_running.borrow_mut().state = State::Finished;
    });

    loop {
        Coroutine::sched();
    }
}

impl Coroutine {
    pub fn empty() -> Rc<RefCell<Coroutine>> {
        Rc::new(RefCell::new(Coroutine {
            current_stack_segment: None,
            saved_context: Context::empty(),
            parent: None,
            state: State::Running,
        }))
    }

    pub fn spawn_opts<F>(f: F, opts: Options) -> Rc<RefCell<Coroutine>>
            where F: FnOnce() + Send + 'static {
        let stack = RefCell::new(unsafe { Stack::dummy_stack() });
        COROUTINE_ENVIRONMENT.with(|env| {
            *stack.borrow_mut() = env.borrow_mut().stack_pool.take_stack(opts.stack_size);
        });

        let mut stack = stack.into_inner();
        let coro = Coroutine::empty();

        let ctx = Context::new(coroutine_initialize,
                               0,
                               f,
                               &mut stack);
        coro.borrow_mut().saved_context = ctx;
        coro.borrow_mut().current_stack_segment = Some(stack);

        Coroutine::resume(&coro);
        coro
    }

    /// Spawn a coroutine with default options
    pub fn spawn<F>(f: F) -> Rc<RefCell<Coroutine>>
            where F: FnOnce() + Send + 'static {
        Coroutine::spawn_opts(f, Default::default())
    }

    pub fn resume(coro: &Rc<RefCell<Coroutine>>) {
        match coro.borrow().state {
            State::Finished => return,
            _ => {}
        }

        COROUTINE_ENVIRONMENT.with(|env| {
            let current_running = env.borrow().current_running.clone();

            let to_coro: &mut Coroutine = unsafe { transmute(coro.borrow().deref()) };
            let from_coro: &mut Coroutine = unsafe { transmute(current_running.borrow().deref()) };

            to_coro.parent = Some(current_running.clone());
            env.borrow_mut().current_running = coro.clone();

            from_coro.state = State::Suspended;
            to_coro.state = State::Running;

            Context::swap(&mut from_coro.saved_context, &to_coro.saved_context);
        })
    }

    pub fn sched() {
        COROUTINE_ENVIRONMENT.with(|env| {
            let from_coro: &mut Coroutine = unsafe { transmute(env.borrow().current_running.deref()) };
            let to_coro: &mut Coroutine = match from_coro.parent {
                Some(ref parent) => unsafe {
                    env.borrow_mut().current_running = parent.clone();
                    transmute(parent.borrow().deref())
                },
                None => return,
            };

            match from_coro.state {
                State::Finished => {},
                _ => from_coro.state = State::Suspended,
            }
            to_coro.state = State::Running;

            Context::swap(&mut from_coro.saved_context, &to_coro.saved_context);
        });
    }

    pub fn current() -> Rc<RefCell<Coroutine>> {
        let opt = RefCell::new(None);
        COROUTINE_ENVIRONMENT.with(|env| {
            *opt.borrow_mut() = Some(env.borrow().current_running.clone());
        });
        opt.into_inner().unwrap()
    }

    pub fn state(&self) -> State {
        self.state
    }
}

thread_local!(static COROUTINE_ENVIRONMENT: RefCell<Environment> = RefCell::new(Environment::new()));

/// Coroutine managing environment
#[allow(raw_pointer_derive)]
#[derive(Debug)]
struct Environment {
    stack_pool: StackPool,
    current_running: Rc<RefCell<Coroutine>>,
}

impl Environment {
    /// Initialize a new environment
    fn new() -> Environment {
        Environment {
            stack_pool: StackPool::new(),
            current_running: Coroutine::empty(),
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::mpsc::channel;

    use test::Bencher;

    use coroutine::Coroutine;

    #[test]
    fn test_coroutine_basic() {
        let (tx, rx) = channel();
        Coroutine::spawn(move|| {
            tx.send(1).unwrap();
        });

        assert_eq!(rx.recv().unwrap(), 1);
    }

    #[test]
    fn test_coroutine_yield() {
        let (tx, rx) = channel();
        let coro = Coroutine::spawn(move|| {
            tx.send(1).unwrap();

            Coroutine::sched();

            tx.send(2).unwrap();
        });

        assert_eq!(rx.recv().unwrap(), 1);
        assert!(rx.try_recv().is_err());

        Coroutine::resume(&coro);

        assert_eq!(rx.recv().unwrap(), 2);
    }

    #[test]
    fn test_coroutine_spawn_inside() {
        let (tx, rx) = channel();
        Coroutine::spawn(move|| {
            tx.send(1).unwrap();

            Coroutine::spawn(move|| {
                tx.send(2).unwrap();
            });
        });

        assert_eq!(rx.recv().unwrap(), 1);
        assert_eq!(rx.recv().unwrap(), 2);
    }

    #[test]
    #[should_fail]
    fn test_coroutine_panic() {
        Coroutine::spawn(move|| {
            panic!("Panic inside a coroutine!!");
        });

        unreachable!();
    }

    #[test]
    #[should_fail]
    fn test_coroutine_child_panic() {
        Coroutine::spawn(move|| {

            Coroutine::spawn(move|| {
                panic!("Panic inside a coroutine's child!!");
            });

        });

        unreachable!();
    }

    #[test]
    fn test_coroutine_resume_after_finished() {
        let coro = Coroutine::spawn(move|| {});

        // It is already finished, but we try to resume it
        // Idealy it would come back immediately
        Coroutine::resume(&coro);

        // Again?
        Coroutine::resume(&coro);
    }

    #[bench]
    fn bench_coroutine_spawning_with_recycle(b: &mut Bencher) {
        b.iter(|| {
            Coroutine::spawn(move|| {});
        });
    }

    #[bench]
    fn bench_coroutine_counting(b: &mut Bencher) {
        b.iter(|| {
            const MAX_NUMBER: usize = 100;
            let (tx, rx) = channel();

            let coro = Coroutine::spawn(move|| {
                for _ in 0..MAX_NUMBER {
                    tx.send(1).unwrap();
                    Coroutine::sched();
                }
            });

            let result = rx.iter().fold(0, |a, b| {
                Coroutine::resume(&coro);
                a + b
            });
            assert_eq!(result, MAX_NUMBER);
        });
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2017 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

use gtk;
use gtk::{ContainerExt, IsA, Object, WidgetExt};

use super::{Component, DisplayVariant, Relm, create_component, create_widget, init_component};
use widget::Widget;

/// Struct for relm containers to add GTK+ and relm `Widget`s.
pub struct ContainerComponent<WIDGET: Container + Widget> {
    component: Component<WIDGET>,
    /// The default container of this component.
    pub container: WIDGET::Container,
    /// Additional containers used for multi-containers. This can be () if not needed.
    pub containers: WIDGET::Containers,
}

impl<WIDGET: Container + Widget> ContainerComponent<WIDGET> {
    #[doc(hidden)]
    pub fn new(component: Component<WIDGET>, container: WIDGET::Container, containers: WIDGET::Containers) -> Self {
        ContainerComponent {
            component,
            container,
            containers,
        }
    }

    /// Add a GTK+ widget to a relm container.
    pub fn add<CHILDWIDGET: IsA<gtk::Widget>>(&self, widget: &CHILDWIDGET) {
        self.container.add(widget);
    }

    /// Add a relm widget to a relm container.
    pub fn add_widget<CHILDWIDGET, PARENTWIDGET>(&self, relm: &Relm<PARENTWIDGET>, model_param: CHILDWIDGET::ModelParam)
        -> Component<CHILDWIDGET>
        where CHILDWIDGET: Widget + 'static,
              PARENTWIDGET: Widget,
              WIDGET::Container: ContainerExt + IsA<gtk::Widget> + IsA<Object>,
    {
        let component = create_component::<CHILDWIDGET, _>(relm, model_param);
        WIDGET::add_widget(&self, &component);
        component
    }

    // TODO: add delete methods?

    /// Get the widget of the component.
    pub fn widget(&self) -> &WIDGET::Root {
        self.component.widget()
    }
}

/// Trait to implement relm container widget.
pub trait Container: Widget {
    /// The type of the containing widget, i.e. where the child widgets will be added.
    type Container: Clone + IsA<gtk::Container> + IsA<Object> + IsA<gtk::Widget>;
    /// Type to contain the additional container widgets.
    // TODO: put that in yet another trait?
    type Containers;

    /// Add a relm widget to this container.
    fn add_widget<WIDGET: Widget>(container: &ContainerComponent<Self>, component: &Component<WIDGET>) {
        container.container.add(component.widget());
    }

    /// Get the containing widget, i.e. the widget where the children will be added.
    fn container(&self) -> &Self::Container;

    /// Get additional container widgets.
    /// This is useful to create a multi-container.
    fn other_containers(&self) -> Self::Containers;
}

/// Extension trait for GTK+ containers to add and remove relm `Widget`s.
pub trait ContainerWidget {
    /// Add a relm `Container` to the current GTK+ container.
    ///
    /// # Note
    ///
    /// The returned `ContainerComponent` must be stored in a `Widget`. If it is not stored, a
    /// communication receiver will be droped which will cause events to be ignored for this
    /// widget.
    fn add_container<CHILDWIDGET, WIDGET>(&self, relm: &Relm<WIDGET>, model_param: CHILDWIDGET::ModelParam)
            -> ContainerComponent<CHILDWIDGET>
        where CHILDWIDGET: Container + Widget + 'static,
              CHILDWIDGET::Msg: DisplayVariant + 'static,
              CHILDWIDGET::Root: IsA<gtk::Widget> + IsA<Object> + WidgetExt,
              WIDGET: Widget;

    /// Add a relm `Widget` to the current GTK+ container.
    ///
    /// # Note
    ///
    /// The returned `Component` must be stored in a `Widget`. If it is not stored, a communication
    /// receiver will be droped which will cause events to be ignored for this widget.
    fn add_widget<CHILDWIDGET, WIDGET>(&self, relm: &Relm<WIDGET>, model_param: CHILDWIDGET::ModelParam)
            -> Component<CHILDWIDGET>
        where CHILDWIDGET: Widget + 'static,
              CHILDWIDGET::Msg: DisplayVariant + 'static,
              CHILDWIDGET::Root: IsA<gtk::Widget> + IsA<Object> + WidgetExt,
              WIDGET: Widget;

    /// Remove a relm `Widget` from the current GTK+ container.
    fn remove_widget<CHILDWIDGET>(&self, component: Component<CHILDWIDGET>)
        where CHILDWIDGET: Widget,
              CHILDWIDGET::Root: IsA<gtk::Widget>;
}

impl<W: Clone + ContainerExt + IsA<gtk::Widget> + IsA<Object>> ContainerWidget for W {
    fn add_container<CHILDWIDGET, WIDGET>(&self, relm: &Relm<WIDGET>, model_param: CHILDWIDGET::ModelParam)
            -> ContainerComponent<CHILDWIDGET>
        where CHILDWIDGET: Container + Widget + 'static,
              CHILDWIDGET::Msg: DisplayVariant + 'static,
              CHILDWIDGET::Root: IsA<gtk::Widget> + IsA<Object> + WidgetExt,
              WIDGET: Widget,
    {
        let (widget, component, child_relm) = create_widget::<CHILDWIDGET>(relm.context(), model_param);
        let container = component.container().clone();
        let containers = component.other_containers();
        let root = component.root().clone();
        self.add(&root);
        component.on_add(self.clone());
        init_component::<CHILDWIDGET>(widget.stream(), component, relm.context(), &child_relm);
        ContainerComponent::new(widget, container, containers)
    }

    fn add_widget<CHILDWIDGET, WIDGET>(&self, relm: &Relm<WIDGET>, model_param: CHILDWIDGET::ModelParam)
            -> Component<CHILDWIDGET>
        where CHILDWIDGET: Widget + 'static,
              CHILDWIDGET::Msg: DisplayVariant + 'static,
              CHILDWIDGET::Root: IsA<gtk::Widget> + IsA<Object> + WidgetExt,
    {
        let (widget, component, child_relm) = create_widget::<CHILDWIDGET>(relm.context(), model_param);
        self.add(widget.widget());
        component.on_add(self.clone());
        init_component::<CHILDWIDGET>(widget.stream(), component, relm.context(), &child_relm);
        widget
    }

    fn remove_widget<WIDGET>(&self, component: Component<WIDGET>)
        where WIDGET: Widget,
              WIDGET::Root: IsA<gtk::Widget>,
    {
        self.remove(component.widget());
    }
}
<commit_msg>Add stream() method for ContainerComponent<commit_after>/*
 * Copyright (c) 2017 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

use gtk;
use gtk::{ContainerExt, IsA, Object, WidgetExt};

use relm_state::EventStream;
use super::{Component, DisplayVariant, Relm, create_component, create_widget, init_component};
use widget::Widget;

/// Struct for relm containers to add GTK+ and relm `Widget`s.
pub struct ContainerComponent<WIDGET: Container + Widget> {
    component: Component<WIDGET>,
    /// The default container of this component.
    pub container: WIDGET::Container,
    /// Additional containers used for multi-containers. This can be () if not needed.
    pub containers: WIDGET::Containers,
}

impl<WIDGET: Container + Widget> ContainerComponent<WIDGET> {
    #[doc(hidden)]
    pub fn new(component: Component<WIDGET>, container: WIDGET::Container, containers: WIDGET::Containers) -> Self {
        ContainerComponent {
            component,
            container,
            containers,
        }
    }

    /// Add a GTK+ widget to a relm container.
    pub fn add<CHILDWIDGET: IsA<gtk::Widget>>(&self, widget: &CHILDWIDGET) {
        self.container.add(widget);
    }

    /// Add a relm widget to a relm container.
    pub fn add_widget<CHILDWIDGET, PARENTWIDGET>(&self, relm: &Relm<PARENTWIDGET>, model_param: CHILDWIDGET::ModelParam)
        -> Component<CHILDWIDGET>
        where CHILDWIDGET: Widget + 'static,
              PARENTWIDGET: Widget,
              WIDGET::Container: ContainerExt + IsA<gtk::Widget> + IsA<Object>,
    {
        let component = create_component::<CHILDWIDGET, _>(relm, model_param);
        WIDGET::add_widget(&self, &component);
        component
    }

    /// Get the event stream of the component.
    /// This is used internally by the library.
    pub fn stream(&self) -> &EventStream<WIDGET::Msg> {
        &self.component.stream()
    }

    // TODO: add delete methods?

    /// Get the widget of the component.
    pub fn widget(&self) -> &WIDGET::Root {
        self.component.widget()
    }
}

/// Trait to implement relm container widget.
pub trait Container: Widget {
    /// The type of the containing widget, i.e. where the child widgets will be added.
    type Container: Clone + IsA<gtk::Container> + IsA<Object> + IsA<gtk::Widget>;
    /// Type to contain the additional container widgets.
    // TODO: put that in yet another trait?
    type Containers;

    /// Add a relm widget to this container.
    fn add_widget<WIDGET: Widget>(container: &ContainerComponent<Self>, component: &Component<WIDGET>) {
        container.container.add(component.widget());
    }

    /// Get the containing widget, i.e. the widget where the children will be added.
    fn container(&self) -> &Self::Container;

    /// Get additional container widgets.
    /// This is useful to create a multi-container.
    fn other_containers(&self) -> Self::Containers;
}

/// Extension trait for GTK+ containers to add and remove relm `Widget`s.
pub trait ContainerWidget {
    /// Add a relm `Container` to the current GTK+ container.
    ///
    /// # Note
    ///
    /// The returned `ContainerComponent` must be stored in a `Widget`. If it is not stored, a
    /// communication receiver will be droped which will cause events to be ignored for this
    /// widget.
    fn add_container<CHILDWIDGET, WIDGET>(&self, relm: &Relm<WIDGET>, model_param: CHILDWIDGET::ModelParam)
            -> ContainerComponent<CHILDWIDGET>
        where CHILDWIDGET: Container + Widget + 'static,
              CHILDWIDGET::Msg: DisplayVariant + 'static,
              CHILDWIDGET::Root: IsA<gtk::Widget> + IsA<Object> + WidgetExt,
              WIDGET: Widget;

    /// Add a relm `Widget` to the current GTK+ container.
    ///
    /// # Note
    ///
    /// The returned `Component` must be stored in a `Widget`. If it is not stored, a communication
    /// receiver will be droped which will cause events to be ignored for this widget.
    fn add_widget<CHILDWIDGET, WIDGET>(&self, relm: &Relm<WIDGET>, model_param: CHILDWIDGET::ModelParam)
            -> Component<CHILDWIDGET>
        where CHILDWIDGET: Widget + 'static,
              CHILDWIDGET::Msg: DisplayVariant + 'static,
              CHILDWIDGET::Root: IsA<gtk::Widget> + IsA<Object> + WidgetExt,
              WIDGET: Widget;

    /// Remove a relm `Widget` from the current GTK+ container.
    fn remove_widget<CHILDWIDGET>(&self, component: Component<CHILDWIDGET>)
        where CHILDWIDGET: Widget,
              CHILDWIDGET::Root: IsA<gtk::Widget>;
}

impl<W: Clone + ContainerExt + IsA<gtk::Widget> + IsA<Object>> ContainerWidget for W {
    fn add_container<CHILDWIDGET, WIDGET>(&self, relm: &Relm<WIDGET>, model_param: CHILDWIDGET::ModelParam)
            -> ContainerComponent<CHILDWIDGET>
        where CHILDWIDGET: Container + Widget + 'static,
              CHILDWIDGET::Msg: DisplayVariant + 'static,
              CHILDWIDGET::Root: IsA<gtk::Widget> + IsA<Object> + WidgetExt,
              WIDGET: Widget,
    {
        let (widget, component, child_relm) = create_widget::<CHILDWIDGET>(relm.context(), model_param);
        let container = component.container().clone();
        let containers = component.other_containers();
        let root = component.root().clone();
        self.add(&root);
        component.on_add(self.clone());
        init_component::<CHILDWIDGET>(widget.stream(), component, relm.context(), &child_relm);
        ContainerComponent::new(widget, container, containers)
    }

    fn add_widget<CHILDWIDGET, WIDGET>(&self, relm: &Relm<WIDGET>, model_param: CHILDWIDGET::ModelParam)
            -> Component<CHILDWIDGET>
        where CHILDWIDGET: Widget + 'static,
              CHILDWIDGET::Msg: DisplayVariant + 'static,
              CHILDWIDGET::Root: IsA<gtk::Widget> + IsA<Object> + WidgetExt,
    {
        let (widget, component, child_relm) = create_widget::<CHILDWIDGET>(relm.context(), model_param);
        self.add(widget.widget());
        component.on_add(self.clone());
        init_component::<CHILDWIDGET>(widget.stream(), component, relm.context(), &child_relm);
        widget
    }

    fn remove_widget<WIDGET>(&self, component: Component<WIDGET>)
        where WIDGET: Widget,
              WIDGET::Root: IsA<gtk::Widget>,
    {
        self.remove(component.widget());
    }
}
<|endoftext|>


<commit_before>//! Return value handling

use crate::prelude::*;

use rustc_middle::ty::layout::FnAbiExt;
use rustc_target::abi::call::{ArgAbi, FnAbi, PassMode};
use smallvec::{smallvec, SmallVec};

/// Can the given type be returned into an ssa var or does it need to be returned on the stack.
pub(crate) fn can_return_to_ssa_var<'tcx>(
    fx: &FunctionCx<'_, '_, 'tcx>,
    func: &mir::Operand<'tcx>,
    args: &[mir::Operand<'tcx>],
) -> bool {
    let fn_ty = fx.monomorphize(func.ty(fx.mir, fx.tcx));
    let fn_sig =
        fx.tcx.normalize_erasing_late_bound_regions(ParamEnv::reveal_all(), fn_ty.fn_sig(fx.tcx));

    // Handle special calls like instrinsics and empty drop glue.
    let instance = if let ty::FnDef(def_id, substs) = *fn_ty.kind() {
        let instance = ty::Instance::resolve(fx.tcx, ty::ParamEnv::reveal_all(), def_id, substs)
            .unwrap()
            .unwrap()
            .polymorphize(fx.tcx);

        match instance.def {
            InstanceDef::Intrinsic(_) | InstanceDef::DropGlue(_, _) => {
                return true;
            }
            _ => Some(instance),
        }
    } else {
        None
    };

    let extra_args = &args[fn_sig.inputs().len()..];
    let extra_args = extra_args
        .iter()
        .map(|op_arg| fx.monomorphize(op_arg.ty(fx.mir, fx.tcx)))
        .collect::<Vec<_>>();
    let fn_abi = if let Some(instance) = instance {
        FnAbi::of_instance(&RevealAllLayoutCx(fx.tcx), instance, &extra_args)
    } else {
        FnAbi::of_fn_ptr(&RevealAllLayoutCx(fx.tcx), fn_ty.fn_sig(fx.tcx), &extra_args)
    };
    match fn_abi.ret.mode {
        PassMode::Ignore | PassMode::Direct(_) | PassMode::Pair(_, _) => true,
        // FIXME Make it possible to return Cast and Indirect to an ssa var.
        PassMode::Cast(_) | PassMode::Indirect { .. } => false,
    }
}

/// Return a place where the return value of the current function can be written to. If necessary
/// this adds an extra parameter pointing to where the return value needs to be stored.
pub(super) fn codegen_return_param<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    ssa_analyzed: &rustc_index::vec::IndexVec<Local, crate::analyze::SsaKind>,
    block_params_iter: &mut impl Iterator<Item = Value>,
) -> CPlace<'tcx> {
    let (ret_place, ret_param): (_, SmallVec<[_; 2]>) = match fx.fn_abi.as_ref().unwrap().ret.mode {
        PassMode::Ignore => (CPlace::no_place(fx.fn_abi.as_ref().unwrap().ret.layout), smallvec![]),
        PassMode::Direct(_) | PassMode::Pair(_, _) | PassMode::Cast(_) => {
            let is_ssa = ssa_analyzed[RETURN_PLACE] == crate::analyze::SsaKind::Ssa;
            (
                super::make_local_place(
                    fx,
                    RETURN_PLACE,
                    fx.fn_abi.as_ref().unwrap().ret.layout,
                    is_ssa,
                ),
                smallvec![],
            )
        }
        PassMode::Indirect { attrs: _, extra_attrs: None, on_stack: _ } => {
            let ret_param = block_params_iter.next().unwrap();
            assert_eq!(fx.bcx.func.dfg.value_type(ret_param), fx.pointer_type);
            (
                CPlace::for_ptr(Pointer::new(ret_param), fx.fn_abi.as_ref().unwrap().ret.layout),
                smallvec![ret_param],
            )
        }
        PassMode::Indirect { attrs: _, extra_attrs: Some(_), on_stack: _ } => {
            unreachable!("unsized return value")
        }
    };

    crate::abi::comments::add_arg_comment(
        fx,
        "ret",
        Some(RETURN_PLACE),
        None,
        &ret_param,
        fx.fn_abi.as_ref().unwrap().ret.mode,
        fx.fn_abi.as_ref().unwrap().ret.layout,
    );

    ret_place
}

/// Invokes the closure with if necessary a value representing the return pointer. When the closure
/// returns the call return value(s) if any are written to the correct place.
pub(super) fn codegen_with_call_return_arg<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    ret_arg_abi: &ArgAbi<'tcx, Ty<'tcx>>,
    ret_place: Option<CPlace<'tcx>>,
    f: impl FnOnce(&mut FunctionCx<'_, '_, 'tcx>, Option<Value>) -> Inst,
) {
    let return_ptr = match ret_arg_abi.mode {
        PassMode::Ignore => None,
        PassMode::Indirect { attrs: _, extra_attrs: None, on_stack: _ } => match ret_place {
            Some(ret_place) => Some(ret_place.to_ptr().get_addr(fx)),
            None => Some(fx.bcx.ins().iconst(fx.pointer_type, 43)), // FIXME allocate temp stack slot
        },
        PassMode::Indirect { attrs: _, extra_attrs: Some(_), on_stack: _ } => {
            unreachable!("unsized return value")
        }
        PassMode::Direct(_) | PassMode::Pair(_, _) | PassMode::Cast(_) => None,
    };

    let call_inst = f(fx, return_ptr);

    match ret_arg_abi.mode {
        PassMode::Ignore => {}
        PassMode::Direct(_) => {
            if let Some(ret_place) = ret_place {
                let ret_val = fx.bcx.inst_results(call_inst)[0];
                ret_place.write_cvalue(fx, CValue::by_val(ret_val, ret_arg_abi.layout));
            }
        }
        PassMode::Pair(_, _) => {
            if let Some(ret_place) = ret_place {
                let ret_val_a = fx.bcx.inst_results(call_inst)[0];
                let ret_val_b = fx.bcx.inst_results(call_inst)[1];
                ret_place.write_cvalue(
                    fx,
                    CValue::by_val_pair(ret_val_a, ret_val_b, ret_arg_abi.layout),
                );
            }
        }
        PassMode::Cast(cast) => {
            if let Some(ret_place) = ret_place {
                let results = fx
                    .bcx
                    .inst_results(call_inst)
                    .iter()
                    .copied()
                    .collect::<SmallVec<[Value; 2]>>();
                let result =
                    super::pass_mode::from_casted_value(fx, &results, ret_place.layout(), cast);
                ret_place.write_cvalue(fx, result);
            }
        }
        PassMode::Indirect { attrs: _, extra_attrs: None, on_stack: _ } => {}
        PassMode::Indirect { attrs: _, extra_attrs: Some(_), on_stack: _ } => {
            unreachable!("unsized return value")
        }
    }
}

/// Codegen a return instruction with the right return value(s) if any.
pub(crate) fn codegen_return(fx: &mut FunctionCx<'_, '_, '_>) {
    match fx.fn_abi.as_ref().unwrap().ret.mode {
        PassMode::Ignore | PassMode::Indirect { attrs: _, extra_attrs: None, on_stack: _ } => {
            fx.bcx.ins().return_(&[]);
        }
        PassMode::Indirect { attrs: _, extra_attrs: Some(_), on_stack: _ } => {
            unreachable!("unsized return value")
        }
        PassMode::Direct(_) => {
            let place = fx.get_local_place(RETURN_PLACE);
            let ret_val = place.to_cvalue(fx).load_scalar(fx);
            fx.bcx.ins().return_(&[ret_val]);
        }
        PassMode::Pair(_, _) => {
            let place = fx.get_local_place(RETURN_PLACE);
            let (ret_val_a, ret_val_b) = place.to_cvalue(fx).load_scalar_pair(fx);
            fx.bcx.ins().return_(&[ret_val_a, ret_val_b]);
        }
        PassMode::Cast(cast) => {
            let place = fx.get_local_place(RETURN_PLACE);
            let ret_val = place.to_cvalue(fx);
            let ret_vals = super::pass_mode::to_casted_value(fx, ret_val, cast);
            fx.bcx.ins().return_(&ret_vals);
        }
    }
}
<commit_msg>Allow returning PassMode::Cast directly to an ssa var<commit_after>//! Return value handling

use crate::prelude::*;

use rustc_middle::ty::layout::FnAbiExt;
use rustc_target::abi::call::{ArgAbi, FnAbi, PassMode};
use smallvec::{smallvec, SmallVec};

/// Can the given type be returned into an ssa var or does it need to be returned on the stack.
pub(crate) fn can_return_to_ssa_var<'tcx>(
    fx: &FunctionCx<'_, '_, 'tcx>,
    func: &mir::Operand<'tcx>,
    args: &[mir::Operand<'tcx>],
) -> bool {
    let fn_ty = fx.monomorphize(func.ty(fx.mir, fx.tcx));
    let fn_sig =
        fx.tcx.normalize_erasing_late_bound_regions(ParamEnv::reveal_all(), fn_ty.fn_sig(fx.tcx));

    // Handle special calls like instrinsics and empty drop glue.
    let instance = if let ty::FnDef(def_id, substs) = *fn_ty.kind() {
        let instance = ty::Instance::resolve(fx.tcx, ty::ParamEnv::reveal_all(), def_id, substs)
            .unwrap()
            .unwrap()
            .polymorphize(fx.tcx);

        match instance.def {
            InstanceDef::Intrinsic(_) | InstanceDef::DropGlue(_, _) => {
                return true;
            }
            _ => Some(instance),
        }
    } else {
        None
    };

    let extra_args = &args[fn_sig.inputs().len()..];
    let extra_args = extra_args
        .iter()
        .map(|op_arg| fx.monomorphize(op_arg.ty(fx.mir, fx.tcx)))
        .collect::<Vec<_>>();
    let fn_abi = if let Some(instance) = instance {
        FnAbi::of_instance(&RevealAllLayoutCx(fx.tcx), instance, &extra_args)
    } else {
        FnAbi::of_fn_ptr(&RevealAllLayoutCx(fx.tcx), fn_ty.fn_sig(fx.tcx), &extra_args)
    };
    match fn_abi.ret.mode {
        PassMode::Ignore | PassMode::Direct(_) | PassMode::Pair(_, _) | PassMode::Cast(_) => true,
        // FIXME Make it possible to return Indirect to an ssa var.
        PassMode::Indirect { .. } => false,
    }
}

/// Return a place where the return value of the current function can be written to. If necessary
/// this adds an extra parameter pointing to where the return value needs to be stored.
pub(super) fn codegen_return_param<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    ssa_analyzed: &rustc_index::vec::IndexVec<Local, crate::analyze::SsaKind>,
    block_params_iter: &mut impl Iterator<Item = Value>,
) -> CPlace<'tcx> {
    let (ret_place, ret_param): (_, SmallVec<[_; 2]>) = match fx.fn_abi.as_ref().unwrap().ret.mode {
        PassMode::Ignore => (CPlace::no_place(fx.fn_abi.as_ref().unwrap().ret.layout), smallvec![]),
        PassMode::Direct(_) | PassMode::Pair(_, _) | PassMode::Cast(_) => {
            let is_ssa = ssa_analyzed[RETURN_PLACE] == crate::analyze::SsaKind::Ssa;
            (
                super::make_local_place(
                    fx,
                    RETURN_PLACE,
                    fx.fn_abi.as_ref().unwrap().ret.layout,
                    is_ssa,
                ),
                smallvec![],
            )
        }
        PassMode::Indirect { attrs: _, extra_attrs: None, on_stack: _ } => {
            let ret_param = block_params_iter.next().unwrap();
            assert_eq!(fx.bcx.func.dfg.value_type(ret_param), fx.pointer_type);
            (
                CPlace::for_ptr(Pointer::new(ret_param), fx.fn_abi.as_ref().unwrap().ret.layout),
                smallvec![ret_param],
            )
        }
        PassMode::Indirect { attrs: _, extra_attrs: Some(_), on_stack: _ } => {
            unreachable!("unsized return value")
        }
    };

    crate::abi::comments::add_arg_comment(
        fx,
        "ret",
        Some(RETURN_PLACE),
        None,
        &ret_param,
        fx.fn_abi.as_ref().unwrap().ret.mode,
        fx.fn_abi.as_ref().unwrap().ret.layout,
    );

    ret_place
}

/// Invokes the closure with if necessary a value representing the return pointer. When the closure
/// returns the call return value(s) if any are written to the correct place.
pub(super) fn codegen_with_call_return_arg<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    ret_arg_abi: &ArgAbi<'tcx, Ty<'tcx>>,
    ret_place: Option<CPlace<'tcx>>,
    f: impl FnOnce(&mut FunctionCx<'_, '_, 'tcx>, Option<Value>) -> Inst,
) {
    let return_ptr = match ret_arg_abi.mode {
        PassMode::Ignore => None,
        PassMode::Indirect { attrs: _, extra_attrs: None, on_stack: _ } => match ret_place {
            Some(ret_place) => Some(ret_place.to_ptr().get_addr(fx)),
            None => Some(fx.bcx.ins().iconst(fx.pointer_type, 43)), // FIXME allocate temp stack slot
        },
        PassMode::Indirect { attrs: _, extra_attrs: Some(_), on_stack: _ } => {
            unreachable!("unsized return value")
        }
        PassMode::Direct(_) | PassMode::Pair(_, _) | PassMode::Cast(_) => None,
    };

    let call_inst = f(fx, return_ptr);

    match ret_arg_abi.mode {
        PassMode::Ignore => {}
        PassMode::Direct(_) => {
            if let Some(ret_place) = ret_place {
                let ret_val = fx.bcx.inst_results(call_inst)[0];
                ret_place.write_cvalue(fx, CValue::by_val(ret_val, ret_arg_abi.layout));
            }
        }
        PassMode::Pair(_, _) => {
            if let Some(ret_place) = ret_place {
                let ret_val_a = fx.bcx.inst_results(call_inst)[0];
                let ret_val_b = fx.bcx.inst_results(call_inst)[1];
                ret_place.write_cvalue(
                    fx,
                    CValue::by_val_pair(ret_val_a, ret_val_b, ret_arg_abi.layout),
                );
            }
        }
        PassMode::Cast(cast) => {
            if let Some(ret_place) = ret_place {
                let results = fx
                    .bcx
                    .inst_results(call_inst)
                    .iter()
                    .copied()
                    .collect::<SmallVec<[Value; 2]>>();
                let result =
                    super::pass_mode::from_casted_value(fx, &results, ret_place.layout(), cast);
                ret_place.write_cvalue(fx, result);
            }
        }
        PassMode::Indirect { attrs: _, extra_attrs: None, on_stack: _ } => {}
        PassMode::Indirect { attrs: _, extra_attrs: Some(_), on_stack: _ } => {
            unreachable!("unsized return value")
        }
    }
}

/// Codegen a return instruction with the right return value(s) if any.
pub(crate) fn codegen_return(fx: &mut FunctionCx<'_, '_, '_>) {
    match fx.fn_abi.as_ref().unwrap().ret.mode {
        PassMode::Ignore | PassMode::Indirect { attrs: _, extra_attrs: None, on_stack: _ } => {
            fx.bcx.ins().return_(&[]);
        }
        PassMode::Indirect { attrs: _, extra_attrs: Some(_), on_stack: _ } => {
            unreachable!("unsized return value")
        }
        PassMode::Direct(_) => {
            let place = fx.get_local_place(RETURN_PLACE);
            let ret_val = place.to_cvalue(fx).load_scalar(fx);
            fx.bcx.ins().return_(&[ret_val]);
        }
        PassMode::Pair(_, _) => {
            let place = fx.get_local_place(RETURN_PLACE);
            let (ret_val_a, ret_val_b) = place.to_cvalue(fx).load_scalar_pair(fx);
            fx.bcx.ins().return_(&[ret_val_a, ret_val_b]);
        }
        PassMode::Cast(cast) => {
            let place = fx.get_local_place(RETURN_PLACE);
            let ret_val = place.to_cvalue(fx);
            let ret_vals = super::pass_mode::to_casted_value(fx, ret_val, cast);
            fx.bcx.ins().return_(&ret_vals);
        }
    }
}
<|endoftext|>