<commit_before>//! Implements syntax highlighting.

mod tags;
mod html;
#[cfg(test)]
mod tests;

use hir::{Name, Semantics};
use ra_ide_db::{
    defs::{classify_name, NameClass, NameDefinition},
    RootDatabase,
};
use ra_prof::profile;
use ra_syntax::{
    ast::{self, HasQuotes, HasStringValue},
    AstNode, AstToken, Direction, NodeOrToken, SyntaxElement,
    SyntaxKind::*,
    SyntaxToken, TextRange, WalkEvent, T,
};
use rustc_hash::FxHashMap;

use crate::{
    call_info::call_info_for_token,
    references::{classify_name_ref, NameRefClass},
    Analysis, FileId,
};

pub(crate) use html::highlight_as_html;
pub use tags::{Highlight, HighlightModifier, HighlightModifiers, HighlightTag};

#[derive(Debug)]
pub struct HighlightedRange {
    pub range: TextRange,
    pub highlight: Highlight,
    pub binding_hash: Option<u64>,
}

pub(crate) fn highlight(
    db: &RootDatabase,
    file_id: FileId,
    range_to_highlight: Option<TextRange>,
) -> Vec<HighlightedRange> {
    let _p = profile("highlight");
    let sema = Semantics::new(db);

    // Determine the root based on the given range.
    let (root, range_to_highlight) = {
        let source_file = sema.parse(file_id);
        match range_to_highlight {
            Some(range) => {
                let node = match source_file.syntax().covering_element(range) {
                    NodeOrToken::Node(it) => it,
                    NodeOrToken::Token(it) => it.parent(),
                };
                (node, range)
            }
            None => (source_file.syntax().clone(), source_file.syntax().text_range()),
        }
    };

    let mut bindings_shadow_count: FxHashMap<Name, u32> = FxHashMap::default();
    let mut res = Vec::new();

    let mut current_macro_call: Option<ast::MacroCall> = None;

    // Walk all nodes, keeping track of whether we are inside a macro or not.
    // If in macro, expand it first and highlight the expanded code.
    for event in root.preorder_with_tokens() {
        let event_range = match &event {
            WalkEvent::Enter(it) => it.text_range(),
            WalkEvent::Leave(it) => it.text_range(),
        };

        // Element outside of the viewport, no need to highlight
        if range_to_highlight.intersection(&event_range).is_none() {
            continue;
        }

        // Track "inside macro" state
        match event.clone().map(|it| it.into_node().and_then(ast::MacroCall::cast)) {
            WalkEvent::Enter(Some(mc)) => {
                current_macro_call = Some(mc.clone());
                if let Some(range) = macro_call_range(&mc) {
                    res.push(HighlightedRange {
                        range,
                        highlight: HighlightTag::Macro.into(),
                        binding_hash: None,
                    });
                }
                continue;
            }
            WalkEvent::Leave(Some(mc)) => {
                assert!(current_macro_call == Some(mc));
                current_macro_call = None;
                continue;
            }
            _ => (),
        }

        let element = match event {
            WalkEvent::Enter(it) => it,
            WalkEvent::Leave(_) => continue,
        };

        let range = element.text_range();

        let element_to_highlight = if current_macro_call.is_some() {
            // Inside a macro -- expand it first
            let token = match element.clone().into_token() {
                Some(it) if it.parent().kind() == TOKEN_TREE => it,
                _ => continue,
            };
            let token = sema.descend_into_macros(token.clone());
            let parent = token.parent();
            // We only care Name and Name_ref
            match (token.kind(), parent.kind()) {
                (IDENT, NAME) | (IDENT, NAME_REF) => parent.into(),
                _ => token.into(),
            }
        } else {
            element.clone()
        };

        if let Some(token) = element.as_token().cloned().and_then(ast::RawString::cast) {
            let expanded = element_to_highlight.as_token().unwrap().clone();
            if highlight_injection(&mut res, &sema, token, expanded).is_some() {
                continue;
            }
        }

        if let Some((highlight, binding_hash)) =
            highlight_element(&sema, &mut bindings_shadow_count, element_to_highlight)
        {
            res.push(HighlightedRange { range, highlight, binding_hash });
        }
    }

    res
}

fn macro_call_range(macro_call: &ast::MacroCall) -> Option<TextRange> {
    let path = macro_call.path()?;
    let name_ref = path.segment()?.name_ref()?;

    let range_start = name_ref.syntax().text_range().start();
    let mut range_end = name_ref.syntax().text_range().end();
    for sibling in path.syntax().siblings_with_tokens(Direction::Next) {
        match sibling.kind() {
            T![!] | IDENT => range_end = sibling.text_range().end(),
            _ => (),
        }
    }

    Some(TextRange::from_to(range_start, range_end))
}

fn highlight_element(
    sema: &Semantics<RootDatabase>,
    bindings_shadow_count: &mut FxHashMap<Name, u32>,
    element: SyntaxElement,
) -> Option<(Highlight, Option<u64>)> {
    let db = sema.db;
    let mut binding_hash = None;
    let highlight: Highlight = match element.kind() {
        FN_DEF => {
            bindings_shadow_count.clear();
            return None;
        }

        // Highlight definitions depending on the "type" of the definition.
        NAME => {
            let name = element.into_node().and_then(ast::Name::cast).unwrap();
            let name_kind = classify_name(sema, &name);

            if let Some(NameClass::NameDefinition(NameDefinition::Local(local))) = &name_kind {
                if let Some(name) = local.name(db) {
                    let shadow_count = bindings_shadow_count.entry(name.clone()).or_default();
                    *shadow_count += 1;
                    binding_hash = Some(calc_binding_hash(&name, *shadow_count))
                }
            };

            match name_kind {
                Some(NameClass::NameDefinition(def)) => {
                    highlight_name(db, def) | HighlightModifier::Definition
                }
                Some(NameClass::ConstReference(def)) => highlight_name(db, def),
                None => highlight_name_by_syntax(name) | HighlightModifier::Definition,
            }
        }

        // Highlight references like the definitions they resolve to
        NAME_REF if element.ancestors().any(|it| it.kind() == ATTR) => return None,
        NAME_REF => {
            let name_ref = element.into_node().and_then(ast::NameRef::cast).unwrap();
            let name_kind = classify_name_ref(sema, &name_ref)?;

            match name_kind {
                NameRefClass::NameDefinition(def) => {
                    if let NameDefinition::Local(local) = &def {
                        if let Some(name) = local.name(db) {
                            let shadow_count =
                                bindings_shadow_count.entry(name.clone()).or_default();
                            binding_hash = Some(calc_binding_hash(&name, *shadow_count))
                        }
                    };
                    highlight_name(db, def)
                }
                NameRefClass::FieldShorthand { .. } => HighlightTag::Field.into(),
            }
        }

        // Simple token-based highlighting
        COMMENT => HighlightTag::Comment.into(),
        STRING | RAW_STRING | RAW_BYTE_STRING | BYTE_STRING => HighlightTag::StringLiteral.into(),
        ATTR => HighlightTag::Attribute.into(),
        INT_NUMBER | FLOAT_NUMBER => HighlightTag::NumericLiteral.into(),
        BYTE => HighlightTag::ByteLiteral.into(),
        CHAR => HighlightTag::CharLiteral.into(),
        LIFETIME => {
            let h = Highlight::new(HighlightTag::Lifetime);
            dbg!(match element.parent().map(|it| it.kind()) {
                Some(LIFETIME_PARAM) | Some(LABEL) => h | HighlightModifier::Definition,
                _ => h,
            })
        }

        k if k.is_keyword() => {
            let h = Highlight::new(HighlightTag::Keyword);
            match k {
                T![break]
                | T![continue]
                | T![else]
                | T![for]
                | T![if]
                | T![loop]
                | T![match]
                | T![return]
                | T![while] => h | HighlightModifier::Control,
                T![unsafe] => h | HighlightModifier::Unsafe,
                _ => h,
            }
        }

        _ => return None,
    };

    return Some((highlight, binding_hash));

    fn calc_binding_hash(name: &Name, shadow_count: u32) -> u64 {
        fn hash<T: std::hash::Hash + std::fmt::Debug>(x: T) -> u64 {
            use std::{collections::hash_map::DefaultHasher, hash::Hasher};

            let mut hasher = DefaultHasher::new();
            x.hash(&mut hasher);
            hasher.finish()
        }

        hash((name, shadow_count))
    }
}

