<commit_before>//! Provides a way to attach fixes to the diagnostics.
//! The same module also has all curret custom fixes for the diagnostics implemented.
use hir::{
    db::AstDatabase,
    diagnostics::{
        Diagnostic, IncorrectCase, MissingFields, MissingOkOrSomeInTailExpr, NoSuchField,
        RemoveThisSemicolon, ReplaceFilterMapNextWithFindMap, UnresolvedModule,
    },
    HasSource, HirDisplay, InFile, Semantics, VariantDef,
};
use ide_db::{
    base_db::{AnchoredPathBuf, FileId},
    source_change::{FileSystemEdit, SourceChange},
    RootDatabase,
};
use syntax::{
    algo,
    ast::{self, edit::IndentLevel, make, ArgListOwner},
    AstNode, TextRange,
};
use text_edit::TextEdit;

use crate::{diagnostics::Fix, references::rename::rename_with_semantics, FilePosition};

/// A [Diagnostic] that potentially has a fix available.
///
/// [Diagnostic]: hir::diagnostics::Diagnostic
pub(crate) trait DiagnosticWithFix: Diagnostic {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix>;
}

impl DiagnosticWithFix for UnresolvedModule {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        let root = sema.db.parse_or_expand(self.file)?;
        let unresolved_module = self.decl.to_node(&root);
        Some(Fix::new(
            "Create module",
            FileSystemEdit::CreateFile {
                dst: AnchoredPathBuf {
                    anchor: self.file.original_file(sema.db),
                    path: self.candidate.clone(),
                },
                initial_contents: "".to_string(),
            }
            .into(),
            unresolved_module.syntax().text_range(),
        ))
    }
}

impl DiagnosticWithFix for NoSuchField {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        let root = sema.db.parse_or_expand(self.file)?;
        missing_record_expr_field_fix(
            &sema,
            self.file.original_file(sema.db),
            &self.field.to_node(&root),
        )
    }
}

impl DiagnosticWithFix for MissingFields {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        // Note that although we could add a diagnostics to
        // fill the missing tuple field, e.g :
        // `struct A(usize);`
        // `let a = A { 0: () }`
        // but it is uncommon usage and it should not be encouraged.
        if self.missed_fields.iter().any(|it| it.as_tuple_index().is_some()) {
            return None;
        }

        let root = sema.db.parse_or_expand(self.file)?;
        let field_list_parent = self.field_list_parent.to_node(&root);
        let old_field_list = field_list_parent.record_expr_field_list()?;
        let mut new_field_list = old_field_list.clone();
        for f in self.missed_fields.iter() {
            let field =
                make::record_expr_field(make::name_ref(&f.to_string()), Some(make::expr_unit()));
            new_field_list = new_field_list.append_field(&field);
        }

        let edit = {
            let mut builder = TextEdit::builder();
            algo::diff(&old_field_list.syntax(), &new_field_list.syntax())
                .into_text_edit(&mut builder);
            builder.finish()
        };
        Some(Fix::new(
            "Fill struct fields",
            SourceChange::from_text_edit(self.file.original_file(sema.db), edit),
            sema.original_range(&field_list_parent.syntax()).range,
        ))
    }
}

impl DiagnosticWithFix for MissingOkOrSomeInTailExpr {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        let root = sema.db.parse_or_expand(self.file)?;
        let tail_expr = self.expr.to_node(&root);
        let tail_expr_range = tail_expr.syntax().text_range();
        let replacement = format!("{}({})", self.required, tail_expr.syntax());
        let edit = TextEdit::replace(tail_expr_range, replacement);
        let source_change = SourceChange::from_text_edit(self.file.original_file(sema.db), edit);
        let name = if self.required == "Ok" { "Wrap with Ok" } else { "Wrap with Some" };
        Some(Fix::new(name, source_change, tail_expr_range))
    }
}

impl DiagnosticWithFix for RemoveThisSemicolon {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        let root = sema.db.parse_or_expand(self.file)?;

        let semicolon = self
            .expr
            .to_node(&root)
            .syntax()
            .parent()
            .and_then(ast::ExprStmt::cast)
            .and_then(|expr| expr.semicolon_token())?
            .text_range();

        let edit = TextEdit::delete(semicolon);
        let source_change = SourceChange::from_text_edit(self.file.original_file(sema.db), edit);

        Some(Fix::new("Remove this semicolon", source_change, semicolon))
    }
}

impl DiagnosticWithFix for IncorrectCase {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        let root = sema.db.parse_or_expand(self.file)?;
        let name_node = self.ident.to_node(&root);

        let name_node = InFile::new(self.file, name_node.syntax());
        let frange = name_node.original_file_range(sema.db);
        let file_position = FilePosition { file_id: frange.file_id, offset: frange.range.start() };

        let rename_changes =
            rename_with_semantics(sema, file_position, &self.suggested_text).ok()?;

        let label = format!("Rename to {}", self.suggested_text);
        Some(Fix::new(&label, rename_changes, frange.range))
    }
}

impl DiagnosticWithFix for ReplaceFilterMapNextWithFindMap {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        let root = sema.db.parse_or_expand(self.file)?;
        let next_expr = self.next_expr.to_node(&root);
        let next_call = ast::MethodCallExpr::cast(next_expr.syntax().clone())?;

        let filter_map_call = ast::MethodCallExpr::cast(next_call.receiver()?.syntax().clone())?;
        let filter_map_name_range = filter_map_call.name_ref()?.ident_token()?.text_range();
        let filter_map_args = filter_map_call.arg_list()?;

        let range_to_replace =
            TextRange::new(filter_map_name_range.start(), next_expr.syntax().text_range().end());
        let replacement = format!("find_map{}", filter_map_args.syntax().text());
        let trigger_range = next_expr.syntax().text_range();

        let edit = TextEdit::replace(range_to_replace, replacement);

        let source_change =
            SourceFileEdit { file_id: self.file.original_file(sema.db), edit }.into();

        Some(Fix::new(
            "Replace filter_map(..).next() with find_map()",
            source_change,
            trigger_range,
        ))
    }
}

fn missing_record_expr_field_fix(
    sema: &Semantics<RootDatabase>,
    usage_file_id: FileId,
    record_expr_field: &ast::RecordExprField,
) -> Option<Fix> {
    let record_lit = ast::RecordExpr::cast(record_expr_field.syntax().parent()?.parent()?)?;
    let def_id = sema.resolve_variant(record_lit)?;
    let module;
    let def_file_id;
    let record_fields = match VariantDef::from(def_id) {
        VariantDef::Struct(s) => {
            module = s.module(sema.db);
            let source = s.source(sema.db)?;
            def_file_id = source.file_id;
            let fields = source.value.field_list()?;
            record_field_list(fields)?
        }
        VariantDef::Union(u) => {
            module = u.module(sema.db);
            let source = u.source(sema.db)?;
            def_file_id = source.file_id;
            source.value.record_field_list()?
        }
        VariantDef::Variant(e) => {
            module = e.module(sema.db);
            let source = e.source(sema.db)?;
            def_file_id = source.file_id;
            let fields = source.value.field_list()?;
            record_field_list(fields)?
        }
    };
    let def_file_id = def_file_id.original_file(sema.db);

    let new_field_type = sema.type_of_expr(&record_expr_field.expr()?)?;
    if new_field_type.is_unknown() {
        return None;
    }
    let new_field = make::record_field(
        None,
        make::name(record_expr_field.field_name()?.text()),
        make::ty(&new_field_type.display_source_code(sema.db, module.into()).ok()?),
    );

    let last_field = record_fields.fields().last()?;
    let last_field_syntax = last_field.syntax();
    let indent = IndentLevel::from_node(last_field_syntax);

    let mut new_field = new_field.to_string();
    if usage_file_id != def_file_id {
        new_field = format!("pub(crate) {}", new_field);
    }
    new_field = format!("\n{}{}", indent, new_field);

    let needs_comma = !last_field_syntax.to_string().ends_with(',');
    if needs_comma {
        new_field = format!(",{}", new_field);
    }

    let source_change = SourceChange::from_text_edit(
        def_file_id,
        TextEdit::insert(last_field_syntax.text_range().end(), new_field),
    );
    return Some(Fix::new("Create field", source_change, record_expr_field.syntax().text_range()));

    fn record_field_list(field_def_list: ast::FieldList) -> Option<ast::RecordFieldList> {
        match field_def_list {
            ast::FieldList::RecordFieldList(it) => Some(it),
            ast::FieldList::TupleFieldList(_) => None,
        }
    }
}
<commit_msg>Remove use of SourceFileEdit<commit_after>//! Provides a way to attach fixes to the diagnostics.
//! The same module also has all curret custom fixes for the diagnostics implemented.
use hir::{
    db::AstDatabase,
    diagnostics::{
        Diagnostic, IncorrectCase, MissingFields, MissingOkOrSomeInTailExpr, NoSuchField,
        RemoveThisSemicolon, ReplaceFilterMapNextWithFindMap, UnresolvedModule,
    },
    HasSource, HirDisplay, InFile, Semantics, VariantDef,
};
use ide_db::{
    base_db::{AnchoredPathBuf, FileId},
    source_change::{FileSystemEdit, SourceChange},
    RootDatabase,
};
use syntax::{
    algo,
    ast::{self, edit::IndentLevel, make, ArgListOwner},
    AstNode, TextRange,
};
use text_edit::TextEdit;

use crate::{diagnostics::Fix, references::rename::rename_with_semantics, FilePosition};

/// A [Diagnostic] that potentially has a fix available.
///
/// [Diagnostic]: hir::diagnostics::Diagnostic
pub(crate) trait DiagnosticWithFix: Diagnostic {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix>;
}

impl DiagnosticWithFix for UnresolvedModule {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        let root = sema.db.parse_or_expand(self.file)?;
        let unresolved_module = self.decl.to_node(&root);
        Some(Fix::new(
            "Create module",
            FileSystemEdit::CreateFile {
                dst: AnchoredPathBuf {
                    anchor: self.file.original_file(sema.db),
                    path: self.candidate.clone(),
                },
                initial_contents: "".to_string(),
            }
            .into(),
            unresolved_module.syntax().text_range(),
        ))
    }
}

