<commit_before>use std::cell::{Cell, RefCell};
use std::ops::Deref;
use std::time::Instant;

use cairo::{self, MatrixTrait};
use owning_ref::RcRef;

use crate::attributes::Attribute;
use crate::bbox::BoundingBox;
use crate::coord_units::CoordUnits;
use crate::drawing_ctx::DrawingCtx;
use crate::error::{RenderingError, ValueErrorKind};
use crate::length::{LengthHorizontal, LengthUnit, LengthVertical};
use crate::node::{NodeResult, NodeTrait, NodeType, RsvgNode};
use crate::parsers::{ParseError, ParseValue};
use crate::properties::{ColorInterpolationFilters, ComputedValues};
use crate::property_bag::PropertyBag;
use crate::surface_utils::shared_surface::{SharedImageSurface, SurfaceType};

mod bounds;
use self::bounds::BoundsBuilder;

pub mod context;
use self::context::{FilterContext, FilterInput, FilterResult};

mod error;
use self::error::FilterError;

mod input;
use self::input::Input;

pub mod node;
use self::node::NodeFilter;

pub mod blend;
pub mod color_matrix;
pub mod component_transfer;
pub mod composite;
pub mod convolve_matrix;
pub mod displacement_map;
pub mod flood;
pub mod gaussian_blur;
pub mod image;
pub mod light;
pub mod merge;
pub mod morphology;
pub mod offset;
pub mod tile;
pub mod turbulence;

/// A filter primitive interface.
trait Filter: NodeTrait {
    /// Renders this filter primitive.
    ///
    /// If this filter primitive can't be rendered for whatever reason (for instance, a required
    /// property hasn't been provided), an error is returned.
    fn render(
        &self,
        node: &RsvgNode,
        ctx: &FilterContext,
        draw_ctx: &mut DrawingCtx,
    ) -> Result<FilterResult, FilterError>;

    /// Returns `true` if this filter primitive is affected by the `color-interpolation-filters`
    /// property.
    ///
    /// Primitives that do color blending (like `feComposite` or `feBlend`) should return `true`
    /// here, whereas primitives that don't (like `feOffset`) should return `false`.
    fn is_affected_by_color_interpolation_filters(&self) -> bool;
}

/// The base filter primitive node containing common properties.
struct Primitive {
    x: Cell<Option<LengthHorizontal>>,
    y: Cell<Option<LengthVertical>>,
    width: Cell<Option<LengthHorizontal>>,
    height: Cell<Option<LengthVertical>>,
    result: RefCell<Option<String>>,
}

/// The base node for filter primitives which accept input.
struct PrimitiveWithInput {
    base: Primitive,
    in_: RefCell<Option<Input>>,
}

impl Primitive {
    /// Constructs a new `Primitive` with empty properties.
    #[inline]
    fn new<T: Filter>() -> Primitive {
        Primitive {
            x: Cell::new(None),
            y: Cell::new(None),
            width: Cell::new(None),
            height: Cell::new(None),
            result: RefCell::new(None),
        }
    }

    /// Returns the `BoundsBuilder` for bounds computation.
    #[inline]
    fn get_bounds<'a>(&self, ctx: &'a FilterContext) -> BoundsBuilder<'a> {
        BoundsBuilder::new(
            ctx,
            self.x.get(),
            self.y.get(),
            self.width.get(),
            self.height.get(),
        )
    }
}

impl NodeTrait for Primitive {
    fn set_atts(&self, node: &RsvgNode, pbag: &PropertyBag<'_>) -> NodeResult {
        // With ObjectBoundingBox, only fractions and percents are allowed.
        let primitiveunits = node
            .get_parent()
            .and_then(|parent| {
                if parent.get_type() == NodeType::Filter {
                    Some(parent.with_impl(|f: &NodeFilter| f.primitiveunits.get()))
                } else {
                    None
                }
            })
            .unwrap_or(CoordUnits::UserSpaceOnUse);

        let no_units_allowed = primitiveunits == CoordUnits::ObjectBoundingBox;

        let check_units_horizontal = |length: LengthHorizontal| {
            if !no_units_allowed {
                return Ok(length);
            }

            match length.unit() {
                LengthUnit::Px | LengthUnit::Percent => Ok(length),
                _ => Err(ValueErrorKind::Parse(ParseError::new(
                    "unit identifiers are not allowed with primitiveUnits set to objectBoundingBox",
                ))),
            }
        };

        let check_units_vertical = |length: LengthVertical| {
            if !no_units_allowed {
                return Ok(length);
            }

            match length.unit() {
                LengthUnit::Px | LengthUnit::Percent => Ok(length),
                _ => Err(ValueErrorKind::Parse(ParseError::new(
                    "unit identifiers are not allowed with primitiveUnits set to objectBoundingBox",
                ))),
            }
        };

        let check_units_horizontal_and_ensure_nonnegative = |length: LengthHorizontal| {
            check_units_horizontal(length).and_then(LengthHorizontal::check_nonnegative)
        };

        let check_units_vertical_and_ensure_nonnegative = |length: LengthVertical| {
            check_units_vertical(length).and_then(LengthVertical::check_nonnegative)
        };

        for (attr, value) in pbag.iter() {
            match attr {
                Attribute::X => self.x.set(Some(
                    attr.parse_and_validate(value, check_units_horizontal)?,
                )),
                Attribute::Y => self
                    .y
                    .set(Some(attr.parse_and_validate(value, check_units_vertical)?)),
                Attribute::Width => self.width.set(Some(
                    attr.parse_and_validate(value, check_units_horizontal_and_ensure_nonnegative)?,
                )),
                Attribute::Height => self.height.set(Some(
                    attr.parse_and_validate(value, check_units_vertical_and_ensure_nonnegative)?,
                )),
                Attribute::Result => *self.result.borrow_mut() = Some(value.to_string()),
                _ => (),
            }
        }

        Ok(())
    }
}

impl PrimitiveWithInput {
    /// Constructs a new `PrimitiveWithInput` with empty properties.
    #[inline]
    fn new<T: Filter>() -> PrimitiveWithInput {
        PrimitiveWithInput {
            base: Primitive::new::<T>(),
            in_: RefCell::new(None),
        }
    }

    /// Returns the input Cairo surface for this filter primitive.
    #[inline]
    fn get_input(
        &self,
        ctx: &FilterContext,
        draw_ctx: &mut DrawingCtx,
    ) -> Result<FilterInput, FilterError> {
        ctx.get_input(draw_ctx, self.in_.borrow().as_ref())
    }
}

impl NodeTrait for PrimitiveWithInput {
    fn set_atts(&self, node: &RsvgNode, pbag: &PropertyBag<'_>) -> NodeResult {
        self.base.set_atts(node, pbag)?;

        for (attr, value) in pbag.iter() {
            match attr {
                Attribute::In => drop(self.in_.replace(Some(Input::parse(Attribute::In, value)?))),
                _ => (),
            }
        }

        Ok(())
    }
}

impl Deref for PrimitiveWithInput {
    type Target = Primitive;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.base
    }
}

/// Applies a filter and returns the resulting surface.
pub fn render(
    filter_node: &RsvgNode,
    computed_from_node_being_filtered: &ComputedValues,
    source: &cairo::ImageSurface,
    draw_ctx: &mut DrawingCtx,
    node_bbox: BoundingBox,
) -> Result<cairo::ImageSurface, RenderingError> {
    let filter_node = &*filter_node;
    assert_eq!(filter_node.get_type(), NodeType::Filter);
    assert!(!filter_node.is_in_error());

    // The source surface has multiple references. We need to copy it to a new surface to have a
    // unique reference to be able to safely access the pixel data.
    let source_surface = cairo::ImageSurface::create(
        cairo::Format::ARgb32,
        source.get_width(),
        source.get_height(),
    )?;
    {
        let cr = cairo::Context::new(&source_surface);
        cr.set_source_surface(source, 0f64, 0f64);
        cr.paint();
    }
    let source_surface = SharedImageSurface::new(source_surface, SurfaceType::SRgb)?;

    let mut filter_ctx = FilterContext::new(
        filter_node,
        computed_from_node_being_filtered,
        source_surface,
        draw_ctx,
        node_bbox,
    );

    // If paffine is non-invertible, we won't draw anything. Also bbox combining in bounds
    // computations will panic due to non-invertible martrix.
    if filter_ctx.paffine().try_invert().is_err() {
        return Ok(filter_ctx.into_output()?.into_image_surface()?);
    }

    let primitives = filter_node
        .children()
        // Skip nodes in error.
        .filter(|c| {
            let in_error = c.is_in_error();

            if in_error {
                rsvg_log!(
                    "(ignoring filter primitive {} because it is in error)",
                    c.get_human_readable_name()
                );
            }

            !in_error
        })
        // Check if the node wants linear RGB.
        .map(|c| {
            let linear_rgb = {
                let cascaded = c.get_cascaded_values();
                let values = cascaded.get();

                values.color_interpolation_filters == ColorInterpolationFilters::LinearRgb
            };

            (c, linear_rgb)
        })
        .filter_map(|(c, linear_rgb)| {
            let rr = RcRef::new(c)
                .try_map(|c| {
                    // Go through the filter primitives and see if the node is one of them.
                    #[inline]
                    fn as_filter<T: Filter>(x: &T) -> &Filter {
                        x
                    }

                    // Unfortunately it's not possible to downcast to a trait object. If we could
                    // attach arbitrary data to nodes it would really help here. Previously that
                    // arbitrary data was in form of RsvgCNodeImpl, but that was heavily tuned for
                    // storing C pointers with all subsequent downsides.
                    macro_rules! try_downcasting_to_filter {
                        ($c:expr; $($t:ty),+$(,)*) => ({
                            let mut filter = None;
                            $(
                                filter = filter.or_else(|| $c.get_impl::<$t>().map(as_filter));
                            )+
                            filter
                        })
                    }

                    let filter = try_downcasting_to_filter!(c;
                        blend::Blend,
                        color_matrix::ColorMatrix,
                        component_transfer::ComponentTransfer,
                        composite::Composite,
                        convolve_matrix::ConvolveMatrix,
                        displacement_map::DisplacementMap,
                        flood::Flood,
                        gaussian_blur::GaussianBlur,
                        image::Image,
                        light::lighting::Lighting,
                        merge::Merge,
                        morphology::Morphology,
                        offset::Offset,
                        tile::Tile,
                        turbulence::Turbulence,
                    );
                    filter.ok_or(())
                })
                .ok();

            rr.map(|rr| (rr, linear_rgb))
        });

    for (rr, linear_rgb) in primitives {
        let mut render = |filter_ctx: &mut FilterContext| {
            if let Err(err) = rr
                .render(rr.as_owner(), filter_ctx, draw_ctx)
                .and_then(|result| filter_ctx.store_result(result))
            {
                rsvg_log!(
                    "(filter primitive {} returned an error: {})",
                    rr.as_owner().get_human_readable_name(),
                    err
                );

                // Exit early on Cairo errors. Continue rendering otherwise.
                if let FilterError::CairoError(status) = err {
                    return Err(status);
                }
            }

            Ok(())
        };

        let start = Instant::now();

        if rr.is_affected_by_color_interpolation_filters() && linear_rgb {
            filter_ctx.with_linear_rgb(render)?;
        } else {
            render(&mut filter_ctx)?;
        }

        let elapsed = start.elapsed();
        rsvg_log!(
            "(rendered filter primitive {} in\n    {} seconds)",
            rr.as_owner().get_human_readable_name(),
            elapsed.as_secs() as f64 + f64::from(elapsed.subsec_nanos()) / 1e9
        );
    }

    Ok(filter_ctx.into_output()?.into_image_surface()?)
}
<commit_msg>Extract function to copy a surface to a SharedImageSurface<commit_after>use std::cell::{Cell, RefCell};
use std::ops::Deref;
use std::time::Instant;