fn highlight_name(db: &RootDatabase, def: NameDefinition) -> Highlight {
    match def {
        NameDefinition::Macro(_) => HighlightTag::Macro,
        NameDefinition::StructField(_) => HighlightTag::Field,
        NameDefinition::ModuleDef(def) => match def {
            hir::ModuleDef::Module(_) => HighlightTag::Module,
            hir::ModuleDef::Function(_) => HighlightTag::Function,
            hir::ModuleDef::Adt(hir::Adt::Struct(_)) => HighlightTag::Struct,
            hir::ModuleDef::Adt(hir::Adt::Enum(_)) => HighlightTag::Enum,
            hir::ModuleDef::Adt(hir::Adt::Union(_)) => HighlightTag::Union,
            hir::ModuleDef::EnumVariant(_) => HighlightTag::EnumVariant,
            hir::ModuleDef::Const(_) => HighlightTag::Constant,
            hir::ModuleDef::Static(_) => HighlightTag::Static,
            hir::ModuleDef::Trait(_) => HighlightTag::Trait,
            hir::ModuleDef::TypeAlias(_) => HighlightTag::TypeAlias,
            hir::ModuleDef::BuiltinType(_) => HighlightTag::BuiltinType,
        },
        NameDefinition::SelfType(_) => HighlightTag::SelfType,
        NameDefinition::TypeParam(_) => HighlightTag::TypeParam,
        // FIXME: distinguish between locals and parameters
        NameDefinition::Local(local) => {
            let mut h = Highlight::new(HighlightTag::Local);
            if local.is_mut(db) || local.ty(db).is_mutable_reference() {
                h |= HighlightModifier::Mutable;
            }
            return h;
        }
    }
    .into()
}

fn highlight_name_by_syntax(name: ast::Name) -> Highlight {
    let default = HighlightTag::Function.into();

    let parent = match name.syntax().parent() {
        Some(it) => it,
        _ => return default,
    };

    match parent.kind() {
        STRUCT_DEF => HighlightTag::Struct.into(),
        ENUM_DEF => HighlightTag::Enum.into(),
        UNION_DEF => HighlightTag::Union.into(),
        TRAIT_DEF => HighlightTag::Trait.into(),
        TYPE_ALIAS_DEF => HighlightTag::TypeAlias.into(),
        TYPE_PARAM => HighlightTag::TypeParam.into(),
        RECORD_FIELD_DEF => HighlightTag::Field.into(),
        _ => default,
    }
}

fn highlight_injection(
    acc: &mut Vec<HighlightedRange>,
    sema: &Semantics<RootDatabase>,
    literal: ast::RawString,
    expanded: SyntaxToken,
) -> Option<()> {
    let call_info = call_info_for_token(&sema, expanded)?;
    let idx = call_info.active_parameter?;
    let name = call_info.signature.parameter_names.get(idx)?;
    if name != "ra_fixture" {
        return None;
    }
    let value = literal.value()?;
    let (analysis, tmp_file_id) = Analysis::from_single_file(value);

    if let Some(range) = literal.open_quote_text_range() {
        acc.push(HighlightedRange {
            range,
            highlight: HighlightTag::StringLiteral.into(),
            binding_hash: None,
        })
    }

    for mut h in analysis.highlight(tmp_file_id).unwrap() {
        if let Some(r) = literal.map_range_up(h.range) {
            h.range = r;
            acc.push(h)
        }
    }

    if let Some(range) = literal.close_quote_text_range() {
        acc.push(HighlightedRange {
            range,
            highlight: HighlightTag::StringLiteral.into(),
            binding_hash: None,
        })
    }

    Some(())
}
<commit_msg>Remove stray dbg<commit_after>//! Implements syntax highlighting.

mod tags;
mod html;
#[cfg(test)]
mod tests;

use hir::{Name, Semantics};
use ra_ide_db::{
    defs::{classify_name, NameClass, NameDefinition},
    RootDatabase,
};
use ra_prof::profile;
use ra_syntax::{
    ast::{self, HasQuotes, HasStringValue},
    AstNode, AstToken, Direction, NodeOrToken, SyntaxElement,
    SyntaxKind::*,
    SyntaxToken, TextRange, WalkEvent, T,
};
use rustc_hash::FxHashMap;

use crate::{
    call_info::call_info_for_token,
    references::{classify_name_ref, NameRefClass},
    Analysis, FileId,
};

pub(crate) use html::highlight_as_html;
pub use tags::{Highlight, HighlightModifier, HighlightModifiers, HighlightTag};

#[derive(Debug)]
pub struct HighlightedRange {
    pub range: TextRange,
    pub highlight: Highlight,
    pub binding_hash: Option<u64>,
}

pub(crate) fn highlight(
    db: &RootDatabase,
    file_id: FileId,
    range_to_highlight: Option<TextRange>,
) -> Vec<HighlightedRange> {
    let _p = profile("highlight");
    let sema = Semantics::new(db);

    // Determine the root based on the given range.
    let (root, range_to_highlight) = {
        let source_file = sema.parse(file_id);
        match range_to_highlight {
            Some(range) => {
                let node = match source_file.syntax().covering_element(range) {
                    NodeOrToken::Node(it) => it,
                    NodeOrToken::Token(it) => it.parent(),
                };
                (node, range)
            }
            None => (source_file.syntax().clone(), source_file.syntax().text_range()),
        }
    };

    let mut bindings_shadow_count: FxHashMap<Name, u32> = FxHashMap::default();
    let mut res = Vec::new();

    let mut current_macro_call: Option<ast::MacroCall> = None;

    // Walk all nodes, keeping track of whether we are inside a macro or not.
    // If in macro, expand it first and highlight the expanded code.
    for event in root.preorder_with_tokens() {
        let event_range = match &event {
            WalkEvent::Enter(it) => it.text_range(),
            WalkEvent::Leave(it) => it.text_range(),
        };

        // Element outside of the viewport, no need to highlight
        if range_to_highlight.intersection(&event_range).is_none() {
            continue;
        }

        // Track "inside macro" state
        match event.clone().map(|it| it.into_node().and_then(ast::MacroCall::cast)) {
            WalkEvent::Enter(Some(mc)) => {
                current_macro_call = Some(mc.clone());
                if let Some(range) = macro_call_range(&mc) {
                    res.push(HighlightedRange {
                        range,
                        highlight: HighlightTag::Macro.into(),
                        binding_hash: None,
                    });
                }
                continue;
            }
            WalkEvent::Leave(Some(mc)) => {
                assert!(current_macro_call == Some(mc));
                current_macro_call = None;
                continue;
            }
            _ => (),
        }

        let element = match event {
            WalkEvent::Enter(it) => it,
            WalkEvent::Leave(_) => continue,
        };

        let range = element.text_range();

        let element_to_highlight = if current_macro_call.is_some() {
            // Inside a macro -- expand it first
            let token = match element.clone().into_token() {
                Some(it) if it.parent().kind() == TOKEN_TREE => it,
                _ => continue,
            };
            let token = sema.descend_into_macros(token.clone());
            let parent = token.parent();
            // We only care Name and Name_ref
            match (token.kind(), parent.kind()) {
                (IDENT, NAME) | (IDENT, NAME_REF) => parent.into(),
                _ => token.into(),
            }
        } else {
            element.clone()
        };

        if let Some(token) = element.as_token().cloned().and_then(ast::RawString::cast) {
            let expanded = element_to_highlight.as_token().unwrap().clone();
            if highlight_injection(&mut res, &sema, token, expanded).is_some() {
                continue;
            }
        }

        if let Some((highlight, binding_hash)) =
            highlight_element(&sema, &mut bindings_shadow_count, element_to_highlight)
        {
            res.push(HighlightedRange { range, highlight, binding_hash });
        }
    }

    res
}

fn macro_call_range(macro_call: &ast::MacroCall) -> Option<TextRange> {
    let path = macro_call.path()?;
    let name_ref = path.segment()?.name_ref()?;

    let range_start = name_ref.syntax().text_range().start();
    let mut range_end = name_ref.syntax().text_range().end();
    for sibling in path.syntax().siblings_with_tokens(Direction::Next) {
        match sibling.kind() {
            T![!] | IDENT => range_end = sibling.text_range().end(),
            _ => (),
        }
    }

    Some(TextRange::from_to(range_start, range_end))
}

