<commit_before>//! The element nodes.

use std::fmt;

use node::{Attributes, Children, Node, Value};

#[doc(hidden)]
pub struct Element {
    name: String,
    attributes: Attributes,
    children: Children,
}

impl Element {
    pub fn new<T>(name: T) -> Self where T: Into<String> {
        Element { name: name.into(), attributes: Attributes::new(), children: Children::new() }
    }
}

impl fmt::Display for Element {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        try!(write!(formatter, "<{}", self.name));
        let mut attributes = self.attributes.iter().collect::<Vec<_>>();
        attributes.sort_by_key(|pair| pair.0.as_str());
        for (name, value) in attributes {
            match (value.contains("'"), value.contains('"')) {
                (true, false) | (false, false) => {
                    try!(write!(formatter, r#" {}="{}""#, name, value));
                },
                (false, true) => {
                    try!(write!(formatter, r#" {}='{}'"#, name, value));
                },
                _ => {},
            }
        }
        if self.children.is_empty() {
            return write!(formatter, "/>");
        }
        try!(write!(formatter, ">"));
        for child in self.children.iter() {
            try!(write!(formatter, "\n{}", child));
        }
        write!(formatter, "\n</{}>", self.name)
    }
}

impl Node for Element {
    #[inline]
    fn append<T>(&mut self, node: T) where T: Node {
        self.children.push(Box::new(node));
    }

    #[inline]
    fn assign<T, U>(&mut self, name: T, value: U) where T: Into<String>, U: Into<Value> {
        self.attributes.insert(name.into(), value.into());
    }
}

macro_rules! implement {
    ($(#[$doc:meta] struct $struct_name:ident($tag_name:expr))*) => (
        $(
            #[$doc]
            pub struct $struct_name {
                inner: Element,
            }

            impl $struct_name {
                /// Create a node.
                #[inline]
                pub fn new() -> Self {
                    $struct_name { inner: Element::new($tag_name) }
                }
            }

            node! { $struct_name::inner }
        )*
    );
}

implement! {
    #[doc = "An [`animate`][1] element.
    [1]: https://www.w3.org/TR/SVG/animate.html#AnimateElement"]
    struct Animate("animate")

    #[doc = "An [`animateColor`][1] element.
    [1]: https://www.w3.org/TR/SVG/animate.html#AnimateColorElement"]
    struct AnimateColor("animateColor")

    #[doc = "An [`animateMotion`][1] element.
    [1]: https://www.w3.org/TR/SVG/animate.html#AnimateMotionElement"]
    struct AnimateMotion("animateMotion")

    #[doc = "An [`animateTransform`][1] element.
    [1]: https://www.w3.org/TR/SVG/animate.html#AnimateTransformElement"]
    struct AnimateTransform("animateTransform")

    #[doc = "A [`circle`][1] element.
    [1]: https://www.w3.org/TR/SVG/shapes.html#CircleElement"]
    struct Circle("circle")

    #[doc = "A [`clipPath`][1] element.
    [1]: https://www.w3.org/TR/SVG/masking.html#ClipPathElement"]
    struct ClipPath("clipPath")

    #[doc = "A [`defs`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#DefsElement"]
    struct Definitions("defs")

    #[doc = "A [`desc`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#DescElement"]
    struct Description("desc")

    #[doc = "An [`ellipse`][1] element.
    [1]: https://www.w3.org/TR/SVG/shapes.html#EllipseElement"]
    struct Ellipse("ellipse")

    #[doc = "A [`filter`][1] element.
    [1]: https://www.w3.org/TR/SVG/filters.html#FilterElement"]
    struct Filter("filter")

    #[doc = "A [`g`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#GElement"]
    struct Group("g")

    #[doc = "An [`image`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#ImageElement"]
    struct Image("image")

    #[doc = "A [`line`][1] element.
    [1]: https://www.w3.org/TR/SVG/shapes.html#LineElement"]
    struct Line("line")

    #[doc = "A [`linearGradient`][1] element.
    [1]: https://www.w3.org/TR/SVG/pservers.html#LinearGradientElement"]
    struct LinearGradient("linearGradient")

    #[doc = "An [`a`][1] element.
    [1]: https://www.w3.org/TR/SVG/linking.html#AElement"]
    struct Link("a")

    #[doc = "A [`marker`][1] element.
    [1]: https://www.w3.org/TR/SVG/painting.html#MarkerElement"]
    struct Marker("marker")

    #[doc = "A [`mask`][1] element.
    [1]: https://www.w3.org/TR/SVG/masking.html#MaskElement"]
    struct Mask("mask")

    #[doc = "An [`mpath`][1] element.
    [1]: https://www.w3.org/TR/SVG/animate.html#MPathElement"]
    struct MotionPath("mpath")

    #[doc = "A [`path`][1] element.
    [1]: https://www.w3.org/TR/SVG/paths.html#PathElement"]
    struct Path("path")

    #[doc = "A [`polygon`][1] element.
    [1]: https://www.w3.org/TR/SVG/shapes.html#PolygonElement"]
    struct Polygon("polygon")

    #[doc = "A [`polyline`][1] element.
    [1]: https://www.w3.org/TR/SVG/shapes.html#PolylineElement"]
    struct Polyline("polyline")

    #[doc = "A [`radialGradient`][1] element.
    [1]: https://www.w3.org/TR/SVG/pservers.html#RadialGradientElement"]
    struct RadialGradient("radialGradient")

    #[doc = "A [`rect`][1] element.
    [1]: https://www.w3.org/TR/SVG/shapes.html#RectElement"]
    struct Rectangle("rect")

    #[doc = "A [`stop`][1] element.
    [1]: https://www.w3.org/TR/SVG/pservers.html#StopElement"]
    struct Stop("stop")

    #[doc = "A [`symbol`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#SymbolElement"]
    struct Symbol("symbol")

    #[doc = "A [`text`][1] element.
    [1]: https://www.w3.org/TR/SVG/text.html#TextElement"]
    struct Text("text")

    #[doc = "A [`textPath`][1] element.
    [1]: https://www.w3.org/TR/SVG/text.html#TextPathElement"]
    struct TextPath("textPath")

    #[doc = "A [`title`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#TitleElement"]
    struct Title("title")

    #[doc = "A [`use`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#UseElement"]
    struct Use("use")
}

macro_rules! implement {
    (@itemize $i:item) => ($i);
    (
        $(
            #[$doc:meta]
            struct $struct_name:ident($tag_name:expr)
            [$($pn:ident: $($pt:tt)*),*] [$inner:ident $(,$an:ident: $at:ty)*] $body:block
        )*
    ) => (
        $(
            #[$doc]
            pub struct $struct_name {
                inner: Element,
            }

            implement! { @itemize
                impl $struct_name {
                    /// Create a node.
                    #[inline]
                    pub fn new<$($pn: $($pt)*),*>($($an: $at),*) -> Self {
                        #[inline(always)]
                        fn initialize<$($pn: $($pt)*),*>($inner: &mut Element $(, $an: $at)*) $body
                        let mut inner = Element::new($tag_name);
                        initialize(&mut inner $(, $an)*);
                        $struct_name { inner: inner }
                    }
                }
            }

            node! { $struct_name::inner }
        )*
    );
}

implement! {
    #[doc = "An [`svg`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#SVGElement"]
    struct SVG("svg") [] [inner] {
        inner.assign("xmlns", "http://www.w3.org/2000/svg");
    }

    #[doc = "A [`script`][1] element.
    [1]: https://www.w3.org/TR/SVG/script.html#ScriptElement"]
    struct Script("script") [T: Into<String>] [inner, content: T] {
        inner.append(::node::Text::new(content));
    }

    #[doc = "A [`style`][1] element.
    [1]: https://www.w3.org/TR/SVG/styling.html#StyleElement"]
    struct Style("style") [T: Into<String>] [inner, content: T] {
        inner.append(::node::Text::new(content));
    }
}

pub mod path;

#[cfg(test)]
mod tests {
    use node::Node;
    use super::{Element, Style};

    #[test]
    fn element_display() {
        let mut element = Element::new("foo");
        element.assign("x", -10);
        element.assign("y", "10px");
        element.assign("s", (12.5, 13.0));
        element.assign("c", "green");
        element.append(Element::new("bar"));
        assert_eq!(element.to_string(), "\
            <foo c=\"green\" s=\"12.5 13\" x=\"-10\" y=\"10px\">\n\
            <bar/>\n\
            </foo>\
        ");
    }

    #[test]
    fn element_display_qoutes() {
        let mut element = Element::new("foo");
        element.assign("s", "'single'");
        element.assign("d", r#""double""#);
        element.assign("m", r#""mixed'"#);
        assert_eq!(element.to_string(), r#"<foo d='"double"' s="'single'"/>"#);
    }

    #[test]
    fn text_display() {
        let element = Style::new("* { font-family: foo; }");
        assert_eq!(element.to_string(), "\
            <style>\n\
            * { font-family: foo; }\n\
            </style>\
        ");
    }
}
<commit_msg>Fix a typo<commit_after>//! The element nodes.

use std::fmt;

use node::{Attributes, Children, Node, Value};

#[doc(hidden)]
pub struct Element {
    name: String,
    attributes: Attributes,
    children: Children,
}

impl Element {
    pub fn new<T>(name: T) -> Self where T: Into<String> {
        Element { name: name.into(), attributes: Attributes::new(), children: Children::new() }
    }
}

impl fmt::Display for Element {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        try!(write!(formatter, "<{}", self.name));
        let mut attributes = self.attributes.iter().collect::<Vec<_>>();
        attributes.sort_by_key(|pair| pair.0.as_str());
        for (name, value) in attributes {
            match (value.contains("'"), value.contains('"')) {
                (true, false) | (false, false) => {
                    try!(write!(formatter, r#" {}="{}""#, name, value));
                },
                (false, true) => {
                    try!(write!(formatter, r#" {}='{}'"#, name, value));
                },
                _ => {},
            }
        }
        if self.children.is_empty() {
            return write!(formatter, "/>");
        }
        try!(write!(formatter, ">"));
        for child in self.children.iter() {
            try!(write!(formatter, "\n{}", child));
        }
        write!(formatter, "\n</{}>", self.name)
    }
}

impl Node for Element {
    #[inline]
    fn append<T>(&mut self, node: T) where T: Node {
        self.children.push(Box::new(node));
    }

    #[inline]
    fn assign<T, U>(&mut self, name: T, value: U) where T: Into<String>, U: Into<Value> {
        self.attributes.insert(name.into(), value.into());
    }
}

macro_rules! implement {
    ($(#[$doc:meta] struct $struct_name:ident($tag_name:expr))*) => (
        $(
            #[$doc]
            pub struct $struct_name {
                inner: Element,
            }

            impl $struct_name {
                /// Create a node.
                #[inline]
                pub fn new() -> Self {
                    $struct_name { inner: Element::new($tag_name) }
                }
            }

            node! { $struct_name::inner }
        )*
    );
}

implement! {
    #[doc = "An [`animate`][1] element.
    [1]: https://www.w3.org/TR/SVG/animate.html#AnimateElement"]
    struct Animate("animate")

    #[doc = "An [`animateColor`][1] element.
    [1]: https://www.w3.org/TR/SVG/animate.html#AnimateColorElement"]
    struct AnimateColor("animateColor")

    #[doc = "An [`animateMotion`][1] element.
    [1]: https://www.w3.org/TR/SVG/animate.html#AnimateMotionElement"]
    struct AnimateMotion("animateMotion")

    #[doc = "An [`animateTransform`][1] element.
    [1]: https://www.w3.org/TR/SVG/animate.html#AnimateTransformElement"]
    struct AnimateTransform("animateTransform")

    #[doc = "A [`circle`][1] element.
    [1]: https://www.w3.org/TR/SVG/shapes.html#CircleElement"]
    struct Circle("circle")

    #[doc = "A [`clipPath`][1] element.
    [1]: https://www.w3.org/TR/SVG/masking.html#ClipPathElement"]
    struct ClipPath("clipPath")

    #[doc = "A [`defs`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#DefsElement"]
    struct Definitions("defs")

    #[doc = "A [`desc`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#DescElement"]
    struct Description("desc")

    #[doc = "An [`ellipse`][1] element.
    [1]: https://www.w3.org/TR/SVG/shapes.html#EllipseElement"]
    struct Ellipse("ellipse")

    #[doc = "A [`filter`][1] element.
    [1]: https://www.w3.org/TR/SVG/filters.html#FilterElement"]
    struct Filter("filter")

    #[doc = "A [`g`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#GElement"]
    struct Group("g")

    #[doc = "An [`image`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#ImageElement"]
    struct Image("image")

    #[doc = "A [`line`][1] element.
    [1]: https://www.w3.org/TR/SVG/shapes.html#LineElement"]
    struct Line("line")

    #[doc = "A [`linearGradient`][1] element.
    [1]: https://www.w3.org/TR/SVG/pservers.html#LinearGradientElement"]
    struct LinearGradient("linearGradient")

    #[doc = "An [`a`][1] element.
    [1]: https://www.w3.org/TR/SVG/linking.html#AElement"]
    struct Link("a")

    #[doc = "A [`marker`][1] element.
    [1]: https://www.w3.org/TR/SVG/painting.html#MarkerElement"]
    struct Marker("marker")

    #[doc = "A [`mask`][1] element.
    [1]: https://www.w3.org/TR/SVG/masking.html#MaskElement"]
    struct Mask("mask")

    #[doc = "An [`mpath`][1] element.
    [1]: https://www.w3.org/TR/SVG/animate.html#MPathElement"]
    struct MotionPath("mpath")

    #[doc = "A [`path`][1] element.
    [1]: https://www.w3.org/TR/SVG/paths.html#PathElement"]
    struct Path("path")

    #[doc = "A [`polygon`][1] element.
    [1]: https://www.w3.org/TR/SVG/shapes.html#PolygonElement"]
    struct Polygon("polygon")

    #[doc = "A [`polyline`][1] element.
    [1]: https://www.w3.org/TR/SVG/shapes.html#PolylineElement"]
    struct Polyline("polyline")

    #[doc = "A [`radialGradient`][1] element.
    [1]: https://www.w3.org/TR/SVG/pservers.html#RadialGradientElement"]
    struct RadialGradient("radialGradient")

    #[doc = "A [`rect`][1] element.
    [1]: https://www.w3.org/TR/SVG/shapes.html#RectElement"]
    struct Rectangle("rect")

    #[doc = "A [`stop`][1] element.
    [1]: https://www.w3.org/TR/SVG/pservers.html#StopElement"]
    struct Stop("stop")

    #[doc = "A [`symbol`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#SymbolElement"]
    struct Symbol("symbol")

    #[doc = "A [`text`][1] element.
    [1]: https://www.w3.org/TR/SVG/text.html#TextElement"]
    struct Text("text")

    #[doc = "A [`textPath`][1] element.
    [1]: https://www.w3.org/TR/SVG/text.html#TextPathElement"]
    struct TextPath("textPath")

    #[doc = "A [`title`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#TitleElement"]
    struct Title("title")

    #[doc = "A [`use`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#UseElement"]
    struct Use("use")
}

macro_rules! implement {
    (@itemize $i:item) => ($i);
    (
        $(
            #[$doc:meta]
            struct $struct_name:ident($tag_name:expr)
            [$($pn:ident: $($pt:tt)*),*] [$inner:ident $(,$an:ident: $at:ty)*] $body:block
        )*
    ) => (
        $(
            #[$doc]
            pub struct $struct_name {
                inner: Element,
            }

            implement! { @itemize
                impl $struct_name {
                    /// Create a node.
                    #[inline]
                    pub fn new<$($pn: $($pt)*),*>($($an: $at),*) -> Self {
                        #[inline(always)]
                        fn initialize<$($pn: $($pt)*),*>($inner: &mut Element $(, $an: $at)*) $body
                        let mut inner = Element::new($tag_name);
                        initialize(&mut inner $(, $an)*);
                        $struct_name { inner: inner }
                    }
                }
            }

            node! { $struct_name::inner }
        )*
    );
}

implement! {
    #[doc = "An [`svg`][1] element.
    [1]: https://www.w3.org/TR/SVG/struct.html#SVGElement"]
    struct SVG("svg") [] [inner] {
        inner.assign("xmlns", "http://www.w3.org/2000/svg");
    }

    #[doc = "A [`script`][1] element.
    [1]: https://www.w3.org/TR/SVG/script.html#ScriptElement"]
    struct Script("script") [T: Into<String>] [inner, content: T] {
        inner.append(::node::Text::new(content));
    }

    #[doc = "A [`style`][1] element.
    [1]: https://www.w3.org/TR/SVG/styling.html#StyleElement"]
    struct Style("style") [T: Into<String>] [inner, content: T] {
        inner.append(::node::Text::new(content));
    }
}

pub mod path;

#[cfg(test)]
mod tests {
    use node::Node;
    use super::{Element, Style};

    #[test]
    fn element_display() {
        let mut element = Element::new("foo");
        element.assign("x", -10);
        element.assign("y", "10px");
        element.assign("s", (12.5, 13.0));
        element.assign("c", "green");
        element.append(Element::new("bar"));
        assert_eq!(element.to_string(), "\
            <foo c=\"green\" s=\"12.5 13\" x=\"-10\" y=\"10px\">\n\
            <bar/>\n\
            </foo>\
        ");
    }

    #[test]
    fn element_display_quotes() {
        let mut element = Element::new("foo");
        element.assign("s", "'single'");
        element.assign("d", r#""double""#);
        element.assign("m", r#""mixed'"#);
        assert_eq!(element.to_string(), r#"<foo d='"double"' s="'single'"/>"#);
    }

    #[test]
    fn text_display() {
        let element = Style::new("* { font-family: foo; }");
        assert_eq!(element.to_string(), "\
            <style>\n\
            * { font-family: foo; }\n\
            </style>\
        ");
    }
}
<|endoftext|>


<commit_before>use base64::write::EncoderWriter;
use indexmap::IndexMap;
use prometheus::proto::{MetricFamily, MetricType};
use prometheus::{Encoder, Error};
use std::io::Write;

/// The `LogEncoder` struct encodes Prometheus metrics in the format [`crates-io-heroku-metrics`]
/// expects metrics to be logged. This can be used to forward instance metrics to it, allowing them
/// to be scraped by the Rust infrastructure monitoring system.
///
/// The metrics are encoded in the format [Vector expects them][vector], and the list of metrics is
/// json-encoded and then base64-encoded. The whole thing is prefixed with a predefined string to
/// let [`crates-io-heroku-metrics`] find it easily.
///
/// This is needed mostly for crates.io hosted on Heroku. Deployments of crates.io on other
/// platforms shouldn't need this.
///
/// [`crates-io-heroku-metrics`]: https://github.com/rust-lang/crates-io-heroku-metrics
/// [vector]: https://vector.dev/docs/about/under-the-hood/architecture/data-model/metric/
#[derive(Debug, Clone, Copy)]
pub struct LogEncoder(());

impl LogEncoder {
    pub fn new() -> Self {
        Self(())
    }
}

impl Encoder for LogEncoder {
    fn encode<W: Write>(
        &self,
        families: &[MetricFamily],
        mut dest: &mut W,
    ) -> prometheus::Result<()> {
        let events = families_to_json_events(families);

        dest.write_all(b"crates-io-heroku-metrics:ingest ")?;
        let base64_dest = EncoderWriter::new(&mut dest, base64::STANDARD);
        serde_json::to_writer(base64_dest, &events).map_err(|e| Error::Msg(e.to_string()))?;
        dest.write_all(b"\n")?;

        Ok(())
    }

    fn format_type(&self) -> &str {
        "crates-io-heroku-metrics log encoding"
    }
}

fn families_to_json_events(families: &[MetricFamily]) -> Vec<VectorEvent<'_>> {
    let mut events = Vec::new();
    for family in families {
        for metric in family.get_metric() {
            let data = match family.get_field_type() {
                MetricType::COUNTER => VectorMetricData::Counter {
                    value: metric.get_counter().get_value(),
                },
                MetricType::GAUGE => VectorMetricData::Gauge {
                    value: metric.get_gauge().get_value(),
                },
                MetricType::HISTOGRAM => {
                    let histogram = metric.get_histogram();

                    // We need to convert from cumulative counts (used by the Prometheus library)
                    // to plain counts (used by Vector).
                    let mut buckets = Vec::new();
                    let mut last_cumulative_count = 0;
                    for bucket in histogram.get_bucket() {
                        buckets.push(VectorHistogramBucket {
                            upper_limit: bucket.get_upper_bound(),
                            count: bucket.get_cumulative_count() - last_cumulative_count,
                        });
                        last_cumulative_count = bucket.get_cumulative_count();
                    }

                    VectorMetricData::AggregatedHistogram {
                        count: histogram.get_sample_count(),
                        sum: histogram.get_sample_sum(),
                        buckets,
                    }
                }
                other => {
                    panic!("unsupported metric type: {:?}", other)
                }
            };
            events.push(VectorEvent {
                metric: VectorMetric {
                    data,
                    kind: "absolute",
                    name: family.get_name(),
                    tags: metric
                        .get_label()
                        .iter()
                        .map(|p| (p.get_name(), p.get_value()))
                        .collect(),
                },
            });
        }
    }
    events
}

#[derive(Serialize, Debug, PartialEq)]
struct VectorEvent<'a> {
    metric: VectorMetric<'a>,
}

#[derive(Serialize, Debug, PartialEq)]
struct VectorMetric<'a> {
    #[serde(flatten)]
    data: VectorMetricData,
    kind: &'a str,
    name: &'a str,
    tags: IndexMap<&'a str, &'a str>,
}

#[derive(Serialize, Debug, PartialEq)]
#[serde(rename_all = "snake_case")]
enum VectorMetricData {
    AggregatedHistogram {
        buckets: Vec<VectorHistogramBucket>,
        count: u64,
        sum: f64,
    },
    Counter {
        value: f64,
    },
    Gauge {
        value: f64,
    },
}

#[derive(Serialize, Debug, PartialEq)]
struct VectorHistogramBucket {
    upper_limit: f64,
    count: u64,
}

#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Error;
    use prometheus::{Histogram, HistogramOpts, IntCounter, IntGauge, IntGaugeVec, Opts, Registry};

    #[test]
    fn test_counter_to_json() -> Result<(), Error> {
        let counter =
            IntCounter::with_opts(Opts::new("sample_counter", "sample_counter help message"))?;
        let registry = Registry::new();
        registry.register(Box::new(counter.clone()))?;

        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Counter { value: 0.0 },
                    kind: "absolute",
                    name: "sample_counter",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        counter.inc_by(42);
        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Counter { value: 42.0 },
                    kind: "absolute",
                    name: "sample_counter",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_gauge_to_json() -> Result<(), Error> {
        let gauge = IntGauge::with_opts(Opts::new("sample_gauge", "sample_gauge help message"))?;
        let registry = Registry::new();
        registry.register(Box::new(gauge.clone()))?;

        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Gauge { value: 0.0 },
                    kind: "absolute",
                    name: "sample_gauge",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        gauge.set(42);
        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Gauge { value: 42.0 },
                    kind: "absolute",
                    name: "sample_gauge",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_histogram_to_json() -> Result<(), Error> {
        let histogram = Histogram::with_opts(HistogramOpts::new(
            "sample_histogram",
            "sample_histogram help message",
        ))?;
        let registry = Registry::new();
        registry.register(Box::new(histogram.clone()))?;

        let mut value = 0.0;
        while value < 11.0 {
            histogram.observe(value);
            value += 0.001;
        }

        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::AggregatedHistogram {
                        buckets: vec![
                            VectorHistogramBucket {
                                upper_limit: 0.005,
                                count: 6,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.01,
                                count: 4,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.025,
                                count: 15,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.05,
                                count: 25,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.1,
                                count: 50,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.25,
                                count: 150,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.5,
                                count: 250,
                            },
                            VectorHistogramBucket {
                                upper_limit: 1.0,
                                count: 500,
                            },
                            VectorHistogramBucket {
                                upper_limit: 2.5,
                                count: 1501,
                            },
                            VectorHistogramBucket {
                                upper_limit: 5.0,
                                count: 2499,
                            },
                            VectorHistogramBucket {
                                upper_limit: 10.0,
                                count: 5001,
                            },
                        ],
                        count: 11001,
                        sum: 60505.50000000138,
                    },
                    kind: "absolute",
                    name: "sample_histogram",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_metric_with_tags_to_json() -> Result<(), Error> {
        let gauge_vec = IntGaugeVec::new(
            Opts::new("sample_gauge", "sample_gauge help message"),
            &["label1", "label2"],
        )?;
        let registry = Registry::new();
        registry.register(Box::new(gauge_vec.clone()))?;

        gauge_vec.with_label_values(&["foo", "1"]).set(42);
        gauge_vec.with_label_values(&["bar", "2"]).set(98);

        assert_eq!(
            vec![
                VectorEvent {
                    metric: VectorMetric {
                        data: VectorMetricData::Gauge { value: 98.0 },
                        kind: "absolute",
                        name: "sample_gauge",
                        tags: [("label1", "bar"), ("label2", "2")]
                            .iter()
                            .copied()
                            .collect(),
                    }
                },
                VectorEvent {
                    metric: VectorMetric {
                        data: VectorMetricData::Gauge { value: 42.0 },
                        kind: "absolute",
                        name: "sample_gauge",
                        tags: [("label1", "foo"), ("label2", "1")]
                            .iter()
                            .copied()
                            .collect(),
                    }
                },
            ],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_encoding() -> Result<(), Error> {
        let gauge = IntGauge::with_opts(Opts::new("sample_gauge", "sample_gauge help message"))?;
        let registry = Registry::new();
        registry.register(Box::new(gauge.clone()))?;

        let mut output = Vec::new();
        LogEncoder::new().encode(&registry.gather(), &mut output)?;

        assert_eq!(
            b"crates-io-heroku-metrics:ingest W3sibWV0cmljIjp7ImdhdWdlIjp7InZhbHVlIjowLjB9LCJraW5kIjoiYWJzb2x1dGUiLCJuYW1lIjoic2FtcGxlX2dhdWdlIiwidGFncyI6e319fV0=\n",
            output.as_slice(),
        );

        Ok(())
    }
}
<commit_msg>Auto merge of #3690 - pietroalbini:max-metrics-size, r=jtgeibel<commit_after>use base64::write::EncoderWriter;
use indexmap::IndexMap;
use prometheus::proto::{MetricFamily, MetricType};
use prometheus::{Encoder, Error};
use serde::ser::SerializeSeq;
use serde::{Serialize, Serializer as _};
use serde_json::Serializer;
use std::cell::Cell;
use std::io::Write;
use std::rc::Rc;

const CHUNKS_MAX_SIZE_BYTES: usize = 5000;

/// The `LogEncoder` struct encodes Prometheus metrics in the format [`crates-io-heroku-metrics`]
/// expects metrics to be logged. This can be used to forward instance metrics to it, allowing them
/// to be scraped by the Rust infrastructure monitoring system.
///
/// The metrics are encoded in the format [Vector expects them][vector], and the list of metrics is
/// json-encoded and then base64-encoded. The whole thing is prefixed with a predefined string to
/// let [`crates-io-heroku-metrics`] find it easily.
///
/// This is needed mostly for crates.io hosted on Heroku. Deployments of crates.io on other
/// platforms shouldn't need this.
///
/// [`crates-io-heroku-metrics`]: https://github.com/rust-lang/crates-io-heroku-metrics
/// [vector]: https://vector.dev/docs/about/under-the-hood/architecture/data-model/metric/
#[derive(Debug, Clone, Copy)]
pub struct LogEncoder(());

impl LogEncoder {
    pub fn new() -> Self {
        Self(())
    }
}

impl Encoder for LogEncoder {
    fn encode<W: Write>(&self, families: &[MetricFamily], dest: &mut W) -> prometheus::Result<()> {
        let events = families_to_json_events(families);

        let chunks = serialize_and_split_list(events.iter(), CHUNKS_MAX_SIZE_BYTES)
            .map_err(|e| Error::Msg(e.to_string()))?;

        for chunk in chunks {
            dest.write_all(b"crates-io-heroku-metrics:ingest ")?;
            dest.write_all(&chunk)?;
            dest.write_all(b"\n")?;
        }

        Ok(())
    }

    fn format_type(&self) -> &str {
        "crates-io-heroku-metrics log encoding"
    }
}

fn families_to_json_events(families: &[MetricFamily]) -> Vec<VectorEvent<'_>> {
    let mut events = Vec::new();
    for family in families {
        for metric in family.get_metric() {
            let data = match family.get_field_type() {
                MetricType::COUNTER => VectorMetricData::Counter {
                    value: metric.get_counter().get_value(),
                },
                MetricType::GAUGE => VectorMetricData::Gauge {
                    value: metric.get_gauge().get_value(),
                },
                MetricType::HISTOGRAM => {
                    let histogram = metric.get_histogram();

                    // We need to convert from cumulative counts (used by the Prometheus library)
                    // to plain counts (used by Vector).
                    let mut buckets = Vec::new();
                    let mut last_cumulative_count = 0;
                    for bucket in histogram.get_bucket() {
                        buckets.push(VectorHistogramBucket {
                            upper_limit: bucket.get_upper_bound(),
                            count: bucket.get_cumulative_count() - last_cumulative_count,
                        });
                        last_cumulative_count = bucket.get_cumulative_count();
                    }

                    VectorMetricData::AggregatedHistogram {
                        count: histogram.get_sample_count(),
                        sum: histogram.get_sample_sum(),
                        buckets,
                    }
                }
                other => {
                    panic!("unsupported metric type: {:?}", other)
                }
            };
            events.push(VectorEvent {
                metric: VectorMetric {
                    data,
                    kind: "absolute",
                    name: family.get_name(),
                    tags: metric
                        .get_label()
                        .iter()
                        .map(|p| (p.get_name(), p.get_value()))
                        .collect(),
                },
            });
        }
    }
    events
}

/// Serialize a list of items into multiple Base64-encoded JSON chunks.
///
/// Our hosting platform (Heroku) limits the size of log lines, arbitrarily splitting them once
/// they reach a threshold. We can't let Heroku do the split as it doesn't know where to properly
/// do that, so we need to do the splitting ourselves.
///
/// This function takes an iterator of serializable items and returns the serialized version,
/// possibly split into multiple chunks. Each chunk is *at least* `max_size_hint` long, as the
/// function stops serializing new items in the same chunk only when the size limit is reached
/// after serializing an item.
///
/// Because of that `max_size_hint` should be lower than the upper bound we can't cross.
fn serialize_and_split_list<'a, S: Serialize + 'a>(
    items: impl Iterator<Item = &'a S>,
    max_size_hint: usize,
) -> Result<Vec<Vec<u8>>, serde_json::Error> {
    let mut items = items.peekable();

    let mut result = Vec::new();
    while items.peek().is_some() {
        let mut writer = TrackedWriter::new();
        let written_count = writer.written_count.clone();
        let mut serializer = Serializer::new(EncoderWriter::new(&mut writer, base64::STANDARD));

        let mut seq = serializer.serialize_seq(None)?;
        while let Some(next) = items.next() {
            seq.serialize_element(next)?;
            if written_count.get() >= max_size_hint {
                break;
            }
        }
        seq.end()?;
        drop(serializer);

        result.push(writer.buffer);
    }

    Ok(result)
}

struct TrackedWriter {
    buffer: Vec<u8>,
    written_count: Rc<Cell<usize>>,
}

impl TrackedWriter {
    fn new() -> Self {
        Self {
            buffer: Vec::new(),
            written_count: Rc::new(Cell::new(0)),
        }
    }
}

impl Write for TrackedWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let written = self.buffer.write(buf)?;
        self.written_count.set(self.written_count.get() + written);
        Ok(written)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.buffer.flush()
    }
}

#[derive(Serialize, Debug, PartialEq)]
struct VectorEvent<'a> {
    metric: VectorMetric<'a>,
}

#[derive(Serialize, Debug, PartialEq)]
struct VectorMetric<'a> {
    #[serde(flatten)]
    data: VectorMetricData,
    kind: &'a str,
    name: &'a str,
    tags: IndexMap<&'a str, &'a str>,
}

#[derive(Serialize, Debug, PartialEq)]
#[serde(rename_all = "snake_case")]
enum VectorMetricData {
    AggregatedHistogram {
        buckets: Vec<VectorHistogramBucket>,
        count: u64,
        sum: f64,
    },
    Counter {
        value: f64,
    },
    Gauge {
        value: f64,
    },
}

#[derive(Serialize, Debug, PartialEq)]
struct VectorHistogramBucket {
    upper_limit: f64,
    count: u64,
}

#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Error;
    use prometheus::{Histogram, HistogramOpts, IntCounter, IntGauge, IntGaugeVec, Opts, Registry};

    #[test]
    fn test_counter_to_json() -> Result<(), Error> {
        let counter =
            IntCounter::with_opts(Opts::new("sample_counter", "sample_counter help message"))?;
        let registry = Registry::new();
        registry.register(Box::new(counter.clone()))?;

        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Counter { value: 0.0 },
                    kind: "absolute",
                    name: "sample_counter",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        counter.inc_by(42);
        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Counter { value: 42.0 },
                    kind: "absolute",
                    name: "sample_counter",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_gauge_to_json() -> Result<(), Error> {
        let gauge = IntGauge::with_opts(Opts::new("sample_gauge", "sample_gauge help message"))?;
        let registry = Registry::new();
        registry.register(Box::new(gauge.clone()))?;

        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Gauge { value: 0.0 },
                    kind: "absolute",
                    name: "sample_gauge",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        gauge.set(42);
        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::Gauge { value: 42.0 },
                    kind: "absolute",
                    name: "sample_gauge",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_histogram_to_json() -> Result<(), Error> {
        let histogram = Histogram::with_opts(HistogramOpts::new(
            "sample_histogram",
            "sample_histogram help message",
        ))?;
        let registry = Registry::new();
        registry.register(Box::new(histogram.clone()))?;

        let mut value = 0.0;
        while value < 11.0 {
            histogram.observe(value);
            value += 0.001;
        }

        assert_eq!(
            vec![VectorEvent {
                metric: VectorMetric {
                    data: VectorMetricData::AggregatedHistogram {
                        buckets: vec![
                            VectorHistogramBucket {
                                upper_limit: 0.005,
                                count: 6,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.01,
                                count: 4,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.025,
                                count: 15,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.05,
                                count: 25,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.1,
                                count: 50,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.25,
                                count: 150,
                            },
                            VectorHistogramBucket {
                                upper_limit: 0.5,
                                count: 250,
                            },
                            VectorHistogramBucket {
                                upper_limit: 1.0,
                                count: 500,
                            },
                            VectorHistogramBucket {
                                upper_limit: 2.5,
                                count: 1501,
                            },
                            VectorHistogramBucket {
                                upper_limit: 5.0,
                                count: 2499,
                            },
                            VectorHistogramBucket {
                                upper_limit: 10.0,
                                count: 5001,
                            },
                        ],
                        count: 11001,
                        sum: 60505.50000000138,
                    },
                    kind: "absolute",
                    name: "sample_histogram",
                    tags: IndexMap::new(),
                }
            }],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_metric_with_tags_to_json() -> Result<(), Error> {
        let gauge_vec = IntGaugeVec::new(
            Opts::new("sample_gauge", "sample_gauge help message"),
            &["label1", "label2"],
        )?;
        let registry = Registry::new();
        registry.register(Box::new(gauge_vec.clone()))?;

        gauge_vec.with_label_values(&["foo", "1"]).set(42);
        gauge_vec.with_label_values(&["bar", "2"]).set(98);

        assert_eq!(
            vec![
                VectorEvent {
                    metric: VectorMetric {
                        data: VectorMetricData::Gauge { value: 98.0 },
                        kind: "absolute",
                        name: "sample_gauge",
                        tags: [("label1", "bar"), ("label2", "2")]
                            .iter()
                            .copied()
                            .collect(),
                    }
                },
                VectorEvent {
                    metric: VectorMetric {
                        data: VectorMetricData::Gauge { value: 42.0 },
                        kind: "absolute",
                        name: "sample_gauge",
                        tags: [("label1", "foo"), ("label2", "1")]
                            .iter()
                            .copied()
                            .collect(),
                    }
                },
            ],
            families_to_json_events(&registry.gather())
        );

        Ok(())
    }

    #[test]
    fn test_encoding() -> Result<(), Error> {
        let gauge = IntGauge::with_opts(Opts::new("sample_gauge", "sample_gauge help message"))?;
        let registry = Registry::new();
        registry.register(Box::new(gauge.clone()))?;

        let mut output = Vec::new();
        LogEncoder::new().encode(&registry.gather(), &mut output)?;

        assert_eq!(
            b"crates-io-heroku-metrics:ingest W3sibWV0cmljIjp7ImdhdWdlIjp7InZhbHVlIjowLjB9LCJraW5kIjoiYWJzb2x1dGUiLCJuYW1lIjoic2FtcGxlX2dhdWdlIiwidGFncyI6e319fV0=\n",
            output.as_slice(),
        );

        Ok(())
    }

    #[test]
    fn test_serialize_and_split_list_small() -> Result<(), Error> {
        let small = (0..10).collect::<Vec<_>>();
        let chunks = serialize_and_split_list(small.iter(), 256)?;

        assert_eq!(chunks.len(), 1);
        assert!(chunks[0].len() <= 256);
        assert_eq!(
            serde_json::from_slice::<Vec<usize>>(&base64::decode(&chunks[0])?)?,
            small,
        );

        Ok(())
    }

    #[test]
    fn test_serialize_and_split_list_long() -> Result<(), Error> {
        let small = (0..100).collect::<Vec<_>>();
        let chunks = serialize_and_split_list(small.iter(), 256)?;

        assert_eq!(chunks.len(), 2);
        assert!(chunks[0].len() >= 256);
        assert!(chunks[1].len() <= 256);
        assert_eq!(
            serde_json::from_slice::<Vec<usize>>(&base64::decode(&chunks[0])?)?,
            (0..=67).collect::<Vec<_>>(),
        );
        assert_eq!(
            serde_json::from_slice::<Vec<usize>>(&base64::decode(&chunks[1])?)?,
            (68..100).collect::<Vec<_>>(),
        );

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


<commit_before>use super::{Index, Range};
use small;
use std::{cmp::Ordering, u8};

fn numeric_range<'a>(
    start: isize,
    end: isize,
    step: isize,
    inclusive: bool,
    nb_digits: usize,
) -> Option<Box<Iterator<Item = small::String> + 'a>> {
    let end = if start < end && inclusive {
        end + 1
    } else if start > end && inclusive {
        end - 1
    } else {
        end
    };

    if step == 0 || (start < end && step < 0) || (start > end && step > 0) {
        None
    } else {
        let (x, y, ordering) = if start < end {
            (start, end, Ordering::Greater)
        } else {
            (end, start, Ordering::Less)
        };

        let iter = (x..y).scan(start, move |index, _| {
            if end.cmp(index) == ordering {
                let index_holder = *index;
                *index += step; // This step adds
                Some(format!("{:0width$}", index_holder, width = nb_digits).into())
            } else {
                None
            }
        });

        Some(Box::new(iter))
    }
}

#[inline]
fn byte_is_valid_range(b: u8) -> bool { (b >= b'a' && b <= b'z') || (b >= b'A' && b <= b'Z') }

fn char_range<'a>(
    start: u8,
    end: u8,
    step: isize,
    inclusive: bool,
) -> Option<Box<Iterator<Item = small::String> + 'a>> {
    let char_step = step.checked_abs()?;
    if !byte_is_valid_range(start)
        || !byte_is_valid_range(end)
        || step == 0
        || char_step > u8::MAX as isize
    {
        return None;
    }

    let end = if start < end && inclusive {
        end + 1
    } else if start > end && inclusive {
        end - 1
    } else {
        end
    };
    let (x, y, ordering) =
        if start < end { (start, end, Ordering::Greater) } else { (end, start, Ordering::Less) };

    let iter = (x..y).scan(start, move |index, _| {
        if end.cmp(index) == ordering {
            let index_holder = *index;
            *index = match ordering {
                Ordering::Greater => index.wrapping_add(char_step as u8),
                Ordering::Less => index.wrapping_sub(char_step as u8),
                _ => unreachable!(),
            };
            Some((index_holder as char).to_string().into())
        } else {
            None
        }
    });

    Some(Box::new(iter))
}

fn count_minimum_digits(a: &str) -> usize {
    for c in a.chars() {
        match c {
            '-' => (),
            '0' => return a.len(),
            '1'...'9' => break,
            _ => panic!("count_minimum_digits should only be called for a valid number."),
        }
    }
    0
}

fn finish(
    inclusive: bool,
    start_str: &str,
    end_str: &str,
    step: isize,
) -> Option<Box<Iterator<Item = small::String>>> {
    if let (Ok(start), Ok(end)) = (start_str.parse::<isize>(), end_str.parse::<isize>()) {
        let step = if step == 1 && start >= end { -step } else { step };
        let nb_digits = usize::max(count_minimum_digits(start_str), count_minimum_digits(end_str));
        numeric_range(start, end, step, inclusive, nb_digits)
    } else if start_str.len() != 1 || end_str.len() != 1 {
        None
    } else {
        char_range(start_str.as_bytes()[0], end_str.as_bytes()[0], step, inclusive)
    }
}

// TODO: Make this an iterator structure.
// In a range we allow the following syntax:
//      Exclusive nonstepped: {start..end}
//      Inclusive nonstepped: {start...end}
//      Exclusive stepped: {start..step..end}
//      Inclusive stepped: {start..step...end}
pub fn parse_range(input: &str) -> Option<Box<Iterator<Item = small::String>>> {
    let mut parts = input.split("..").collect::<Vec<_>>();
    let len = parts.len();

    // if the last separator contains three dots, this can only be an inclusive range
    let inclusive = parts.last()?.starts_with(|c| c == '.' || c == '=');
    if inclusive {
        parts[len - 1] = parts[len - 1].trim_start_matches(|c| c == '.' || c == '=');
    }

    match len {
        // two parts means unstepped range
        2 => finish(inclusive, parts[0], parts[1], 1),
        // middle string contains the step size
        3 => finish(inclusive, parts[0], parts[2], parts[1].parse::<isize>().ok()?),
        // not a valid byte for ranges
        _ => None,
    }
}

pub fn parse_index_range(input: &str) -> Option<Range> {
    let mut bytes_iterator = input.bytes().enumerate();
    while let Some((id, byte)) = bytes_iterator.next() {
        match byte {
            b'0'...b'9' | b'-' => continue,
            b'.' => {
                let first = &input[..id];

                let mut dots = 1;
                let mut last_byte = 0;

                for (_, byte) in bytes_iterator {
                    last_byte = byte;
                    if byte == b'.' {
                        dots += 1
                    } else {
                        break;
                    }
                }

                if dots < 2 || dots > 3 {
                    break;
                }

                let inclusive = dots == 3 || (dots == 2 && last_byte == b'=');

                let end = &input[id + if inclusive { 3 } else { 2 }..];

                return if first.is_empty() && !end.is_empty() {
                    end.parse::<isize>().map(|end| Range::to(Index::new(end))).ok()
                } else if end.is_empty() {
                    first.parse::<isize>().map(|start| Range::from(Index::new(start))).ok()
                } else {
                    first
                        .parse::<isize>()
                        .and_then(|start| end.parse::<isize>().map(|end| (start, end)))
                        .map(|(start, end)| {
                            if inclusive {
                                Range::inclusive(Index::new(start), Index::new(end))
                            } else {
                                Range::exclusive(Index::new(start), Index::new(end))
                            }
                        })
                        .ok()
                };
            }
            _ => break,
        }
    }

    None
}
<commit_msg>Refactor range parsing<commit_after>use super::{Index, Range};
use small;
use std::{cmp::Ordering, u8};

fn numeric_range<'a>(
    start: isize,
    end: isize,
    step: isize,
    inclusive: bool,
    nb_digits: usize,
) -> Option<Box<Iterator<Item = small::String> + 'a>> {
    let end = if start < end && inclusive {
        end + 1
    } else if start > end && inclusive {
        end - 1
    } else {
        end
    };

    if step == 0 || (start < end && step < 0) || (start > end && step > 0) {
        None
    } else {
        let (x, y, ordering) = if start < end {
            (start, end, Ordering::Greater)
        } else {
            (end, start, Ordering::Less)
        };

        let iter = (x..y).scan(start, move |index, _| {
            if end.cmp(index) == ordering {
                let index_holder = *index;
                *index += step; // This step adds
                Some(format!("{:0width$}", index_holder, width = nb_digits).into())
            } else {
                None
            }
        });

        Some(Box::new(iter))
    }
}

fn char_range<'a>(
    start: u8,
    mut end: u8,
    step: isize,
    inclusive: bool,
) -> Option<Box<Iterator<Item = small::String> + 'a>> {
    if !start.is_ascii_alphabetic()
        || !end.is_ascii_alphabetic()
        || step == 0
        || step.checked_abs()? > u8::MAX as isize
    {
        return None;
    }

    if (start < end && inclusive) || (start > end && !inclusive) {
        end += 1;
    }

    if start < end {
        Some(Box::new((start..end).step_by(step as usize).map(|x| (x as char).to_string().into())))
    } else {
        Some(Box::new(
            (end..=start).rev().step_by(step as usize).map(|x| (x as char).to_string().into()),
        ))
    }
}

fn count_minimum_digits(a: &str) -> usize {
    match a.bytes().find(|&c| c != b'-') {
        Some(b'0') => a.len(),
        Some(b'1'...b'9') => 0,
        Some(_) => panic!("count_minimum_digits should only be called for a valid number."),
        None => 0,
    }
}

fn finish(
    inclusive: bool,
    start_str: &str,
    end_str: &str,
    step: isize,
) -> Option<Box<Iterator<Item = small::String>>> {
    if let (Ok(start), Ok(end)) = (start_str.parse::<isize>(), end_str.parse::<isize>()) {
        let step = if step == 1 && start >= end { -step } else { step };
        let nb_digits = usize::max(count_minimum_digits(start_str), count_minimum_digits(end_str));
        numeric_range(start, end, step, inclusive, nb_digits)
    } else if start_str.len() != 1 || end_str.len() != 1 {
        None
    } else {
        char_range(start_str.as_bytes()[0], end_str.as_bytes()[0], step, inclusive)
    }
}

// TODO: Make this an iterator structure.
// In a range we allow the following syntax:
//      Exclusive nonstepped: {start..end}
//      Inclusive nonstepped: {start...end}
//      Exclusive stepped: {start..step..end}
//      Inclusive stepped: {start..step...end}
pub fn parse_range(input: &str) -> Option<Box<Iterator<Item = small::String>>> {
    let mut parts = input.split("..").collect::<Vec<_>>();
    let len = parts.len();

    // if the last separator contains three dots, this can only be an inclusive range
    let inclusive = parts.last()?.starts_with(|c| c == '.' || c == '=');
    if inclusive {
        parts[len - 1] = parts[len - 1].trim_start_matches(|c| c == '.' || c == '=');
    }

    match len {
        // two parts means unstepped range
        2 => finish(inclusive, parts[0], parts[1], 1),
        // middle string contains the step size
        3 => finish(inclusive, parts[0], parts[2], parts[1].parse::<isize>().ok()?),
        // not a valid byte for ranges
        _ => None,
    }
}

pub fn parse_index_range(input: &str) -> Option<Range> {
    let mut parts = input.splitn(2, "..");
    let first = parts.next()?;
    let mut end = parts.next()?;

    if first.is_empty() && !end.is_empty() {
        end.parse::<isize>().map(|end| Range::to(Index::new(end))).ok()
    } else if !end.is_empty() {
        let inclusive = end.starts_with('.') || end.starts_with('=');
        if inclusive {
            end = &end[1..];
        }

        let start = first.parse::<isize>().ok()?;
        let end = end.parse::<isize>().ok()?;
        if inclusive {
            Some(Range::inclusive(Index::new(start), Index::new(end)))
        } else {
            Some(Range::exclusive(Index::new(start), Index::new(end)))
        }
    } else {
        first.parse::<isize>().map(|start| Range::from(Index::new(start))).ok()
    }
}
<|endoftext|>


<commit_before>/// Logs a message at the info level.
///
/// # Examples
///
/// ```edition2018
/// use log::info;
///
/// # fn main() {
/// # struct Connection { port: u32, speed: f32 }
/// let conn_info = Connection { port: 40, speed: 3.20 };
///
/// info!("Connected to port {} at {} Mb/s", conn_info.port, conn_info.speed);
/// info!(target: "connection_events", "Successfull connection, port: {}, speed: {}",
///       conn_info.port, conn_info.speed);
/// # }
/// ```
#[macro_export]
macro_rules! info {
    ($($args:tt)+) => {
        log::info!("{}", paris::formatter::colorize_string(format!($($args)*)));
    };
}

/// Logs a message at the debug level.
///
/// # Examples
///
/// ```edition2018
/// use log::debug;
///
/// # fn main() {
/// # struct Position { x: f32, y: f32 }
/// let pos = Position { x: 3.234, y: -1.223 };
///
/// debug!("New position: x: {}, y: {}", pos.x, pos.y);
/// debug!(target: "app_events", "New position: x: {}, y: {}", pos.x, pos.y);
/// # }
/// ```
#[macro_export]
macro_rules! debug {
    ($($args:tt)+) => {
        log::debug!("{}", paris::formatter::colorize_string(format!($($args)*)));
    };
}

/// Logs a message at the trace level.
///
/// # Examples
///
/// ```edition2018
/// use log::trace;
///
/// # fn main() {
/// # struct Position { x: f32, y: f32 }
/// let pos = Position { x: 3.234, y: -1.223 };
///
/// trace!("Position is: x: {}, y: {}", pos.x, pos.y);
/// trace!(target: "app_events", "x is {} and y is {}",
///        if pos.x >= 0.0 { "positive" } else { "negative" },
///        if pos.y >= 0.0 { "positive" } else { "negative" });
/// # }
/// ```
#[macro_export]
macro_rules! trace {
    ($($args:tt)+) => {
        log::trace!("{}", paris::formatter::colorize_string(format!($($args)*)));
    };
}

/// Logs a message at the warn level.
///
/// # Examples
///
/// ```edition2018
/// use log::warn;
///
/// # fn main() {
/// let warn_description = "Invalid Input";
///
/// warn!("Warning! {}!", warn_description);
/// warn!(target: "input_events", "App received warning: {}", warn_description);
/// # }
/// ```
#[macro_export]
macro_rules! warn {
    ($($args:tt)+) => {
        log::warn!("{}", paris::formatter::colorize_string(format!($($args)*)));
    };
}

/// Logs a message at the error level.
///
/// # Examples
///
/// ```edition2018
/// use log::error;
///
/// # fn main() {
/// let (err_info, port) = ("No connection", 22);
///
/// error!("Error: {} on port {}", err_info, port);
/// error!(target: "app_events", "App Error: {}, Port: {}", err_info, 22);
/// # }
/// ```
#[macro_export]
macro_rules! error {
    ($($args:tt)+) => {
        log::error!("{}", paris::formatter::colorize_string(format!($($args)*)));
    };
}
<commit_msg>paris_macros: add information about colorize_string<commit_after>/// Logs a message at the info level.
///
/// Passed data uses a colorize_string formatter from a `paris` crate, so it can
/// contains special tags for controlling ANSI colors and styles
/// More info: https://docs.rs/paris/1.5.7/paris/formatter/fn.colorize_string.html
///
/// # Examples
///
/// ```edition2018
/// use log::info;
///
/// # fn main() {
/// # struct Connection { port: u32, speed: f32 }
/// let conn_info = Connection { port: 40, speed: 3.20 };
///
/// info!("Connected to port {} at {} Mb/s", conn_info.port, conn_info.speed);
/// info!(target: "connection_events", "Successfull connection, port: {}, speed: {}",
///       conn_info.port, conn_info.speed);
/// # }
/// ```
#[macro_export]
macro_rules! info {
    ($($args:tt)+) => {
        log::info!("{}", paris::formatter::colorize_string(format!($($args)*)));
    };
}

/// Logs a message at the debug level.
///
/// Passed data uses a colorize_string formatter from a `paris` crate, so it can
/// contains special tags for controlling ANSI colors and styles
/// More info: https://docs.rs/paris/1.5.7/paris/formatter/fn.colorize_string.html
///
/// # Examples
///
/// ```edition2018
/// use log::debug;
///
/// # fn main() {
/// # struct Position { x: f32, y: f32 }
/// let pos = Position { x: 3.234, y: -1.223 };
///
/// debug!("New position: x: {}, y: {}", pos.x, pos.y);
/// debug!(target: "app_events", "New position: x: {}, y: {}", pos.x, pos.y);
/// # }
/// ```
#[macro_export]
macro_rules! debug {
    ($($args:tt)+) => {
        log::debug!("{}", paris::formatter::colorize_string(format!($($args)*)));
    };
}

/// Logs a message at the trace level.
///
/// Passed data uses a colorize_string formatter from a `paris` crate, so it can
/// contains special tags for controlling ANSI colors and styles
/// More info: https://docs.rs/paris/1.5.7/paris/formatter/fn.colorize_string.html
///
/// # Examples
///
/// ```edition2018
/// use log::trace;
///
/// # fn main() {
/// # struct Position { x: f32, y: f32 }
/// let pos = Position { x: 3.234, y: -1.223 };
///
/// trace!("Position is: x: {}, y: {}", pos.x, pos.y);
/// trace!(target: "app_events", "x is {} and y is {}",
///        if pos.x >= 0.0 { "positive" } else { "negative" },
///        if pos.y >= 0.0 { "positive" } else { "negative" });
/// # }
/// ```
#[macro_export]
macro_rules! trace {
    ($($args:tt)+) => {
        log::trace!("{}", paris::formatter::colorize_string(format!($($args)*)));
    };
}

/// Logs a message at the warn level.
///
/// Passed data uses a colorize_string formatter from a `paris` crate, so it can
/// contains special tags for controlling ANSI colors and styles
/// More info: https://docs.rs/paris/1.5.7/paris/formatter/fn.colorize_string.html
///
/// # Examples
///
/// ```edition2018
/// use log::warn;
///
/// # fn main() {
/// let warn_description = "Invalid Input";
///
/// warn!("Warning! {}!", warn_description);
/// warn!(target: "input_events", "App received warning: {}", warn_description);
/// # }
/// ```
#[macro_export]
macro_rules! warn {
    ($($args:tt)+) => {
        log::warn!("{}", paris::formatter::colorize_string(format!($($args)*)));
    };
}

/// Logs a message at the error level.
///
/// Passed data uses a colorize_string formatter from a `paris` crate, so it can
/// contains special tags for controlling ANSI colors and styles
/// More info: https://docs.rs/paris/1.5.7/paris/formatter/fn.colorize_string.html
///
/// # Examples
///
/// ```edition2018
/// use log::error;
///
/// # fn main() {
/// let (err_info, port) = ("No connection", 22);
///
/// error!("Error: {} on port {}", err_info, port);
/// error!(target: "app_events", "App Error: {}, Port: {}", err_info, 22);
/// # }
/// ```
#[macro_export]
macro_rules! error {
    ($($args:tt)+) => {
        log::error!("{}", paris::formatter::colorize_string(format!($($args)*)));
    };
}
<|endoftext|>


<commit_before>
use {Incoming};
use super::{IntoNeighbors, IntoNeighborsDirected, Visitable, VisitMap};
use super::{GraphRef, Reversed, IntoNodeIdentifiers};
use std::collections::VecDeque;

/// Visit nodes of a graph in a depth-first-search (DFS) emitting nodes in
/// preorder (when they are first discovered).
///
/// The traversal starts at a given node and only traverses nodes reachable
/// from it.
///
/// `Dfs` is not recursive.
///
/// `Dfs` does not itself borrow the graph, and because of this you can run
/// a traversal over a graph while still retaining mutable access to it, if you
/// use it like the following example:
///
/// ```
/// use petgraph::Graph;
/// use petgraph::visit::Dfs;
///
/// let mut graph = Graph::<_,()>::new();
/// let a = graph.add_node(0);
///
/// let mut dfs = Dfs::new(&graph, a);
/// while let Some(nx) = dfs.next(&graph) {
///     // we can access `graph` mutably here still
///     graph[nx] += 1;
/// }
///
/// assert_eq!(graph[a], 1);
/// ```
///
/// **Note:** The algorithm may not behave correctly if nodes are removed
/// during iteration. It may not necessarily visit added nodes or edges.
#[derive(Clone, Debug)]
pub struct Dfs<N, VM> {
    /// The stack of nodes to visit
    pub stack: Vec<N>,
    /// The map of discovered nodes
    pub discovered: VM,
}

impl<N, VM> Dfs<N, VM>
    where N: Copy + PartialEq,
          VM: VisitMap<N>,
{
    /// Create a new **Dfs**, using the graph's visitor map, and put **start**
    /// in the stack of nodes to visit.
    pub fn new<G>(graph: G, start: N) -> Self
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        let mut dfs = Dfs::empty(graph);
        dfs.move_to(start);
        dfs
    }

    /// Create a `Dfs` from a vector and a visit map
    pub fn from_parts(stack: Vec<N>, discovered: VM) -> Self {
        Dfs {
            stack: stack,
            discovered: discovered,
        }
    }

    /// Clear the visit state
    pub fn reset<G>(&mut self, graph: G)
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        graph.reset_map(&mut self.discovered);
        self.stack.clear();
    }

    /// Create a new **Dfs** using the graph's visitor map, and no stack.
    pub fn empty<G>(graph: G) -> Self
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        Dfs {
            stack: Vec::new(),
            discovered: graph.visit_map(),
        }
    }

    /// Keep the discovered map, but clear the visit stack and restart
    /// the dfs from a particular node.
    pub fn move_to(&mut self, start: N)
    {
        self.discovered.visit(start);
        self.stack.clear();
        self.stack.push(start);
    }

    /// Return the next node in the dfs, or **None** if the traversal is done.
    pub fn next<G>(&mut self, graph: G) -> Option<N>
        where G: IntoNeighbors<NodeId=N>,
    {
        while let Some(node) = self.stack.pop() {
            for succ in graph.neighbors(node) {
                if self.discovered.visit(succ) {
                    self.stack.push(succ);
                }
            }

            return Some(node);
        }
        None
    }
}

/// Visit nodes in a depth-first-search (DFS) emitting nodes in postorder
/// (each node after all its decendants have been emitted).
///
/// `DfsPostOrder` is not recursive.
///
/// The traversal starts at a given node and only traverses nodes reachable
/// from it.
#[derive(Clone, Debug)]
pub struct DfsPostOrder<N, VM> {
    /// The stack of nodes to visit
    pub stack: Vec<N>,
    /// The map of discovered nodes
    pub discovered: VM,
    /// The map of finished nodes
    pub finished: VM,
}

impl<N, VM> DfsPostOrder<N, VM>
    where N: Copy + PartialEq,
          VM: VisitMap<N>,
{
    /// Create a new `DfsPostOrder` using the graph's visitor map, and put
    /// `start` in the stack of nodes to visit.
    pub fn new<G>(graph: G, start: N) -> Self
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        let mut dfs = Self::empty(graph);
        dfs.move_to(start);
        dfs
    }

    /// Create a new `DfsPostOrder` using the graph's visitor map, and no stack.
    pub fn empty<G>(graph: G) -> Self
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        DfsPostOrder {
            stack: Vec::new(),
            discovered: graph.visit_map(),
            finished: graph.visit_map(),
        }
    }

    /// Clear the visit state
    pub fn reset<G>(&mut self, graph: G)
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        graph.reset_map(&mut self.discovered);
        graph.reset_map(&mut self.finished);
        self.stack.clear();
    }

    /// Keep the discovered and finished map, but clear the visit stack and restart
    /// the dfs from a particular node.
    pub fn move_to(&mut self, start: N)
    {
        self.stack.clear();
        self.stack.push(start);
    }

    /// Return the next node in the traversal, or `None` if the traversal is done.
    pub fn next<G>(&mut self, graph: G) -> Option<N>
        where G: IntoNeighbors<NodeId=N>,
    {
        while let Some(&nx) = self.stack.last() {
            if self.discovered.visit(nx) {
                // First time visiting `nx`: Push neighbors, don't pop `nx`
                for succ in graph.neighbors(nx) {
                    if !self.discovered.is_visited(&succ) {
                        self.stack.push(succ);
                    }
                }
            } else {
                self.stack.pop();
                if self.finished.visit(nx) {
                    // Second time: All reachable nodes must have been finished
                    return Some(nx);
                }
            }
        }
        None
    }
}

/// A breadth first search (BFS) of a graph.
///
/// The traversal starts at a given node and only traverses nodes reachable
/// from it.
///
/// `Bfs` is not recursive.
///
/// `Bfs` does not itself borrow the graph, and because of this you can run
/// a traversal over a graph while still retaining mutable access to it, if you
/// use it like the following example:
///
/// ```
/// use petgraph::Graph;
/// use petgraph::visit::Bfs;
///
/// let mut graph = Graph::<_,()>::new();
/// let a = graph.add_node(0);
///
/// let mut bfs = Bfs::new(&graph, a);
/// while let Some(nx) = bfs.next(&graph) {
///     // we can access `graph` mutably here still
///     graph[nx] += 1;
/// }
///
/// assert_eq!(graph[a], 1);
/// ```
///
/// **Note:** The algorithm may not behave correctly if nodes are removed
/// during iteration. It may not necessarily visit added nodes or edges.
#[derive(Clone)]
pub struct Bfs<N, VM> {
    /// The queue of nodes to visit
    pub stack: VecDeque<N>,
    /// The map of discovered nodes
    pub discovered: VM,
}

impl<N, VM> Bfs<N, VM>
    where N: Copy + PartialEq,
          VM: VisitMap<N>,
{
    /// Create a new **Bfs**, using the graph's visitor map, and put **start**
    /// in the stack of nodes to visit.
    pub fn new<G>(graph: G, start: N) -> Self
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        let mut discovered = graph.visit_map();
        discovered.visit(start);
        let mut stack = VecDeque::new();
        stack.push_front(start);
        Bfs {
            stack: stack,
            discovered: discovered,
        }
    }

