<commit_before>// Copyright 2013 The CGMath Developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// 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.

use std::num::{zero, one, cast};

use angle::{Angle, tan, cot};
use frustum::Frustum;
use matrix::{Mat4, ToMat4};
use plane::Plane;

/// Create a perspective projection matrix.
///
/// This is the equivalent to the [gluPerspective]
/// (http://www.opengl.org/sdk/docs/man2/xhtml/gluPerspective.xml) function.
pub fn perspective<S: Float, A: Angle<S>>(fovy: A, aspect: S, near: S, far: S) -> Mat4<S> {
    PerspectiveFov {
        fovy:   fovy,
        aspect: aspect,
        near:   near,
        far:    far,
    }.to_mat4()
}

/// Create a perspective matrix from a view frustrum.
///
/// This is the equivalent of the now deprecated [glFrustrum]
/// (http://www.opengl.org/sdk/docs/man2/xhtml/glFrustum.xml) function.
pub fn frustum<S: Float>(left: S, right: S, bottom: S, top: S, near: S, far: S) -> Mat4<S> {
    Perspective {
        left:   left,
        right:  right,
        bottom: bottom,
        top:    top,
        near:   near,
        far:    far,
    }.to_mat4()
}

/// Create an orthographic projection matrix.
///
/// This is the equivalent of the now deprecated [glOrtho]
/// (http://www.opengl.org/sdk/docs/man2/xhtml/glOrtho.xml) function.
pub fn ortho<S: Float>(left: S, right: S, bottom: S, top: S, near: S, far: S) -> Mat4<S> {
    Ortho {
        left:   left,
        right:  right,
        bottom: bottom,
        top:    top,
        near:   near,
        far:    far,
    }.to_mat4()
}

pub trait Projection<S>: ToMat4<S> {
    fn to_frustum(&self) -> Frustum<S>;
}

/// A perspective projection based on a vertical field-of-view angle.
#[deriving(Clone, Eq)]
pub struct PerspectiveFov<S, A> {
    fovy:   A,
    aspect: S,
    near:   S,
    far:    S,
}

impl<S: Float, A: Angle<S>> PerspectiveFov<S, A> {
    pub fn to_perspective(&self) -> Perspective<S> {
        let angle = self.fovy.div_s(cast(2).unwrap());
        let ymax = self.near * tan(angle.to_rad());
        let xmax = ymax * self.aspect;

        Perspective {
            left:   -xmax,
            right:   xmax,
            bottom: -ymax,
            top:     ymax,
            near:    self.near.clone(),
            far:     self.far.clone(),
        }
    }
}

impl<S: Float, A: Angle<S>> Projection<S> for PerspectiveFov<S, A> {
    fn to_frustum(&self) -> Frustum<S> {
        // TODO: Could this be faster?
        Frustum::from_mat4(self.to_mat4())
    }
}

impl<S: Float, A: Angle<S>> ToMat4<S> for PerspectiveFov<S, A> {
    fn to_mat4(&self) -> Mat4<S> {
        let half_turn: A = Angle::turn_div_2();

        assert!(self.fovy   > zero(),    "The vertical field of view cannot be below zero, found: {:?}", self.fovy);
        assert!(self.fovy   < half_turn, "The vertical field of view cannot be greater than a half turn, found: {:?}", self.fovy);
        assert!(self.aspect > zero(),    "The aspect ratio cannot be below zero, found: {:?}", self.aspect);
        assert!(self.near   > zero(),    "The near plane distance cannot be below zero, found: {:?}", self.near);
        assert!(self.far    > zero(),    "The far plane distance cannot be below zero, found: {:?}", self.far);
        assert!(self.far    > self.near, "The far plane cannot be closer than the near plane, found: far: {:?}, near: {:?}", self.far, self.near);

        let f = cot(self.fovy.div_s(cast(2).unwrap()).to_rad());
        let two: S = cast(2).unwrap();

        let c0r0 = f / self.aspect;
        let c0r1 = zero();
        let c0r2 = zero();
        let c0r3 = zero();

        let c1r0 = zero();
        let c1r1 = f;
        let c1r2 = zero();
        let c1r3 = zero();

        let c2r0 = zero();
        let c2r1 = zero();
        let c2r2 = (self.far + self.near) / (self.near - self.far);
        let c2r3 = -one::<S>();

        let c3r0 = zero();
        let c3r1 = zero();
        let c3r2 = (two * self.far * self.near) / (self.near - self.far);
        let c3r3 = zero();

        Mat4::new(c0r0, c0r1, c0r2, c0r3,
                  c1r0, c1r1, c1r2, c1r3,
                  c2r0, c2r1, c2r2, c2r3,
                  c3r0, c3r1, c3r2, c3r3)
    }
}

/// A perspective projection with arbitrary left/right/bottom/top distances
#[deriving(Clone, Eq)]
pub struct Perspective<S> {
    left:   S,  right:  S,
    bottom: S,  top:    S,
    near:   S,  far:    S,
}

impl<S: Float> Projection<S> for Perspective<S> {
    fn to_frustum(&self) -> Frustum<S> {
        // TODO: Could this be faster?
        Frustum::from_mat4(self.to_mat4())
    }
}

impl<S: Float> ToMat4<S> for Perspective<S> {
    fn to_mat4(&self) -> Mat4<S> {
        assert!(self.left   > self.right, "`left` cannot be greater than `right`, found: left: {:?} right: {:?}", self.left, self.right);
        assert!(self.bottom > self.top,   "`bottom` cannot be greater than `top`, found: bottom: {:?} top: {:?}", self.bottom, self.top);
        assert!(self.near   > self.far,   "`near` cannot be greater than `far`, found: near: {:?} far: {:?}", self.near, self.far);

        let two: S = cast(2).unwrap();

        let c0r0 = (two * self.near) / (self.right - self.left);
        let c0r1 = zero();
        let c0r2 = zero();
        let c0r3 = zero();

        let c1r0 = zero();
        let c1r1 = (two * self.near) / (self.top - self.bottom);
        let c1r2 = zero();
        let c1r3 = zero();

        let c2r0 = (self.right + self.left) / (self.right - self.left);
        let c2r1 = (self.top + self.bottom) / (self.top - self.bottom);
        let c2r2 = -(self.far + self.near) / (self.far - self.near);
        let c2r3 = -one::<S>();

        let c3r0 = zero();
        let c3r1 = zero();
        let c3r2 = -(two * self.far * self.near) / (self.far - self.near);
        let c3r3 = zero();

        Mat4::new(c0r0, c0r1, c0r2, c0r3,
                  c1r0, c1r1, c1r2, c1r3,
                  c2r0, c2r1, c2r2, c2r3,
                  c3r0, c3r1, c3r2, c3r3)
    }
}

/// An orthographic projection with arbitrary left/right/bottom/top distances
#[deriving(Clone, Eq)]
pub struct Ortho<S> {
    left:   S,  right:  S,
    bottom: S,  top:    S,
    near:   S,  far:    S,
}

impl<S: Float> Projection<S> for Ortho<S> {
    fn to_frustum(&self) -> Frustum<S> {
        Frustum {
            left:   Plane::from_abcd( one::<S>(), zero::<S>(), zero::<S>(), self.left.clone()),
            right:  Plane::from_abcd(-one::<S>(), zero::<S>(), zero::<S>(), self.right.clone()),
            bottom: Plane::from_abcd(zero::<S>(),  one::<S>(), zero::<S>(), self.bottom.clone()),
            top:    Plane::from_abcd(zero::<S>(), -one::<S>(), zero::<S>(), self.top.clone()),
            near:   Plane::from_abcd(zero::<S>(), zero::<S>(), -one::<S>(), self.near.clone()),
            far:    Plane::from_abcd(zero::<S>(), zero::<S>(),  one::<S>(), self.far.clone()),
        }
    }
}

impl<S: Float> ToMat4<S> for Ortho<S> {
    fn to_mat4(&self) -> Mat4<S> {
        assert!(self.left   > self.right, "`left` cannot be greater than `right`, found: left: {:?} right: {:?}", self.left, self.right);
        assert!(self.bottom > self.top,   "`bottom` cannot be greater than `top`, found: bottom: {:?} top: {:?}", self.bottom, self.top);
        assert!(self.near   > self.far,   "`near` cannot be greater than `far`, found: near: {:?} far: {:?}", self.near, self.far);

        let two: S = cast(2).unwrap();

        let c0r0 = two / (self.right - self.left);
        let c0r1 = zero();
        let c0r2 = zero();
        let c0r3 = zero();

        let c1r0 = zero();
        let c1r1 = two / (self.top - self.bottom);
        let c1r2 = zero();
        let c1r3 = zero();

        let c2r0 = zero();
        let c2r1 = zero();
        let c2r2 = -two / (self.far - self.near);
        let c2r3 = -one::<S>();

        let c3r0 = -(self.right + self.left) / (self.right - self.left);
        let c3r1 = -(self.top + self.bottom) / (self.top - self.bottom);
        let c3r2 = -(self.far + self.near) / (self.far - self.near);
        let c3r3 = one::<S>();

        Mat4::new(c0r0, c0r1, c0r2, c0r3,
                  c1r0, c1r1, c1r2, c1r3,
                  c2r0, c2r1, c2r2, c2r3,
                  c3r0, c3r1, c3r2, c3r3)
    }
}
<commit_msg>Fix ortho asserts<commit_after>// Copyright 2013 The CGMath Developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// 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.

use std::num::{zero, one, cast};

use angle::{Angle, tan, cot};
use frustum::Frustum;
use matrix::{Mat4, ToMat4};
use plane::Plane;

/// Create a perspective projection matrix.
///
/// This is the equivalent to the [gluPerspective]
/// (http://www.opengl.org/sdk/docs/man2/xhtml/gluPerspective.xml) function.
pub fn perspective<S: Float, A: Angle<S>>(fovy: A, aspect: S, near: S, far: S) -> Mat4<S> {
    PerspectiveFov {
        fovy:   fovy,
        aspect: aspect,
        near:   near,
        far:    far,
    }.to_mat4()
}

