<commit_before>use std::{mem, ptr};
use math::*;
use brush::{self, Brush};
use geometry::Geometry;
use error::D2D1Error;
use stroke_style::StrokeStyle;
use factory::Factory;
use comptr::ComPtr;
use helpers::{FromRaw, GetRaw, ToWide};
use directwrite::{TextFormat, TextLayout};

use winapi::shared::winerror::*;
use winapi::um::d2d1::*;
use winapi::um::d2d1_1::*;
use winapi::um::dcommon::*;

/// This trait is intended to be implemented by external APIs who would
/// like to allow a Direct2D interface for drawing onto them. Since the
/// Direct2D creation APIs require more information, this will likely be
/// implemented on Builder objects which contain enough context about
/// the actual backing structure and what needs to be done to set it up.
pub unsafe trait RenderTargetBacking {
    /// The ID2D1RenderTarget's ownership is passed out of the function and as such
    /// the caller is now responsible for ensuring the pointer will receive its
    /// Release() call.
    fn create_target(self, factory: &mut ID2D1Factory1) -> Result<*mut ID2D1RenderTarget, HRESULT>;
}

#[derive(Clone, Debug)]
pub struct RenderTarget {
    ptr: ComPtr<ID2D1RenderTarget>,
}

#[derive(Copy, Clone, Debug)]
pub struct RenderTag {
    pub file_line: &'static str,
}

struct RenderTagRaw(usize, usize);

#[macro_export]
macro_rules! make_render_tag {
    () => {
        RenderTag {
            file: concat!(file!(), ':', line!())
        }
    }
}

#[macro_export]
macro_rules! set_render_tag {
    ($rt:ident) => {
        $rt.set_tag(make_render_tag!());
    }
}

impl RenderTarget {
    unsafe fn rt<'a>(&self) -> &'a mut ID2D1RenderTarget {
        &mut *self.ptr.raw_value()
    }

    pub unsafe fn hwnd_rt(&self) -> Option<ComPtr<ID2D1HwndRenderTarget>> {
        self.ptr.query_interface().ok()
    }

    unsafe fn make_tag(tag1: D2D1_TAG, tag2: D2D1_TAG) -> Option<RenderTag> {
        if tag1 == 0 {
            None
        } else {
            let raw = RenderTagRaw(tag1 as usize, tag2 as usize);
            let tag = mem::transmute(raw);
            Some(tag)
        }
    }

    pub fn get_factory(&mut self) -> Factory {
        unsafe {
            let mut factory = ComPtr::<ID2D1Factory>::new();
            self.rt().GetFactory(factory.raw_addr());

            Factory::from_ptr(factory.query_interface().unwrap())
        }
    }

    pub fn get_size(&self) -> SizeF {
        unsafe { SizeF(self.rt().GetSize()) }
    }

    pub fn create_solid_color_brush<C: Into<ColorF>>(
        &self,
        color: C,
        props: &BrushProperties,
    ) -> Result<brush::SolidColor, D2D1Error> {
        unsafe {
            let mut ptr = ComPtr::<ID2D1SolidColorBrush>::new();
            let result = self.rt()
                .CreateSolidColorBrush(&color.into().0, &props.0, ptr.raw_addr());

            if SUCCEEDED(result) {
                Ok(FromRaw::from_raw(ptr.raw_value()))
            } else {
                Err(From::from(result))
            }
        }
    }

    pub fn begin_draw(&mut self) {
        unsafe {
            self.rt().BeginDraw();
        }
    }

    pub fn end_draw(&mut self) -> Result<(), (D2D1Error, Option<RenderTag>)> {
        let mut tag1 = 0;
        let mut tag2 = 0;
        unsafe {
            let result = self.rt().EndDraw(&mut tag1, &mut tag2);

            if SUCCEEDED(result) {
                Ok(())
            } else {
                let tag = RenderTarget::make_tag(tag1, tag2);
                Err((From::from(result), tag))
            }
        }
    }

    pub fn set_tag(&mut self, tag: RenderTag) {
        unsafe {
            let RenderTagRaw(tag1, tag2) = mem::transmute(tag);
            self.rt().SetTags(tag1 as u64, tag2 as u64)
        };
    }

    pub fn get_tag(&mut self) -> Option<RenderTag> {
        let mut tag1 = 0;
        let mut tag2 = 0;
        unsafe {
            self.rt().GetTags(&mut tag1, &mut tag2);
            RenderTarget::make_tag(tag1, tag2)
        }
    }

    pub fn flush(&mut self) -> Result<(), (D2D1Error, Option<RenderTag>)> {
        let mut tag1 = 0;
        let mut tag2 = 0;
        unsafe {
            let result = self.rt().Flush(&mut tag1, &mut tag2);

            if SUCCEEDED(result) {
                Ok(())
            } else {
                let tag = RenderTarget::make_tag(tag1, tag2);
                Err((From::from(result), tag))
            }
        }
    }

    pub fn clear(&mut self, color: &ColorF) {
        unsafe {
            self.rt().Clear(&color.0);
        }
    }

    pub fn draw_line<B: Brush>(
        &mut self,
        p0: &Point2F,
        p1: &Point2F,
        brush: &B,
        stroke_width: f32,
        stroke_style: Option<&StrokeStyle>,
    ) {
        unsafe {
            let stroke_style = match stroke_style {
                Some(s) => s.get_ptr() as *mut _,
                None => ptr::null_mut(),
            };

            self.rt()
                .DrawLine(p0.0, p1.0, brush.get_ptr(), stroke_width, stroke_style)
        }
    }

    pub fn draw_rectangle<B: Brush>(
        &mut self,
        rect: &RectF,
        brush: &B,
        stroke_width: f32,
        stroke_style: Option<&StrokeStyle>,
    ) {
        unsafe {
            let stroke_style = match stroke_style {
                Some(s) => s.get_ptr() as *mut _,
                None => ptr::null_mut(),
            };

            self.rt()
                .DrawRectangle(&rect.0, brush.get_ptr(), stroke_width, stroke_style);
        }
    }

    pub fn fill_rectangle<B: Brush>(&mut self, rect: &RectF, brush: &B) {
        unsafe {
            self.rt().FillRectangle(&rect.0, brush.get_ptr());
        }
    }

    pub fn draw_rounded_rectangle<B: Brush>(
        &mut self,
        rect: &RoundedRect,
        brush: &B,
        stroke_width: f32,
        stroke_style: Option<&StrokeStyle>,
    ) {
        unsafe {
            let stroke_style = match stroke_style {
                Some(s) => s.get_ptr() as *mut _,
                None => ptr::null_mut(),
            };

            self.rt()
                .DrawRoundedRectangle(&rect.0, brush.get_ptr(), stroke_width, stroke_style);
        }
    }

    pub fn fill_rounded_rectangle<B: Brush>(&mut self, rect: &RoundedRect, brush: &B) {
        unsafe {
            self.rt().FillRoundedRectangle(&rect.0, brush.get_ptr());
        }
    }

    pub fn draw_ellipse<B: Brush>(
        &mut self,
        ellipse: &Ellipse,
        brush: &B,
        stroke_width: f32,
        stroke_style: Option<&StrokeStyle>,
    ) {
        unsafe {
            let stroke_style = match stroke_style {
                Some(s) => s.get_ptr() as *mut _,
                None => ptr::null_mut(),
            };

            self.rt()
                .DrawEllipse(&ellipse.0, brush.get_ptr(), stroke_width, stroke_style);
        }
    }

    pub fn fill_ellipse<B: Brush>(&mut self, ellipse: &Ellipse, brush: &B) {
        unsafe {
            self.rt().FillEllipse(&ellipse.0, brush.get_ptr());
        }
    }

    pub fn draw_geometry<G: Geometry, B: Brush>(
        &mut self,
        geometry: &G,
        brush: &B,
        stroke_width: f32,
        stroke_style: Option<&StrokeStyle>,
    ) {
        unsafe {
            let stroke_style = match stroke_style {
                Some(s) => s.get_ptr() as *mut _,
                None => ptr::null_mut(),
            };

            self.rt().DrawGeometry(
                geometry.get_ptr(),
                brush.get_ptr(),
                stroke_width,
                stroke_style,
            );
        }
    }

    pub fn fill_geometry<G: Geometry, B: Brush>(&mut self, geometry: &G, brush: &B) {
        unsafe {
            self.rt()
                .FillGeometry(geometry.get_ptr(), brush.get_ptr(), ptr::null_mut());
        }
    }

    pub fn fill_geometry_with_opacity<G: Geometry, B: Brush, OB: Brush>(
        &mut self,
        geometry: &G,
        brush: &B,
        opacity_brush: &OB,
    ) {
        unsafe {
            self.rt()
                .FillGeometry(geometry.get_ptr(), brush.get_ptr(), opacity_brush.get_ptr());
        }
    }

    pub fn draw_text<B: Brush>(
        &mut self,
        text: &str,
        format: &TextFormat,
        layout_rect: &RectF,
        foreground_brush: &B,
        options: &[DrawTextOption],
    ) {
        let text = text.to_wide_null();
        let mut draw_options = D2D1_DRAW_TEXT_OPTIONS_NONE;
        for &option in options {
            draw_options |= option as u32;
        }

        unsafe {
            let format = format.get_raw();
            self.rt().DrawText(
                text.as_ptr(),
                text.len() as u32,
                format,
                &layout_rect.0,
                foreground_brush.get_ptr(),
                (draw_options),
                DWRITE_MEASURING_MODE_NATURAL,
            );
        }
    }

    pub fn draw_text_layout<B: Brush>(
        &mut self,
        origin: &Point2F,
        layout: &TextLayout,
        brush: &B,
        options: &[DrawTextOption],
    ) {
        let mut draw_options = D2D1_DRAW_TEXT_OPTIONS_NONE;
        for &option in options {
            draw_options |= option as u32;
        }

        unsafe {
            let layout = layout.get_raw();
            self.rt()
                .DrawTextLayout(origin.0, layout, brush.get_ptr(), (draw_options));
        }
    }
}

impl GetRaw for RenderTarget {
    type Raw = ID2D1RenderTarget;
    unsafe fn get_raw(&self) -> *mut ID2D1RenderTarget {
        self.ptr.raw_value()
    }
}