    /// Return the next node in the dfs, or **None** if the traversal is done.
    pub fn next<G>(&mut self, graph: G) -> Option<N>
        where G: IntoNeighbors<NodeId=N>
    {
        while let Some(node) = self.stack.pop_front() {
            for succ in graph.neighbors(node) {
                if self.discovered.visit(succ) {
                    self.stack.push_back(succ);
                }
            }

            return Some(node);
        }
        None
    }

}

/// A topological order traversal for a graph.
///
/// **Note** that `Topo` only visits nodes that are not part of cycles,
/// i.e. nodes in a true DAG. Use other visitors like `DfsPostOrder` or
/// algorithms like kosaraju_scc to handle graphs with possible cycles.
#[derive(Clone)]
pub struct Topo<N, VM> {
    tovisit: Vec<N>,
    ordered: VM,
}

impl<N, VM> Topo<N, VM>
    where N: Copy + PartialEq,
          VM: VisitMap<N>,
{
    /// Create a new `Topo`, using the graph's visitor map, and put all
    /// initial nodes in the to visit list.
    pub fn new<G>(graph: G) -> Self
        where G: IntoNodeIdentifiers + IntoNeighborsDirected + Visitable<NodeId=N, Map=VM>,
    {
        let mut topo = Self::empty(graph);
        topo.extend_with_initials(graph);
        topo
    }

    fn extend_with_initials<G>(&mut self, g: G)
        where G: IntoNodeIdentifiers + IntoNeighborsDirected<NodeId=N>,
    {
        // find all initial nodes (nodes without incoming edges)
        self.tovisit.extend(g.node_identifiers()
                             .filter(move |&a| g.neighbors_directed(a, Incoming).next().is_none()));
    }

    /* Private until it has a use */
    /// Create a new `Topo`, using the graph's visitor map with *no* starting
    /// index specified.
    fn empty<G>(graph: G) -> Self
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        Topo {
            ordered: graph.visit_map(),
            tovisit: Vec::new(),
        }
    }

    /// Clear visited state, and put all initial nodes in the to visit list.
    pub fn reset<G>(&mut self, graph: G)
        where G: IntoNodeIdentifiers + IntoNeighborsDirected + Visitable<NodeId=N, Map=VM>,
    {
        graph.reset_map(&mut self.ordered);
        self.tovisit.clear();
        self.extend_with_initials(graph);
    }

    /// Return the next node in the current topological order traversal, or
    /// `None` if the traversal is at the end.
    ///
    /// *Note:* The graph may not have a complete topological order, and the only
    /// way to know is to run the whole traversal and make sure it visits every node.
    pub fn next<G>(&mut self, g: G) -> Option<N>
        where G: IntoNeighborsDirected + Visitable<NodeId=N, Map=VM>,
    {
        // Take an unvisited element and find which of its neighbors are next
        while let Some(nix) = self.tovisit.pop() {
            if self.ordered.is_visited(&nix) {
                continue;
            }
            self.ordered.visit(nix);
            for neigh in g.neighbors(nix) {
                // Look at each neighbor, and those that only have incoming edges
                // from the already ordered list, they are the next to visit.
                if Reversed(g).neighbors(neigh).all(|b| self.ordered.is_visited(&b)) {
                    self.tovisit.push(neigh);
                }
            }
            return Some(nix);
        }
        None
    }
}