impl DiagnosticWithFix for NoSuchField {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        let root = sema.db.parse_or_expand(self.file)?;
        missing_record_expr_field_fix(
            &sema,
            self.file.original_file(sema.db),
            &self.field.to_node(&root),
        )
    }
}

impl DiagnosticWithFix for MissingFields {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        // Note that although we could add a diagnostics to
        // fill the missing tuple field, e.g :
        // `struct A(usize);`
        // `let a = A { 0: () }`
        // but it is uncommon usage and it should not be encouraged.
        if self.missed_fields.iter().any(|it| it.as_tuple_index().is_some()) {
            return None;
        }

        let root = sema.db.parse_or_expand(self.file)?;
        let field_list_parent = self.field_list_parent.to_node(&root);
        let old_field_list = field_list_parent.record_expr_field_list()?;
        let mut new_field_list = old_field_list.clone();
        for f in self.missed_fields.iter() {
            let field =
                make::record_expr_field(make::name_ref(&f.to_string()), Some(make::expr_unit()));
            new_field_list = new_field_list.append_field(&field);
        }

        let edit = {
            let mut builder = TextEdit::builder();
            algo::diff(&old_field_list.syntax(), &new_field_list.syntax())
                .into_text_edit(&mut builder);
            builder.finish()
        };
        Some(Fix::new(
            "Fill struct fields",
            SourceChange::from_text_edit(self.file.original_file(sema.db), edit),
            sema.original_range(&field_list_parent.syntax()).range,
        ))
    }
}

impl DiagnosticWithFix for MissingOkOrSomeInTailExpr {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        let root = sema.db.parse_or_expand(self.file)?;
        let tail_expr = self.expr.to_node(&root);
        let tail_expr_range = tail_expr.syntax().text_range();
        let replacement = format!("{}({})", self.required, tail_expr.syntax());
        let edit = TextEdit::replace(tail_expr_range, replacement);
        let source_change = SourceChange::from_text_edit(self.file.original_file(sema.db), edit);
        let name = if self.required == "Ok" { "Wrap with Ok" } else { "Wrap with Some" };
        Some(Fix::new(name, source_change, tail_expr_range))
    }
}

impl DiagnosticWithFix for RemoveThisSemicolon {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        let root = sema.db.parse_or_expand(self.file)?;

        let semicolon = self
            .expr
            .to_node(&root)
            .syntax()
            .parent()
            .and_then(ast::ExprStmt::cast)
            .and_then(|expr| expr.semicolon_token())?
            .text_range();

        let edit = TextEdit::delete(semicolon);
        let source_change = SourceChange::from_text_edit(self.file.original_file(sema.db), edit);

        Some(Fix::new("Remove this semicolon", source_change, semicolon))
    }
}

impl DiagnosticWithFix for IncorrectCase {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        let root = sema.db.parse_or_expand(self.file)?;
        let name_node = self.ident.to_node(&root);

        let name_node = InFile::new(self.file, name_node.syntax());
        let frange = name_node.original_file_range(sema.db);
        let file_position = FilePosition { file_id: frange.file_id, offset: frange.range.start() };

        let rename_changes =
            rename_with_semantics(sema, file_position, &self.suggested_text).ok()?;

        let label = format!("Rename to {}", self.suggested_text);
        Some(Fix::new(&label, rename_changes, frange.range))
    }
}

impl DiagnosticWithFix for ReplaceFilterMapNextWithFindMap {
    fn fix(&self, sema: &Semantics<RootDatabase>) -> Option<Fix> {
        let root = sema.db.parse_or_expand(self.file)?;
        let next_expr = self.next_expr.to_node(&root);
        let next_call = ast::MethodCallExpr::cast(next_expr.syntax().clone())?;

        let filter_map_call = ast::MethodCallExpr::cast(next_call.receiver()?.syntax().clone())?;
        let filter_map_name_range = filter_map_call.name_ref()?.ident_token()?.text_range();
        let filter_map_args = filter_map_call.arg_list()?;

        let range_to_replace =
            TextRange::new(filter_map_name_range.start(), next_expr.syntax().text_range().end());
        let replacement = format!("find_map{}", filter_map_args.syntax().text());
        let trigger_range = next_expr.syntax().text_range();

        let edit = TextEdit::replace(range_to_replace, replacement);

        let source_change = SourceChange::from_text_edit(self.file.original_file(sema.db), edit);

        Some(Fix::new(
            "Replace filter_map(..).next() with find_map()",
            source_change,
            trigger_range,
        ))
    }
}

fn missing_record_expr_field_fix(
    sema: &Semantics<RootDatabase>,
    usage_file_id: FileId,
    record_expr_field: &ast::RecordExprField,
) -> Option<Fix> {
    let record_lit = ast::RecordExpr::cast(record_expr_field.syntax().parent()?.parent()?)?;
    let def_id = sema.resolve_variant(record_lit)?;
    let module;
    let def_file_id;
    let record_fields = match VariantDef::from(def_id) {
        VariantDef::Struct(s) => {
            module = s.module(sema.db);
            let source = s.source(sema.db)?;
            def_file_id = source.file_id;
            let fields = source.value.field_list()?;
            record_field_list(fields)?
        }
        VariantDef::Union(u) => {
            module = u.module(sema.db);
            let source = u.source(sema.db)?;
            def_file_id = source.file_id;
            source.value.record_field_list()?
        }
        VariantDef::Variant(e) => {
            module = e.module(sema.db);
            let source = e.source(sema.db)?;
            def_file_id = source.file_id;
            let fields = source.value.field_list()?;
            record_field_list(fields)?
        }
    };
    let def_file_id = def_file_id.original_file(sema.db);

    let new_field_type = sema.type_of_expr(&record_expr_field.expr()?)?;
    if new_field_type.is_unknown() {
        return None;
    }
    let new_field = make::record_field(
        None,
        make::name(record_expr_field.field_name()?.text()),
        make::ty(&new_field_type.display_source_code(sema.db, module.into()).ok()?),
    );

    let last_field = record_fields.fields().last()?;
    let last_field_syntax = last_field.syntax();
    let indent = IndentLevel::from_node(last_field_syntax);

    let mut new_field = new_field.to_string();
    if usage_file_id != def_file_id {
        new_field = format!("pub(crate) {}", new_field);
    }
    new_field = format!("\n{}{}", indent, new_field);

    let needs_comma = !last_field_syntax.to_string().ends_with(',');
    if needs_comma {
        new_field = format!(",{}", new_field);
    }

    let source_change = SourceChange::from_text_edit(
        def_file_id,
        TextEdit::insert(last_field_syntax.text_range().end(), new_field),
    );
    return Some(Fix::new("Create field", source_change, record_expr_field.syntax().text_range()));

    fn record_field_list(field_def_list: ast::FieldList) -> Option<ast::RecordFieldList> {
        match field_def_list {
            ast::FieldList::RecordFieldList(it) => Some(it),
            ast::FieldList::TupleFieldList(_) => None,
        }
    }
}
<|endoftext|>


<commit_before>pub(crate) mod tags;

mod highlights;
mod injector;

mod highlight;
mod format;
mod macro_;
mod inject;

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

use hir::{InFile, Name, Semantics};
use ide_db::{RootDatabase, SymbolKind};
use rustc_hash::FxHashMap;
use syntax::{
    ast::{self, HasFormatSpecifier},
    AstNode, AstToken, Direction, NodeOrToken,
    SyntaxKind::*,
    SyntaxNode, TextRange, WalkEvent, T,
};

use crate::{
    syntax_highlighting::{
        format::highlight_format_string, highlights::Highlights, macro_::MacroHighlighter,
        tags::Highlight,
    },
    FileId, HlMod, HlTag,
};

pub(crate) use html::highlight_as_html;

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

// Feature: Semantic Syntax Highlighting
//
// rust-analyzer highlights the code semantically.
// For example, `bar` in `foo::Bar` might be colored differently depending on whether `Bar` is an enum or a trait.
// rust-analyzer does not specify colors directly, instead it assigns tag (like `struct`) and a set of modifiers (like `declaration`) to each token.
// It's up to the client to map those to specific colors.
//
// The general rule is that a reference to an entity gets colored the same way as the entity itself.
// We also give special modifier for `mut` and `&mut` local variables.
//
// image::https://user-images.githubusercontent.com/48062697/113164457-06cfb980-9239-11eb-819b-0f93e646acf8.png[]
pub(crate) fn highlight(
    db: &RootDatabase,
    file_id: FileId,
    range_to_highlight: Option<TextRange>,
    syntactic_name_ref_highlighting: bool,
) -> Vec<HlRange> {
    let _p = profile::span("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().unwrap(),
                };
                (node, range)
            }
            None => (source_file.syntax().clone(), source_file.syntax().text_range()),
        }
    };

    let mut hl = highlights::Highlights::new(root.text_range());
    traverse(
        &mut hl,
        &sema,
        InFile::new(file_id.into(), &root),
        range_to_highlight,
        syntactic_name_ref_highlighting,
    );
    hl.to_vec()
}