fn highlight_element(
    sema: &Semantics<RootDatabase>,
    bindings_shadow_count: &mut FxHashMap<Name, u32>,
    element: SyntaxElement,
) -> Option<(Highlight, Option<u64>)> {
    let db = sema.db;
    let mut binding_hash = None;
    let highlight: Highlight = match element.kind() {
        FN_DEF => {
            bindings_shadow_count.clear();
            return None;
        }

        // Highlight definitions depending on the "type" of the definition.
        NAME => {
            let name = element.into_node().and_then(ast::Name::cast).unwrap();
            let name_kind = classify_name(sema, &name);

            if let Some(NameClass::NameDefinition(NameDefinition::Local(local))) = &name_kind {
                if let Some(name) = local.name(db) {
                    let shadow_count = bindings_shadow_count.entry(name.clone()).or_default();
                    *shadow_count += 1;
                    binding_hash = Some(calc_binding_hash(&name, *shadow_count))
                }
            };

            match name_kind {
                Some(NameClass::NameDefinition(def)) => {
                    highlight_name(db, def) | HighlightModifier::Definition
                }
                Some(NameClass::ConstReference(def)) => highlight_name(db, def),
                None => highlight_name_by_syntax(name) | HighlightModifier::Definition,
            }
        }

        // Highlight references like the definitions they resolve to
        NAME_REF if element.ancestors().any(|it| it.kind() == ATTR) => return None,
        NAME_REF => {
            let name_ref = element.into_node().and_then(ast::NameRef::cast).unwrap();
            let name_kind = classify_name_ref(sema, &name_ref)?;

            match name_kind {
                NameRefClass::NameDefinition(def) => {
                    if let NameDefinition::Local(local) = &def {
                        if let Some(name) = local.name(db) {
                            let shadow_count =
                                bindings_shadow_count.entry(name.clone()).or_default();
                            binding_hash = Some(calc_binding_hash(&name, *shadow_count))
                        }
                    };
                    highlight_name(db, def)
                }
                NameRefClass::FieldShorthand { .. } => HighlightTag::Field.into(),
            }
        }

        // Simple token-based highlighting
        COMMENT => HighlightTag::Comment.into(),
        STRING | RAW_STRING | RAW_BYTE_STRING | BYTE_STRING => HighlightTag::StringLiteral.into(),
        ATTR => HighlightTag::Attribute.into(),
        INT_NUMBER | FLOAT_NUMBER => HighlightTag::NumericLiteral.into(),
        BYTE => HighlightTag::ByteLiteral.into(),
        CHAR => HighlightTag::CharLiteral.into(),
        LIFETIME => {
            let h = Highlight::new(HighlightTag::Lifetime);
            match element.parent().map(|it| it.kind()) {
                Some(LIFETIME_PARAM) | Some(LABEL) => h | HighlightModifier::Definition,
                _ => h,
            }
        }

        k if k.is_keyword() => {
            let h = Highlight::new(HighlightTag::Keyword);
            match k {
                T![break]
                | T![continue]
                | T![else]
                | T![for]
                | T![if]
                | T![loop]
                | T![match]
                | T![return]
                | T![while] => h | HighlightModifier::Control,
                T![unsafe] => h | HighlightModifier::Unsafe,
                _ => h,
            }
        }

        _ => return None,
    };

    return Some((highlight, binding_hash));

    fn calc_binding_hash(name: &Name, shadow_count: u32) -> u64 {
        fn hash<T: std::hash::Hash + std::fmt::Debug>(x: T) -> u64 {
            use std::{collections::hash_map::DefaultHasher, hash::Hasher};

            let mut hasher = DefaultHasher::new();
            x.hash(&mut hasher);
            hasher.finish()
        }

        hash((name, shadow_count))
    }
}

fn highlight_name(db: &RootDatabase, def: NameDefinition) -> Highlight {
    match def {
        NameDefinition::Macro(_) => HighlightTag::Macro,
        NameDefinition::StructField(_) => HighlightTag::Field,
        NameDefinition::ModuleDef(def) => match def {
            hir::ModuleDef::Module(_) => HighlightTag::Module,
            hir::ModuleDef::Function(_) => HighlightTag::Function,
            hir::ModuleDef::Adt(hir::Adt::Struct(_)) => HighlightTag::Struct,
            hir::ModuleDef::Adt(hir::Adt::Enum(_)) => HighlightTag::Enum,
            hir::ModuleDef::Adt(hir::Adt::Union(_)) => HighlightTag::Union,
            hir::ModuleDef::EnumVariant(_) => HighlightTag::EnumVariant,
            hir::ModuleDef::Const(_) => HighlightTag::Constant,
            hir::ModuleDef::Static(_) => HighlightTag::Static,
            hir::ModuleDef::Trait(_) => HighlightTag::Trait,
            hir::ModuleDef::TypeAlias(_) => HighlightTag::TypeAlias,
            hir::ModuleDef::BuiltinType(_) => HighlightTag::BuiltinType,
        },
        NameDefinition::SelfType(_) => HighlightTag::SelfType,
        NameDefinition::TypeParam(_) => HighlightTag::TypeParam,
        // FIXME: distinguish between locals and parameters
        NameDefinition::Local(local) => {
            let mut h = Highlight::new(HighlightTag::Local);
            if local.is_mut(db) || local.ty(db).is_mutable_reference() {
                h |= HighlightModifier::Mutable;
            }
            return h;
        }
    }
    .into()
}

fn highlight_name_by_syntax(name: ast::Name) -> Highlight {
    let default = HighlightTag::Function.into();

    let parent = match name.syntax().parent() {
        Some(it) => it,
        _ => return default,
    };

    match parent.kind() {
        STRUCT_DEF => HighlightTag::Struct.into(),
        ENUM_DEF => HighlightTag::Enum.into(),
        UNION_DEF => HighlightTag::Union.into(),
        TRAIT_DEF => HighlightTag::Trait.into(),
        TYPE_ALIAS_DEF => HighlightTag::TypeAlias.into(),
        TYPE_PARAM => HighlightTag::TypeParam.into(),
        RECORD_FIELD_DEF => HighlightTag::Field.into(),
        _ => default,
    }
}

fn highlight_injection(
    acc: &mut Vec<HighlightedRange>,
    sema: &Semantics<RootDatabase>,
    literal: ast::RawString,
    expanded: SyntaxToken,
) -> Option<()> {
    let call_info = call_info_for_token(&sema, expanded)?;
    let idx = call_info.active_parameter?;
    let name = call_info.signature.parameter_names.get(idx)?;
    if name != "ra_fixture" {
        return None;
    }
    let value = literal.value()?;
    let (analysis, tmp_file_id) = Analysis::from_single_file(value);

    if let Some(range) = literal.open_quote_text_range() {
        acc.push(HighlightedRange {
            range,
            highlight: HighlightTag::StringLiteral.into(),
            binding_hash: None,
        })
    }

    for mut h in analysis.highlight(tmp_file_id).unwrap() {
        if let Some(r) = literal.map_range_up(h.range) {
            h.range = r;
            acc.push(h)
        }
    }

    if let Some(range) = literal.close_quote_text_range() {
        acc.push(HighlightedRange {
            range,
            highlight: HighlightTag::StringLiteral.into(),
            binding_hash: None,
        })
    }

    Some(())
}
<|endoftext|>


<commit_before>use base64::write::EncoderWriter;
use indexmap::IndexMap;
use prometheus::proto::{MetricFamily, MetricType};
use prometheus::{Encoder, Error};
use serde::ser::SerializeSeq;
use serde::{Serialize, Serializer as _};
use serde_json::Serializer;
use std::cell::Cell;
use std::io::Write;
use std::rc::Rc;

const CHUNKS_MAX_SIZE_BYTES: usize = 5000;

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

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

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

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

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

        Ok(())
    }

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

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

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

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

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

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

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

        result.push(writer.buffer);
    }

    Ok(result)
}

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

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

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

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

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

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

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

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

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

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

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

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

        Ok(())
    }

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

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

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

        Ok(())
    }

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

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

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

        Ok(())
    }

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

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

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

        Ok(())
    }

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

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

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

        Ok(())
    }

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

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

        Ok(())
    }

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

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

        Ok(())
    }
}
<commit_msg>fix vector.dev's histograms representation<commit_after>use base64::write::EncoderWriter;
use indexmap::IndexMap;
use prometheus::proto::{MetricFamily, MetricType};
use prometheus::{Encoder, Error};
use serde::ser::SerializeSeq;
use serde::{Serialize, Serializer as _};
use serde_json::Serializer;
use std::cell::Cell;
use std::io::Write;
use std::rc::Rc;

const CHUNKS_MAX_SIZE_BYTES: usize = 5000;

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

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

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

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

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

        Ok(())
    }

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

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

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

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

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

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

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

        result.push(writer.buffer);
    }

    Ok(result)
}

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

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

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

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

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

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

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

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

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

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

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

        Ok(())
    }

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

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

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

        Ok(())
    }

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

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

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

        Ok(())
    }

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

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

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

        Ok(())
    }

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

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

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

        Ok(())
    }

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

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

        Ok(())
    }

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

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

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


<commit_before>import std::cmp::fuzzy_eq;
import Num = num::num;

class Matrix4<T:copy fuzzy_eq Num> {
    let m11: T; let m12: T; let m13: T; let m14: T;
    let m21: T; let m22: T; let m23: T; let m24: T;
    let m31: T; let m32: T; let m33: T; let m34: T;
    let m41: T; let m42: T; let m43: T; let m44: T;