use cairo::{self, MatrixTrait};
use owning_ref::RcRef;

use crate::attributes::Attribute;
use crate::bbox::BoundingBox;
use crate::coord_units::CoordUnits;
use crate::drawing_ctx::DrawingCtx;
use crate::error::{RenderingError, ValueErrorKind};
use crate::length::{LengthHorizontal, LengthUnit, LengthVertical};
use crate::node::{NodeResult, NodeTrait, NodeType, RsvgNode};
use crate::parsers::{ParseError, ParseValue};
use crate::properties::{ColorInterpolationFilters, ComputedValues};
use crate::property_bag::PropertyBag;
use crate::surface_utils::shared_surface::{SharedImageSurface, SurfaceType};

mod bounds;
use self::bounds::BoundsBuilder;

pub mod context;
use self::context::{FilterContext, FilterInput, FilterResult};

mod error;
use self::error::FilterError;

mod input;
use self::input::Input;

pub mod node;
use self::node::NodeFilter;

pub mod blend;
pub mod color_matrix;
pub mod component_transfer;
pub mod composite;
pub mod convolve_matrix;
pub mod displacement_map;
pub mod flood;
pub mod gaussian_blur;
pub mod image;
pub mod light;
pub mod merge;
pub mod morphology;
pub mod offset;
pub mod tile;
pub mod turbulence;

/// A filter primitive interface.
trait Filter: NodeTrait {
    /// Renders this filter primitive.
    ///
    /// If this filter primitive can't be rendered for whatever reason (for instance, a required
    /// property hasn't been provided), an error is returned.
    fn render(
        &self,
        node: &RsvgNode,
        ctx: &FilterContext,
        draw_ctx: &mut DrawingCtx,
    ) -> Result<FilterResult, FilterError>;

    /// Returns `true` if this filter primitive is affected by the `color-interpolation-filters`
    /// property.
    ///
    /// Primitives that do color blending (like `feComposite` or `feBlend`) should return `true`
    /// here, whereas primitives that don't (like `feOffset`) should return `false`.
    fn is_affected_by_color_interpolation_filters(&self) -> bool;
}

/// The base filter primitive node containing common properties.
struct Primitive {
    x: Cell<Option<LengthHorizontal>>,
    y: Cell<Option<LengthVertical>>,
    width: Cell<Option<LengthHorizontal>>,
    height: Cell<Option<LengthVertical>>,
    result: RefCell<Option<String>>,
}

/// The base node for filter primitives which accept input.
struct PrimitiveWithInput {
    base: Primitive,
    in_: RefCell<Option<Input>>,
}

impl Primitive {
    /// Constructs a new `Primitive` with empty properties.
    #[inline]
    fn new<T: Filter>() -> Primitive {
        Primitive {
            x: Cell::new(None),
            y: Cell::new(None),
            width: Cell::new(None),
            height: Cell::new(None),
            result: RefCell::new(None),
        }
    }

    /// Returns the `BoundsBuilder` for bounds computation.
    #[inline]
    fn get_bounds<'a>(&self, ctx: &'a FilterContext) -> BoundsBuilder<'a> {
        BoundsBuilder::new(
            ctx,
            self.x.get(),
            self.y.get(),
            self.width.get(),
            self.height.get(),
        )
    }
}

impl NodeTrait for Primitive {
    fn set_atts(&self, node: &RsvgNode, pbag: &PropertyBag<'_>) -> NodeResult {
        // With ObjectBoundingBox, only fractions and percents are allowed.
        let primitiveunits = node
            .get_parent()
            .and_then(|parent| {
                if parent.get_type() == NodeType::Filter {
                    Some(parent.with_impl(|f: &NodeFilter| f.primitiveunits.get()))
                } else {
                    None
                }
            })
            .unwrap_or(CoordUnits::UserSpaceOnUse);

        let no_units_allowed = primitiveunits == CoordUnits::ObjectBoundingBox;

        let check_units_horizontal = |length: LengthHorizontal| {
            if !no_units_allowed {
                return Ok(length);
            }

            match length.unit() {
                LengthUnit::Px | LengthUnit::Percent => Ok(length),
                _ => Err(ValueErrorKind::Parse(ParseError::new(
                    "unit identifiers are not allowed with primitiveUnits set to objectBoundingBox",
                ))),
            }
        };

        let check_units_vertical = |length: LengthVertical| {
            if !no_units_allowed {
                return Ok(length);
            }

            match length.unit() {
                LengthUnit::Px | LengthUnit::Percent => Ok(length),
                _ => Err(ValueErrorKind::Parse(ParseError::new(
                    "unit identifiers are not allowed with primitiveUnits set to objectBoundingBox",
                ))),
            }
        };

        let check_units_horizontal_and_ensure_nonnegative = |length: LengthHorizontal| {
            check_units_horizontal(length).and_then(LengthHorizontal::check_nonnegative)
        };

        let check_units_vertical_and_ensure_nonnegative = |length: LengthVertical| {
            check_units_vertical(length).and_then(LengthVertical::check_nonnegative)
        };

        for (attr, value) in pbag.iter() {
            match attr {
                Attribute::X => self.x.set(Some(
                    attr.parse_and_validate(value, check_units_horizontal)?,
                )),
                Attribute::Y => self
                    .y
                    .set(Some(attr.parse_and_validate(value, check_units_vertical)?)),
                Attribute::Width => self.width.set(Some(
                    attr.parse_and_validate(value, check_units_horizontal_and_ensure_nonnegative)?,
                )),
                Attribute::Height => self.height.set(Some(
                    attr.parse_and_validate(value, check_units_vertical_and_ensure_nonnegative)?,
                )),
                Attribute::Result => *self.result.borrow_mut() = Some(value.to_string()),
                _ => (),
            }
        }

        Ok(())
    }
}

impl PrimitiveWithInput {
    /// Constructs a new `PrimitiveWithInput` with empty properties.
    #[inline]
    fn new<T: Filter>() -> PrimitiveWithInput {
        PrimitiveWithInput {
            base: Primitive::new::<T>(),
            in_: RefCell::new(None),
        }
    }

    /// Returns the input Cairo surface for this filter primitive.
    #[inline]
    fn get_input(
        &self,
        ctx: &FilterContext,
        draw_ctx: &mut DrawingCtx,
    ) -> Result<FilterInput, FilterError> {
        ctx.get_input(draw_ctx, self.in_.borrow().as_ref())
    }
}

impl NodeTrait for PrimitiveWithInput {
    fn set_atts(&self, node: &RsvgNode, pbag: &PropertyBag<'_>) -> NodeResult {
        self.base.set_atts(node, pbag)?;

        for (attr, value) in pbag.iter() {
            match attr {
                Attribute::In => drop(self.in_.replace(Some(Input::parse(Attribute::In, value)?))),
                _ => (),
            }
        }

        Ok(())
    }
}

impl Deref for PrimitiveWithInput {
    type Target = Primitive;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.base
    }
}

/// Creates a `SharedImageSurface` from an `ImageSurface`, even if the former
/// does not have a reference count of 1.
fn copy_to_shared_surface(
    surface: &cairo::ImageSurface,
) -> Result<SharedImageSurface, cairo::Status> {
    let copy = cairo::ImageSurface::create(
        cairo::Format::ARgb32,
        surface.get_width(),
        surface.get_height(),
    )?;
    {
        let cr = cairo::Context::new(&copy);
        cr.set_source_surface(surface, 0f64, 0f64);
        cr.paint();
    }
    SharedImageSurface::new(copy, SurfaceType::SRgb)
}