/// A walker is a traversal state, but where part of the traversal
/// information is supplied manually to each next call.
///
/// This for example allows graph traversals that don't hold a borrow of the
/// graph they are traversing.
pub trait Walker<Context> {
    type Item;
    /// Advance to the next item
    fn walk_next(&mut self, context: Context) -> Option<Self::Item>;

    /// Create an iterator out of the walker and given `context`.
    fn iter(self, context: Context) -> WalkerIter<Self, Context>
        where Self: Sized,
              Context: Clone,
    {
        WalkerIter {
            walker: self,
            context: context,
        }
    }
}

/// A walker and its context wrapped into an iterator.
#[derive(Clone, Debug)]
pub struct WalkerIter<W, C> {
    walker: W,
    context: C,
}

impl<W, C> WalkerIter<W, C>
    where W: Walker<C>,
          C: Clone,
{
    pub fn context(&self) -> C {
        self.context.clone()
    }

    pub fn inner_ref(&self) -> &W {
        &self.walker
    }

    pub fn inner_mut(&mut self) -> &mut W {
        &mut self.walker
    }
}

impl<W, C> Iterator for WalkerIter<W, C>
    where W: Walker<C>,
          C: Clone,
{
    type Item = W::Item;
    fn next(&mut self) -> Option<Self::Item> {
        self.walker.walk_next(self.context.clone())
    }
}

