<commit_before>//! Completion for associated items in a trait implementation.
//!
//! This module adds the completion items related to implementing associated
//! items within a `impl Trait for Struct` block. The current context node
//! must be within either a `FN_DEF`, `TYPE_ALIAS_DEF`, or `CONST_DEF` node
//! and an direct child of an `IMPL_DEF`.
//!
//! # Examples
//!
//! Considering the following trait `impl`:
//!
//! ```ignore
//! trait SomeTrait {
//!     fn foo();
//! }
//!
//! impl SomeTrait for () {
//!     fn f<|>
//! }
//! ```
//!
//! may result in the completion of the following method:
//!
//! ```ignore
//! # trait SomeTrait {
//! #    fn foo();
//! # }
//!
//! impl SomeTrait for () {
//!     fn foo() {}<|>
//! }
//! ```

use hir::{self, Docs, HasSource};
use ra_assists::utils::get_missing_impl_items;
use ra_syntax::{
    ast::{self, edit},
    AstNode, SyntaxKind, SyntaxNode, TextRange,
};
use ra_text_edit::TextEdit;

use crate::{
    completion::{
        CompletionContext, CompletionItem, CompletionItemKind, CompletionKind, Completions,
    },
    display::FunctionSignature,
};

pub(crate) fn complete_trait_impl(acc: &mut Completions, ctx: &CompletionContext) {
    let trigger = ctx.token.ancestors().find(|p| match p.kind() {
        SyntaxKind::FN_DEF
        | SyntaxKind::TYPE_ALIAS_DEF
        | SyntaxKind::CONST_DEF
        | SyntaxKind::BLOCK_EXPR => true,
        _ => false,
    });

    let impl_def = trigger
        .as_ref()
        .and_then(|node| node.parent())
        .and_then(|node| node.parent())
        .and_then(ast::ImplDef::cast);

    if let (Some(trigger), Some(impl_def)) = (trigger, impl_def) {
        match trigger.kind() {
            SyntaxKind::FN_DEF => {
                for missing_fn in get_missing_impl_items(&ctx.sema, &impl_def).iter().filter_map(
                    |item| match item {
                        hir::AssocItem::Function(fn_item) => Some(fn_item),
                        _ => None,
                    },
                ) {
                    add_function_impl(&trigger, acc, ctx, &missing_fn);
                }
            }

            SyntaxKind::TYPE_ALIAS_DEF => {
                for missing_fn in get_missing_impl_items(&ctx.sema, &impl_def).iter().filter_map(
                    |item| match item {
                        hir::AssocItem::TypeAlias(type_item) => Some(type_item),
                        _ => None,
                    },
                ) {
                    add_type_alias_impl(&trigger, acc, ctx, &missing_fn);
                }
            }

            SyntaxKind::CONST_DEF => {
                for missing_fn in get_missing_impl_items(&ctx.sema, &impl_def).iter().filter_map(
                    |item| match item {
                        hir::AssocItem::Const(const_item) => Some(const_item),
                        _ => None,
                    },
                ) {
                    add_const_impl(&trigger, acc, ctx, &missing_fn);
                }
            }

            _ => {}
        }
    }
}

fn add_function_impl(
    fn_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    func: &hir::Function,
) {
    let display = FunctionSignature::from_hir(ctx.db, *func);

    let fn_name = func.name(ctx.db).to_string();

    let label = if !func.params(ctx.db).is_empty() {
        format!("fn {}(..)", fn_name)
    } else {
        format!("fn {}()", fn_name)
    };

    let builder = CompletionItem::new(CompletionKind::Magic, ctx.source_range(), label)
        .lookup_by(fn_name)
        .set_documentation(func.docs(ctx.db));

    let completion_kind = if func.has_self_param(ctx.db) {
        CompletionItemKind::Method
    } else {
        CompletionItemKind::Function
    };

    let snippet = format!("{} {{}}", display);

    let range = TextRange::from_to(fn_def_node.text_range().start(), ctx.source_range().end());

    builder.text_edit(TextEdit::replace(range, snippet)).kind(completion_kind).add_to(acc);
}

fn add_type_alias_impl(
    type_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    type_alias: &hir::TypeAlias,
) {
    let alias_name = type_alias.name(ctx.db).to_string();

    let snippet = format!("type {} = ", alias_name);

    let range = TextRange::from_to(type_def_node.text_range().start(), ctx.source_range().end());

    CompletionItem::new(CompletionKind::Magic, ctx.source_range(), snippet.clone())
        .text_edit(TextEdit::replace(range, snippet))
        .lookup_by(alias_name)
        .kind(CompletionItemKind::TypeAlias)
        .set_documentation(type_alias.docs(ctx.db))
        .add_to(acc);
}

fn add_const_impl(
    const_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    const_: &hir::Const,
) {
    let const_name = const_.name(ctx.db).map(|n| n.to_string());

    if let Some(const_name) = const_name {
        let snippet = make_const_compl_syntax(&const_.source(ctx.db).value);

        let range =
            TextRange::from_to(const_def_node.text_range().start(), ctx.source_range().end());

        CompletionItem::new(CompletionKind::Magic, ctx.source_range(), snippet.clone())
            .text_edit(TextEdit::replace(range, snippet))
            .lookup_by(const_name)
            .kind(CompletionItemKind::Const)
            .set_documentation(const_.docs(ctx.db))
            .add_to(acc);
    }
}

fn make_const_compl_syntax(const_: &ast::ConstDef) -> String {
    let const_ = edit::strip_attrs_and_docs(const_);

    let const_start = const_.syntax().text_range().start();
    let const_end = const_.syntax().text_range().end();

    let start =
        const_.syntax().first_child_or_token().map_or(const_start, |f| f.text_range().start());

    let end = const_
        .syntax()
        .children_with_tokens()
        .find(|s| s.kind() == SyntaxKind::SEMI || s.kind() == SyntaxKind::EQ)
        .map_or(const_end, |f| f.text_range().start());

    let len = end - start;
    let range = TextRange::from_to(0.into(), len);

    let syntax = const_.syntax().text().slice(range).to_string();

    format!("{} = ", syntax.trim_end())
}

#[cfg(test)]
mod tests {
    use crate::completion::{do_completion, CompletionItem, CompletionKind};
    use insta::assert_debug_snapshot;

    fn complete(code: &str) -> Vec<CompletionItem> {
        do_completion(code, CompletionKind::Magic)
    }

    #[test]
    fn single_function() {
        let completions = complete(
            r"
            trait Test {
                fn foo();
            }

            struct T1;

            impl Test for T1 {
                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo()",
                source_range: [141; 142),
                delete: [138; 142),
                insert: "fn foo() {}",
                kind: Function,
                lookup: "foo",
            },
        ]
        "###);
    }

    #[test]
    fn hide_implemented_fn() {
        let completions = complete(
            r"
            trait Test {
                fn foo();
                fn foo_bar();
            }

            struct T1;

            impl Test for T1 {
                fn foo() {}

                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo_bar()",
                source_range: [200; 201),
                delete: [197; 201),
                insert: "fn foo_bar() {}",
                kind: Function,
                lookup: "foo_bar",
            },
        ]
        "###);
    }