/// Applies a filter and returns the resulting surface.
pub fn render(
    filter_node: &RsvgNode,
    computed_from_node_being_filtered: &ComputedValues,
    source: &cairo::ImageSurface,
    draw_ctx: &mut DrawingCtx,
    node_bbox: BoundingBox,
) -> Result<cairo::ImageSurface, RenderingError> {
    let filter_node = &*filter_node;
    assert_eq!(filter_node.get_type(), NodeType::Filter);
    assert!(!filter_node.is_in_error());

    // The source surface has multiple references. We need to copy it to a new surface to have a
    // unique reference to be able to safely access the pixel data.
    let source_surface = copy_to_shared_surface(source)?;

    let mut filter_ctx = FilterContext::new(
        filter_node,
        computed_from_node_being_filtered,
        source_surface,
        draw_ctx,
        node_bbox,
    );

    // If paffine is non-invertible, we won't draw anything. Also bbox combining in bounds
    // computations will panic due to non-invertible martrix.
    if filter_ctx.paffine().try_invert().is_err() {
        return Ok(filter_ctx.into_output()?.into_image_surface()?);
    }

    let primitives = filter_node
        .children()
        // Skip nodes in error.
        .filter(|c| {
            let in_error = c.is_in_error();

            if in_error {
                rsvg_log!(
                    "(ignoring filter primitive {} because it is in error)",
                    c.get_human_readable_name()
                );
            }

            !in_error
        })
        // Check if the node wants linear RGB.
        .map(|c| {
            let linear_rgb = {
                let cascaded = c.get_cascaded_values();
                let values = cascaded.get();

                values.color_interpolation_filters == ColorInterpolationFilters::LinearRgb
            };

            (c, linear_rgb)
        })
        .filter_map(|(c, linear_rgb)| {
            let rr = RcRef::new(c)
                .try_map(|c| {
                    // Go through the filter primitives and see if the node is one of them.
                    #[inline]
                    fn as_filter<T: Filter>(x: &T) -> &Filter {
                        x
                    }

                    // Unfortunately it's not possible to downcast to a trait object. If we could
                    // attach arbitrary data to nodes it would really help here. Previously that
                    // arbitrary data was in form of RsvgCNodeImpl, but that was heavily tuned for
                    // storing C pointers with all subsequent downsides.
                    macro_rules! try_downcasting_to_filter {
                        ($c:expr; $($t:ty),+$(,)*) => ({
                            let mut filter = None;
                            $(
                                filter = filter.or_else(|| $c.get_impl::<$t>().map(as_filter));
                            )+
                            filter
                        })
                    }

                    let filter = try_downcasting_to_filter!(c;
                        blend::Blend,
                        color_matrix::ColorMatrix,
                        component_transfer::ComponentTransfer,
                        composite::Composite,
                        convolve_matrix::ConvolveMatrix,
                        displacement_map::DisplacementMap,
                        flood::Flood,
                        gaussian_blur::GaussianBlur,
                        image::Image,
                        light::lighting::Lighting,
                        merge::Merge,
                        morphology::Morphology,
                        offset::Offset,
                        tile::Tile,
                        turbulence::Turbulence,
                    );
                    filter.ok_or(())
                })
                .ok();

            rr.map(|rr| (rr, linear_rgb))
        });

    for (rr, linear_rgb) in primitives {
        let mut render = |filter_ctx: &mut FilterContext| {
            if let Err(err) = rr
                .render(rr.as_owner(), filter_ctx, draw_ctx)
                .and_then(|result| filter_ctx.store_result(result))
            {
                rsvg_log!(
                    "(filter primitive {} returned an error: {})",
                    rr.as_owner().get_human_readable_name(),
                    err
                );

                // Exit early on Cairo errors. Continue rendering otherwise.
                if let FilterError::CairoError(status) = err {
                    return Err(status);
                }
            }

            Ok(())
        };

        let start = Instant::now();

        if rr.is_affected_by_color_interpolation_filters() && linear_rgb {
            filter_ctx.with_linear_rgb(render)?;
        } else {
            render(&mut filter_ctx)?;
        }

        let elapsed = start.elapsed();
        rsvg_log!(
            "(rendered filter primitive {} in\n    {} seconds)",
            rr.as_owner().get_human_readable_name(),
            elapsed.as_secs() as f64 + f64::from(elapsed.subsec_nanos()) / 1e9
        );
    }

    Ok(filter_ctx.into_output()?.into_image_surface()?)
}
<|endoftext|>


<commit_before>#[derive(Debug)]
struct Point<T, U> {
    x: T,
    y: U
}

// <T, U> after impl only go with Point
impl<T, U> Point<T, U> {
    // V, W after mixup only to do with the method
    // if add T, will compile err: T be shadowed
    fn mixup<V, W, T>(self, other: Point<V, W>) -> Point<T, W> {
        Point {
            x: self.x,
            y: other.y
        }
    }
}

fn main() {
    let p1 = Point { x: 5, y: 10.4 };
    let p2 = Point { x: "Hello", y: 'c' };

    let p3 = p1.mixup(p2);

    println!("p3.x = {}, p3.y = {}", p3.x, p3.y);
}
<commit_msg>Fix generics.<commit_after>#[derive(Debug)]
struct Point<T, U> {
    x: T,
    y: U
}

// <T, U> after impl only go with Point
impl<T, U> Point<T, U> {
    // V, W after mixup only to do with the method
    // if add T, will compile err: T be shadowed
    fn mixup<V, W>(self, other: Point<V, W>) -> Point<T, W> {
        Point {
            x: self.x,
            y: other.y
        }
    }
}