impl<G> Walker<G> for Dfs<G::NodeId, G::Map>
    where G: IntoNeighbors + Visitable
{
    type Item = G::NodeId;
    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
        self.next(context)
    }
}

impl<G> Walker<G> for DfsPostOrder<G::NodeId, G::Map>
    where G: IntoNeighbors + Visitable
{
    type Item = G::NodeId;
    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
        self.next(context)
    }
}

impl<G> Walker<G> for Bfs<G::NodeId, G::Map>
    where G: IntoNeighbors + Visitable
{
    type Item = G::NodeId;
    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
        self.next(context)
    }
}

impl<G> Walker<G> for Topo<G::NodeId, G::Map>
    where G: IntoNeighborsDirected + Visitable,
{
    type Item = G::NodeId;
    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
        self.next(context)
    }
}
<commit_msg>Fix doc for Bfs::next() (#134)<commit_after>
use {Incoming};
use super::{IntoNeighbors, IntoNeighborsDirected, Visitable, VisitMap};
use super::{GraphRef, Reversed, IntoNodeIdentifiers};
use std::collections::VecDeque;

/// Visit nodes of a graph in a depth-first-search (DFS) emitting nodes in
/// preorder (when they are first discovered).
///
/// The traversal starts at a given node and only traverses nodes reachable
/// from it.
///
/// `Dfs` is not recursive.
///
/// `Dfs` does not itself borrow the graph, and because of this you can run
/// a traversal over a graph while still retaining mutable access to it, if you
/// use it like the following example:
///
/// ```
/// use petgraph::Graph;
/// use petgraph::visit::Dfs;
///
/// let mut graph = Graph::<_,()>::new();
/// let a = graph.add_node(0);
///
/// let mut dfs = Dfs::new(&graph, a);
/// while let Some(nx) = dfs.next(&graph) {
///     // we can access `graph` mutably here still
///     graph[nx] += 1;
/// }
///
/// assert_eq!(graph[a], 1);
/// ```
///
/// **Note:** The algorithm may not behave correctly if nodes are removed
/// during iteration. It may not necessarily visit added nodes or edges.
#[derive(Clone, Debug)]
pub struct Dfs<N, VM> {
    /// The stack of nodes to visit
    pub stack: Vec<N>,
    /// The map of discovered nodes
    pub discovered: VM,
}

impl<N, VM> Dfs<N, VM>
    where N: Copy + PartialEq,
          VM: VisitMap<N>,
{
    /// Create a new **Dfs**, using the graph's visitor map, and put **start**
    /// in the stack of nodes to visit.
    pub fn new<G>(graph: G, start: N) -> Self
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        let mut dfs = Dfs::empty(graph);
        dfs.move_to(start);
        dfs
    }

    /// Create a `Dfs` from a vector and a visit map
    pub fn from_parts(stack: Vec<N>, discovered: VM) -> Self {
        Dfs {
            stack: stack,
            discovered: discovered,
        }
    }

    /// Clear the visit state
    pub fn reset<G>(&mut self, graph: G)
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        graph.reset_map(&mut self.discovered);
        self.stack.clear();
    }

    /// Create a new **Dfs** using the graph's visitor map, and no stack.
    pub fn empty<G>(graph: G) -> Self
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        Dfs {
            stack: Vec::new(),
            discovered: graph.visit_map(),
        }
    }

    /// Keep the discovered map, but clear the visit stack and restart
    /// the dfs from a particular node.
    pub fn move_to(&mut self, start: N)
    {
        self.discovered.visit(start);
        self.stack.clear();
        self.stack.push(start);
    }

    /// Return the next node in the dfs, or **None** if the traversal is done.
    pub fn next<G>(&mut self, graph: G) -> Option<N>
        where G: IntoNeighbors<NodeId=N>,
    {
        while let Some(node) = self.stack.pop() {
            for succ in graph.neighbors(node) {
                if self.discovered.visit(succ) {
                    self.stack.push(succ);
                }
            }

            return Some(node);
        }
        None
    }
}