impl FromRaw for RenderTarget {
    type Raw = ID2D1RenderTarget;
    unsafe fn from_raw(raw: *mut ID2D1RenderTarget) -> Self {
        RenderTarget {
            ptr: ComPtr::attach(raw),
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum DrawTextOption {
    NoSnap = 1,
    Clip = 2,
    EnableColorFont = 4,
}
<commit_msg>Add set_transform and get_transform to render target<commit_after>use std::{mem, ptr};
use math::*;
use brush::{self, Brush};
use geometry::Geometry;
use error::D2D1Error;
use stroke_style::StrokeStyle;
use factory::Factory;
use comptr::ComPtr;
use helpers::{FromRaw, GetRaw, ToWide};
use directwrite::{TextFormat, TextLayout};

use winapi::shared::winerror::*;
use winapi::um::d2d1::*;
use winapi::um::d2d1_1::*;
use winapi::um::dcommon::*;

/// This trait is intended to be implemented by external APIs who would
/// like to allow a Direct2D interface for drawing onto them. Since the
/// Direct2D creation APIs require more information, this will likely be
/// implemented on Builder objects which contain enough context about
/// the actual backing structure and what needs to be done to set it up.
pub unsafe trait RenderTargetBacking {
    /// The ID2D1RenderTarget's ownership is passed out of the function and as such
    /// the caller is now responsible for ensuring the pointer will receive its
    /// Release() call.
    fn create_target(self, factory: &mut ID2D1Factory1) -> Result<*mut ID2D1RenderTarget, HRESULT>;
}

#[derive(Clone, Debug)]
pub struct RenderTarget {
    ptr: ComPtr<ID2D1RenderTarget>,
}

#[derive(Copy, Clone, Debug)]
pub struct RenderTag {
    pub file_line: &'static str,
}

struct RenderTagRaw(usize, usize);

#[macro_export]
macro_rules! make_render_tag {
    () => {
        RenderTag {
            file: concat!(file!(), ':', line!())
        }
    }
}

#[macro_export]
macro_rules! set_render_tag {
    ($rt:ident) => {
        $rt.set_tag(make_render_tag!());
    }
}

impl RenderTarget {
    unsafe fn rt<'a>(&self) -> &'a mut ID2D1RenderTarget {
        &mut *self.ptr.raw_value()
    }

    pub unsafe fn hwnd_rt(&self) -> Option<ComPtr<ID2D1HwndRenderTarget>> {
        self.ptr.query_interface().ok()
    }

    unsafe fn make_tag(tag1: D2D1_TAG, tag2: D2D1_TAG) -> Option<RenderTag> {
        if tag1 == 0 {
            None
        } else {
            let raw = RenderTagRaw(tag1 as usize, tag2 as usize);
            let tag = mem::transmute(raw);
            Some(tag)
        }
    }

    pub fn get_factory(&mut self) -> Factory {
        unsafe {
            let mut factory = ComPtr::<ID2D1Factory>::new();
            self.rt().GetFactory(factory.raw_addr());

            Factory::from_ptr(factory.query_interface().unwrap())
        }
    }

    pub fn get_size(&self) -> SizeF {
        unsafe { SizeF(self.rt().GetSize()) }
    }

    pub fn create_solid_color_brush<C: Into<ColorF>>(
        &self,
        color: C,
        props: &BrushProperties,
    ) -> Result<brush::SolidColor, D2D1Error> {
        unsafe {
            let mut ptr = ComPtr::<ID2D1SolidColorBrush>::new();
            let result = self.rt()
                .CreateSolidColorBrush(&color.into().0, &props.0, ptr.raw_addr());

            if SUCCEEDED(result) {
                Ok(FromRaw::from_raw(ptr.raw_value()))
            } else {
                Err(From::from(result))
            }
        }
    }

    pub fn begin_draw(&mut self) {
        unsafe {
            self.rt().BeginDraw();
        }
    }

    pub fn end_draw(&mut self) -> Result<(), (D2D1Error, Option<RenderTag>)> {
        let mut tag1 = 0;
        let mut tag2 = 0;
        unsafe {
            let result = self.rt().EndDraw(&mut tag1, &mut tag2);

            if SUCCEEDED(result) {
                Ok(())
            } else {
                let tag = RenderTarget::make_tag(tag1, tag2);
                Err((From::from(result), tag))
            }
        }
    }

    pub fn set_tag(&mut self, tag: RenderTag) {
        unsafe {
            let RenderTagRaw(tag1, tag2) = mem::transmute(tag);
            self.rt().SetTags(tag1 as u64, tag2 as u64)
        };
    }

    pub fn get_tag(&mut self) -> Option<RenderTag> {
        let mut tag1 = 0;
        let mut tag2 = 0;
        unsafe {
            self.rt().GetTags(&mut tag1, &mut tag2);
            RenderTarget::make_tag(tag1, tag2)
        }
    }

    pub fn flush(&mut self) -> Result<(), (D2D1Error, Option<RenderTag>)> {
        let mut tag1 = 0;
        let mut tag2 = 0;
        unsafe {
            let result = self.rt().Flush(&mut tag1, &mut tag2);

            if SUCCEEDED(result) {
                Ok(())
            } else {
                let tag = RenderTarget::make_tag(tag1, tag2);
                Err((From::from(result), tag))
            }
        }
    }

    pub fn clear(&mut self, color: &ColorF) {
        unsafe {
            self.rt().Clear(&color.0);
        }
    }

    pub fn draw_line<B: Brush>(
        &mut self,
        p0: &Point2F,
        p1: &Point2F,
        brush: &B,
        stroke_width: f32,
        stroke_style: Option<&StrokeStyle>,
    ) {
        unsafe {
            let stroke_style = match stroke_style {
                Some(s) => s.get_ptr() as *mut _,
                None => ptr::null_mut(),
            };

            self.rt()
                .DrawLine(p0.0, p1.0, brush.get_ptr(), stroke_width, stroke_style)
        }
    }

    pub fn draw_rectangle<B: Brush>(
        &mut self,
        rect: &RectF,
        brush: &B,
        stroke_width: f32,
        stroke_style: Option<&StrokeStyle>,
    ) {
        unsafe {
            let stroke_style = match stroke_style {
                Some(s) => s.get_ptr() as *mut _,
                None => ptr::null_mut(),
            };

            self.rt()
                .DrawRectangle(&rect.0, brush.get_ptr(), stroke_width, stroke_style);
        }
    }

    pub fn fill_rectangle<B: Brush>(&mut self, rect: &RectF, brush: &B) {
        unsafe {
            self.rt().FillRectangle(&rect.0, brush.get_ptr());
        }
    }

    pub fn draw_rounded_rectangle<B: Brush>(
        &mut self,
        rect: &RoundedRect,
        brush: &B,
        stroke_width: f32,
        stroke_style: Option<&StrokeStyle>,
    ) {
        unsafe {
            let stroke_style = match stroke_style {
                Some(s) => s.get_ptr() as *mut _,
                None => ptr::null_mut(),
            };

            self.rt()
                .DrawRoundedRectangle(&rect.0, brush.get_ptr(), stroke_width, stroke_style);
        }
    }

    pub fn fill_rounded_rectangle<B: Brush>(&mut self, rect: &RoundedRect, brush: &B) {
        unsafe {
            self.rt().FillRoundedRectangle(&rect.0, brush.get_ptr());
        }
    }

    pub fn draw_ellipse<B: Brush>(
        &mut self,
        ellipse: &Ellipse,
        brush: &B,
        stroke_width: f32,
        stroke_style: Option<&StrokeStyle>,
    ) {
        unsafe {
            let stroke_style = match stroke_style {
                Some(s) => s.get_ptr() as *mut _,
                None => ptr::null_mut(),
            };

            self.rt()
                .DrawEllipse(&ellipse.0, brush.get_ptr(), stroke_width, stroke_style);
        }
    }

    pub fn fill_ellipse<B: Brush>(&mut self, ellipse: &Ellipse, brush: &B) {
        unsafe {
            self.rt().FillEllipse(&ellipse.0, brush.get_ptr());
        }
    }

    pub fn draw_geometry<G: Geometry, B: Brush>(
        &mut self,
        geometry: &G,
        brush: &B,
        stroke_width: f32,
        stroke_style: Option<&StrokeStyle>,
    ) {
        unsafe {
            let stroke_style = match stroke_style {
                Some(s) => s.get_ptr() as *mut _,
                None => ptr::null_mut(),
            };

            self.rt().DrawGeometry(
                geometry.get_ptr(),
                brush.get_ptr(),
                stroke_width,
                stroke_style,
            );
        }
    }

    pub fn fill_geometry<G: Geometry, B: Brush>(&mut self, geometry: &G, brush: &B) {
        unsafe {
            self.rt()
                .FillGeometry(geometry.get_ptr(), brush.get_ptr(), ptr::null_mut());
        }
    }

    pub fn fill_geometry_with_opacity<G: Geometry, B: Brush, OB: Brush>(
        &mut self,
        geometry: &G,
        brush: &B,
        opacity_brush: &OB,
    ) {
        unsafe {
            self.rt()
                .FillGeometry(geometry.get_ptr(), brush.get_ptr(), opacity_brush.get_ptr());
        }
    }

    pub fn draw_text<B: Brush>(
        &mut self,
        text: &str,
        format: &TextFormat,
        layout_rect: &RectF,
        foreground_brush: &B,
        options: &[DrawTextOption],
    ) {
        let text = text.to_wide_null();
        let mut draw_options = D2D1_DRAW_TEXT_OPTIONS_NONE;
        for &option in options {
            draw_options |= option as u32;
        }

        unsafe {
            let format = format.get_raw();
            self.rt().DrawText(
                text.as_ptr(),
                text.len() as u32,
                format,
                &layout_rect.0,
                foreground_brush.get_ptr(),
                (draw_options),
                DWRITE_MEASURING_MODE_NATURAL,
            );
        }
    }

    pub fn draw_text_layout<B: Brush>(
        &mut self,
        origin: &Point2F,
        layout: &TextLayout,
        brush: &B,
        options: &[DrawTextOption],
    ) {
        let mut draw_options = D2D1_DRAW_TEXT_OPTIONS_NONE;
        for &option in options {
            draw_options |= option as u32;
        }

        unsafe {
            let layout = layout.get_raw();
            self.rt()
                .DrawTextLayout(origin.0, layout, brush.get_ptr(), (draw_options));
        }
    }

    pub fn set_transform(&mut self, transform: &Matrix3x2F) {
        unsafe {
            self.rt().SetTransform(&transform.0)
        }
    }

    pub fn get_transform(&self) -> Matrix3x2F {
        unsafe {
            let mut mat: Matrix3x2F = mem::uninitialized();
            self.rt().GetTransform(&mut mat.0);
            mat
        }
    }

}

impl GetRaw for RenderTarget {
    type Raw = ID2D1RenderTarget;
    unsafe fn get_raw(&self) -> *mut ID2D1RenderTarget {
        self.ptr.raw_value()
    }
}

impl FromRaw for RenderTarget {
    type Raw = ID2D1RenderTarget;
    unsafe fn from_raw(raw: *mut ID2D1RenderTarget) -> Self {
        RenderTarget {
            ptr: ComPtr::attach(raw),
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum DrawTextOption {
    NoSnap = 1,
    Clip = 2,
    EnableColorFont = 4,
}
<|endoftext|>


<commit_before>//! Atom-class implementation.

use std::ops::{BitXor, Sub, Range, RangeInclusive};
use std::iter::FromIterator;
use std::fmt::{self, Debug};

use regex_syntax;
use itertools::Itertools;

use super::{Atom, ByteOrChar, Element, Repetition, Sequence, Union};
use super::{Distance, Step};
use util::set::{self, Contains};

// ----------------------------------------------------------------
// Class

/// Any member of a class of atoms
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Eq, Ord)]
pub enum ClassMember<T: Atom> {
    /// A single atom
    Atom(T),

    /// A range of atoms
    Range(T, T)
}

impl<T: Atom> ClassMember<T> {
    /// Map the class-member's atoms to a different type
    pub fn map_atoms<U, F>(self, f: F) -> ClassMember<U>
        where F: Fn(T) -> U,
              U: Atom
    {
        match self {
            ClassMember::Atom(a) => ClassMember::Atom(f(a)),
            ClassMember::Range(a, b) => ClassMember::Range(f(a), f(b))
        }
    }

    /// Get the number of individual atoms that this class member represents.
    pub fn len(&self) -> usize
        where T: Distance
    {
        match self {
            &ClassMember::Atom(_) => 1,
            &ClassMember::Range(first, last) => last.distance(&first) + 1,
        }
    }

    /// Fetch an iterator over each discrete value contained by this member.
    #[inline]
     pub fn iter(&self) -> impl Iterator<Item=T>
        where T: Step {
        match self {
            &ClassMember::Atom(a) => AtomIter(a...a),
            &ClassMember::Range(a, b) => AtomIter(a...b),
        }
    }
}

impl<T: Atom> set::Contains<T> for ClassMember<T> {
    /// Check if the member is or contains a particular atom.
    fn contains(&self, x: T) -> bool {
        match self {
            &ClassMember::Atom(a) => x == a,
            &ClassMember::Range(a, b) => (a...b).contains(x),
        }
    }
}

/// Iterator over the atoms specified by a class member.
struct AtomIter<T: Atom>(RangeInclusive<T>);

impl<T: Atom> Iterator for AtomIter<T>
        where T: Step
{
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        match self.0 {
            RangeInclusive::NonEmpty{start, end} => {
                let next_start = start.increment();
                if next_start > end {
                    self.0 = RangeInclusive::Empty{at: end};
                } else {
                    self.0 = RangeInclusive::NonEmpty{start: next_start, end: end};
                }
                Some(start) },
            RangeInclusive::Empty{..} => None
        }
    }
}

/*impl<T: Atom> Debug for ClassMember<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ClassMember::Atom(a) => write!(f, "{:?}", a),
            ClassMember::Range(first, last) => write!(f, "{:?}-{:?}", first, last)
        }
    }
}*/

impl<T: Atom> From<T> for ClassMember<T> {
    fn from(a: T) -> Self { ClassMember::Atom(a) }
}

impl From<regex_syntax::ClassRange> for ClassMember<ByteOrChar> {
    fn from(cr: regex_syntax::ClassRange) -> Self {
        ClassMember::Range(cr.start.into(), cr.end.into())
    }
}

impl From<regex_syntax::ByteRange> for ClassMember<ByteOrChar> {
    fn from(cr: regex_syntax::ByteRange) -> Self {
        ClassMember::Range(cr.start.into(), cr.end.into())
    }
}

impl From<regex_syntax::ClassRange> for ClassMember<char> {
    fn from(cr: regex_syntax::ClassRange) -> Self {
        ClassMember::Range(cr.start, cr.end)
    }
}

impl From<regex_syntax::ByteRange> for ClassMember<u8> {
    fn from(cr: regex_syntax::ByteRange) -> Self {
        ClassMember::Range(cr.start, cr.end)
    }
}

impl<T: Atom> From<Range<T>> for ClassMember<T>
    where T: Sub<usize,Output=T> + Sub<T,Output=usize>
{
    fn from(r: Range<T>) -> Self {
        if r.end - r.start > 1 { ClassMember::Range(r.start, r.end - 1usize) }
        else { ClassMember::Atom(r.start) }
    }
}
// ----------------------------------------------------------------

/// A set of atoms and/or ranges of atoms.
#[derive(Clone, PartialEq, PartialOrd, Eq, Ord)]
pub struct Class<T: Atom> {
    polarity: Polarity,
    members: Vec<ClassMember<T>>
}

// FIXME: [optimize] All operations on a Class<T>, including observer and
// predicate methods, currently require at least O(N) time.  We could probably
// improve this by changing the storage representation.
impl<T: Atom> Class<T> {
    /// Create a new Class instance with specified polarity and members taken
    /// from the given iterator.
    pub fn new<I,U>(p: Polarity, items: I) -> Self
        where I: IntoIterator<Item=U>,
              ClassMember<T>: From<U> {
        Class{polarity: p, members: Vec::from_iter(items.into_iter().map(|x| x.into()))}
    }

    /// Map the atoms in the class to a different atom type.
    pub fn map_atoms<U, F>(self, f: F) -> Class<U>
        where F: Fn(T) -> U,
              U: Atom
    {
        Class::new(self.polarity, self.members.into_iter().map(|m| m.map_atoms(&f)))
    }

    /// Check whether the class would match a particular atom.
    pub fn matches(&self, x: T) -> bool {
        self.has_member(x).bitxor(self.polarity == Polarity::NORMAL)
    }

    /// Check whether a particular atom is a member of this class.
    /// Unlike `matches`, this method does not take the class's polarity
    /// into account.
    pub fn has_member(&self, x: T) -> bool {
        self.members.iter().any(|&range| range.contains(x))
    }

    /// Get the number of members (atoms) in the class.
    ///
    /// Ranges of atoms count as the number of atoms in each range.
    // FIXME: [bug] handle inverted polarity?
    pub fn len(&self) -> usize {
        self.members.iter().map(|m| m.len()).sum()
    }

    /// Get the class's polarity.
    #[inline]
    pub fn polarity(&self) -> Polarity {
        self.polarity
    }

    /// Check if the class contains zero members.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.members.is_empty()
    }

    /// Fetch an iterator over the members specified in the class.
    ///
    /// Note that if the class has inverted polarity, the atoms in these
    /// members will be those which are **not** matched by the class.
    pub fn iter_members<'a>(&'a self) -> impl Iterator<Item=&'a ClassMember<T>> {
        self.members.iter()
    }
    }
}