fn traverse(
    hl: &mut Highlights,
    sema: &Semantics<RootDatabase>,
    root: InFile<&SyntaxNode>,
    range_to_highlight: TextRange,
    syntactic_name_ref_highlighting: bool,
) {
    let mut bindings_shadow_count: FxHashMap<Name, u32> = FxHashMap::default();

    let mut current_macro_call: Option<ast::MacroCall> = None;
    let mut current_macro: Option<ast::Macro> = None;
    let mut macro_highlighter = MacroHighlighter::default();
    let mut inside_attribute = false;

    // 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.value.preorder_with_tokens() {
        let event_range = match &event {
            WalkEvent::Enter(it) | WalkEvent::Leave(it) => it.text_range(),
        };

        // Element outside of the viewport, no need to highlight
        if range_to_highlight.intersect(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)) => {
                if let Some(range) = macro_call_range(&mc) {
                    hl.add(HlRange {
                        range,
                        highlight: HlTag::Symbol(SymbolKind::Macro).into(),
                        binding_hash: None,
                    });
                }
                current_macro_call = Some(mc.clone());
                continue;
            }
            WalkEvent::Leave(Some(mc)) => {
                assert_eq!(current_macro_call, Some(mc));
                current_macro_call = None;
            }
            _ => (),
        }

        match event.clone().map(|it| it.into_node().and_then(ast::Macro::cast)) {
            WalkEvent::Enter(Some(mac)) => {
                macro_highlighter.init();
                current_macro = Some(mac);
                continue;
            }
            WalkEvent::Leave(Some(mac)) => {
                assert_eq!(current_macro, Some(mac));
                current_macro = None;
                macro_highlighter = MacroHighlighter::default();
            }
            _ => (),
        }
        match &event {
            WalkEvent::Enter(NodeOrToken::Node(node)) if ast::Attr::can_cast(node.kind()) => {
                inside_attribute = true
            }
            WalkEvent::Leave(NodeOrToken::Node(node)) if ast::Attr::can_cast(node.kind()) => {
                inside_attribute = false
            }
            _ => (),
        }

        let element = match event {
            WalkEvent::Enter(it) => it,
            WalkEvent::Leave(it) => {
                if let Some(node) = it.as_node() {
                    inject::doc_comment(hl, sema, root.with_value(node));
                }
                continue;
            }
        };

        let range = element.text_range();

        if current_macro.is_some() {
            if let Some(tok) = element.as_token() {
                macro_highlighter.advance(tok);
            }
        }

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

        if let Some(token) = element.as_token().cloned().and_then(ast::String::cast) {
            if token.is_raw() {
                let expanded = element_to_highlight.as_token().unwrap().clone();
                if inject::ra_fixture(hl, &sema, token, expanded).is_some() {
                    continue;
                }
            }
        }

        if let Some(_) = macro_highlighter.highlight(element_to_highlight.clone()) {
            continue;
        }

        if let Some((mut highlight, binding_hash)) = highlight::element(
            &sema,
            &mut bindings_shadow_count,
            syntactic_name_ref_highlighting,
            element_to_highlight.clone(),
        ) {
            if inside_attribute {
                highlight = highlight | HlMod::Attribute;
            }

            hl.add(HlRange { range, highlight, binding_hash });
        }

        if let Some(string) = element_to_highlight.as_token().cloned().and_then(ast::String::cast) {
            highlight_format_string(hl, &string, range);
            // Highlight escape sequences
            if let Some(char_ranges) = string.char_ranges() {
                for (piece_range, _) in char_ranges.iter().filter(|(_, char)| char.is_ok()) {
                    if string.text()[piece_range.start().into()..].starts_with('\\') {
                        hl.add(HlRange {
                            range: piece_range + range.start(),
                            highlight: HlTag::EscapeSequence.into(),
                            binding_hash: None,
                        });
                    }
                }
            }
        }
    }
}

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::new(range_start, range_end))
}
<commit_msg>Additional syntax highlighting screenshot in manual<commit_after>pub(crate) mod tags;

mod highlights;
mod injector;

mod highlight;
mod format;
mod macro_;
mod inject;

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

use hir::{InFile, Name, Semantics};
use ide_db::{RootDatabase, SymbolKind};
use rustc_hash::FxHashMap;
use syntax::{
    ast::{self, HasFormatSpecifier},
    AstNode, AstToken, Direction, NodeOrToken,
    SyntaxKind::*,
    SyntaxNode, TextRange, WalkEvent, T,
};

use crate::{
    syntax_highlighting::{
        format::highlight_format_string, highlights::Highlights, macro_::MacroHighlighter,
        tags::Highlight,
    },
    FileId, HlMod, HlTag,
};

pub(crate) use html::highlight_as_html;

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

// Feature: Semantic Syntax Highlighting
//
// rust-analyzer highlights the code semantically.
// For example, `bar` in `foo::Bar` might be colored differently depending on whether `Bar` is an enum or a trait.
// rust-analyzer does not specify colors directly, instead it assigns tag (like `struct`) and a set of modifiers (like `declaration`) to each token.
// It's up to the client to map those to specific colors.
//
// The general rule is that a reference to an entity gets colored the same way as the entity itself.
// We also give special modifier for `mut` and `&mut` local variables.
//
// image::https://user-images.githubusercontent.com/48062697/113164457-06cfb980-9239-11eb-819b-0f93e646acf8.png[]
// image::https://user-images.githubusercontent.com/48062697/113187625-f7f50100-9250-11eb-825e-91c58f236071.png[]
pub(crate) fn highlight(
    db: &RootDatabase,
    file_id: FileId,
    range_to_highlight: Option<TextRange>,
    syntactic_name_ref_highlighting: bool,
) -> Vec<HlRange> {
    let _p = profile::span("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().unwrap(),
                };
                (node, range)
            }
            None => (source_file.syntax().clone(), source_file.syntax().text_range()),
        }
    };

    let mut hl = highlights::Highlights::new(root.text_range());
    traverse(
        &mut hl,
        &sema,
        InFile::new(file_id.into(), &root),
        range_to_highlight,
        syntactic_name_ref_highlighting,
    );
    hl.to_vec()
}

fn traverse(
    hl: &mut Highlights,
    sema: &Semantics<RootDatabase>,
    root: InFile<&SyntaxNode>,
    range_to_highlight: TextRange,
    syntactic_name_ref_highlighting: bool,
) {
    let mut bindings_shadow_count: FxHashMap<Name, u32> = FxHashMap::default();

    let mut current_macro_call: Option<ast::MacroCall> = None;
    let mut current_macro: Option<ast::Macro> = None;
    let mut macro_highlighter = MacroHighlighter::default();
    let mut inside_attribute = false;

    // 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.value.preorder_with_tokens() {
        let event_range = match &event {
            WalkEvent::Enter(it) | WalkEvent::Leave(it) => it.text_range(),
        };

        // Element outside of the viewport, no need to highlight
        if range_to_highlight.intersect(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)) => {
                if let Some(range) = macro_call_range(&mc) {
                    hl.add(HlRange {
                        range,
                        highlight: HlTag::Symbol(SymbolKind::Macro).into(),
                        binding_hash: None,
                    });
                }
                current_macro_call = Some(mc.clone());
                continue;
            }
            WalkEvent::Leave(Some(mc)) => {
                assert_eq!(current_macro_call, Some(mc));
                current_macro_call = None;
            }
            _ => (),
        }

        match event.clone().map(|it| it.into_node().and_then(ast::Macro::cast)) {
            WalkEvent::Enter(Some(mac)) => {
                macro_highlighter.init();
                current_macro = Some(mac);
                continue;
            }
            WalkEvent::Leave(Some(mac)) => {
                assert_eq!(current_macro, Some(mac));
                current_macro = None;
                macro_highlighter = MacroHighlighter::default();
            }
            _ => (),
        }
        match &event {
            WalkEvent::Enter(NodeOrToken::Node(node)) if ast::Attr::can_cast(node.kind()) => {
                inside_attribute = true
            }
            WalkEvent::Leave(NodeOrToken::Node(node)) if ast::Attr::can_cast(node.kind()) => {
                inside_attribute = false
            }
            _ => (),
        }

        let element = match event {
            WalkEvent::Enter(it) => it,
            WalkEvent::Leave(it) => {
                if let Some(node) = it.as_node() {
                    inject::doc_comment(hl, sema, root.with_value(node));
                }
                continue;
            }
        };

        let range = element.text_range();

        if current_macro.is_some() {
            if let Some(tok) = element.as_token() {
                macro_highlighter.advance(tok);
            }
        }

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

        if let Some(token) = element.as_token().cloned().and_then(ast::String::cast) {
            if token.is_raw() {
                let expanded = element_to_highlight.as_token().unwrap().clone();
                if inject::ra_fixture(hl, &sema, token, expanded).is_some() {
                    continue;
                }
            }
        }

        if let Some(_) = macro_highlighter.highlight(element_to_highlight.clone()) {
            continue;
        }

        if let Some((mut highlight, binding_hash)) = highlight::element(
            &sema,
            &mut bindings_shadow_count,
            syntactic_name_ref_highlighting,
            element_to_highlight.clone(),
        ) {
            if inside_attribute {
                highlight = highlight | HlMod::Attribute;
            }

            hl.add(HlRange { range, highlight, binding_hash });
        }

        if let Some(string) = element_to_highlight.as_token().cloned().and_then(ast::String::cast) {
            highlight_format_string(hl, &string, range);
            // Highlight escape sequences
            if let Some(char_ranges) = string.char_ranges() {
                for (piece_range, _) in char_ranges.iter().filter(|(_, char)| char.is_ok()) {
                    if string.text()[piece_range.start().into()..].starts_with('\\') {
                        hl.add(HlRange {
                            range: piece_range + range.start(),
                            highlight: HlTag::EscapeSequence.into(),
                            binding_hash: None,
                        });
                    }
                }
            }
        }
    }
}

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::new(range_start, range_end))
}
<|endoftext|>


<commit_before>//! Handles dynamic library loading for proc macro

use crate::{proc_macro::bridge, rustc_server::TokenStream};
use std::path::Path;

use goblin::{mach::Mach, Object};
use libloading::Library;
use ra_proc_macro::ProcMacroKind;

static NEW_REGISTRAR_SYMBOL: &str = "_rustc_proc_macro_decls_";

fn get_symbols_from_lib(file: &Path) -> Option<Vec<String>> {
    let buffer = std::fs::read(file).ok()?;
    let object = Object::parse(&buffer).ok()?;

    return match object {
        Object::Elf(elf) => {
            let symbols = elf.dynstrtab.to_vec().ok()?;
            let names = symbols.iter().map(|s| s.to_string()).collect();
            Some(names)
        }
        Object::PE(pe) => {
            let symbol_names =
                pe.exports.iter().flat_map(|s| s.name).map(|n| n.to_string()).collect();
            Some(symbol_names)
        }
        Object::Mach(mach) => match mach {
            Mach::Binary(binary) => {
                let exports = binary.exports().ok()?;
                let names = exports
                    .iter()
                    .map(|s| {
                        // In macos doc:
                        // https://developer.apple.com/library/archive/documentation/System/Conceptual/ManPages_iPhoneOS/man3/dlsym.3.html
                        // Unlike other dyld API's, the symbol name passed to dlsym() must NOT be
                        // prepended with an underscore.
                        if s.name.starts_with("_") {
                            s.name[1..].to_string()
                        } else {
                            s.name.to_string()
                        }
                    })
                    .collect();
                Some(names)
            }
            Mach::Fat(_) => None,
        },
        Object::Archive(_) | Object::Unknown(_) => None,
    };
}