fn main() {
    let p1 = Point { x: 5, y: 10.4 };
    let p2 = Point { x: "Hello", y: 'c' };

    let p3 = p1.mixup(p2);

    println!("p3.x = {}, p3.y = {}", p3.x, p3.y);
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

//! SQL Parser
//!
//! Declares a SQL parser based on sqlparser that handles custom formats that we need.

use sqlparser::{
    ast::{ColumnDef, Statement as SQLStatement, TableConstraint},
    dialect::{keywords::Keyword, GenericDialect},
    parser::{Parser, ParserError},
    tokenizer::{Token, Tokenizer},
};

// Use `Parser::expected` instead, if possible
macro_rules! parser_err {
    ($MSG:expr) => {
        Err(ParserError::ParserError($MSG.to_string()))
    };
}

/// Types of files to parse as DataFrames
#[derive(Debug, Clone, PartialEq)]
pub enum FileType {
    /// Newline-delimited JSON
    NdJson,
    /// Apache Parquet columnar storage
    Parquet,
    /// Comma separated values
    CSV,
}

/// DataFusion extension DDL for `CREATE EXTERNAL TABLE`
#[derive(Debug, Clone, PartialEq)]
pub struct CreateExternalTable {
    /// Table name
    pub name: String,
    /// Optional schema
    pub columns: Vec<ColumnDef>,
    /// File type (Parquet, NDJSON, CSV)
    pub file_type: FileType,
    /// CSV Header row?
    pub has_header: bool,
    /// Path to file
    pub location: String,
}

/// DataFusion extension DDL for `EXPLAIN` and `EXPLAIN VERBOSE`
#[derive(Debug, Clone, PartialEq)]
pub struct ExplainPlan {
    /// If true, dumps more intermediate plans and results of optimizaton passes
    pub verbose: bool,
    /// The statement for which to generate an planning explination
    pub statement: Box<Statement>,
}

/// DataFusion Statement representations.
///
/// Tokens parsed by `DFParser` are converted into these values.
#[derive(Debug, Clone, PartialEq)]
pub enum Statement {
    /// ANSI SQL AST node
    Statement(SQLStatement),
    /// Extension: `CREATE EXTERNAL TABLE`
    CreateExternalTable(CreateExternalTable),
    /// Extension: `EXPLAIN <SQL>`
    Explain(ExplainPlan),
}

/// SQL Parser
pub struct DFParser {
    parser: Parser,
}

impl DFParser {
    /// Parse the specified tokens
    pub fn new(sql: &str) -> Result<Self, ParserError> {
        let dialect = GenericDialect {};
        let mut tokenizer = Tokenizer::new(&dialect, sql);
        let tokens = tokenizer.tokenize()?;
        Ok(DFParser {
            parser: Parser::new(tokens),
        })
    }

    /// Parse a SQL statement and produce a set of statements
    pub fn parse_sql(sql: &str) -> Result<Vec<Statement>, ParserError> {
        let mut tokenizer = Tokenizer::new(&GenericDialect {}, &sql);
        tokenizer.tokenize()?;
        let mut parser = DFParser::new(sql)?;
        let mut stmts = Vec::new();
        let mut expecting_statement_delimiter = false;
        loop {
            // ignore empty statements (between successive statement delimiters)
            while parser.parser.consume_token(&Token::SemiColon) {
                expecting_statement_delimiter = false;
            }

            if parser.parser.peek_token() == Token::EOF {
                break;
            }
            if expecting_statement_delimiter {
                return parser.expected("end of statement", parser.parser.peek_token());
            }

            let statement = parser.parse_statement()?;
            stmts.push(statement);
            expecting_statement_delimiter = true;
        }
        Ok(stmts)
    }

    /// Report unexpected token
    fn expected<T>(&self, expected: &str, found: Token) -> Result<T, ParserError> {
        parser_err!(format!("Expected {}, found: {}", expected, found))
    }

    /// Parse a new expression
    pub fn parse_statement(&mut self) -> Result<Statement, ParserError> {
        match self.parser.peek_token() {
            Token::Word(w) => {
                match w.keyword {
                    Keyword::CREATE => {
                        // move one token forward
                        self.parser.next_token();
                        // use custom parsing
                        self.parse_create()
                    }
                    Keyword::NoKeyword if w.value.to_uppercase() == "EXPLAIN" => {
                        self.parser.next_token();
                        self.parse_explain()
                    }
                    _ => {
                        // use the native parser
                        Ok(Statement::Statement(self.parser.parse_statement()?))
                    }
                }
            }
            _ => {
                // use the native parser
                Ok(Statement::Statement(self.parser.parse_statement()?))
            }
        }
    }

    /// Parse a SQL CREATE statement
    pub fn parse_create(&mut self) -> Result<Statement, ParserError> {
        if self.parser.parse_keyword(Keyword::EXTERNAL) {
            self.parse_create_external_table()
        } else {
            Ok(Statement::Statement(self.parser.parse_create()?))
        }
    }

    /// Parse an SQL EXPLAIN statement.
    pub fn parse_explain(&mut self) -> Result<Statement, ParserError> {
        // Parser is at the token immediately after EXPLAIN
        // Check for EXPLAIN VERBOSE
        let verbose = match self.parser.peek_token() {
            Token::Word(w) => match w.keyword {
                Keyword::NoKeyword if w.value.to_uppercase() == "VERBOSE" => {
                    self.parser.next_token();
                    true
                }
                _ => false,
            },
            _ => false,
        };

        let statement = Box::new(self.parse_statement()?);
        let explain_plan = ExplainPlan { statement, verbose };
        Ok(Statement::Explain(explain_plan))
    }

    // This is a copy of the equivalent implementation in sqlparser.
    fn parse_columns(
        &mut self,
    ) -> Result<(Vec<ColumnDef>, Vec<TableConstraint>), ParserError> {
        let mut columns = vec![];
        let mut constraints = vec![];
        if !self.parser.consume_token(&Token::LParen)
            || self.parser.consume_token(&Token::RParen)
        {
            return Ok((columns, constraints));
        }

        loop {
            if let Some(constraint) = self.parser.parse_optional_table_constraint()? {
                constraints.push(constraint);
            } else if let Token::Word(_) = self.parser.peek_token() {
                let column_def = self.parse_column_def()?;
                columns.push(column_def);
            } else {
                return self.expected(
                    "column name or constraint definition",
                    self.parser.peek_token(),
                );
            }
            let comma = self.parser.consume_token(&Token::Comma);
            if self.parser.consume_token(&Token::RParen) {
                // allow a trailing comma, even though it's not in standard
                break;
            } else if !comma {
                return self.expected(
                    "',' or ')' after column definition",
                    self.parser.peek_token(),
                );
            }
        }

        Ok((columns, constraints))
    }

    fn parse_column_def(&mut self) -> Result<ColumnDef, ParserError> {
        let name = self.parser.parse_identifier()?;
        let data_type = self.parser.parse_data_type()?;
        let collation = if self.parser.parse_keyword(Keyword::COLLATE) {
            Some(self.parser.parse_object_name()?)
        } else {
            None
        };
        let mut options = vec![];
        loop {
            match self.parser.peek_token() {
                Token::EOF | Token::Comma | Token::RParen => break,
                _ => options.push(self.parser.parse_column_option_def()?),
            }
        }
        Ok(ColumnDef {
            name,
            data_type,
            collation,
            options,
        })
    }

    fn parse_create_external_table(&mut self) -> Result<Statement, ParserError> {
        self.parser.expect_keyword(Keyword::TABLE)?;
        let table_name = self.parser.parse_object_name()?;
        let (columns, _) = self.parse_columns()?;
        self.parser
            .expect_keywords(&[Keyword::STORED, Keyword::AS])?;

        // THIS is the main difference: we parse a different file format.
        let file_type = self.parse_file_format()?;

        let has_header = self.parse_csv_has_header();

        self.parser.expect_keyword(Keyword::LOCATION)?;
        let location = self.parser.parse_literal_string()?;

        let create = CreateExternalTable {
            name: table_name.to_string(),
            columns,
            file_type,
            has_header,
            location,
        };
        Ok(Statement::CreateExternalTable(create))
    }

    /// Parses the set of valid formats
    fn parse_file_format(&mut self) -> Result<FileType, ParserError> {
        match self.parser.next_token() {
            Token::Word(w) => match &*w.value {
                "PARQUET" => Ok(FileType::Parquet),
                "NDJSON" => Ok(FileType::NdJson),
                "CSV" => Ok(FileType::CSV),
                _ => self.expected("one of PARQUET, NDJSON, or CSV", Token::Word(w)),
            },
            unexpected => self.expected("one of PARQUET, NDJSON, or CSV", unexpected),
        }
    }

    fn consume_token(&mut self, expected: &str) -> bool {
        if self.parser.peek_token().to_string() == *expected {
            self.parser.next_token();
            true
        } else {
            false
        }
    }

    fn parse_csv_has_header(&mut self) -> bool {
        self.consume_token("WITH")
            & self.consume_token("HEADER")
            & self.consume_token("ROW")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use sqlparser::ast::{DataType, Ident};

    fn expect_parse_ok(sql: &str, expected: Statement) -> Result<(), ParserError> {
        let statements = DFParser::parse_sql(sql)?;
        assert_eq!(
            statements.len(),
            1,
            "Expected to parse exactly one statement"
        );
        assert_eq!(statements[0], expected);
        Ok(())
    }

    /// Parses sql and asserts that the expected error message was found
    fn expect_parse_error(sql: &str, expected_error: &str) -> Result<(), ParserError> {
        match DFParser::parse_sql(sql) {
            Ok(statements) => {
                assert!(
                    false,
                    "Expected parse error for '{}', but was successful: {:?}",
                    sql, statements
                );
            }
            Err(e) => {
                let error_message = e.to_string();
                assert!(
                    error_message.contains(expected_error),
                    "Expected error '{}' not found in actual error '{}'",
                    expected_error,
                    error_message
                );
            }
        }
        Ok(())
    }

    fn make_column_def(name: impl Into<String>, data_type: DataType) -> ColumnDef {
        ColumnDef {
            name: Ident {
                value: name.into(),
                quote_style: None,
            },
            data_type,
            collation: None,
            options: vec![],
        }
    }

    #[test]
    fn create_external_table() -> Result<(), ParserError> {
        // positive case
        let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV LOCATION 'foo.csv'";
        let expected = Statement::CreateExternalTable(CreateExternalTable {
            name: "t".into(),
            columns: vec![make_column_def("c1", DataType::Int)],
            file_type: FileType::CSV,
            has_header: false,
            location: "foo.csv".into(),
        });
        expect_parse_ok(sql, expected)?;

        // positive case: it is ok for parquet files not to have columns specified
        let sql = "CREATE EXTERNAL TABLE t STORED AS PARQUET LOCATION 'foo.parquet'";
        let expected = Statement::CreateExternalTable(CreateExternalTable {
            name: "t".into(),
            columns: vec![],
            file_type: FileType::Parquet,
            has_header: false,
            location: "foo.parquet".into(),
        });
        expect_parse_ok(sql, expected)?;

        // Error cases: Invalid type
        let sql =
            "CREATE EXTERNAL TABLE t(c1 int) STORED AS UNKNOWN_TYPE LOCATION 'foo.csv'";
        expect_parse_error(
            sql,
            "Expected one of PARQUET, NDJSON, or CSV, found: UNKNOWN_TYPE",
        )?;

        Ok(())
    }
}
<commit_msg>ARROW-9673: [Rust] [DataFusion] Add a param "dialect" for DFParser::parse_sql<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

//! SQL Parser
//!
//! Declares a SQL parser based on sqlparser that handles custom formats that we need.

use sqlparser::{
    ast::{ColumnDef, Statement as SQLStatement, TableConstraint},
    dialect::{keywords::Keyword, Dialect, GenericDialect},
    parser::{Parser, ParserError},
    tokenizer::{Token, Tokenizer},
};

// Use `Parser::expected` instead, if possible
macro_rules! parser_err {
    ($MSG:expr) => {
        Err(ParserError::ParserError($MSG.to_string()))
    };
}

/// Types of files to parse as DataFrames
#[derive(Debug, Clone, PartialEq)]
pub enum FileType {
    /// Newline-delimited JSON
    NdJson,
    /// Apache Parquet columnar storage
    Parquet,
    /// Comma separated values
    CSV,
}

/// DataFusion extension DDL for `CREATE EXTERNAL TABLE`
#[derive(Debug, Clone, PartialEq)]
pub struct CreateExternalTable {
    /// Table name
    pub name: String,
    /// Optional schema
    pub columns: Vec<ColumnDef>,
    /// File type (Parquet, NDJSON, CSV)
    pub file_type: FileType,
    /// CSV Header row?
    pub has_header: bool,
    /// Path to file
    pub location: String,
}

/// DataFusion extension DDL for `EXPLAIN` and `EXPLAIN VERBOSE`
#[derive(Debug, Clone, PartialEq)]
pub struct ExplainPlan {
    /// If true, dumps more intermediate plans and results of optimizaton passes
    pub verbose: bool,
    /// The statement for which to generate an planning explination
    pub statement: Box<Statement>,
}

/// DataFusion Statement representations.
///
/// Tokens parsed by `DFParser` are converted into these values.
#[derive(Debug, Clone, PartialEq)]
pub enum Statement {
    /// ANSI SQL AST node
    Statement(SQLStatement),
    /// Extension: `CREATE EXTERNAL TABLE`
    CreateExternalTable(CreateExternalTable),
    /// Extension: `EXPLAIN <SQL>`
    Explain(ExplainPlan),
}

/// SQL Parser
pub struct DFParser {
    parser: Parser,
}

impl DFParser {
    /// Parse the specified tokens
    pub fn new(sql: &str) -> Result<Self, ParserError> {
        let dialect = &GenericDialect {};
        DFParser::new_with_dialect(sql, dialect)
    }

    /// Parse the specified tokens with dialect
    pub fn new_with_dialect(
        sql: &str,
        dialect: &dyn Dialect,
    ) -> Result<Self, ParserError> {
        let mut tokenizer = Tokenizer::new(dialect, sql);
        let tokens = tokenizer.tokenize()?;
        Ok(DFParser {
            parser: Parser::new(tokens),
        })
    }

    /// Parse a SQL statement and produce a set of statements with dialect
    pub fn parse_sql(sql: &str) -> Result<Vec<Statement>, ParserError> {
        let dialect = &GenericDialect {};
        DFParser::parse_sql_with_dialect(sql, dialect)
    }

    /// Parse a SQL statement and produce a set of statements
    pub fn parse_sql_with_dialect(
        sql: &str,
        dialect: &dyn Dialect,
    ) -> Result<Vec<Statement>, ParserError> {
        let mut parser = DFParser::new_with_dialect(sql, dialect)?;
        let mut stmts = Vec::new();
        let mut expecting_statement_delimiter = false;
        loop {
            // ignore empty statements (between successive statement delimiters)
            while parser.parser.consume_token(&Token::SemiColon) {
                expecting_statement_delimiter = false;
            }

            if parser.parser.peek_token() == Token::EOF {
                break;
            }
            if expecting_statement_delimiter {
                return parser.expected("end of statement", parser.parser.peek_token());
            }

            let statement = parser.parse_statement()?;
            stmts.push(statement);
            expecting_statement_delimiter = true;
        }
        Ok(stmts)
    }

    /// Report unexpected token
    fn expected<T>(&self, expected: &str, found: Token) -> Result<T, ParserError> {
        parser_err!(format!("Expected {}, found: {}", expected, found))
    }

    /// Parse a new expression
    pub fn parse_statement(&mut self) -> Result<Statement, ParserError> {
        match self.parser.peek_token() {
            Token::Word(w) => {
                match w.keyword {
                    Keyword::CREATE => {
                        // move one token forward
                        self.parser.next_token();
                        // use custom parsing
                        self.parse_create()
                    }
                    Keyword::NoKeyword if w.value.to_uppercase() == "EXPLAIN" => {
                        self.parser.next_token();
                        self.parse_explain()
                    }
                    _ => {
                        // use the native parser
                        Ok(Statement::Statement(self.parser.parse_statement()?))
                    }
                }
            }
            _ => {
                // use the native parser
                Ok(Statement::Statement(self.parser.parse_statement()?))
            }
        }
    }

    /// Parse a SQL CREATE statement
    pub fn parse_create(&mut self) -> Result<Statement, ParserError> {
        if self.parser.parse_keyword(Keyword::EXTERNAL) {
            self.parse_create_external_table()
        } else {
            Ok(Statement::Statement(self.parser.parse_create()?))
        }
    }

    /// Parse an SQL EXPLAIN statement.
    pub fn parse_explain(&mut self) -> Result<Statement, ParserError> {
        // Parser is at the token immediately after EXPLAIN
        // Check for EXPLAIN VERBOSE
        let verbose = match self.parser.peek_token() {
            Token::Word(w) => match w.keyword {
                Keyword::NoKeyword if w.value.to_uppercase() == "VERBOSE" => {
                    self.parser.next_token();
                    true
                }
                _ => false,
            },
            _ => false,
        };

        let statement = Box::new(self.parse_statement()?);
        let explain_plan = ExplainPlan { statement, verbose };
        Ok(Statement::Explain(explain_plan))
    }

    // This is a copy of the equivalent implementation in sqlparser.
    fn parse_columns(
        &mut self,
    ) -> Result<(Vec<ColumnDef>, Vec<TableConstraint>), ParserError> {
        let mut columns = vec![];
        let mut constraints = vec![];
        if !self.parser.consume_token(&Token::LParen)
            || self.parser.consume_token(&Token::RParen)
        {
            return Ok((columns, constraints));
        }

        loop {
            if let Some(constraint) = self.parser.parse_optional_table_constraint()? {
                constraints.push(constraint);
            } else if let Token::Word(_) = self.parser.peek_token() {
                let column_def = self.parse_column_def()?;
                columns.push(column_def);
            } else {
                return self.expected(
                    "column name or constraint definition",
                    self.parser.peek_token(),
                );
            }
            let comma = self.parser.consume_token(&Token::Comma);
            if self.parser.consume_token(&Token::RParen) {
                // allow a trailing comma, even though it's not in standard
                break;
            } else if !comma {
                return self.expected(
                    "',' or ')' after column definition",
                    self.parser.peek_token(),
                );
            }
        }

        Ok((columns, constraints))
    }

    fn parse_column_def(&mut self) -> Result<ColumnDef, ParserError> {
        let name = self.parser.parse_identifier()?;
        let data_type = self.parser.parse_data_type()?;
        let collation = if self.parser.parse_keyword(Keyword::COLLATE) {
            Some(self.parser.parse_object_name()?)
        } else {
            None
        };
        let mut options = vec![];
        loop {
            match self.parser.peek_token() {
                Token::EOF | Token::Comma | Token::RParen => break,
                _ => options.push(self.parser.parse_column_option_def()?),
            }
        }
        Ok(ColumnDef {
            name,
            data_type,
            collation,
            options,
        })
    }

    fn parse_create_external_table(&mut self) -> Result<Statement, ParserError> {
        self.parser.expect_keyword(Keyword::TABLE)?;
        let table_name = self.parser.parse_object_name()?;
        let (columns, _) = self.parse_columns()?;
        self.parser
            .expect_keywords(&[Keyword::STORED, Keyword::AS])?;

        // THIS is the main difference: we parse a different file format.
        let file_type = self.parse_file_format()?;

        let has_header = self.parse_csv_has_header();

        self.parser.expect_keyword(Keyword::LOCATION)?;
        let location = self.parser.parse_literal_string()?;

        let create = CreateExternalTable {
            name: table_name.to_string(),
            columns,
            file_type,
            has_header,
            location,
        };
        Ok(Statement::CreateExternalTable(create))
    }

    /// Parses the set of valid formats
    fn parse_file_format(&mut self) -> Result<FileType, ParserError> {
        match self.parser.next_token() {
            Token::Word(w) => match &*w.value {
                "PARQUET" => Ok(FileType::Parquet),
                "NDJSON" => Ok(FileType::NdJson),
                "CSV" => Ok(FileType::CSV),
                _ => self.expected("one of PARQUET, NDJSON, or CSV", Token::Word(w)),
            },
            unexpected => self.expected("one of PARQUET, NDJSON, or CSV", unexpected),
        }
    }

    fn consume_token(&mut self, expected: &str) -> bool {
        if self.parser.peek_token().to_string() == *expected {
            self.parser.next_token();
            true
        } else {
            false
        }
    }

    fn parse_csv_has_header(&mut self) -> bool {
        self.consume_token("WITH")
            & self.consume_token("HEADER")
            & self.consume_token("ROW")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use sqlparser::ast::{DataType, Ident};

    fn expect_parse_ok(sql: &str, expected: Statement) -> Result<(), ParserError> {
        let statements = DFParser::parse_sql(sql)?;
        assert_eq!(
            statements.len(),
            1,
            "Expected to parse exactly one statement"
        );
        assert_eq!(statements[0], expected);
        Ok(())
    }

    /// Parses sql and asserts that the expected error message was found
    fn expect_parse_error(sql: &str, expected_error: &str) -> Result<(), ParserError> {
        match DFParser::parse_sql(sql) {
            Ok(statements) => {
                assert!(
                    false,
                    "Expected parse error for '{}', but was successful: {:?}",
                    sql, statements
                );
            }
            Err(e) => {
                let error_message = e.to_string();
                assert!(
                    error_message.contains(expected_error),
                    "Expected error '{}' not found in actual error '{}'",
                    expected_error,
                    error_message
                );
            }
        }
        Ok(())
    }

    fn make_column_def(name: impl Into<String>, data_type: DataType) -> ColumnDef {
        ColumnDef {
            name: Ident {
                value: name.into(),
                quote_style: None,
            },
            data_type,
            collation: None,
            options: vec![],
        }
    }

    #[test]
    fn create_external_table() -> Result<(), ParserError> {
        // positive case
        let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV LOCATION 'foo.csv'";
        let expected = Statement::CreateExternalTable(CreateExternalTable {
            name: "t".into(),
            columns: vec![make_column_def("c1", DataType::Int)],
            file_type: FileType::CSV,
            has_header: false,
            location: "foo.csv".into(),
        });
        expect_parse_ok(sql, expected)?;

        // positive case: it is ok for parquet files not to have columns specified
        let sql = "CREATE EXTERNAL TABLE t STORED AS PARQUET LOCATION 'foo.parquet'";
        let expected = Statement::CreateExternalTable(CreateExternalTable {
            name: "t".into(),
            columns: vec![],
            file_type: FileType::Parquet,
            has_header: false,
            location: "foo.parquet".into(),
        });
        expect_parse_ok(sql, expected)?;

        // Error cases: Invalid type
        let sql =
            "CREATE EXTERNAL TABLE t(c1 int) STORED AS UNKNOWN_TYPE LOCATION 'foo.csv'";
        expect_parse_error(
            sql,
            "Expected one of PARQUET, NDJSON, or CSV, found: UNKNOWN_TYPE",
        )?;

        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::io;

use super::datagram::{DatagramReceiver, ReadAdapter};
use super::ipv4_header::IPv4Header;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::transport_header::TransportHeader;

/// Convert from level 5 to level 3 by appending correct IP and transport headers.
pub struct Packetizer {
    buffer: Box<[u8; MAX_PACKET_LENGTH]>,
    transport_index: usize,
    payload_index: usize,
    ipv4_header: IPv4Header,
    transport_header: TransportHeader,
}

impl Packetizer {
    pub fn new(raw: &[u8], mut ipv4_header: IPv4Header, mut transport_header: TransportHeader) -> Self {
        let mut buffer = Box::new([0; MAX_PACKET_LENGTH]);

        let mut ipv4_header = ipv4_header.clone();
        let mut transport_header = transport_header.clone();

        let transport_index = ipv4_header.header_length() as usize;
        let payload_index = transport_index + transport_header.header_length() as usize;
        &mut buffer[..payload_index].copy_from_slice(&raw[..payload_index]);

        ipv4_header.swap_source_and_destination(&mut buffer[..]);
        transport_header.swap_source_and_destination(&mut buffer[transport_index..]);

        Self {
            buffer: buffer,
            transport_index: transport_index,
            payload_index: payload_index,
            ipv4_header: ipv4_header,
            transport_header: transport_header,
        }
    }

    pub fn packetize_empty_payload(&mut self) -> IPv4Packet {
        self.inflate(0)
    }

    pub fn packetize<R: DatagramReceiver>(&mut self, source: &mut R) -> io::Result<IPv4Packet> {
        let r = source.recv(&mut self.buffer[self.payload_index..])?;
        let ipv4_packet = self.inflate(r as u16);
        Ok(ipv4_packet)
    }

    pub fn packetize_read<R: io::Read>(&mut self, source: &mut R, max_chunk_size: Option<usize>) -> io::Result<IPv4Packet> {
        let mut adapter = ReadAdapter::new(source, max_chunk_size);
        self.packetize(&mut adapter)
    }

    fn inflate(&mut self, payload_length: u16) -> IPv4Packet {
        let total_length = self.payload_index as u16 + payload_length;

        self.ipv4_header.set_total_length(&mut self.buffer[..], total_length);
        self.transport_header.set_payload_length(&mut self.buffer[self.transport_index..], payload_length);

        let mut ipv4_packet = IPv4Packet::new(&mut self.buffer[..total_length as usize], self.ipv4_header.clone(), self.transport_header.clone());
        ipv4_packet.compute_checksums();
        ipv4_packet
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io;
    use byteorder::{BigEndian, WriteBytesExt};
    use relay::datagram::tests::MockDatagramSocket;

    fn create_packet() -> Vec<u8> {
        let mut raw = Vec::new();
        raw.write_u8(4u8 << 4 | 5).unwrap();
        raw.write_u8(0).unwrap(); // ToS
        raw.write_u16::<BigEndian>(32).unwrap(); // total length 20 + 8 + 4
        raw.write_u32::<BigEndian>(0).unwrap(); // id_flags_fragment_offset
        raw.write_u8(0).unwrap(); // TTL
        raw.write_u8(17).unwrap(); // protocol (UDP)
        raw.write_u16::<BigEndian>(0).unwrap(); // checksum
        raw.write_u32::<BigEndian>(0x12345678).unwrap(); // source address
        raw.write_u32::<BigEndian>(0x42424242).unwrap(); // destination address

        raw.write_u16::<BigEndian>(1234).unwrap(); // source port
        raw.write_u16::<BigEndian>(5678).unwrap(); // destination port
        raw.write_u16::<BigEndian>(12).unwrap(); // length
        raw.write_u16::<BigEndian>(0).unwrap(); // checksum

        raw.write_u32::<BigEndian>(0x11223344).unwrap(); // payload
        raw
    }

    #[test]
    fn merge_headers_and_payload() {
        let mut raw = &mut create_packet()[..];
        let reference_packet = IPv4Packet::parse(raw);

        let data = [ 0x11u8, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88 ];
        let mut mock = MockDatagramSocket::from_data(&data);

        let ipv4_header = reference_packet.ipv4_header().clone();
        let transport_header = reference_packet.transport_header().as_ref().unwrap().clone();
        let mut packetizer = Packetizer::new(reference_packet.raw(), ipv4_header, transport_header);

        let packet = packetizer.packetize(&mut mock).unwrap();
        assert_eq!(36, packet.ipv4_header().total_length());
    }

    #[test]
    fn packetize_chunks() {
        let mut raw = &mut create_packet()[..];
        let reference_packet = IPv4Packet::parse(raw);

        let data = [ 0x11u8, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88 ];
        let mut cursor = io::Cursor::new(&data);

        let ipv4_header = reference_packet.ipv4_header().clone();
        let transport_header = reference_packet.transport_header().as_ref().unwrap().clone();
        let mut packetizer = Packetizer::new(reference_packet.raw(), ipv4_header, transport_header);

        {
            let packet = packetizer.packetize_read(&mut cursor, Some(2)).unwrap();
            assert_eq!(30, packet.ipv4_header().total_length());
            assert_eq!(2, packet.payload_length().unwrap());
            assert_eq!([0x11, 0x22], packet.raw()[packet.payload_index().unwrap() as usize..]);
        }

        {
            let packet = packetizer.packetize_read(&mut cursor, Some(3)).unwrap();
            assert_eq!(31, packet.ipv4_header().total_length());
            assert_eq!(3, packet.payload_length().unwrap());
            assert_eq!([0x33, 0x44, 0x55], packet.raw()[packet.payload_index().unwrap() as usize..]);
        }

        {
            let packet = packetizer.packetize_read(&mut cursor, Some(1024)).unwrap();
            assert_eq!(31, packet.ipv4_header().total_length());
            assert_eq!(3, packet.payload_length().unwrap());
            assert_eq!([0x66, 0x77, 0x88], packet.raw()[packet.payload_index().unwrap() as usize..]);
        }
    }
}
<commit_msg>Remove useless clone()<commit_after>use std::io;

use super::datagram::{DatagramReceiver, ReadAdapter};
use super::ipv4_header::IPv4Header;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::transport_header::TransportHeader;

/// Convert from level 5 to level 3 by appending correct IP and transport headers.
pub struct Packetizer {
    buffer: Box<[u8; MAX_PACKET_LENGTH]>,
    transport_index: usize,
    payload_index: usize,
    ipv4_header: IPv4Header,
    transport_header: TransportHeader,
}

impl Packetizer {
    pub fn new(raw: &[u8], mut ipv4_header: IPv4Header, mut transport_header: TransportHeader) -> Self {
        let mut buffer = Box::new([0; MAX_PACKET_LENGTH]);

        let transport_index = ipv4_header.header_length() as usize;
        let payload_index = transport_index + transport_header.header_length() as usize;
        &mut buffer[..payload_index].copy_from_slice(&raw[..payload_index]);

        ipv4_header.swap_source_and_destination(&mut buffer[..]);
        transport_header.swap_source_and_destination(&mut buffer[transport_index..]);

        Self {
            buffer: buffer,
            transport_index: transport_index,
            payload_index: payload_index,
            ipv4_header: ipv4_header,
            transport_header: transport_header,
        }
    }

    pub fn packetize_empty_payload(&mut self) -> IPv4Packet {
        self.inflate(0)
    }

    pub fn packetize<R: DatagramReceiver>(&mut self, source: &mut R) -> io::Result<IPv4Packet> {
        let r = source.recv(&mut self.buffer[self.payload_index..])?;
        let ipv4_packet = self.inflate(r as u16);
        Ok(ipv4_packet)
    }

    pub fn packetize_read<R: io::Read>(&mut self, source: &mut R, max_chunk_size: Option<usize>) -> io::Result<IPv4Packet> {
        let mut adapter = ReadAdapter::new(source, max_chunk_size);
        self.packetize(&mut adapter)
    }

    fn inflate(&mut self, payload_length: u16) -> IPv4Packet {
        let total_length = self.payload_index as u16 + payload_length;

        self.ipv4_header.set_total_length(&mut self.buffer[..], total_length);
        self.transport_header.set_payload_length(&mut self.buffer[self.transport_index..], payload_length);

        let mut ipv4_packet = IPv4Packet::new(&mut self.buffer[..total_length as usize], self.ipv4_header.clone(), self.transport_header.clone());
        ipv4_packet.compute_checksums();
        ipv4_packet
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io;
    use byteorder::{BigEndian, WriteBytesExt};
    use relay::datagram::tests::MockDatagramSocket;

    fn create_packet() -> Vec<u8> {
        let mut raw = Vec::new();
        raw.write_u8(4u8 << 4 | 5).unwrap();
        raw.write_u8(0).unwrap(); // ToS
        raw.write_u16::<BigEndian>(32).unwrap(); // total length 20 + 8 + 4
        raw.write_u32::<BigEndian>(0).unwrap(); // id_flags_fragment_offset
        raw.write_u8(0).unwrap(); // TTL
        raw.write_u8(17).unwrap(); // protocol (UDP)
        raw.write_u16::<BigEndian>(0).unwrap(); // checksum
        raw.write_u32::<BigEndian>(0x12345678).unwrap(); // source address
        raw.write_u32::<BigEndian>(0x42424242).unwrap(); // destination address

        raw.write_u16::<BigEndian>(1234).unwrap(); // source port
        raw.write_u16::<BigEndian>(5678).unwrap(); // destination port
        raw.write_u16::<BigEndian>(12).unwrap(); // length
        raw.write_u16::<BigEndian>(0).unwrap(); // checksum

        raw.write_u32::<BigEndian>(0x11223344).unwrap(); // payload
        raw
    }

    #[test]
    fn merge_headers_and_payload() {
        let mut raw = &mut create_packet()[..];
        let reference_packet = IPv4Packet::parse(raw);

        let data = [ 0x11u8, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88 ];
        let mut mock = MockDatagramSocket::from_data(&data);

        let ipv4_header = reference_packet.ipv4_header().clone();
        let transport_header = reference_packet.transport_header().as_ref().unwrap().clone();
        let mut packetizer = Packetizer::new(reference_packet.raw(), ipv4_header, transport_header);

        let packet = packetizer.packetize(&mut mock).unwrap();
        assert_eq!(36, packet.ipv4_header().total_length());
    }

    #[test]
    fn packetize_chunks() {
        let mut raw = &mut create_packet()[..];
        let reference_packet = IPv4Packet::parse(raw);

        let data = [ 0x11u8, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88 ];
        let mut cursor = io::Cursor::new(&data);

        let ipv4_header = reference_packet.ipv4_header().clone();
        let transport_header = reference_packet.transport_header().as_ref().unwrap().clone();
        let mut packetizer = Packetizer::new(reference_packet.raw(), ipv4_header, transport_header);

        {
            let packet = packetizer.packetize_read(&mut cursor, Some(2)).unwrap();
            assert_eq!(30, packet.ipv4_header().total_length());
            assert_eq!(2, packet.payload_length().unwrap());
            assert_eq!([0x11, 0x22], packet.raw()[packet.payload_index().unwrap() as usize..]);
        }

        {
            let packet = packetizer.packetize_read(&mut cursor, Some(3)).unwrap();
            assert_eq!(31, packet.ipv4_header().total_length());
            assert_eq!(3, packet.payload_length().unwrap());
            assert_eq!([0x33, 0x44, 0x55], packet.raw()[packet.payload_index().unwrap() as usize..]);
        }

        {
            let packet = packetizer.packetize_read(&mut cursor, Some(1024)).unwrap();
            assert_eq!(31, packet.ipv4_header().total_length());
            assert_eq!(3, packet.payload_length().unwrap());
            assert_eq!([0x66, 0x77, 0x88], packet.raw()[packet.payload_index().unwrap() as usize..]);
        }
    }
}
<|endoftext|>


<commit_before>use consts::*;
use proto::*;
use io::ZkIo;
use mio;
use num::FromPrimitive;
use watch::{Watch, Watcher, WatchType, ZkWatch};
use std::io::{Read, Result};
use std::net::{SocketAddr, ToSocketAddrs};
use std::result;
use std::sync::atomic::{AtomicIsize, Ordering};
use std::sync::mpsc::{Sender, sync_channel, SyncSender};
use std::time::Duration;
use std::thread;

pub type ZkResult<T> = result::Result<T, ZkError>;

pub struct RawRequest {
    pub opcode: OpCode,
    pub data: ByteBuf,
    pub listener: Option<SyncSender<RawResponse>>,
    pub watch: Option<Watch>
}

pub struct RawResponse {
    pub header: ReplyHeader,
    pub data: ByteBuf
}

pub struct ZooKeeper {
    chroot: Option<String>,
    xid: AtomicIsize,
    io: mio::Sender<RawRequest>
}

impl ZooKeeper {

    fn zk_thread<F>(name: &str, task: F) -> ZkResult<thread::JoinHandle<()>>
    where F: FnOnce() + Send + 'static {
        thread::Builder::new()
            .name(name.to_owned())
            .spawn(task)
            .map_err(|_|ZkError::SystemError)
    }

    pub fn connect<W>(connect_string: &str, timeout: Duration, watcher: W) -> ZkResult<ZooKeeper>
    where W: Watcher + 'static {

        let (addrs, chroot) = try!(Self::parse_connect_string(connect_string));

        debug!("Initiating connection to {}", connect_string);

        let watch = ZkWatch::new(watcher);
        let io = ZkIo::new(addrs.clone(), timeout, watch.sender());
        let sender = io.sender();

        try!(Self::zk_thread("event", move || { watch.run().unwrap() }));
        try!(Self::zk_thread("io", move || { io.run().unwrap() }));

        Ok(ZooKeeper{chroot: chroot,
                     xid: AtomicIsize::new(1),
                     io: sender})
    }

    fn parse_connect_string(connect_string: &str) -> ZkResult<(Vec<SocketAddr>, Option<String>)> {
        let (chroot, end) = match connect_string.find('/') {
            Some(start) => {
                match &connect_string[start..connect_string.len()] {
                    "" | "/" => (None, start),
                    chroot => (Some(try!(Self::validate_path(chroot)).to_owned()), start)
                }
            },
            None => (None, connect_string.len())
        };

        let mut addrs = Vec::new();
        for addr_str in connect_string[..end].split(',') {
            let addr = match addr_str.to_socket_addrs() {
                Ok(mut addrs) => match addrs.nth(0) {
                    Some(addr) => addr,
                    None => return Err(ZkError::BadArguments)
                },
                Err(_) => return Err(ZkError::BadArguments)
            };
            addrs.push(addr);
        }

        Ok((addrs, chroot))
    }

    fn xid(&self) -> i32 {
        self.xid.fetch_add(1, Ordering::Relaxed) as i32
    }

    fn request<Req: WriteTo, Resp: ReadFrom>(&self, opcode: OpCode, xid: i32, req: Req, watch: Option<Watch>) -> ZkResult<Resp> {
        let rh = RequestHeader{xid: xid, opcode: opcode};
        let buf = try!(to_len_prefixed_buf(rh, req).map_err(|_|ZkError::MarshallingError));

        let (resp_tx, resp_rx) = sync_channel(0);
        let request = RawRequest{opcode: opcode, data: buf, listener: Some(resp_tx), watch: watch};

        try!(self.io.send(request).map_err(|_|ZkError::ConnectionLoss));

        let mut response = try!(resp_rx.recv().map_err(|_|ZkError::ConnectionLoss));

        match response.header.err {
            0 => Ok(try!(ReadFrom::read_from(&mut response.data).map_err(|_|ZkError::MarshallingError))),
            e => Err(FromPrimitive::from_i32(e).unwrap())
        }
    }

    fn validate_path(path: &str) -> ZkResult<&str> {
        match path {
            "" => Err(ZkError::BadArguments),
            path => if path.len() > 1 && path.chars().last() == Some('/') {
                    Err(ZkError::BadArguments)
                } else {
                    Ok(path)
                }
        }
    }

    fn path(&self, path: &str) -> ZkResult<String> {
        match self.chroot {
            Some(ref chroot) => match path {
                "/" => Ok(chroot.clone()),
                path => Ok(chroot.clone() + try!(Self::validate_path(path)))
            },
            None => Ok(try!(Self::validate_path(path)).to_owned())
        }
    }

    pub fn add_auth(&self, scheme: &str, auth: Vec<u8>) -> ZkResult<()> {
        let req = AuthRequest{typ: 0, scheme: scheme.to_owned(), auth: auth};

        let _: EmptyResponse = try!(self.request(OpCode::Auth, -4, req, None));

        Ok(())
    }

    pub fn create(&self, path: &str, data: Vec<u8>, acl: Vec<Acl>, mode: CreateMode) -> ZkResult<String> {
        let req = CreateRequest{path: try!(self.path(path)), data: data, acl: acl, flags: mode as i32};

        let response: CreateResponse = try!(self.request(OpCode::Create, self.xid(), req, None));

        Ok(response.path)
    }

    pub fn delete(&self, path: &str, version: i32) -> ZkResult<()> {
        let req = DeleteRequest{path: try!(self.path(path)), version: version};

        let _: EmptyResponse = try!(self.request(OpCode::Delete, self.xid(), req, None));

        Ok(())
    }

    pub fn exists(&self, path: &str, watch: bool) -> ZkResult<Option<Stat>> {
        let req = ExistsRequest{path: try!(self.path(path)), watch: watch};

        match self.request::<ExistsRequest, ExistsResponse>(OpCode::Exists, self.xid(), req, None) {
            Ok(response) => Ok(Some(response.stat)),
            Err(ZkError::NoNode) => Ok(None),
            Err(e) => Err(e)
        }
    }

    pub fn exists_w<W: FnOnce(&WatchedEvent) + Send + 'static>(&self, path: &str, watcher: W) -> ZkResult<Stat> {
        let req = ExistsRequest{path: try!(self.path(path)), watch: true};

        let watch = Watch{path: path.to_owned(),
                          watch_type: WatchType::Exist,
                          watcher: Box::new(watcher)};

        let response: ExistsResponse = try!(self.request(OpCode::Exists, self.xid(), req, Some(watch)));

        Ok(response.stat)
    }

    pub fn get_acl(&self, path: &str) -> ZkResult<(Vec<Acl>, Stat)> {
        let req = GetAclRequest{path: try!(self.path(path))};

        let response: GetAclResponse = try!(self.request(OpCode::GetAcl, self.xid(), req, None));

        Ok(response.acl_stat)
    }

    pub fn get_children_w<W: FnOnce(&WatchedEvent) + Send + 'static>(&self, path: &str, watcher: W) -> ZkResult<Vec<String>> {
        let req = GetChildrenRequest{path: try!(self.path(path)), watch: true};

        let watch = Watch{path: path.to_owned(),
                          watch_type: WatchType::Child,
                          watcher: Box::new(watcher)};

        let response: GetChildrenResponse = try!(self.request(OpCode::GetChildren, self.xid(), req, Some(watch)));

        Ok(response.children)
    }

    pub fn get_children(&self, path: &str, watch: bool) -> ZkResult<Vec<String>> {
        let req = GetChildrenRequest{path: try!(self.path(path)), watch: watch};

        let response: GetChildrenResponse = try!(self.request(OpCode::GetChildren, self.xid(), req, None));

        Ok(response.children)
    }

    pub fn get_data(&self, path: &str, watch: bool) -> ZkResult<(Vec<u8>, Stat)> {
        let req = GetDataRequest{path: try!(self.path(path)), watch: watch};

        let response: GetDataResponse = try!(self.request(OpCode::GetData, self.xid(), req, None));

        Ok(response.data_stat)
    }

    pub fn get_data_w<W: FnOnce(&WatchedEvent) + Send + 'static>(&self, path: &str, watcher: W) -> ZkResult<(Vec<u8>, Stat)> {
        let req = GetDataRequest{path: try!(self.path(path)), watch: true};

        let watch = Watch{path: path.to_owned(),
                          watch_type: WatchType::Data,
                          watcher: Box::new(watcher)};

        let response: GetDataResponse = try!(self.request(OpCode::GetData, self.xid(), req, Some(watch)));

        Ok(response.data_stat)
    }

    pub fn set_acl(&self, path: &str, acl: Vec<Acl>, version: i32) -> ZkResult<Stat> {
        let req = SetAclRequest{path: try!(self.path(path)), acl: acl, version: version};

        let response: SetAclResponse = try!(self.request(OpCode::SetAcl, self.xid(), req, None));

        Ok(response.stat)
    }

    pub fn set_data(&self, path: &str, data: Vec<u8>, version: i32) -> ZkResult<Stat> {
        let req = SetDataRequest{path: try!(self.path(path)), data: data, version: version};

        let response: SetDataResponse = try!(self.request(OpCode::SetData, self.xid(), req, None));

        Ok(response.stat)
    }

    pub fn close(&self) -> ZkResult<()> {
        let _: EmptyResponse = try!(self.request(OpCode::CloseSession, 0, EmptyRequest, None));

        Ok(())
    }
}

impl Drop for ZooKeeper {
    fn drop(&mut self) {
        // TODO First check if state is closed
        self.close().unwrap();
    }
}

#[cfg(test)]
mod tests {
    use super::ZooKeeper;
    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};

    // TODO This is flaky, host resolving happens in a hectic order, sort?
    #[test]
    fn parse_connect_string() {
        let (addrs, chroot) = ZooKeeper::parse_connect_string("127.0.0.1:2181,::1:2181/mesos").unwrap();
        assert_eq!(addrs,
                   vec![SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127,0,0,1), 2181)),
                        SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0,0,0,0,0,0,0,1), 2181, 0, 0))]);
        assert_eq!(chroot.unwrap(), "/mesos");

        let (addrs, chroot) = ZooKeeper::parse_connect_string("::1:2181").unwrap();
        assert_eq!(addrs, vec![SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0,0,0,0,0,0,0,1), 2181, 0, 0))]);
        assert_eq!(chroot, None);

        let (addrs, chroot) = ZooKeeper::parse_connect_string("::1:2181/").unwrap();
        assert_eq!(addrs, vec![SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0,0,0,0,0,0,0,1), 2181, 0, 0))]);
        assert_eq!(chroot, None);
    }

    #[test]
    #[should_panic(expected = "BadArguments")]
    fn parse_connect_string_fails() {
        // This fails with ZooKeeper.java: Path must not end with / character
        ZooKeeper::parse_connect_string("127.0.0.1:2181/mesos/").unwrap();
    }
}
<commit_msg>Update zookeeper.rs<commit_after>use consts::*;
use proto::*;
use io::ZkIo;
use mio;
use num::FromPrimitive;
use watch::{Watch, Watcher, WatchType, ZkWatch};
use std::io::{Read, Result};
use std::net::{SocketAddr, ToSocketAddrs};
use std::result;
use std::sync::atomic::{AtomicIsize, Ordering};
use std::sync::mpsc::{Sender, sync_channel, SyncSender};
use std::time::Duration;
use std::thread;

