<commit_before>//! This pass checks HIR bodies that may be evaluated at compile-time (e.g., `const`, `static`,
//! `const fn`) for structured control flow (e.g. `if`, `while`), which is forbidden in a const
//! context.
//!
//! By the time the MIR const-checker runs, these high-level constructs have been lowered to
//! control-flow primitives (e.g., `Goto`, `SwitchInt`), making it tough to properly attribute
//! errors. We still look for those primitives in the MIR const-checker to ensure nothing slips
//! through, but errors for structured control flow in a `const` should be emitted here.

use rustc::hir::def_id::DefId;
use rustc::hir::intravisit::{Visitor, NestedVisitorMap};
use rustc::hir::map::Map;
use rustc::hir;
use rustc::ty::TyCtxt;
use rustc::ty::query::Providers;
use syntax::ast::Mutability;
use syntax::span_err;
use syntax_pos::Span;
use rustc_error_codes::*;

use std::fmt;

#[derive(Copy, Clone)]
enum ConstKind {
    Static,
    StaticMut,
    ConstFn,
    Const,
    AnonConst,
}

impl ConstKind {
    fn for_body(body: &hir::Body, hir_map: &Map<'_>) -> Option<Self> {
        let is_const_fn = |id| hir_map.fn_sig_by_hir_id(id).unwrap().header.is_const();

        let owner = hir_map.body_owner(body.id());
        let const_kind = match hir_map.body_owner_kind(owner) {
            hir::BodyOwnerKind::Const => Self::Const,
            hir::BodyOwnerKind::Static(Mutability::Mutable) => Self::StaticMut,
            hir::BodyOwnerKind::Static(Mutability::Immutable) => Self::Static,

            hir::BodyOwnerKind::Fn if is_const_fn(owner) => Self::ConstFn,
            hir::BodyOwnerKind::Fn | hir::BodyOwnerKind::Closure => return None,
        };

        Some(const_kind)
    }
}

impl fmt::Display for ConstKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let s = match self {
            Self::Static => "static",
            Self::StaticMut => "static mut",
            Self::Const | Self::AnonConst => "const",
            Self::ConstFn => "const fn",
        };

        write!(f, "{}", s)
    }
}

fn check_mod_const_bodies(tcx: TyCtxt<'_>, module_def_id: DefId) {
    let mut vis = CheckConstVisitor::new(tcx);
    tcx.hir().visit_item_likes_in_module(module_def_id, &mut vis.as_deep_visitor());
}

pub(crate) fn provide(providers: &mut Providers<'_>) {
    *providers = Providers {
        check_mod_const_bodies,
        ..*providers
    };
}

#[derive(Copy, Clone)]
struct CheckConstVisitor<'tcx> {
    tcx: TyCtxt<'tcx>,
    const_kind: Option<ConstKind>,
}

impl<'tcx> CheckConstVisitor<'tcx> {
    fn new(tcx: TyCtxt<'tcx>) -> Self {
        CheckConstVisitor {
            tcx,
            const_kind: None,
        }
    }

    /// Emits an error when an unsupported expression is found in a const context.
    fn const_check_violated(&self, bad_op: &str, span: Span) {
        if self.tcx.sess.opts.debugging_opts.unleash_the_miri_inside_of_you {
            self.tcx.sess.span_warn(span, "skipping const checks");
            return;
        }

        let const_kind = self.const_kind
            .expect("`const_check_violated` may only be called inside a const context");

        span_err!(self.tcx.sess, span, E0744, "`{}` is not allowed in a `{}`", bad_op, const_kind);
    }

    /// Saves the parent `const_kind` before calling `f` and restores it afterwards.
    fn recurse_into(&mut self, kind: Option<ConstKind>, f: impl FnOnce(&mut Self)) {
        let parent_kind = self.const_kind;
        self.const_kind = kind;
        f(self);
        self.const_kind = parent_kind;
    }
}

impl<'tcx> Visitor<'tcx> for CheckConstVisitor<'tcx> {
    fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
        NestedVisitorMap::OnlyBodies(&self.tcx.hir())
    }

    fn visit_anon_const(&mut self, anon: &'tcx hir::AnonConst) {
        let kind = Some(ConstKind::AnonConst);
        self.recurse_into(kind, |this| hir::intravisit::walk_anon_const(this, anon));
    }

    fn visit_body(&mut self, body: &'tcx hir::Body) {
        let kind = ConstKind::for_body(body, self.tcx.hir());
        self.recurse_into(kind, |this| hir::intravisit::walk_body(this, body));
    }

    fn visit_expr(&mut self, e: &'tcx hir::Expr) {
        match &e.kind {
            // Skip the following checks if we are not currently in a const context.
            _ if self.const_kind.is_none() => {}

            hir::ExprKind::Loop(_, _, source) => {
                self.const_check_violated(source.name(), e.span);
            }

            hir::ExprKind::Match(_, _, source) if !self.tcx.features().const_if_match => {
                use hir::MatchSource::*;

                let op = match source {
                    Normal => Some("match"),
                    IfDesugar { .. } | IfLetDesugar { .. } => Some("if"),
                    TryDesugar => Some("?"),
                    AwaitDesugar => Some(".await"),

                    // These are handled by `ExprKind::Loop` above.
                    WhileDesugar | WhileLetDesugar | ForLoopDesugar => None,
                };

                if let Some(op) = op {
                    self.const_check_violated(op, e.span);
                }
            }

            _ => {},
        }

        hir::intravisit::walk_expr(self, e);
    }
}
<commit_msg>Suggest `const_if_match` on nightly<commit_after>//! This pass checks HIR bodies that may be evaluated at compile-time (e.g., `const`, `static`,
//! `const fn`) for structured control flow (e.g. `if`, `while`), which is forbidden in a const
//! context.
//!
//! By the time the MIR const-checker runs, these high-level constructs have been lowered to
//! control-flow primitives (e.g., `Goto`, `SwitchInt`), making it tough to properly attribute
//! errors. We still look for those primitives in the MIR const-checker to ensure nothing slips
//! through, but errors for structured control flow in a `const` should be emitted here.

use rustc::hir::def_id::DefId;
use rustc::hir::intravisit::{Visitor, NestedVisitorMap};
use rustc::hir::map::Map;
use rustc::hir;
use rustc::ty::TyCtxt;
use rustc::ty::query::Providers;
use syntax::ast::Mutability;
use syntax::feature_gate::{emit_feature_err, Features, GateIssue};
use syntax::span_err;
use syntax_pos::{sym, Span};
use rustc_error_codes::*;

use std::fmt;

/// An expression that is not *always* legal in a const context.
#[derive(Clone, Copy)]
enum NonConstExpr {
    Loop(hir::LoopSource),
    Match(hir::MatchSource),
}

impl NonConstExpr {
    fn name(self) -> &'static str {
        match self {
            Self::Loop(src) => src.name(),
            Self::Match(src) => src.name(),
        }
    }

    /// Returns `true` if all feature gates required to enable this expression are turned on, or
    /// `None` if there is no feature gate corresponding to this expression.
    fn is_feature_gate_enabled(self, features: &Features) -> Option<bool> {
        use hir::MatchSource::*;
        match self {
            | Self::Match(Normal)
            | Self::Match(IfDesugar { .. })
            | Self::Match(IfLetDesugar { .. })
            => Some(features.const_if_match),

            _ => None,
        }
    }
}

#[derive(Copy, Clone)]
enum ConstKind {
    Static,
    StaticMut,
    ConstFn,
    Const,
    AnonConst,
}

impl ConstKind {
    fn for_body(body: &hir::Body, hir_map: &Map<'_>) -> Option<Self> {
        let is_const_fn = |id| hir_map.fn_sig_by_hir_id(id).unwrap().header.is_const();

        let owner = hir_map.body_owner(body.id());
        let const_kind = match hir_map.body_owner_kind(owner) {
            hir::BodyOwnerKind::Const => Self::Const,
            hir::BodyOwnerKind::Static(Mutability::Mutable) => Self::StaticMut,
            hir::BodyOwnerKind::Static(Mutability::Immutable) => Self::Static,

            hir::BodyOwnerKind::Fn if is_const_fn(owner) => Self::ConstFn,
            hir::BodyOwnerKind::Fn | hir::BodyOwnerKind::Closure => return None,
        };

        Some(const_kind)
    }
}

impl fmt::Display for ConstKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let s = match self {
            Self::Static => "static",
            Self::StaticMut => "static mut",
            Self::Const | Self::AnonConst => "const",
            Self::ConstFn => "const fn",
        };

        write!(f, "{}", s)
    }
}

fn check_mod_const_bodies(tcx: TyCtxt<'_>, module_def_id: DefId) {
    let mut vis = CheckConstVisitor::new(tcx);
    tcx.hir().visit_item_likes_in_module(module_def_id, &mut vis.as_deep_visitor());
}

pub(crate) fn provide(providers: &mut Providers<'_>) {
    *providers = Providers {
        check_mod_const_bodies,
        ..*providers
    };
}

#[derive(Copy, Clone)]
struct CheckConstVisitor<'tcx> {
    tcx: TyCtxt<'tcx>,
    const_kind: Option<ConstKind>,
}

impl<'tcx> CheckConstVisitor<'tcx> {
    fn new(tcx: TyCtxt<'tcx>) -> Self {
        CheckConstVisitor {
            tcx,
            const_kind: None,
        }
    }

    /// Emits an error when an unsupported expression is found in a const context.
    fn const_check_violated(&self, expr: NonConstExpr, span: Span) {
        match expr.is_feature_gate_enabled(self.tcx.features()) {
            // Don't emit an error if the user has enabled the requisite feature gates.
            Some(true) => return,

            // Users of `-Zunleash-the-miri-inside-of-you` must use feature gates when possible.
            None if self.tcx.sess.opts.debugging_opts.unleash_the_miri_inside_of_you => {
                self.tcx.sess.span_warn(span, "skipping const checks");
                return;
            }

            _ => {}
        }

        let const_kind = self.const_kind
            .expect("`const_check_violated` may only be called inside a const context");

        let msg = format!("`{}` is not allowed in a `{}`", expr.name(), const_kind);
        match expr {
            | NonConstExpr::Match(hir::MatchSource::Normal)
            | NonConstExpr::Match(hir::MatchSource::IfDesugar { .. })
            | NonConstExpr::Match(hir::MatchSource::IfLetDesugar { .. })
            => emit_feature_err(
                &self.tcx.sess.parse_sess,
                sym::const_if_match,
                span,
                GateIssue::Language,
                &msg
            ),

            _ => span_err!(self.tcx.sess, span, E0744, "{}", msg),
        }
    }

    /// Saves the parent `const_kind` before calling `f` and restores it afterwards.
    fn recurse_into(&mut self, kind: Option<ConstKind>, f: impl FnOnce(&mut Self)) {
        let parent_kind = self.const_kind;
        self.const_kind = kind;
        f(self);
        self.const_kind = parent_kind;
    }
}