fn is_derive_registrar_symbol(symbol: &str) -> bool {
    symbol.contains(NEW_REGISTRAR_SYMBOL)
}

fn find_registrar_symbol(file: &Path) -> Option<String> {
    let symbols = get_symbols_from_lib(file)?;

    symbols.iter().find(|s| is_derive_registrar_symbol(s)).map(|s| s.clone())
}

/// Loads dynamic library in platform dependent manner.
///
/// For unix, you have to use RTLD_DEEPBIND flag to escape problems described
/// [here](https://github.com/fedochet/rust-proc-macro-panic-inside-panic-expample)
/// and [here](https://github.com/rust-lang/rust/issues/60593).
///
/// Usage of RTLD_DEEPBIND
/// [here](https://github.com/fedochet/rust-proc-macro-panic-inside-panic-expample/issues/1)
///
/// It seems that on Windows that behaviour is default, so we do nothing in that case.
#[cfg(windows)]
fn load_library(file: &Path) -> Result<Library, libloading::Error> {
    Library::new(file)
}

#[cfg(unix)]
fn load_library(file: &Path) -> Result<Library, libloading::Error> {
    use libloading::os::unix::Library as UnixLibrary;
    use std::os::raw::c_int;

    const RTLD_NOW: c_int = 0x00002;
    const RTLD_DEEPBIND: c_int = 0x00008;

    UnixLibrary::open(Some(file), RTLD_NOW | RTLD_DEEPBIND).map(|lib| lib.into())
}

struct ProcMacroLibraryLibloading {
    // Hold the dylib to prevent it for unloadeding
    _lib: Library,
    exported_macros: Vec<bridge::client::ProcMacro>,
}

impl ProcMacroLibraryLibloading {
    fn open(file: &Path) -> Result<Self, String> {
        let symbol_name = find_registrar_symbol(file)
            .ok_or(format!("Cannot find registrar symbol in file {:?}", file))?;

        let lib = load_library(file).map_err(|e| e.to_string())?;
        let exported_macros = {
            let macros: libloading::Symbol<&&[bridge::client::ProcMacro]> =
                unsafe { lib.get(symbol_name.as_bytes()) }.map_err(|e| e.to_string())?;
            macros.to_vec()
        };

        Ok(ProcMacroLibraryLibloading { _lib: lib, exported_macros })
    }
}

type ProcMacroLibraryImpl = ProcMacroLibraryLibloading;

pub struct Expander {
    libs: Vec<ProcMacroLibraryImpl>,
}

impl Expander {
    pub fn new<P: AsRef<Path>>(lib: &P) -> Result<Expander, String> {
        let mut libs = vec![];
        /* Some libraries for dynamic loading require canonicalized path (even when it is
        already absolute
        */
        let lib =
            lib.as_ref().canonicalize().expect(&format!("Cannot canonicalize {:?}", lib.as_ref()));

        let library = ProcMacroLibraryImpl::open(&lib)?;
        libs.push(library);

        Ok(Expander { libs })
    }

    pub fn expand(
        &self,
        macro_name: &str,
        macro_body: &ra_tt::Subtree,
        attributes: Option<&ra_tt::Subtree>,
    ) -> Result<ra_tt::Subtree, bridge::PanicMessage> {
        let parsed_body = TokenStream::with_subtree(macro_body.clone());

        let parsed_attributes = attributes
            .map_or(crate::rustc_server::TokenStream::new(), |attr| {
                TokenStream::with_subtree(attr.clone())
            });

        for lib in &self.libs {
            for proc_macro in &lib.exported_macros {
                match proc_macro {
                    bridge::client::ProcMacro::CustomDerive { trait_name, client, .. }
                        if *trait_name == macro_name =>
                    {
                        let res = client.run(
                            &crate::proc_macro::bridge::server::SameThread,
                            crate::rustc_server::Rustc::default(),
                            parsed_body,
                        );
                        return res.map(|it| it.subtree);
                    }
                    bridge::client::ProcMacro::Bang { name, client } if *name == macro_name => {
                        let res = client.run(
                            &crate::proc_macro::bridge::server::SameThread,
                            crate::rustc_server::Rustc::default(),
                            parsed_body,
                        );
                        return res.map(|it| it.subtree);
                    }
                    bridge::client::ProcMacro::Attr { name, client } if *name == macro_name => {
                        let res = client.run(
                            &crate::proc_macro::bridge::server::SameThread,
                            crate::rustc_server::Rustc::default(),
                            parsed_attributes,
                            parsed_body,
                        );

                        return res.map(|it| it.subtree);
                    }
                    _ => continue,
                }
            }
        }

        Err(bridge::PanicMessage::String("Nothing to expand".to_string()))
    }

    pub fn list_macros(&self) -> Result<Vec<(String, ProcMacroKind)>, bridge::PanicMessage> {
        let mut result = vec![];

        for lib in &self.libs {
            for proc_macro in &lib.exported_macros {
                let res = match proc_macro {
                    bridge::client::ProcMacro::CustomDerive { trait_name, .. } => {
                        (trait_name.to_string(), ProcMacroKind::CustomDerive)
                    }
                    bridge::client::ProcMacro::Bang { name, .. } => {
                        (name.to_string(), ProcMacroKind::FuncLike)
                    }
                    bridge::client::ProcMacro::Attr { name, .. } => {
                        (name.to_string(), ProcMacroKind::Attr)
                    }
                };
                result.push(res);
            }
        }

        Ok(result)
    }
}
<commit_msg>Improve error handling<commit_after>//! Handles dynamic library loading for proc macro

use crate::{proc_macro::bridge, rustc_server::TokenStream};
use std::path::Path;

use goblin::{mach::Mach, Object};
use libloading::Library;
use ra_proc_macro::ProcMacroKind;

use std::io::Error as IoError;
use std::io::ErrorKind as IoErrorKind;

const NEW_REGISTRAR_SYMBOL: &str = "_rustc_proc_macro_decls_";