    new(m11: T, m12: T, m13: T, m14: T,
        m21: T, m22: T, m23: T, m24: T,
        m31: T, m32: T, m33: T, m34: T,
        m41: T, m42: T, m43: T, m44: T) {

        self.m11 = m11; self.m12 = m12; self.m13 = m13; self.m14 = m14;
        self.m21 = m21; self.m22 = m22; self.m23 = m23; self.m24 = m24;
        self.m31 = m31; self.m32 = m32; self.m33 = m33; self.m34 = m34;
        self.m41 = m41; self.m42 = m42; self.m43 = m43; self.m44 = m44;
    }

    fn fuzzy_eq(&&other: Matrix4<T>) -> bool {
        self.m11.fuzzy_eq(other.m11) && self.m12.fuzzy_eq(other.m12) &&
        self.m13.fuzzy_eq(other.m13) && self.m14.fuzzy_eq(other.m14) &&
        self.m21.fuzzy_eq(other.m21) && self.m22.fuzzy_eq(other.m22) &&
        self.m23.fuzzy_eq(other.m23) && self.m24.fuzzy_eq(other.m24) &&
        self.m31.fuzzy_eq(other.m31) && self.m32.fuzzy_eq(other.m32) &&
        self.m33.fuzzy_eq(other.m33) && self.m34.fuzzy_eq(other.m34) &&
        self.m41.fuzzy_eq(other.m41) && self.m42.fuzzy_eq(other.m42) &&
        self.m43.fuzzy_eq(other.m43) && self.m44.fuzzy_eq(other.m44)
    }

    fn mul_s(&&x: T) -> Matrix4<T> {
        ret Matrix4(self.m11.mul(x), self.m12.mul(x), self.m13.mul(x), self.m14.mul(x),
                    self.m21.mul(x), self.m22.mul(x), self.m23.mul(x), self.m24.mul(x),
                    self.m31.mul(x), self.m32.mul(x), self.m33.mul(x), self.m34.mul(x),
                    self.m41.mul(x), self.m42.mul(x), self.m43.mul(x), self.m44.mul(x));
    }

    fn scale(&&x: T, &&y: T, &&z: T) -> Matrix4<T> {
        ret Matrix4(self.m11.mul(x), self.m12,        self.m13,        self.m14,
                    self.m21,        self.m22.mul(y), self.m23,        self.m24,
                    self.m31,        self.m32,        self.m33.mul(z), self.m34,
                    self.m41,        self.m42,        self.m43,        self.m44);
    }

    pure fn to_array() -> ~[T] {
        ret ~[
            self.m11, self.m12, self.m13, self.m14,
            self.m21, self.m22, self.m23, self.m24,
            self.m31, self.m32, self.m33, self.m34,
            self.m41, self.m42, self.m43, self.m44
        ];
    }
}

fn ortho<T:copy fuzzy_eq Num>(+left: T, +right: T, +bottom: T, +top: T, +near: T, +far: T)
                           -> Matrix4<T> {

    let two = left.from_int(2);
    let one = left.from_int(1);
    let zero = left.from_int(0);
    let minus_two = left.from_int(-2);

    let tx = right.add(left).div(right.sub(left)).neg();
    let ty = top.add(bottom).div(top.sub(bottom)).neg();
    let tz = far.add(near).div(far.sub(near)).neg();

    ret Matrix4(two.div(right.sub(left)), zero, zero, zero,
                zero, two.div(top.sub(bottom)), zero, zero,
                zero, zero, minus_two.div(far.sub(near)), zero,
                tx, ty, tz, one);
}

fn identity<T:copy fuzzy_eq Num>(_0: T) -> Matrix4<T> {
    let _1 = _0.from_int(1);
    ret Matrix4(_1, _0, _0, _0,
                _0, _1, _0, _0,
                _0, _0, _1, _0,
                _0, _0, _0, _1);
}

#[test]
fn test_ortho() {
    let (left, right, bottom, top) = (0.0, 1.0, 0.1, 1.0);
    let (near, far) = (-1.0, 1.0);
    let result = ortho(left, right, bottom, top, near, far);
    let expected = Matrix4(2.0,  0.0,         0.0,  0.0,
                           0.0,  2.22222222,  0.0,  0.0,
                           0.0,  0.0,         -1.0, 0.0,
                           -1.0, -1.22222222, -0.0, 1.0);
    #debug("result=%? expected=%?", result, expected);
    assert result.fuzzy_eq(expected);
}

<commit_msg>Add matrix multiplication and scale<commit_after>import std::cmp::fuzzy_eq;
import Num = num::num;

pure fn Matrix4<T:copy fuzzy_eq Num>(m11: T, m12: T, m13: T, m14: T,
                                     m21: T, m22: T, m23: T, m24: T,
                                     m31: T, m32: T, m33: T, m34: T,
                                     m41: T, m42: T, m43: T, m44: T)
                                  -> Matrix4<T> {

    Matrix4 {
        m11: m11, m12: m12, m13: m13, m14: m14,
        m21: m21, m22: m22, m23: m23, m24: m24,
        m31: m31, m32: m32, m33: m33, m34: m34,
        m41: m41, m42: m42, m43: m43, m44: m44
    }
}

struct Matrix4<T:copy fuzzy_eq Num> {
    let m11: T; let m12: T; let m13: T; let m14: T;
    let m21: T; let m22: T; let m23: T; let m24: T;
    let m31: T; let m32: T; let m33: T; let m34: T;
    let m41: T; let m42: T; let m43: T; let m44: T;

    pure fn fuzzy_eq(&&other: Matrix4<T>) -> bool {
        self.m11.fuzzy_eq(other.m11) && self.m12.fuzzy_eq(other.m12) &&
        self.m13.fuzzy_eq(other.m13) && self.m14.fuzzy_eq(other.m14) &&
        self.m21.fuzzy_eq(other.m21) && self.m22.fuzzy_eq(other.m22) &&
        self.m23.fuzzy_eq(other.m23) && self.m24.fuzzy_eq(other.m24) &&
        self.m31.fuzzy_eq(other.m31) && self.m32.fuzzy_eq(other.m32) &&
        self.m33.fuzzy_eq(other.m33) && self.m34.fuzzy_eq(other.m34) &&
        self.m41.fuzzy_eq(other.m41) && self.m42.fuzzy_eq(other.m42) &&
        self.m43.fuzzy_eq(other.m43) && self.m44.fuzzy_eq(other.m44)
    }

    pure fn mul(&&m: Matrix4<T>) -> Matrix4<T> {
        ret Matrix4(m.m11.mul(self.m11).add(m.m12.mul(self.m21))
                                       .add(m.m13.mul(self.m31))
                                       .add(m.m14.mul(self.m41)),
                    m.m11.mul(self.m12).add(m.m12.mul(self.m22))
                                       .add(m.m13.mul(self.m32))
                                       .add(m.m14.mul(self.m42)),
                    m.m11.mul(self.m13).add(m.m12.mul(self.m23))
                                       .add(m.m13.mul(self.m33))
                                       .add(m.m14.mul(self.m43)),
                    m.m11.mul(self.m14).add(m.m12.mul(self.m24))
                                       .add(m.m13.mul(self.m34))
                                       .add(m.m14.mul(self.m44)),
                    m.m21.mul(self.m11).add(m.m22.mul(self.m21))
                                       .add(m.m23.mul(self.m31))
                                       .add(m.m24.mul(self.m41)),
                    m.m21.mul(self.m12).add(m.m22.mul(self.m22))
                                       .add(m.m23.mul(self.m32))
                                       .add(m.m24.mul(self.m42)),
                    m.m21.mul(self.m13).add(m.m22.mul(self.m23))
                                       .add(m.m23.mul(self.m33))
                                       .add(m.m24.mul(self.m43)),
                    m.m21.mul(self.m14).add(m.m22.mul(self.m24))
                                       .add(m.m23.mul(self.m34))
                                       .add(m.m24.mul(self.m44)),
                    m.m31.mul(self.m11).add(m.m32.mul(self.m21))
                                       .add(m.m33.mul(self.m31))
                                       .add(m.m34.mul(self.m41)),
                    m.m31.mul(self.m12).add(m.m32.mul(self.m22))
                                       .add(m.m33.mul(self.m32))
                                       .add(m.m34.mul(self.m42)),
                    m.m31.mul(self.m13).add(m.m32.mul(self.m23))
                                       .add(m.m33.mul(self.m33))
                                       .add(m.m34.mul(self.m43)),
                    m.m31.mul(self.m14).add(m.m32.mul(self.m24))
                                       .add(m.m33.mul(self.m34))
                                       .add(m.m34.mul(self.m44)),
                    m.m41.mul(self.m11).add(m.m42.mul(self.m21))
                                       .add(m.m43.mul(self.m31))
                                       .add(m.m44.mul(self.m41)),
                    m.m41.mul(self.m12).add(m.m42.mul(self.m22))
                                       .add(m.m43.mul(self.m32))
                                       .add(m.m44.mul(self.m42)),
                    m.m41.mul(self.m13).add(m.m42.mul(self.m23))
                                       .add(m.m43.mul(self.m33))
                                       .add(m.m44.mul(self.m43)),
                    m.m41.mul(self.m14).add(m.m42.mul(self.m24))
                                       .add(m.m43.mul(self.m34))
                                       .add(m.m44.mul(self.m44)));
    }