impl<T: Atom> Debug for Class<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.polarity == Polarity::NORMAL {
            write!(f, "[{}]", self.members.iter().map(|m| format!("{:?}", m)).join(","))
        } else {
            write!(f, "[^{}]", self.members.iter().map(|m| format!("{:?}", m)).join(","))
        }
    }
}

impl<'a, T: 'a + Atom> set::Contains<&'a ClassMember<T>> for Class<T> {
    fn contains(&self, cm: &'a ClassMember<T>) -> bool {
        self.iter_members().any(|m| cm == m)
    }
}
impl<T: Atom> set::Contains<T> for Class<T> {
    fn contains(&self, x: T) -> bool {
        self.matches(x)
    }
}

// FIXME [optimize] I suspect that we need some sort of specialized data
// structure or clever algorithm in order to perform this test in a time better
// than O(N·M).
impl<T: Atom> set::IsSubsetOf<Class<T>> for Class<T> {
    fn is_subset_of(&self, other: &Self) -> bool {
        self.iter_members().flat_map(|m| m.iter()).all(|a| other.contains(a))
    }
}
impl<T: Atom> set::IsSubsetOf<Sequence<T>> for Class<T> {
    /// Any single-atom pattern like Class cannot match the same set of inputs
    /// as a Sequence, because a Sequence is guaranteed to contain at least
    /// two elements.
    #[inline]
    fn is_subset_of(&self, _: &Sequence<T>) -> bool {
        false
    }
}

impl<T: Atom> set::IsSubsetOf<Union<T>> for Class<T> {
    /// A class is a subset of a union if all members of the class are subsets
    /// of the union.
    fn is_subset_of(&self, _: &Union<T>) -> bool {
        unimplemented!()
    }
}

impl<T: Atom> set::IsSubsetOf<Repetition<T>> for Class<T> {
    fn is_subset_of(&self, _: &Repetition<T>) -> bool {
        unimplemented!()
    }
}

impl<T: Atom> set::IsSubsetOf<Element<T>> for Class<T> {
    fn is_subset_of(&self, other: &Element<T>) -> bool {
        match other {
            // Anything we can do, Wildcard can do betteeeeer...
            &Element::Wildcard => true,

            // A class *can* be a subset of an atom -- as long as that atom is
            // the only member of the class.
            &Element::Atom(a) => self.contains(a) && self.polarity == Polarity::NORMAL && self.len() == 1,
            &Element::Class(ref d) => self.is_subset_of(d),
            &Element::Tagged{ref element, ..} => self.is_subset_of(&**element),
            &Element::Sequence(ref seq) => self.is_subset_of(seq),
            &Element::Union(ref union) => self.is_subset_of(union),
            &Element::Repeat(ref rep) => self.is_subset_of(rep),
            &Element::Anchor(_) => false,
        }
    }
}

impl<T: Atom, U> FromIterator<U> for Class<T>
where ClassMember<T>: From<U> {
    fn from_iter<I>(iter: I) -> Self
        where I: IntoIterator<Item=U> {
        Class{polarity: Polarity::NORMAL, members: Vec::from_iter(iter.into_iter().map(|x| x.into()))}
    }
}

macro_rules! impl_class_from {
    ($T: ty, $from: ty) => {
        impl<'a> From<$from> for Class<$T> {
            fn from(c: $from) -> Self {
                Self::new(Polarity::NORMAL,
                          c.into_iter().map(|cr| if cr.end != cr.start { ClassMember::Range(cr.start.into(), cr.end.into()) }
                                            else { ClassMember::Atom(cr.start.into()) }))
            }
        }
    };
}

impl_class_from!(char, regex_syntax::CharClass);
impl_class_from!(u8, regex_syntax::ByteClass);
impl_class_from!(ByteOrChar, regex_syntax::CharClass);
impl_class_from!(ByteOrChar, regex_syntax::ByteClass);

// ----------------------------------------------------------------

/** Flag type used to differentiate between normal and inverting matches on
 * pattern elements. */
#[derive(Debug,Copy,Clone,PartialEq, PartialOrd, Eq, Ord)]
pub enum Polarity {
    /// Normal match behavior.
    NORMAL,
    /// Invert the match result.
    INVERTED
}
<commit_msg>pattern/class: fix polarity check in `matches`<commit_after>//! Atom-class implementation.

use std::ops::{BitXor, Sub, Range, RangeInclusive};
use std::iter::FromIterator;
use std::fmt::{self, Debug};

use regex_syntax;
use itertools::Itertools;

use super::{Atom, ByteOrChar, Element, Repetition, Sequence, Union};
use super::{Distance, Step};
use util::set::{self, Contains};

// ----------------------------------------------------------------
// Class

/// Any member of a class of atoms
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Eq, Ord)]
pub enum ClassMember<T: Atom> {
    /// A single atom
    Atom(T),

    /// A range of atoms
    Range(T, T)
}

impl<T: Atom> ClassMember<T> {
    /// Map the class-member's atoms to a different type
    pub fn map_atoms<U, F>(self, f: F) -> ClassMember<U>
        where F: Fn(T) -> U,
              U: Atom
    {
        match self {
            ClassMember::Atom(a) => ClassMember::Atom(f(a)),
            ClassMember::Range(a, b) => ClassMember::Range(f(a), f(b))
        }
    }

    /// Get the number of individual atoms that this class member represents.
    pub fn len(&self) -> usize
        where T: Distance
    {
        match self {
            &ClassMember::Atom(_) => 1,
            &ClassMember::Range(first, last) => last.distance(&first) + 1,
        }
    }

    /// Fetch an iterator over each discrete value contained by this member.
    #[inline]
     pub fn iter(&self) -> impl Iterator<Item=T>
        where T: Step {
        match self {
            &ClassMember::Atom(a) => AtomIter(a...a),
            &ClassMember::Range(a, b) => AtomIter(a...b),
        }
    }
}

impl<T: Atom> set::Contains<T> for ClassMember<T> {
    /// Check if the member is or contains a particular atom.
    fn contains(&self, x: T) -> bool {
        match self {
            &ClassMember::Atom(a) => x == a,
            &ClassMember::Range(a, b) => (a...b).contains(x),
        }
    }
}

/// Iterator over the atoms specified by a class member.
struct AtomIter<T: Atom>(RangeInclusive<T>);

impl<T: Atom> Iterator for AtomIter<T>
        where T: Step
{
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        match self.0 {
            RangeInclusive::NonEmpty{start, end} => {
                let next_start = start.increment();
                if next_start > end {
                    self.0 = RangeInclusive::Empty{at: end};
                } else {
                    self.0 = RangeInclusive::NonEmpty{start: next_start, end: end};
                }
                Some(start) },
            RangeInclusive::Empty{..} => None
        }
    }
}

/*impl<T: Atom> Debug for ClassMember<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ClassMember::Atom(a) => write!(f, "{:?}", a),
            ClassMember::Range(first, last) => write!(f, "{:?}-{:?}", first, last)
        }
    }
}*/

impl<T: Atom> From<T> for ClassMember<T> {
    fn from(a: T) -> Self { ClassMember::Atom(a) }
}

impl From<regex_syntax::ClassRange> for ClassMember<ByteOrChar> {
    fn from(cr: regex_syntax::ClassRange) -> Self {
        ClassMember::Range(cr.start.into(), cr.end.into())
    }
}

impl From<regex_syntax::ByteRange> for ClassMember<ByteOrChar> {
    fn from(cr: regex_syntax::ByteRange) -> Self {
        ClassMember::Range(cr.start.into(), cr.end.into())
    }
}

impl From<regex_syntax::ClassRange> for ClassMember<char> {
    fn from(cr: regex_syntax::ClassRange) -> Self {
        ClassMember::Range(cr.start, cr.end)
    }
}

impl From<regex_syntax::ByteRange> for ClassMember<u8> {
    fn from(cr: regex_syntax::ByteRange) -> Self {
        ClassMember::Range(cr.start, cr.end)
    }
}

impl<T: Atom> From<Range<T>> for ClassMember<T>
    where T: Sub<usize,Output=T> + Sub<T,Output=usize>
{
    fn from(r: Range<T>) -> Self {
        if r.end - r.start > 1 { ClassMember::Range(r.start, r.end - 1usize) }
        else { ClassMember::Atom(r.start) }
    }
}
// ----------------------------------------------------------------

/// A set of atoms and/or ranges of atoms.
#[derive(Clone, PartialEq, PartialOrd, Eq, Ord)]
pub struct Class<T: Atom> {
    polarity: Polarity,
    members: Vec<ClassMember<T>>
}

// FIXME: [optimize] All operations on a Class<T>, including observer and
// predicate methods, currently require at least O(N) time.  We could probably
// improve this by changing the storage representation.
impl<T: Atom> Class<T> {
    /// Create a new Class instance with specified polarity and members taken
    /// from the given iterator.
    pub fn new<I,U>(p: Polarity, items: I) -> Self
        where I: IntoIterator<Item=U>,
              ClassMember<T>: From<U> {
        Class{polarity: p, members: Vec::from_iter(items.into_iter().map(|x| x.into()))}
    }