pub type ZkResult<T> = result::Result<T, ZkError>;

pub struct RawRequest {
    pub opcode: OpCode,
    pub data: ByteBuf,
    pub listener: Option<SyncSender<RawResponse>>,
    pub watch: Option<Watch>
}

pub struct RawResponse {
    pub header: ReplyHeader,
    pub data: ByteBuf
}

pub struct ZooKeeper {
    chroot: Option<String>,
    xid: AtomicIsize,
    io: mio::Sender<RawRequest>
}

impl ZooKeeper {

    fn zk_thread<F>(name: &str, task: F) -> ZkResult<thread::JoinHandle<()>>
    where F: FnOnce() + Send + 'static {
        thread::Builder::new()
            .name(name.to_owned())
            .spawn(task)
            .map_err(|_|ZkError::SystemError)
    }

    pub fn connect<W>(connect_string: &str, timeout: Duration, watcher: W) -> ZkResult<ZooKeeper>
    where W: Watcher + 'static {

        let (addrs, chroot) = try!(Self::parse_connect_string(connect_string));

        debug!("Initiating connection to {}", connect_string);

        let watch = ZkWatch::new(watcher);
        let io = ZkIo::new(addrs.clone(), timeout, watch.sender());
        let sender = io.sender();

        try!(Self::zk_thread("event", move || { watch.run().unwrap() }));
        try!(Self::zk_thread("io", move || { io.run().unwrap() }));

        Ok(ZooKeeper{chroot: chroot,
                     xid: AtomicIsize::new(1),
                     io: sender})
    }

    fn parse_connect_string(connect_string: &str) -> ZkResult<(Vec<SocketAddr>, Option<String>)> {
        let (chroot, end) = match connect_string.find('/') {
            Some(start) => {
                match &connect_string[start..connect_string.len()] {
                    "" | "/" => (None, start),
                    chroot => (Some(try!(Self::validate_path(chroot)).to_owned()), start)
                }
            },
            None => (None, connect_string.len())
        };

        let mut addrs = Vec::new();
        for addr_str in connect_string[..end].split(',') {
            let addr = match addr_str.to_socket_addrs() {
                Ok(mut addrs) => match addrs.nth(0) {
                    Some(addr) => addr,
                    None => return Err(ZkError::BadArguments)
                },
                Err(_) => return Err(ZkError::BadArguments)
            };
            addrs.push(addr);
        }

        Ok((addrs, chroot))
    }

    fn xid(&self) -> i32 {
        self.xid.fetch_add(1, Ordering::Relaxed) as i32
    }

    fn request<Req: WriteTo, Resp: ReadFrom>(&self, opcode: OpCode, xid: i32, req: Req, watch: Option<Watch>) -> ZkResult<Resp> {
        let rh = RequestHeader{xid: xid, opcode: opcode};
        let buf = try!(to_len_prefixed_buf(rh, req).map_err(|_|ZkError::MarshallingError));

        let (resp_tx, resp_rx) = sync_channel(0);
        let request = RawRequest{opcode: opcode, data: buf, listener: Some(resp_tx), watch: watch};

        try!(self.io.send(request).map_err(|_|ZkError::ConnectionLoss));

        let mut response = try!(resp_rx.recv().map_err(|_|ZkError::ConnectionLoss));

        match response.header.err {
            0 => Ok(try!(ReadFrom::read_from(&mut response.data).map_err(|_|ZkError::MarshallingError))),
            e => Err(FromPrimitive::from_i32(e).unwrap())
        }
    }

    fn validate_path(path: &str) -> ZkResult<&str> {
        match path {
            "" => Err(ZkError::BadArguments),
            path => if path.len() > 1 && path.chars().last() == Some('/') {
                    Err(ZkError::BadArguments)
                } else {
                    Ok(path)
                }
        }
    }

    fn path(&self, path: &str) -> ZkResult<String> {
        match self.chroot {
            Some(ref chroot) => match path {
                "/" => Ok(chroot.clone()),
                path => Ok(chroot.clone() + try!(Self::validate_path(path)))
            },
            None => Ok(try!(Self::validate_path(path)).to_owned())
        }
    }

    pub fn add_auth(&self, scheme: &str, auth: Vec<u8>) -> ZkResult<()> {
        let req = AuthRequest{typ: 0, scheme: scheme.to_owned(), auth: auth};

        let _: EmptyResponse = try!(self.request(OpCode::Auth, -4, req, None));

        Ok(())
    }

    pub fn create(&self, path: &str, data: Vec<u8>, acl: Vec<Acl>, mode: CreateMode) -> ZkResult<String> {
        let req = CreateRequest{path: try!(self.path(path)), data: data, acl: acl, flags: mode as i32};

        let response: CreateResponse = try!(self.request(OpCode::Create, self.xid(), req, None));

        Ok(response.path)
    }

    pub fn delete(&self, path: &str, version: i32) -> ZkResult<()> {
        let req = DeleteRequest{path: try!(self.path(path)), version: version};

        let _: EmptyResponse = try!(self.request(OpCode::Delete, self.xid(), req, None));

        Ok(())
    }

    pub fn exists(&self, path: &str, watch: bool) -> ZkResult<Option<Stat>> {
        let req = ExistsRequest{path: try!(self.path(path)), watch: watch};

        match self.request::<ExistsRequest, ExistsResponse>(OpCode::Exists, self.xid(), req, None) {
            Ok(response) => Ok(Some(response.stat)),
            Err(ZkError::NoNode) => Ok(None),
            Err(e) => Err(e)
        }
    }

    pub fn exists_w<W: FnOnce(&WatchedEvent) + Send + 'static>(&self, path: &str, watcher: W) -> ZkResult<Stat> {
        let req = ExistsRequest{path: try!(self.path(path)), watch: true};

        let watch = Watch{path: path.to_owned(),
                          watch_type: WatchType::Exist,
                          watcher: Box::new(watcher)};

        let response: ExistsResponse = try!(self.request(OpCode::Exists, self.xid(), req, Some(watch)));

        Ok(response.stat)
    }

    pub fn get_acl(&self, path: &str) -> ZkResult<(Vec<Acl>, Stat)> {
        let req = GetAclRequest{path: try!(self.path(path))};

        let response: GetAclResponse = try!(self.request(OpCode::GetAcl, self.xid(), req, None));

        Ok(response.acl_stat)
    }

    pub fn get_children_w<W: FnOnce(&WatchedEvent) + Send + 'static>(&self, path: &str, watcher: W) -> ZkResult<Vec<String>> {
        let req = GetChildrenRequest{path: try!(self.path(path)), watch: true};

        let watch = Watch{path: path.to_owned(),
                          watch_type: WatchType::Child,
                          watcher: Box::new(watcher)};

        let response: GetChildrenResponse = try!(self.request(OpCode::GetChildren, self.xid(), req, Some(watch)));

        Ok(response.children)
    }

    pub fn get_children(&self, path: &str, watch: bool) -> ZkResult<Vec<String>> {
        let req = GetChildrenRequest{path: try!(self.path(path)), watch: watch};

        let response: GetChildrenResponse = try!(self.request(OpCode::GetChildren, self.xid(), req, None));

        Ok(response.children)
    }

    pub fn get_data(&self, path: &str, watch: bool) -> ZkResult<(Vec<u8>, Stat)> {
        let req = GetDataRequest{path: try!(self.path(path)), watch: watch};

        let response: GetDataResponse = try!(self.request(OpCode::GetData, self.xid(), req, None));

        Ok(response.data_stat)
    }

    pub fn get_data_w<W: FnOnce(&WatchedEvent) + Send + 'static>(&self, path: &str, watcher: W) -> ZkResult<(Vec<u8>, Stat)> {
        let req = GetDataRequest{path: try!(self.path(path)), watch: true};

        let watch = Watch{path: path.to_owned(),
                          watch_type: WatchType::Data,
                          watcher: Box::new(watcher)};

        let response: GetDataResponse = try!(self.request(OpCode::GetData, self.xid(), req, Some(watch)));

        Ok(response.data_stat)
    }

    pub fn set_acl(&self, path: &str, acl: Vec<Acl>, version: i32) -> ZkResult<Stat> {
        let req = SetAclRequest{path: try!(self.path(path)), acl: acl, version: version};

        let response: SetAclResponse = try!(self.request(OpCode::SetAcl, self.xid(), req, None));

        Ok(response.stat)
    }

    pub fn set_data(&self, path: &str, data: Vec<u8>, version: i32) -> ZkResult<Stat> {
        let req = SetDataRequest{path: try!(self.path(path)), data: data, version: version};

        let response: SetDataResponse = try!(self.request(OpCode::SetData, self.xid(), req, None));

        Ok(response.stat)
    }

    pub fn close(&self) -> ZkResult<()> {
        let _: EmptyResponse = try!(self.request(OpCode::CloseSession, 0, EmptyRequest, None));

        Ok(())
    }
}