fn invalid_data_err(e: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> IoError {
    IoError::new(IoErrorKind::InvalidData, e)
}

fn get_symbols_from_lib(file: &Path) -> Result<Vec<String>, IoError> {
    let buffer = std::fs::read(file)?;
    let object = Object::parse(&buffer).map_err(invalid_data_err)?;

    match object {
        Object::Elf(elf) => {
            let symbols = elf.dynstrtab.to_vec().map_err(invalid_data_err)?;
            let names = symbols.iter().map(|s| s.to_string()).collect();
            Ok(names)
        }
        Object::PE(pe) => {
            let symbol_names =
                pe.exports.iter().flat_map(|s| s.name).map(|n| n.to_string()).collect();
            Ok(symbol_names)
        }
        Object::Mach(mach) => match mach {
            Mach::Binary(binary) => {
                let exports = binary.exports().map_err(invalid_data_err)?;
                let names = exports
                    .into_iter()
                    .map(|s| {
                        // In macos doc:
                        // https://developer.apple.com/library/archive/documentation/System/Conceptual/ManPages_iPhoneOS/man3/dlsym.3.html
                        // Unlike other dyld API's, the symbol name passed to dlsym() must NOT be
                        // prepended with an underscore.
                        if s.name.starts_with("_") {
                            s.name[1..].to_string()
                        } else {
                            s.name
                        }
                    })
                    .collect();
                Ok(names)
            }
            Mach::Fat(_) => Ok(vec![]),
        },
        Object::Archive(_) | Object::Unknown(_) => Ok(vec![]),
    }
}

fn is_derive_registrar_symbol(symbol: &str) -> bool {
    symbol.contains(NEW_REGISTRAR_SYMBOL)
}

fn find_registrar_symbol(file: &Path) -> Result<Option<String>, IoError> {
    let symbols = get_symbols_from_lib(file)?;
    Ok(symbols.into_iter().find(|s| is_derive_registrar_symbol(s)))
}

/// Loads dynamic library in platform dependent manner.
///
/// For unix, you have to use RTLD_DEEPBIND flag to escape problems described
/// [here](https://github.com/fedochet/rust-proc-macro-panic-inside-panic-expample)
/// and [here](https://github.com/rust-lang/rust/issues/60593).
///
/// Usage of RTLD_DEEPBIND
/// [here](https://github.com/fedochet/rust-proc-macro-panic-inside-panic-expample/issues/1)
///
/// It seems that on Windows that behaviour is default, so we do nothing in that case.
#[cfg(windows)]
fn load_library(file: &Path) -> Result<Library, libloading::Error> {
    Library::new(file)
}

#[cfg(unix)]
fn load_library(file: &Path) -> Result<Library, libloading::Error> {
    use libloading::os::unix::Library as UnixLibrary;
    use std::os::raw::c_int;

    const RTLD_NOW: c_int = 0x00002;
    const RTLD_DEEPBIND: c_int = 0x00008;

    UnixLibrary::open(Some(file), RTLD_NOW | RTLD_DEEPBIND).map(|lib| lib.into())
}

struct ProcMacroLibraryLibloading {
    // Hold the dylib to prevent it for unloadeding
    _lib: Library,
    exported_macros: Vec<bridge::client::ProcMacro>,
}

impl ProcMacroLibraryLibloading {
    fn open(file: &Path) -> Result<Self, IoError> {
        let symbol_name = find_registrar_symbol(file)?
            .ok_or(invalid_data_err(format!("Cannot find registrar symbol in file {:?}", file)))?;

        let lib = load_library(file).map_err(invalid_data_err)?;
        let exported_macros = {
            let macros: libloading::Symbol<&&[bridge::client::ProcMacro]> =
                unsafe { lib.get(symbol_name.as_bytes()) }.map_err(invalid_data_err)?;
            macros.to_vec()
        };

        Ok(ProcMacroLibraryLibloading { _lib: lib, exported_macros })
    }
}

type ProcMacroLibraryImpl = ProcMacroLibraryLibloading;

pub struct Expander {
    libs: Vec<ProcMacroLibraryImpl>,
}

impl Expander {
    pub fn new<P: AsRef<Path>>(lib: &P) -> Result<Expander, String> {
        let mut libs = vec![];
        /* Some libraries for dynamic loading require canonicalized path (even when it is
        already absolute
        */
        let lib =
            lib.as_ref().canonicalize().expect(&format!("Cannot canonicalize {:?}", lib.as_ref()));

        let library = ProcMacroLibraryImpl::open(&lib).map_err(|e| e.to_string())?;
        libs.push(library);

        Ok(Expander { libs })
    }

    pub fn expand(
        &self,
        macro_name: &str,
        macro_body: &ra_tt::Subtree,
        attributes: Option<&ra_tt::Subtree>,
    ) -> Result<ra_tt::Subtree, bridge::PanicMessage> {
        let parsed_body = TokenStream::with_subtree(macro_body.clone());

        let parsed_attributes = attributes
            .map_or(crate::rustc_server::TokenStream::new(), |attr| {
                TokenStream::with_subtree(attr.clone())
            });

        for lib in &self.libs {
            for proc_macro in &lib.exported_macros {
                match proc_macro {
                    bridge::client::ProcMacro::CustomDerive { trait_name, client, .. }
                        if *trait_name == macro_name =>
                    {
                        let res = client.run(
                            &crate::proc_macro::bridge::server::SameThread,
                            crate::rustc_server::Rustc::default(),
                            parsed_body,
                        );
                        return res.map(|it| it.subtree);
                    }
                    bridge::client::ProcMacro::Bang { name, client } if *name == macro_name => {
                        let res = client.run(
                            &crate::proc_macro::bridge::server::SameThread,
                            crate::rustc_server::Rustc::default(),
                            parsed_body,
                        );
                        return res.map(|it| it.subtree);
                    }
                    bridge::client::ProcMacro::Attr { name, client } if *name == macro_name => {
                        let res = client.run(
                            &crate::proc_macro::bridge::server::SameThread,
                            crate::rustc_server::Rustc::default(),
                            parsed_attributes,
                            parsed_body,
                        );

                        return res.map(|it| it.subtree);
                    }
                    _ => continue,
                }
            }
        }

        Err(bridge::PanicMessage::String("Nothing to expand".to_string()))
    }

    pub fn list_macros(&self) -> Result<Vec<(String, ProcMacroKind)>, bridge::PanicMessage> {
        let mut result = vec![];

        for lib in &self.libs {
            for proc_macro in &lib.exported_macros {
                let res = match proc_macro {
                    bridge::client::ProcMacro::CustomDerive { trait_name, .. } => {
                        (trait_name.to_string(), ProcMacroKind::CustomDerive)
                    }
                    bridge::client::ProcMacro::Bang { name, .. } => {
                        (name.to_string(), ProcMacroKind::FuncLike)
                    }
                    bridge::client::ProcMacro::Attr { name, .. } => {
                        (name.to_string(), ProcMacroKind::Attr)
                    }
                };
                result.push(res);
            }
        }

        Ok(result)
    }
}
<|endoftext|>


<commit_before>use super::*;
use ssz::{Decode, DecodeError};

fn get_block_bytes<T: Store>(store: &T, root: Hash256) -> Result<Option<Vec<u8>>, Error> {
    store.get_bytes(BeaconBlock::db_column().into(), &root[..])
}

fn read_slot_from_block_bytes(bytes: &[u8]) -> Result<Slot, DecodeError> {
    let end = std::cmp::min(Slot::ssz_fixed_len(), bytes.len());

    Slot::from_ssz_bytes(&bytes[0..end])
}

fn read_previous_block_root_from_block_bytes(bytes: &[u8]) -> Result<Hash256, DecodeError> {
    let previous_bytes = Slot::ssz_fixed_len();
    let slice = bytes
        .get(previous_bytes..previous_bytes + Hash256::ssz_fixed_len())
        .ok_or_else(|| DecodeError::BytesInvalid("Not enough bytes.".to_string()))?;

    Hash256::from_ssz_bytes(slice)
}

pub fn get_block_at_preceeding_slot<T: Store>(
    store: &T,
    slot: Slot,
    start_root: Hash256,
) -> Result<Option<(Hash256, BeaconBlock)>, Error> {
    let mut root = start_root;

    loop {
        if let Some(bytes) = get_block_bytes(store, root)? {
            let this_slot = read_slot_from_block_bytes(&bytes)?;

            if this_slot == slot {
                let block = BeaconBlock::from_ssz_bytes(&bytes)?;
                break Ok(Some((root, block)));
            } else if this_slot < slot {
                break Ok(None);
            } else {
                root = read_previous_block_root_from_block_bytes(&bytes)?;
            }
        } else {
            break Ok(None);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ssz::Encode;
    use tree_hash::TreeHash;

    #[test]
    fn read_slot() {
        let spec = FewValidatorsEthSpec::spec();

        let test_slot = |slot: Slot| {
            let mut block = BeaconBlock::empty(&spec);
            block.slot = slot;
            let bytes = block.as_ssz_bytes();
            assert_eq!(read_slot_from_block_bytes(&bytes).unwrap(), slot);
        };

        test_slot(Slot::new(0));
        test_slot(Slot::new(1));
        test_slot(Slot::new(42));
        test_slot(Slot::new(u64::max_value()));
    }

    #[test]
    fn bad_slot() {
        for i in 0..8 {
            assert!(read_slot_from_block_bytes(&vec![0; i]).is_err());
        }
    }

    #[test]
    fn read_previous_block_root() {
        let spec = FewValidatorsEthSpec::spec();

        let test_root = |root: Hash256| {
            let mut block = BeaconBlock::empty(&spec);
            block.previous_block_root = root;
            let bytes = block.as_ssz_bytes();
            assert_eq!(
                read_previous_block_root_from_block_bytes(&bytes).unwrap(),
                root
            );
        };

        test_root(Hash256::random());
        test_root(Hash256::random());
        test_root(Hash256::random());
    }

    fn build_chain(
        store: &impl Store,
        slots: &[usize],
        spec: &ChainSpec,
    ) -> Vec<(Hash256, BeaconBlock)> {
        let mut blocks_and_roots: Vec<(Hash256, BeaconBlock)> = vec![];

        for (i, slot) in slots.iter().enumerate() {
            let mut block = BeaconBlock::empty(spec);
            block.slot = Slot::from(*slot);

            if i > 0 {
                block.previous_block_root = blocks_and_roots[i - 1].0;
            }

            let root = Hash256::from_slice(&block.tree_hash_root());

            store.put(&root, &block).unwrap();
            blocks_and_roots.push((root, block));
        }

        blocks_and_roots
    }

    #[test]
    fn chain_without_skips() {
        let n: usize = 10;
        let store = MemoryStore::open();
        let spec = FewValidatorsEthSpec::spec();

        let slots: Vec<usize> = (0..n).collect();
        let blocks_and_roots = build_chain(&store, &slots, &spec);

        for source in 1..n {
            for target in 0..=source {
                let (source_root, _source_block) = &blocks_and_roots[source];
                let (target_root, target_block) = &blocks_and_roots[target];

                let (found_root, found_block) = store
                    .get_block_at_preceeding_slot(*source_root, target_block.slot)
                    .unwrap()
                    .unwrap();

                assert_eq!(found_root, *target_root);
                assert_eq!(found_block, *target_block);
            }
        }
    }

    #[test]
    fn chain_with_skips() {
        let store = MemoryStore::open();
        let spec = FewValidatorsEthSpec::spec();

        let slots = vec![0, 1, 2, 5];

        let blocks_and_roots = build_chain(&store, &slots, &spec);

        // Valid slots
        for target in 0..3 {
            let (source_root, _source_block) = &blocks_and_roots[3];
            let (target_root, target_block) = &blocks_and_roots[target];

            let (found_root, found_block) = store
                .get_block_at_preceeding_slot(*source_root, target_block.slot)
                .unwrap()
                .unwrap();

            assert_eq!(found_root, *target_root);
            assert_eq!(found_block, *target_block);
        }

        // Slot that doesn't exist
        let (source_root, _source_block) = &blocks_and_roots[3];
        assert!(store
            .get_block_at_preceeding_slot(*source_root, Slot::new(3))
            .unwrap()
            .is_none());

        // Slot too high
        let (source_root, _source_block) = &blocks_and_roots[3];
        assert!(store
            .get_block_at_preceeding_slot(*source_root, Slot::new(3))
            .unwrap()
            .is_none());
    }
}
<commit_msg>Break `store` fns into smaller pieces<commit_after>use super::*;
use ssz::{Decode, DecodeError};

fn get_block_bytes<T: Store>(store: &T, root: Hash256) -> Result<Option<Vec<u8>>, Error> {
    store.get_bytes(BeaconBlock::db_column().into(), &root[..])
}

fn read_slot_from_block_bytes(bytes: &[u8]) -> Result<Slot, DecodeError> {
    let end = std::cmp::min(Slot::ssz_fixed_len(), bytes.len());

    Slot::from_ssz_bytes(&bytes[0..end])
}

fn read_previous_block_root_from_block_bytes(bytes: &[u8]) -> Result<Hash256, DecodeError> {
    let previous_bytes = Slot::ssz_fixed_len();
    let slice = bytes
        .get(previous_bytes..previous_bytes + Hash256::ssz_fixed_len())
        .ok_or_else(|| DecodeError::BytesInvalid("Not enough bytes.".to_string()))?;

    Hash256::from_ssz_bytes(slice)
}

pub fn get_block_at_preceeding_slot<T: Store>(
    store: &T,
    slot: Slot,
    start_root: Hash256,
) -> Result<Option<(Hash256, BeaconBlock)>, Error> {
    Ok(match get_at_preceeding_slot(store, slot, start_root)? {
        Some((hash, bytes)) => Some((hash, BeaconBlock::from_ssz_bytes(&bytes)?)),
        None => None,
    })
}

fn get_at_preceeding_slot<T: Store>(
    store: &T,
    slot: Slot,
    mut root: Hash256,
) -> Result<Option<(Hash256, Vec<u8>)>, Error> {
    loop {
        if let Some(bytes) = get_block_bytes(store, root)? {
            let this_slot = read_slot_from_block_bytes(&bytes)?;

            if this_slot == slot {
                break Ok(Some((root, bytes)));
            } else if this_slot < slot {
                break Ok(None);
            } else {
                root = read_previous_block_root_from_block_bytes(&bytes)?;
            }
        } else {
            break Ok(None);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ssz::Encode;
    use tree_hash::TreeHash;

    #[test]
    fn read_slot() {
        let spec = FewValidatorsEthSpec::spec();

        let test_slot = |slot: Slot| {
            let mut block = BeaconBlock::empty(&spec);
            block.slot = slot;
            let bytes = block.as_ssz_bytes();
            assert_eq!(read_slot_from_block_bytes(&bytes).unwrap(), slot);
        };

        test_slot(Slot::new(0));
        test_slot(Slot::new(1));
        test_slot(Slot::new(42));
        test_slot(Slot::new(u64::max_value()));
    }

    #[test]
    fn bad_slot() {
        for i in 0..8 {
            assert!(read_slot_from_block_bytes(&vec![0; i]).is_err());
        }
    }

    #[test]
    fn read_previous_block_root() {
        let spec = FewValidatorsEthSpec::spec();

        let test_root = |root: Hash256| {
            let mut block = BeaconBlock::empty(&spec);
            block.previous_block_root = root;
            let bytes = block.as_ssz_bytes();
            assert_eq!(
                read_previous_block_root_from_block_bytes(&bytes).unwrap(),
                root
            );
        };

        test_root(Hash256::random());
        test_root(Hash256::random());
        test_root(Hash256::random());
    }

    fn build_chain(
        store: &impl Store,
        slots: &[usize],
        spec: &ChainSpec,
    ) -> Vec<(Hash256, BeaconBlock)> {
        let mut blocks_and_roots: Vec<(Hash256, BeaconBlock)> = vec![];

        for (i, slot) in slots.iter().enumerate() {
            let mut block = BeaconBlock::empty(spec);
            block.slot = Slot::from(*slot);

            if i > 0 {
                block.previous_block_root = blocks_and_roots[i - 1].0;
            }

            let root = Hash256::from_slice(&block.tree_hash_root());

            store.put(&root, &block).unwrap();
            blocks_and_roots.push((root, block));
        }

        blocks_and_roots
    }

    #[test]
    fn chain_without_skips() {
        let n: usize = 10;
        let store = MemoryStore::open();
        let spec = FewValidatorsEthSpec::spec();

        let slots: Vec<usize> = (0..n).collect();
        let blocks_and_roots = build_chain(&store, &slots, &spec);

        for source in 1..n {
            for target in 0..=source {
                let (source_root, _source_block) = &blocks_and_roots[source];
                let (target_root, target_block) = &blocks_and_roots[target];

                let (found_root, found_block) = store
                    .get_block_at_preceeding_slot(*source_root, target_block.slot)
                    .unwrap()
                    .unwrap();

                assert_eq!(found_root, *target_root);
                assert_eq!(found_block, *target_block);
            }
        }
    }

    #[test]
    fn chain_with_skips() {
        let store = MemoryStore::open();
        let spec = FewValidatorsEthSpec::spec();

        let slots = vec![0, 1, 2, 5];

        let blocks_and_roots = build_chain(&store, &slots, &spec);

        // Valid slots
        for target in 0..3 {
            let (source_root, _source_block) = &blocks_and_roots[3];
            let (target_root, target_block) = &blocks_and_roots[target];

            let (found_root, found_block) = store
                .get_block_at_preceeding_slot(*source_root, target_block.slot)
                .unwrap()
                .unwrap();

            assert_eq!(found_root, *target_root);
            assert_eq!(found_block, *target_block);
        }

        // Slot that doesn't exist
        let (source_root, _source_block) = &blocks_and_roots[3];
        assert!(store
            .get_block_at_preceeding_slot(*source_root, Slot::new(3))
            .unwrap()
            .is_none());

        // Slot too high
        let (source_root, _source_block) = &blocks_and_roots[3];
        assert!(store
            .get_block_at_preceeding_slot(*source_root, Slot::new(3))
            .unwrap()
            .is_none());
    }
}
<|endoftext|>


<commit_before>//! An example that timelines the specified folder.
//!
//! The timeline action is one of the filesystem tools that the GRR agent
//! exposes. Its description is pretty simple: given a folder, it returns a
//! stream of stat entries for all the files in it (including all subfolders).
//!
//! This example is a tiny wrapper around this action that allows to execute it
//! as a standalone binary without all the RRG setup required. A primary reason
//! for its existence is to compare the efficiency Rust implementation against
//! the existing Python one.

use std::fs::File;
use std::path::{Path, PathBuf};

use rrg::action::timeline;
use rrg::session::Sink;
use structopt::StructOpt;

/// A type for the timelining example command-line arguments.
#[derive(StructOpt)]
#[structopt(name = "timeline", about = "A binary for the timeline action.")]
struct Opts {
    /// A path to the root directory to timeline.
    #[structopt(long = "root", name = "FILE", default_value = "/",
                help = "Root directory to timeline")]
    root: PathBuf,

    /// A path to a file to dump the results into.
    #[structopt(long = "output", name = "DIRECTORY",
                help = "File to dump the results into")]
    output: PathBuf,
}

/// A session type of the timeline example.
///
/// Every action handler needs a session object. "Real" sessions handle comms
/// with the GRR server, error reporting and alike. However, this is not needed
/// here and the example session simply dumps everything into the specified file
/// encoding data in the gzchunked format.
struct Session {
    /// A file into which the data chunks of the output are dumped.
    output: File,
}

impl Session {

    /// Creates a session object that will write to the given `output` path.
    fn open<P: AsRef<Path>>(output: P) -> Session {
        let output = File::create(output)
            .expect("failed to create the output file");

        Session {
            output: output,
        }
    }
}

impl rrg::session::Session for Session {

    fn reply<R>(&mut self, response: R) -> rrg::session::Result<()>
    where
        R: rrg::action::Response + 'static,
    {
        // For now we are not interested in doing anything useful with chunk ids
        // since everything is dumped into one file and there is no need to
        // refer to a particular chunk.
        //
        // In the future they might be useful for printing some statistics about
        // the collected files.
        drop(response);

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> rrg::session::Result<()>
    where
        R: rrg::action::Response + 'static,
    {
        use std::io::Write as _;
        use byteorder::{BigEndian, WriteBytesExt as _};

        // Just a sanity check in case the implementation of the timeline action
        // starts sending data to other sinks.
        assert_eq!(sink, Sink::TRANSFER_STORE);

        let response = (&response as &dyn std::any::Any)
            .downcast_ref::<timeline::Chunk>()
            .expect("unexpected response type");

        self.output.write_u64::<BigEndian>(response.data.len() as u64)
            .expect("failed to write chunk size tag");
        self.output.write_all(&response.data[..])
            .expect("failed to write chunk data");

        Ok(())
    }
}

fn main() {
    let opts = Opts::from_args();

    timeline::handle(&mut Session::open(opts.output), timeline::Request {
        root: opts.root,
    }).expect("failed to execute the action");
}
<commit_msg>Migrate the timeline example to `argh`.<commit_after>//! An example that timelines the specified folder.
//!
//! The timeline action is one of the filesystem tools that the GRR agent
//! exposes. Its description is pretty simple: given a folder, it returns a
//! stream of stat entries for all the files in it (including all subfolders).
//!
//! This example is a tiny wrapper around this action that allows to execute it
//! as a standalone binary without all the RRG setup required. A primary reason
//! for its existence is to compare the efficiency Rust implementation against
//! the existing Python one.

use std::fs::File;
use std::path::{Path, PathBuf};

use rrg::action::timeline;
use rrg::session::Sink;

/// A binary for the timeline action.
#[derive(argh::FromArgs)]
struct Opts {
    /// A path to the root directory to timeline.
    #[argh(option,
           long = "root",
           arg_name = "FILE",
           default = "::std::path::PathBuf::from(\"/\")",
           description = "root directory to timeline")]
    root: PathBuf,

    /// A path to a file to dump the results into.
    #[argh(positional,
           arg_name = "OUTPUT",
           description = "path to dump the results into")]
    output: PathBuf,
}