impl<'tcx> Visitor<'tcx> for CheckConstVisitor<'tcx> {
    fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
        NestedVisitorMap::OnlyBodies(&self.tcx.hir())
    }

    fn visit_anon_const(&mut self, anon: &'tcx hir::AnonConst) {
        let kind = Some(ConstKind::AnonConst);
        self.recurse_into(kind, |this| hir::intravisit::walk_anon_const(this, anon));
    }

    fn visit_body(&mut self, body: &'tcx hir::Body) {
        let kind = ConstKind::for_body(body, self.tcx.hir());
        self.recurse_into(kind, |this| hir::intravisit::walk_body(this, body));
    }

    fn visit_expr(&mut self, e: &'tcx hir::Expr) {
        match &e.kind {
            // Skip the following checks if we are not currently in a const context.
            _ if self.const_kind.is_none() => {}

            hir::ExprKind::Loop(_, _, source) => {
                self.const_check_violated(NonConstExpr::Loop(*source), e.span);
            }

            hir::ExprKind::Match(_, _, source) => {
                let non_const_expr = match source {
                    // These are handled by `ExprKind::Loop` above.
                    | hir::MatchSource::WhileDesugar
                    | hir::MatchSource::WhileLetDesugar
                    | hir::MatchSource::ForLoopDesugar
                    => None,

                    _ => Some(NonConstExpr::Match(*source)),
                };

                if let Some(expr) = non_const_expr {
                    self.const_check_violated(expr, e.span);
                }
            }

            _ => {},
        }

        hir::intravisit::walk_expr(self, e);
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::iter;
use std::cmp;
use super::stacks_and_queues::bag::Bag;
use super::stacks_and_queues::Queue;
use super::stacks_and_queues::resizing_array_queue::ResizingArrayQueue;
use super::priority_queues::MinPQ;
use super::priority_queues::binary_heaps::BinaryHeapMinPQ;
use super::union_find::UF;
use super::union_find::weighted_quick_union::UnionFind;

/// a weighted edge
#[derive(Clone, Copy)]
pub struct Edge {
    v: usize,
    w: usize,
    weight: f64
}

impl Edge {
    pub fn new(v: usize, w: usize, weight: f64) -> Edge {
        Edge { v: v, w: w, weight: weight }
    }

    pub fn weight(&self) -> f64 {
        self.weight
    }

    pub fn either(&self) -> usize {
        self.v
    }

    pub fn other(&self, vertex: usize) -> usize {
        if vertex == self.v {
            self.w
        } else if vertex == self.w {
            self.v
        } else {
            panic!("illegal endppint")
        }
    }

    fn swap(mut self) -> Edge {
        let v = self.v;
        self.v = self.w;
        self.w =v;
        self
    }
}

impl PartialEq for Edge {
    fn eq(&self, other: &Edge) -> bool {
        self.weight == other.weight
    }
}

impl PartialOrd for Edge {
    fn partial_cmp(&self, other: &Edge) -> Option<cmp::Ordering> {
        self.weight.partial_cmp(&other.weight)
    }
}

impl fmt::Display for Edge {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}-{} {:.5}", self.v, self.w, self.weight)
    }
}

pub struct EdgeWeightedGraph {
    v: usize,
    e: usize,
    adj: Vec<Bag<Edge>>
}

/// an edge-weighted graph
impl EdgeWeightedGraph {
    pub fn new(v: usize) -> EdgeWeightedGraph {
        EdgeWeightedGraph {
            v: v,
            e: 0,
            adj: iter::repeat(Bag::new()).take(v).collect::<Vec<Bag<Edge>>>()
        }
    }

    pub fn v(&self) -> usize {
        self.v
    }

    pub fn e(&self) -> usize {
        self.e
    }

    #[inline]
    fn validate_vertex(&self, v: usize) {
        assert!(v < self.v, "vertex is not between 0 and max V")
    }

    pub fn add_edge(&mut self, e: Edge) {
        let v = e.either();
        let w = e.other(v);
        self.validate_vertex(v);
        self.validate_vertex(w);
        self.adj[v].add(e.clone());
        if v != w {
            self.adj[w].add(e.swap());
        }
        self.e += 1
    }

    pub fn adj(&self, v: usize) -> ::std::vec::IntoIter<Edge> {
        self.validate_vertex(v);
        self.adj[v].iter().map(|e| e.clone()).collect::<Vec<Edge>>().into_iter()
    }

    pub fn degree(&self, v: usize) -> usize {
        self.validate_vertex(v);
        self.adj[v].len()
    }

    pub fn edges(&self) -> ::std::vec::IntoIter<Edge> {
        self.adj
            .iter()
            .flat_map(|adj| {
                adj.iter().map(|e| e.clone()).collect::<Vec<Edge>>().into_iter()
            })
            .filter(|ref e| e.either() <= e.other(e.either()) )
            .collect::<Vec<Edge>>()
            .into_iter()
    }

    // use fdp or neato
    pub fn to_dot(&self) -> String {
        let mut dot = String::new();

        dot.push_str("graph G {\n");
        for i in 0 .. self.v {
            dot.push_str(&format!("  {};\n", i));
        }

        for e in self.edges() {
            let v = e.either();
            let w = e.other(v);
            dot.push_str(&format!("  {} -- {} [len={}];\n",
                                  v, w, e.weight))
        }
        dot.push_str("}\n");
        dot

    }
}

#[test]
fn test_edge_weighted_graph() {
    let mut g = EdgeWeightedGraph::new(6);
    g.add_edge(Edge::new(0, 1, 7.0));
    g.add_edge(Edge::new(1, 2, 10.0));
    g.add_edge(Edge::new(0, 2, 9.0));
    g.add_edge(Edge::new(0, 5, 14.0));
    g.add_edge(Edge::new(1, 3, 15.0));
    g.add_edge(Edge::new(2, 5, 2.0));
    g.add_edge(Edge::new(2, 3, 11.0));
    g.add_edge(Edge::new(4, 5, 9.0));
    g.add_edge(Edge::new(3, 4, 6.0));
    g.add_edge(Edge::new(2, 2, 1.0));

    assert_eq!(10, g.edges().count());
    assert!(!g.to_dot().is_empty());
}

#[allow(dead_code)]
pub struct KruskalMST<'a> {
    graph: &'a EdgeWeightedGraph,
    weight: f64,
    mst: ResizingArrayQueue<Edge>
}

impl<'a> KruskalMST<'a> {
    fn new<'b>(graph: &'b EdgeWeightedGraph) -> KruskalMST<'b> {
        let n = graph.v();
        let mut weight = 0f64;
        let mut mst = ResizingArrayQueue::<Edge>::new();
        let mut pq = BinaryHeapMinPQ::<Edge>::new();
        for e in graph.edges() {
            pq.insert(e);
        }
        let mut uf = UnionFind::new(n);

        while !pq.is_empty() && mst.size() < n - 1 {
            let e = pq.del_min().unwrap();
            let v = e.either();
            let w = e.other(v);
            if !uf.connected(v, w) {
                uf.union(v, w);
                weight += e.weight();
                mst.enqueue(e);
            }
        }

        KruskalMST {
            graph: graph,
            weight: weight,
            mst: mst
        }
    }

    pub fn edges(&self) -> ::std::vec::IntoIter<Edge> {
        self.mst.clone().into_iter().collect::<Vec<Edge>>().into_iter()
    }
}

impl EdgeWeightedGraph {
    pub fn kruskal_mst<'a>(&'a self) -> KruskalMST<'a> {
        KruskalMST::new(self)
    }
}

#[test]
fn test_edge_weighted_graph_kruskal_mst() {
    let mut g = EdgeWeightedGraph::new(6);
    g.add_edge(Edge::new(0, 1, 7.0));
    g.add_edge(Edge::new(1, 2, 10.0));
    g.add_edge(Edge::new(0, 2, 9.0));
    g.add_edge(Edge::new(0, 5, 14.0));
    g.add_edge(Edge::new(1, 3, 15.0));
    g.add_edge(Edge::new(2, 5, 2.0));
    g.add_edge(Edge::new(2, 3, 11.0));
    g.add_edge(Edge::new(4, 5, 9.0));
    g.add_edge(Edge::new(3, 4, 6.0));
    g.add_edge(Edge::new(2, 2, 1.0));

    assert_eq!(33.0, g.kruskal_mst().weight);
    assert_eq!(33.0, g.kruskal_mst().edges().map(|e| e.weight).sum());

}

/// data type for computing a minimum spanning tree in an edge-weighted graph
pub struct LazyPrimMST<'a> {
    graph: &'a EdgeWeightedGraph,
    weight: f64,
    mst: ResizingArrayQueue<Edge>,
    marked: Vec<bool>,
    pq: BinaryHeapMinPQ<Edge>
}
<commit_msg>minimum_spanning_trees: add LazyPrimMST<commit_after>use std::fmt;
use std::iter;
use std::cmp;
use super::stacks_and_queues::bag::Bag;
use super::stacks_and_queues::Queue;
use super::stacks_and_queues::resizing_array_queue::ResizingArrayQueue;
use super::priority_queues::MinPQ;
use super::priority_queues::binary_heaps::BinaryHeapMinPQ;
use super::union_find::UF;
use super::union_find::weighted_quick_union::UnionFind;

/// a weighted edge
#[derive(Clone, Copy)]
pub struct Edge {
    v: usize,
    w: usize,
    weight: f64
}

impl Edge {
    pub fn new(v: usize, w: usize, weight: f64) -> Edge {
        Edge { v: v, w: w, weight: weight }
    }

    pub fn weight(&self) -> f64 {
        self.weight
    }

    pub fn either(&self) -> usize {
        self.v
    }

    pub fn other(&self, vertex: usize) -> usize {
        if vertex == self.v {
            self.w
        } else if vertex == self.w {
            self.v
        } else {
            panic!("illegal endppint")
        }
    }

    fn swap(mut self) -> Edge {
        let v = self.v;
        self.v = self.w;
        self.w =v;
        self
    }
}

impl PartialEq for Edge {
    fn eq(&self, other: &Edge) -> bool {
        self.weight == other.weight
    }
}

impl PartialOrd for Edge {
    fn partial_cmp(&self, other: &Edge) -> Option<cmp::Ordering> {
        self.weight.partial_cmp(&other.weight)
    }
}

impl fmt::Display for Edge {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}-{} {:.5}", self.v, self.w, self.weight)
    }
}

pub struct EdgeWeightedGraph {
    v: usize,
    e: usize,
    adj: Vec<Bag<Edge>>
}

/// an edge-weighted graph
impl EdgeWeightedGraph {
    pub fn new(v: usize) -> EdgeWeightedGraph {
        EdgeWeightedGraph {
            v: v,
            e: 0,
            adj: iter::repeat(Bag::new()).take(v).collect::<Vec<Bag<Edge>>>()
        }
    }

    pub fn v(&self) -> usize {
        self.v
    }

    pub fn e(&self) -> usize {
        self.e
    }

    #[inline]
    fn validate_vertex(&self, v: usize) {
        assert!(v < self.v, "vertex is not between 0 and max V")
    }

    pub fn add_edge(&mut self, e: Edge) {
        let v = e.either();
        let w = e.other(v);
        self.validate_vertex(v);
        self.validate_vertex(w);
        self.adj[v].add(e.clone());
        if v != w {
            self.adj[w].add(e.swap());
        }
        self.e += 1
    }

    pub fn adj(&self, v: usize) -> ::std::vec::IntoIter<Edge> {
        self.validate_vertex(v);
        self.adj[v].iter().map(|e| e.clone()).collect::<Vec<Edge>>().into_iter()
    }

    pub fn degree(&self, v: usize) -> usize {
        self.validate_vertex(v);
        self.adj[v].len()
    }

    pub fn edges(&self) -> ::std::vec::IntoIter<Edge> {
        self.adj
            .iter()
            .flat_map(|adj| {
                adj.iter().map(|e| e.clone()).collect::<Vec<Edge>>().into_iter()
            })
            .filter(|ref e| e.either() <= e.other(e.either()) )
            .collect::<Vec<Edge>>()
            .into_iter()
    }

    // use fdp or neato
    pub fn to_dot(&self) -> String {
        let mut dot = String::new();

        dot.push_str("graph G {\n");
        for i in 0 .. self.v {
            dot.push_str(&format!("  {};\n", i));
        }

        for e in self.edges() {
            let v = e.either();
            let w = e.other(v);
            dot.push_str(&format!("  {} -- {} [len={}];\n",
                                  v, w, e.weight))
        }
        dot.push_str("}\n");
        dot

    }
}