impl Drop for ZooKeeper {
    fn drop(&mut self) {
        // TODO First check if state is closed
        self.close();
    }
}

#[cfg(test)]
mod tests {
    use super::ZooKeeper;
    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};

    // TODO This is flaky, host resolving happens in a hectic order, sort?
    #[test]
    fn parse_connect_string() {
        let (addrs, chroot) = ZooKeeper::parse_connect_string("127.0.0.1:2181,::1:2181/mesos").unwrap();
        assert_eq!(addrs,
                   vec![SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127,0,0,1), 2181)),
                        SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0,0,0,0,0,0,0,1), 2181, 0, 0))]);
        assert_eq!(chroot.unwrap(), "/mesos");

        let (addrs, chroot) = ZooKeeper::parse_connect_string("::1:2181").unwrap();
        assert_eq!(addrs, vec![SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0,0,0,0,0,0,0,1), 2181, 0, 0))]);
        assert_eq!(chroot, None);

        let (addrs, chroot) = ZooKeeper::parse_connect_string("::1:2181/").unwrap();
        assert_eq!(addrs, vec![SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0,0,0,0,0,0,0,1), 2181, 0, 0))]);
        assert_eq!(chroot, None);
    }

    #[test]
    #[should_panic(expected = "BadArguments")]
    fn parse_connect_string_fails() {
        // This fails with ZooKeeper.java: Path must not end with / character
        ZooKeeper::parse_connect_string("127.0.0.1:2181/mesos/").unwrap();
    }
}
<|endoftext|>