/// Create a perspective matrix from a view frustrum.
///
/// This is the equivalent of the now deprecated [glFrustrum]
/// (http://www.opengl.org/sdk/docs/man2/xhtml/glFrustum.xml) function.
pub fn frustum<S: Float>(left: S, right: S, bottom: S, top: S, near: S, far: S) -> Mat4<S> {
    Perspective {
        left:   left,
        right:  right,
        bottom: bottom,
        top:    top,
        near:   near,
        far:    far,
    }.to_mat4()
}

/// Create an orthographic projection matrix.
///
/// This is the equivalent of the now deprecated [glOrtho]
/// (http://www.opengl.org/sdk/docs/man2/xhtml/glOrtho.xml) function.
pub fn ortho<S: Float>(left: S, right: S, bottom: S, top: S, near: S, far: S) -> Mat4<S> {
    Ortho {
        left:   left,
        right:  right,
        bottom: bottom,
        top:    top,
        near:   near,
        far:    far,
    }.to_mat4()
}

pub trait Projection<S>: ToMat4<S> {
    fn to_frustum(&self) -> Frustum<S>;
}

/// A perspective projection based on a vertical field-of-view angle.
#[deriving(Clone, Eq)]
pub struct PerspectiveFov<S, A> {
    fovy:   A,
    aspect: S,
    near:   S,
    far:    S,
}

impl<S: Float, A: Angle<S>> PerspectiveFov<S, A> {
    pub fn to_perspective(&self) -> Perspective<S> {
        let angle = self.fovy.div_s(cast(2).unwrap());
        let ymax = self.near * tan(angle.to_rad());
        let xmax = ymax * self.aspect;

        Perspective {
            left:   -xmax,
            right:   xmax,
            bottom: -ymax,
            top:     ymax,
            near:    self.near.clone(),
            far:     self.far.clone(),
        }
    }
}

impl<S: Float, A: Angle<S>> Projection<S> for PerspectiveFov<S, A> {
    fn to_frustum(&self) -> Frustum<S> {
        // TODO: Could this be faster?
        Frustum::from_mat4(self.to_mat4())
    }
}

impl<S: Float, A: Angle<S>> ToMat4<S> for PerspectiveFov<S, A> {
    fn to_mat4(&self) -> Mat4<S> {
        let half_turn: A = Angle::turn_div_2();

        assert!(self.fovy   > zero(),    "The vertical field of view cannot be below zero, found: {:?}", self.fovy);
        assert!(self.fovy   < half_turn, "The vertical field of view cannot be greater than a half turn, found: {:?}", self.fovy);
        assert!(self.aspect > zero(),    "The aspect ratio cannot be below zero, found: {:?}", self.aspect);
        assert!(self.near   > zero(),    "The near plane distance cannot be below zero, found: {:?}", self.near);
        assert!(self.far    > zero(),    "The far plane distance cannot be below zero, found: {:?}", self.far);
        assert!(self.far    > self.near, "The far plane cannot be closer than the near plane, found: far: {:?}, near: {:?}", self.far, self.near);

        let f = cot(self.fovy.div_s(cast(2).unwrap()).to_rad());
        let two: S = cast(2).unwrap();

        let c0r0 = f / self.aspect;
        let c0r1 = zero();
        let c0r2 = zero();
        let c0r3 = zero();

        let c1r0 = zero();
        let c1r1 = f;
        let c1r2 = zero();
        let c1r3 = zero();

        let c2r0 = zero();
        let c2r1 = zero();
        let c2r2 = (self.far + self.near) / (self.near - self.far);
        let c2r3 = -one::<S>();

        let c3r0 = zero();
        let c3r1 = zero();
        let c3r2 = (two * self.far * self.near) / (self.near - self.far);
        let c3r3 = zero();

        Mat4::new(c0r0, c0r1, c0r2, c0r3,
                  c1r0, c1r1, c1r2, c1r3,
                  c2r0, c2r1, c2r2, c2r3,
                  c3r0, c3r1, c3r2, c3r3)
    }
}

/// A perspective projection with arbitrary left/right/bottom/top distances
#[deriving(Clone, Eq)]
pub struct Perspective<S> {
    left:   S,  right:  S,
    bottom: S,  top:    S,
    near:   S,  far:    S,
}

impl<S: Float> Projection<S> for Perspective<S> {
    fn to_frustum(&self) -> Frustum<S> {
        // TODO: Could this be faster?
        Frustum::from_mat4(self.to_mat4())
    }
}

impl<S: Float> ToMat4<S> for Perspective<S> {
    fn to_mat4(&self) -> Mat4<S> {
        assert!(self.left   > self.right, "`left` cannot be greater than `right`, found: left: {:?} right: {:?}", self.left, self.right);
        assert!(self.bottom > self.top,   "`bottom` cannot be greater than `top`, found: bottom: {:?} top: {:?}", self.bottom, self.top);
        assert!(self.near   > self.far,   "`near` cannot be greater than `far`, found: near: {:?} far: {:?}", self.near, self.far);

        let two: S = cast(2).unwrap();

        let c0r0 = (two * self.near) / (self.right - self.left);
        let c0r1 = zero();
        let c0r2 = zero();
        let c0r3 = zero();

        let c1r0 = zero();
        let c1r1 = (two * self.near) / (self.top - self.bottom);
        let c1r2 = zero();
        let c1r3 = zero();

        let c2r0 = (self.right + self.left) / (self.right - self.left);
        let c2r1 = (self.top + self.bottom) / (self.top - self.bottom);
        let c2r2 = -(self.far + self.near) / (self.far - self.near);
        let c2r3 = -one::<S>();

        let c3r0 = zero();
        let c3r1 = zero();
        let c3r2 = -(two * self.far * self.near) / (self.far - self.near);
        let c3r3 = zero();

        Mat4::new(c0r0, c0r1, c0r2, c0r3,
                  c1r0, c1r1, c1r2, c1r3,
                  c2r0, c2r1, c2r2, c2r3,
                  c3r0, c3r1, c3r2, c3r3)
    }
}

/// An orthographic projection with arbitrary left/right/bottom/top distances
#[deriving(Clone, Eq)]
pub struct Ortho<S> {
    left:   S,  right:  S,
    bottom: S,  top:    S,
    near:   S,  far:    S,
}

impl<S: Float> Projection<S> for Ortho<S> {
    fn to_frustum(&self) -> Frustum<S> {
        Frustum {
            left:   Plane::from_abcd( one::<S>(), zero::<S>(), zero::<S>(), self.left.clone()),
            right:  Plane::from_abcd(-one::<S>(), zero::<S>(), zero::<S>(), self.right.clone()),
            bottom: Plane::from_abcd(zero::<S>(),  one::<S>(), zero::<S>(), self.bottom.clone()),
            top:    Plane::from_abcd(zero::<S>(), -one::<S>(), zero::<S>(), self.top.clone()),
            near:   Plane::from_abcd(zero::<S>(), zero::<S>(), -one::<S>(), self.near.clone()),
            far:    Plane::from_abcd(zero::<S>(), zero::<S>(),  one::<S>(), self.far.clone()),
        }
    }
}

impl<S: Float> ToMat4<S> for Ortho<S> {
    fn to_mat4(&self) -> Mat4<S> {
        assert!(self.left   < self.right, "`left` cannot be greater than `right`, found: left: {:?} right: {:?}", self.left, self.right);
        assert!(self.bottom < self.top,   "`bottom` cannot be greater than `top`, found: bottom: {:?} top: {:?}", self.bottom, self.top);
        assert!(self.near   < self.far,   "`near` cannot be greater than `far`, found: near: {:?} far: {:?}", self.near, self.far);

        let two: S = cast(2).unwrap();

        let c0r0 = two / (self.right - self.left);
        let c0r1 = zero();
        let c0r2 = zero();
        let c0r3 = zero();

        let c1r0 = zero();
        let c1r1 = two / (self.top - self.bottom);
        let c1r2 = zero();
        let c1r3 = zero();

        let c2r0 = zero();
        let c2r1 = zero();
        let c2r2 = -two / (self.far - self.near);
        let c2r3 = -one::<S>();

        let c3r0 = -(self.right + self.left) / (self.right - self.left);
        let c3r1 = -(self.top + self.bottom) / (self.top - self.bottom);
        let c3r2 = -(self.far + self.near) / (self.far - self.near);
        let c3r3 = one::<S>();

        Mat4::new(c0r0, c0r1, c0r2, c0r3,
                  c1r0, c1r1, c1r2, c1r3,
                  c2r0, c2r1, c2r2, c2r3,
                  c3r0, c3r1, c3r2, c3r3)
    }
}
<|endoftext|>


<commit_before>//! Massage and work the auzten file to make sure we can deal with real data.

extern crate las;

macro_rules! autzen {
    ($name:ident, $major:expr, $minor:expr) => {
        mod $name {
            use std::io::Cursor;

            use las::{Reader, Header, Writer};

            const VERSION: (u8, u8) = ($major, $minor);

            #[test]
            fn read_write() {
                let mut reader = Reader::from_path("tests/data/autzen.las").unwrap();
                let mut writer = Writer::new(Cursor::new(Vec::new()), Header {
                    version: VERSION,
                    ..Default::default()
                }).unwrap();
                for point in reader.points() {
                    writer.write(&point.unwrap()).unwrap();
                }
                writer.close().unwrap();
            }
        }
    }
}

autzen!(las_1_0, 1, 0);
autzen!(las_1_1, 1, 1);
autzen!(las_1_2, 1, 2);
autzen!(las_1_3, 1, 3);
autzen!(las_1_4, 1, 4);
<commit_msg>Remove trailing whitespace<commit_after>//! Massage and work the auzten file to make sure we can deal with real data.

extern crate las;

macro_rules! autzen {
    ($name:ident, $major:expr, $minor:expr) => {
        mod $name {
            use std::io::Cursor;

            use las::{Reader, Header, Writer};

            const VERSION: (u8, u8) = ($major, $minor);

            #[test]
            fn read_write() {
                let mut reader = Reader::from_path("tests/data/autzen.las").unwrap();
                let mut writer = Writer::new(Cursor::new(Vec::new()), Header {
                    version: VERSION,
                    ..Default::default()
                }).unwrap();
                for point in reader.points() {
                    writer.write(&point.unwrap()).unwrap();
                }
                writer.close().unwrap();
            }
        }
    }
}