/// Visit nodes in a depth-first-search (DFS) emitting nodes in postorder
/// (each node after all its decendants have been emitted).
///
/// `DfsPostOrder` is not recursive.
///
/// The traversal starts at a given node and only traverses nodes reachable
/// from it.
#[derive(Clone, Debug)]
pub struct DfsPostOrder<N, VM> {
    /// The stack of nodes to visit
    pub stack: Vec<N>,
    /// The map of discovered nodes
    pub discovered: VM,
    /// The map of finished nodes
    pub finished: VM,
}

impl<N, VM> DfsPostOrder<N, VM>
    where N: Copy + PartialEq,
          VM: VisitMap<N>,
{
    /// Create a new `DfsPostOrder` using the graph's visitor map, and put
    /// `start` in the stack of nodes to visit.
    pub fn new<G>(graph: G, start: N) -> Self
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        let mut dfs = Self::empty(graph);
        dfs.move_to(start);
        dfs
    }

    /// Create a new `DfsPostOrder` using the graph's visitor map, and no stack.
    pub fn empty<G>(graph: G) -> Self
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        DfsPostOrder {
            stack: Vec::new(),
            discovered: graph.visit_map(),
            finished: graph.visit_map(),
        }
    }

    /// Clear the visit state
    pub fn reset<G>(&mut self, graph: G)
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        graph.reset_map(&mut self.discovered);
        graph.reset_map(&mut self.finished);
        self.stack.clear();
    }

    /// Keep the discovered and finished map, but clear the visit stack and restart
    /// the dfs from a particular node.
    pub fn move_to(&mut self, start: N)
    {
        self.stack.clear();
        self.stack.push(start);
    }

    /// Return the next node in the traversal, or `None` if the traversal is done.
    pub fn next<G>(&mut self, graph: G) -> Option<N>
        where G: IntoNeighbors<NodeId=N>,
    {
        while let Some(&nx) = self.stack.last() {
            if self.discovered.visit(nx) {
                // First time visiting `nx`: Push neighbors, don't pop `nx`
                for succ in graph.neighbors(nx) {
                    if !self.discovered.is_visited(&succ) {
                        self.stack.push(succ);
                    }
                }
            } else {
                self.stack.pop();
                if self.finished.visit(nx) {
                    // Second time: All reachable nodes must have been finished
                    return Some(nx);
                }
            }
        }
        None
    }
}