/// A session type of the timeline example.
///
/// Every action handler needs a session object. "Real" sessions handle comms
/// with the GRR server, error reporting and alike. However, this is not needed
/// here and the example session simply dumps everything into the specified file
/// encoding data in the gzchunked format.
struct Session {
    /// A file into which the data chunks of the output are dumped.
    output: File,
}

impl Session {

    /// Creates a session object that will write to the given `output` path.
    fn open<P: AsRef<Path>>(output: P) -> Session {
        let output = File::create(output)
            .expect("failed to create the output file");

        Session {
            output: output,
        }
    }
}

impl rrg::session::Session for Session {

    fn reply<R>(&mut self, response: R) -> rrg::session::Result<()>
    where
        R: rrg::action::Response + 'static,
    {
        // For now we are not interested in doing anything useful with chunk ids
        // since everything is dumped into one file and there is no need to
        // refer to a particular chunk.
        //
        // In the future they might be useful for printing some statistics about
        // the collected files.
        drop(response);

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> rrg::session::Result<()>
    where
        R: rrg::action::Response + 'static,
    {
        use std::io::Write as _;
        use byteorder::{BigEndian, WriteBytesExt as _};

        // Just a sanity check in case the implementation of the timeline action
        // starts sending data to other sinks.
        assert_eq!(sink, Sink::TRANSFER_STORE);

        let response = (&response as &dyn std::any::Any)
            .downcast_ref::<timeline::Chunk>()
            .expect("unexpected response type");

        self.output.write_u64::<BigEndian>(response.data.len() as u64)
            .expect("failed to write chunk size tag");
        self.output.write_all(&response.data[..])
            .expect("failed to write chunk data");

        Ok(())
    }
}

fn main() {
    let opts: Opts = argh::from_env();

    timeline::handle(&mut Session::open(opts.output), timeline::Request {
        root: opts.root,
    }).expect("failed to execute the action");
}
<|endoftext|>


<commit_before>extern crate "rust-crypto" as rust_crypto;
extern crate serialize;

use std::io::{TcpListener, TcpStream};
use std::io::{Acceptor, Listener};
use std::io::{File, BufferedStream};
use std::string::String;
use std::os;

use rust_crypto::sha1::Sha1;
use rust_crypto::digest::Digest;
use serialize::base64::{ToBase64, STANDARD};

pub const CR: u8 = b'\r';
pub const LF: u8 = b'\n';
pub const SP: u8 = b' ';
pub const COLON: u8 = b':';
pub const CRLF: &'static [u8] = &[CR,LF];

struct RequestedRoute {
    method: String,
    pathname: String
}

struct ClientHeader {
    key: String,
    value: String
}

enum Payload {
    Text(String)
}

enum Opcode {
    TextOp = 0x1
}

struct Message {
    payload: Payload,
    opcode: Opcode
 }

impl Message {
    fn send (msg: &[u8],
             mut stream: BufferedStream<TcpStream>,
             headers: Vec<ClientHeader>) {
        let (payload, opcode) = match msg {
            String => (Payload::Text(msg.to_string()), Opcode::TextOp)
        };
        let length = msg.len() as u8;

        stream.write_u8(0b1000_0000 | opcode as u8).unwrap();
        stream.write_u8(length).unwrap();
        stream.write(msg).unwrap();
    }