autzen!(las_1_0, 1, 0);
autzen!(las_1_1, 1, 1);
autzen!(las_1_2, 1, 2);
autzen!(las_1_3, 1, 3);
autzen!(las_1_4, 1, 4);
<|endoftext|>


<commit_before>#![feature(phase)]
#![allow(uppercase_variables)]
#![allow(unused_imports)]

extern crate test;
extern crate ndarray;

use ndarray::{Array, S, Si};
use ndarray::{arr0, arr1, arr2};

use test::black_box;

#[bench]
fn bench_std_add(bench: &mut test::Bencher)
{
    let mut a = arr2::<f32>([[1., 2., 2.],
                             [3., 4., 4.],
                             [3., 4., 4.],
                             [3., 4., 4.],
                             [5., 6., 6.]]);
    let b = a.clone();
    bench.iter(|| a.iadd(&b));
}

#[bench]
fn bench_std_iter_1d(bench: &mut test::Bencher)
{
    let a = arr1::<f32>([1., 2., 2.,
                         3., 4., 4.,
                         3., 4., 4.,
                         3., 4., 4.,
                         5., 6., 6.]);
    bench.iter(|| for elt in a.iter() { black_box(elt) })
}

#[bench]
fn bench_std_iter_2d(bench: &mut test::Bencher)
{
    let a = arr2::<f32>([[1., 2., 2.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [5., 6., 6.]]);
    bench.iter(|| for elt in a.iter() { black_box(elt) })
}
<commit_msg>Add mat mul bench<commit_after>#![feature(phase)]
#![allow(uppercase_variables)]
#![allow(unused_imports)]

extern crate test;
extern crate ndarray;

use ndarray::{Array, S, Si};
use ndarray::{arr0, arr1, arr2};

use test::black_box;

#[bench]
fn bench_std_add(bench: &mut test::Bencher)
{
    let mut a = arr2::<f32>([[1., 2., 2.],
                             [3., 4., 4.],
                             [3., 4., 4.],
                             [3., 4., 4.],
                             [5., 6., 6.]]);
    let b = a.clone();
    bench.iter(|| a.iadd(&b));
}

#[bench]
fn bench_std_iter_1d(bench: &mut test::Bencher)
{
    let a = arr1::<f32>([1., 2., 2.,
                         3., 4., 4.,
                         3., 4., 4.,
                         3., 4., 4.,
                         5., 6., 6.]);
    bench.iter(|| for elt in a.iter() { black_box(elt) })
}

#[bench]
fn bench_std_iter_2d(bench: &mut test::Bencher)
{
    let a = arr2::<f32>([[1., 2., 2.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [5., 6., 6.]]);
    bench.iter(|| for elt in a.iter() { black_box(elt) })
}

#[bench]
fn bench_mat_mul(bench: &mut test::Bencher)
{
    let a = arr2::<f32>([[1., 2., 2.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [3., 4., 4.],
                         [5., 6., 6.]]);
    let mut at = a.clone();
    at.swap_axes(0, 1);
    bench.iter(|| at.mat_mul(&a));
}
<|endoftext|>


<commit_before>use common::conn::Conn;
use common::stream::HttpStreamCommon;
use common::stream::HttpStreamData;
use common::types::Types;
use solicit::StreamId;

use data_or_headers_with_flag::DataOrHeadersWithFlag;

use bytes::Bytes;
use common::conn::ConnStateSnapshot;
use common::conn_read::ConnReadSideCustom;
use common::iteration_exit::IterationExit;
use common::stream::HttpStreamCommand;
use error;
use futures::sync::oneshot;
use futures::task;
use futures::Async;
use futures::Poll;
use result;
use solicit::end_stream::EndStream;
use solicit::frame::continuation::ContinuationFlag;
use solicit::frame::ContinuationFrame;
use solicit::frame::DataFlag;
use solicit::frame::DataFrame;
use solicit::frame::GoawayFrame;
use solicit::frame::HeadersFlag;
use solicit::frame::HeadersFrame;
use solicit::frame::HttpFrame;
use solicit::frame::RstStreamFrame;
use solicit::frame::SettingsFrame;
use std::cmp;
use ErrorCode;

pub trait ConnWriteSideCustom {
    type Types: Types;

    fn process_message(
        &mut self,
        message: <Self::Types as Types>::ToWriteMessage,
    ) -> result::Result<()>;
}

impl<T> Conn<T>
where
    T: Types,
    Self: ConnReadSideCustom<Types = T>,
    Self: ConnWriteSideCustom<Types = T>,
    HttpStreamCommon<T>: HttpStreamData<Types = T>,
{
    fn write_part(&mut self, stream_id: StreamId, part: HttpStreamCommand) {
        let max_frame_size = self.peer_settings.max_frame_size as usize;

        match part {
            HttpStreamCommand::Data(data, end_stream) => {
                // if client requested end of stream,
                // we must send at least one frame with end stream flag
                if end_stream == EndStream::Yes && data.len() == 0 {
                    let mut frame = DataFrame::with_data(stream_id, Bytes::new());
                    frame.set_flag(DataFlag::EndStream);

                    debug!("sending frame {:?}", frame);

                    self.queued_write.queue(frame);

                    return;
                }

                let mut pos = 0;
                while pos < data.len() {
                    let end = cmp::min(data.len(), pos + max_frame_size);

                    let end_stream_in_frame = if end == data.len() && end_stream == EndStream::Yes {
                        EndStream::Yes
                    } else {
                        EndStream::No
                    };

                    let mut frame = DataFrame::with_data(stream_id, data.slice(pos, end));
                    if end_stream_in_frame == EndStream::Yes {
                        frame.set_flag(DataFlag::EndStream);
                    }

                    debug!("sending frame {:?}", frame);

                    self.queued_write.queue(frame);

                    pos = end;
                }
            }
            HttpStreamCommand::Headers(headers, end_stream) => {
                let headers_fragment = self
                    .encoder
                    .encode(headers.0.iter().map(|h| (h.name(), h.value())));

                let mut pos = 0;
                while pos < headers_fragment.len() || pos == 0 {
                    let end = cmp::min(headers_fragment.len(), pos + max_frame_size);

                    let chunk = headers_fragment.slice(pos, end);

                    if pos == 0 {
                        let mut frame = HeadersFrame::new(chunk, stream_id);
                        if end_stream == EndStream::Yes {
                            frame.set_flag(HeadersFlag::EndStream);
                        }

                        if end == headers_fragment.len() {
                            frame.set_flag(HeadersFlag::EndHeaders);
                        }

                        debug!("sending frame {:?}", frame);

                        self.queued_write.queue(frame);
                    } else {
                        let mut frame = ContinuationFrame::new(chunk, stream_id);

                        if end == headers_fragment.len() {
                            frame.set_flag(ContinuationFlag::EndHeaders);
                        }

                        debug!("sending frame {:?}", frame);

                        self.queued_write.queue(frame);
                    }

                    pos = end;
                }
            }
            HttpStreamCommand::Rst(error_code) => {
                let frame = RstStreamFrame::new(stream_id, error_code);

                debug!("sending frame {:?}", frame);

                self.queued_write.queue(frame);
            }
        }
    }

    fn has_write_buffer_capacity(&self) -> bool {
        self.queued_write.queued_bytes_len() < 0x8000
    }

    fn pop_outg_for_stream(
        &mut self,
        stream_id: StreamId,
    ) -> Option<(StreamId, HttpStreamCommand, bool)> {
        let stream = self.streams.get_mut(stream_id).unwrap();
        if let (Some(command), stream) = stream.pop_outg_maybe_remove(&mut self.out_window_size) {
            return Some((stream_id, command, stream.is_some()));
        }

        None
    }

    pub fn buffer_outg_conn(&mut self) -> result::Result<bool> {
        let mut updated = false;

        // shortcut
        if !self.has_write_buffer_capacity() {
            return Ok(updated);
        }

        let writable_stream_ids = self.streams.writable_stream_ids();

        for &stream_id in &writable_stream_ids {
            loop {
                if !self.has_write_buffer_capacity() {
                    return Ok(updated);
                }

                if let Some((stream_id, part, cont)) = self.pop_outg_for_stream(stream_id) {
                    self.write_part(stream_id, part);
                    updated = true;

                    // Stream is removed from map, need to continue to the next stream
                    if !cont {
                        break;
                    }
                } else {
                    break;
                }
            }
        }

        Ok(updated)
    }

    pub fn send_frame_and_notify<F: Into<HttpFrame>>(&mut self, frame: F) {
        // TODO: some of frames should not be in front of GOAWAY
        self.queued_write.queue(frame);
        // Notify the task to make sure write loop is called again
        // to flush the buffer
        task::current().notify();
    }

    /// Sends an SETTINGS Frame with ack set to acknowledge seeing a SETTINGS frame from the peer.
    pub fn send_ack_settings(&mut self) -> result::Result<()> {
        let settings = SettingsFrame::new_ack();
        self.send_frame_and_notify(settings);
        Ok(())
    }

    fn process_stream_end(
        &mut self,
        stream_id: StreamId,
        error_code: ErrorCode,
    ) -> result::Result<()> {
        let stream = self.streams.get_mut(stream_id);
        if let Some(mut stream) = stream {
            stream.close_outgoing(error_code);
        }
        Ok(())
    }

    fn process_stream_enqueue(
        &mut self,
        stream_id: StreamId,
        part: DataOrHeadersWithFlag,
    ) -> result::Result<()> {
        let stream = self.streams.get_mut(stream_id);
        if let Some(mut stream) = stream {
            stream.push_back_part(part);
        }
        Ok(())
    }

    pub fn process_common_message(&mut self, common: CommonToWriteMessage) -> result::Result<()> {
        match common {
            CommonToWriteMessage::StreamEnd(stream_id, error_code) => {
                self.process_stream_end(stream_id, error_code)?;
            }
            CommonToWriteMessage::StreamEnqueue(stream_id, part) => {
                self.process_stream_enqueue(stream_id, part)?;
            }
            CommonToWriteMessage::IncreaseInWindow(stream_id, increase) => {
                self.increase_in_window(stream_id, increase)?;
            }
            CommonToWriteMessage::DumpState(sender) => {
                self.process_dump_state(sender)?;
            }
        }
        Ok(())
    }

    pub fn send_goaway(&mut self, error_code: ErrorCode) -> result::Result<()> {
        debug!("requesting to send GOAWAY with code {:?}", error_code);
        let frame = GoawayFrame::new(self.last_peer_stream_id, error_code);
        self.queued_write.queue_goaway(frame);
        task::current().notify();
        Ok(())
    }

    pub fn process_goaway_state(&mut self) -> result::Result<IterationExit> {
        Ok(if self.queued_write.goaway_queued() {
            self.queued_write.poll()?;
            if self.queued_write.queued_empty() {
                IterationExit::ExitEarly
            } else {
                IterationExit::NotReady
            }
        } else {
            IterationExit::Continue
        })
    }

    fn process_write_queue(&mut self) -> Poll<(), error::Error> {
        loop {
            let message = match self.write_rx.poll()? {
                Async::NotReady => return Ok(Async::NotReady),
                Async::Ready(Some(message)) => message,
                Async::Ready(None) => return Ok(Async::Ready(())), // Add some diagnostics maybe?
            };

            self.process_message(message)?;
        }
    }

    fn poll_flush(&mut self) -> result::Result<()> {
        self.buffer_outg_conn()?;
        loop {
            self.queued_write.poll()?;
            let updated = self.buffer_outg_conn()?;
            if !updated {
                return Ok(());
            }
        }
    }

    pub fn poll_write(&mut self) -> Poll<(), error::Error> {
        if let Async::Ready(()) = self.process_write_queue()? {
            return Ok(Async::Ready(()));
        }

        self.poll_flush()?;

        Ok(Async::NotReady)
    }
}

// Message sent to write loop.
// Processed while write loop is not handling network I/O.
pub enum CommonToWriteMessage {
    IncreaseInWindow(StreamId, u32),
    StreamEnqueue(StreamId, DataOrHeadersWithFlag),
    StreamEnd(StreamId, ErrorCode), // send when user provided handler completed the stream
    DumpState(oneshot::Sender<ConnStateSnapshot>),
}
<commit_msg>Cleanup<commit_after>use common::conn::Conn;
use common::stream::HttpStreamCommon;
use common::stream::HttpStreamData;
use common::types::Types;
use solicit::StreamId;

use data_or_headers_with_flag::DataOrHeadersWithFlag;

use bytes::Bytes;
use common::conn::ConnStateSnapshot;
use common::conn_read::ConnReadSideCustom;
use common::iteration_exit::IterationExit;
use common::stream::HttpStreamCommand;
use error;
use futures::sync::oneshot;
use futures::task;
use futures::Async;
use futures::Poll;
use result;
use solicit::end_stream::EndStream;
use solicit::frame::continuation::ContinuationFlag;
use solicit::frame::ContinuationFrame;
use solicit::frame::DataFlag;
use solicit::frame::DataFrame;
use solicit::frame::GoawayFrame;
use solicit::frame::HeadersFlag;
use solicit::frame::HeadersFrame;
use solicit::frame::HttpFrame;
use solicit::frame::RstStreamFrame;
use solicit::frame::SettingsFrame;
use std::cmp;
use ErrorCode;
use Headers;

pub trait ConnWriteSideCustom {
    type Types: Types;

    fn process_message(
        &mut self,
        message: <Self::Types as Types>::ToWriteMessage,
    ) -> result::Result<()>;
}

impl<T> Conn<T>
where
    T: Types,
    Self: ConnReadSideCustom<Types = T>,
    Self: ConnWriteSideCustom<Types = T>,
    HttpStreamCommon<T>: HttpStreamData<Types = T>,
{
    fn write_part_data(&mut self, stream_id: StreamId, data: Bytes, end_stream: EndStream) {
        let max_frame_size = self.peer_settings.max_frame_size as usize;

        // if client requested end of stream,
        // we must send at least one frame with end stream flag
        if end_stream == EndStream::Yes && data.len() == 0 {
            let mut frame = DataFrame::with_data(stream_id, Bytes::new());
            frame.set_flag(DataFlag::EndStream);

            debug!("sending frame {:?}", frame);

            self.queued_write.queue(frame);

            return;
        }

        let mut pos = 0;
        while pos < data.len() {
            let end = cmp::min(data.len(), pos + max_frame_size);

            let end_stream_in_frame = if end == data.len() && end_stream == EndStream::Yes {
                EndStream::Yes
            } else {
                EndStream::No
            };

            let mut frame = DataFrame::with_data(stream_id, data.slice(pos, end));
            if end_stream_in_frame == EndStream::Yes {
                frame.set_flag(DataFlag::EndStream);
            }

            debug!("sending frame {:?}", frame);

            self.queued_write.queue(frame);

            pos = end;
        }
    }

    fn write_part_headers(&mut self, stream_id: StreamId, headers: Headers, end_stream: EndStream) {
        let max_frame_size = self.peer_settings.max_frame_size as usize;

        // TODO: avoid allocation
        let headers_fragment = self
            .encoder
            .encode(headers.0.iter().map(|h| (h.name(), h.value())));

        let mut pos = 0;
        while pos < headers_fragment.len() || pos == 0 {
            let end = cmp::min(headers_fragment.len(), pos + max_frame_size);

            let chunk = headers_fragment.slice(pos, end);

            if pos == 0 {
                let mut frame = HeadersFrame::new(chunk, stream_id);
                if end_stream == EndStream::Yes {
                    frame.set_flag(HeadersFlag::EndStream);
                }

                if end == headers_fragment.len() {
                    frame.set_flag(HeadersFlag::EndHeaders);
                }

                debug!("sending frame {:?}", frame);

                self.queued_write.queue(frame);
            } else {
                let mut frame = ContinuationFrame::new(chunk, stream_id);

                if end == headers_fragment.len() {
                    frame.set_flag(ContinuationFlag::EndHeaders);
                }

                debug!("sending frame {:?}", frame);

                self.queued_write.queue(frame);
            }

            pos = end;
        }
    }

    fn write_part_rst(&mut self, stream_id: StreamId, error_code: ErrorCode) {
        let frame = RstStreamFrame::new(stream_id, error_code);

        debug!("sending frame {:?}", frame);

        self.queued_write.queue(frame);
    }

    fn write_part(&mut self, stream_id: StreamId, part: HttpStreamCommand) {
        match part {
            HttpStreamCommand::Data(data, end_stream) => {
                self.write_part_data(stream_id, data, end_stream);
            }
            HttpStreamCommand::Headers(headers, end_stream) => {
                self.write_part_headers(stream_id, headers, end_stream);
            }
            HttpStreamCommand::Rst(error_code) => {
                self.write_part_rst(stream_id, error_code);
            }
        }
    }

    fn has_write_buffer_capacity(&self) -> bool {
        self.queued_write.queued_bytes_len() < 0x8000
    }

    fn pop_outg_for_stream(
        &mut self,
        stream_id: StreamId,
    ) -> Option<(StreamId, HttpStreamCommand, bool)> {
        let stream = self.streams.get_mut(stream_id).unwrap();
        if let (Some(command), stream) = stream.pop_outg_maybe_remove(&mut self.out_window_size) {
            return Some((stream_id, command, stream.is_some()));
        }

        None
    }

    pub fn buffer_outg_conn(&mut self) -> result::Result<bool> {
        let mut updated = false;

        // shortcut
        if !self.has_write_buffer_capacity() {
            return Ok(updated);
        }

        let writable_stream_ids = self.streams.writable_stream_ids();

        for &stream_id in &writable_stream_ids {
            loop {
                if !self.has_write_buffer_capacity() {
                    return Ok(updated);
                }

                if let Some((stream_id, part, cont)) = self.pop_outg_for_stream(stream_id) {
                    self.write_part(stream_id, part);
                    updated = true;

                    // Stream is removed from map, need to continue to the next stream
                    if !cont {
                        break;
                    }
                } else {
                    break;
                }
            }
        }

        Ok(updated)
    }

    pub fn send_frame_and_notify<F: Into<HttpFrame>>(&mut self, frame: F) {
        // TODO: some of frames should not be in front of GOAWAY
        self.queued_write.queue(frame);
        // Notify the task to make sure write loop is called again
        // to flush the buffer
        task::current().notify();
    }

    /// Sends an SETTINGS Frame with ack set to acknowledge seeing a SETTINGS frame from the peer.
    pub fn send_ack_settings(&mut self) -> result::Result<()> {
        let settings = SettingsFrame::new_ack();
        self.send_frame_and_notify(settings);
        Ok(())
    }

    fn process_stream_end(
        &mut self,
        stream_id: StreamId,
        error_code: ErrorCode,
    ) -> result::Result<()> {
        let stream = self.streams.get_mut(stream_id);
        if let Some(mut stream) = stream {
            stream.close_outgoing(error_code);
        }
        Ok(())
    }

    fn process_stream_enqueue(
        &mut self,
        stream_id: StreamId,
        part: DataOrHeadersWithFlag,
    ) -> result::Result<()> {
        let stream = self.streams.get_mut(stream_id);
        if let Some(mut stream) = stream {
            stream.push_back_part(part);
        }
        Ok(())
    }

    pub fn process_common_message(&mut self, common: CommonToWriteMessage) -> result::Result<()> {
        match common {
            CommonToWriteMessage::StreamEnd(stream_id, error_code) => {
                self.process_stream_end(stream_id, error_code)?;
            }
            CommonToWriteMessage::StreamEnqueue(stream_id, part) => {
                self.process_stream_enqueue(stream_id, part)?;
            }
            CommonToWriteMessage::IncreaseInWindow(stream_id, increase) => {
                self.increase_in_window(stream_id, increase)?;
            }
            CommonToWriteMessage::DumpState(sender) => {
                self.process_dump_state(sender)?;
            }
        }
        Ok(())
    }

    pub fn send_goaway(&mut self, error_code: ErrorCode) -> result::Result<()> {
        debug!("requesting to send GOAWAY with code {:?}", error_code);
        let frame = GoawayFrame::new(self.last_peer_stream_id, error_code);
        self.queued_write.queue_goaway(frame);
        task::current().notify();
        Ok(())
    }

    pub fn process_goaway_state(&mut self) -> result::Result<IterationExit> {
        Ok(if self.queued_write.goaway_queued() {
            self.queued_write.poll()?;
            if self.queued_write.queued_empty() {
                IterationExit::ExitEarly
            } else {
                IterationExit::NotReady
            }
        } else {
            IterationExit::Continue
        })
    }

    fn process_write_queue(&mut self) -> Poll<(), error::Error> {
        loop {
            let message = match self.write_rx.poll()? {
                Async::NotReady => return Ok(Async::NotReady),
                Async::Ready(Some(message)) => message,
                Async::Ready(None) => return Ok(Async::Ready(())), // Add some diagnostics maybe?
            };

            self.process_message(message)?;
        }
    }

    fn poll_flush(&mut self) -> result::Result<()> {
        self.buffer_outg_conn()?;
        loop {
            self.queued_write.poll()?;
            let updated = self.buffer_outg_conn()?;
            if !updated {
                return Ok(());
            }
        }
    }

    pub fn poll_write(&mut self) -> Poll<(), error::Error> {
        if let Async::Ready(()) = self.process_write_queue()? {
            return Ok(Async::Ready(()));
        }

        self.poll_flush()?;

        Ok(Async::NotReady)
    }
}

// Message sent to write loop.
// Processed while write loop is not handling network I/O.
pub enum CommonToWriteMessage {
    IncreaseInWindow(StreamId, u32),
    StreamEnqueue(StreamId, DataOrHeadersWithFlag),
    StreamEnd(StreamId, ErrorCode), // send when user provided handler completed the stream
    DumpState(oneshot::Sender<ConnStateSnapshot>),
}
<|endoftext|>


<commit_before>extern crate futures;
extern crate tokio;
extern crate tokio_io;
extern crate env_logger;

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

use futures::prelude::*;
use tokio::net::{TcpStream, TcpListener};
use tokio::runtime::Runtime;

macro_rules! t {
    ($e:expr) => (match $e {
        Ok(e) => e,
        Err(e) => panic!("{} failed with {:?}", stringify!($e), e),
    })
}

#[test]
fn hammer_old() {
    let _ = env_logger::init();

    let threads = (0..10).map(|_| {
        thread::spawn(|| {
            let srv = t!(TcpListener::bind(&"127.0.0.1:0".parse().unwrap()));
            let addr = t!(srv.local_addr());
            let mine = TcpStream::connect(&addr);
            let theirs = srv.incoming().into_future()
                .map(|(s, _)| s.unwrap())
                .map_err(|(s, _)| s);
            let (mine, theirs) = t!(mine.join(theirs).wait());

            assert_eq!(t!(mine.local_addr()), t!(theirs.peer_addr()));
            assert_eq!(t!(theirs.local_addr()), t!(mine.peer_addr()));
        })
    }).collect::<Vec<_>>();
    for thread in threads {
        thread.join().unwrap();
    }
}

struct Rd(Arc<TcpStream>);
struct Wr(Arc<TcpStream>);

impl io::Read for Rd {
    fn read(&mut self, dst: &mut [u8]) -> io::Result<usize> {
        <&TcpStream>::read(&mut &*self.0, dst)
    }
}

impl tokio_io::AsyncRead for Rd {
}

impl io::Write for Wr {
    fn write(&mut self, src: &[u8]) -> io::Result<usize> {
        <&TcpStream>::write(&mut &*self.0, src)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl tokio_io::AsyncWrite for Wr {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        Ok(().into())
    }
}

#[test]
fn hammer_split() {
    use tokio_io::io;

    const N: usize = 100;
    const ITER: usize = 100;

    let _ = env_logger::init();

    for _ in 0..ITER {
        let srv = t!(TcpListener::bind(&"127.0.0.1:0".parse().unwrap()));
        let addr = t!(srv.local_addr());

        let cnt = Arc::new(AtomicUsize::new(0));

        let mut rt = Runtime::new().unwrap();

        fn split(socket: TcpStream, cnt: Arc<AtomicUsize>) {
            let socket = Arc::new(socket);
            let rd = Rd(socket.clone());
            let wr = Wr(socket);

            let cnt2 = cnt.clone();

            let rd = io::read(rd, vec![0; 1])
                .map(move |_| {
                    cnt2.fetch_add(1, Relaxed);
                })
                .map_err(|e| panic!("read error = {:?}", e));

            let wr = io::write_all(wr, b"1")
                .map(move |_| {
                    cnt.fetch_add(1, Relaxed);
                })
                .map_err(move |e| panic!("write error = {:?}", e));

            tokio::spawn(rd);
            tokio::spawn(wr);
        }

        rt.spawn({
            let cnt = cnt.clone();
            srv.incoming()
                .map_err(|e| panic!("accept error = {:?}", e))
                .take(N as u64)
                .for_each(move |socket| {
                    split(socket, cnt.clone());
                    Ok(())
                })
        });

        for _ in 0..N {
            rt.spawn({
                let cnt = cnt.clone();
                TcpStream::connect(&addr)
                    .map_err(move |e| panic!("connect error = {:?}", e))
                    .map(move |socket| split(socket, cnt))
            });
        }

        rt.shutdown_on_idle().wait().unwrap();
        assert_eq!(N * 4, cnt.load(Relaxed));
    }
}
<commit_msg>Change hammer test consts (#255)<commit_after>extern crate futures;
extern crate tokio;
extern crate tokio_io;
extern crate env_logger;

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

use futures::prelude::*;
use tokio::net::{TcpStream, TcpListener};
use tokio::runtime::Runtime;

macro_rules! t {
    ($e:expr) => (match $e {
        Ok(e) => e,
        Err(e) => panic!("{} failed with {:?}", stringify!($e), e),
    })
}

#[test]
fn hammer_old() {
    let _ = env_logger::init();

    let threads = (0..10).map(|_| {
        thread::spawn(|| {
            let srv = t!(TcpListener::bind(&"127.0.0.1:0".parse().unwrap()));
            let addr = t!(srv.local_addr());
            let mine = TcpStream::connect(&addr);
            let theirs = srv.incoming().into_future()
                .map(|(s, _)| s.unwrap())
                .map_err(|(s, _)| s);
            let (mine, theirs) = t!(mine.join(theirs).wait());

            assert_eq!(t!(mine.local_addr()), t!(theirs.peer_addr()));
            assert_eq!(t!(theirs.local_addr()), t!(mine.peer_addr()));
        })
    }).collect::<Vec<_>>();
    for thread in threads {
        thread.join().unwrap();
    }
}

struct Rd(Arc<TcpStream>);
struct Wr(Arc<TcpStream>);

impl io::Read for Rd {
    fn read(&mut self, dst: &mut [u8]) -> io::Result<usize> {
        <&TcpStream>::read(&mut &*self.0, dst)
    }
}

impl tokio_io::AsyncRead for Rd {
}

impl io::Write for Wr {
    fn write(&mut self, src: &[u8]) -> io::Result<usize> {
        <&TcpStream>::write(&mut &*self.0, src)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl tokio_io::AsyncWrite for Wr {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        Ok(().into())
    }
}

#[test]
fn hammer_split() {
    use tokio_io::io;

    const N: usize = 100;
    const ITER: usize = 10;

    let _ = env_logger::init();

    for _ in 0..ITER {
        let srv = t!(TcpListener::bind(&"127.0.0.1:0".parse().unwrap()));
        let addr = t!(srv.local_addr());

        let cnt = Arc::new(AtomicUsize::new(0));

        let mut rt = Runtime::new().unwrap();

        fn split(socket: TcpStream, cnt: Arc<AtomicUsize>) {
            let socket = Arc::new(socket);
            let rd = Rd(socket.clone());
            let wr = Wr(socket);

            let cnt2 = cnt.clone();

            let rd = io::read(rd, vec![0; 1])
                .map(move |_| {
                    cnt2.fetch_add(1, Relaxed);
                })
                .map_err(|e| panic!("read error = {:?}", e));

            let wr = io::write_all(wr, b"1")
                .map(move |_| {
                    cnt.fetch_add(1, Relaxed);
                })
                .map_err(move |e| panic!("write error = {:?}", e));

            tokio::spawn(rd);
            tokio::spawn(wr);
        }

        rt.spawn({
            let cnt = cnt.clone();
            srv.incoming()
                .map_err(|e| panic!("accept error = {:?}", e))
                .take(N as u64)
                .for_each(move |socket| {
                    split(socket, cnt.clone());
                    Ok(())
                })
        });

        for _ in 0..N {
            rt.spawn({
                let cnt = cnt.clone();
                TcpStream::connect(&addr)
                    .map_err(move |e| panic!("connect error = {:?}", e))
                    .map(move |socket| split(socket, cnt))
            });
        }

        rt.shutdown_on_idle().wait().unwrap();
        assert_eq!(N * 4, cnt.load(Relaxed));
    }
}
<|endoftext|>


<commit_before>#![warn(missing_copy_implementations)]
#![warn(missing_debug_implementations)]
#![warn(trivial_casts)]
#![warn(trivial_numeric_casts)]
#![warn(unused_extern_crates)]
#![warn(unused_import_braces)]
#![warn(unused_qualifications)]
#![warn(unused_results)]

#![feature(plugin)]
#![plugin(power_assert)]

use std::thread;

#[derive(Debug)]
struct Foo {
    val: u32
}

#[derive(Debug)]
struct Bar {
    val: u32,
    foo: Foo
}

fn run_panic_test<F>(f: F) -> String
    where F: FnOnce(), F: Send + 'static
{
    thread::spawn(f).join().unwrap_err().downcast_ref::<String>().unwrap().clone()
}

#[test]
fn test_member_access() {
    assert_eq!(run_panic_test(|| {
        let bar = Bar { val: 3, foo: Foo { val: 2 }};
        power_assert!(bar.val == bar.foo.val);
    }), "assertion failed: bar.val == bar.foo.val
power_assert!(bar.val == bar.foo.val)
              |   |   |  |   |   |
              |   3   |  |   |   2
              |       |  |   Foo { val: 2 }
              |       |  Bar { val: 3, foo: Foo { val: 2 } }
              |       false
              Bar { val: 3, foo: Foo { val: 2 } }
");

    assert_eq!(run_panic_test(|| {
        let bar = Bar { val: 3, foo: Foo { val: 2 }};
        power_assert_eq!(bar.val, bar.foo.val);
    }), "assertion failed: `(left == right)` (left: `3`, right: `2`)
power_assert_eq!(bar.val, bar.foo.val)
left: bar.val
      |   |
      |   3
      Bar { val: 3, foo: Foo { val: 2 } }
right: bar.foo.val
       |   |   |
       |   |   2
       |   Foo { val: 2 }
       Bar { val: 3, foo: Foo { val: 2 } }
");
}

#[test]
fn issue4() {
    fn check(_: &str) -> Result<(), ()> {
        Ok(())
    }

    let s = "hello".to_owned();

    assert!(check(&s) == Ok(()));
    // power_assert!(check(&s) == Ok(())); // FIXME(#4)
}

#[test]
fn issue5() {
    assert_eq!(2*2, 4);
    power_assert_eq!(2*2, 4);
}
<commit_msg>Uncomment issue4 power_assert! now that https://github.com/rust-lang/rust/issues/26978 is fixed<commit_after>#![warn(missing_copy_implementations)]
#![warn(missing_debug_implementations)]
#![warn(trivial_casts)]
#![warn(trivial_numeric_casts)]
#![warn(unused_extern_crates)]
#![warn(unused_import_braces)]
#![warn(unused_qualifications)]
#![warn(unused_results)]

#![feature(plugin)]
#![plugin(power_assert)]

use std::thread;

#[derive(Debug)]
struct Foo {
    val: u32
}

#[derive(Debug)]
struct Bar {
    val: u32,
    foo: Foo
}

fn run_panic_test<F>(f: F) -> String
    where F: FnOnce(), F: Send + 'static
{
    thread::spawn(f).join().unwrap_err().downcast_ref::<String>().unwrap().clone()
}

#[test]
fn test_member_access() {
    assert_eq!(run_panic_test(|| {
        let bar = Bar { val: 3, foo: Foo { val: 2 }};
        power_assert!(bar.val == bar.foo.val);
    }), "assertion failed: bar.val == bar.foo.val
power_assert!(bar.val == bar.foo.val)
              |   |   |  |   |   |
              |   3   |  |   |   2
              |       |  |   Foo { val: 2 }
              |       |  Bar { val: 3, foo: Foo { val: 2 } }
              |       false
              Bar { val: 3, foo: Foo { val: 2 } }
");

    assert_eq!(run_panic_test(|| {
        let bar = Bar { val: 3, foo: Foo { val: 2 }};
        power_assert_eq!(bar.val, bar.foo.val);
    }), "assertion failed: `(left == right)` (left: `3`, right: `2`)
power_assert_eq!(bar.val, bar.foo.val)
left: bar.val
      |   |
      |   3
      Bar { val: 3, foo: Foo { val: 2 } }
right: bar.foo.val
       |   |   |
       |   |   2
       |   Foo { val: 2 }
       Bar { val: 3, foo: Foo { val: 2 } }
");
}

#[test]
fn issue4() {
    fn check(_: &str) -> Result<(), ()> {
        Ok(())
    }

    let s = "hello".to_owned();

    assert!(check(&s) == Ok(()));
    power_assert!(check(&s) == Ok(()));
}

#[test]
fn issue5() {
    assert_eq!(2*2, 4);
    power_assert_eq!(2*2, 4);
}
<|endoftext|>


<commit_before>use index::btree::dump::dump_tree;
use index::btree::external::ExtNode;
use index::btree::internal::InNode;
use index::btree::node::is_node_locked;
use index::btree::node::read_node;
use index::btree::node::read_unchecked;
use index::btree::node::write_node;
use index::btree::node::write_non_empty;
use index::btree::node::write_targeted;
use index::btree::node::EmptyNode;
use index::btree::node::Node;
use index::btree::node::NodeData;
use index::btree::node::NodeWriteGuard;
use index::btree::search::mut_search;
use index::btree::search::MutSearchResult;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
use index::btree::{external, BPlusTree};
use index::lsmtree::tree::{LEVEL_2, LEVEL_3, LEVEL_PAGE_DIFF_MULTIPLIER};
use index::EntryKey;
use index::Slice;
use itertools::free::all;
use itertools::Itertools;
use smallvec::SmallVec;
use std::cmp::min;
use std::collections::{BTreeSet, HashMap, HashSet};
use std::fmt::Debug;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::Arc;
use std::cell::RefCell;
use core::borrow::{Borrow, BorrowMut};
use index::btree::node::NodeData::Empty;
use serde_json::to_vec;

enum Selection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    Selected(Vec<NodeWriteGuard<KS, PS>>),
    Innode(NodeCellRef),
}

struct KeyAltered {
    new_key: EntryKey,
    ptr: NodeCellRef,
}

fn select<KS, PS>(node: &NodeCellRef) -> Vec<NodeWriteGuard<KS, PS>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let search = mut_search::<KS, PS>(node, &smallvec!());
    match search {
        MutSearchResult::External => {
            let first_node = write_node(node);
            let mut collected_keys = first_node.len();
            let mut collected = vec![first_node];
            let target_keys = min(KS::slice_len() * LEVEL_PAGE_DIFF_MULTIPLIER, LEVEL_3);
            let target_guards = LEVEL_2;
            while collected_keys < target_keys && collected.len() < target_guards {
                let right = write_node(
                    collected
                        .last_mut()
                        .unwrap()
                        .right_ref_mut_no_empty()
                        .unwrap(),
                );
                if right.is_none() {
                    break;
                } else {
                    collected_keys += right.len();
                    collected.push(right);
                }
            }
            return collected;
        }
        MutSearchResult::Internal(node) => select::<KS, PS>(&node),
    }
}

fn prune_selected<'a, KS, PS>(node: &NodeCellRef, bound: &EntryKey, level: usize) -> Vec<KeyAltered>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let mut level_page_altered = vec![];
    let first_search = mut_search::<KS, PS>(node, &smallvec!());
    let altered_keys = match first_search {
        MutSearchResult::Internal(sub_node_ref) => {
            let sub_node = read_unchecked::<KS, PS>(&sub_node_ref);
            if !sub_node.is_empty_node() && !sub_node.is_ext() {
                prune_selected::<KS, PS>(&sub_node_ref, bound, level + 1)
            } else {
                vec![]
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    let mut all_pages = vec![write_node::<KS, PS>(node)];
    // collect all pages in bound and in this level
    loop {
        let (right_ref, right_bound) = {
            let last_page = all_pages.last().unwrap().borrow();
            if last_page.is_none() {
                break;
            }
            (
                last_page.right_ref().unwrap().clone(),
                last_page.right_bound().clone(),
            )
        };
        let right = write_node::<KS, PS>(&right_ref);
        if right.is_none() {
            break;
        }
        all_pages.push(right);
        if &right_bound > bound {
            break;
        }
    }
    let select_live = |page: &NodeWriteGuard<KS, PS>| {
        page.innode().ptrs.as_slice_immute()[..page.len() + 1]
            .iter()
            .enumerate()
            .filter_map(|(i, sub_level)| {
                if !read_unchecked::<KS, PS>(sub_level).is_empty_node() {
                    Some((i, sub_level.clone()))
                } else {
                    None
                }
            })
            .collect_vec()
    };
    let page_lives = all_pages.iter().map(|p| select_live(p)).collect_vec();
    all_pages
        .iter_mut()
        .zip(page_lives)
        .for_each(|(mut page, live_ptrs)| {
            if live_ptrs.len() == 0 {
                **page = NodeData::Empty(box EmptyNode {
                    left: None,
                    right: Default::default(),
                })
            } else {
                let mut new_keys = KS::init();
                let mut new_ptrs = PS::init();
                let ptr_len = live_ptrs.len();
                for (i, &(oi, _)) in live_ptrs.iter().take(live_ptrs.len() - 1).enumerate() {
                    new_keys.as_slice()[i] = page.keys()[oi].clone();
                }
                for (i, (_, ptr)) in live_ptrs.into_iter().enumerate() {
                    new_ptrs.as_slice()[i] = ptr;
                }
                let innode = page.innode_mut();
                innode.len = ptr_len - 1;
                innode.keys = new_keys;
                innode.ptrs = new_ptrs;
            }
        });

    let non_empty_right_node = |i: usize, pages: &Vec<NodeWriteGuard<KS, PS>>| {
        let mut i = i + 1;
        loop {
            if i == pages.len() {
                return pages[i - 1].borrow().right_ref().unwrap().clone();
            } else if i < pages.len() {
                let p = pages[i].borrow();
                if !p.is_empty_node() {
                    return p.node_ref().clone();
                } else {
                    i += 1;
                }
            } else {
                unreachable!()
            }
        }
    };

    let right_ptrs = all_pages
        .iter()
        .enumerate()
        .map(|(i, _)| non_empty_right_node(i, &all_pages))
        .collect_vec();
    all_pages
        .iter_mut()
        .zip(right_ptrs.into_iter())
        .for_each(|(p, r)| *p.right_ref_mut().unwrap() = r);

    all_pages = all_pages
        .into_iter()
        .filter(|p| !p.borrow().is_empty_node())
        .collect_vec();

    let mut index = 0;
    while index < all_pages.len() {
        if all_pages[index].len() == 0 {
            // current page have one ptr and no keys
            // need to merge with the right page
            // if the right page is full, partial of the right page will be moved to the current page
            // merging right page will also been cleaned
            let mut next_from_ptr = if index + 1 >= all_pages.len() {
                Some(write_node::<KS, PS>(all_pages[index].right_ref().unwrap()))
            } else {
                None
            };
            let (keys, ptrs, right_bound, right_ref) = {
                let mut next = next_from_ptr.as_mut().unwrap_or_else(|| &mut all_pages[index + 1]);
                if next.len() < KS::slice_len() - 1 {
                    // Merge next node with current node
                    let tuple = {
                        let next_innode = next.innode();
                        let len = next_innode.len;
                        (
                            next_innode.keys.as_slice_immute()[..len].to_vec(),
                            next_innode.ptrs.as_slice_immute()[..len + 1].to_vec(),
                            next_innode.right_bound.clone(),
                            next_innode.right.clone()
                        )
                    };
                    **next = NodeData::Empty(box EmptyNode {
                        left: None,
                        right: NodeCellRef::default()
                    });
                    tuple
                } else {
                    // Rebalance next node with current node
                    // Next node left bound need to be updated in upper level
                    let right_ref = next.node_ref().clone();
                    let next_innode = next.innode_mut();
                    let next_len = next_innode.len;
                    let next_mid = next_len / 2;
                    let right_left_bound = next_innode.keys.as_slice_immute()[next_mid].clone();
                    let tuple = (
                        next_innode.keys.as_slice_immute()[..next_mid].to_vec(),
                        next_innode.ptrs.as_slice_immute()[..next_mid + 1].to_vec(),
                        right_left_bound.clone(),
                        right_ref.clone()
                    );
                    let mut keys = KS::init();
                    let mut ptrs = PS::init();

                    for (i, key) in next_innode.keys.as_slice_immute()[next_mid + 1..next_len].iter().enumerate() {
                        keys.as_slice()[i] = key.clone();
                    }
                    for (i, ptr) in next_innode.ptrs.as_slice_immute()[next_mid + 1..next_len + 1].iter().enumerate() {
                        ptrs.as_slice()[i] = ptr.clone();
                    }
                    next_innode.keys = keys;
                    next_innode.ptrs = ptrs;
                    next_innode.len = next_len - next_mid - 1;

                    level_page_altered.push(KeyAltered {
                        new_key: right_left_bound,
                        ptr: right_ref
                    });
                    tuple
                }
            };
            let page_innode = all_pages[index].innode_mut();
            page_innode.keys.as_slice()[0] = page_innode.right_bound.clone();
            page_innode.right_bound = right_bound;
            page_innode.right = right_ref;
            page_innode.len = 1 + keys.len();
            for (i, key) in keys.into_iter().enumerate() {
                page_innode.keys.as_slice()[i + 1] = key;
            }
            for (i, ptr) in ptrs.into_iter().enumerate() {
                page_innode.ptrs.as_slice()[i + 1] = ptr;
            }
            index += 1;
        }
        index += 1;
    }

    // TODO: update right ptr after dealt with emptying nodes, cleanup empty nodes

    level_page_altered
}

pub fn level_merge<KS, PS>(src_tree: &BPlusTree<KS, PS>, dest_tree: &LevelTree) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let mut left_most_leaf_guards = select::<KS, PS>(&src_tree.get_root());
    let merge_page_len = left_most_leaf_guards.len();
    let mut num_keys_removed = 0;
    let left_most_id = left_most_leaf_guards.first().unwrap().ext_id();
    let prune_bound = left_most_leaf_guards.last().unwrap().right_bound().clone();

    debug!("Merge selected {} pages", left_most_leaf_guards.len());
    if cfg!(debug_assertions) {
        if left_most_id != src_tree.head_page_id {
            dump_tree(src_tree, "level_lsm_merge_failure_dump.json");
        }
    }
    debug_assert_eq!(
        left_most_id,
        src_tree.head_page_id,
        "{}",
        left_most_leaf_guards.first().unwrap().type_name()
    );

    // merge to dest_tree
    {
        let mut deleted_keys = src_tree.deleted.write();
        let mut merged_deleted_keys = vec![];
        let keys: Vec<EntryKey> = left_most_leaf_guards
            .iter()
            .filter(|&g| !g.is_empty_node())
            .map(|g| &g.keys()[..g.len()])
            .flatten()
            .filter(|&k| {
                if deleted_keys.contains(k) {
                    merged_deleted_keys.push(k.clone());
                    false
                } else {
                    true
                }
            })
            .cloned()
            .collect_vec();
        num_keys_removed = keys.len();
        debug!("Merge selected keys {:?}", keys.len());
        dest_tree.merge_with_keys(box keys);
        for rk in &merged_deleted_keys {
            deleted_keys.remove(rk);
        }
    }

    // adjust leaf left, right references
    {
        let right_right_most = left_most_leaf_guards
            .last()
            .unwrap()
            .right_ref()
            .unwrap()
            .clone();

        let left_left_most = left_most_leaf_guards
            .first()
            .unwrap()
            .left_ref()
            .unwrap()
            .clone();

        debug_assert!(read_unchecked::<KS, PS>(&left_left_most).is_none());
        debug_assert!(read_unchecked::<KS, PS>(&right_right_most).is_ext());
        debug_assert!(!right_right_most.ptr_eq(left_most_leaf_guards.last().unwrap().node_ref()));

        for mut g in &mut left_most_leaf_guards {
            external::make_deleted(&g.ext_id());
            debug!("make deleted {:?}", g.ext_id());
            **g = NodeData::Empty(box EmptyNode {
                left: Some(left_left_most.clone()),
                right: right_right_most.clone(),
            });
        }

        let mut new_first_node = write_node::<KS, PS>(&right_right_most);
        let mut new_first_node_ext = new_first_node.extnode_mut();
        debug!(
            "Right most original id is {:?}, now is {:?}",
            new_first_node_ext.id, src_tree.head_page_id
        );
        debug!(
            "New first node right is {:?}",
            read_unchecked::<KS, PS>(&new_first_node_ext.next).ext_id()
        );
        new_first_node_ext.id = src_tree.head_page_id;
        new_first_node_ext.prev = left_left_most;

        debug_assert!(&new_first_node_ext.right_bound > &prune_bound);

        ExtNode::<KS, PS>::make_changed(&right_right_most, src_tree);
    }

    // cleanup upper level references
    prune_selected::<KS, PS>(&src_tree.get_root(), &prune_bound, 0);

    src_tree.len.fetch_sub(num_keys_removed, Relaxed);

    debug!("Merge completed");

    merge_page_len
}
<commit_msg>final touches<commit_after>use core::borrow::{Borrow, BorrowMut};
use index::btree::dump::dump_tree;
use index::btree::external::ExtNode;
use index::btree::internal::InNode;
use index::btree::node::is_node_locked;
use index::btree::node::read_node;
use index::btree::node::read_unchecked;
use index::btree::node::write_node;
use index::btree::node::write_non_empty;
use index::btree::node::write_targeted;
use index::btree::node::EmptyNode;
use index::btree::node::Node;
use index::btree::node::NodeData;
use index::btree::node::NodeData::Empty;
use index::btree::node::NodeWriteGuard;
use index::btree::search::mut_search;
use index::btree::search::MutSearchResult;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
use index::btree::{external, BPlusTree};
use index::lsmtree::tree::{LEVEL_2, LEVEL_3, LEVEL_PAGE_DIFF_MULTIPLIER};
use index::EntryKey;
use index::Slice;
use itertools::free::all;
use itertools::Itertools;
use serde_json::to_vec;
use smallvec::SmallVec;
use std::cell::RefCell;
use std::cmp::min;
use std::collections::{BTreeSet, HashMap, HashSet};
use std::fmt::Debug;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::Arc;

enum Selection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    Selected(Vec<NodeWriteGuard<KS, PS>>),
    Innode(NodeCellRef),
}

struct KeyAltered {
    new_key: EntryKey,
    ptr: NodeCellRef,
}

fn select<KS, PS>(node: &NodeCellRef) -> Vec<NodeWriteGuard<KS, PS>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let search = mut_search::<KS, PS>(node, &smallvec!());
    match search {
        MutSearchResult::External => {
            let first_node = write_node(node);
            let mut collected_keys = first_node.len();
            let mut collected = vec![first_node];
            let target_keys = min(KS::slice_len() * LEVEL_PAGE_DIFF_MULTIPLIER, LEVEL_3);
            let target_guards = LEVEL_2;
            while collected_keys < target_keys && collected.len() < target_guards {
                let right = write_node(
                    collected
                        .last_mut()
                        .unwrap()
                        .right_ref_mut_no_empty()
                        .unwrap(),
                );
                if right.is_none() {
                    break;
                } else {
                    collected_keys += right.len();
                    collected.push(right);
                }
            }
            return collected;
        }
        MutSearchResult::Internal(node) => select::<KS, PS>(&node),
    }
}

fn prune_selected<'a, KS, PS>(node: &NodeCellRef, bound: &EntryKey, level: usize) -> Vec<KeyAltered>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let mut level_page_altered = vec![];
    let first_search = mut_search::<KS, PS>(node, &smallvec!());
    let altered_keys = match first_search {
        MutSearchResult::Internal(sub_node_ref) => {
            let sub_node = read_unchecked::<KS, PS>(&sub_node_ref);
            if !sub_node.is_empty_node() && !sub_node.is_ext() {
                prune_selected::<KS, PS>(&sub_node_ref, bound, level + 1)
            } else {
                vec![]
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    let mut all_pages = vec![write_node::<KS, PS>(node)];
    // collect all pages in bound and in this level
    loop {
        let (right_ref, right_bound) = {
            let last_page = all_pages.last().unwrap().borrow();
            if last_page.is_none() {
                break;
            }
            (
                last_page.right_ref().unwrap().clone(),
                last_page.right_bound().clone(),
            )
        };
        let right = write_node::<KS, PS>(&right_ref);
        if right.is_none() {
            break;
        }
        all_pages.push(right);
        if &right_bound > bound {
            break;
        }
    }
    let select_live = |page: &NodeWriteGuard<KS, PS>| {
        page.innode().ptrs.as_slice_immute()[..page.len() + 1]
            .iter()
            .enumerate()
            .filter_map(|(i, sub_level)| {
                if !read_unchecked::<KS, PS>(sub_level).is_empty_node() {
                    Some((i, sub_level.clone()))
                } else {
                    None
                }
            })
            .collect_vec()
    };
    let page_lives = all_pages.iter().map(|p| select_live(p)).collect_vec();
    all_pages
        .iter_mut()
        .zip(page_lives)
        .for_each(|(mut page, live_ptrs)| {
            if live_ptrs.len() == 0 {
                **page = NodeData::Empty(box EmptyNode {
                    left: None,
                    right: Default::default(),
                })
            } else {
                let mut new_keys = KS::init();
                let mut new_ptrs = PS::init();
                let ptr_len = live_ptrs.len();
                for (i, &(oi, _)) in live_ptrs.iter().take(live_ptrs.len() - 1).enumerate() {
                    new_keys.as_slice()[i] = page.keys()[oi].clone();
                }
                for (i, (_, ptr)) in live_ptrs.into_iter().enumerate() {
                    new_ptrs.as_slice()[i] = ptr;
                }
                let innode = page.innode_mut();
                innode.len = ptr_len - 1;
                innode.keys = new_keys;
                innode.ptrs = new_ptrs;
            }
        });

    let non_empty_right_node = |i: usize, pages: &Vec<NodeWriteGuard<KS, PS>>| {
        let mut i = i + 1;
        loop {
            if i == pages.len() {
                return pages[i - 1].borrow().right_ref().unwrap().clone();
            } else if i < pages.len() {
                let p = pages[i].borrow();
                if !p.is_empty_node() {
                    return p.node_ref().clone();
                } else {
                    i += 1;
                }
            } else {
                unreachable!()
            }
        }
    };

    let update_and_mark_altered_keys =
        |page: &mut NodeWriteGuard<KS, PS>, altered: &mut Vec<KeyAltered>| {
            let mut innode = page.innode_mut();
            let marked_ptrs = innode
                .ptrs
                .as_slice_immute()
                .iter()
                .enumerate()
                .filter_map(|(i, p)| {
                    altered_keys
                        .iter()
                        .find(|ak| ak.ptr.ptr_eq(p))
                        .map(|ak| (i, ak.new_key.clone()))
                })
                .collect_vec();
            for (i, new_key) in marked_ptrs {
                if i == 0 {
                    // cannot update the key in current page
                    // will postpone to upper level
                    altered.push(KeyAltered {
                        new_key: new_key,
                        ptr: innode.ptrs.as_slice_immute()[i].clone(),
                    });
                } else {
                    innode.keys.as_slice()[i - 1] = new_key;
                }
            }
        };

    let update_right_nodes = |all_pages: &mut Vec<NodeWriteGuard<KS, PS>>| {
        let right_ptrs = all_pages
            .iter()
            .enumerate()
            .map(|(i, _)| non_empty_right_node(i, &all_pages))
            .collect_vec();
        all_pages
            .iter_mut()
            .zip(right_ptrs.into_iter())
            .for_each(|(p, r)| *p.right_ref_mut().unwrap() = r);
    };

    update_right_nodes(&mut all_pages);

    all_pages = all_pages
        .into_iter()
        .filter(|p| !p.borrow().is_empty_node())
        .map(|mut p| {
            update_and_mark_altered_keys(&mut p, &mut level_page_altered);
            p
        })
        .collect_vec();

    let mut index = 0;
    while index < all_pages.len() {
        if all_pages[index].len() == 0 {
            // current page have one ptr and no keys
            // need to merge with the right page
            // if the right page is full, partial of the right page will be moved to the current page
            // merging right page will also been cleaned
            let mut next_from_ptr = if index + 1 >= all_pages.len() {
                Some(write_node::<KS, PS>(all_pages[index].right_ref().unwrap()))
            } else {
                None
            };
            let (keys, ptrs, right_bound, right_ref) = {
                let mut next = next_from_ptr
                    .as_mut()
                    .unwrap_or_else(|| &mut all_pages[index + 1]);
                if next.len() < KS::slice_len() - 1 {
                    // Merge next node with current node
                    let tuple = {
                        let next_innode = next.innode();
                        let len = next_innode.len;
                        (
                            next_innode.keys.as_slice_immute()[..len].to_vec(),
                            next_innode.ptrs.as_slice_immute()[..len + 1].to_vec(),
                            next_innode.right_bound.clone(),
                            next_innode.right.clone(),
                        )
                    };
                    **next = NodeData::Empty(box EmptyNode {
                        left: None,
                        right: NodeCellRef::default(),
                    });
                    tuple
                } else {
                    // Rebalance next node with current node
                    // Next node left bound need to be updated in upper level
                    let right_ref = next.node_ref().clone();
                    let next_innode = next.innode_mut();
                    let next_len = next_innode.len;
                    let next_mid = next_len / 2;
                    let right_left_bound = next_innode.keys.as_slice_immute()[next_mid].clone();
                    let tuple = (
                        next_innode.keys.as_slice_immute()[..next_mid].to_vec(),
                        next_innode.ptrs.as_slice_immute()[..next_mid + 1].to_vec(),
                        right_left_bound.clone(),
                        right_ref.clone(),
                    );
                    let mut keys = KS::init();
                    let mut ptrs = PS::init();

                    for (i, key) in next_innode.keys.as_slice_immute()[next_mid + 1..next_len]
                        .iter()
                        .enumerate()
                    {
                        keys.as_slice()[i] = key.clone();
                    }
                    for (i, ptr) in next_innode.ptrs.as_slice_immute()[next_mid + 1..next_len + 1]
                        .iter()
                        .enumerate()
                    {
                        ptrs.as_slice()[i] = ptr.clone();
                    }
                    next_innode.keys = keys;
                    next_innode.ptrs = ptrs;
                    next_innode.len = next_len - next_mid - 1;

                    level_page_altered.push(KeyAltered {
                        new_key: right_left_bound,
                        ptr: right_ref,
                    });
                    tuple
                }
            };
            let page_innode = all_pages[index].innode_mut();
            page_innode.keys.as_slice()[0] = page_innode.right_bound.clone();
            page_innode.right_bound = right_bound;
            page_innode.right = right_ref;
            page_innode.len = 1 + keys.len();
            for (i, key) in keys.into_iter().enumerate() {
                page_innode.keys.as_slice()[i + 1] = key;
            }
            for (i, ptr) in ptrs.into_iter().enumerate() {
                page_innode.ptrs.as_slice()[i + 1] = ptr;
            }
            index += 1;
        }
        index += 1;
    }

    update_right_nodes(&mut all_pages);

    level_page_altered
}

pub fn level_merge<KS, PS>(src_tree: &BPlusTree<KS, PS>, dest_tree: &LevelTree) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let mut left_most_leaf_guards = select::<KS, PS>(&src_tree.get_root());
    let merge_page_len = left_most_leaf_guards.len();
    let mut num_keys_removed = 0;
    let left_most_id = left_most_leaf_guards.first().unwrap().ext_id();
    let prune_bound = left_most_leaf_guards.last().unwrap().right_bound().clone();

    debug!("Merge selected {} pages", left_most_leaf_guards.len());
    if cfg!(debug_assertions) {
        if left_most_id != src_tree.head_page_id {
            dump_tree(src_tree, "level_lsm_merge_failure_dump.json");
        }
    }
    debug_assert_eq!(
        left_most_id,
        src_tree.head_page_id,
        "{}",
        left_most_leaf_guards.first().unwrap().type_name()
    );

    // merge to dest_tree
    {
        let mut deleted_keys = src_tree.deleted.write();
        let mut merged_deleted_keys = vec![];
        let keys: Vec<EntryKey> = left_most_leaf_guards
            .iter()
            .filter(|&g| !g.is_empty_node())
            .map(|g| &g.keys()[..g.len()])
            .flatten()
            .filter(|&k| {
                if deleted_keys.contains(k) {
                    merged_deleted_keys.push(k.clone());
                    false
                } else {
                    true
                }
            })
            .cloned()
            .collect_vec();
        num_keys_removed = keys.len();
        debug!("Merge selected keys {:?}", keys.len());
        dest_tree.merge_with_keys(box keys);
        for rk in &merged_deleted_keys {
            deleted_keys.remove(rk);
        }
    }

    // adjust leaf left, right references
    {
        let right_right_most = left_most_leaf_guards
            .last()
            .unwrap()
            .right_ref()
            .unwrap()
            .clone();

        let left_left_most = left_most_leaf_guards
            .first()
            .unwrap()
            .left_ref()
            .unwrap()
            .clone();

        debug_assert!(read_unchecked::<KS, PS>(&left_left_most).is_none());
        debug_assert!(read_unchecked::<KS, PS>(&right_right_most).is_ext());
        debug_assert!(!right_right_most.ptr_eq(left_most_leaf_guards.last().unwrap().node_ref()));

        for mut g in &mut left_most_leaf_guards {
            external::make_deleted(&g.ext_id());
            debug!("make deleted {:?}", g.ext_id());
            **g = NodeData::Empty(box EmptyNode {
                left: Some(left_left_most.clone()),
                right: right_right_most.clone(),
            });
        }

        let mut new_first_node = write_node::<KS, PS>(&right_right_most);
        let mut new_first_node_ext = new_first_node.extnode_mut();
        debug!(
            "Right most original id is {:?}, now is {:?}",
            new_first_node_ext.id, src_tree.head_page_id
        );
        debug!(
            "New first node right is {:?}",
            read_unchecked::<KS, PS>(&new_first_node_ext.next).ext_id()
        );
        new_first_node_ext.id = src_tree.head_page_id;
        new_first_node_ext.prev = left_left_most;

        debug_assert!(&new_first_node_ext.right_bound > &prune_bound);

        ExtNode::<KS, PS>::make_changed(&right_right_most, src_tree);
    }

    // cleanup upper level references
    prune_selected::<KS, PS>(&src_tree.get_root(), &prune_bound, 0);

    src_tree.len.fetch_sub(num_keys_removed, Relaxed);

    debug!("Merge completed");

    merge_page_len
}
<|endoftext|>