    pure fn mul_s(&&x: T) -> Matrix4<T> {
        ret Matrix4(self.m11.mul(x), self.m12.mul(x), self.m13.mul(x), self.m14.mul(x),
                    self.m21.mul(x), self.m22.mul(x), self.m23.mul(x), self.m24.mul(x),
                    self.m31.mul(x), self.m32.mul(x), self.m33.mul(x), self.m34.mul(x),
                    self.m41.mul(x), self.m42.mul(x), self.m43.mul(x), self.m44.mul(x));
    }

    pure fn scale(&&x: T, &&y: T, &&z: T) -> Matrix4<T> {
        ret Matrix4(self.m11.mul(x), self.m12,        self.m13,        self.m14,
                    self.m21,        self.m22.mul(y), self.m23,        self.m24,
                    self.m31,        self.m32,        self.m33.mul(z), self.m34,
                    self.m41,        self.m42,        self.m43,        self.m44);
    }

    pure fn to_array() -> ~[T] {
        ret ~[
            self.m11, self.m12, self.m13, self.m14,
            self.m21, self.m22, self.m23, self.m24,
            self.m31, self.m32, self.m33, self.m34,
            self.m41, self.m42, self.m43, self.m44
        ];
    }

    pure fn translate(&&x: T, &&y: T, &&z: T) -> Matrix4<T> {
        let _0 = x.from_int(0);
        let _1 = x.from_int(1);
        let matrix = Matrix4(_1, _0, _0, _0,
                             _0, _1, _0, _0,
                             _0, _0, _1, _0,
                              x,  y,  z, _1);

        ret self.mul(matrix);
    }
}

fn ortho<T:copy fuzzy_eq Num>(+left: T, +right: T, +bottom: T, +top: T, +near: T, +far: T)
                           -> Matrix4<T> {

    let two = left.from_int(2);
    let one = left.from_int(1);
    let zero = left.from_int(0);
    let minus_two = left.from_int(-2);

    let tx = right.add(left).div(right.sub(left)).neg();
    let ty = top.add(bottom).div(top.sub(bottom)).neg();
    let tz = far.add(near).div(far.sub(near)).neg();

    ret Matrix4(two.div(right.sub(left)), zero, zero, zero,
                zero, two.div(top.sub(bottom)), zero, zero,
                zero, zero, minus_two.div(far.sub(near)), zero,
                tx, ty, tz, one);
}

fn identity<T:copy fuzzy_eq Num>(_0: T) -> Matrix4<T> {
    let _1 = _0.from_int(1);
    ret Matrix4(_1, _0, _0, _0,
                _0, _1, _0, _0,
                _0, _0, _1, _0,
                _0, _0, _0, _1);
}

#[test]
fn test_ortho() {
    let (left, right, bottom, top) = (0.0, 1.0, 0.1, 1.0);
    let (near, far) = (-1.0, 1.0);
    let result = ortho(left, right, bottom, top, near, far);
    let expected = Matrix4(2.0,  0.0,         0.0,  0.0,
                           0.0,  2.22222222,  0.0,  0.0,
                           0.0,  0.0,         -1.0, 0.0,
                           -1.0, -1.22222222, -0.0, 1.0);
    #debug("result=%? expected=%?", result, expected);
    assert result.fuzzy_eq(expected);
}

<|endoftext|>


<commit_before>#![allow(non_camel_case_types)]

extern crate rand;
extern crate rayon;
extern crate time;

use rand::{Rng, SeedableRng, XorShiftRng};

trait Joiner {
    fn join<A,R_A,B,R_B>(oper_a: A, oper_b: B) -> (R_A, R_B)
        where A: FnOnce() -> R_A + Send, B: FnOnce() -> R_B + Send;
}

struct Parallel;
impl Joiner for Parallel {
    fn join<A,R_A,B,R_B>(oper_a: A, oper_b: B) -> (R_A, R_B)
        where A: FnOnce() -> R_A + Send, B: FnOnce() -> R_B + Send
    {
        rayon::join(oper_a, oper_b)
    }
}

struct Sequential;
impl Joiner for Sequential {
    fn join<A,R_A,B,R_B>(oper_a: A, oper_b: B) -> (R_A, R_B)
        where A: FnOnce() -> R_A + Send, B: FnOnce() -> R_B + Send
    {
        let a = oper_a();
        let b = oper_b();
        (a, b)
    }
}

fn quick_sort<J:Joiner, T:PartialOrd+Send>(v: &mut [T]) {
    if v.len() <= 1 {
        return;
    }

    let mid = partition(v);
    let (lo, hi) = v.split_at_mut(mid);
    J::join(|| quick_sort::<J,T>(lo),
            || quick_sort::<J,T>(hi));
}

fn partition<T:PartialOrd+Send>(v: &mut [T]) -> usize {
    let pivot = v.len() - 1;
    let mut i = 0;
    for j in 0..pivot {
        if v[j] <= v[pivot] {
            v.swap(i, j);
            i += 1;
        }
    }
    v.swap(i, pivot);
    i
}

fn main() {
    let mut rng = XorShiftRng::from_seed([0, 1, 2, 3]);
    let data: Vec<_> = (0..18*1024).map(|_| rng.next_u32()).collect();

    rayon::initialize();

    let par_time = {
        let mut data = data.clone();
        let par_start = time::precise_time_ns();
        quick_sort::<Parallel, u32>(&mut data);
        time::precise_time_ns() - par_start
    };

    let seq_time = {
        //let mut data = data.clone();
        let seq_start = time::precise_time_ns();
        //quick_sort::<Sequential, u32>(&mut data);
        time::precise_time_ns() - seq_start
    };

    println!("Parallel time   : {}", par_time);
    println!("Sequential time : {}", seq_time);
    println!("Parallel speedup: {}", (seq_time as f64) / (par_time as f64));
}
<commit_msg>improve demo<commit_after>#![allow(non_camel_case_types)]

extern crate rand;
extern crate rayon;
extern crate time;

use rand::{Rng, SeedableRng, XorShiftRng};
use std::env;
use std::str::FromStr;

trait Joiner {
    fn join<A,R_A,B,R_B>(oper_a: A, oper_b: B) -> (R_A, R_B)
        where A: FnOnce() -> R_A + Send, B: FnOnce() -> R_B + Send;
}

struct Parallel;
impl Joiner for Parallel {
    fn join<A,R_A,B,R_B>(oper_a: A, oper_b: B) -> (R_A, R_B)
        where A: FnOnce() -> R_A + Send, B: FnOnce() -> R_B + Send
    {
        rayon::join(oper_a, oper_b)
    }
}

struct Sequential;
impl Joiner for Sequential {
    fn join<A,R_A,B,R_B>(oper_a: A, oper_b: B) -> (R_A, R_B)
        where A: FnOnce() -> R_A + Send, B: FnOnce() -> R_B + Send
    {
        let a = oper_a();
        let b = oper_b();
        (a, b)
    }
}

fn quick_sort<J:Joiner, T:PartialOrd+Send>(v: &mut [T]) {
    if v.len() <= 1 {
        return;
    }

    let mid = partition(v);
    let (lo, hi) = v.split_at_mut(mid);
    J::join(|| quick_sort::<J,T>(lo),
            || quick_sort::<J,T>(hi));
}

fn partition<T:PartialOrd+Send>(v: &mut [T]) -> usize {
    let pivot = v.len() - 1;
    let mut i = 0;
    for j in 0..pivot {
        if v[j] <= v[pivot] {
            v.swap(i, j);
            i += 1;
        }
    }
    v.swap(i, pivot);
    i
}