    /// Map the atoms in the class to a different atom type.
    pub fn map_atoms<U, F>(self, f: F) -> Class<U>
        where F: Fn(T) -> U,
              U: Atom
    {
        Class::new(self.polarity, self.members.into_iter().map(|m| m.map_atoms(&f)))
    }

    /// Check whether the class would match a particular atom.
    pub fn matches(&self, x: T) -> bool {
        self.has_member(x).bitxor(self.polarity == Polarity::INVERTED)
    }

    /// Check whether a particular atom is a member of this class.
    /// Unlike `matches`, this method does not take the class's polarity
    /// into account.
    pub fn has_member(&self, x: T) -> bool {
        self.members.iter().any(|&range| range.contains(x))
    }

    /// Get the number of members (atoms) in the class.
    ///
    /// Ranges of atoms count as the number of atoms in each range.
    // FIXME: [bug] handle inverted polarity?
    pub fn len(&self) -> usize {
        self.members.iter().map(|m| m.len()).sum()
    }

    /// Get the class's polarity.
    #[inline]
    pub fn polarity(&self) -> Polarity {
        self.polarity
    }

    /// Check if the class contains zero members.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.members.is_empty()
    }

    /// Fetch an iterator over the members specified in the class.
    ///
    /// Note that if the class has inverted polarity, the atoms in these
    /// members will be those which are **not** matched by the class.
    pub fn iter_members<'a>(&'a self) -> impl Iterator<Item=&'a ClassMember<T>> {
        self.members.iter()
    }
    }
}

impl<T: Atom> Debug for Class<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.polarity == Polarity::NORMAL {
            write!(f, "[{}]", self.members.iter().map(|m| format!("{:?}", m)).join(","))
        } else {
            write!(f, "[^{}]", self.members.iter().map(|m| format!("{:?}", m)).join(","))
        }
    }
}

impl<'a, T: 'a + Atom> set::Contains<&'a ClassMember<T>> for Class<T> {
    fn contains(&self, cm: &'a ClassMember<T>) -> bool {
        self.iter_members().any(|m| cm == m)
    }
}
impl<T: Atom> set::Contains<T> for Class<T> {
    fn contains(&self, x: T) -> bool {
        self.matches(x)
    }
}

// FIXME [optimize] I suspect that we need some sort of specialized data
// structure or clever algorithm in order to perform this test in a time better
// than O(N·M).
impl<T: Atom> set::IsSubsetOf<Class<T>> for Class<T> {
    fn is_subset_of(&self, other: &Self) -> bool {
        self.iter_members().flat_map(|m| m.iter()).all(|a| other.contains(a))
    }
}
impl<T: Atom> set::IsSubsetOf<Sequence<T>> for Class<T> {
    /// Any single-atom pattern like Class cannot match the same set of inputs
    /// as a Sequence, because a Sequence is guaranteed to contain at least
    /// two elements.
    #[inline]
    fn is_subset_of(&self, _: &Sequence<T>) -> bool {
        false
    }
}

impl<T: Atom> set::IsSubsetOf<Union<T>> for Class<T> {
    /// A class is a subset of a union if all members of the class are subsets
    /// of the union.
    fn is_subset_of(&self, _: &Union<T>) -> bool {
        unimplemented!()
    }
}

impl<T: Atom> set::IsSubsetOf<Repetition<T>> for Class<T> {
    fn is_subset_of(&self, _: &Repetition<T>) -> bool {
        unimplemented!()
    }
}

impl<T: Atom> set::IsSubsetOf<Element<T>> for Class<T> {
    fn is_subset_of(&self, other: &Element<T>) -> bool {
        match other {
            // Anything we can do, Wildcard can do betteeeeer...
            &Element::Wildcard => true,

            // A class *can* be a subset of an atom -- as long as that atom is
            // the only member of the class.
            &Element::Atom(a) => self.contains(a) && self.polarity == Polarity::NORMAL && self.len() == 1,
            &Element::Class(ref d) => self.is_subset_of(d),
            &Element::Tagged{ref element, ..} => self.is_subset_of(&**element),
            &Element::Sequence(ref seq) => self.is_subset_of(seq),
            &Element::Union(ref union) => self.is_subset_of(union),
            &Element::Repeat(ref rep) => self.is_subset_of(rep),
            &Element::Anchor(_) => false,
        }
    }
}

impl<T: Atom, U> FromIterator<U> for Class<T>
where ClassMember<T>: From<U> {
    fn from_iter<I>(iter: I) -> Self
        where I: IntoIterator<Item=U> {
        Class{polarity: Polarity::NORMAL, members: Vec::from_iter(iter.into_iter().map(|x| x.into()))}
    }
}

macro_rules! impl_class_from {
    ($T: ty, $from: ty) => {
        impl<'a> From<$from> for Class<$T> {
            fn from(c: $from) -> Self {
                Self::new(Polarity::NORMAL,
                          c.into_iter().map(|cr| if cr.end != cr.start { ClassMember::Range(cr.start.into(), cr.end.into()) }
                                            else { ClassMember::Atom(cr.start.into()) }))
            }
        }
    };
}

impl_class_from!(char, regex_syntax::CharClass);
impl_class_from!(u8, regex_syntax::ByteClass);
impl_class_from!(ByteOrChar, regex_syntax::CharClass);
impl_class_from!(ByteOrChar, regex_syntax::ByteClass);

// ----------------------------------------------------------------

/** Flag type used to differentiate between normal and inverting matches on
 * pattern elements. */
#[derive(Debug,Copy,Clone,PartialEq, PartialOrd, Eq, Ord)]
pub enum Polarity {
    /// Normal match behavior.
    NORMAL,
    /// Invert the match result.
    INVERTED
}
<|endoftext|>


<commit_before>use num_::Integer;

fn wrapping_pow(mut base: u64, mut exp: u32) -> u64 {
    let mut acc: u64 = 1;
    while exp > 0 {
        if exp % 2 == 1 {
            acc = acc.wrapping_mul(base)
        }
        base = base.wrapping_mul(base);
        exp /= 2;
    }
    acc
}

/// Returns integers `(y, k)` such that `x = y^k` with `k` maximised
/// (other than for `x = 0, 1`, in which case `y = x`, `k = 1`).
///
/// # Examples
///
/// ```rust
/// assert_eq!(slow_primes::as_perfect_power(2), (2, 1));
/// assert_eq!(slow_primes::as_perfect_power(4), (2, 2));
/// assert_eq!(slow_primes::as_perfect_power(8), (2, 3));
/// assert_eq!(slow_primes::as_perfect_power(1024), (2, 10));
///
/// assert_eq!(slow_primes::as_perfect_power(1000), (10, 3));
///
/// assert_eq!(slow_primes::as_perfect_power(15), (15, 1));
/// ```
pub fn as_perfect_power(x: u64) -> (u64, u8) {
    if x == 0 || x == 1 {
        return (x, 1)
    }

    let floor_log_2 = 64 - x.leading_zeros() as u32 - 1;

    let x_ = x as f64;
    let mut last = (x, 1);
    // TODO: we could be smarter about this: we know all the possible
    // primes that can divide the exponent (since we have a list up to
    // 251 >= 64), so we really only need to check them.
    let mut expn: u32 = 2;
    let mut step = 1;
    while expn <= floor_log_2 {
        let factor = x_.powf(1.0/expn as f64).round() as u64;
        // the only case this will wrap is if x is close to 2^64 and
        // the round() rounds up, pushing this calculation over the
        // edge, however, the overflow will be well away from x, so we
        // still correctly don't take this branch. (x can't be a
        // perfect power if the result rounds away.)
        if wrapping_pow(factor, expn) == x {
            last = (factor, expn as u8);
            // if x is a 2nd and 5th power, it's going to be a 10th
            // power too, meaning we can search faster.
            // TODO: check if this is actually saving work
            step = step.lcm(&expn);
        }

        expn += step;
    }
    last
}

/// Return `Some((p, k))` if `x = p^k` for some prime `p` and `k >= 1`
/// (that is, including when `x` is itself a prime).
///
/// Returns `None` if `x` not a perfect power.
pub fn as_prime_power(x: u64) -> Option<(u64, u8)> {
    let (y, k) = as_perfect_power(x);
    if ::is_prime_miller_rabin(y) {
        Some((y, k))
    } else {
        None
    }
}

#[cfg(test)]
mod tests {
    use std::num::{Float, Int};
    use Primes;

    use super::{as_perfect_power, as_prime_power};

    #[test]
    fn perfect_and_prime_power() {
        let tests = [
            (0, (0, 1), false),
            (1, (1, 1), false),
            (2, (2, 1), true),
            (3, (3, 1), true),
            (4, (2, 2), true),
            (5, (5, 1), true),
            (6, (6, 1), false),
            (8, (2, 3), true),
            (9, (3, 2), true),
            (16, (2, 4), true),
            (25, (5, 2), true),
            (32, (2, 5), true),
            (36, (6, 2), false),
            (100, (10, 2), false),
            (1000, (10, 3), false),
                ];

        for &(x, expected, is_prime) in tests.iter() {
            assert_eq!(as_perfect_power(x), expected);
            assert_eq!(as_prime_power(x),
                       if is_prime { Some((expected))} else { None })
        }

        let sieve = Primes::sieve(200);
        let mut primes = sieve.primes();
        const MAX: f64 = 0xFFFF_FFFF_FFFF_FFFFu64 as f64;
        // test a whole pile of (semi)primes
        loop {
            let p = match primes.next() {
                Some(p) => p as u64,
                None => break
            };

            let subprimes = primes.clone().map(|x| (x, false));
            // include 1 to test p itself.
            for (q, is_prime) in Some((1, true)).into_iter().chain(subprimes) {
                let pq = p * q as u64;
                for n in (1..MAX.log(pq as f64) as u32) {
                    let x = pq.pow(n);

                    let expected = (pq, n as u8);
                    assert_eq!(as_perfect_power(x), expected);
                    assert_eq!(as_prime_power(x),
                               if is_prime { Some(expected) } else { None });
                }
            }
        }
    }
}
<commit_msg>Update to (90cc83015 2015-04-22) (built 2015-04-23).<commit_after>use num_::Integer;

fn wrapping_pow(mut base: u64, mut exp: u32) -> u64 {
    let mut acc: u64 = 1;
    while exp > 0 {
        if exp % 2 == 1 {
            acc = acc.wrapping_mul(base)
        }
        base = base.wrapping_mul(base);
        exp /= 2;
    }
    acc
}

/// Returns integers `(y, k)` such that `x = y^k` with `k` maximised
/// (other than for `x = 0, 1`, in which case `y = x`, `k = 1`).
///
/// # Examples
///
/// ```rust
/// assert_eq!(slow_primes::as_perfect_power(2), (2, 1));
/// assert_eq!(slow_primes::as_perfect_power(4), (2, 2));
/// assert_eq!(slow_primes::as_perfect_power(8), (2, 3));
/// assert_eq!(slow_primes::as_perfect_power(1024), (2, 10));
///
/// assert_eq!(slow_primes::as_perfect_power(1000), (10, 3));
///
/// assert_eq!(slow_primes::as_perfect_power(15), (15, 1));
/// ```
pub fn as_perfect_power(x: u64) -> (u64, u8) {
    if x == 0 || x == 1 {
        return (x, 1)
    }

    let floor_log_2 = 64 - x.leading_zeros() as u32 - 1;

    let x_ = x as f64;
    let mut last = (x, 1);
    // TODO: we could be smarter about this: we know all the possible
    // primes that can divide the exponent (since we have a list up to
    // 251 >= 64), so we really only need to check them.
    let mut expn: u32 = 2;
    let mut step = 1;
    while expn <= floor_log_2 {
        let factor = x_.powf(1.0/expn as f64).round() as u64;
        // the only case this will wrap is if x is close to 2^64 and
        // the round() rounds up, pushing this calculation over the
        // edge, however, the overflow will be well away from x, so we
        // still correctly don't take this branch. (x can't be a
        // perfect power if the result rounds away.)
        if wrapping_pow(factor, expn) == x {
            last = (factor, expn as u8);
            // if x is a 2nd and 5th power, it's going to be a 10th
            // power too, meaning we can search faster.
            // TODO: check if this is actually saving work
            step = step.lcm(&expn);
        }

        expn += step;
    }
    last
}