/// A breadth first search (BFS) of a graph.
///
/// The traversal starts at a given node and only traverses nodes reachable
/// from it.
///
/// `Bfs` is not recursive.
///
/// `Bfs` does not itself borrow the graph, and because of this you can run
/// a traversal over a graph while still retaining mutable access to it, if you
/// use it like the following example:
///
/// ```
/// use petgraph::Graph;
/// use petgraph::visit::Bfs;
///
/// let mut graph = Graph::<_,()>::new();
/// let a = graph.add_node(0);
///
/// let mut bfs = Bfs::new(&graph, a);
/// while let Some(nx) = bfs.next(&graph) {
///     // we can access `graph` mutably here still
///     graph[nx] += 1;
/// }
///
/// assert_eq!(graph[a], 1);
/// ```
///
/// **Note:** The algorithm may not behave correctly if nodes are removed
/// during iteration. It may not necessarily visit added nodes or edges.
#[derive(Clone)]
pub struct Bfs<N, VM> {
    /// The queue of nodes to visit
    pub stack: VecDeque<N>,
    /// The map of discovered nodes
    pub discovered: VM,
}

impl<N, VM> Bfs<N, VM>
    where N: Copy + PartialEq,
          VM: VisitMap<N>,
{
    /// Create a new **Bfs**, using the graph's visitor map, and put **start**
    /// in the stack of nodes to visit.
    pub fn new<G>(graph: G, start: N) -> Self
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        let mut discovered = graph.visit_map();
        discovered.visit(start);
        let mut stack = VecDeque::new();
        stack.push_front(start);
        Bfs {
            stack: stack,
            discovered: discovered,
        }
    }

    /// Return the next node in the bfs, or **None** if the traversal is done.
    pub fn next<G>(&mut self, graph: G) -> Option<N>
        where G: IntoNeighbors<NodeId=N>
    {
        while let Some(node) = self.stack.pop_front() {
            for succ in graph.neighbors(node) {
                if self.discovered.visit(succ) {
                    self.stack.push_back(succ);
                }
            }

            return Some(node);
        }
        None
    }

}