fn main() {
    let mut kilo = 24;
    let mut run_par = true;
    let mut run_seq = true;

    let args = env::args();
    for arg in args.skip(1) {
        if arg == "--no-par" {
            run_par = false;
        } else if arg == "--no-seq" {
            run_seq = false;
        } else {
            match usize::from_str(&arg) {
                Ok(v) => kilo = v,
                Err(_) => {
                    println!("invalid argument `{}`, expected integer", arg);
                }
            }
        }
    }

    let mut rng = XorShiftRng::from_seed([0, 1, 2, 3]);
    let data: Vec<_> = (0..kilo*1024).map(|_| rng.next_u32()).collect();

    rayon::initialize();

    let par_time = if run_par {
        let mut data = data.clone();
        let par_start = time::precise_time_ns();
        quick_sort::<Parallel, u32>(&mut data);
        time::precise_time_ns() - par_start
    } else {
        0
    };

    let seq_time = if run_seq {
        let mut data = data.clone();
        let seq_start = time::precise_time_ns();
        quick_sort::<Sequential, u32>(&mut data);
        time::precise_time_ns() - seq_start
    } else {
        0
    };

    println!("Array length    : {}K 32-bit integers", kilo);
    if run_par { println!("Parallel time   : {}", par_time); }
    if run_seq { println!("Sequential time : {}", seq_time); }
    if run_par && run_seq {
        println!("Parallel speedup: {}", (seq_time as f64) / (par_time as f64));
    }
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate sysinfo;
extern crate test;

use sysinfo::get_current_pid;
use sysinfo::SystemExt;

#[bench]
fn bench_new(b: &mut test::Bencher) {
    b.iter(|| {
        sysinfo::System::new();
    });
}

#[bench]
fn bench_new_all(b: &mut test::Bencher) {
    b.iter(|| {
        sysinfo::System::new_all();
    });
}

#[bench]
fn bench_refresh_all(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    b.iter(move || {
        s.refresh_all();
    });
}

#[bench]
fn bench_refresh_system(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    s.refresh_system();
    b.iter(move || {
        s.refresh_system();
    });
}

#[bench]
fn bench_refresh_processes(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new();

    s.refresh_processes(); // to load the whole processes list a first time.
    b.iter(move || {
        s.refresh_processes();
    });
}

#[bench]
fn bench_refresh_process(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new();

    s.refresh_all();
    // to be sure it'll exist for at least as long as we run
    let pid = get_current_pid().expect("failed to get current pid");
    b.iter(move || {
        s.refresh_process(pid);
    });
}

#[bench]
fn bench_refresh_disks(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    b.iter(move || {
        s.refresh_disks();
    });
}

#[bench]
fn bench_refresh_disks_lists(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new();

    b.iter(move || {
        s.refresh_disks_list();
    });
}

#[bench]
fn bench_refresh_networks(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    b.iter(move || {
        s.refresh_networks();
    });
}

#[bench]
fn bench_refresh_networks_list(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new();

    b.iter(move || {
        s.refresh_networks_list();
    });
}

#[bench]
fn bench_refresh_memory(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new();

    b.iter(move || {
        s.refresh_memory();
    });
}

#[bench]
fn bench_refresh_cpu(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new();

    b.iter(move || {
        s.refresh_cpu();
    });
}

#[bench]
fn bench_refresh_components(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    b.iter(move || {
        s.refresh_components();
    });
}

#[bench]
fn bench_refresh_components_list(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    b.iter(move || {
        s.refresh_components_list();
    });
}

#[bench]
fn bench_refresh_users_list(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    b.iter(move || {
        s.refresh_users_list();
    });
}
<commit_msg>Fix bench typo<commit_after>#![feature(test)]

extern crate sysinfo;
extern crate test;

use sysinfo::get_current_pid;
use sysinfo::SystemExt;

#[bench]
fn bench_new(b: &mut test::Bencher) {
    b.iter(|| {
        sysinfo::System::new();
    });
}

#[bench]
fn bench_new_all(b: &mut test::Bencher) {
    b.iter(|| {
        sysinfo::System::new_all();
    });
}

#[bench]
fn bench_refresh_all(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    b.iter(move || {
        s.refresh_all();
    });
}

#[bench]
fn bench_refresh_system(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    s.refresh_system();
    b.iter(move || {
        s.refresh_system();
    });
}

#[bench]
fn bench_refresh_processes(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new();

    s.refresh_processes(); // to load the whole processes list a first time.
    b.iter(move || {
        s.refresh_processes();
    });
}

#[bench]
fn bench_refresh_process(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new();

    s.refresh_all();
    // to be sure it'll exist for at least as long as we run
    let pid = get_current_pid().expect("failed to get current pid");
    b.iter(move || {
        s.refresh_process(pid);
    });
}

#[bench]
fn bench_refresh_disks(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    b.iter(move || {
        s.refresh_disks();
    });
}

#[bench]
fn bench_refresh_disks_list(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new();

    b.iter(move || {
        s.refresh_disks_list();
    });
}

#[bench]
fn bench_refresh_networks(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    b.iter(move || {
        s.refresh_networks();
    });
}

#[bench]
fn bench_refresh_networks_list(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new();

    b.iter(move || {
        s.refresh_networks_list();
    });
}

#[bench]
fn bench_refresh_memory(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new();

    b.iter(move || {
        s.refresh_memory();
    });
}

#[bench]
fn bench_refresh_cpu(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new();

    b.iter(move || {
        s.refresh_cpu();
    });
}

#[bench]
fn bench_refresh_components(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    b.iter(move || {
        s.refresh_components();
    });
}

#[bench]
fn bench_refresh_components_list(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    b.iter(move || {
        s.refresh_components_list();
    });
}

#[bench]
fn bench_refresh_users_list(b: &mut test::Bencher) {
    let mut s = sysinfo::System::new_all();

    b.iter(move || {
        s.refresh_users_list();
    });
}
<|endoftext|>


<commit_before>extern crate time;

use driver::Driver;
use telemetry::Telemetry;
use telemetry_message::TelemetryMessage;

type MilliSeconds = u64;

#[derive(PartialEq)]
enum ControlState {
    WAITING_FOR_START,
    RUNNING,
    COLLISION_RECOVERY,
}

pub struct Control<'a, 'b> {
    state: ControlState,
    run: bool,
    telemetry: &'a (Telemetry + 'a),
    driver: &'b mut (Driver + 'b),
    collision_time_ms: MilliSeconds,
}

impl <'a, 'b> Control <'a, 'b> {
    pub fn new(telemetry: &'a Telemetry, driver: &'b mut Driver) -> Control <'a, 'b> {
        Control {
            state: ControlState::WAITING_FOR_START,
            run: false,
            telemetry: telemetry,
            driver: driver,
            collision_time_ms: 0,
        }
    }

    /**
     * Decides what to do and commands the vehicle for this time slice.
     */
    fn run_incremental(&mut self) {
        // Halting the car supercedes all other states
        if !self.run {
            self.state = ControlState::WAITING_FOR_START;
        } else if (self.telemetry.is_stopped()) {
            // We want to drive for at least one second between collisions
            self.collision_time_ms = time::now().to_milliseconds();
            self.state = ControlState::COLLISION_RECOVERY;
        }

        match self.state {
            ControlState::WAITING_FOR_START => self.waiting_for_start(),
            ControlState::RUNNING => self.running(),
            ControlState::COLLISION_RECOVERY => {
                let now_ms = time::now().to_milliseconds();
                self.collision_recovery(now_ms);
            }
        }
    }

    pub fn handle_message(&self, message: &str) {
        // TODO
    }

    fn waiting_for_start(&mut self) {
        if !self.run {
            return;
        }
        self.state = ControlState::RUNNING;
    }

    fn running(&self) {
        // TODO: Drive around
    }

    fn collision_recovery(&mut self, now_ms: MilliSeconds) {
        // Stop the motor for .5 seconds, then back up for 1 second, then pause
        // for .5 seconds
        let stop_ms = 500 as MilliSeconds;
        let back_up_ms = 1000 as MilliSeconds;
        let pause_ms = 500 as MilliSeconds;
        if now_ms < self.collision_time_ms + stop_ms {
            self.driver.drive(0.0f32, 0.0f32);
        } else if now_ms < self.collision_time_ms + stop_ms + back_up_ms {
            // TODO Choose a random direction
            self.driver.drive(-0.5f32, -0.5f32);
        } else if now_ms < self.collision_time_ms + stop_ms + back_up_ms + pause_ms {
            self.driver.drive(0.0f32, 0.0f32);
        } else {
            self.state = ControlState::RUNNING;
        }
    }
}

trait ToMilliseconds {
    fn to_milliseconds(&self) -> MilliSeconds;
}

impl ToMilliseconds for time::Tm {
    fn to_milliseconds(&self) -> MilliSeconds {
        let timespec = self.to_timespec();
        (timespec.sec * 1000 + timespec.nsec as i64 / 1000000) as MilliSeconds
    }
}

#[cfg(test)]
mod tests {
    use driver::Driver;
    use driver::Percentage;
    use telemetry::Telemetry;
    use telemetry_message::CompassMessage;
    use telemetry_message::GpsMessage;
    use telemetry_message::TelemetryMessage;
    use super::Control;
    use super::ControlState;
    use super::MilliSeconds;

    struct TestDriver {
        pub throttle: Percentage,
        pub steering: Percentage,
    }

    impl TestDriver {
        fn new() -> TestDriver {
            TestDriver {
                throttle: 0.0f32,
                steering: 0.0f32,
            }
        }
    }

    impl Driver for TestDriver {
        fn drive(&mut self, throttle: Percentage, steering: Percentage) {
            self.throttle = throttle;
            self.steering = steering;
        }