#[test]
fn test_edge_weighted_graph() {
    let mut g = EdgeWeightedGraph::new(6);
    g.add_edge(Edge::new(0, 1, 7.0));
    g.add_edge(Edge::new(1, 2, 10.0));
    g.add_edge(Edge::new(0, 2, 9.0));
    g.add_edge(Edge::new(0, 5, 14.0));
    g.add_edge(Edge::new(1, 3, 15.0));
    g.add_edge(Edge::new(2, 5, 2.0));
    g.add_edge(Edge::new(2, 3, 11.0));
    g.add_edge(Edge::new(4, 5, 9.0));
    g.add_edge(Edge::new(3, 4, 6.0));
    g.add_edge(Edge::new(2, 2, 1.0));

    assert_eq!(10, g.edges().count());
    assert!(!g.to_dot().is_empty());
}

#[allow(dead_code)]
pub struct KruskalMST<'a> {
    graph: &'a EdgeWeightedGraph,
    weight: f64,
    mst: ResizingArrayQueue<Edge>
}

impl<'a> KruskalMST<'a> {
    fn new<'b>(graph: &'b EdgeWeightedGraph) -> KruskalMST<'b> {
        let n = graph.v();
        let mut weight = 0f64;
        let mut mst = ResizingArrayQueue::<Edge>::new();
        let mut pq = BinaryHeapMinPQ::<Edge>::new();
        for e in graph.edges() {
            pq.insert(e);
        }
        let mut uf = UnionFind::new(n);

        while !pq.is_empty() && mst.size() < n - 1 {
            let e = pq.del_min().unwrap();
            let v = e.either();
            let w = e.other(v);
            if !uf.connected(v, w) {
                uf.union(v, w);
                weight += e.weight();
                mst.enqueue(e);
            }
        }

        KruskalMST {
            graph: graph,
            weight: weight,
            mst: mst
        }
    }

    pub fn edges(&self) -> ::std::vec::IntoIter<Edge> {
        self.mst.clone().into_iter().collect::<Vec<Edge>>().into_iter()
    }
}

impl EdgeWeightedGraph {
    pub fn kruskal_mst<'a>(&'a self) -> KruskalMST<'a> {
        KruskalMST::new(self)
    }
}

#[test]
fn test_edge_weighted_graph_kruskal_mst() {
    let mut g = EdgeWeightedGraph::new(6);
    g.add_edge(Edge::new(0, 1, 7.0));
    g.add_edge(Edge::new(1, 2, 10.0));
    g.add_edge(Edge::new(0, 2, 9.0));
    g.add_edge(Edge::new(0, 5, 14.0));
    g.add_edge(Edge::new(1, 3, 15.0));
    g.add_edge(Edge::new(2, 5, 2.0));
    g.add_edge(Edge::new(2, 3, 11.0));
    g.add_edge(Edge::new(4, 5, 9.0));
    g.add_edge(Edge::new(3, 4, 6.0));
    g.add_edge(Edge::new(2, 2, 1.0));

    assert_eq!(33.0, g.kruskal_mst().weight);
    assert_eq!(33.0, g.kruskal_mst().edges().map(|e| e.weight).sum());
}

/// data type for computing a minimum spanning tree in an edge-weighted graph
pub struct LazyPrimMST<'a> {
    graph: &'a EdgeWeightedGraph,
    weight: f64,
    mst: ResizingArrayQueue<Edge>,
    marked: Vec<bool>,
    pq: BinaryHeapMinPQ<Edge>
}

impl<'a> LazyPrimMST<'a> {
    fn new<'b>(graph: &'b EdgeWeightedGraph) -> LazyPrimMST<'b> {
        let n = graph.v();
        let marked = iter::repeat(false).take(n).collect();
        let pq = BinaryHeapMinPQ::new();
        let mst = ResizingArrayQueue::new();

        let mut ret = LazyPrimMST {
            graph: graph,
            weight: 0.0,
            mst: mst,
            marked: marked,
            pq: pq
        };
        // run Prim for all vertices
        // get a minimum spanning forest
        for v in 0 .. n {
            if !ret.marked[v] {
                ret.prim(v);
            }
        }
        ret
    }

    // Prim's algorithm
    fn prim(&mut self, s: usize) {
        self.visit(s);

        while !self.pq.is_empty() {
            let e = self.pq.del_min().unwrap();
            let v = e.either();
            let w = e.other(v);

            if self.marked[v] && self.marked[w] {
                continue;
            }
            self.weight += e.weight();
            self.mst.enqueue(e);
            if !self.marked[v] {
                self.visit(v);
            }
            if !self.marked[w] {
                self.visit(w)
            }
        }
    }

    fn visit(&mut self, v: usize) {
        self.marked[v] = true;
        for e in self.graph.adj(v) {
            if !self.marked[e.other(v)] {
                self.pq.insert(e)
            }
        }
    }

    pub fn edges(&self) -> ::std::vec::IntoIter<Edge> {
        self.mst.clone().into_iter().collect::<Vec<Edge>>().into_iter()
    }
}

impl EdgeWeightedGraph {
     pub fn lazy_prim_mst<'a>(&'a self) -> LazyPrimMST<'a> {
         LazyPrimMST::new(self)
    }
}

#[test]
fn test_edge_weighted_graph_lazy_prim_mst() {
    let mut g = EdgeWeightedGraph::new(6);
    g.add_edge(Edge::new(0, 1, 7.0));
    g.add_edge(Edge::new(1, 2, 10.0));
    g.add_edge(Edge::new(0, 2, 9.0));
    g.add_edge(Edge::new(0, 5, 14.0));
    g.add_edge(Edge::new(1, 3, 15.0));
    g.add_edge(Edge::new(2, 5, 2.0));
    g.add_edge(Edge::new(2, 3, 11.0));
    g.add_edge(Edge::new(4, 5, 9.0));
    g.add_edge(Edge::new(3, 4, 6.0));
    g.add_edge(Edge::new(2, 2, 1.0));

    assert_eq!(33.0, g.lazy_prim_mst().weight);
    assert_eq!(33.0, g.lazy_prim_mst().edges().map(|e| e.weight).sum());
}
<|endoftext|>


<commit_before>//! This module is for statically evaluating a position.

use std::fmt;
use std::ops::{Add, Sub, Neg};
use std::cmp::Ordering;

use moves::NumMoves;
use {Color, Piece, PieceType};
use {Pawn, King, Queen, Bishop, Knight, Rook};
use super::Position;

#[derive(PartialEq, Eq, PartialOrd, Ord, Copy, Clone, Debug)]
pub struct ScoreUnit(pub i32);
impl Add for ScoreUnit {
    type Output = Self;
    fn add(self, rhs: Self) -> Self { ScoreUnit(self.0 + rhs.0) }
}
impl Sub for ScoreUnit {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self { ScoreUnit(self.0 - rhs.0) }
}
impl Neg for ScoreUnit {
    type Output = Self;
    fn neg(self) -> Self { ScoreUnit(-self.0) }
}

/// An assessment of the position.
#[derive(PartialEq, Eq, Copy, Clone, Debug)]
pub enum Score {
    // Positive: advantage for side to move.
    // Negative: disadvantage for side to move.
    Value(ScoreUnit),
    // Side to move can checkmate in x moves.
    // WinIn(NumMoves(1)): can checkmate now.
    // WinIn(NumMoves(2)): can checkmate next move.
    WinIn(NumMoves),
    // Side to move will be checkmated in x moves.
    // WinIn(NumMoves(0)): already checkmated.
    // WinIn(NumMoves(1)): Will be immediately checkmated after any move.
    LoseIn(NumMoves),
}
impl Score {
    pub fn increment(self) -> Score {
        match self {
            Score::Value(val) => Score::Value(-val),
            Score::WinIn(val) => Score::LoseIn(val),
            Score::LoseIn(val) => Score::WinIn(NumMoves(val.0+1)),
        }
    }
    pub fn decrement(self) -> Score {
        match self {
            Score::Value(val) => Score::Value(-val),
            Score::WinIn(val) => Score::LoseIn(NumMoves(
                if val.0 == 0 { 0 } else { val.0 - 1 } )),
            Score::LoseIn(val) => Score::WinIn(val),
        }
    }
}
impl fmt::Display for Score {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Score::Value(val)   => write!(f, "cp {}", val.0),
            Score::WinIn(val)   => write!(f, "mate {}", val.0 as i32),
            Score::LoseIn(val)  => write!(f, "mate {}", (val.0 as i32) * -1),
        }
    }
}
impl Ord for Score {
    fn cmp(&self, other: &Score) -> Ordering {
        match *self {
            Score::WinIn(val1) => match *other {
                Score::WinIn(val2) => val2.cmp(&val1),
                _ => Ordering::Greater,
            },
            Score::LoseIn(val1) => match *other {
                Score::LoseIn(val2) => val1.cmp(&val2),
                _ => Ordering::Less,
            },
            Score::Value(val1) => match *other {
                Score::WinIn(_) => Ordering::Less,
                Score::LoseIn(_) => Ordering::Greater,
                Score::Value(val2) => val1.cmp(&val2),
            },
        }
    }
}
impl PartialOrd for Score {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

/// Evaluates the position without searching.
pub fn eval(p: &mut Position, draw_val: ScoreUnit) -> Score {
    if p.is_checkmated() {
        Score::LoseIn(NumMoves(0))
    } else if p.is_stalemated() {
        Score::Value(draw_val)
    } else {
        let c = p.side_to_move();
        // TODO change fold() to sum() when possible
        Score::Value(p.piece_iter().map( |(piece, _pos)| val_for_color(piece, c) ).fold(ScoreUnit(0), |x, y| x+y))
    }
}

fn val_for_color(piece: Piece, c: Color) -> ScoreUnit {
    let val = val_of_piece_type(piece.piece_type());
    if piece.color() == c {
        val
    } else {
        -val
    }
}

fn val_of_piece_type(piece_type: PieceType) -> ScoreUnit {
    ScoreUnit(match piece_type {
        King   => 100000,
        Pawn   =>    100,
        Queen  =>    900,
        Bishop =>    300,
        Knight =>    300,
        Rook   =>    500,
    })
}
<commit_msg>Add mobility in evaluation<commit_after>//! This module is for statically evaluating a position.

use std::fmt;
use std::ops::{Add, Sub, Neg, Mul};
use std::cmp::Ordering;

use moves::NumMoves;
use {Color, Piece, PieceType};
use {Pawn, King, Queen, Bishop, Knight, Rook};
use super::Position;

#[derive(PartialEq, Eq, PartialOrd, Ord, Copy, Clone, Debug)]
pub struct ScoreUnit(pub i32);
impl Add for ScoreUnit {
    type Output = Self;
    fn add(self, rhs: Self) -> Self { ScoreUnit(self.0 + rhs.0) }
}
impl Sub for ScoreUnit {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self { ScoreUnit(self.0 - rhs.0) }
}
impl Neg for ScoreUnit {
    type Output = Self;
    fn neg(self) -> Self { ScoreUnit(-self.0) }
}
impl Mul<i32> for ScoreUnit {
    type Output = Self;
    fn mul(self, rhs: i32) -> Self { ScoreUnit(self.0 * rhs) }
}

/// An assessment of the position.
#[derive(PartialEq, Eq, Copy, Clone, Debug)]
pub enum Score {
    // Positive: advantage for side to move.
    // Negative: disadvantage for side to move.
    Value(ScoreUnit),
    // Side to move can checkmate in x moves.
    // WinIn(NumMoves(1)): can checkmate now.
    // WinIn(NumMoves(2)): can checkmate next move.
    WinIn(NumMoves),
    // Side to move will be checkmated in x moves.
    // WinIn(NumMoves(0)): already checkmated.
    // WinIn(NumMoves(1)): Will be immediately checkmated after any move.
    LoseIn(NumMoves),
}
impl Score {
    pub fn increment(self) -> Score {
        match self {
            Score::Value(val) => Score::Value(-val),
            Score::WinIn(val) => Score::LoseIn(val),
            Score::LoseIn(val) => Score::WinIn(NumMoves(val.0+1)),
        }
    }
    pub fn decrement(self) -> Score {
        match self {
            Score::Value(val) => Score::Value(-val),
            Score::WinIn(val) => Score::LoseIn(NumMoves(
                if val.0 == 0 { 0 } else { val.0 - 1 } )),
            Score::LoseIn(val) => Score::WinIn(val),
        }
    }
}
impl fmt::Display for Score {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Score::Value(val)   => write!(f, "cp {}", val.0),
            Score::WinIn(val)   => write!(f, "mate {}", val.0 as i32),
            Score::LoseIn(val)  => write!(f, "mate {}", (val.0 as i32) * -1),
        }
    }
}
impl Ord for Score {
    fn cmp(&self, other: &Score) -> Ordering {
        match *self {
            Score::WinIn(val1) => match *other {
                Score::WinIn(val2) => val2.cmp(&val1),
                _ => Ordering::Greater,
            },
            Score::LoseIn(val1) => match *other {
                Score::LoseIn(val2) => val1.cmp(&val2),
                _ => Ordering::Less,
            },
            Score::Value(val1) => match *other {
                Score::WinIn(_) => Ordering::Less,
                Score::LoseIn(_) => Ordering::Greater,
                Score::Value(val2) => val1.cmp(&val2),
            },
        }
    }
}
impl PartialOrd for Score {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

/// Evaluates the position without searching.
pub fn eval(p: &mut Position, draw_val: ScoreUnit) -> Score {
    if p.is_checkmated() {
        Score::LoseIn(NumMoves(0))
    } else if p.is_stalemated() {
        Score::Value(draw_val)
    } else {
        let c = p.side_to_move();
        // TODO change fold() to sum() when possible
        let piece_eval = p.piece_iter()
                          .map( |(piece, _pos)| val_for_color(piece, c) )
                          .fold(ScoreUnit(0), |x, y| x+y);
        let our_mobility = p.psudo_legal_iter().count();
        p.swap_side_to_move();
        let his_mobility = p.psudo_legal_iter().count();
        p.swap_side_to_move();
        Score::Value(piece_eval + VALUE_PER_MOBILITY * (our_mobility as i32- his_mobility as i32))
    }
}

fn val_for_color(piece: Piece, c: Color) -> ScoreUnit {
    let val = val_of_piece_type(piece.piece_type());
    if piece.color() == c {
        val
    } else {
        -val
    }
}

const VALUE_PER_MOBILITY: ScoreUnit = ScoreUnit(5);

fn val_of_piece_type(piece_type: PieceType) -> ScoreUnit {
    ScoreUnit(match piece_type {
        King   => 100000,
        Pawn   =>    100,
        Queen  =>    900,
        Bishop =>    300,
        Knight =>    300,
        Rook   =>    500,
    })
}
<|endoftext|>


<commit_before>use super::V2_COOKIE;
use super::super::{Counter, Histogram, RestatState};
use super::super::num::ToPrimitive;
use std::io::{self, Cursor, ErrorKind, Read};
use std::marker::PhantomData;
use std;
use super::byteorder::{BigEndian, ReadBytesExt};

/// Errors that can happen during deserialization.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum DeserializeError {
    /// An i/o operation failed.
    IoError(ErrorKind),
    /// The cookie (first 4 bytes) did not match that for any supported format.
    InvalidCookie,
    /// The histogram uses features that this implementation doesn't support (yet), so it cannot
    /// be deserialized correctly.
    UnsupportedFeature,
    /// A count exceeded what can be represented in the chosen counter type.
    UnsuitableCounterType,
    /// The histogram instance could not be created because the serialized parameters were invalid
    /// (e.g. lowest value, highest value, etc.)
    InvalidParameters,
    /// The current system's pointer width cannot represent the encoded histogram.
    UsizeTypeTooSmall,
    /// The encoded array is longer than it should be for the histogram's value range.
    EncodedArrayTooLong
}