    fn new (msg: &String) -> Message {

        let (payload, opcode) = match msg {
            String => (Payload::Text(msg.to_string()), Opcode::TextOp)
        };

        return Message{payload: payload, opcode: opcode};
    }

    fn snd (self,
            mut stream: BufferedStream<TcpStream>,
            headers: Vec<ClientHeader>) {

        let msg = match self.payload {
            Payload::Text(ref s) => s.as_bytes()
        };

        let length = msg.len() as u8;

        println!("msg: {}, opcode: {}", msg, self.opcode as u8);

        stream.write_u8(0b1000_0000 | self.opcode as u8).unwrap();
        stream.write_u8(length).unwrap();
        stream.write(msg).unwrap();

    }
}

fn get_header_by_name (header: &[u8], headers: &Vec<ClientHeader>) -> String {
    let mut result = String::new();

    for h in headers.iter() {
        if h.key.as_bytes() == header {
            result = h.value.clone();
            break;
        }
    }

    return result;
}

fn parse_request_line (header: &str) -> RequestedRoute {
    let mut pathname = String::new();
    let mut method = String::new();
    let mut starting = true;
    let mut found_method = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];

        if starting && cbyte != SP {
            method.push_str(ch);
        }
        if starting && cbyte == SP {
            starting = false;
            found_method = true;
            continue;
        }
        if found_method && cbyte != SP {
            pathname.push_str(ch);
        }
        if found_method && cbyte == SP {
            break;
        }

    }

    return RequestedRoute {method: method, pathname: pathname};
}

fn parse_normal_header (header: &str) -> ClientHeader {
    let mut lhs = String::new(); // lhs of ':' in header
    let mut rhs = String::new(); // rhs of ':' in header
    let mut found_colon = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];
        if !found_colon && cbyte == COLON {
            found_colon = true;
            continue;
        }

        if cbyte == CR || cbyte == LF {
            continue;
        }

        if !found_colon {
            lhs.push_str(ch);
        } else {
            rhs.push_str(ch);
        }
    }

    return ClientHeader {
        key:lhs.trim().to_string(),
        value:rhs.trim().to_string()
    };
}

fn get_normal_body (path_on_disk: &str) -> String {
    let path = Path::new(path_on_disk);
    let display = path.display();
    let mut file = match File::open(&path) {
        Ok(f) => f,
        Err(err) => panic!("file error: {}", err)
    };

    let content = match file.read_to_end() {
        Ok(c) => c,
        Err(err) => panic!("{}", err)
    };

    let s = String::from_utf8(content).unwrap();
    return s;
}

fn sec_handshake (from_server: &[u8]) -> String {

    let guid = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

    let mut sha = Sha1::new();

    sha.input(from_server);
    sha.input(guid);
    let mut out = [0u8, ..20];
    sha.result(out.as_mut_slice());

    return out.to_base64(STANDARD);
}

fn ws_handshake (mut stream: BufferedStream<TcpStream>,
                 headers: Vec<ClientHeader>) {

    let header_sec_key = b"Sec-WebSocket-Key";

    let from_server = get_header_by_name(header_sec_key, &headers);
    let accept = sec_handshake(from_server.as_bytes());

    let sec_header = format!("Sec-WebSocket-Accept: {}\r\n", accept);

    stream.write(b"HTTP/1.1 101 Switching Protocols\r\n").unwrap();
    stream.write(b"Upgrade: websocket\r\n").unwrap();
    stream.write(b"Connection: Upgrade\r\n").unwrap();
    stream.write(b"Sec-WebSocket-Version: 13\r\n").unwrap();
    stream.write(b"Sec-WebSocket-Protocol: protocolOne\r\n").unwrap();
    stream.write(sec_header.as_bytes()).unwrap();
    stream.write(b"\r\n");

    ws_listen(stream, headers);
}

fn ws_listen(mut stream: BufferedStream<TcpStream>,
             headers: Vec<ClientHeader>) {
    let string_msg = "this is a message".to_string();

    let msg = box Message::new(&string_msg);
    msg.snd(stream, headers);
}

fn main () {

    for argument in os::args().iter() {
        println!("arg: {}", argument);
    }

    let addr = "127.0.0.1:8099";

    let listener = TcpListener::bind(addr).unwrap();
    let mut acceptor = listener.listen().unwrap();
    println!("Listening on {}", addr);

    fn handle_client(mut stream: BufferedStream<TcpStream>) {
        let mut body: String = "
You sent it!
".to_string();

        let mut cur_line: String;
        let mut ii = 0u;
        let mut req: RequestedRoute = RequestedRoute {
            pathname: String::new(),
            method: String::new()
        };

        let mut headers: Vec<ClientHeader> = Vec::new();
        loop {
            match stream.read_line() { // XXX: strange unwrap like thing
                Ok(line) => cur_line = line,
                Err(_) => break
            }

            body = body + cur_line + "
";
            if cur_line.as_bytes() == CRLF {
                break;
            }

            if ii == 0u { // the Request-Line is always the first line
                req = parse_request_line(cur_line.as_slice());
                println!("{}", req.pathname);
                ii += 1;
            } else { // this should be  "normal" key-value type header
                headers.push(parse_normal_header(cur_line.as_slice()));
            }

        }

        println!("pathname {} method {}", req.pathname, req.method);

        if req.pathname.as_bytes() == b"/" {
            body = get_normal_body("../html/ws1.html");
        } else if req.pathname.as_bytes() == b"/ws1.js" {
            body = get_normal_body("../html/ws1.js");
        } else if req.pathname.as_bytes() == b"/ws" {
            ws_handshake(stream, headers);
            return;
        }

        let body_length = format!("Content-length: {}", body.len());

        stream.write(b"HTTP/1.1 200 OK\r\n").unwrap();
        stream.write(b"Content-type: text/html\r\n").unwrap();
        stream.write(b"X-header: from bytes\r\n").unwrap();
        stream.write(body_length.as_bytes()).unwrap();
        stream.write(b"\r\n\r\n").unwrap();
        stream.write(body.as_bytes()).unwrap();
        stream.flush().unwrap();

        println!("Done handling acceptor");
    }

    for stream in acceptor.incoming() {
        match stream {
            Err(e) => { println!("connection failed: {}", e) }
            Ok(stream) => spawn(proc() {
                let bs = BufferedStream::new(stream);
                handle_client(bs)
            })
        }
    }

    drop(acceptor);
}
<commit_msg>clean up<commit_after>extern crate "rust-crypto" as rust_crypto;
extern crate serialize;

use std::io::{TcpListener, TcpStream};
use std::io::{Acceptor, Listener};
use std::io::{File, BufferedStream};
use std::string::String;
use std::os;

use rust_crypto::sha1::Sha1;
use rust_crypto::digest::Digest;
use serialize::base64::{ToBase64, STANDARD};

pub const CR: u8 = b'\r';
pub const LF: u8 = b'\n';
pub const SP: u8 = b' ';
pub const COLON: u8 = b':';
pub const CRLF: &'static [u8] = &[CR,LF];

struct RequestedRoute {
    method: String,
    pathname: String
}

struct ClientHeader {
    key: String,
    value: String
}

enum Payload {
    Text(String),
    Empty
}

enum Opcode {
    TextOp = 0x1,
    BinaryOp = 0x2
}

struct Message {
    payload: Payload,
    opcode: Opcode
 }

impl Message {

    fn new (msg: &String) -> Message {

        let (payload, opcode) = match msg {
            String => (Payload::Text(msg.to_string()), Opcode::TextOp)
        };

        return Message{payload: payload, opcode: opcode};
    }

    fn send (self,
            mut stream: BufferedStream<TcpStream>,
            headers: Vec<ClientHeader>) {

        let msg = match self.payload {
            Payload::Text(ref s) => s.as_bytes(),
            Payload::Empty => unimplemented!()
        };

        let length = msg.len() as u8;

        println!("msg: {}, opcode: {}", msg, self.opcode as u8);

        stream.write_u8(0b1000_0000 | self.opcode as u8).unwrap();
        stream.write_u8(length).unwrap();
        stream.write(msg).unwrap();

    }
}

fn get_header_by_name (header: &[u8], headers: &Vec<ClientHeader>) -> String {
    let mut result = String::new();

    for h in headers.iter() {
        if h.key.as_bytes() == header {
            result = h.value.clone();
            break;
        }
    }

    return result;
}

fn parse_request_line (header: &str) -> RequestedRoute {
    let mut pathname = String::new();
    let mut method = String::new();
    let mut starting = true;
    let mut found_method = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];

        if starting && cbyte != SP {
            method.push_str(ch);
        }
        if starting && cbyte == SP {
            starting = false;
            found_method = true;
            continue;
        }
        if found_method && cbyte != SP {
            pathname.push_str(ch);
        }
        if found_method && cbyte == SP {
            break;
        }

    }

    return RequestedRoute {method: method, pathname: pathname};
}

fn parse_normal_header (header: &str) -> ClientHeader {
    let mut lhs = String::new(); // lhs of ':' in header
    let mut rhs = String::new(); // rhs of ':' in header
    let mut found_colon = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];
        if !found_colon && cbyte == COLON {
            found_colon = true;
            continue;
        }

        if cbyte == CR || cbyte == LF {
            continue;
        }

        if !found_colon {
            lhs.push_str(ch);
        } else {
            rhs.push_str(ch);
        }
    }

    return ClientHeader {
        key:lhs.trim().to_string(),
        value:rhs.trim().to_string()
    };
}