/// A topological order traversal for a graph.
///
/// **Note** that `Topo` only visits nodes that are not part of cycles,
/// i.e. nodes in a true DAG. Use other visitors like `DfsPostOrder` or
/// algorithms like kosaraju_scc to handle graphs with possible cycles.
#[derive(Clone)]
pub struct Topo<N, VM> {
    tovisit: Vec<N>,
    ordered: VM,
}

impl<N, VM> Topo<N, VM>
    where N: Copy + PartialEq,
          VM: VisitMap<N>,
{
    /// Create a new `Topo`, using the graph's visitor map, and put all
    /// initial nodes in the to visit list.
    pub fn new<G>(graph: G) -> Self
        where G: IntoNodeIdentifiers + IntoNeighborsDirected + Visitable<NodeId=N, Map=VM>,
    {
        let mut topo = Self::empty(graph);
        topo.extend_with_initials(graph);
        topo
    }

    fn extend_with_initials<G>(&mut self, g: G)
        where G: IntoNodeIdentifiers + IntoNeighborsDirected<NodeId=N>,
    {
        // find all initial nodes (nodes without incoming edges)
        self.tovisit.extend(g.node_identifiers()
                             .filter(move |&a| g.neighbors_directed(a, Incoming).next().is_none()));
    }

    /* Private until it has a use */
    /// Create a new `Topo`, using the graph's visitor map with *no* starting
    /// index specified.
    fn empty<G>(graph: G) -> Self
        where G: GraphRef + Visitable<NodeId=N, Map=VM>
    {
        Topo {
            ordered: graph.visit_map(),
            tovisit: Vec::new(),
        }
    }

    /// Clear visited state, and put all initial nodes in the to visit list.
    pub fn reset<G>(&mut self, graph: G)
        where G: IntoNodeIdentifiers + IntoNeighborsDirected + Visitable<NodeId=N, Map=VM>,
    {
        graph.reset_map(&mut self.ordered);
        self.tovisit.clear();
        self.extend_with_initials(graph);
    }

    /// Return the next node in the current topological order traversal, or
    /// `None` if the traversal is at the end.
    ///
    /// *Note:* The graph may not have a complete topological order, and the only
    /// way to know is to run the whole traversal and make sure it visits every node.
    pub fn next<G>(&mut self, g: G) -> Option<N>
        where G: IntoNeighborsDirected + Visitable<NodeId=N, Map=VM>,
    {
        // Take an unvisited element and find which of its neighbors are next
        while let Some(nix) = self.tovisit.pop() {
            if self.ordered.is_visited(&nix) {
                continue;
            }
            self.ordered.visit(nix);
            for neigh in g.neighbors(nix) {
                // Look at each neighbor, and those that only have incoming edges
                // from the already ordered list, they are the next to visit.
                if Reversed(g).neighbors(neigh).all(|b| self.ordered.is_visited(&b)) {
                    self.tovisit.push(neigh);
                }
            }
            return Some(nix);
        }
        None
    }
}