/// Return `Some((p, k))` if `x = p^k` for some prime `p` and `k >= 1`
/// (that is, including when `x` is itself a prime).
///
/// Returns `None` if `x` not a perfect power.
pub fn as_prime_power(x: u64) -> Option<(u64, u8)> {
    let (y, k) = as_perfect_power(x);
    if ::is_prime_miller_rabin(y) {
        Some((y, k))
    } else {
        None
    }
}

#[cfg(test)]
mod tests {
    use Primes;

    use super::{as_perfect_power, as_prime_power};

    #[test]
    fn perfect_and_prime_power() {
        let tests = [
            (0, (0, 1), false),
            (1, (1, 1), false),
            (2, (2, 1), true),
            (3, (3, 1), true),
            (4, (2, 2), true),
            (5, (5, 1), true),
            (6, (6, 1), false),
            (8, (2, 3), true),
            (9, (3, 2), true),
            (16, (2, 4), true),
            (25, (5, 2), true),
            (32, (2, 5), true),
            (36, (6, 2), false),
            (100, (10, 2), false),
            (1000, (10, 3), false),
                ];

        for &(x, expected, is_prime) in tests.iter() {
            assert_eq!(as_perfect_power(x), expected);
            assert_eq!(as_prime_power(x),
                       if is_prime { Some((expected))} else { None })
        }

        let sieve = Primes::sieve(200);
        let mut primes = sieve.primes();
        const MAX: f64 = 0xFFFF_FFFF_FFFF_FFFFu64 as f64;
        // test a whole pile of (semi)primes
        loop {
            let p = match primes.next() {
                Some(p) => p as u64,
                None => break
            };

            let subprimes = primes.clone().map(|x| (x, false));
            // include 1 to test p itself.
            for (q, is_prime) in Some((1, true)).into_iter().chain(subprimes) {
                let pq = p * q as u64;
                for n in (1..MAX.log(pq as f64) as u32) {
                    let x = pq.pow(n);

                    let expected = (pq, n as u8);
                    assert_eq!(as_perfect_power(x), expected);
                    assert_eq!(as_prime_power(x),
                               if is_prime { Some(expected) } else { None });
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2013-2014 Jeffery Olson
//
// Licensed under the 3-Clause BSD License, see LICENSE.txt
// at the top-level of this repository.
// This file may not be copied, modified, or distributed
// except according to those terms.
#![crate_id="http://olsonjeffery.github.io#schemers:0.0.1"]
#![desc = "Simple Scheme Interpreter"]
#![license = "3-Clause BSD"]

use std::from_str::FromStr;

// this is starting out as a straight port of Peter Norvig's
// lis.py (the first iteration) to Rust

#[cfg(not(test))]
fn main() {
    use std::os::args;
    let args: Vec<~str> = args().move_iter().skip(1).collect();
    if args.len() == 0 {
        fail!("no program provided");
    }
    let program: ~str = args.move_iter()
        .fold(~"", |memo, arg| memo + arg + " ").trim().to_owned();
    let tokens = &mut tokenize(program);
    let parsed_items = parse(tokens);
    println!("{}", parsed_items.print());
}

// parser impl
fn pad_input(input: ~str) -> ~str {
    input.replace("(", " ( ").replace(")", " ) ")
}

fn tokenize(input: ~str) -> Vec<~str> {
    pad_input(input).split_str(" ")
        .filter(|i| *i != "").map(|i| i.to_owned()).collect()
}

#[deriving(Eq, Show)]
pub enum Expr {
    Atom(AtomType),
    List(Vec<~Expr>)
}
impl Expr {
    pub fn new_atom(input: ~str) -> Expr {
        let first_char = input.char_at(0);
        match first_char {
            '0' | '1' | '2' | '3' | '4' | '5' |
            '6' | '7' | '8' | '9' => {
                match input.contains(".") {
                    true => {
                        match FromStr::from_str(input) {
                            Some(val) => Atom(Float(val)),
                            None => fail!("Cannot parse number-like input of: {}", input)
                        }
                    },
                    false => {
                        match FromStr::from_str(input) {
                            Some(val) => Atom(Integer(val)),
                            None => fail!("Cannot parse number-like input of: {}", input)
                        }
                    }
                }
            },
            _ => Atom(Symbol(input))
        }
    }

    pub fn is_atom(&self) -> bool {
        match *self {
            Atom(_) => true,
            _ => false
        }
    }

    pub fn print(&self) -> ~str {
        match self {
            &List(ref items) => {
                let out = items.iter().map(|i| i.print())
                    .fold(~"(", |m, v| m + v + " ");
                out.trim() + ")"
            },
            &Atom(ref v) => v.print()
        }
    }
}

#[deriving(Eq, Show)]
pub enum AtomType {
    Symbol(~str),
    Integer(i64),
    Float(f64)
}

impl AtomType {
    pub fn print(&self) -> ~str {
        match self {
            &Symbol(ref v) => v.to_owned(),
            &Integer(ref v) => v.to_str(),
            &Float(ref v) => v.to_str()
        }
    }
}

fn parse(tokens: &mut Vec<~str>) -> Expr {
    let current_token =
        tokens.shift().expect("calling parse() w/ empty token list; shouldn't happen");
    match current_token {
        ref x if *x == ~"(" => {
            let mut list = Vec::new();
            while *tokens.get(0) != ~")" {
                list.push(box parse(tokens));
            }
            tokens.shift();
            List(list)
        },
        ref x if *x == ~")" => fail!("hit ) token; shouldn't happen"),
        x => Expr::new_atom(x)
    }
}

#[cfg(test)]
mod parser_test {
    use super::{pad_input, tokenize, parse, Atom, List, Symbol, Integer, Float};

    #[test]
    fn pad_input_should_insert_spaces_before_and_after_parens() {
        assert_eq!(pad_input(~"(x 1 2 3)"), ~" ( x 1 2 3 ) ");
    }

    #[test]
    fn tokenize_string_into_tokens() {
        let tokens = tokenize(~"(foo 12 3)");
        assert_eq!(tokens.len(), 5);
        assert_eq!(tokens.get(0).to_owned(), ~"(");
        assert_eq!(tokens.get(1).to_owned(), ~"foo");
        assert_eq!(tokens.get(2).to_owned(), ~"12");
        assert_eq!(tokens.get(3).to_owned(), ~"3");
        assert_eq!(tokens.get(4).to_owned(), ~")");
    }

    #[test]
    fn should_parse_tokens_consisting_of_a_str_atom_and_convert_it_to_a_parse_item() {
        let tokens = &mut tokenize(~"bar");
        let parsed_item = parse(tokens);
        println!("{:?}", parsed_item);
        match parsed_item {
            Atom(Symbol(val)) => {
                assert_eq!(val, ~"bar")
            },
            _ => assert!(false)
        }
    }
    #[test]
    fn should_parse_tokens_consisting_of_an_int_atom_and_convert_it_to_a_parse_item() {
        let tokens = &mut tokenize(~"42");
        let parsed_item = parse(tokens);
        println!("{:?}", parsed_item);
        match parsed_item {
            Atom(Integer(val)) => {
                assert_eq!(val, 42)
            },
            _ => assert!(false)
        }
    }
    #[test]
    fn should_parse_tokens_consisting_of_a_float_atom_and_convert_it_to_a_parse_item() {
        let tokens = &mut tokenize(~"1.1");
        let parsed_item = parse(tokens);
        println!("{:?}", parsed_item);
        match parsed_item {
            Atom(Float(val)) => {
                assert_eq!(val, 1.1)
            },
            _ => assert!(false)
        }
    }

    #[test]
    fn should_parse_tokens_consisting_of_a_list_of_atoms_and_parse_it() {
        let tokens = &mut tokenize(~"(bar 12 45)");
        let parsed_item = parse(tokens);
        assert_eq!(parsed_item.is_atom(), false);
        match parsed_item {
            List(items) => {
                let items_len = items.len();
                assert!(3 == items_len);
                assert_eq!(items.get(0).is_atom(), true);
                assert_eq!(*items.get(0), ~Atom(Symbol(~"bar")));
                assert_eq!(items.get(1).is_atom(), true);
                assert_eq!(*items.get(1), ~Atom(Integer(12)));
                assert_eq!(items.get(2).is_atom(), true);
                assert_eq!(*items.get(2), ~Atom(Integer(45)));
            },
            _ => fail!("got back an atom, it seems")
        }
    }
    #[test]
    fn should_parse_tokens_consisting_of_a_list_of_atoms_and_a_nested_list() {
        let tokens = &mut tokenize(~"((2 3) bar 12 (hee (hah)))");
        let parsed_item = parse(tokens);
        assert_eq!(parsed_item.is_atom(), false);
        match parsed_item {
            List(items) => {
                let items_len = items.len();
                assert!(4 == items_len);
                match *items.get(0) {
                    ~List(ref items) => {
                        assert_eq!(items.len(), 2);
                        assert_eq!(*items.get(0), ~Atom(Integer(2)));
                        assert_eq!(*items.get(1), ~Atom(Integer(3)));
                    },
                    _ => fail!("shoulda got a list")
                }
                assert_eq!(*items.get(1), ~Atom(Symbol(~"bar")));
                assert_eq!(*items.get(2), ~Atom(Integer(12)));
                match *items.get(3) {
                    ~List(ref items) => {
                        assert_eq!(items.len(), 2);
                        assert_eq!(*items.get(0), ~Atom(Symbol(~"hee")));
                        match *items.get(1) {
                            ~List(ref items) => {
                                assert_eq!(items.len(), 1);
                                assert_eq!(*items.get(0), ~Atom(Symbol(~"hah")));
                            },
                            _ => fail!("shoulda got a list")
                        }
                    },
                    _ => fail!("shoulda got a list")
                }
            },
            _ => fail!("got back an atom, it seems")
        }
    }
    #[test]
    fn parsed_tokens_of_a_list_of_atoms_and_a_nested_list_should_to_str_correctly() {
        let input = ~"((2 3) bar 12 (hee (hah)))";
        let tokens = &mut tokenize(input.to_owned());
        let parsed_item = parse(tokens);
        let output = parsed_item.print();
        assert_eq!(input, output);
    }
}
<commit_msg>impl+tests for parse_str and un_cons()<commit_after>// Copyright 2013-2014 Jeffery Olson
//
// Licensed under the 3-Clause BSD License, see LICENSE.txt
// at the top-level of this repository.
// This file may not be copied, modified, or distributed
// except according to those terms.
#![crate_id="http://olsonjeffery.github.io#schemers:0.0.1"]
#![desc = "Simple Scheme Interpreter"]
#![license = "3-Clause BSD"]

use std::from_str::FromStr;

// this is starting out as a straight port of Peter Norvig's
// lis.py (the first iteration) to Rust

#[cfg(not(test))]
fn main() {
    use std::os::args;
    let args: Vec<~str> = args().move_iter().skip(1).collect();
    if args.len() == 0 {
        fail!("no program provided");
    }
    let program: ~str = args.move_iter()
        .fold(~"", |memo, arg| memo + arg + " ").trim().to_owned();
    let tokens = &mut tokenize(program);
    let parsed_items = parse(tokens);
    println!("{}", parsed_items.print());
}

// parser impl
fn parse_str(input: ~str) -> Expr {
    parse(&mut tokenize(input))
}

fn pad_input(input: ~str) -> ~str {
    input.replace("(", " ( ").replace(")", " ) ")
}

fn tokenize(input: ~str) -> Vec<~str> {
    pad_input(input).split_str(" ")
        .filter(|i| *i != "").map(|i| i.to_owned()).collect()
}

#[deriving(Eq, Show)]
pub enum Expr {
    Atom(AtomType),
    List(Vec<~Expr>)
}
impl Expr {
    pub fn new_atom(input: ~str) -> Expr {
        let first_char = input.char_at(0);
        match first_char {
            '0' | '1' | '2' | '3' | '4' | '5' |
            '6' | '7' | '8' | '9' => {
                match input.contains(".") {
                    true => {
                        match FromStr::from_str(input) {
                            Some(val) => Atom(Float(val)),
                            None => fail!("Cannot parse number-like input of: {}", input)
                        }
                    },
                    false => {
                        match FromStr::from_str(input) {
                            Some(val) => Atom(Integer(val)),
                            None => fail!("Cannot parse number-like input of: {}", input)
                        }
                    }
                }
            },
            _ => Atom(Symbol(input))
        }
    }

    pub fn is_atom(&self) -> bool {
        match *self {
            Atom(_) => true,
            _ => false
        }
    }

    pub fn print(&self) -> ~str {
        match self {
            &List(ref items) => {
                let out = items.iter().map(|i| i.print())
                    .fold(~"(", |m, v| m + v + " ");
                out.trim() + ")"
            },
            &Atom(ref v) => v.print()
        }
    }
    pub fn un_cons(self) -> (Expr, Option<Expr>) {
        match self {
            Atom(v) => (Atom(v), None),
            List(mut items) => {
                if items.len() == 0 {
                    fail!("cannot build car/cdr of empty list");
                }
                let car = items.shift()
                    .expect("at least one item in the least; shouldn't happen");
                (*car, Some(List(items)))
            }
        }
    }
}

#[deriving(Eq, Show)]
pub enum AtomType {
    Symbol(~str),
    Integer(i64),
    Float(f64)
}

impl AtomType {
    pub fn print(&self) -> ~str {
        match self {
            &Symbol(ref v) => v.to_owned(),
            &Integer(ref v) => v.to_str(),
            &Float(ref v) => v.to_str()
        }
    }
}

fn parse(tokens: &mut Vec<~str>) -> Expr {
    let current_token =
        tokens.shift().expect("calling parse() w/ empty token list; shouldn't happen");
    match current_token {
        ref x if *x == ~"(" => {
            let mut list = Vec::new();
            while *tokens.get(0) != ~")" {
                list.push(box parse(tokens));
            }
            tokens.shift();
            List(list)
        },
        ref x if *x == ~")" => fail!("hit ) token; shouldn't happen"),
        x => Expr::new_atom(x)
    }
}

#[cfg(test)]
mod parser_test {
    use super::{pad_input, tokenize, parse, Atom, List, Symbol, Integer, Float};

    #[test]
    fn pad_input_should_insert_spaces_before_and_after_parens() {
        assert_eq!(pad_input(~"(x 1 2 3)"), ~" ( x 1 2 3 ) ");
    }

    #[test]
    fn tokenize_string_into_tokens() {
        let tokens = tokenize(~"(foo 12 3)");
        assert_eq!(tokens.len(), 5);
        assert_eq!(tokens.get(0).to_owned(), ~"(");
        assert_eq!(tokens.get(1).to_owned(), ~"foo");
        assert_eq!(tokens.get(2).to_owned(), ~"12");
        assert_eq!(tokens.get(3).to_owned(), ~"3");
        assert_eq!(tokens.get(4).to_owned(), ~")");
    }

    #[test]
    fn should_parse_tokens_consisting_of_a_str_atom_and_convert_it_to_a_parse_item() {
        let tokens = &mut tokenize(~"bar");
        let parsed_item = parse(tokens);
        println!("{:?}", parsed_item);
        match parsed_item {
            Atom(Symbol(val)) => {
                assert_eq!(val, ~"bar")
            },
            _ => assert!(false)
        }
    }
    #[test]
    fn should_parse_tokens_consisting_of_an_int_atom_and_convert_it_to_a_parse_item() {
        let tokens = &mut tokenize(~"42");
        let parsed_item = parse(tokens);
        println!("{:?}", parsed_item);
        match parsed_item {
            Atom(Integer(val)) => {
                assert_eq!(val, 42)
            },
            _ => assert!(false)
        }
    }
    #[test]
    fn should_parse_tokens_consisting_of_a_float_atom_and_convert_it_to_a_parse_item() {
        let tokens = &mut tokenize(~"1.1");
        let parsed_item = parse(tokens);
        println!("{:?}", parsed_item);
        match parsed_item {
            Atom(Float(val)) => {
                assert_eq!(val, 1.1)
            },
            _ => assert!(false)
        }
    }

    #[test]
    fn should_parse_tokens_consisting_of_a_list_of_atoms_and_parse_it() {
        let tokens = &mut tokenize(~"(bar 12 45)");
        let parsed_item = parse(tokens);
        assert_eq!(parsed_item.is_atom(), false);
        match parsed_item {
            List(items) => {
                let items_len = items.len();
                assert!(3 == items_len);
                assert_eq!(items.get(0).is_atom(), true);
                assert_eq!(*items.get(0), ~Atom(Symbol(~"bar")));
                assert_eq!(items.get(1).is_atom(), true);
                assert_eq!(*items.get(1), ~Atom(Integer(12)));
                assert_eq!(items.get(2).is_atom(), true);
                assert_eq!(*items.get(2), ~Atom(Integer(45)));
            },
            _ => fail!("got back an atom, it seems")
        }
    }
    #[test]
    fn should_parse_tokens_consisting_of_a_list_of_atoms_and_a_nested_list() {
        let tokens = &mut tokenize(~"((2 3) bar 12 (hee (hah)))");
        let parsed_item = parse(tokens);
        assert_eq!(parsed_item.is_atom(), false);
        match parsed_item {
            List(items) => {
                let items_len = items.len();
                assert!(4 == items_len);
                match *items.get(0) {
                    ~List(ref items) => {
                        assert_eq!(items.len(), 2);
                        assert_eq!(*items.get(0), ~Atom(Integer(2)));
                        assert_eq!(*items.get(1), ~Atom(Integer(3)));
                    },
                    _ => fail!("shoulda got a list")
                }
                assert_eq!(*items.get(1), ~Atom(Symbol(~"bar")));
                assert_eq!(*items.get(2), ~Atom(Integer(12)));
                match *items.get(3) {
                    ~List(ref items) => {
                        assert_eq!(items.len(), 2);
                        assert_eq!(*items.get(0), ~Atom(Symbol(~"hee")));
                        match *items.get(1) {
                            ~List(ref items) => {
                                assert_eq!(items.len(), 1);
                                assert_eq!(*items.get(0), ~Atom(Symbol(~"hah")));
                            },
                            _ => fail!("shoulda got a list")
                        }
                    },
                    _ => fail!("shoulda got a list")
                }
            },
            _ => fail!("got back an atom, it seems")
        }
    }
    #[test]
    fn parsed_tokens_of_a_list_of_atoms_and_a_nested_list_should_to_str_correctly() {
        let input = ~"((2 3) bar 12 (hee (hah)))";
        let tokens = &mut tokenize(input.to_owned());
        let parsed_item = parse(tokens);
        let output = parsed_item.print();
        assert_eq!(input, output);
    }
}

#[cfg(test)]
mod eval_test {
    mod un_cons {
        use super::super::{parse_str, Atom, List, Symbol};
        #[test]
        fn an_atom_expr_returns_the_atom_in_the_car_with_none_in_the_cdr() {
            let expr = parse_str(~"x");
            let (car, cdr) = expr.un_cons();
            assert_eq!(car, Atom(Symbol(~"x")));
            assert_eq!(cdr, None);
        }

        #[test]
        fn a_list_with_one_atom_element_returns_an_atom_car_and_an_empty_list_in_the_cdr() {
            let expr = parse_str(~"(x)");
            let (car, cdr) = expr.un_cons();
            assert_eq!(car, Atom(Symbol(~"x")));
            match cdr.expect("should be a value here") {
                List(items) => assert_eq!(items.len(), 0),
                _ => fail!("expected a list")
            }
        }

        #[test]
        fn a_list_with_multiple_elements_returns_the_first_car_and_the_rust_in_the_cdr() {
            let expr = parse_str(~"(y 2 3)");
            let (car, cdr) = expr.un_cons();
            assert_eq!(car, Atom(Symbol(~"y")));
            match cdr.expect("should be a value here") {
                List(items) => {
                    assert_eq!(items.len(), 2)
                },
                _ => fail!("expected a list")
            }
        }

        #[test]
        #[should_fail]
        fn calling_extract_on_an_empty_list_should_fail() {
            let expr = parse_str(~"()");
            expr.un_cons();
        }
    }
}
<|endoftext|>


<commit_before>/// HTTP2 Ping usage
///
/// hyper uses HTTP2 pings for two purposes:
///
/// 1. Adaptive flow control using BDP
/// 2. Connection keep-alive
///
/// Both cases are optional.
///
/// # BDP Algorithm
///
/// 1. When receiving a DATA frame, if a BDP ping isn't outstanding:
///   1a. Record current time.
///   1b. Send a BDP ping.
/// 2. Increment the number of received bytes.
/// 3. When the BDP ping ack is received:
///   3a. Record duration from sent time.
///   3b. Merge RTT with a running average.
///   3c. Calculate bdp as bytes/rtt.
///   3d. If bdp is over 2/3 max, set new max to bdp and update windows.

#[cfg(feature = "runtime")]
use std::fmt;
#[cfg(feature = "runtime")]
use std::future::Future;
#[cfg(feature = "runtime")]
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{self, Poll};
use std::time::Duration;
#[cfg(not(feature = "runtime"))]
use std::time::Instant;

use h2::{Ping, PingPong};
#[cfg(feature = "runtime")]
use tokio::time::{Delay, Instant};

type WindowSize = u32;

pub(super) fn disabled() -> Recorder {
    Recorder { shared: None }
}

pub(super) fn channel(ping_pong: PingPong, config: Config) -> (Recorder, Ponger) {
    debug_assert!(
        config.is_enabled(),
        "ping channel requires bdp or keep-alive config",
    );

    let bdp = config.bdp_initial_window.map(|wnd| Bdp {
        bdp: wnd,
        max_bandwidth: 0.0,
        rtt: 0.0,
    });

    let bytes = bdp.as_ref().map(|_| 0);

    #[cfg(feature = "runtime")]
    let keep_alive = config.keep_alive_interval.map(|interval| KeepAlive {
        interval,
        timeout: config.keep_alive_timeout,
        while_idle: config.keep_alive_while_idle,
        timer: tokio::time::delay_for(interval),
        state: KeepAliveState::Init,
    });

    #[cfg(feature = "runtime")]
    let last_read_at = keep_alive.as_ref().map(|_| Instant::now());

    let shared = Arc::new(Mutex::new(Shared {
        bytes,
        #[cfg(feature = "runtime")]
        last_read_at,
        #[cfg(feature = "runtime")]
        is_keep_alive_timed_out: false,
        ping_pong,
        ping_sent_at: None,
    }));

    (
        Recorder {
            shared: Some(shared.clone()),
        },
        Ponger {
            bdp,
            #[cfg(feature = "runtime")]
            keep_alive,
            shared,
        },
    )
}

#[derive(Clone)]
pub(super) struct Config {
    pub(super) bdp_initial_window: Option<WindowSize>,
    /// If no frames are received in this amount of time, a PING frame is sent.
    #[cfg(feature = "runtime")]
    pub(super) keep_alive_interval: Option<Duration>,
    /// After sending a keepalive PING, the connection will be closed if
    /// a pong is not received in this amount of time.
    #[cfg(feature = "runtime")]
    pub(super) keep_alive_timeout: Duration,
    /// If true, sends pings even when there are no active streams.
    #[cfg(feature = "runtime")]
    pub(super) keep_alive_while_idle: bool,
}

#[derive(Clone)]
pub(crate) struct Recorder {
    shared: Option<Arc<Mutex<Shared>>>,
}

pub(super) struct Ponger {
    bdp: Option<Bdp>,
    #[cfg(feature = "runtime")]
    keep_alive: Option<KeepAlive>,
    shared: Arc<Mutex<Shared>>,
}

struct Shared {
    ping_pong: PingPong,
    ping_sent_at: Option<Instant>,

    // bdp
    /// If `Some`, bdp is enabled, and this tracks how many bytes have been
    /// read during the current sample.
    bytes: Option<usize>,

    // keep-alive
    /// If `Some`, keep-alive is enabled, and the Instant is how long ago
    /// the connection read the last frame.
    #[cfg(feature = "runtime")]
    last_read_at: Option<Instant>,

    #[cfg(feature = "runtime")]
    is_keep_alive_timed_out: bool,
}

struct Bdp {
    /// Current BDP in bytes
    bdp: u32,
    /// Largest bandwidth we've seen so far.
    max_bandwidth: f64,
    /// Round trip time in seconds
    rtt: f64,
}

#[cfg(feature = "runtime")]
struct KeepAlive {
    /// If no frames are received in this amount of time, a PING frame is sent.
    interval: Duration,
    /// After sending a keepalive PING, the connection will be closed if
    /// a pong is not received in this amount of time.
    timeout: Duration,
    /// If true, sends pings even when there are no active streams.
    while_idle: bool,

    state: KeepAliveState,
    timer: Delay,
}

#[cfg(feature = "runtime")]
enum KeepAliveState {
    Init,
    Scheduled,
    PingSent,
}

pub(super) enum Ponged {
    SizeUpdate(WindowSize),
    #[cfg(feature = "runtime")]
    KeepAliveTimedOut,
}

#[cfg(feature = "runtime")]
#[derive(Debug)]
pub(super) struct KeepAliveTimedOut;

// ===== impl Config =====

impl Config {
    pub(super) fn is_enabled(&self) -> bool {
        #[cfg(feature = "runtime")]
        {
            self.bdp_initial_window.is_some() || self.keep_alive_interval.is_some()
        }

        #[cfg(not(feature = "runtime"))]
        {
            self.bdp_initial_window.is_some()
        }
    }
}

// ===== impl Recorder =====

impl Recorder {
    pub(crate) fn record_data(&self, len: usize) {
        let shared = if let Some(ref shared) = self.shared {
            shared
        } else {
            return;
        };

        let mut locked = shared.lock().unwrap();

        #[cfg(feature = "runtime")]
        locked.update_last_read_at();

        if let Some(ref mut bytes) = locked.bytes {
            *bytes += len;
        } else {
            // no need to send bdp ping if bdp is disabled
            return;
        }

        if !locked.is_ping_sent() {
            locked.send_ping();
        }
    }

    pub(crate) fn record_non_data(&self) {
        #[cfg(feature = "runtime")]
        {
            let shared = if let Some(ref shared) = self.shared {
                shared
            } else {
                return;
            };

            let mut locked = shared.lock().unwrap();

            locked.update_last_read_at();
        }
    }

    /// If the incoming stream is already closed, convert self into
    /// a disabled reporter.
    pub(super) fn for_stream(self, stream: &h2::RecvStream) -> Self {
        if stream.is_end_stream() {
            disabled()
        } else {
            self
        }
    }

    pub(super) fn ensure_not_timed_out(&self) -> crate::Result<()> {
        #[cfg(feature = "runtime")]
        {
            if let Some(ref shared) = self.shared {
                let locked = shared.lock().unwrap();
                if locked.is_keep_alive_timed_out {
                    return Err(KeepAliveTimedOut.crate_error());
                }
            }
        }

        // else
        Ok(())
    }
}

// ===== impl Ponger =====

impl Ponger {
    pub(super) fn poll(&mut self, cx: &mut task::Context<'_>) -> Poll<Ponged> {
        let mut locked = self.shared.lock().unwrap();
        #[cfg(feature = "runtime")]
        let is_idle = self.is_idle();

        #[cfg(feature = "runtime")]
        {
            if let Some(ref mut ka) = self.keep_alive {
                ka.schedule(is_idle, &locked);
                ka.maybe_ping(cx, &mut locked);
            }
        }

        if !locked.is_ping_sent() {
            // XXX: this doesn't register a waker...?
            return Poll::Pending;
        }

        let (bytes, rtt) = match locked.ping_pong.poll_pong(cx) {
            Poll::Ready(Ok(_pong)) => {
                let rtt = locked
                    .ping_sent_at
                    .expect("pong received implies ping_sent_at")
                    .elapsed();
                locked.ping_sent_at = None;
                trace!("recv pong");

                #[cfg(feature = "runtime")]
                {
                    if let Some(ref mut ka) = self.keep_alive {
                        locked.update_last_read_at();
                        ka.schedule(is_idle, &locked);
                    }
                }

                if self.bdp.is_some() {
                    let bytes = locked.bytes.expect("bdp enabled implies bytes");
                    locked.bytes = Some(0); // reset
                    trace!("received BDP ack; bytes = {}, rtt = {:?}", bytes, rtt);
                    (bytes, rtt)
                } else {
                    // no bdp, done!
                    return Poll::Pending;
                }
            }
            Poll::Ready(Err(e)) => {
                debug!("pong error: {}", e);
                return Poll::Pending;
            }
            Poll::Pending => {
                #[cfg(feature = "runtime")]
                {
                    if let Some(ref mut ka) = self.keep_alive {
                        if let Err(KeepAliveTimedOut) = ka.maybe_timeout(cx) {
                            self.keep_alive = None;
                            locked.is_keep_alive_timed_out = true;
                            return Poll::Ready(Ponged::KeepAliveTimedOut);
                        }
                    }
                }

                return Poll::Pending;
            }
        };

        drop(locked);

        if let Some(bdp) = self.bdp.as_mut().and_then(|bdp| bdp.calculate(bytes, rtt)) {
            Poll::Ready(Ponged::SizeUpdate(bdp))
        } else {
            // XXX: this doesn't register a waker...?
            Poll::Pending
        }
    }

    #[cfg(feature = "runtime")]
    fn is_idle(&self) -> bool {
        Arc::strong_count(&self.shared) <= 2
    }
}

// ===== impl Shared =====

impl Shared {
    fn send_ping(&mut self) {
        match self.ping_pong.send_ping(Ping::opaque()) {
            Ok(()) => {
                self.ping_sent_at = Some(Instant::now());
                trace!("sent ping");
            }
            Err(err) => {
                debug!("error sending ping: {}", err);
            }
        }
    }

    fn is_ping_sent(&self) -> bool {
        self.ping_sent_at.is_some()
    }

    #[cfg(feature = "runtime")]
    fn update_last_read_at(&mut self) {
        if self.last_read_at.is_some() {
            self.last_read_at = Some(Instant::now());
        }
    }

    #[cfg(feature = "runtime")]
    fn last_read_at(&self) -> Instant {
        self.last_read_at.expect("keep_alive expects last_read_at")
    }
}

// ===== impl Bdp =====

/// Any higher than this likely will be hitting the TCP flow control.
const BDP_LIMIT: usize = 1024 * 1024 * 16;

impl Bdp {
    fn calculate(&mut self, bytes: usize, rtt: Duration) -> Option<WindowSize> {
        // No need to do any math if we're at the limit.
        if self.bdp as usize == BDP_LIMIT {
            return None;
        }

        // average the rtt
        let rtt = seconds(rtt);
        if self.rtt == 0.0 {
            // First sample means rtt is first rtt.
            self.rtt = rtt;
        } else {
            // Weigh this rtt as 1/8 for a moving average.
            self.rtt += (rtt - self.rtt) * 0.125;
        }

        // calculate the current bandwidth
        let bw = (bytes as f64) / (self.rtt * 1.5);
        trace!("current bandwidth = {:.1}B/s", bw);

        if bw < self.max_bandwidth {
            // not a faster bandwidth, so don't update
            return None;
        } else {
            self.max_bandwidth = bw;
        }

        // if the current `bytes` sample is at least 2/3 the previous
        // bdp, increase to double the current sample.
        if bytes >= self.bdp as usize * 2 / 3 {
            self.bdp = (bytes * 2).min(BDP_LIMIT) as WindowSize;
            trace!("BDP increased to {}", self.bdp);
            Some(self.bdp)
        } else {
            None
        }
    }
}

fn seconds(dur: Duration) -> f64 {
    const NANOS_PER_SEC: f64 = 1_000_000_000.0;
    let secs = dur.as_secs() as f64;
    secs + (dur.subsec_nanos() as f64) / NANOS_PER_SEC
}

// ===== impl KeepAlive =====

#[cfg(feature = "runtime")]
impl KeepAlive {
    fn schedule(&mut self, is_idle: bool, shared: &Shared) {
        match self.state {
            KeepAliveState::Init => {
                if !self.while_idle && is_idle {
                    return;
                }

                self.state = KeepAliveState::Scheduled;
                let interval = shared.last_read_at() + self.interval;
                self.timer.reset(interval);
            }
            KeepAliveState::Scheduled | KeepAliveState::PingSent => (),
        }
    }

    fn maybe_ping(&mut self, cx: &mut task::Context<'_>, shared: &mut Shared) {
        match self.state {
            KeepAliveState::Scheduled => {
                if Pin::new(&mut self.timer).poll(cx).is_pending() {
                    return;
                }
                // check if we've received a frame while we were scheduled
                if shared.last_read_at() + self.interval > self.timer.deadline() {
                    self.state = KeepAliveState::Init;
                    cx.waker().wake_by_ref(); // schedule us again
                    return;
                }
                trace!("keep-alive interval ({:?}) reached", self.interval);
                shared.send_ping();
                self.state = KeepAliveState::PingSent;
                let timeout = Instant::now() + self.timeout;
                self.timer.reset(timeout);
            }
            KeepAliveState::Init | KeepAliveState::PingSent => (),
        }
    }

    fn maybe_timeout(&mut self, cx: &mut task::Context<'_>) -> Result<(), KeepAliveTimedOut> {
        match self.state {
            KeepAliveState::PingSent => {
                if Pin::new(&mut self.timer).poll(cx).is_pending() {
                    return Ok(());
                }
                trace!("keep-alive timeout ({:?}) reached", self.timeout);
                Err(KeepAliveTimedOut)
            }
            KeepAliveState::Init | KeepAliveState::Scheduled => Ok(()),
        }
    }
}

// ===== impl KeepAliveTimedOut =====

#[cfg(feature = "runtime")]
impl KeepAliveTimedOut {
    pub(super) fn crate_error(self) -> crate::Error {
        crate::Error::new(crate::error::Kind::Http2).with(self)
    }
}

#[cfg(feature = "runtime")]
impl fmt::Display for KeepAliveTimedOut {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("keep-alive timed out")
    }
}

#[cfg(feature = "runtime")]
impl std::error::Error for KeepAliveTimedOut {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        Some(&crate::error::TimedOut)
    }
}
<commit_msg>fix(server): reschedule keep alive interval timer once a pong is received<commit_after>/// HTTP2 Ping usage
///
/// hyper uses HTTP2 pings for two purposes:
///
/// 1. Adaptive flow control using BDP
/// 2. Connection keep-alive
///
/// Both cases are optional.
///
/// # BDP Algorithm
///
/// 1. When receiving a DATA frame, if a BDP ping isn't outstanding:
///   1a. Record current time.
///   1b. Send a BDP ping.
/// 2. Increment the number of received bytes.
/// 3. When the BDP ping ack is received:
///   3a. Record duration from sent time.
///   3b. Merge RTT with a running average.
///   3c. Calculate bdp as bytes/rtt.
///   3d. If bdp is over 2/3 max, set new max to bdp and update windows.

#[cfg(feature = "runtime")]
use std::fmt;
#[cfg(feature = "runtime")]
use std::future::Future;
#[cfg(feature = "runtime")]
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{self, Poll};
use std::time::Duration;
#[cfg(not(feature = "runtime"))]
use std::time::Instant;

use h2::{Ping, PingPong};
#[cfg(feature = "runtime")]
use tokio::time::{Delay, Instant};

type WindowSize = u32;

pub(super) fn disabled() -> Recorder {
    Recorder { shared: None }
}

pub(super) fn channel(ping_pong: PingPong, config: Config) -> (Recorder, Ponger) {
    debug_assert!(
        config.is_enabled(),
        "ping channel requires bdp or keep-alive config",
    );

    let bdp = config.bdp_initial_window.map(|wnd| Bdp {
        bdp: wnd,
        max_bandwidth: 0.0,
        rtt: 0.0,
    });

    let bytes = bdp.as_ref().map(|_| 0);

    #[cfg(feature = "runtime")]
    let keep_alive = config.keep_alive_interval.map(|interval| KeepAlive {
        interval,
        timeout: config.keep_alive_timeout,
        while_idle: config.keep_alive_while_idle,
        timer: tokio::time::delay_for(interval),
        state: KeepAliveState::Init,
    });

    #[cfg(feature = "runtime")]
    let last_read_at = keep_alive.as_ref().map(|_| Instant::now());

    let shared = Arc::new(Mutex::new(Shared {
        bytes,
        #[cfg(feature = "runtime")]
        last_read_at,
        #[cfg(feature = "runtime")]
        is_keep_alive_timed_out: false,
        ping_pong,
        ping_sent_at: None,
    }));

    (
        Recorder {
            shared: Some(shared.clone()),
        },
        Ponger {
            bdp,
            #[cfg(feature = "runtime")]
            keep_alive,
            shared,
        },
    )
}

#[derive(Clone)]
pub(super) struct Config {
    pub(super) bdp_initial_window: Option<WindowSize>,
    /// If no frames are received in this amount of time, a PING frame is sent.
    #[cfg(feature = "runtime")]
    pub(super) keep_alive_interval: Option<Duration>,
    /// After sending a keepalive PING, the connection will be closed if
    /// a pong is not received in this amount of time.
    #[cfg(feature = "runtime")]
    pub(super) keep_alive_timeout: Duration,
    /// If true, sends pings even when there are no active streams.
    #[cfg(feature = "runtime")]
    pub(super) keep_alive_while_idle: bool,
}

#[derive(Clone)]
pub(crate) struct Recorder {
    shared: Option<Arc<Mutex<Shared>>>,
}

pub(super) struct Ponger {
    bdp: Option<Bdp>,
    #[cfg(feature = "runtime")]
    keep_alive: Option<KeepAlive>,
    shared: Arc<Mutex<Shared>>,
}

struct Shared {
    ping_pong: PingPong,
    ping_sent_at: Option<Instant>,

    // bdp
    /// If `Some`, bdp is enabled, and this tracks how many bytes have been
    /// read during the current sample.
    bytes: Option<usize>,

    // keep-alive
    /// If `Some`, keep-alive is enabled, and the Instant is how long ago
    /// the connection read the last frame.
    #[cfg(feature = "runtime")]
    last_read_at: Option<Instant>,

    #[cfg(feature = "runtime")]
    is_keep_alive_timed_out: bool,
}

struct Bdp {
    /// Current BDP in bytes
    bdp: u32,
    /// Largest bandwidth we've seen so far.
    max_bandwidth: f64,
    /// Round trip time in seconds
    rtt: f64,
}

#[cfg(feature = "runtime")]
struct KeepAlive {
    /// If no frames are received in this amount of time, a PING frame is sent.
    interval: Duration,
    /// After sending a keepalive PING, the connection will be closed if
    /// a pong is not received in this amount of time.
    timeout: Duration,
    /// If true, sends pings even when there are no active streams.
    while_idle: bool,

    state: KeepAliveState,
    timer: Delay,
}

#[cfg(feature = "runtime")]
enum KeepAliveState {
    Init,
    Scheduled,
    PingSent,
}

pub(super) enum Ponged {
    SizeUpdate(WindowSize),
    #[cfg(feature = "runtime")]
    KeepAliveTimedOut,
}

#[cfg(feature = "runtime")]
#[derive(Debug)]
pub(super) struct KeepAliveTimedOut;

// ===== impl Config =====

impl Config {
    pub(super) fn is_enabled(&self) -> bool {
        #[cfg(feature = "runtime")]
        {
            self.bdp_initial_window.is_some() || self.keep_alive_interval.is_some()
        }

        #[cfg(not(feature = "runtime"))]
        {
            self.bdp_initial_window.is_some()
        }
    }
}

// ===== impl Recorder =====

impl Recorder {
    pub(crate) fn record_data(&self, len: usize) {
        let shared = if let Some(ref shared) = self.shared {
            shared
        } else {
            return;
        };

        let mut locked = shared.lock().unwrap();

        #[cfg(feature = "runtime")]
        locked.update_last_read_at();

        if let Some(ref mut bytes) = locked.bytes {
            *bytes += len;
        } else {
            // no need to send bdp ping if bdp is disabled
            return;
        }

        if !locked.is_ping_sent() {
            locked.send_ping();
        }
    }

    pub(crate) fn record_non_data(&self) {
        #[cfg(feature = "runtime")]
        {
            let shared = if let Some(ref shared) = self.shared {
                shared
            } else {
                return;
            };

            let mut locked = shared.lock().unwrap();

            locked.update_last_read_at();
        }
    }

    /// If the incoming stream is already closed, convert self into
    /// a disabled reporter.
    pub(super) fn for_stream(self, stream: &h2::RecvStream) -> Self {
        if stream.is_end_stream() {
            disabled()
        } else {
            self
        }
    }

    pub(super) fn ensure_not_timed_out(&self) -> crate::Result<()> {
        #[cfg(feature = "runtime")]
        {
            if let Some(ref shared) = self.shared {
                let locked = shared.lock().unwrap();
                if locked.is_keep_alive_timed_out {
                    return Err(KeepAliveTimedOut.crate_error());
                }
            }
        }

        // else
        Ok(())
    }
}

// ===== impl Ponger =====

impl Ponger {
    pub(super) fn poll(&mut self, cx: &mut task::Context<'_>) -> Poll<Ponged> {
        let mut locked = self.shared.lock().unwrap();
        #[cfg(feature = "runtime")]
        let is_idle = self.is_idle();

        #[cfg(feature = "runtime")]
        {
            if let Some(ref mut ka) = self.keep_alive {
                ka.schedule(is_idle, &locked);
                ka.maybe_ping(cx, &mut locked);
            }
        }

        if !locked.is_ping_sent() {
            // XXX: this doesn't register a waker...?
            return Poll::Pending;
        }

        let (bytes, rtt) = match locked.ping_pong.poll_pong(cx) {
            Poll::Ready(Ok(_pong)) => {
                let rtt = locked
                    .ping_sent_at
                    .expect("pong received implies ping_sent_at")
                    .elapsed();
                locked.ping_sent_at = None;
                trace!("recv pong");

                #[cfg(feature = "runtime")]
                {
                    if let Some(ref mut ka) = self.keep_alive {
                        locked.update_last_read_at();
                        ka.schedule(is_idle, &locked);
                    }
                }

                if self.bdp.is_some() {
                    let bytes = locked.bytes.expect("bdp enabled implies bytes");
                    locked.bytes = Some(0); // reset
                    trace!("received BDP ack; bytes = {}, rtt = {:?}", bytes, rtt);
                    (bytes, rtt)
                } else {
                    // no bdp, done!
                    return Poll::Pending;
                }
            }
            Poll::Ready(Err(e)) => {
                debug!("pong error: {}", e);
                return Poll::Pending;
            }
            Poll::Pending => {
                #[cfg(feature = "runtime")]
                {
                    if let Some(ref mut ka) = self.keep_alive {
                        if let Err(KeepAliveTimedOut) = ka.maybe_timeout(cx) {
                            self.keep_alive = None;
                            locked.is_keep_alive_timed_out = true;
                            return Poll::Ready(Ponged::KeepAliveTimedOut);
                        }
                    }
                }

                return Poll::Pending;
            }
        };

        drop(locked);

        if let Some(bdp) = self.bdp.as_mut().and_then(|bdp| bdp.calculate(bytes, rtt)) {
            Poll::Ready(Ponged::SizeUpdate(bdp))
        } else {
            // XXX: this doesn't register a waker...?
            Poll::Pending
        }
    }

    #[cfg(feature = "runtime")]
    fn is_idle(&self) -> bool {
        Arc::strong_count(&self.shared) <= 2
    }
}

// ===== impl Shared =====

impl Shared {
    fn send_ping(&mut self) {
        match self.ping_pong.send_ping(Ping::opaque()) {
            Ok(()) => {
                self.ping_sent_at = Some(Instant::now());
                trace!("sent ping");
            }
            Err(err) => {
                debug!("error sending ping: {}", err);
            }
        }
    }

    fn is_ping_sent(&self) -> bool {
        self.ping_sent_at.is_some()
    }

    #[cfg(feature = "runtime")]
    fn update_last_read_at(&mut self) {
        if self.last_read_at.is_some() {
            self.last_read_at = Some(Instant::now());
        }
    }

    #[cfg(feature = "runtime")]
    fn last_read_at(&self) -> Instant {
        self.last_read_at.expect("keep_alive expects last_read_at")
    }
}

// ===== impl Bdp =====

/// Any higher than this likely will be hitting the TCP flow control.
const BDP_LIMIT: usize = 1024 * 1024 * 16;

impl Bdp {
    fn calculate(&mut self, bytes: usize, rtt: Duration) -> Option<WindowSize> {
        // No need to do any math if we're at the limit.
        if self.bdp as usize == BDP_LIMIT {
            return None;
        }

        // average the rtt
        let rtt = seconds(rtt);
        if self.rtt == 0.0 {
            // First sample means rtt is first rtt.
            self.rtt = rtt;
        } else {
            // Weigh this rtt as 1/8 for a moving average.
            self.rtt += (rtt - self.rtt) * 0.125;
        }

        // calculate the current bandwidth
        let bw = (bytes as f64) / (self.rtt * 1.5);
        trace!("current bandwidth = {:.1}B/s", bw);

        if bw < self.max_bandwidth {
            // not a faster bandwidth, so don't update
            return None;
        } else {
            self.max_bandwidth = bw;
        }

        // if the current `bytes` sample is at least 2/3 the previous
        // bdp, increase to double the current sample.
        if bytes >= self.bdp as usize * 2 / 3 {
            self.bdp = (bytes * 2).min(BDP_LIMIT) as WindowSize;
            trace!("BDP increased to {}", self.bdp);
            Some(self.bdp)
        } else {
            None
        }
    }
}

fn seconds(dur: Duration) -> f64 {
    const NANOS_PER_SEC: f64 = 1_000_000_000.0;
    let secs = dur.as_secs() as f64;
    secs + (dur.subsec_nanos() as f64) / NANOS_PER_SEC
}

// ===== impl KeepAlive =====

#[cfg(feature = "runtime")]
impl KeepAlive {
    fn schedule(&mut self, is_idle: bool, shared: &Shared) {
        match self.state {
            KeepAliveState::Init => {
                if !self.while_idle && is_idle {
                    return;
                }

                self.state = KeepAliveState::Scheduled;
                let interval = shared.last_read_at() + self.interval;
                self.timer.reset(interval);
            }
            KeepAliveState::PingSent => {
                if shared.is_ping_sent() {
                    return;
                }

                self.state = KeepAliveState::Scheduled;
                let interval = shared.last_read_at() + self.interval;
                self.timer.reset(interval);
            }
            KeepAliveState::Scheduled => (),
        }
    }

    fn maybe_ping(&mut self, cx: &mut task::Context<'_>, shared: &mut Shared) {
        match self.state {
            KeepAliveState::Scheduled => {
                if Pin::new(&mut self.timer).poll(cx).is_pending() {
                    return;
                }
                // check if we've received a frame while we were scheduled
                if shared.last_read_at() + self.interval > self.timer.deadline() {
                    self.state = KeepAliveState::Init;
                    cx.waker().wake_by_ref(); // schedule us again
                    return;
                }
                trace!("keep-alive interval ({:?}) reached", self.interval);
                shared.send_ping();
                self.state = KeepAliveState::PingSent;
                let timeout = Instant::now() + self.timeout;
                self.timer.reset(timeout);
            }
            KeepAliveState::Init | KeepAliveState::PingSent => (),
        }
    }

    fn maybe_timeout(&mut self, cx: &mut task::Context<'_>) -> Result<(), KeepAliveTimedOut> {
        match self.state {
            KeepAliveState::PingSent => {
                if Pin::new(&mut self.timer).poll(cx).is_pending() {
                    return Ok(());
                }
                trace!("keep-alive timeout ({:?}) reached", self.timeout);
                Err(KeepAliveTimedOut)
            }
            KeepAliveState::Init | KeepAliveState::Scheduled => Ok(()),
        }
    }
}

// ===== impl KeepAliveTimedOut =====

#[cfg(feature = "runtime")]
impl KeepAliveTimedOut {
    pub(super) fn crate_error(self) -> crate::Error {
        crate::Error::new(crate::error::Kind::Http2).with(self)
    }
}

#[cfg(feature = "runtime")]
impl fmt::Display for KeepAliveTimedOut {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("keep-alive timed out")
    }
}

#[cfg(feature = "runtime")]
impl std::error::Error for KeepAliveTimedOut {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        Some(&crate::error::TimedOut)
    }
}
<|endoftext|>