impl std::convert::From<std::io::Error> for DeserializeError {
    fn from(e: std::io::Error) -> Self {
        DeserializeError::IoError(e.kind())
    }
}

/// Deserializer for all supported formats.
///
/// Since the serialization formats all include some magic bytes that allow reliable identification
/// of the different formats, only one Deserializer implementation is needed.
pub struct Deserializer {
    payload_buf: Vec<u8>
}

impl Deserializer {
    /// Create a new deserializer.
    pub fn new() -> Deserializer {
        Deserializer {
            payload_buf: Vec::new()
        }
    }

    /// Deserialize an encoded histogram from the provided reader.
    ///
    /// Note that `&[u8]` and `Cursor` are convenient implementations of `Read` if you have some
    /// bytes already in slice or `Vec` form.
    pub fn deserialize<T: Counter, R: Read>(&mut self, reader: &mut R)
                                            -> Result<Histogram<T>, DeserializeError> {
        // TODO benchmark minimizing read calls by reading into a fixed-size header buffer

        let cookie = reader.read_u32::<BigEndian>()?;

        if cookie != V2_COOKIE {
            return Err(DeserializeError::InvalidCookie);
        }

        let payload_len = reader.read_u32::<BigEndian>()?.to_usize()
            .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        let normalizing_offset = reader.read_u32::<BigEndian>()?;
        if normalizing_offset != 0 {
            return Err(DeserializeError::UnsupportedFeature);
        }
        let num_digits = reader.read_u32::<BigEndian>()?.to_u8()
            .ok_or(DeserializeError::InvalidParameters)?;
        let low = reader.read_u64::<BigEndian>()?;
        let high = reader.read_u64::<BigEndian>()?;
        let int_double_ratio = reader.read_f64::<BigEndian>()?;
        if int_double_ratio != 1.0 {
            return Err(DeserializeError::UnsupportedFeature);
        }

        let mut h = Histogram::new_with_bounds(low, high, num_digits)
            .map_err(|_| DeserializeError::InvalidParameters)?;

        if payload_len > self.payload_buf.len() {
            self.payload_buf.resize(payload_len, 0);
        }

        let mut payload_slice = &mut self.payload_buf[0..payload_len];
        reader.read_exact(&mut payload_slice)?;

        let mut payload_index: usize = 0;
        let mut restat_state = RestatState::new();
        let mut decode_state = DecodeLoopState::new();

        while payload_index < payload_len.saturating_sub(9) {
            // Read with fast loop until we are within 9 of the end. Fast loop can't handle EOF,
            // so bail to slow version for the last few bytes.

            let (zz_num, bytes_read) = varint_read_slice(
                &payload_slice[payload_index..(payload_index + 9)]);
            payload_index += bytes_read;

            let count_or_zeros = zig_zag_decode(zz_num);

            decode_state.on_decoded_num(count_or_zeros, &mut restat_state, &mut h)?;
        }

        // Now read the leftovers
        let leftover_slice = &payload_slice[payload_index..];
        let mut cursor = Cursor::new(&leftover_slice);
        while cursor.position() < leftover_slice.len() as u64 {
            let count_or_zeros = zig_zag_decode(varint_read(&mut cursor)?);

            decode_state.on_decoded_num(count_or_zeros, &mut restat_state, &mut h)?;
        }

        restat_state.update_histogram(&mut h);

        Ok(h)
    }
}

// Only public for testing.
/// Read from a slice that must be 9 bytes long or longer. Returns the decoded number and how many
/// bytes were consumed.
#[inline]
pub fn varint_read_slice(slice: &[u8]) -> (u64, usize) {
    let mut b = slice[0];

    // take low 7 bits
    let mut value: u64 = low_7_bits(b);
    if !is_high_bit_set(b) {
        return (value, 1);
    }
    // high bit set, keep reading
    b = slice[1];
    value |= low_7_bits(b) << 7;
    if !is_high_bit_set(b) {
        return (value, 2);
    }
    b = slice[2];
    value |= low_7_bits(b) << 7 * 2;
    if !is_high_bit_set(b) {
        return (value, 3);
    }
    b = slice[3];
    value |= low_7_bits(b) << 7 * 3;
    if !is_high_bit_set(b) {
        return (value, 4);
    }
    b = slice[4];
    value |= low_7_bits(b) << 7 * 4;
    if !is_high_bit_set(b) {
        return (value, 5);
    }
    b = slice[5];
    value |= low_7_bits(b) << 7 * 5;
    if !is_high_bit_set(b) {
        return (value, 6);
    }
    b = slice[6];
    value |= low_7_bits(b) << 7 * 6;
    if !is_high_bit_set(b) {
        return (value, 7);
    }
    b = slice[7];
    value |= low_7_bits(b) << 7 * 7;
    if !is_high_bit_set(b) {
        return (value, 8);
    }

    b = slice[8];
    // special case: use last byte as is
    value |= (b as u64) << 7 * 8;

    (value, 9)
}