/// A walker is a traversal state, but where part of the traversal
/// information is supplied manually to each next call.
///
/// This for example allows graph traversals that don't hold a borrow of the
/// graph they are traversing.
pub trait Walker<Context> {
    type Item;
    /// Advance to the next item
    fn walk_next(&mut self, context: Context) -> Option<Self::Item>;

    /// Create an iterator out of the walker and given `context`.
    fn iter(self, context: Context) -> WalkerIter<Self, Context>
        where Self: Sized,
              Context: Clone,
    {
        WalkerIter {
            walker: self,
            context: context,
        }
    }
}

/// A walker and its context wrapped into an iterator.
#[derive(Clone, Debug)]
pub struct WalkerIter<W, C> {
    walker: W,
    context: C,
}

impl<W, C> WalkerIter<W, C>
    where W: Walker<C>,
          C: Clone,
{
    pub fn context(&self) -> C {
        self.context.clone()
    }

    pub fn inner_ref(&self) -> &W {
        &self.walker
    }

    pub fn inner_mut(&mut self) -> &mut W {
        &mut self.walker
    }
}

impl<W, C> Iterator for WalkerIter<W, C>
    where W: Walker<C>,
          C: Clone,
{
    type Item = W::Item;
    fn next(&mut self) -> Option<Self::Item> {
        self.walker.walk_next(self.context.clone())
    }
}

impl<G> Walker<G> for Dfs<G::NodeId, G::Map>
    where G: IntoNeighbors + Visitable
{
    type Item = G::NodeId;
    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
        self.next(context)
    }
}

impl<G> Walker<G> for DfsPostOrder<G::NodeId, G::Map>
    where G: IntoNeighbors + Visitable
{
    type Item = G::NodeId;
    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
        self.next(context)
    }
}

impl<G> Walker<G> for Bfs<G::NodeId, G::Map>
    where G: IntoNeighbors + Visitable
{
    type Item = G::NodeId;
    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
        self.next(context)
    }
}

impl<G> Walker<G> for Topo<G::NodeId, G::Map>
    where G: IntoNeighborsDirected + Visitable,
{
    type Item = G::NodeId;
    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
        self.next(context)
    }
}
<|endoftext|>


<commit_before>use tokens::*;

pub struct ParseTree<'a>
{
	pub root: Node<'a>,
}

impl<'a> ParseTree<'a>
{

}

pub struct Node<'a>
{
	pub token: Token<'a>,
	pub children: Vec<Node<'a>>
}

impl<'a> Node<'a>
{
	fn new(token: Token<'a>) -> Node<'a>
	{
		Node
		{
			token: token,
			children: Vec::new()
		}
	}

	fn has_children(&self) -> bool
	{
		self.children.len() > 0
	}

	fn add_child(&mut self, child: Node<'a>)
	{
		self.children.push(child);
	}

	fn value(&self) -> String
	{
		self.token.value.clone()
	}
}
<commit_msg>Type implementation for tree<commit_after>use tokens::*;

pub struct ParseTree<'a>
{
	pub root: Node<'a>,
}

impl<'a> ParseTree<'a>
{

}

pub struct Node<'a>
{
	pub token: Token<'a>,
	pub children: Vec<Node<'a>>
}

impl<'a> Node<'a>
{
	fn new(token: Token<'a>) -> Node<'a>
	{
		Node
		{
			token: token,
			children: Vec::new()
		}
	}

	fn has_children(&self) -> bool
	{
		self.children.len() > 0
	}

	fn add_child(&mut self, child: Node<'a>)
	{
		self.children.push(child);
	}

	fn value(&self) -> String
	{
		self.token.value.clone()
	}

	fn token_type(&self) -> String
	{
		self.token.name.to_string().clone()
	}
}
<|endoftext|>