fn get_normal_body (path_on_disk: &str) -> String {
    let path = Path::new(path_on_disk);
    let display = path.display();
    let mut file = match File::open(&path) {
        Ok(f) => f,
        Err(err) => panic!("file error: {}", err)
    };

    let content = match file.read_to_end() {
        Ok(c) => c,
        Err(err) => panic!("{}", err)
    };

    let s = String::from_utf8(content).unwrap();
    return s;
}

fn sec_handshake (from_server: &[u8]) -> String {

    let guid = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

    let mut sha = Sha1::new();

    sha.input(from_server);
    sha.input(guid);
    let mut out = [0u8, ..20];
    sha.result(out.as_mut_slice());

    return out.to_base64(STANDARD);
}

fn ws_handshake (mut stream: BufferedStream<TcpStream>,
                 headers: Vec<ClientHeader>) {

    let header_sec_key = b"Sec-WebSocket-Key";

    let from_server = get_header_by_name(header_sec_key, &headers);
    let accept = sec_handshake(from_server.as_bytes());

    let sec_header = format!("Sec-WebSocket-Accept: {}\r\n", accept);

    stream.write(b"HTTP/1.1 101 Switching Protocols\r\n").unwrap();
    stream.write(b"Upgrade: websocket\r\n").unwrap();
    stream.write(b"Connection: Upgrade\r\n").unwrap();
    stream.write(b"Sec-WebSocket-Version: 13\r\n").unwrap();
    stream.write(b"Sec-WebSocket-Protocol: protocolOne\r\n").unwrap();
    stream.write(sec_header.as_bytes()).unwrap();
    stream.write(b"\r\n");

    ws_listen(stream, headers);
}

fn ws_listen(mut stream: BufferedStream<TcpStream>,
             headers: Vec<ClientHeader>) {
    let string_msg = "this is a message".to_string();
    let msg = box Message::new(&string_msg);
    msg.send(stream, headers);
}

fn main () {

    for argument in os::args().iter() {
        println!("arg: {}", argument);
    }

    let addr = "127.0.0.1:8099";

    let listener = TcpListener::bind(addr).unwrap();
    let mut acceptor = listener.listen().unwrap();
    println!("Listening on {}", addr);

    fn handle_client(mut stream: BufferedStream<TcpStream>) {
        let mut body: String = "
You sent it!
".to_string();

        let mut cur_line: String;
        let mut ii = 0u;
        let mut req: RequestedRoute = RequestedRoute {
            pathname: String::new(),
            method: String::new()
        };

        let mut headers: Vec<ClientHeader> = Vec::new();
        loop {
            match stream.read_line() { // XXX: strange unwrap like thing
                Ok(line) => cur_line = line,
                Err(_) => break
            }

            body = body + cur_line + "
";
            if cur_line.as_bytes() == CRLF {
                break;
            }

            if ii == 0u { // the Request-Line is always the first line
                req = parse_request_line(cur_line.as_slice());
                println!("{}", req.pathname);
                ii += 1;
            } else { // this should be  "normal" key-value type header
                headers.push(parse_normal_header(cur_line.as_slice()));
            }

        }

        println!("pathname {} method {}", req.pathname, req.method);

        if req.pathname.as_bytes() == b"/" {
            body = get_normal_body("../html/ws1.html");
        } else if req.pathname.as_bytes() == b"/ws1.js" {
            body = get_normal_body("../html/ws1.js");
        } else if req.pathname.as_bytes() == b"/ws" {
            ws_handshake(stream, headers);
            return;
        }

        let body_length = format!("Content-length: {}", body.len());

        stream.write(b"HTTP/1.1 200 OK\r\n").unwrap();
        stream.write(b"Content-type: text/html\r\n").unwrap();
        stream.write(b"X-header: from bytes\r\n").unwrap();
        stream.write(body_length.as_bytes()).unwrap();
        stream.write(b"\r\n\r\n").unwrap();
        stream.write(body.as_bytes()).unwrap();
        stream.flush().unwrap();

        println!("Done handling acceptor");
    }

    for stream in acceptor.incoming() {
        match stream {
            Err(e) => { println!("connection failed: {}", e) }
            Ok(stream) => spawn(proc() {
                let bs = BufferedStream::new(stream);
                handle_client(bs)
            })
        }
    }

    drop(acceptor);
}
<|endoftext|>


<commit_before>use std::ffi::NulError;
use std::str::Utf8Error;
use std::string::FromUtf8Error;
use fitserror::FitsError;

#[derive(Debug, PartialEq)]
pub enum Error {
    Fits(FitsError),
    Message(String),
    Null(NulError),
    Utf8(Utf8Error),
}

pub type Result<T> = ::std::result::Result<T, Error>;

impl ::std::convert::From<FitsError> for Error {
    fn from(error: FitsError) -> Self {
        Error::Fits(error)
    }
}

impl<'a> ::std::convert::From<&'a str> for Error {
    fn from(error: &'a str) -> Self {
        Error::Message(error.to_string())
    }
}

impl ::std::convert::From<NulError> for Error {
    fn from(error: NulError) -> Self {
        Error::Null(error)
    }
}

impl ::std::convert::From<FromUtf8Error> for Error {
    fn from(error: FromUtf8Error) -> Self {
        Error::Utf8(error.utf8_error())
    }
}

impl ::std::convert::From<Utf8Error> for Error {
    fn from(error: Utf8Error) -> Self {
        Error::Utf8(error)
    }
}

impl ::std::convert::From<Box<::std::error::Error>> for Error {
    fn from(error: Box<::std::error::Error>) -> Self {
        let description = error.description();
        let message = match error.cause() {
            Some(msg) => format!("Error: {} caused by {}", description, msg),
            None => format!("Error: {}", description),
        };
        Error::Message(message)
    }
}
<commit_msg>Add extra error implementations<commit_after>use std::ffi::NulError;
use std::str::Utf8Error;
use std::string::FromUtf8Error;
use fitserror::FitsError;

#[derive(Debug, PartialEq)]
pub enum Error {
    Fits(FitsError),
    Message(String),
    Null(NulError),
    Utf8(Utf8Error),
}

pub type Result<T> = ::std::result::Result<T, Error>;

impl ::std::convert::From<FitsError> for Error {
    fn from(error: FitsError) -> Self {
        Error::Fits(error)
    }
}

impl<'a> ::std::convert::From<&'a str> for Error {
    fn from(error: &'a str) -> Self {
        Error::Message(error.to_string())
    }
}

impl ::std::convert::From<NulError> for Error {
    fn from(error: NulError) -> Self {
        Error::Null(error)
    }
}

impl ::std::convert::From<FromUtf8Error> for Error {
    fn from(error: FromUtf8Error) -> Self {
        Error::Utf8(error.utf8_error())
    }
}

impl ::std::convert::From<Utf8Error> for Error {
    fn from(error: Utf8Error) -> Self {
        Error::Utf8(error)
    }
}

impl ::std::convert::From<Box<::std::error::Error>> for Error {
    fn from(error: Box<::std::error::Error>) -> Self {
        let description = error.description();
        let message = match error.cause() {
            Some(msg) => format!("Error: {} caused by {}", description, msg),
            None => format!("Error: {}", description),
        };
        Error::Message(message)
    }
}

impl ::std::fmt::Display for Error {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
        match *self {
            Error::Fits(ref e) => write!(f, "Fits error: {:?}", e),
            Error::Message(ref s) => write!(f, "Error: {}", s),
            Error::Null(ref e) => write!(f, "Error: {}", e),
            Error::Utf8(ref e) => write!(f, "Error: {}", e),
            _ => Ok(()),
        }
    }
}

impl ::std::error::Error for Error {
    fn description(&self) -> &str {
        "fitsio error"
    }
}
<|endoftext|>


<commit_before>// Rust on bare metal - Tested with Rust 1.0.0
// Uses Pure64 to get into a 64-bit SMP state
// rustc -O --crate-type lib -o kernel64.o --emit obj kernel64.rs
// ld -T app.ld -o kernel64.sys kernel64.o

// declare the address of the text buffer
const BUF_TEXT: *mut u16 = 0xb8000 as *mut u16;

#[no_mangle]
pub fn main() {
    clear_screen(Color::LightBlue);
    loop {
        // Loop forever
    }
}

#[derive(Copy,Clone)]
pub enum Color {
    Black       = 0,
    Blue        = 1,
    Green       = 2,
    Cyan        = 3,
    Red         = 4,
    Pink        = 5,
    Brown       = 6,
    LightGray   = 7,
    DarkGray    = 8,
    LightBlue   = 9,
    LightGreen  = 10,
    LightCyan   = 11,
    LightRed    = 12,
    LightPink   = 13,
    Yello       = 14,
    White       = 15,
}

fn clear_screen(background: Color) {
    for x in 0 .. 80*25 {
        unsafe {
            // one should probably use the `core` feature (which is only
            // available outside the stable release channel)
            std::ptr::write(BUF_TEXT.offset(x),(background as u16) << 12);
        }
    }
}
<commit_msg>now using the `core` library<commit_after>// Rust on bare metal - Tested with Rust 1.0.0
// Uses Pure64 to get into a 64-bit SMP state
// rustc -O --crate-type lib -o kernel64.o --emit obj kernel64.rs
// ld -T app.ld -o kernel64.sys kernel64.o

// do not include the standard library
#![feature(no_std)]
#![no_std]

// declare the address of the text buffer
const BUF_TEXT: *mut u16 = 0xb8000 as *mut u16;

#[no_mangle]
pub fn main() {
    clear_screen(Color::LightBlue);
    loop {
        // Loop forever
    }
}

#[derive(Copy,Clone)]
pub enum Color {
    Black       = 0,
    Blue        = 1,
    Green       = 2,
    Cyan        = 3,
    Red         = 4,
    Pink        = 5,
    Brown       = 6,
    LightGray   = 7,
    DarkGray    = 8,
    LightBlue   = 9,
    LightGreen  = 10,
    LightCyan   = 11,
    LightRed    = 12,
    LightPink   = 13,
    Yello       = 14,
    White       = 15,
}

fn clear_screen(background: Color) {
    for x in 0 .. 80*25 {
        unsafe {
			// use the `core` library because we disabled `std`
            core::ptr::write(BUF_TEXT.offset(x),(background as u16) << 12);
        }
    }
}
<|endoftext|>