        fn get_throttle(&self) -> f32 {
            self.throttle
        }

        fn get_steering(&self) -> f32 {
            self.steering
        }
    }

    struct TestTelemetry {
        pub gps: GpsMessage,
        pub compass: CompassMessage,
        pub stopped: bool,
    }

    impl TestTelemetry {
        fn new() -> TestTelemetry {
            TestTelemetry {
                gps: GpsMessage {
                    x_m: 0f32,
                    y_m: 0f32,
                    heading_d: 0f32,
                    speed_m_s: 0f32,
                    ms_since_midnight: 0i32,
                },
                compass: CompassMessage {
                    heading_d: 0f32,
                    magnetometer: (0f32, 0f32, 0f32),
                    ms_since_midnight: 0i32,
                },
                stopped: false,
            }
        }
    }

    impl Telemetry for TestTelemetry {
        fn get_raw_gps(&self) -> &GpsMessage { &self.gps }
        fn get_raw_compass(&self) -> &CompassMessage { &self.compass }
        fn get_data(&self) -> &GpsMessage { &self.gps }
        fn process_drive_command(&mut self, throttle:f32, steering:f32) { }
        fn handle_message(&mut self, message: &TelemetryMessage) { }
        fn is_stopped(&self) -> bool { self.stopped }
    }

    #[test]
    fn test_collision_recovery() {
        let telemetry = TestTelemetry::new();
        let mut driver = TestDriver::new();
        let control = Control::new(&telemetry, &mut driver);

        // TODO: Figure out how to do these assertions after the driver has
        // been lent to control
        assert!(false);
        /*
        control.collision_recovery(0 as MilliSeconds);
        assert!(control.state == ControlState::COLLISION_RECOVERY);
        assert!(driver.get_throttle() == 0.0 as Percentage);
        assert!(driver.get_steering() == 0.0 as Percentage);

        control.collision_recovery(400 as MilliSeconds);
        assert!(control.state == ControlState::COLLISION_RECOVERY);
        assert!(driver.get_throttle() == 0.0 as Percentage);
        assert!(driver.get_steering() == 0.0 as Percentage);

        control.collision_recovery(600 as MilliSeconds);
        assert!(control.state == ControlState::COLLISION_RECOVERY);
        assert!(driver.get_throttle() < 0.0 as Percentage);
        assert!(driver.get_steering() != 0.0 as Percentage);

        control.collision_recovery(1400 as MilliSeconds);
        assert!(control.state == ControlState::COLLISION_RECOVERY);
        assert!(driver.get_throttle() < 0.0 as Percentage);
        assert!(driver.get_steering() != 0.0 as Percentage);

        control.collision_recovery(1600 as MilliSeconds);
        assert!(control.state == ControlState::COLLISION_RECOVERY);
        assert!(driver.get_throttle() == 0.0 as Percentage);
        assert!(driver.get_steering() == 0.0 as Percentage);

        control.collision_recovery(1900 as MilliSeconds);
        assert!(control.state == ControlState::COLLISION_RECOVERY);
        assert!(driver.get_throttle() == 0.0 as Percentage);
        assert!(driver.get_steering() == 0.0 as Percentage);

        control.collision_recovery(2100 as MilliSeconds);
        assert!(control.state != ControlState::COLLISION_RECOVERY);
        */
    }
}
<commit_msg>Fix some warnings<commit_after>extern crate time;

use driver::Driver;
use telemetry::Telemetry;
use telemetry_message::TelemetryMessage;

type MilliSeconds = u64;

#[derive(PartialEq)]
enum ControlState {
    WaitingForStart,
    Running,
    CollisionRecovery,
}

pub struct Control<'a, 'b> {
    state: ControlState,
    run: bool,
    telemetry: &'a (Telemetry + 'a),
    driver: &'b mut (Driver + 'b),
    collision_time_ms: MilliSeconds,
}

impl <'a, 'b> Control <'a, 'b> {
    pub fn new(telemetry: &'a Telemetry, driver: &'b mut Driver) -> Control <'a, 'b> {
        Control {
            state: ControlState::WaitingForStart,
            run: false,
            telemetry: telemetry,
            driver: driver,
            collision_time_ms: 0,
        }
    }

    /**
     * Decides what to do and commands the vehicle for this time slice.
     */
    fn run_incremental(&mut self) {
        // Halting the car supercedes all other states
        if !self.run {
            self.state = ControlState::WaitingForStart;
        } else if (self.telemetry.is_stopped()) {
            // We want to drive for at least one second between collisions
            self.collision_time_ms = time::now().to_milliseconds();
            self.state = ControlState::CollisionRecovery;
        }

        match self.state {
            ControlState::WaitingForStart => self.waiting_for_start(),
            ControlState::Running => self.running(),
            ControlState::CollisionRecovery => {
                let now_ms = time::now().to_milliseconds();
                self.collision_recovery(now_ms);
            }
        }
    }

    pub fn handle_message(&self, message: &str) {
        // TODO
    }

    fn waiting_for_start(&mut self) {
        if !self.run {
            return;
        }
        self.state = ControlState::Running;
    }

    fn running(&self) {
        // TODO: Drive around
    }

    fn collision_recovery(&mut self, now_ms: MilliSeconds) {
        // Stop the motor for .5 seconds, then back up for 1 second, then pause
        // for .5 seconds
        let stop_ms = 500 as MilliSeconds;
        let back_up_ms = 1000 as MilliSeconds;
        let pause_ms = 500 as MilliSeconds;
        if now_ms < self.collision_time_ms + stop_ms {
            self.driver.drive(0.0f32, 0.0f32);
        } else if now_ms < self.collision_time_ms + stop_ms + back_up_ms {
            // TODO Choose a random direction
            self.driver.drive(-0.5f32, -0.5f32);
        } else if now_ms < self.collision_time_ms + stop_ms + back_up_ms + pause_ms {
            self.driver.drive(0.0f32, 0.0f32);
        } else {
            self.state = ControlState::Running;
        }
    }
}

trait ToMilliseconds {
    fn to_milliseconds(&self) -> MilliSeconds;
}

impl ToMilliseconds for time::Tm {
    fn to_milliseconds(&self) -> MilliSeconds {
        let timespec = self.to_timespec();
        (timespec.sec * 1000 + timespec.nsec as i64 / 1000000) as MilliSeconds
    }
}

#[cfg(test)]
mod tests {
    use driver::Driver;
    use driver::Percentage;
    use telemetry::Telemetry;
    use telemetry_message::CompassMessage;
    use telemetry_message::GpsMessage;
    use telemetry_message::TelemetryMessage;
    use super::Control;
    use super::ControlState;
    use super::MilliSeconds;

    struct TestDriver {
        pub throttle: Percentage,
        pub steering: Percentage,
    }

    impl TestDriver {
        fn new() -> TestDriver {
            TestDriver {
                throttle: 0.0f32,
                steering: 0.0f32,
            }
        }
    }

    impl Driver for TestDriver {
        fn drive(&mut self, throttle: Percentage, steering: Percentage) {
            self.throttle = throttle;
            self.steering = steering;
        }

        fn get_throttle(&self) -> f32 {
            self.throttle
        }

        fn get_steering(&self) -> f32 {
            self.steering
        }
    }

    struct TestTelemetry {
        pub gps: GpsMessage,
        pub compass: CompassMessage,
        pub stopped: bool,
    }

    impl TestTelemetry {
        fn new() -> TestTelemetry {
            TestTelemetry {
                gps: GpsMessage {
                    x_m: 0f32,
                    y_m: 0f32,
                    heading_d: 0f32,
                    speed_m_s: 0f32,
                    ms_since_midnight: 0i32,
                },
                compass: CompassMessage {
                    heading_d: 0f32,
                    magnetometer: (0f32, 0f32, 0f32),
                    ms_since_midnight: 0i32,
                },
                stopped: false,
            }
        }
    }

    impl Telemetry for TestTelemetry {
        fn get_raw_gps(&self) -> &GpsMessage { &self.gps }
        fn get_raw_compass(&self) -> &CompassMessage { &self.compass }
        fn get_data(&self) -> &GpsMessage { &self.gps }
        fn process_drive_command(&mut self, throttle:f32, steering:f32) { }
        fn handle_message(&mut self, message: &TelemetryMessage) { }
        fn is_stopped(&self) -> bool { self.stopped }
    }