// Only public for testing.
/// Read a LEB128-64b9B from the buffer
pub fn varint_read<R: Read>(reader: &mut R) -> io::Result<u64> {
    let mut b = reader.read_u8()?;

    // take low 7 bits
    let mut value: u64 = low_7_bits(b);

    if is_high_bit_set(b) {
        // high bit set, keep reading
        b = reader.read_u8()?;
        value |= low_7_bits(b) << 7;
        if is_high_bit_set(b) {
            b = reader.read_u8()?;
            value |= low_7_bits(b) << 7 * 2;
            if is_high_bit_set(b) {
                b = reader.read_u8()?;
                value |= low_7_bits(b) << 7 * 3;
                if is_high_bit_set(b) {
                    b = reader.read_u8()?;
                    value |= low_7_bits(b) << 7 * 4;
                    if is_high_bit_set(b) {
                        b = reader.read_u8()?;
                        value |= low_7_bits(b) << 7 * 5;
                        if is_high_bit_set(b) {
                            b = reader.read_u8()?;
                            value |= low_7_bits(b) << 7 * 6;
                            if is_high_bit_set(b) {
                                b = reader.read_u8()?;
                                value |= low_7_bits(b) << 7 * 7;
                                if is_high_bit_set(b) {
                                    b = reader.read_u8()?;
                                    // special case: use last byte as is
                                    value |= (b as u64) << 7 * 8;
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    Ok(value)
}

/// truncate byte to low 7 bits, cast to u64
#[inline]
fn low_7_bits(b: u8) -> u64 {
    (b & 0x7F) as u64
}

#[inline]
fn is_high_bit_set(b: u8) -> bool {
    // TODO benchmark leading zeros rather than masking
    (b & 0x80) != 0
}

// Only public for testing.
#[inline]
pub fn zig_zag_decode(encoded: u64) -> i64 {
    ((encoded >> 1) as i64) ^ -((encoded & 1) as i64)
}

/// We need to perform the same logic in two different decode loops while carrying over a modicum
/// of state.
struct DecodeLoopState<T: Counter> {
    dest_index:usize,
    phantom: PhantomData<T>
}

impl <T: Counter> DecodeLoopState<T> {

    fn new() -> DecodeLoopState<T> {
        DecodeLoopState {
            dest_index: 0,
            phantom: PhantomData
        }
    }

    #[inline]
    fn on_decoded_num(&mut self, count_or_zeros: i64, restat_state: &mut RestatState<T>,
                      h: &mut Histogram<T>) -> Result<(), DeserializeError> {
        if count_or_zeros < 0 {
            let zero_count = (-count_or_zeros).to_usize()
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
            // skip the zeros
            self.dest_index = self.dest_index.checked_add(zero_count)
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        } else {
            let count: T = T::from_i64(count_or_zeros)
                .ok_or(DeserializeError::UnsuitableCounterType)?;

            h.set_count_at_index(self.dest_index, count)
                .map_err(|_| DeserializeError::EncodedArrayTooLong)?;

            restat_state.on_nonzero_count(self.dest_index, count);

            self.dest_index = self.dest_index.checked_add(1)
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        }

        Ok(())
    }
}
<commit_msg>Remove TODO about using leading zero count instead of mask. Leading zeros are unambiguously slower.<commit_after>use super::V2_COOKIE;
use super::super::{Counter, Histogram, RestatState};
use super::super::num::ToPrimitive;
use std::io::{self, Cursor, ErrorKind, Read};
use std::marker::PhantomData;
use std;
use super::byteorder::{BigEndian, ReadBytesExt};

/// Errors that can happen during deserialization.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum DeserializeError {
    /// An i/o operation failed.
    IoError(ErrorKind),
    /// The cookie (first 4 bytes) did not match that for any supported format.
    InvalidCookie,
    /// The histogram uses features that this implementation doesn't support (yet), so it cannot
    /// be deserialized correctly.
    UnsupportedFeature,
    /// A count exceeded what can be represented in the chosen counter type.
    UnsuitableCounterType,
    /// The histogram instance could not be created because the serialized parameters were invalid
    /// (e.g. lowest value, highest value, etc.)
    InvalidParameters,
    /// The current system's pointer width cannot represent the encoded histogram.
    UsizeTypeTooSmall,
    /// The encoded array is longer than it should be for the histogram's value range.
    EncodedArrayTooLong
}

impl std::convert::From<std::io::Error> for DeserializeError {
    fn from(e: std::io::Error) -> Self {
        DeserializeError::IoError(e.kind())
    }
}

/// Deserializer for all supported formats.
///
/// Since the serialization formats all include some magic bytes that allow reliable identification
/// of the different formats, only one Deserializer implementation is needed.
pub struct Deserializer {
    payload_buf: Vec<u8>
}

impl Deserializer {
    /// Create a new deserializer.
    pub fn new() -> Deserializer {
        Deserializer {
            payload_buf: Vec::new()
        }
    }

    /// Deserialize an encoded histogram from the provided reader.
    ///
    /// Note that `&[u8]` and `Cursor` are convenient implementations of `Read` if you have some
    /// bytes already in slice or `Vec` form.
    pub fn deserialize<T: Counter, R: Read>(&mut self, reader: &mut R)
                                            -> Result<Histogram<T>, DeserializeError> {
        // TODO benchmark minimizing read calls by reading into a fixed-size header buffer

        let cookie = reader.read_u32::<BigEndian>()?;

        if cookie != V2_COOKIE {
            return Err(DeserializeError::InvalidCookie);
        }

        let payload_len = reader.read_u32::<BigEndian>()?.to_usize()
            .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        let normalizing_offset = reader.read_u32::<BigEndian>()?;
        if normalizing_offset != 0 {
            return Err(DeserializeError::UnsupportedFeature);
        }
        let num_digits = reader.read_u32::<BigEndian>()?.to_u8()
            .ok_or(DeserializeError::InvalidParameters)?;
        let low = reader.read_u64::<BigEndian>()?;
        let high = reader.read_u64::<BigEndian>()?;
        let int_double_ratio = reader.read_f64::<BigEndian>()?;
        if int_double_ratio != 1.0 {
            return Err(DeserializeError::UnsupportedFeature);
        }

        let mut h = Histogram::new_with_bounds(low, high, num_digits)
            .map_err(|_| DeserializeError::InvalidParameters)?;

        if payload_len > self.payload_buf.len() {
            self.payload_buf.resize(payload_len, 0);
        }

        let mut payload_slice = &mut self.payload_buf[0..payload_len];
        reader.read_exact(&mut payload_slice)?;

        let mut payload_index: usize = 0;
        let mut restat_state = RestatState::new();
        let mut decode_state = DecodeLoopState::new();

        while payload_index < payload_len.saturating_sub(9) {
            // Read with fast loop until we are within 9 of the end. Fast loop can't handle EOF,
            // so bail to slow version for the last few bytes.

            let (zz_num, bytes_read) = varint_read_slice(
                &payload_slice[payload_index..(payload_index + 9)]);
            payload_index += bytes_read;

            let count_or_zeros = zig_zag_decode(zz_num);

            decode_state.on_decoded_num(count_or_zeros, &mut restat_state, &mut h)?;
        }

        // Now read the leftovers
        let leftover_slice = &payload_slice[payload_index..];
        let mut cursor = Cursor::new(&leftover_slice);
        while cursor.position() < leftover_slice.len() as u64 {
            let count_or_zeros = zig_zag_decode(varint_read(&mut cursor)?);

            decode_state.on_decoded_num(count_or_zeros, &mut restat_state, &mut h)?;
        }

        restat_state.update_histogram(&mut h);

        Ok(h)
    }
}

// Only public for testing.
/// Read from a slice that must be 9 bytes long or longer. Returns the decoded number and how many
/// bytes were consumed.
#[inline]
pub fn varint_read_slice(slice: &[u8]) -> (u64, usize) {
    let mut b = slice[0];

    // take low 7 bits
    let mut value: u64 = low_7_bits(b);
    if !is_high_bit_set(b) {
        return (value, 1);
    }
    // high bit set, keep reading
    b = slice[1];
    value |= low_7_bits(b) << 7;
    if !is_high_bit_set(b) {
        return (value, 2);
    }
    b = slice[2];
    value |= low_7_bits(b) << 7 * 2;
    if !is_high_bit_set(b) {
        return (value, 3);
    }
    b = slice[3];
    value |= low_7_bits(b) << 7 * 3;
    if !is_high_bit_set(b) {
        return (value, 4);
    }
    b = slice[4];
    value |= low_7_bits(b) << 7 * 4;
    if !is_high_bit_set(b) {
        return (value, 5);
    }
    b = slice[5];
    value |= low_7_bits(b) << 7 * 5;
    if !is_high_bit_set(b) {
        return (value, 6);
    }
    b = slice[6];
    value |= low_7_bits(b) << 7 * 6;
    if !is_high_bit_set(b) {
        return (value, 7);
    }
    b = slice[7];
    value |= low_7_bits(b) << 7 * 7;
    if !is_high_bit_set(b) {
        return (value, 8);
    }

    b = slice[8];
    // special case: use last byte as is
    value |= (b as u64) << 7 * 8;

    (value, 9)
}

// Only public for testing.
/// Read a LEB128-64b9B from the buffer
pub fn varint_read<R: Read>(reader: &mut R) -> io::Result<u64> {
    let mut b = reader.read_u8()?;

    // take low 7 bits
    let mut value: u64 = low_7_bits(b);

    if is_high_bit_set(b) {
        // high bit set, keep reading
        b = reader.read_u8()?;
        value |= low_7_bits(b) << 7;
        if is_high_bit_set(b) {
            b = reader.read_u8()?;
            value |= low_7_bits(b) << 7 * 2;
            if is_high_bit_set(b) {
                b = reader.read_u8()?;
                value |= low_7_bits(b) << 7 * 3;
                if is_high_bit_set(b) {
                    b = reader.read_u8()?;
                    value |= low_7_bits(b) << 7 * 4;
                    if is_high_bit_set(b) {
                        b = reader.read_u8()?;
                        value |= low_7_bits(b) << 7 * 5;
                        if is_high_bit_set(b) {
                            b = reader.read_u8()?;
                            value |= low_7_bits(b) << 7 * 6;
                            if is_high_bit_set(b) {
                                b = reader.read_u8()?;
                                value |= low_7_bits(b) << 7 * 7;
                                if is_high_bit_set(b) {
                                    b = reader.read_u8()?;
                                    // special case: use last byte as is
                                    value |= (b as u64) << 7 * 8;
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    Ok(value)
}

/// truncate byte to low 7 bits, cast to u64
#[inline]
fn low_7_bits(b: u8) -> u64 {
    (b & 0x7F) as u64
}

#[inline]
fn is_high_bit_set(b: u8) -> bool {
    (b & 0x80) != 0
}

// Only public for testing.
#[inline]
pub fn zig_zag_decode(encoded: u64) -> i64 {
    ((encoded >> 1) as i64) ^ -((encoded & 1) as i64)
}

/// We need to perform the same logic in two different decode loops while carrying over a modicum
/// of state.
struct DecodeLoopState<T: Counter> {
    dest_index:usize,
    phantom: PhantomData<T>
}

impl <T: Counter> DecodeLoopState<T> {

    fn new() -> DecodeLoopState<T> {
        DecodeLoopState {
            dest_index: 0,
            phantom: PhantomData
        }
    }

    #[inline]
    fn on_decoded_num(&mut self, count_or_zeros: i64, restat_state: &mut RestatState<T>,
                      h: &mut Histogram<T>) -> Result<(), DeserializeError> {
        if count_or_zeros < 0 {
            let zero_count = (-count_or_zeros).to_usize()
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
            // skip the zeros
            self.dest_index = self.dest_index.checked_add(zero_count)
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        } else {
            let count: T = T::from_i64(count_or_zeros)
                .ok_or(DeserializeError::UnsuitableCounterType)?;

            h.set_count_at_index(self.dest_index, count)
                .map_err(|_| DeserializeError::EncodedArrayTooLong)?;

            restat_state.on_nonzero_count(self.dest_index, count);

            self.dest_index = self.dest_index.checked_add(1)
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        }

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


<commit_before>use super::V2_COOKIE;
use super::super::{Counter, Histogram, RestatState};
use super::super::num::ToPrimitive;
use std::io::{self, Cursor, ErrorKind, Read};
use std;
use super::byteorder::{BigEndian, ReadBytesExt};

/// Errors that can happen during deserialization.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum DeserializeError {
    /// An i/o operation failed.
    IoError(ErrorKind),
    /// The cookie (first 4 bytes) did not match that for any supported format.
    InvalidCookie,
    /// The histogram uses features that this implementation doesn't support (yet), so it cannot
    /// be deserialized correctly.
    UnsupportedFeature,
    /// A count exceeded what can be represented in the chosen counter type.
    UnsuitableCounterType,
    /// The histogram instance could not be created because the serialized parameters were invalid
    /// (e.g. lowest value, highest value, etc.)
    InvalidParameters,
    /// The current system's pointer width cannot represent the encoded histogram.
    UsizeTypeTooSmall,
    /// The encoded array is longer than it should be for the histogram's value range.
    EncodedArrayTooLong
}

impl std::convert::From<std::io::Error> for DeserializeError {
    fn from(e: std::io::Error) -> Self {
        DeserializeError::IoError(e.kind())
    }
}

/// Deserializer for all supported formats.
///
/// Since the serialization formats all include some magic bytes that allow reliable identification
/// of the different formats, only one Deserializer implementation is needed.
pub struct Deserializer {
    payload_buf: Vec<u8>
}

impl Deserializer {
    /// Create a new deserializer.
    pub fn new() -> Deserializer {
        Deserializer {
            payload_buf: Vec::new()
        }
    }

    /// Deserialize an encoded histogram from the provided reader.
    ///
    /// Note that `&[u8]` and `Cursor` are convenient implementations of `Read` if you have some
    /// bytes already in slice or `Vec` form.
    pub fn deserialize<T: Counter, R: Read>(&mut self, reader: &mut R)
                                            -> Result<Histogram<T>, DeserializeError> {
        // TODO benchmark minimizing read calls by reading into a fixed-size header buffer

        let cookie = reader.read_u32::<BigEndian>()?;

        if cookie != V2_COOKIE {
            return Err(DeserializeError::InvalidCookie);
        }

        let payload_len = reader.read_u32::<BigEndian>()?.to_usize()
            .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        let normalizing_offset = reader.read_u32::<BigEndian>()?;
        if normalizing_offset != 0 {
            return Err(DeserializeError::UnsupportedFeature);
        }
        let num_digits = reader.read_u32::<BigEndian>()?.to_u8()
            .ok_or(DeserializeError::InvalidParameters)?;
        let low = reader.read_u64::<BigEndian>()?;
        let high = reader.read_u64::<BigEndian>()?;
        let int_double_ratio = reader.read_f64::<BigEndian>()?;
        if int_double_ratio != 1.0 {
            return Err(DeserializeError::UnsupportedFeature);
        }

        let mut h = Histogram::new_with_bounds(low, high, num_digits)
            .map_err(|_| DeserializeError::InvalidParameters)?;

        if payload_len > self.payload_buf.len() {
            self.payload_buf.resize(payload_len, 0);
        }

        let mut payload_slice = &mut self.payload_buf[0..payload_len];
        reader.read_exact(&mut payload_slice)?;

        let mut payload_index: usize = 0;
        let mut dest_index: usize = 0;
        let mut restat_state = RestatState::new();

        while payload_index < payload_len.saturating_sub(9) {
            // Read with fast loop until we are within 9 of the end. Fast loop can't handle EOF,
            // so bail to slow version for the last few bytes.

            let (zz_num, bytes_read) = varint_read_slice(
                &payload_slice[payload_index..(payload_index + 9)]);
            payload_index += bytes_read;

            let count_or_zeros = zig_zag_decode(zz_num);

            // TODO this logic is exactly the same as the contents of the lower loop.
            if count_or_zeros < 0 {
                let zero_count = (-count_or_zeros).to_usize()
                    .ok_or(DeserializeError::UsizeTypeTooSmall)?;
                // skip the zeros
                dest_index = dest_index.checked_add(zero_count)
                    .ok_or(DeserializeError::UsizeTypeTooSmall)?;
                continue;
            } else {
                let count: T = T::from_i64(count_or_zeros)
                    .ok_or(DeserializeError::UnsuitableCounterType)?;

                h.set_count_at_index(dest_index, count)
                    .map_err(|_| DeserializeError::EncodedArrayTooLong)?;

                restat_state.on_nonzero_count(dest_index, count);

                dest_index = dest_index.checked_add(1)
                    .ok_or(DeserializeError::UsizeTypeTooSmall)?;
            }
        }

        // Now read the leftovers
        let leftover_slice = &payload_slice[payload_index..];
        let mut cursor = Cursor::new(&leftover_slice);
        while cursor.position() < leftover_slice.len() as u64 {
            let count_or_zeros = zig_zag_decode(varint_read(&mut cursor)?);

            if count_or_zeros < 0 {
                let zero_count = (-count_or_zeros).to_usize()
                    .ok_or(DeserializeError::UsizeTypeTooSmall)?;
                // skip the zeros
                dest_index = dest_index.checked_add(zero_count)
                    .ok_or(DeserializeError::UsizeTypeTooSmall)?;
                continue;
            } else {
                let count: T = T::from_i64(count_or_zeros)
                    .ok_or(DeserializeError::UnsuitableCounterType)?;

                h.set_count_at_index(dest_index, count)
                    .map_err(|_| DeserializeError::EncodedArrayTooLong)?;

                restat_state.on_nonzero_count(dest_index, count);

                dest_index = dest_index.checked_add(1)
                    .ok_or(DeserializeError::UsizeTypeTooSmall)?;
            }
        }

        restat_state.update_histogram(&mut h);

        Ok(h)
    }
}

// Only public for testing.
/// Read from a slice that must be 9 bytes long or longer. Returns the decoded number and how many
/// bytes were consumed.
#[inline]
pub fn varint_read_slice(slice: &[u8]) -> (u64, usize) {
    let mut b = slice[0];

    // take low 7 bits
    let mut value: u64 = low_7_bits(b);
    if !is_high_bit_set(b) {
        return (value, 1);
    }
    // high bit set, keep reading
    b = slice[1];
    value |= low_7_bits(b) << 7;
    if !is_high_bit_set(b) {
        return (value, 2);
    }
    b = slice[2];
    value |= low_7_bits(b) << 7 * 2;
    if !is_high_bit_set(b) {
        return (value, 3);
    }
    b = slice[3];
    value |= low_7_bits(b) << 7 * 3;
    if !is_high_bit_set(b) {
        return (value, 4);
    }
    b = slice[4];
    value |= low_7_bits(b) << 7 * 4;
    if !is_high_bit_set(b) {
        return (value, 5);
    }
    b = slice[5];
    value |= low_7_bits(b) << 7 * 5;
    if !is_high_bit_set(b) {
        return (value, 6);
    }
    b = slice[6];
    value |= low_7_bits(b) << 7 * 6;
    if !is_high_bit_set(b) {
        return (value, 7);
    }
    b = slice[7];
    value |= low_7_bits(b) << 7 * 7;
    if !is_high_bit_set(b) {
        return (value, 8);
    }

    b = slice[8];
    // special case: use last byte as is
    value |= (b as u64) << 7 * 8;

    (value, 9)
}

// Only public for testing.
/// Read a LEB128-64b9B from the buffer
pub fn varint_read<R: Read>(reader: &mut R) -> io::Result<u64> {
    let mut b = reader.read_u8()?;

    // take low 7 bits
    let mut value: u64 = low_7_bits(b);

    if is_high_bit_set(b) {
        // high bit set, keep reading
        b = reader.read_u8()?;
        value |= low_7_bits(b) << 7;
        if is_high_bit_set(b) {
            b = reader.read_u8()?;
            value |= low_7_bits(b) << 7 * 2;
            if is_high_bit_set(b) {
                b = reader.read_u8()?;
                value |= low_7_bits(b) << 7 * 3;
                if is_high_bit_set(b) {
                    b = reader.read_u8()?;
                    value |= low_7_bits(b) << 7 * 4;
                    if is_high_bit_set(b) {
                        b = reader.read_u8()?;
                        value |= low_7_bits(b) << 7 * 5;
                        if is_high_bit_set(b) {
                            b = reader.read_u8()?;
                            value |= low_7_bits(b) << 7 * 6;
                            if is_high_bit_set(b) {
                                b = reader.read_u8()?;
                                value |= low_7_bits(b) << 7 * 7;
                                if is_high_bit_set(b) {
                                    b = reader.read_u8()?;
                                    // special case: use last byte as is
                                    value |= (b as u64) << 7 * 8;
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    Ok(value)
}

/// truncate byte to low 7 bits, cast to u64
#[inline]
fn low_7_bits(b: u8) -> u64 {
    (b & 0x7F) as u64
}

#[inline]
fn is_high_bit_set(b: u8) -> bool {
    // TODO benchmark leading zeros rather than masking
    (b & 0x80) != 0
}

// Only public for testing.
#[inline]
pub fn zig_zag_decode(encoded: u64) -> i64 {
    ((encoded >> 1) as i64) ^ -((encoded & 1) as i64)
}
<commit_msg>Collapse duplicated loop bodies into a helper struct; no performance impact<commit_after>use super::V2_COOKIE;
use super::super::{Counter, Histogram, RestatState};
use super::super::num::ToPrimitive;
use std::io::{self, Cursor, ErrorKind, Read};
use std::marker::PhantomData;
use std;
use super::byteorder::{BigEndian, ReadBytesExt};

/// Errors that can happen during deserialization.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum DeserializeError {
    /// An i/o operation failed.
    IoError(ErrorKind),
    /// The cookie (first 4 bytes) did not match that for any supported format.
    InvalidCookie,
    /// The histogram uses features that this implementation doesn't support (yet), so it cannot
    /// be deserialized correctly.
    UnsupportedFeature,
    /// A count exceeded what can be represented in the chosen counter type.
    UnsuitableCounterType,
    /// The histogram instance could not be created because the serialized parameters were invalid
    /// (e.g. lowest value, highest value, etc.)
    InvalidParameters,
    /// The current system's pointer width cannot represent the encoded histogram.
    UsizeTypeTooSmall,
    /// The encoded array is longer than it should be for the histogram's value range.
    EncodedArrayTooLong
}

impl std::convert::From<std::io::Error> for DeserializeError {
    fn from(e: std::io::Error) -> Self {
        DeserializeError::IoError(e.kind())
    }
}

/// Deserializer for all supported formats.
///
/// Since the serialization formats all include some magic bytes that allow reliable identification
/// of the different formats, only one Deserializer implementation is needed.
pub struct Deserializer {
    payload_buf: Vec<u8>
}

impl Deserializer {
    /// Create a new deserializer.
    pub fn new() -> Deserializer {
        Deserializer {
            payload_buf: Vec::new()
        }
    }

    /// Deserialize an encoded histogram from the provided reader.
    ///
    /// Note that `&[u8]` and `Cursor` are convenient implementations of `Read` if you have some
    /// bytes already in slice or `Vec` form.
    pub fn deserialize<T: Counter, R: Read>(&mut self, reader: &mut R)
                                            -> Result<Histogram<T>, DeserializeError> {
        // TODO benchmark minimizing read calls by reading into a fixed-size header buffer

        let cookie = reader.read_u32::<BigEndian>()?;

        if cookie != V2_COOKIE {
            return Err(DeserializeError::InvalidCookie);
        }

        let payload_len = reader.read_u32::<BigEndian>()?.to_usize()
            .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        let normalizing_offset = reader.read_u32::<BigEndian>()?;
        if normalizing_offset != 0 {
            return Err(DeserializeError::UnsupportedFeature);
        }
        let num_digits = reader.read_u32::<BigEndian>()?.to_u8()
            .ok_or(DeserializeError::InvalidParameters)?;
        let low = reader.read_u64::<BigEndian>()?;
        let high = reader.read_u64::<BigEndian>()?;
        let int_double_ratio = reader.read_f64::<BigEndian>()?;
        if int_double_ratio != 1.0 {
            return Err(DeserializeError::UnsupportedFeature);
        }

        let mut h = Histogram::new_with_bounds(low, high, num_digits)
            .map_err(|_| DeserializeError::InvalidParameters)?;

        if payload_len > self.payload_buf.len() {
            self.payload_buf.resize(payload_len, 0);
        }

        let mut payload_slice = &mut self.payload_buf[0..payload_len];
        reader.read_exact(&mut payload_slice)?;

        let mut payload_index: usize = 0;
        let mut restat_state = RestatState::new();
        let mut decode_state = DecodeLoopState::new();

        while payload_index < payload_len.saturating_sub(9) {
            // Read with fast loop until we are within 9 of the end. Fast loop can't handle EOF,
            // so bail to slow version for the last few bytes.

            let (zz_num, bytes_read) = varint_read_slice(
                &payload_slice[payload_index..(payload_index + 9)]);
            payload_index += bytes_read;

            let count_or_zeros = zig_zag_decode(zz_num);

            decode_state.on_decoded_num(count_or_zeros, &mut restat_state, &mut h)?;
        }

        // Now read the leftovers
        let leftover_slice = &payload_slice[payload_index..];
        let mut cursor = Cursor::new(&leftover_slice);
        while cursor.position() < leftover_slice.len() as u64 {
            let count_or_zeros = zig_zag_decode(varint_read(&mut cursor)?);

            decode_state.on_decoded_num(count_or_zeros, &mut restat_state, &mut h)?;
        }

        restat_state.update_histogram(&mut h);

        Ok(h)
    }
}

// Only public for testing.
/// Read from a slice that must be 9 bytes long or longer. Returns the decoded number and how many
/// bytes were consumed.
#[inline]
pub fn varint_read_slice(slice: &[u8]) -> (u64, usize) {
    let mut b = slice[0];

    // take low 7 bits
    let mut value: u64 = low_7_bits(b);
    if !is_high_bit_set(b) {
        return (value, 1);
    }
    // high bit set, keep reading
    b = slice[1];
    value |= low_7_bits(b) << 7;
    if !is_high_bit_set(b) {
        return (value, 2);
    }
    b = slice[2];
    value |= low_7_bits(b) << 7 * 2;
    if !is_high_bit_set(b) {
        return (value, 3);
    }
    b = slice[3];
    value |= low_7_bits(b) << 7 * 3;
    if !is_high_bit_set(b) {
        return (value, 4);
    }
    b = slice[4];
    value |= low_7_bits(b) << 7 * 4;
    if !is_high_bit_set(b) {
        return (value, 5);
    }
    b = slice[5];
    value |= low_7_bits(b) << 7 * 5;
    if !is_high_bit_set(b) {
        return (value, 6);
    }
    b = slice[6];
    value |= low_7_bits(b) << 7 * 6;
    if !is_high_bit_set(b) {
        return (value, 7);
    }
    b = slice[7];
    value |= low_7_bits(b) << 7 * 7;
    if !is_high_bit_set(b) {
        return (value, 8);
    }

    b = slice[8];
    // special case: use last byte as is
    value |= (b as u64) << 7 * 8;

    (value, 9)
}

// Only public for testing.
/// Read a LEB128-64b9B from the buffer
pub fn varint_read<R: Read>(reader: &mut R) -> io::Result<u64> {
    let mut b = reader.read_u8()?;

    // take low 7 bits
    let mut value: u64 = low_7_bits(b);

    if is_high_bit_set(b) {
        // high bit set, keep reading
        b = reader.read_u8()?;
        value |= low_7_bits(b) << 7;
        if is_high_bit_set(b) {
            b = reader.read_u8()?;
            value |= low_7_bits(b) << 7 * 2;
            if is_high_bit_set(b) {
                b = reader.read_u8()?;
                value |= low_7_bits(b) << 7 * 3;
                if is_high_bit_set(b) {
                    b = reader.read_u8()?;
                    value |= low_7_bits(b) << 7 * 4;
                    if is_high_bit_set(b) {
                        b = reader.read_u8()?;
                        value |= low_7_bits(b) << 7 * 5;
                        if is_high_bit_set(b) {
                            b = reader.read_u8()?;
                            value |= low_7_bits(b) << 7 * 6;
                            if is_high_bit_set(b) {
                                b = reader.read_u8()?;
                                value |= low_7_bits(b) << 7 * 7;
                                if is_high_bit_set(b) {
                                    b = reader.read_u8()?;
                                    // special case: use last byte as is
                                    value |= (b as u64) << 7 * 8;
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    Ok(value)
}

/// truncate byte to low 7 bits, cast to u64
#[inline]
fn low_7_bits(b: u8) -> u64 {
    (b & 0x7F) as u64
}

#[inline]
fn is_high_bit_set(b: u8) -> bool {
    // TODO benchmark leading zeros rather than masking
    (b & 0x80) != 0
}

// Only public for testing.
#[inline]
pub fn zig_zag_decode(encoded: u64) -> i64 {
    ((encoded >> 1) as i64) ^ -((encoded & 1) as i64)
}

/// We need to perform the same logic in two different decode loops while carrying over a modicum
/// of state.
struct DecodeLoopState<T: Counter> {
    dest_index:usize,
    phantom: PhantomData<T>
}

impl <T: Counter> DecodeLoopState<T> {

    fn new() -> DecodeLoopState<T> {
        DecodeLoopState {
            dest_index: 0,
            phantom: PhantomData
        }
    }

    #[inline]
    fn on_decoded_num(&mut self, count_or_zeros: i64, restat_state: &mut RestatState<T>,
                      h: &mut Histogram<T>) -> Result<(), DeserializeError> {
        if count_or_zeros < 0 {
            let zero_count = (-count_or_zeros).to_usize()
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
            // skip the zeros
            self.dest_index = self.dest_index.checked_add(zero_count)
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        } else {
            let count: T = T::from_i64(count_or_zeros)
                .ok_or(DeserializeError::UnsuitableCounterType)?;

            h.set_count_at_index(self.dest_index, count)
                .map_err(|_| DeserializeError::EncodedArrayTooLong)?;

            restat_state.on_nonzero_count(self.dest_index, count);

            self.dest_index = self.dest_index.checked_add(1)
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        }

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


<commit_before>//! Insertion sort algorithm.
//!
//! The `insertion_sort` module contains the simple sorting algorithm.

use std::fmt::Debug;

pub fn sort<T: Debug + PartialOrd>(input: &Vec<T>) -> Vec<usize> /*&[T]*/ {
    let n = input.len();

    let mut output_index: Vec<usize> = Vec::with_capacity(n);

    // Place the first element's index (0) as the first element of the output_index vector.
    output_index.push(0);

    println!("n = {:?}", n);
    println!("input:        {:?}", input);
    println!("output_index: {:?}", output_index);

    if n > 1 {
        // Loop over input vector, skipping the first element as it's already inserted as the first
        // element of output_index.
        for (i_start_at_0,elem) in input.iter().skip(1).enumerate() {
            let i = i_start_at_0 + 1;
            println!("  i={:?} Finding where to put element {:?} (i={:?}) in output_index: {:?}", i, elem, i, output_index);

            // Verify next element if it is smaller
            // If next element is larger than the last stored index, just place that element at
            // the end.
            if *elem >= input[*output_index.last().unwrap()] {
                // Next element "elem" is already larger than the previous stored index.
                // Just use its index.
                output_index.push(i);

                println!("    1. Pushing element '{:?}' at the end.", elem);
            } else {
                println!("    2. elem ?< input[output_index[0]]");
                println!("           {:?} ?< {:?}", elem, input[output_index[0]]);
                // The next element of the input vector is smaller than the last stored on.
                // Let's find where it should go by looping back on the stored indices.
                // NOTE: we must use a loop on indices as to prevent problems with the borrow-checker.
                if *elem < input[output_index[0]] {
                    println!("        2a. Yes: --> putting i in front");
                    output_index.insert(0, i);
                } else {
                    println!("        2b. No: --> Looping back");
                    for j in (0..i).rev() {
                        println!("            j={:?}   elem={:?}   input[output_index[j]]={:?}", j, elem, input[output_index[j]]);
                        if *elem > input[output_index[j]] {
                            output_index.insert(j+1, i);
                            break;
                        }
                    }
                }
                println!("        3. output_index: {:?}", output_index);
            }

            print!("           output: [");
            for j in output_index.iter() {
                    print!("{:?}, ", input[*j]);
            }
            println!("]");
        }
    }

    println!("END: output_index: {:?}", output_index);
    print!("END: output: [");
    for j in output_index.iter() {
            print!("{:?}, ", input[*j]);
    }
    println!("]");

    assert_eq!(output_index.len(), n);

    return output_index;
}
<commit_msg>Cleanup<commit_after>//! Insertion sort algorithm.
//!
//! The `insertion_sort` module contains the simple sorting algorithm.

use std::fmt::Debug;

pub fn sort<T: Debug + PartialOrd>(input: &Vec<T>) -> Vec<usize> {
    let n = input.len();

    let mut output_index: Vec<usize> = Vec::with_capacity(n);

    // Place the first element's index (0) as the first element of the output_index vector.
    output_index.push(0);

    println!("n = {:?}", n);
    println!("input:        {:?}", input);
    println!("output_index: {:?}", output_index);

    if n > 1 {
        // Loop over input vector, skipping the first element as it's already inserted as the first
        // element of output_index.
        for (i_start_at_0,elem) in input.iter().skip(1).enumerate() {
            let i = i_start_at_0 + 1;
            println!("  i={:?} Finding where to put element {:?} (i={:?}) in output_index: {:?}", i, elem, i, output_index);

            // Verify next element if it is smaller
            // If next element is larger than the last stored index, just place that element at
            // the end.
            if *elem >= input[*output_index.last().unwrap()] {
                // Next element "elem" is already larger than the previous stored index.
                // Just use its index.
                output_index.push(i);

                println!("    1. Pushing element '{:?}' at the end.", elem);
            } else {
                println!("    2. elem ?< input[output_index[0]]");
                println!("           {:?} ?< {:?}", elem, input[output_index[0]]);
                // The next element of the input vector is smaller than the last stored on.
                // Let's find where it should go by looping back on the stored indices.
                // NOTE: we must use a loop on indices as to prevent problems with the borrow-checker.
                if *elem < input[output_index[0]] {
                    println!("        2a. Yes: --> putting i in front");
                    output_index.insert(0, i);
                } else {
                    println!("        2b. No: --> Looping back");
                    for j in (0..i).rev() {
                        println!("            j={:?}   elem={:?}   input[output_index[j]]={:?}", j, elem, input[output_index[j]]);
                        if *elem > input[output_index[j]] {
                            output_index.insert(j+1, i);
                            break;
                        }
                    }
                }
                println!("        3. output_index: {:?}", output_index);
            }

            print!("           output: [");
            for j in output_index.iter() {
                    print!("{:?}, ", input[*j]);
            }
            println!("]");
        }
    }

    println!("END: output_index: {:?}", output_index);
    print!("END: output: [");
    for j in output_index.iter() {
            print!("{:?}, ", input[*j]);
    }
    println!("]");

    assert_eq!(output_index.len(), n);

    return output_index;
}
<|endoftext|>


<commit_before>// FIXME This is the simplest most naive implementation of an allocator. I
// basically envisioned it in C in my head and translated it on the fly. I hope
// to revisit this someday when I know the language better, but in the meantime
// don't take this as an example for anything ever.

#![allocator]

use core::ptr;
use core::mem;

use self::spin::Mutex;

extern crate spin;

const SPLIT_FUDGE_FACTOR: usize = 32;

#[repr(C)]
struct Region {
    size: usize,
    prev: Option<*mut Region>,
    next: Option<*mut Region>,
}

pub struct Allocator {
    heap_start: *mut u8,
    heap_size: usize,
    free_list: Option<*mut Region>,
}

static mut ALLOCATOR: Mutex<Allocator> = Mutex::new(Allocator {
    heap_start: 0 as *mut u8,
    heap_size: 0,
    free_list: None,
});

pub fn init_global_allocator(size: u64, raw_bytes: *mut u8) {
    static mut CALLED: bool = false;

    unsafe {
        assert!(!CALLED);
        init_allocator(&mut ALLOCATOR.lock(), size as usize, raw_bytes);
        CALLED = true;
    }
}

pub fn init_allocator(allocator: &mut Allocator, size: usize, raw_bytes: *mut u8) {

    assert!(!raw_bytes.is_null());

    let mut region = raw_bytes as *mut Region;
    unsafe {
        (*region).size = size;
        (*region).next = None;
        (*region).prev = None;
    }
    allocator.heap_size = size;
    allocator.heap_start = raw_bytes;
    allocator.free_list = Some(raw_bytes as *mut Region);
}

impl Allocator {
    pub fn free(&mut self, bytes: *mut u8) {

        let mut region: *mut Region;

        if bytes.is_null() {
            return;
        }

        unsafe {
            let offset = -(mem::size_of::<Region>() as isize);
            region = bytes.offset(offset) as *mut Region;

            assert!(region as *mut u8 >= self.heap_start);

            assert!(!region.is_null());

            (*region).prev = None;
            (*region).next = self.free_list;
        }

        self.free_list = Some(region);
    }

    pub fn alloc(&mut self, size: usize) -> *mut u8 {

        let region_offset = mem::size_of::<Region>() as isize;
        let real_size = size + mem::size_of::<Region>();
        let mut current_region = self.free_list;

        while let Some(region) = current_region {

            unsafe {
                if (*region).size >= real_size {
                    let new_next: Option<*mut Region>;
                    let new_prev: Option<*mut Region>;
                    let split = (*region).size - real_size;
                    if split > (region_offset as usize) + SPLIT_FUDGE_FACTOR {
                        let split_region =
                            (region as *mut u8)
                                .offset((size as isize) +
                                        region_offset) as *mut Region;
                        new_next = Some(split_region);
                        new_prev = Some(split_region);
                    } else {
                        new_next = (*region).next;
                        new_prev = (*region).prev;
                    }
                    match (*region).prev {
                        None => {}
                        Some(prev) => (*prev).next = new_next,
                    }
                    match (*region).next {
                        None => {}
                        Some(next) => (*next).prev = new_prev,
                    }
                    if new_next == None && new_prev == None {
                        self.free_list = None;
                    }
                    return (region as *mut u8).offset(region_offset);
                }
                current_region = (*region).next;
            }
        }

        ptr::null::<u8>() as *mut u8
    }
}

#[cfg(not(test))]
#[no_mangle]
pub extern "C" fn __rust_allocate(size: usize, _align: usize) -> *mut u8 {
    unsafe { ALLOCATOR.lock().alloc(size) }
}

#[cfg(not(test))]
#[no_mangle]
pub extern "C" fn __rust_deallocate(ptr: *mut u8, _old_size: usize, _align: usize) {
    unsafe {
        ALLOCATOR.lock().free(ptr);
    }
}

#[cfg(not(test))]
#[no_mangle]
pub extern "C" fn __rust_reallocate(ptr: *mut u8,
                                    old_size: usize,
                                    new_size: usize,
                                    align: usize)
                                    -> *mut u8 {
    use core::cmp;
    let new_ptr = __rust_allocate(new_size, align);
    unsafe { ptr::copy(ptr, new_ptr, cmp::min(old_size, new_size)) };
    __rust_deallocate(ptr, old_size, align);
    new_ptr
}

#[cfg(not(test))]
#[no_mangle]
pub extern "C" fn __rust_reallocate_inplace(_ptr: *mut u8,
                                            old_size: usize,
                                            _size: usize,
                                            _align: usize)
                                            -> usize {
    old_size
}

#[cfg(not(test))]
#[no_mangle]
pub extern "C" fn __rust_usable_size(size: usize, _align: usize) -> usize {
    size
}

#[cfg(test)]
mod tests {
    use core::mem;

    use super::Allocator;
    use super::Region;
    use super::init_allocator;

    #[test]
    fn test_init_allocator() {

        let mut al = Allocator {
            heap_size: 0,
            heap_start: 0 as *mut u8,
            free_list: None,
        };

        let mut memory: [u8; 100] = [0; 100];
        let pmemory: *mut u8 = memory.as_mut_ptr();
        // init_allocator(&mut al, mem::size_of::<[u8; 100]>(), pmemory);
        init_allocator(&mut al, 100, pmemory);
        assert_eq!(al.heap_size, 100);
        assert_eq!(al.heap_start, pmemory);

        match al.free_list {
            None => assert!(false),
            Some(fl) => {
                assert!(!fl.is_null());
                unsafe {
                    assert_eq!((*fl).size, 100);
                    assert_eq!((*fl).next, None);
                };
            }
        }
    }

    #[test]
    fn test_alloc_and_free() {

        let mut al = Allocator {
            heap_size: 0,
            heap_start: 0 as *mut u8,
            free_list: None,
        };

        let mut memory: [u8; 1000] = [0; 1000];
        let pmemory: *mut u8 = memory.as_mut_ptr();
        init_allocator(&mut al, mem::size_of::<[u8; 1000]>(), pmemory);

        let m0_size = 10000 - mem::size_of::<Region>();
        let m0 = al.alloc(m0_size);
        assert!(m0.is_null());

        let m1_size = 1000 - mem::size_of::<Region>();
        let m1 = al.alloc(m1_size);
        assert!(!m1.is_null());

        assert_eq!(al.free_list, None);
        al.free(m1);
        assert_ne!(al.free_list, None);

        let m2_size = 100;
        let m2 = al.alloc(m2_size);
        assert!(!m2.is_null());

        let m3_size = 10;
        let m3 = al.alloc(m3_size);
        assert!(!m3.is_null());

        let m4_size = 1;
        let m4 = al.alloc(m4_size);
        assert!(!m4.is_null());

        assert_ne!(al.free_list, None);
        al.free(m3);

        assert_ne!(al.free_list, None);
        al.free(m2);

        assert_ne!(al.free_list, None);
        al.free(m4);
    }
}
<commit_msg>Fix commented line<commit_after>// FIXME This is the simplest most naive implementation of an allocator. I
// basically envisioned it in C in my head and translated it on the fly. I hope
// to revisit this someday when I know the language better, but in the meantime
// don't take this as an example for anything ever.

#![allocator]

use core::ptr;
use core::mem;

use self::spin::Mutex;

extern crate spin;

const SPLIT_FUDGE_FACTOR: usize = 32;

#[repr(C)]
struct Region {
    size: usize,
    prev: Option<*mut Region>,
    next: Option<*mut Region>,
}

pub struct Allocator {
    heap_start: *mut u8,
    heap_size: usize,
    free_list: Option<*mut Region>,
}

static mut ALLOCATOR: Mutex<Allocator> = Mutex::new(Allocator {
    heap_start: 0 as *mut u8,
    heap_size: 0,
    free_list: None,
});

pub fn init_global_allocator(size: u64, raw_bytes: *mut u8) {
    static mut CALLED: bool = false;

    unsafe {
        assert!(!CALLED);
        init_allocator(&mut ALLOCATOR.lock(), size as usize, raw_bytes);
        CALLED = true;
    }
}

pub fn init_allocator(allocator: &mut Allocator, size: usize, raw_bytes: *mut u8) {

    assert!(!raw_bytes.is_null());

    let mut region = raw_bytes as *mut Region;
    unsafe {
        (*region).size = size;
        (*region).next = None;
        (*region).prev = None;
    }
    allocator.heap_size = size;
    allocator.heap_start = raw_bytes;
    allocator.free_list = Some(raw_bytes as *mut Region);
}

impl Allocator {
    pub fn free(&mut self, bytes: *mut u8) {

        let mut region: *mut Region;

        if bytes.is_null() {
            return;
        }

        unsafe {
            let offset = -(mem::size_of::<Region>() as isize);
            region = bytes.offset(offset) as *mut Region;

            assert!(region as *mut u8 >= self.heap_start);

            assert!(!region.is_null());

            (*region).prev = None;
            (*region).next = self.free_list;
        }

        self.free_list = Some(region);
    }

    pub fn alloc(&mut self, size: usize) -> *mut u8 {

        let region_offset = mem::size_of::<Region>() as isize;
        let real_size = size + mem::size_of::<Region>();
        let mut current_region = self.free_list;

        while let Some(region) = current_region {

            unsafe {
                if (*region).size >= real_size {
                    let new_next: Option<*mut Region>;
                    let new_prev: Option<*mut Region>;
                    let split = (*region).size - real_size;
                    if split > (region_offset as usize) + SPLIT_FUDGE_FACTOR {
                        let split_region =
                            (region as *mut u8)
                                .offset((size as isize) +
                                        region_offset) as *mut Region;
                        new_next = Some(split_region);
                        new_prev = Some(split_region);
                    } else {
                        new_next = (*region).next;
                        new_prev = (*region).prev;
                    }
                    match (*region).prev {
                        None => {}
                        Some(prev) => (*prev).next = new_next,
                    }
                    match (*region).next {
                        None => {}
                        Some(next) => (*next).prev = new_prev,
                    }
                    if new_next == None && new_prev == None {
                        self.free_list = None;
                    }
                    return (region as *mut u8).offset(region_offset);
                }
                current_region = (*region).next;
            }
        }

        ptr::null::<u8>() as *mut u8
    }
}

#[cfg(not(test))]
#[no_mangle]
pub extern "C" fn __rust_allocate(size: usize, _align: usize) -> *mut u8 {
    unsafe { ALLOCATOR.lock().alloc(size) }
}

#[cfg(not(test))]
#[no_mangle]
pub extern "C" fn __rust_deallocate(ptr: *mut u8, _old_size: usize, _align: usize) {
    unsafe {
        ALLOCATOR.lock().free(ptr);
    }
}

#[cfg(not(test))]
#[no_mangle]
pub extern "C" fn __rust_reallocate(ptr: *mut u8,
                                    old_size: usize,
                                    new_size: usize,
                                    align: usize)
                                    -> *mut u8 {
    use core::cmp;
    let new_ptr = __rust_allocate(new_size, align);
    unsafe { ptr::copy(ptr, new_ptr, cmp::min(old_size, new_size)) };
    __rust_deallocate(ptr, old_size, align);
    new_ptr
}

#[cfg(not(test))]
#[no_mangle]
pub extern "C" fn __rust_reallocate_inplace(_ptr: *mut u8,
                                            old_size: usize,
                                            _size: usize,
                                            _align: usize)
                                            -> usize {
    old_size
}

#[cfg(not(test))]
#[no_mangle]
pub extern "C" fn __rust_usable_size(size: usize, _align: usize) -> usize {
    size
}

#[cfg(test)]
mod tests {
    use core::mem;

    use super::Allocator;
    use super::Region;
    use super::init_allocator;

    #[test]
    fn test_init_allocator() {

        let mut al = Allocator {
            heap_size: 0,
            heap_start: 0 as *mut u8,
            free_list: None,
        };

        let mut memory: [u8; 100] = [0; 100];
        let pmemory: *mut u8 = memory.as_mut_ptr();
        init_allocator(&mut al, mem::size_of::<[u8; 100]>(), pmemory);
        assert_eq!(al.heap_size, 100);
        assert_eq!(al.heap_start, pmemory);

        match al.free_list {
            None => assert!(false),
            Some(fl) => {
                assert!(!fl.is_null());
                unsafe {
                    assert_eq!((*fl).size, 100);
                    assert_eq!((*fl).next, None);
                };
            }
        }
    }

    #[test]
    fn test_alloc_and_free() {

        let mut al = Allocator {
            heap_size: 0,
            heap_start: 0 as *mut u8,
            free_list: None,
        };

        let mut memory: [u8; 1000] = [0; 1000];
        let pmemory: *mut u8 = memory.as_mut_ptr();
        init_allocator(&mut al, mem::size_of::<[u8; 1000]>(), pmemory);

        let m0_size = 10000 - mem::size_of::<Region>();
        let m0 = al.alloc(m0_size);
        assert!(m0.is_null());

        let m1_size = 1000 - mem::size_of::<Region>();
        let m1 = al.alloc(m1_size);
        assert!(!m1.is_null());

        assert_eq!(al.free_list, None);
        al.free(m1);
        assert_ne!(al.free_list, None);

        let m2_size = 100;
        let m2 = al.alloc(m2_size);
        assert!(!m2.is_null());

        let m3_size = 10;
        let m3 = al.alloc(m3_size);
        assert!(!m3.is_null());

        let m4_size = 1;
        let m4 = al.alloc(m4_size);
        assert!(!m4.is_null());

        assert_ne!(al.free_list, None);
        al.free(m3);

        assert_ne!(al.free_list, None);
        al.free(m2);

        assert_ne!(al.free_list, None);
        al.free(m4);
    }
}
<|endoftext|>


<commit_before>extern crate sysfs_gpio;
extern crate sysfs_pwm;

use self::sysfs_gpio::{Direction, Pin};
use self::sysfs_pwm::{Pwm};
use ::std::error::Error;

#[test]
fn it_works() {
    assert_eq!(1, 0+1);
}
<commit_msg>expensive test<commit_after>extern crate sysfs_gpio;
extern crate sysfs_pwm;

use self::sysfs_gpio::{Direction, Pin};
use self::sysfs_pwm::{Pwm};
use ::std::error::Error;

#[test]
fn it_works() {
    assert_eq!(1, 0+1);
}

const PWM_CHIP: u32 = 0;
const PWM_NUMBER: u32 = 0;

//http://www.jumpnowtek.com/rpi/Using-the-Raspberry-Pi-Hardware-PWM-timers.html
//https://github.com/rust-embedded/rust-sysfs-pwm/blob/master/examples/breathe.rs
#[test]
#[ignore]
fn testPwm() {
    let pwm = Pwm::new(0, 0).unwrap();
    pwm.export().unwrap();
    pwm.enable(true).unwrap();
    pwm.set_period_ns(10000000).unwrap();
    pwm.set_duty_cycle_ns(8000000).unwrap();
    pwm.enable(false).unwrap();
    pwm.unexport().unwrap()
}<|endoftext|>