    #[test]
    fn completes_only_on_top_level() {
        let completions = complete(
            r"
            trait Test {
                fn foo();

                fn foo_bar();
            }

            struct T1;

            impl Test for T1 {
                fn foo() {
                    <|>
                }
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"[]"###);
    }

    #[test]
    fn generic_fn() {
        let completions = complete(
            r"
            trait Test {
                fn foo<T>();
            }

            struct T1;

            impl Test for T1 {
                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo()",
                source_range: [144; 145),
                delete: [141; 145),
                insert: "fn foo<T>() {}",
                kind: Function,
                lookup: "foo",
            },
        ]
        "###);
    }

    #[test]
    fn generic_constrait_fn() {
        let completions = complete(
            r"
            trait Test {
                fn foo<T>() where T: Into<String>;
            }

            struct T1;

            impl Test for T1 {
                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo()",
                source_range: [166; 167),
                delete: [163; 167),
                insert: "fn foo<T>()\nwhere T: Into<String> {}",
                kind: Function,
                lookup: "foo",
            },
        ]
        "###);
    }

    #[test]
    fn associated_type() {
        let completions = complete(
            r"
            trait Test {
                type SomeType;
            }

            impl Test for () {
                type S<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "type SomeType = ",
                source_range: [124; 125),
                delete: [119; 125),
                insert: "type SomeType = ",
                kind: TypeAlias,
                lookup: "SomeType",
            },
        ]
        "###);
    }

    #[test]
    fn associated_const() {
        let completions = complete(
            r"
            trait Test {
                const SOME_CONST: u16;
            }

            impl Test for () {
                const S<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "const SOME_CONST: u16 = ",
                source_range: [133; 134),
                delete: [127; 134),
                insert: "const SOME_CONST: u16 = ",
                kind: Const,
                lookup: "SOME_CONST",
            },
        ]
        "###);
    }

    #[test]
    fn associated_const_with_default() {
        let completions = complete(
            r"
            trait Test {
                const SOME_CONST: u16 = 42;
            }

            impl Test for () {
                const S<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "const SOME_CONST: u16 = ",
                source_range: [138; 139),
                delete: [132; 139),
                insert: "const SOME_CONST: u16 = ",
                kind: Const,
                lookup: "SOME_CONST",
            },
        ]
        "###);
    }
}
<commit_msg>Next steps in assoc item completion #3183<commit_after>//! Completion for associated items in a trait implementation.
//!
//! This module adds the completion items related to implementing associated
//! items within a `impl Trait for Struct` block. The current context node
//! must be within either a `FN_DEF`, `TYPE_ALIAS_DEF`, or `CONST_DEF` node
//! and an direct child of an `IMPL_DEF`.
//!
//! # Examples
//!
//! Considering the following trait `impl`:
//!
//! ```ignore
//! trait SomeTrait {
//!     fn foo();
//! }
//!
//! impl SomeTrait for () {
//!     fn f<|>
//! }
//! ```
//!
//! may result in the completion of the following method:
//!
//! ```ignore
//! # trait SomeTrait {
//! #    fn foo();
//! # }
//!
//! impl SomeTrait for () {
//!     fn foo() {}<|>
//! }
//! ```

use hir::{self, Docs, HasSource};
use ra_assists::utils::get_missing_impl_items;
use ra_syntax::{
    ast::{self, edit},
    AstNode, SyntaxKind, SyntaxNode, TextRange,
};
use ra_text_edit::TextEdit;

use crate::{
    completion::{
        CompletionContext, CompletionItem, CompletionItemKind, CompletionKind, Completions,
    },
    display::FunctionSignature,
};

pub(crate) fn complete_trait_impl(acc: &mut Completions, ctx: &CompletionContext) {
    let mut tokens = ctx.token.ancestors();
    let trigger = tokens.find(|p| match p.kind() {
        SyntaxKind::FN_DEF
        | SyntaxKind::TYPE_ALIAS_DEF
        | SyntaxKind::CONST_DEF
        | SyntaxKind::NAME_REF
        | SyntaxKind::BLOCK_EXPR => true,
        _ => false,
    });

    let impl_def = tokens
        .find(|p| match p.kind() {
            SyntaxKind::IMPL_DEF => true,
            _ => false,
        })
        .and_then(|n| ast::ImplDef::cast(n));

    if let (Some(trigger), Some(impl_def)) = (trigger, impl_def) {
        match trigger.kind() {
            SyntaxKind::NAME_REF => {
                get_missing_impl_items(&ctx.sema, &impl_def).iter().for_each(|item| match item {
                    hir::AssocItem::Function(fn_item) => {
                        add_function_impl(&trigger, acc, ctx, &fn_item)
                    }
                    hir::AssocItem::TypeAlias(type_item) => {
                        add_type_alias_impl(&trigger, acc, ctx, &type_item)
                    }
                    hir::AssocItem::Const(const_item) => {
                        add_const_impl(&trigger, acc, ctx, &const_item)
                    }
                })
            }

            SyntaxKind::FN_DEF => {
                for missing_fn in get_missing_impl_items(&ctx.sema, &impl_def).iter().filter_map(
                    |item| match item {
                        hir::AssocItem::Function(fn_item) => Some(fn_item),
                        _ => None,
                    },
                ) {
                    add_function_impl(&trigger, acc, ctx, &missing_fn);
                }
            }

            SyntaxKind::TYPE_ALIAS_DEF => {
                for missing_fn in get_missing_impl_items(&ctx.sema, &impl_def).iter().filter_map(
                    |item| match item {
                        hir::AssocItem::TypeAlias(type_item) => Some(type_item),
                        _ => None,
                    },
                ) {
                    add_type_alias_impl(&trigger, acc, ctx, &missing_fn);
                }
            }

            SyntaxKind::CONST_DEF => {
                for missing_fn in get_missing_impl_items(&ctx.sema, &impl_def).iter().filter_map(
                    |item| match item {
                        hir::AssocItem::Const(const_item) => Some(const_item),
                        _ => None,
                    },
                ) {
                    add_const_impl(&trigger, acc, ctx, &missing_fn);
                }
            }

            _ => {}
        }
    }
}

fn add_function_impl(
    fn_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    func: &hir::Function,
) {
    let display = FunctionSignature::from_hir(ctx.db, *func);

    let fn_name = func.name(ctx.db).to_string();

    let label = if !func.params(ctx.db).is_empty() {
        format!("fn {}(..)", fn_name)
    } else {
        format!("fn {}()", fn_name)
    };

    let builder = CompletionItem::new(CompletionKind::Magic, ctx.source_range(), label)
        .lookup_by(fn_name)
        .set_documentation(func.docs(ctx.db));

    let completion_kind = if func.has_self_param(ctx.db) {
        CompletionItemKind::Method
    } else {
        CompletionItemKind::Function
    };

    let snippet = format!("{} {{}}", display);

    let range = TextRange::from_to(fn_def_node.text_range().start(), ctx.source_range().end());

    builder.text_edit(TextEdit::replace(range, snippet)).kind(completion_kind).add_to(acc);
}

fn add_type_alias_impl(
    type_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    type_alias: &hir::TypeAlias,
) {
    let alias_name = type_alias.name(ctx.db).to_string();

    let snippet = format!("type {} = ", alias_name);

    let range = TextRange::from_to(type_def_node.text_range().start(), ctx.source_range().end());

    CompletionItem::new(CompletionKind::Magic, ctx.source_range(), snippet.clone())
        .text_edit(TextEdit::replace(range, snippet))
        .lookup_by(alias_name)
        .kind(CompletionItemKind::TypeAlias)
        .set_documentation(type_alias.docs(ctx.db))
        .add_to(acc);
}

fn add_const_impl(
    const_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    const_: &hir::Const,
) {
    let const_name = const_.name(ctx.db).map(|n| n.to_string());

    if let Some(const_name) = const_name {
        let snippet = make_const_compl_syntax(&const_.source(ctx.db).value);

        let range =
            TextRange::from_to(const_def_node.text_range().start(), ctx.source_range().end());

        CompletionItem::new(CompletionKind::Magic, ctx.source_range(), snippet.clone())
            .text_edit(TextEdit::replace(range, snippet))
            .lookup_by(const_name)
            .kind(CompletionItemKind::Const)
            .set_documentation(const_.docs(ctx.db))
            .add_to(acc);
    }
}

fn make_const_compl_syntax(const_: &ast::ConstDef) -> String {
    let const_ = edit::strip_attrs_and_docs(const_);

    let const_start = const_.syntax().text_range().start();
    let const_end = const_.syntax().text_range().end();

    let start =
        const_.syntax().first_child_or_token().map_or(const_start, |f| f.text_range().start());

    let end = const_
        .syntax()
        .children_with_tokens()
        .find(|s| s.kind() == SyntaxKind::SEMI || s.kind() == SyntaxKind::EQ)
        .map_or(const_end, |f| f.text_range().start());

    let len = end - start;
    let range = TextRange::from_to(0.into(), len);

    let syntax = const_.syntax().text().slice(range).to_string();

    format!("{} = ", syntax.trim_end())
}

#[cfg(test)]
mod tests {
    use crate::completion::{do_completion, CompletionItem, CompletionKind};
    use insta::assert_debug_snapshot;

    fn complete(code: &str) -> Vec<CompletionItem> {
        do_completion(code, CompletionKind::Magic)
    }

    #[test]
    fn name_ref_function_type_const() {
        let completions = complete(
            r"
            trait Test {
                type TestType;
                const TEST_CONST: u16;
                fn test();
            }

            struct T1;

            impl Test for T1 {
                t<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "const TEST_CONST: u16 = ",
                source_range: [209; 210),
                delete: [209; 210),
                insert: "const TEST_CONST: u16 = ",
                kind: Const,
                lookup: "TEST_CONST",
            },
            CompletionItem {
                label: "fn test()",
                source_range: [209; 210),
                delete: [209; 210),
                insert: "fn test() {}",
                kind: Function,
                lookup: "test",
            },
            CompletionItem {
                label: "type TestType = ",
                source_range: [209; 210),
                delete: [209; 210),
                insert: "type TestType = ",
                kind: TypeAlias,
                lookup: "TestType",
            },
        ]
        "###);
    }

    #[test]
    fn name_ref_single_function() {
        let completions = complete(
            r"
            trait Test {
                fn test();
            }

            struct T1;

            impl Test for T1 {
                t<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn test()",
                source_range: [139; 140),
                delete: [139; 140),
                insert: "fn test() {}",
                kind: Function,
                lookup: "test",
            },
        ]
        "###);
    }

    #[test]
    fn single_function() {
        let completions = complete(
            r"
            trait Test {
                fn foo();
            }

            struct T1;

            impl Test for T1 {
                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo()",
                source_range: [141; 142),
                delete: [138; 142),
                insert: "fn foo() {}",
                kind: Function,
                lookup: "foo",
            },
        ]
        "###);
    }

    #[test]
    fn hide_implemented_fn() {
        let completions = complete(
            r"
            trait Test {
                fn foo();
                fn foo_bar();
            }

            struct T1;

            impl Test for T1 {
                fn foo() {}

                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo_bar()",
                source_range: [200; 201),
                delete: [197; 201),
                insert: "fn foo_bar() {}",
                kind: Function,
                lookup: "foo_bar",
            },
        ]
        "###);
    }

    #[test]
    fn completes_only_on_top_level() {
        let completions = complete(
            r"
            trait Test {
                fn foo();

                fn foo_bar();
            }

            struct T1;

            impl Test for T1 {
                fn foo() {
                    <|>
                }
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"[]"###);
    }

    #[test]
    fn generic_fn() {
        let completions = complete(
            r"
            trait Test {
                fn foo<T>();
            }

            struct T1;

            impl Test for T1 {
                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo()",
                source_range: [144; 145),
                delete: [141; 145),
                insert: "fn foo<T>() {}",
                kind: Function,
                lookup: "foo",
            },
        ]
        "###);
    }

    #[test]
    fn generic_constrait_fn() {
        let completions = complete(
            r"
            trait Test {
                fn foo<T>() where T: Into<String>;
            }

            struct T1;

            impl Test for T1 {
                fn f<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "fn foo()",
                source_range: [166; 167),
                delete: [163; 167),
                insert: "fn foo<T>()\nwhere T: Into<String> {}",
                kind: Function,
                lookup: "foo",
            },
        ]
        "###);
    }

    #[test]
    fn associated_type() {
        let completions = complete(
            r"
            trait Test {
                type SomeType;
            }

            impl Test for () {
                type S<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "type SomeType = ",
                source_range: [124; 125),
                delete: [119; 125),
                insert: "type SomeType = ",
                kind: TypeAlias,
                lookup: "SomeType",
            },
        ]
        "###);
    }

    #[test]
    fn associated_const() {
        let completions = complete(
            r"
            trait Test {
                const SOME_CONST: u16;
            }

            impl Test for () {
                const S<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "const SOME_CONST: u16 = ",
                source_range: [133; 134),
                delete: [127; 134),
                insert: "const SOME_CONST: u16 = ",
                kind: Const,
                lookup: "SOME_CONST",
            },
        ]
        "###);
    }

    #[test]
    fn associated_const_with_default() {
        let completions = complete(
            r"
            trait Test {
                const SOME_CONST: u16 = 42;
            }

            impl Test for () {
                const S<|>
            }
            ",
        );
        assert_debug_snapshot!(completions, @r###"
        [
            CompletionItem {
                label: "const SOME_CONST: u16 = ",
                source_range: [138; 139),
                delete: [132; 139),
                insert: "const SOME_CONST: u16 = ",
                kind: Const,
                lookup: "SOME_CONST",
            },
        ]
        "###);
    }
}
<|endoftext|>


<commit_before>//! Track which values are live in an EBB with instruction granularity.
//!
//! The `LiveValueTracker` keeps track of the set of live SSA values at each instruction in an EBB.
//! The sets of live values are computed on the fly as the tracker is moved from instruction to
//! instruction, starting at the EBB header.

use dominator_tree::DominatorTree;
use entity_list::{EntityList, ListPool};
use ir::instructions::BranchInfo;
use ir::{Inst, Ebb, Value, DataFlowGraph, ProgramOrder, ExpandedProgramPoint};
use partition_slice::partition_slice;
use regalloc::liveness::Liveness;

use std::collections::HashMap;

type ValueList = EntityList<Value>;

/// Compute and track live values throughout an EBB.
pub struct LiveValueTracker {
    /// The set of values that are live at the current program point.
    live: LiveValueVec,

    /// Saved set of live values for every jump and branch that can potentially be an immediate
    /// dominator of an EBB.
    ///
    /// This is the set of values that are live *before* the branch.
    idom_sets: HashMap<Inst, ValueList>,

    /// Memory pool for the live sets.
    idom_pool: ListPool<Value>,
}

/// Information about a value that is live at the current program point.
pub struct LiveValue {
    /// The live value.
    pub value: Value,

    /// The local ending point of the live range in the current EBB, as returned by
    /// `LiveRange::def_local_end()` or `LiveRange::livein_local_end()`.
    pub endpoint: Inst,
}

struct LiveValueVec {
    /// The set of values that are live at the current program point.
    values: Vec<LiveValue>,

    /// How many values at the front of `values` are known to be live after `inst`?
    ///
    /// This is used to pass a much smaller slice to `partition_slice` when its called a second
    /// time for the same instruction.
    live_prefix: Option<(Inst, usize)>,
}

impl LiveValueVec {
    fn new() -> LiveValueVec {
        LiveValueVec {
            values: Vec::new(),
            live_prefix: None,
        }
    }

    /// Add a new live value to `values`.
    fn push(&mut self, value: Value, endpoint: Inst) {
        self.values.push(LiveValue {
            value: value,
            endpoint: endpoint,
        });
    }

    /// Remove all elements.
    fn clear(&mut self) {
        self.values.clear();
        self.live_prefix = None;
    }

    /// Make sure that the values killed by `next_inst` are moved to the end of the `values`
    /// vector.
    ///
    /// Returns the number of values that will be live after `next_inst`.
    fn live_after(&mut self, next_inst: Inst) -> usize {
        // How many values at the front of the vector are already known to survive `next_inst`?
        // We don't need to pass this prefix to `partition_slice()`
        let keep = match self.live_prefix {
            Some((i, prefix)) if i == next_inst => prefix,
            _ => 0,
        };

        // Move the remaining surviving values to the front partition of the vector.
        let prefix = keep + partition_slice(&mut self.values[keep..], |v| v.endpoint != next_inst);

        // Remember the new prefix length in case we get called again for the same `next_inst`.
        self.live_prefix = Some((next_inst, prefix));
        prefix
    }

    /// Remove the values killed by `next_inst`.
    fn remove_kill_values(&mut self, next_inst: Inst) {
        let keep = self.live_after(next_inst);
        self.values.truncate(keep);
    }
}

impl LiveValueTracker {
    /// Create a new blank tracker.
    pub fn new() -> LiveValueTracker {
        LiveValueTracker {
            live: LiveValueVec::new(),
            idom_sets: HashMap::new(),
            idom_pool: ListPool::new(),
        }
    }

    /// Clear all cached information.
    pub fn clear(&mut self) {
        self.live.clear();
        self.idom_sets.clear();
        self.idom_pool.clear();
    }

    /// Get the set of currently live values.
    ///
    /// Between calls to `process_inst()` and `drop_dead()`, this includes both values killed and
    /// defined by the current instruction.
    pub fn live(&self) -> &[LiveValue] {
        &self.live.values
    }

    /// Move the current position to the top of `ebb`.
    ///
    /// This depends on the stored live value set at `ebb`'s immediate dominator, so that must have
    /// been visited first.
    pub fn ebb_top<PO: ProgramOrder>(&mut self,
                                     ebb: Ebb,
                                     dfg: &DataFlowGraph,
                                     liveness: &Liveness,
                                     program_order: &PO,
                                     domtree: &DominatorTree) {
        // Start over, compute the set of live values at the top of the EBB from two sources:
        //
        // 1. Values that were live before `ebb`'s immediate dominator, filtered for those that are
        //    actually live-in.
        // 2. Arguments to `ebb` that are not dead.
        //
        self.live.clear();

        // Compute the live-in values. Start by filtering the set of values that were live before
        // the immediate dominator. Just use the empty set if there's no immediate dominator (i.e.,
        // the entry block or an unreachable block).
        if let Some(idom) = domtree.idom(ebb) {
            // If the immediate dominator exits, we must have a stored list for it. This is a
            // requirement to the order EBBs are visited: All dominators must have been processed
            // before the current EBB.
            let idom_live_list =
                self.idom_sets.get(&idom).expect("No stored live set for dominator");
            // Get just the values that are live-in to `ebb`.
            for &value in idom_live_list.as_slice(&self.idom_pool) {
                let lr = liveness.get(value)
                    .expect("Immediate dominator value has no live range");

                // Check if this value is live-in here.
                if let Some(endpoint) = lr.livein_local_end(ebb, program_order) {
                    self.live.push(value, endpoint);
                }
            }
        }

        // Now add all the live arguments to `ebb`.
        for value in dfg.ebb_args(ebb) {
            let lr = liveness.get(value).expect("EBB argument value has no live range");
            assert_eq!(lr.def(), ebb.into());
            match lr.def_local_end().into() {
                ExpandedProgramPoint::Inst(endpoint) => {
                    self.live.push(value, endpoint);
                }
                ExpandedProgramPoint::Ebb(local_ebb) => {
                    // This is a dead EBB argument which is not even live into the first
                    // instruction in the EBB. We can ignore it.
                    assert_eq!(local_ebb,
                               ebb,
                               "EBB argument live range ends at wrong EBB header");
                }
            }
        }
    }

    /// Prepare to move past `inst`.
    ///
    /// Determine the set of already live values that are killed by `inst`, and add the new defined
    /// values to the tracked set.
    ///
    /// Returns `(kills, defs)` as a pair of slices. The `defs` slice is guaranteed to be in the
    /// same order as `inst`'s results, and includes dead defines. The order of `kills` is
    /// arbitrary.
    ///
    /// The `drop_dead()` method must be called next to actually remove the dead values from the
    /// tracked set after the two returned slices are no longer needed.
    pub fn process_inst(&mut self,
                        inst: Inst,
                        dfg: &DataFlowGraph,
                        liveness: &Liveness)
                        -> (&[LiveValue], &[LiveValue]) {
        // Save a copy of the live values before any branches or jumps that could be somebody's
        // immediate dominator.
        match dfg[inst].analyze_branch() {
            BranchInfo::NotABranch => {}
            _ => self.save_idom_live_set(inst),
        }

        // Move killed values to the end of the vector.
        // Don't remove them yet, `drop_dead()` will do that.
        let first_kill = self.live.live_after(inst);

        // Add the values defined by `inst`.
        let first_def = self.live.values.len();
        for value in dfg.inst_results(inst) {
            let lr = liveness.get(value).expect("Instruction result has no live range");
            assert_eq!(lr.def(), inst.into());
            match lr.def_local_end().into() {
                ExpandedProgramPoint::Inst(endpoint) => {
                    self.live.push(value, endpoint);
                }
                ExpandedProgramPoint::Ebb(ebb) => {
                    panic!("Instruction result live range can't end at {}", ebb);
                }
            }
        }

        (&self.live.values[first_kill..first_def], &self.live.values[first_def..])
    }

    /// Drop the values that are now dead after moving past `inst`.
    ///
    /// This removes both live values that were killed by `inst` and dead defines on `inst` itself.
    ///
    /// This must be called after `process_inst(inst)` and before proceeding to the next
    /// instruction.
    pub fn drop_dead(&mut self, inst: Inst) {
        // Remove both live values that were killed by `inst` and dead defines from `inst`.
        self.live.remove_kill_values(inst);
    }

    /// Save the current set of live values so it is associated with `idom`.
    fn save_idom_live_set(&mut self, idom: Inst) {
        let values = self.live.values.iter().map(|lv| lv.value);
        let pool = &mut self.idom_pool;
        // If there already is a set saved for `idom`, just keep it.
        self.idom_sets.entry(idom).or_insert_with(|| {
            let mut list = ValueList::default();
            list.extend(values, pool);
            list
        });
    }
}
<commit_msg>Return slices of live-ins and arguments from ebb_top().<commit_after>//! Track which values are live in an EBB with instruction granularity.
//!
//! The `LiveValueTracker` keeps track of the set of live SSA values at each instruction in an EBB.
//! The sets of live values are computed on the fly as the tracker is moved from instruction to
//! instruction, starting at the EBB header.

use dominator_tree::DominatorTree;
use entity_list::{EntityList, ListPool};
use ir::instructions::BranchInfo;
use ir::{Inst, Ebb, Value, DataFlowGraph, ProgramOrder, ExpandedProgramPoint};
use partition_slice::partition_slice;
use regalloc::liveness::Liveness;

use std::collections::HashMap;

type ValueList = EntityList<Value>;

/// Compute and track live values throughout an EBB.
pub struct LiveValueTracker {
    /// The set of values that are live at the current program point.
    live: LiveValueVec,

    /// Saved set of live values for every jump and branch that can potentially be an immediate
    /// dominator of an EBB.
    ///
    /// This is the set of values that are live *before* the branch.
    idom_sets: HashMap<Inst, ValueList>,

    /// Memory pool for the live sets.
    idom_pool: ListPool<Value>,
}

/// Information about a value that is live at the current program point.
pub struct LiveValue {
    /// The live value.
    pub value: Value,

    /// The local ending point of the live range in the current EBB, as returned by
    /// `LiveRange::def_local_end()` or `LiveRange::livein_local_end()`.
    pub endpoint: Inst,
}

struct LiveValueVec {
    /// The set of values that are live at the current program point.
    values: Vec<LiveValue>,

    /// How many values at the front of `values` are known to be live after `inst`?
    ///
    /// This is used to pass a much smaller slice to `partition_slice` when its called a second
    /// time for the same instruction.
    live_prefix: Option<(Inst, usize)>,
}

impl LiveValueVec {
    fn new() -> LiveValueVec {
        LiveValueVec {
            values: Vec::new(),
            live_prefix: None,
        }
    }

    /// Add a new live value to `values`.
    fn push(&mut self, value: Value, endpoint: Inst) {
        self.values.push(LiveValue {
            value: value,
            endpoint: endpoint,
        });
    }

    /// Remove all elements.
    fn clear(&mut self) {
        self.values.clear();
        self.live_prefix = None;
    }

    /// Make sure that the values killed by `next_inst` are moved to the end of the `values`
    /// vector.
    ///
    /// Returns the number of values that will be live after `next_inst`.
    fn live_after(&mut self, next_inst: Inst) -> usize {
        // How many values at the front of the vector are already known to survive `next_inst`?
        // We don't need to pass this prefix to `partition_slice()`
        let keep = match self.live_prefix {
            Some((i, prefix)) if i == next_inst => prefix,
            _ => 0,
        };

        // Move the remaining surviving values to the front partition of the vector.
        let prefix = keep + partition_slice(&mut self.values[keep..], |v| v.endpoint != next_inst);

        // Remember the new prefix length in case we get called again for the same `next_inst`.
        self.live_prefix = Some((next_inst, prefix));
        prefix
    }

    /// Remove the values killed by `next_inst`.
    fn remove_kill_values(&mut self, next_inst: Inst) {
        let keep = self.live_after(next_inst);
        self.values.truncate(keep);
    }
}

impl LiveValueTracker {
    /// Create a new blank tracker.
    pub fn new() -> LiveValueTracker {
        LiveValueTracker {
            live: LiveValueVec::new(),
            idom_sets: HashMap::new(),
            idom_pool: ListPool::new(),
        }
    }

    /// Clear all cached information.
    pub fn clear(&mut self) {
        self.live.clear();
        self.idom_sets.clear();
        self.idom_pool.clear();
    }

    /// Get the set of currently live values.
    ///
    /// Between calls to `process_inst()` and `drop_dead()`, this includes both values killed and
    /// defined by the current instruction.
    pub fn live(&self) -> &[LiveValue] {
        &self.live.values
    }

    /// Move the current position to the top of `ebb`.
    ///
    /// This depends on the stored live value set at `ebb`'s immediate dominator, so that must have
    /// been visited first.
    ///
    /// Returns `(liveins, args)` as a pair or slices. The first slice is the set of live-in values
    /// from the immediate dominator. The second slice is the set of `ebb` arguments that are live.
    /// Dead arguments with no uses are ignored and not added to the set.
    pub fn ebb_top<PO: ProgramOrder>(&mut self,
                                     ebb: Ebb,
                                     dfg: &DataFlowGraph,
                                     liveness: &Liveness,
                                     program_order: &PO,
                                     domtree: &DominatorTree)
                                     -> (&[LiveValue], &[LiveValue]) {
        // Start over, compute the set of live values at the top of the EBB from two sources:
        //
        // 1. Values that were live before `ebb`'s immediate dominator, filtered for those that are
        //    actually live-in.
        // 2. Arguments to `ebb` that are not dead.
        //
        self.live.clear();

        // Compute the live-in values. Start by filtering the set of values that were live before
        // the immediate dominator. Just use the empty set if there's no immediate dominator (i.e.,
        // the entry block or an unreachable block).
        if let Some(idom) = domtree.idom(ebb) {
            // If the immediate dominator exits, we must have a stored list for it. This is a
            // requirement to the order EBBs are visited: All dominators must have been processed
            // before the current EBB.
            let idom_live_list =
                self.idom_sets.get(&idom).expect("No stored live set for dominator");
            // Get just the values that are live-in to `ebb`.
            for &value in idom_live_list.as_slice(&self.idom_pool) {
                let lr = liveness.get(value)
                    .expect("Immediate dominator value has no live range");

                // Check if this value is live-in here.
                if let Some(endpoint) = lr.livein_local_end(ebb, program_order) {
                    self.live.push(value, endpoint);
                }
            }
        }

        // Now add all the live arguments to `ebb`.
        let first_arg = self.live.values.len();
        for value in dfg.ebb_args(ebb) {
            let lr = liveness.get(value).expect("EBB argument value has no live range");
            assert_eq!(lr.def(), ebb.into());
            match lr.def_local_end().into() {
                ExpandedProgramPoint::Inst(endpoint) => {
                    self.live.push(value, endpoint);
                }
                ExpandedProgramPoint::Ebb(local_ebb) => {
                    // This is a dead EBB argument which is not even live into the first
                    // instruction in the EBB. We can ignore it.
                    assert_eq!(local_ebb,
                               ebb,
                               "EBB argument live range ends at wrong EBB header");
                }
            }
        }

        self.live.values.split_at(first_arg)
    }

    /// Prepare to move past `inst`.
    ///
    /// Determine the set of already live values that are killed by `inst`, and add the new defined
    /// values to the tracked set.
    ///
    /// Returns `(kills, defs)` as a pair of slices. The `defs` slice is guaranteed to be in the
    /// same order as `inst`'s results, and includes dead defines. The order of `kills` is
    /// arbitrary.
    ///
    /// The `drop_dead()` method must be called next to actually remove the dead values from the
    /// tracked set after the two returned slices are no longer needed.
    pub fn process_inst(&mut self,
                        inst: Inst,
                        dfg: &DataFlowGraph,
                        liveness: &Liveness)
                        -> (&[LiveValue], &[LiveValue]) {
        // Save a copy of the live values before any branches or jumps that could be somebody's
        // immediate dominator.
        match dfg[inst].analyze_branch() {
            BranchInfo::NotABranch => {}
            _ => self.save_idom_live_set(inst),
        }

        // Move killed values to the end of the vector.
        // Don't remove them yet, `drop_dead()` will do that.
        let first_kill = self.live.live_after(inst);

        // Add the values defined by `inst`.
        let first_def = self.live.values.len();
        for value in dfg.inst_results(inst) {
            let lr = liveness.get(value).expect("Instruction result has no live range");
            assert_eq!(lr.def(), inst.into());
            match lr.def_local_end().into() {
                ExpandedProgramPoint::Inst(endpoint) => {
                    self.live.push(value, endpoint);
                }
                ExpandedProgramPoint::Ebb(ebb) => {
                    panic!("Instruction result live range can't end at {}", ebb);
                }
            }
        }

        (&self.live.values[first_kill..first_def], &self.live.values[first_def..])
    }

    /// Drop the values that are now dead after moving past `inst`.
    ///
    /// This removes both live values that were killed by `inst` and dead defines on `inst` itself.
    ///
    /// This must be called after `process_inst(inst)` and before proceeding to the next
    /// instruction.
    pub fn drop_dead(&mut self, inst: Inst) {
        // Remove both live values that were killed by `inst` and dead defines from `inst`.
        self.live.remove_kill_values(inst);
    }

    /// Save the current set of live values so it is associated with `idom`.
    fn save_idom_live_set(&mut self, idom: Inst) {
        let values = self.live.values.iter().map(|lv| lv.value);
        let pool = &mut self.idom_pool;
        // If there already is a set saved for `idom`, just keep it.
        self.idom_sets.entry(idom).or_insert_with(|| {
            let mut list = ValueList::default();
            list.extend(values, pool);
            list
        });
    }
}
<|endoftext|>


<commit_before>// Copyright (c) <2015> <lummax>
// Licensed under MIT (http://opensource.org/licenses/MIT)

use std::{ptr, os};
use std::collections::RingBuf;

use spaces::immix_space::block_info::BlockInfo;

use constants::{BLOCK_SIZE, HEAP_SIZE};
use gc_object::GCObjectRef;

pub struct BlockAllocator {
    mmap: os::MemoryMap,
    data: *mut u8,
    data_bound: *mut u8,
    free_blocks: Vec<*mut BlockInfo>,
}

impl BlockAllocator {
    pub fn new() -> BlockAllocator {
        let mmap = os::MemoryMap::new(HEAP_SIZE + BLOCK_SIZE,
                                      &[os::MapOption::MapReadable,
                                        os::MapOption::MapWritable]).unwrap();
        let data = unsafe{ mmap.data().offset((BLOCK_SIZE - (mmap.data() as usize) % BLOCK_SIZE) as isize) };
        let data_bound = unsafe{ mmap.data().offset(mmap.len() as isize) };
        debug!("Allocated heap {:p} of size {}, usable range: {:p} - {:p} (size {}, {} blocks)",
                mmap.data(), mmap.len(), data, data_bound,
                (data_bound as usize) - (data  as usize),
                ((data_bound as usize) - (data  as usize)) / BLOCK_SIZE);
        debug_assert!((data as usize) % BLOCK_SIZE == 0,
            "Allocated mmap {:p} is not aligned (offset {})",
            data, (data as usize) % BLOCK_SIZE);
        return BlockAllocator {
            mmap: mmap,
            data: data,
            data_bound: data_bound,
            free_blocks: Vec::with_capacity(HEAP_SIZE / BLOCK_SIZE),
        };
    }

    pub fn get_block(&mut self) -> Option<*mut BlockInfo> {
        return self.free_blocks.pop().or_else(|| self.build_next_block());
    }

    pub fn return_blocks(&mut self, blocks: RingBuf<*mut BlockInfo>) {
        self.free_blocks.extend(blocks.into_iter());
    }

    pub fn total_blocks(&self) -> usize {
        return HEAP_SIZE / BLOCK_SIZE;
    }

    pub fn available_blocks(&self) -> usize {
        return (((self.data_bound as usize) - (self.data as usize)) % BLOCK_SIZE)
            + self.free_blocks.len();
    }

    pub fn is_in_space(&self, object: GCObjectRef) -> bool {
        return self.mmap.data() < (object as *mut u8)
            && (object as *mut u8) < self.data_bound;
    }
}

impl BlockAllocator{
    fn build_next_block(&mut self) -> Option<*mut BlockInfo> {
        let block = unsafe{ self.data.offset(BLOCK_SIZE as isize) };
        if block < self.data_bound {
            self.data = block;
            debug_assert!((block as usize) % BLOCK_SIZE == 0,
                "Allocated block {:p} is not aligned (offset {})",
                block, (block as usize) % BLOCK_SIZE);
            unsafe{ ptr::write(block as *mut BlockInfo, BlockInfo::new()); }
            return Some(block as *mut BlockInfo);
        }
        return None;
    }
}
<commit_msg>Add rustdoc for block_allocator.rs<commit_after>// Copyright (c) <2015> <lummax>
// Licensed under MIT (http://opensource.org/licenses/MIT)

use std::{ptr, os};
use std::collections::RingBuf;

use spaces::immix_space::block_info::BlockInfo;

use constants::{BLOCK_SIZE, HEAP_SIZE};
use gc_object::GCObjectRef;

/// The `BlockAllocator` is the global resource for blocks for the immix
/// space.
///
/// On initialization it will allocate a memory map of `HEAP_SIZE` and align
/// it to `BLOCK_SIZE`. During normal runtime it will allocate blocks on the
/// fly from this memory map and store returned blocks in a list.
///
/// Blocks from this `BlockAllocator` are always aligned to `BLOCK_SIZE`.
///
/// The list of returned free blocks is a stack. The `BlockAllocator` will
/// first exhaust the returned free blocks and then fall back to allocating
/// new blocks from the memory map. This means it will return recently
/// returned blocks first.
pub struct BlockAllocator {
    /// The memory map of `HEAP_SIZE`.
    mmap: os::MemoryMap,

    /// The pointer to the last allocated block.
    data: *mut u8,

    /// The upper bound of the memory map as a pointer.
    data_bound: *mut u8,

    /// A list of returned (free) blocks.
    free_blocks: Vec<*mut BlockInfo>,
}

impl BlockAllocator {
    /// Create a new `BlockAllocator`.
    ///
    /// This will `panic` if no memory map of size `HEAP_SIZE` can be
    /// allocared.
    pub fn new() -> BlockAllocator {
        let mmap = os::MemoryMap::new(HEAP_SIZE + BLOCK_SIZE,
                                      &[os::MapOption::MapReadable,
                                        os::MapOption::MapWritable]).unwrap();
        let data = unsafe{ mmap.data().offset((BLOCK_SIZE - (mmap.data() as usize) % BLOCK_SIZE) as isize) };
        let data_bound = unsafe{ mmap.data().offset(mmap.len() as isize) };
        debug!("Allocated heap {:p} of size {}, usable range: {:p} - {:p} (size {}, {} blocks)",
                mmap.data(), mmap.len(), data, data_bound,
                (data_bound as usize) - (data  as usize),
                ((data_bound as usize) - (data  as usize)) / BLOCK_SIZE);
        debug_assert!((data as usize) % BLOCK_SIZE == 0,
            "Allocated mmap {:p} is not aligned (offset {})",
            data, (data as usize) % BLOCK_SIZE);
        return BlockAllocator {
            mmap: mmap,
            data: data,
            data_bound: data_bound,
            free_blocks: Vec::with_capacity(HEAP_SIZE / BLOCK_SIZE),
        };
    }

    /// Get a new block aligned to `BLOCK_SIZE`.
    pub fn get_block(&mut self) -> Option<*mut BlockInfo> {
        return self.free_blocks.pop().or_else(|| self.build_next_block());
    }

    /// Return a collection of blocks.
    pub fn return_blocks(&mut self, blocks: RingBuf<*mut BlockInfo>) {
        self.free_blocks.extend(blocks.into_iter());
    }

    /// Return the total number of possible blocks.
    pub fn total_blocks(&self) -> usize {
        return HEAP_SIZE / BLOCK_SIZE;
    }

    /// Return the number of unallocated blocks.
    pub fn available_blocks(&self) -> usize {
        return (((self.data_bound as usize) - (self.data as usize)) % BLOCK_SIZE)
            + self.free_blocks.len();
    }

    /// Return if an address is within the bounds of the memory map.
    pub fn is_in_space(&self, object: GCObjectRef) -> bool {
        return self.mmap.data() < (object as *mut u8)
            && (object as *mut u8) < self.data_bound;
    }
}

impl BlockAllocator {
    /// Build a new block from the memory map.
    ///
    /// Returns `None` if the memory map is exhausted.
    fn build_next_block(&mut self) -> Option<*mut BlockInfo> {
        let block = unsafe{ self.data.offset(BLOCK_SIZE as isize) };
        if block < self.data_bound {
            self.data = block;
            debug_assert!((block as usize) % BLOCK_SIZE == 0,
                "Allocated block {:p} is not aligned (offset {})",
                block, (block as usize) % BLOCK_SIZE);
            unsafe{ ptr::write(block as *mut BlockInfo, BlockInfo::new()); }
            return Some(block as *mut BlockInfo);
        }
        return None;
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2019, Cloudflare, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use crate::octets;

use super::Error;
use super::Result;

use crate::h3::Header;

use super::INDEXED;
use super::INDEXED_WITH_POST_BASE;
use super::LITERAL;
use super::LITERAL_WITH_NAME_REF;

#[derive(Clone, Copy, Debug, PartialEq)]
enum Representation {
    Indexed,
    IndexedWithPostBase,
    Literal,
    LiteralWithNameRef,
    LiteralWithPostBase,
}

impl Representation {
    pub fn from_byte(b: u8) -> Representation {
        if b & INDEXED == INDEXED {
            return Representation::Indexed;
        }

        if b & LITERAL_WITH_NAME_REF == LITERAL_WITH_NAME_REF {
            return Representation::LiteralWithNameRef;
        }

        if b & LITERAL == LITERAL {
            return Representation::Literal;
        }

        if b & INDEXED_WITH_POST_BASE == INDEXED_WITH_POST_BASE {
            return Representation::IndexedWithPostBase;
        }

        Representation::LiteralWithPostBase
    }
}

/// A QPACK decoder.
pub struct Decoder {}

impl Default for Decoder {
    fn default() -> Decoder {
        Decoder {}
    }
}

impl Decoder {
    /// Creates a new QPACK decoder.
    pub fn new() -> Decoder {
        Decoder::default()
    }

    /// Processes control instructions from the encoder.
    pub fn control(&mut self, _buf: &mut [u8]) -> Result<()> {
        // TODO: process control instructions
        Ok(())
    }

    /// Decodes a QPACK header block into a list of headers.
    pub fn decode(&mut self, buf: &mut [u8]) -> Result<Vec<Header>> {
        let mut b = octets::Octets::with_slice(buf);

        let mut out = Vec::new();

        let req_insert_count = decode_int(&mut b, 8)?;
        let base = decode_int(&mut b, 7)?;

        trace!("Header count={} base={}", req_insert_count, base);

        while b.cap() > 0 {
            let first = b.peek_u8()?;

            match Representation::from_byte(first) {
                Representation::Indexed => {
                    const STATIC: u8 = 0x40;

                    let s = first & STATIC == STATIC;
                    let index = decode_int(&mut b, 6)?;

                    trace!("Indexed index={} static={}", index, s);

                    if !s {
                        unimplemented!("dynamic table");
                    }

                    let (name, value) = lookup_static(index)?;
                    out.push(Header::new(&name, &value));
                },

                Representation::IndexedWithPostBase => {
                    let index = decode_int(&mut b, 4)?;

                    trace!("Indexed With Post Base index={}", index);

                    unimplemented!("dynamic table");
                },

                Representation::Literal => {
                    let name_huff = b.as_ref()[0] & 0x08 == 0x08;
                    let name_len = decode_int(&mut b, 3)? as usize;

                    let mut name = b.get_bytes(name_len)?;
                    let name = if name_huff {
                        super::huffman::decode(&mut name)?
                    } else {
                        name.to_vec()
                    };

                    let name = String::from_utf8(name)
                        .map_err(|_| Error::InvalidHeaderValue)?;
                    let value = decode_str(&mut b)?;

                    trace!(
                        "Literal Without Name Reference name={:?} value={:?}",
                        name,
                        value
                    );

                    out.push(Header::new(&name, &value));
                },

                Representation::LiteralWithNameRef => {
                    const STATIC: u8 = 0x10;

                    let s = first & STATIC == STATIC;
                    let name_idx = decode_int(&mut b, 4)?;
                    let value = decode_str(&mut b)?;

                    trace!(
                        "Literal name_idx={} static={} value={:?}",
                        name_idx,
                        s,
                        value
                    );

                    if !s {
                        unimplemented!("dynamic table");
                    }

                    let (name, _) = lookup_static(name_idx)?;
                    out.push(Header::new(name, &value));
                },

                Representation::LiteralWithPostBase => {
                    trace!("Literal With Post Base");

                    unimplemented!("dynamic table");
                },
            }
        }

        Ok(out)
    }
}

fn lookup_static(idx: u64) -> Result<(&'static str, &'static str)> {
    let hdr = match idx {
        0 => (":authority", ""),
        1 => (":path", "/"),
        2 => ("age", "0"),
        3 => ("content-disposition", ""),
        4 => ("content-length", "0"),
        5 => ("cookie", ""),
        6 => ("date", ""),
        7 => ("etag", ""),
        8 => ("if-modified-since", ""),
        9 => ("if-none-match", ""),
        10 => ("last-modified", ""),
        11 => ("link", ""),
        12 => ("location", ""),
        13 => ("referer", ""),
        14 => ("set-cookie", ""),
        15 => (":method", "CONNECT"),
        16 => (":method", "DELETE"),
        17 => (":method", "GET"),
        18 => (":method", "HEAD"),
        19 => (":method", "OPTIONS"),
        20 => (":method", "POST"),
        21 => (":method", "PUT"),
        22 => (":scheme", "http"),
        23 => (":scheme", "https"),
        24 => (":status", "103"),
        25 => (":status", "200"),
        26 => (":status", "304"),
        27 => (":status", "404"),
        28 => (":status", "503"),
        29 => ("accept", "*/*"),
        30 => ("accept", "application/dns-message"),
        31 => ("accept-encoding", "gzip, deflate, br"),
        32 => ("accept-ranges", "bytes"),
        33 => ("access-control-allow-headers", "cache-control"),
        34 => ("access-control-allow-headers", "content-type"),
        35 => ("access-control-allow-origin", "*"),
        36 => ("cache-control", "max-age=0"),
        37 => ("cache-control", "max-age=2592000"),
        38 => ("cache-control", "max-age=604800"),
        39 => ("cache-control", "no-cache"),
        40 => ("cache-control", "no-store"),
        41 => ("cache-control", "public, max-age=31536000"),
        42 => ("content-encoding", "br"),
        43 => ("content-encoding", "gzip"),
        44 => ("content-type", "application/dns-message"),
        45 => ("content-type", "application/javascript"),
        46 => ("content-type", "application/json"),
        47 => ("content-type", "application/x-www-form-urlencoded"),
        48 => ("content-type", "image/gif"),
        49 => ("content-type", "image/jpeg"),
        50 => ("content-type", "image/png"),
        51 => ("content-type", "text/css"),
        52 => ("content-type", "text/html; charset=utf-8"),
        53 => ("content-type", "text/plain"),
        54 => ("content-type", "text/plain;charset=utf-8"),
        55 => ("range", "bytes=0-"),
        56 => ("strict-transport-security", "max-age=31536000"),
        57 => (
            "strict-transport-security",
            "max-age=31536000; includesubdomains",
        ),
        58 => (
            "strict-transport-security",
            "max-age=31536000; includesubdomains; preload",
        ),
        59 => ("vary", "accept-encoding"),
        60 => ("vary", "origin"),
        61 => ("x-content-type-options", "nosniff"),
        62 => ("x-xss-protection", "1; mode=block"),
        63 => (":status", "100"),
        64 => (":status", "204"),
        65 => (":status", "206"),
        66 => (":status", "302"),
        67 => (":status", "400"),
        68 => (":status", "403"),
        69 => (":status", "421"),
        70 => (":status", "425"),
        71 => (":status", "500"),
        72 => ("accept-language", ""),
        73 => ("access-control-allow-credentials", "FALSE"),
        74 => ("access-control-allow-credentials", "TRUE"),
        75 => ("access-control-allow-headers", "*"),
        76 => ("access-control-allow-methods", "get"),
        77 => ("access-control-allow-methods", "get, post, options"),
        78 => ("access-control-allow-methods", "options"),
        79 => ("access-control-expose-headers", "content-length"),
        80 => ("access-control-request-headers", "content-type"),
        81 => ("access-control-request-method", "get"),
        82 => ("access-control-request-method", "post"),
        83 => ("alt-svc", "clear"),
        84 => ("authorization", ""),
        85 => (
            "content-security-policy",
            "script-src 'none'; object-src 'none'; base-uri 'none'",
        ),
        86 => ("early-data", "1"),
        87 => ("expect-ct", ""),
        88 => ("forwarded", ""),
        89 => ("if-range", ""),
        90 => ("origin", ""),
        91 => ("purpose", "prefetch"),
        92 => ("server", ""),
        93 => ("timing-allow-origin", "*"),
        94 => ("upgrade-insecure-requests", "1"),
        95 => ("user-agent", ""),
        96 => ("x-forwarded-for", ""),
        97 => ("x-frame-options", "deny"),
        98 => ("x-frame-options", "sameorigin"),

        _ => return Err(Error::InvalidStaticTableIndex),
    };

    Ok(hdr)
}

fn decode_int(b: &mut octets::Octets, prefix: usize) -> Result<u64> {
    let mask = 2u64.pow(prefix as u32) - 1;

    let mut val = u64::from(b.get_u8()?);
    val &= mask;

    if val < mask {
        return Ok(val);
    }

    let mut shift = 0;

    while b.cap() > 0 {
        let byte = b.get_u8()?;

        val += u64::from(byte & 0x7f) << shift;

        shift += 7;

        if byte & 0x80 == 0 {
            return Ok(val);
        }
    }

    Err(Error::BufferTooShort)
}

fn decode_str<'a>(b: &'a mut octets::Octets) -> Result<String> {
    let first = b.peek_u8()?;

    let huff = first & 0x80 == 0x80;

    let len = decode_int(b, 7)? as usize;

    let mut val = b.get_bytes(len)?;

    let val = if huff {
        super::huffman::decode(&mut val)?
    } else {
        val.to_vec()
    };

    let val = String::from_utf8(val).map_err(|_| Error::InvalidHeaderValue)?;
    Ok(val)
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::octets;

    #[test]
    fn decode_int1() {
        let mut encoded = [0b01010, 0x02];
        let mut b = octets::Octets::with_slice(&mut encoded);

        assert_eq!(decode_int(&mut b, 5), Ok(10));
    }

    #[test]
    fn decode_int2() {
        let mut encoded = [0b11111, 0b10011010, 0b00001010];
        let mut b = octets::Octets::with_slice(&mut encoded);

        assert_eq!(decode_int(&mut b, 5), Ok(1337));
    }

    #[test]
    fn decode_int3() {
        let mut encoded = [0b101010];
        let mut b = octets::Octets::with_slice(&mut encoded);

        assert_eq!(decode_int(&mut b, 8), Ok(42));
    }
}
<commit_msg>qpack: return a proper error when dynamic table is requested<commit_after>// Copyright (C) 2019, Cloudflare, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use crate::octets;

use super::Error;
use super::Result;

use crate::h3::Header;

use super::INDEXED;
use super::INDEXED_WITH_POST_BASE;
use super::LITERAL;
use super::LITERAL_WITH_NAME_REF;

#[derive(Clone, Copy, Debug, PartialEq)]
enum Representation {
    Indexed,
    IndexedWithPostBase,
    Literal,
    LiteralWithNameRef,
    LiteralWithPostBase,
}

impl Representation {
    pub fn from_byte(b: u8) -> Representation {
        if b & INDEXED == INDEXED {
            return Representation::Indexed;
        }

        if b & LITERAL_WITH_NAME_REF == LITERAL_WITH_NAME_REF {
            return Representation::LiteralWithNameRef;
        }

        if b & LITERAL == LITERAL {
            return Representation::Literal;
        }

        if b & INDEXED_WITH_POST_BASE == INDEXED_WITH_POST_BASE {
            return Representation::IndexedWithPostBase;
        }

        Representation::LiteralWithPostBase
    }
}

/// A QPACK decoder.
pub struct Decoder {}

impl Default for Decoder {
    fn default() -> Decoder {
        Decoder {}
    }
}

impl Decoder {
    /// Creates a new QPACK decoder.
    pub fn new() -> Decoder {
        Decoder::default()
    }

    /// Processes control instructions from the encoder.
    pub fn control(&mut self, _buf: &mut [u8]) -> Result<()> {
        // TODO: process control instructions
        Ok(())
    }

    /// Decodes a QPACK header block into a list of headers.
    pub fn decode(&mut self, buf: &mut [u8]) -> Result<Vec<Header>> {
        let mut b = octets::Octets::with_slice(buf);

        let mut out = Vec::new();

        let req_insert_count = decode_int(&mut b, 8)?;
        let base = decode_int(&mut b, 7)?;

        trace!("Header count={} base={}", req_insert_count, base);

        while b.cap() > 0 {
            let first = b.peek_u8()?;

            match Representation::from_byte(first) {
                Representation::Indexed => {
                    const STATIC: u8 = 0x40;

                    let s = first & STATIC == STATIC;
                    let index = decode_int(&mut b, 6)?;

                    trace!("Indexed index={} static={}", index, s);

                    if !s {
                        // TODO: implement dynamic table
                        return Err(Error::InvalidHeaderValue);
                    }

                    let (name, value) = lookup_static(index)?;
                    out.push(Header::new(&name, &value));
                },

                Representation::IndexedWithPostBase => {
                    let index = decode_int(&mut b, 4)?;

                    trace!("Indexed With Post Base index={}", index);

                    // TODO: implement dynamic table
                    return Err(Error::InvalidHeaderValue);
                },

                Representation::Literal => {
                    let name_huff = b.as_ref()[0] & 0x08 == 0x08;
                    let name_len = decode_int(&mut b, 3)? as usize;

                    let mut name = b.get_bytes(name_len)?;
                    let name = if name_huff {
                        super::huffman::decode(&mut name)?
                    } else {
                        name.to_vec()
                    };

                    let name = String::from_utf8(name)
                        .map_err(|_| Error::InvalidHeaderValue)?;
                    let value = decode_str(&mut b)?;

                    trace!(
                        "Literal Without Name Reference name={:?} value={:?}",
                        name,
                        value
                    );

                    out.push(Header::new(&name, &value));
                },

                Representation::LiteralWithNameRef => {
                    const STATIC: u8 = 0x10;

                    let s = first & STATIC == STATIC;
                    let name_idx = decode_int(&mut b, 4)?;
                    let value = decode_str(&mut b)?;

                    trace!(
                        "Literal name_idx={} static={} value={:?}",
                        name_idx,
                        s,
                        value
                    );

                    if !s {
                        // TODO: implement dynamic table
                        return Err(Error::InvalidHeaderValue);
                    }

                    let (name, _) = lookup_static(name_idx)?;
                    out.push(Header::new(name, &value));
                },

                Representation::LiteralWithPostBase => {
                    trace!("Literal With Post Base");

                    // TODO: implement dynamic table
                    return Err(Error::InvalidHeaderValue);
                },
            }
        }

        Ok(out)
    }
}

fn lookup_static(idx: u64) -> Result<(&'static str, &'static str)> {
    let hdr = match idx {
        0 => (":authority", ""),
        1 => (":path", "/"),
        2 => ("age", "0"),
        3 => ("content-disposition", ""),
        4 => ("content-length", "0"),
        5 => ("cookie", ""),
        6 => ("date", ""),
        7 => ("etag", ""),
        8 => ("if-modified-since", ""),
        9 => ("if-none-match", ""),
        10 => ("last-modified", ""),
        11 => ("link", ""),
        12 => ("location", ""),
        13 => ("referer", ""),
        14 => ("set-cookie", ""),
        15 => (":method", "CONNECT"),
        16 => (":method", "DELETE"),
        17 => (":method", "GET"),
        18 => (":method", "HEAD"),
        19 => (":method", "OPTIONS"),
        20 => (":method", "POST"),
        21 => (":method", "PUT"),
        22 => (":scheme", "http"),
        23 => (":scheme", "https"),
        24 => (":status", "103"),
        25 => (":status", "200"),
        26 => (":status", "304"),
        27 => (":status", "404"),
        28 => (":status", "503"),
        29 => ("accept", "*/*"),
        30 => ("accept", "application/dns-message"),
        31 => ("accept-encoding", "gzip, deflate, br"),
        32 => ("accept-ranges", "bytes"),
        33 => ("access-control-allow-headers", "cache-control"),
        34 => ("access-control-allow-headers", "content-type"),
        35 => ("access-control-allow-origin", "*"),
        36 => ("cache-control", "max-age=0"),
        37 => ("cache-control", "max-age=2592000"),
        38 => ("cache-control", "max-age=604800"),
        39 => ("cache-control", "no-cache"),
        40 => ("cache-control", "no-store"),
        41 => ("cache-control", "public, max-age=31536000"),
        42 => ("content-encoding", "br"),
        43 => ("content-encoding", "gzip"),
        44 => ("content-type", "application/dns-message"),
        45 => ("content-type", "application/javascript"),
        46 => ("content-type", "application/json"),
        47 => ("content-type", "application/x-www-form-urlencoded"),
        48 => ("content-type", "image/gif"),
        49 => ("content-type", "image/jpeg"),
        50 => ("content-type", "image/png"),
        51 => ("content-type", "text/css"),
        52 => ("content-type", "text/html; charset=utf-8"),
        53 => ("content-type", "text/plain"),
        54 => ("content-type", "text/plain;charset=utf-8"),
        55 => ("range", "bytes=0-"),
        56 => ("strict-transport-security", "max-age=31536000"),
        57 => (
            "strict-transport-security",
            "max-age=31536000; includesubdomains",
        ),
        58 => (
            "strict-transport-security",
            "max-age=31536000; includesubdomains; preload",
        ),
        59 => ("vary", "accept-encoding"),
        60 => ("vary", "origin"),
        61 => ("x-content-type-options", "nosniff"),
        62 => ("x-xss-protection", "1; mode=block"),
        63 => (":status", "100"),
        64 => (":status", "204"),
        65 => (":status", "206"),
        66 => (":status", "302"),
        67 => (":status", "400"),
        68 => (":status", "403"),
        69 => (":status", "421"),
        70 => (":status", "425"),
        71 => (":status", "500"),
        72 => ("accept-language", ""),
        73 => ("access-control-allow-credentials", "FALSE"),
        74 => ("access-control-allow-credentials", "TRUE"),
        75 => ("access-control-allow-headers", "*"),
        76 => ("access-control-allow-methods", "get"),
        77 => ("access-control-allow-methods", "get, post, options"),
        78 => ("access-control-allow-methods", "options"),
        79 => ("access-control-expose-headers", "content-length"),
        80 => ("access-control-request-headers", "content-type"),
        81 => ("access-control-request-method", "get"),
        82 => ("access-control-request-method", "post"),
        83 => ("alt-svc", "clear"),
        84 => ("authorization", ""),
        85 => (
            "content-security-policy",
            "script-src 'none'; object-src 'none'; base-uri 'none'",
        ),
        86 => ("early-data", "1"),
        87 => ("expect-ct", ""),
        88 => ("forwarded", ""),
        89 => ("if-range", ""),
        90 => ("origin", ""),
        91 => ("purpose", "prefetch"),
        92 => ("server", ""),
        93 => ("timing-allow-origin", "*"),
        94 => ("upgrade-insecure-requests", "1"),
        95 => ("user-agent", ""),
        96 => ("x-forwarded-for", ""),
        97 => ("x-frame-options", "deny"),
        98 => ("x-frame-options", "sameorigin"),

        _ => return Err(Error::InvalidStaticTableIndex),
    };

    Ok(hdr)
}

fn decode_int(b: &mut octets::Octets, prefix: usize) -> Result<u64> {
    let mask = 2u64.pow(prefix as u32) - 1;

    let mut val = u64::from(b.get_u8()?);
    val &= mask;

    if val < mask {
        return Ok(val);
    }

    let mut shift = 0;

    while b.cap() > 0 {
        let byte = b.get_u8()?;

        val += u64::from(byte & 0x7f) << shift;

        shift += 7;

        if byte & 0x80 == 0 {
            return Ok(val);
        }
    }

    Err(Error::BufferTooShort)
}

fn decode_str<'a>(b: &'a mut octets::Octets) -> Result<String> {
    let first = b.peek_u8()?;

    let huff = first & 0x80 == 0x80;

    let len = decode_int(b, 7)? as usize;

    let mut val = b.get_bytes(len)?;

    let val = if huff {
        super::huffman::decode(&mut val)?
    } else {
        val.to_vec()
    };

    let val = String::from_utf8(val).map_err(|_| Error::InvalidHeaderValue)?;
    Ok(val)
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::octets;

    #[test]
    fn decode_int1() {
        let mut encoded = [0b01010, 0x02];
        let mut b = octets::Octets::with_slice(&mut encoded);

        assert_eq!(decode_int(&mut b, 5), Ok(10));
    }

    #[test]
    fn decode_int2() {
        let mut encoded = [0b11111, 0b10011010, 0b00001010];
        let mut b = octets::Octets::with_slice(&mut encoded);

        assert_eq!(decode_int(&mut b, 5), Ok(1337));
    }

    #[test]
    fn decode_int3() {
        let mut encoded = [0b101010];
        let mut b = octets::Octets::with_slice(&mut encoded);

        assert_eq!(decode_int(&mut b, 8), Ok(42));
    }
}
<|endoftext|>