    #[test]
    fn test_collision_recovery() {
        let telemetry = TestTelemetry::new();
        let mut driver = TestDriver::new();
        let control = Control::new(&telemetry, &mut driver);

        // TODO: Figure out how to do these assertions after the driver has
        // been lent to control
        assert!(false);
        /*
        control.collision_recovery(0 as MilliSeconds);
        assert!(control.state == ControlState::CollisionRecovery);
        assert!(driver.get_throttle() == 0.0 as Percentage);
        assert!(driver.get_steering() == 0.0 as Percentage);

        control.collision_recovery(400 as MilliSeconds);
        assert!(control.state == ControlState::CollisionRecovery);
        assert!(driver.get_throttle() == 0.0 as Percentage);
        assert!(driver.get_steering() == 0.0 as Percentage);

        control.collision_recovery(600 as MilliSeconds);
        assert!(control.state == ControlState::CollisionRecovery);
        assert!(driver.get_throttle() < 0.0 as Percentage);
        assert!(driver.get_steering() != 0.0 as Percentage);

        control.collision_recovery(1400 as MilliSeconds);
        assert!(control.state == ControlState::CollisionRecovery);
        assert!(driver.get_throttle() < 0.0 as Percentage);
        assert!(driver.get_steering() != 0.0 as Percentage);

        control.collision_recovery(1600 as MilliSeconds);
        assert!(control.state == ControlState::CollisionRecovery);
        assert!(driver.get_throttle() == 0.0 as Percentage);
        assert!(driver.get_steering() == 0.0 as Percentage);

        control.collision_recovery(1900 as MilliSeconds);
        assert!(control.state == ControlState::CollisionRecovery);
        assert!(driver.get_throttle() == 0.0 as Percentage);
        assert!(driver.get_steering() == 0.0 as Percentage);

        control.collision_recovery(2100 as MilliSeconds);
        assert!(control.state != ControlState::CollisionRecovery);
        */
    }
}
<|endoftext|>


<commit_before>use kernel::eyre::Result;
use kernel_micro::{include_file, MicroKernel};

pub async fn new() -> Result<MicroKernel> {
    MicroKernel::new(
        "javascript",
        ("node", "*", &["{{script}}"]),
        include_file!("node-kernel.js"),
        &[include_file!("node-codec.js")],
        "{{name}} = decodeValue('{{json}}')",
        "{{name}}",
    )
    .await
}

#[cfg(test)]
mod tests {
    use super::*;
    use kernel::{stencila_schema::Node, KernelTrait};
    use test_utils::{assert_json_eq, serde_json::json};

    /// Tests of basic functionality
    /// This test is replicated in all the microkernels.
    /// Other test should be written for language specific quirks and regressions.
    #[tokio::test]
    async fn basics() -> Result<()> {
        let mut kernel = new().await?;
        kernel.start().await?;

        // Assign a variable and output it
        let (outputs, messages) = kernel.exec("const a = 2\na").await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [2]);

        // Print the variable twice and then output it
        let (outputs, messages) = kernel.exec("console.log(a)\nconsole.log(a)\na").await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [2, 2, 2]);

        // Syntax error
        let (outputs, messages) = kernel.exec("bad ^ # syntax").await?;
        assert_json_eq!(messages[0].error_type, "SyntaxError");
        assert_json_eq!(messages[0].error_message, "Invalid or unexpected token");
        assert_json_eq!(outputs, json!([]));

        // Runtime error
        let (outputs, messages) = kernel.exec("foo").await?;
        assert_json_eq!(messages[0].error_type, "ReferenceError");
        assert_json_eq!(messages[0].error_message, "foo is not defined");
        assert_json_eq!(outputs, json!([]));

        // Set and get another variable
        kernel.set("b", Node::Integer(3)).await?;
        let b = kernel.get("b").await?;
        assert_json_eq!(b, 3);

        // Use both variables
        let (outputs, messages) = kernel.exec("a*b").await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [6]);

        Ok(())
    }

    // Test that `console.log` arguments are treated as separate outputs
    #[tokio::test]
    async fn console_log() -> Result<()> {
        let mut kernel = new().await?;
        kernel.start().await?;

        let (outputs, messages) = kernel.exec("console.log(1)").await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [1]);

        let (outputs, messages) = kernel.exec("console.log(1, 2, 3, 4)").await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [1, 2, 3, 4]);

        let (outputs, messages) = kernel.exec("console.log([1, 2, 3], 4, 'str')").await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, json!([[1, 2, 3], 4, "str"]));

        Ok(())
    }

    // Test that `console.debug`, `console.warn` etc are treated as separate messages
    #[tokio::test]
    async fn console_messages() -> Result<()> {
        let mut kernel = new().await?;
        kernel.start().await?;

        let (outputs, messages) = kernel
            .exec(
                r#"
console.log(1)
console.debug("Debug message")
console.log(2)
console.info("Info message")
console.log(3)
console.warn("Warn message")
console.log(4)
console.error("Error message")
5
"#,
            )
            .await?;

        assert_json_eq!(
            messages,
            json!([{
                "type": "CodeError",
                "errorType": "CodeDebug",
                "errorMessage": "Debug message",
            }, {
                "type": "CodeError",
                "errorType": "CodeInfo",
                "errorMessage": "Info message",
            }, {
                "type": "CodeError",
                "errorType": "CodeWarn",
                "errorMessage": "Warn message",
            } , {
                "type": "CodeError",
                "errorType": "CodeError",
                "errorMessage": "Error message",
            }])
        );
        assert_json_eq!(outputs, json!([1, 2, 3, 4, 5]));

        Ok(())
    }
}
<commit_msg>test(Node kernel): Check for stack trace<commit_after>use kernel::eyre::Result;
use kernel_micro::{include_file, MicroKernel};

pub async fn new() -> Result<MicroKernel> {
    MicroKernel::new(
        "javascript",
        ("node", "*", &["{{script}}"]),
        include_file!("node-kernel.js"),
        &[include_file!("node-codec.js")],
        "{{name}} = decodeValue('{{json}}')",
        "{{name}}",
    )
    .await
}

#[cfg(test)]
mod tests {
    use super::*;
    use kernel::{stencila_schema::Node, KernelTrait};
    use test_utils::{assert_json_eq, serde_json::json};

    /// Tests of basic functionality
    /// This test is replicated in all the microkernels.
    /// Other test should be written for language specific quirks and regressions.
    #[tokio::test]
    async fn basics() -> Result<()> {
        let mut kernel = new().await?;
        kernel.start().await?;

        // Assign a variable and output it
        let (outputs, messages) = kernel.exec("const a = 2\na").await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [2]);

        // Print the variable twice and then output it
        let (outputs, messages) = kernel.exec("console.log(a)\nconsole.log(a)\na").await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [2, 2, 2]);

        // Syntax error
        let (outputs, messages) = kernel.exec("bad ^ # syntax").await?;
        assert_json_eq!(messages[0].error_type, "SyntaxError");
        assert_json_eq!(messages[0].error_message, "Invalid or unexpected token");
        assert!(messages[0].stack_trace.is_some());
        assert_json_eq!(outputs, json!([]));

        // Runtime error
        let (outputs, messages) = kernel.exec("foo").await?;
        assert_json_eq!(messages[0].error_type, "ReferenceError");
        assert_json_eq!(messages[0].error_message, "foo is not defined");
        assert!(messages[0].stack_trace.is_some());
        assert_json_eq!(outputs, json!([]));

        // Set and get another variable
        kernel.set("b", Node::Integer(3)).await?;
        let b = kernel.get("b").await?;
        assert_json_eq!(b, 3);

        // Use both variables
        let (outputs, messages) = kernel.exec("a*b").await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [6]);

        Ok(())
    }

    // Test that `console.log` arguments are treated as separate outputs
    #[tokio::test]
    async fn console_log() -> Result<()> {
        let mut kernel = new().await?;
        kernel.start().await?;

        let (outputs, messages) = kernel.exec("console.log(1)").await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [1]);

        let (outputs, messages) = kernel.exec("console.log(1, 2, 3, 4)").await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [1, 2, 3, 4]);

        let (outputs, messages) = kernel.exec("console.log([1, 2, 3], 4, 'str')").await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, json!([[1, 2, 3], 4, "str"]));

        Ok(())
    }

    // Test that `console.debug`, `console.warn` etc are treated as separate messages
    #[tokio::test]
    async fn console_messages() -> Result<()> {
        let mut kernel = new().await?;
        kernel.start().await?;

        let (outputs, messages) = kernel
            .exec(
                r#"
console.log(1)
console.debug("Debug message")
console.log(2)
console.info("Info message")
console.log(3)
console.warn("Warn message")
console.log(4)
console.error("Error message")
5
"#,
            )
            .await?;

        assert_json_eq!(
            messages,
            json!([{
                "type": "CodeError",
                "errorType": "CodeDebug",
                "errorMessage": "Debug message",
            }, {
                "type": "CodeError",
                "errorType": "CodeInfo",
                "errorMessage": "Info message",
            }, {
                "type": "CodeError",
                "errorType": "CodeWarn",
                "errorMessage": "Warn message",
            } , {
                "type": "CodeError",
                "errorType": "CodeError",
                "errorMessage": "Error message",
            }])
        );
        assert_json_eq!(outputs, json!([1, 2, 3, 4, 5]));

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