<commit_before>use crate::utils::{match_def_path, match_type, paths, span_help_and_lint, span_lint, walk_ptrs_ty};
use if_chain::if_chain;
use rustc::hir;
use rustc::hir::def::{DefKind, Res};
use rustc::hir::intravisit::{walk_expr, NestedVisitorMap, Visitor};
use rustc::hir::*;
use rustc::lint::{EarlyContext, EarlyLintPass, LateContext, LateLintPass, LintArray, LintPass};
use rustc::{declare_lint_pass, declare_tool_lint, impl_lint_pass};
use rustc_data_structures::fx::{FxHashMap, FxHashSet};
use syntax::ast::{Crate as AstCrate, ItemKind, Name};
use syntax::source_map::Span;
use syntax_pos::symbol::LocalInternedString;

declare_clippy_lint! {
    /// **What it does:** Checks for various things we like to keep tidy in clippy.
    ///
    /// **Why is this bad?** We like to pretend we're an example of tidy code.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:** Wrong ordering of the util::paths constants.
    pub CLIPPY_LINTS_INTERNAL,
    internal,
    "various things that will negatively affect your clippy experience"
}

declare_clippy_lint! {
    /// **What it does:** Ensures every lint is associated to a `LintPass`.
    ///
    /// **Why is this bad?** The compiler only knows lints via a `LintPass`. Without
    /// putting a lint to a `LintPass::get_lints()`'s return, the compiler will not
    /// know the name of the lint.
    ///
    /// **Known problems:** Only checks for lints associated using the
    /// `declare_lint_pass!`, `impl_lint_pass!`, and `lint_array!` macros.
    ///
    /// **Example:**
    /// ```rust
    /// declare_lint! { pub LINT_1, ... }
    /// declare_lint! { pub LINT_2, ... }
    /// declare_lint! { pub FORGOTTEN_LINT, ... }
    /// // ...
    /// declare_lint_pass!(Pass => [LINT_1, LINT_2]);
    /// // missing FORGOTTEN_LINT
    /// ```
    pub LINT_WITHOUT_LINT_PASS,
    internal,
    "declaring a lint without associating it in a LintPass"
}

declare_clippy_lint! {
    /// **What it does:** Checks for calls to `cx.span_lint*` and suggests to use the `utils::*`
    /// variant of the function.
    ///
    /// **Why is this bad?** The `utils::*` variants also add a link to the Clippy documentation to the
    /// warning/error messages.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    /// Bad:
    /// ```rust
    /// cx.span_lint(LINT_NAME, "message");
    /// ```
    ///
    /// Good:
    /// ```rust
    /// utils::span_lint(cx, LINT_NAME, "message");
    /// ```
    pub COMPILER_LINT_FUNCTIONS,
    internal,
    "usage of the lint functions of the compiler instead of the utils::* variant"
}

declare_lint_pass!(ClippyLintsInternal => [CLIPPY_LINTS_INTERNAL]);

impl EarlyLintPass for ClippyLintsInternal {
    fn check_crate(&mut self, cx: &EarlyContext<'_>, krate: &AstCrate) {
        if let Some(utils) = krate
            .module
            .items
            .iter()
            .find(|item| item.ident.name.as_str() == "utils")
        {
            if let ItemKind::Mod(ref utils_mod) = utils.node {
                if let Some(paths) = utils_mod.items.iter().find(|item| item.ident.name.as_str() == "paths") {
                    if let ItemKind::Mod(ref paths_mod) = paths.node {
                        let mut last_name: Option<LocalInternedString> = None;
                        for item in &*paths_mod.items {
                            let name = item.ident.as_str();
                            if let Some(ref last_name) = last_name {
                                if **last_name > *name {
                                    span_lint(
                                        cx,
                                        CLIPPY_LINTS_INTERNAL,
                                        item.span,
                                        "this constant should be before the previous constant due to lexical \
                                         ordering",
                                    );
                                }
                            }
                            last_name = Some(name);
                        }
                    }
                }
            }
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct LintWithoutLintPass {
    declared_lints: FxHashMap<Name, Span>,
    registered_lints: FxHashSet<Name>,
}

impl_lint_pass!(LintWithoutLintPass => [LINT_WITHOUT_LINT_PASS]);

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for LintWithoutLintPass {
    fn check_item(&mut self, cx: &LateContext<'a, 'tcx>, item: &'tcx Item) {
        if let hir::ItemKind::Static(ref ty, MutImmutable, _) = item.node {
            if is_lint_ref_type(cx, ty) {
                self.declared_lints.insert(item.ident.name, item.span);
            }
        } else if let hir::ItemKind::Impl(.., Some(ref trait_ref), _, ref impl_item_refs) = item.node {
            if_chain! {
                if let hir::TraitRef{path, ..} = trait_ref;
                if let Res::Def(DefKind::Trait, def_id) = path.res;
                if match_def_path(cx, def_id, &paths::LINT_PASS);
                then {
                    let mut collector = LintCollector {
                        output: &mut self.registered_lints,
                        cx,
                    };
                    let body_id = cx.tcx.hir().body_owned_by(
                        impl_item_refs
                            .iter()
                            .find(|iiref| iiref.ident.as_str() == "get_lints")
                            .expect("LintPass needs to implement get_lints")
                            .id.hir_id
                    );
                    collector.visit_expr(&cx.tcx.hir().body(body_id).value);
                }
            }
        }
    }

    fn check_crate_post(&mut self, cx: &LateContext<'a, 'tcx>, _: &'tcx Crate) {
        for (lint_name, &lint_span) in &self.declared_lints {
            // When using the `declare_tool_lint!` macro, the original `lint_span`'s
            // file points to "<rustc macros>".
            // `compiletest-rs` thinks that's an error in a different file and
            // just ignores it. This causes the test in compile-fail/lint_pass
            // not able to capture the error.
            // Therefore, we need to climb the macro expansion tree and find the
            // actual span that invoked `declare_tool_lint!`:
            let lint_span = lint_span
                .ctxt()
                .outer_expn_info()
                .map(|ei| ei.call_site)
                .expect("unable to get call_site");

            if !self.registered_lints.contains(lint_name) {
                span_lint(
                    cx,
                    LINT_WITHOUT_LINT_PASS,
                    lint_span,
                    &format!("the lint `{}` is not added to any `LintPass`", lint_name),
                );
            }
        }
    }
}

fn is_lint_ref_type<'tcx>(cx: &LateContext<'_, 'tcx>, ty: &Ty) -> bool {
    if let TyKind::Rptr(
        _,
        MutTy {
            ty: ref inner,
            mutbl: MutImmutable,
        },
    ) = ty.node
    {
        if let TyKind::Path(ref path) = inner.node {
            if let Res::Def(DefKind::Struct, def_id) = cx.tables.qpath_res(path, inner.hir_id) {
                return match_def_path(cx, def_id, &paths::LINT);
            }
        }
    }

    false
}

struct LintCollector<'a, 'tcx: 'a> {
    output: &'a mut FxHashSet<Name>,
    cx: &'a LateContext<'a, 'tcx>,
}

impl<'a, 'tcx: 'a> Visitor<'tcx> for LintCollector<'a, 'tcx> {
    fn visit_expr(&mut self, expr: &'tcx Expr) {
        walk_expr(self, expr);
    }

    fn visit_path(&mut self, path: &'tcx Path, _: HirId) {
        if path.segments.len() == 1 {
            self.output.insert(path.segments[0].ident.name);
        }
    }
    fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
        NestedVisitorMap::All(&self.cx.tcx.hir())
    }
}

#[derive(Clone, Default)]
pub struct CompilerLintFunctions {
    map: FxHashMap<String, String>,
}

impl CompilerLintFunctions {
    pub fn new() -> Self {
        let mut map = FxHashMap::default();
        map.insert("span_lint".to_string(), "utils::span_lint".to_string());
        map.insert("struct_span_lint".to_string(), "utils::span_lint".to_string());
        map.insert("lint".to_string(), "utils::span_lint".to_string());
        map.insert("span_lint_note".to_string(), "utils::span_note_and_lint".to_string());
        map.insert("span_lint_help".to_string(), "utils::span_help_and_lint".to_string());
        Self { map }
    }
}

impl_lint_pass!(CompilerLintFunctions => [COMPILER_LINT_FUNCTIONS]);

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for CompilerLintFunctions {
    fn check_expr(&mut self, cx: &LateContext<'a, 'tcx>, expr: &'tcx Expr) {
        if_chain! {
            if let ExprKind::MethodCall(ref path, _, ref args) = expr.node;
            let fn_name = path.ident.as_str().to_string();
            if let Some(sugg) = self.map.get(&fn_name);
            let ty = walk_ptrs_ty(cx.tables.expr_ty(&args[0]));
            if match_type(cx, ty, &paths::EARLY_CONTEXT)
                || match_type(cx, ty, &paths::LATE_CONTEXT);
            then {
                span_help_and_lint(
                    cx,
                    COMPILER_LINT_FUNCTIONS,
                    path.ident.span,
                    "usage of a compiler lint function",
                    &format!("please use the Clippy variant of this function: `{}`", sugg),
                );
            }
        }
    }
}
<commit_msg>Remove `to_string()`s from CompilerLintFunctions<commit_after>use crate::utils::{match_def_path, match_type, paths, span_help_and_lint, span_lint, walk_ptrs_ty};
use if_chain::if_chain;
use rustc::hir;
use rustc::hir::def::{DefKind, Res};
use rustc::hir::intravisit::{walk_expr, NestedVisitorMap, Visitor};
use rustc::hir::*;
use rustc::lint::{EarlyContext, EarlyLintPass, LateContext, LateLintPass, LintArray, LintPass};
use rustc::{declare_lint_pass, declare_tool_lint, impl_lint_pass};
use rustc_data_structures::fx::{FxHashMap, FxHashSet};
use syntax::ast::{Crate as AstCrate, ItemKind, Name};
use syntax::source_map::Span;
use syntax_pos::symbol::LocalInternedString;

declare_clippy_lint! {
    /// **What it does:** Checks for various things we like to keep tidy in clippy.
    ///
    /// **Why is this bad?** We like to pretend we're an example of tidy code.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:** Wrong ordering of the util::paths constants.
    pub CLIPPY_LINTS_INTERNAL,
    internal,
    "various things that will negatively affect your clippy experience"
}

declare_clippy_lint! {
    /// **What it does:** Ensures every lint is associated to a `LintPass`.
    ///
    /// **Why is this bad?** The compiler only knows lints via a `LintPass`. Without
    /// putting a lint to a `LintPass::get_lints()`'s return, the compiler will not
    /// know the name of the lint.
    ///
    /// **Known problems:** Only checks for lints associated using the
    /// `declare_lint_pass!`, `impl_lint_pass!`, and `lint_array!` macros.
    ///
    /// **Example:**
    /// ```rust
    /// declare_lint! { pub LINT_1, ... }
    /// declare_lint! { pub LINT_2, ... }
    /// declare_lint! { pub FORGOTTEN_LINT, ... }
    /// // ...
    /// declare_lint_pass!(Pass => [LINT_1, LINT_2]);
    /// // missing FORGOTTEN_LINT
    /// ```
    pub LINT_WITHOUT_LINT_PASS,
    internal,
    "declaring a lint without associating it in a LintPass"
}

declare_clippy_lint! {
    /// **What it does:** Checks for calls to `cx.span_lint*` and suggests to use the `utils::*`
    /// variant of the function.
    ///
    /// **Why is this bad?** The `utils::*` variants also add a link to the Clippy documentation to the
    /// warning/error messages.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    /// Bad:
    /// ```rust
    /// cx.span_lint(LINT_NAME, "message");
    /// ```
    ///
    /// Good:
    /// ```rust
    /// utils::span_lint(cx, LINT_NAME, "message");
    /// ```
    pub COMPILER_LINT_FUNCTIONS,
    internal,
    "usage of the lint functions of the compiler instead of the utils::* variant"
}

declare_lint_pass!(ClippyLintsInternal => [CLIPPY_LINTS_INTERNAL]);

impl EarlyLintPass for ClippyLintsInternal {
    fn check_crate(&mut self, cx: &EarlyContext<'_>, krate: &AstCrate) {
        if let Some(utils) = krate
            .module
            .items
            .iter()
            .find(|item| item.ident.name.as_str() == "utils")
        {
            if let ItemKind::Mod(ref utils_mod) = utils.node {
                if let Some(paths) = utils_mod.items.iter().find(|item| item.ident.name.as_str() == "paths") {
                    if let ItemKind::Mod(ref paths_mod) = paths.node {
                        let mut last_name: Option<LocalInternedString> = None;
                        for item in &*paths_mod.items {
                            let name = item.ident.as_str();
                            if let Some(ref last_name) = last_name {
                                if **last_name > *name {
                                    span_lint(
                                        cx,
                                        CLIPPY_LINTS_INTERNAL,
                                        item.span,
                                        "this constant should be before the previous constant due to lexical \
                                         ordering",
                                    );
                                }
                            }
                            last_name = Some(name);
                        }
                    }
                }
            }
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct LintWithoutLintPass {
    declared_lints: FxHashMap<Name, Span>,
    registered_lints: FxHashSet<Name>,
}

impl_lint_pass!(LintWithoutLintPass => [LINT_WITHOUT_LINT_PASS]);

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for LintWithoutLintPass {
    fn check_item(&mut self, cx: &LateContext<'a, 'tcx>, item: &'tcx Item) {
        if let hir::ItemKind::Static(ref ty, MutImmutable, _) = item.node {
            if is_lint_ref_type(cx, ty) {
                self.declared_lints.insert(item.ident.name, item.span);
            }
        } else if let hir::ItemKind::Impl(.., Some(ref trait_ref), _, ref impl_item_refs) = item.node {
            if_chain! {
                if let hir::TraitRef{path, ..} = trait_ref;
                if let Res::Def(DefKind::Trait, def_id) = path.res;
                if match_def_path(cx, def_id, &paths::LINT_PASS);
                then {
                    let mut collector = LintCollector {
                        output: &mut self.registered_lints,
                        cx,
                    };
                    let body_id = cx.tcx.hir().body_owned_by(
                        impl_item_refs
                            .iter()
                            .find(|iiref| iiref.ident.as_str() == "get_lints")
                            .expect("LintPass needs to implement get_lints")
                            .id.hir_id
                    );
                    collector.visit_expr(&cx.tcx.hir().body(body_id).value);
                }
            }
        }
    }

    fn check_crate_post(&mut self, cx: &LateContext<'a, 'tcx>, _: &'tcx Crate) {
        for (lint_name, &lint_span) in &self.declared_lints {
            // When using the `declare_tool_lint!` macro, the original `lint_span`'s
            // file points to "<rustc macros>".
            // `compiletest-rs` thinks that's an error in a different file and
            // just ignores it. This causes the test in compile-fail/lint_pass
            // not able to capture the error.
            // Therefore, we need to climb the macro expansion tree and find the
            // actual span that invoked `declare_tool_lint!`:
            let lint_span = lint_span
                .ctxt()
                .outer_expn_info()
                .map(|ei| ei.call_site)
                .expect("unable to get call_site");

            if !self.registered_lints.contains(lint_name) {
                span_lint(
                    cx,
                    LINT_WITHOUT_LINT_PASS,
                    lint_span,
                    &format!("the lint `{}` is not added to any `LintPass`", lint_name),
                );
            }
        }
    }
}

fn is_lint_ref_type<'tcx>(cx: &LateContext<'_, 'tcx>, ty: &Ty) -> bool {
    if let TyKind::Rptr(
        _,
        MutTy {
            ty: ref inner,
            mutbl: MutImmutable,
        },
    ) = ty.node
    {
        if let TyKind::Path(ref path) = inner.node {
            if let Res::Def(DefKind::Struct, def_id) = cx.tables.qpath_res(path, inner.hir_id) {
                return match_def_path(cx, def_id, &paths::LINT);
            }
        }
    }

    false
}

struct LintCollector<'a, 'tcx: 'a> {
    output: &'a mut FxHashSet<Name>,
    cx: &'a LateContext<'a, 'tcx>,
}

impl<'a, 'tcx: 'a> Visitor<'tcx> for LintCollector<'a, 'tcx> {
    fn visit_expr(&mut self, expr: &'tcx Expr) {
        walk_expr(self, expr);
    }

    fn visit_path(&mut self, path: &'tcx Path, _: HirId) {
        if path.segments.len() == 1 {
            self.output.insert(path.segments[0].ident.name);
        }
    }
    fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
        NestedVisitorMap::All(&self.cx.tcx.hir())
    }
}

#[derive(Clone, Default)]
pub struct CompilerLintFunctions {
    map: FxHashMap<&'static str, &'static str>,
}

impl CompilerLintFunctions {
    pub fn new() -> Self {
        let mut map = FxHashMap::default();
        map.insert("span_lint", "utils::span_lint");
        map.insert("struct_span_lint", "utils::span_lint");
        map.insert("lint", "utils::span_lint");
        map.insert("span_lint_note", "utils::span_note_and_lint");
        map.insert("span_lint_help", "utils::span_help_and_lint");
        Self { map }
    }
}

impl_lint_pass!(CompilerLintFunctions => [COMPILER_LINT_FUNCTIONS]);

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for CompilerLintFunctions {
    fn check_expr(&mut self, cx: &LateContext<'a, 'tcx>, expr: &'tcx Expr) {
        if_chain! {
            if let ExprKind::MethodCall(ref path, _, ref args) = expr.node;
            let fn_name = path.ident;
            if let Some(sugg) = self.map.get(&*fn_name.as_str());
            let ty = walk_ptrs_ty(cx.tables.expr_ty(&args[0]));
            if match_type(cx, ty, &paths::EARLY_CONTEXT)
                || match_type(cx, ty, &paths::LATE_CONTEXT);
            then {
                span_help_and_lint(
                    cx,
                    COMPILER_LINT_FUNCTIONS,
                    path.ident.span,
                    "usage of a compiler lint function",
                    &format!("please use the Clippy variant of this function: `{}`", sugg),
                );
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate llvm_sys;

use std::ptr;
use self::llvm_sys::core::*;
use self::llvm_sys::*;
use llvm::codegen::{Context, c_str_ptr};

// Generate LLVM IR for functions built into the language
pub unsafe fn generate_builtins(context: &mut Context) {
    generate_types(context);
    generate_print(context);

    // Generate a `main` function for the script to be housed
    let i32_type= LLVMInt32TypeInContext(context.get_context());
    let fn_type = LLVMFunctionType(i32_type, ptr::null_mut(), 0, 0);
    let main = LLVMAddFunction(context.get_module(), c_str_ptr("main"), fn_type);
    let bb = LLVMAppendBasicBlockInContext(context.get_context(), main, c_str_ptr("entry"));

    // Goes to the end of the block
    LLVMPositionBuilderAtEnd(context.get_builder(), bb);

    // Adds a return statement
    LLVMBuildRet(context.get_builder(), LLVMConstInt(i32_type, 1, 1));

    // Reposition to the start of the block
    LLVMPositionBuilder(context.get_builder(), bb, LLVMGetFirstInstruction(bb));
}

unsafe fn generate_print(context: &mut Context) {
    // Types
    let void = LLVMVoidTypeInContext(context.get_context());
    let i8_type = LLVMInt8TypeInContext(context.get_context());
    let i32_type = LLVMInt32TypeInContext(context.get_context());
    let i64_type = LLVMInt64TypeInContext(context.get_context());
    let i32_ptr_type = LLVMPointerType(i32_type, 0);
    let string_type_fields = vec![LLVMPointerType(i8_type, 0), i64_type];
    let string_type = LLVMStructTypeInContext(context.get_context(), string_type_fields.as_ptr() as *mut _, 2, 0);

    let args = vec![string_type];

    // 2nd arg is param types (*mut LLVMTypeRef) and 3rd is param count. 4 is variable # of args t/f?
    let fn_type = LLVMFunctionType(void, args.as_ptr() as *mut _, 1, 0);
    let print_fn = LLVMAddFunction(context.get_module(), c_str_ptr("print"), fn_type);

    // Name the param
    let param = LLVMGetFirstParam(print_fn);
    LLVMSetValueName(param, c_str_ptr("str"));

    // Create basic blocks to generate code in
    let start_block = LLVMAppendBasicBlockInContext(context.get_context(), print_fn, c_str_ptr("start"));
    let loop_block = LLVMAppendBasicBlockInContext(context.get_context(), print_fn, c_str_ptr("loop"));
    let end_block = LLVMAppendBasicBlockInContext(context.get_context(), print_fn, c_str_ptr("end"));

    // Start
    LLVMPositionBuilderAtEnd(context.get_builder(), start_block);

    // let i32_zero = LLVMConstInt(i32_type, 0, 0);
    let i64_zero = LLVMConstInt(i64_type, 0, 0);
    let op = LLVMIntPredicate::LLVMIntEQ;
    let offset_ptr = LLVMBuildAlloca(context.get_builder(), i64_type, c_str_ptr("offsetptr"));
    LLVMBuildStore(context.get_builder(), i64_zero, offset_ptr);
    let str_ptr = LLVMBuildExtractValue(context.get_builder(), param, 0, c_str_ptr("strptr"));
    let len = LLVMBuildExtractValue(context.get_builder(), param, 1, c_str_ptr("len"));
    let cmp = LLVMBuildICmp(context.get_builder(), op, len, i64_zero, c_str_ptr("cmp"));

    // Branch to end if string len is 0
    LLVMBuildCondBr(context.get_builder(), cmp, end_block, loop_block);

    // Loop
    LLVMPositionBuilderAtEnd(context.get_builder(), loop_block);

    let offset = LLVMBuildLoad(context.get_builder(), offset_ptr, c_str_ptr("offset"));
    let indices = vec![offset];
    let iter_ptr = LLVMBuildGEP(context.get_builder(), str_ptr, indices.as_ptr() as *mut _, 1, c_str_ptr("iterptr"));

    let op = LLVMOpcode::LLVMBitCast;
    let iterptr32 = LLVMBuildCast(context.get_builder(), op, iter_ptr, i32_ptr_type, c_str_ptr("iterptr32"));

    let iter32 = LLVMBuildLoad(context.get_builder(), iterptr32, c_str_ptr("iter"));

    let mut putchar_fn = LLVMGetNamedFunction(context.get_module(), c_str_ptr("putchar"));
    if putchar_fn.is_null() {
        let fn_type2 = LLVMFunctionType(i32_type, vec![i32_type].as_ptr() as *mut _, 1, 0);
        putchar_fn = LLVMAddFunction(context.get_module(), c_str_ptr("putchar"), fn_type2);
    }

    // Print char at ptr_iter here
    LLVMBuildCall(context.get_builder(), putchar_fn, vec![iter32].as_ptr() as *mut _, 1, c_str_ptr(""));

    // Increment the offset
    let i64_one = LLVMConstInt(i64_type, 1, 0);
    let inc = LLVMBuildAdd(context.get_builder(), offset, i64_one, c_str_ptr("inc"));
    LLVMBuildStore(context.get_builder(), inc, offset_ptr);

    // Branch if offset equals len, else keep looping
    let op = LLVMIntPredicate::LLVMIntEQ;
    let cmp2 = LLVMBuildICmp(context.get_builder(), op, inc, len, c_str_ptr("cmp"));
    LLVMBuildCondBr(context.get_builder(), cmp2, end_block, loop_block);

    // End
    LLVMPositionBuilderAtEnd(context.get_builder(), end_block);
    LLVMBuildRetVoid(context.get_builder());
}

unsafe fn generate_types(context: &mut Context) {
    // Define a string structure that looks like:
    // string {
    //     array: i8*
    //     len: i64,
    // }

    let string_type_fields = vec![LLVMInt32TypeInContext(context.get_context()),
                  LLVMPointerType(LLVMInt8TypeInContext(context.get_context()), 0)];

//    let struct_type = LLVMStructType(string_type_fields.as_ptr() as *mut _, 2, 0);
//    let const_str_var = LLVMAddGlobal(context.get_module(), struct_type, c_str_ptr("struct.String"));

    let string_struct_type = LLVMStructCreateNamed(LLVMGetGlobalContext(), c_str_ptr("string"));
    LLVMStructSetBody(string_struct_type, string_type_fields.as_ptr() as *mut _, 2, 0);
//    let string_struct = LLVMStructType(struct_type, 2, 0);
}
<commit_msg>Changed how string struct is created.<commit_after>extern crate llvm_sys;

use std::ptr;
use self::llvm_sys::core::*;
use self::llvm_sys::*;
use llvm::codegen::{Context, c_str_ptr};

// Generate LLVM IR for functions built into the language
pub unsafe fn generate_builtins(context: &mut Context) {
    generate_types(context);
    generate_print(context);

    // Generate a `main` function for the script to be housed
    let i32_type= LLVMInt32TypeInContext(context.get_context());
    let fn_type = LLVMFunctionType(i32_type, ptr::null_mut(), 0, 0);
    let main = LLVMAddFunction(context.get_module(), c_str_ptr("main"), fn_type);
    let bb = LLVMAppendBasicBlockInContext(context.get_context(), main, c_str_ptr("entry"));

    // Goes to the end of the block
    LLVMPositionBuilderAtEnd(context.get_builder(), bb);

    // Adds a return statement
    LLVMBuildRet(context.get_builder(), LLVMConstInt(i32_type, 1, 1));

    // Reposition to the start of the block
    LLVMPositionBuilder(context.get_builder(), bb, LLVMGetFirstInstruction(bb));
}

unsafe fn generate_print(context: &mut Context) {
    // Types
    let void = LLVMVoidTypeInContext(context.get_context());
    let i8_type = LLVMInt8TypeInContext(context.get_context());
    let i32_type = LLVMInt32TypeInContext(context.get_context());
    let i64_type = LLVMInt64TypeInContext(context.get_context());
    let i32_ptr_type = LLVMPointerType(i32_type, 0);
    let string_type_fields = vec![LLVMPointerType(i8_type, 0), i64_type];
    let string_type = LLVMStructTypeInContext(context.get_context(), string_type_fields.as_ptr() as *mut _, 2, 0);

    let args = vec![string_type];

    // 2nd arg is param types (*mut LLVMTypeRef) and 3rd is param count. 4 is variable # of args t/f?
    let fn_type = LLVMFunctionType(void, args.as_ptr() as *mut _, 1, 0);
    let print_fn = LLVMAddFunction(context.get_module(), c_str_ptr("print"), fn_type);

    // Name the param
    let param = LLVMGetFirstParam(print_fn);
    LLVMSetValueName(param, c_str_ptr("str"));

    // Create basic blocks to generate code in
    let start_block = LLVMAppendBasicBlockInContext(context.get_context(), print_fn, c_str_ptr("start"));
    let loop_block = LLVMAppendBasicBlockInContext(context.get_context(), print_fn, c_str_ptr("loop"));
    let end_block = LLVMAppendBasicBlockInContext(context.get_context(), print_fn, c_str_ptr("end"));

    // Start
    LLVMPositionBuilderAtEnd(context.get_builder(), start_block);

    // let i32_zero = LLVMConstInt(i32_type, 0, 0);
    let i64_zero = LLVMConstInt(i64_type, 0, 0);
    let op = LLVMIntPredicate::LLVMIntEQ;
    let offset_ptr = LLVMBuildAlloca(context.get_builder(), i64_type, c_str_ptr("offsetptr"));
    LLVMBuildStore(context.get_builder(), i64_zero, offset_ptr);
    let str_ptr = LLVMBuildExtractValue(context.get_builder(), param, 0, c_str_ptr("strptr"));
    let len = LLVMBuildExtractValue(context.get_builder(), param, 1, c_str_ptr("len"));
    let cmp = LLVMBuildICmp(context.get_builder(), op, len, i64_zero, c_str_ptr("cmp"));

    // Branch to end if string len is 0
    LLVMBuildCondBr(context.get_builder(), cmp, end_block, loop_block);

    // Loop
    LLVMPositionBuilderAtEnd(context.get_builder(), loop_block);

    let offset = LLVMBuildLoad(context.get_builder(), offset_ptr, c_str_ptr("offset"));
    let indices = vec![offset];
    let iter_ptr = LLVMBuildGEP(context.get_builder(), str_ptr, indices.as_ptr() as *mut _, 1, c_str_ptr("iterptr"));

    let op = LLVMOpcode::LLVMBitCast;
    let iterptr32 = LLVMBuildCast(context.get_builder(), op, iter_ptr, i32_ptr_type, c_str_ptr("iterptr32"));

    let iter32 = LLVMBuildLoad(context.get_builder(), iterptr32, c_str_ptr("iter"));

    let mut putchar_fn = LLVMGetNamedFunction(context.get_module(), c_str_ptr("putchar"));
    if putchar_fn.is_null() {
        let fn_type2 = LLVMFunctionType(i32_type, vec![i32_type].as_ptr() as *mut _, 1, 0);
        putchar_fn = LLVMAddFunction(context.get_module(), c_str_ptr("putchar"), fn_type2);
    }

    // Print char at ptr_iter here
    LLVMBuildCall(context.get_builder(), putchar_fn, vec![iter32].as_ptr() as *mut _, 1, c_str_ptr(""));

    // Increment the offset
    let i64_one = LLVMConstInt(i64_type, 1, 0);
    let inc = LLVMBuildAdd(context.get_builder(), offset, i64_one, c_str_ptr("inc"));
    LLVMBuildStore(context.get_builder(), inc, offset_ptr);

    // Branch if offset equals len, else keep looping
    let op = LLVMIntPredicate::LLVMIntEQ;
    let cmp2 = LLVMBuildICmp(context.get_builder(), op, inc, len, c_str_ptr("cmp"));
    LLVMBuildCondBr(context.get_builder(), cmp2, end_block, loop_block);

    // End
    LLVMPositionBuilderAtEnd(context.get_builder(), end_block);
    LLVMBuildRetVoid(context.get_builder());
}

unsafe fn generate_types(context: &mut Context) {
    // Define a string structure that looks like:
    // string {
    //     array: i8*
    //     len: i64,
    // }

    let string_type_fields = vec![LLVMInt32TypeInContext(context.get_context()),
                  LLVMPointerType(LLVMInt8TypeInContext(context.get_context()), 0)];
    let string_type = LLVMStructType(string_type_fields.as_ptr() as *mut _, 2, 0);

//    let struct_type = LLVMStructType(string_type_fields.as_ptr() as *mut _, 2, 0);
//    let const_str_var = LLVMAddGlobal(context.get_module(), struct_type, c_str_ptr("struct.String"));

    // let string_struct_type = LLVMStructCreateNamed(LLVMGetGlobalContext(), c_str_ptr("string"));
    LLVMStructSetBody(string_type, string_type as *mut _, 2, 0);
//    let string_struct = LLVMStructType(struct_type, 2, 0);
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! See docs in build/expr/mod.rs

use build::expr::category::{Category, RvalueFunc};
use build::{BlockAnd, BlockAndExtension, BlockFrame, Builder};
use hair::*;
use rustc::mir::*;
use rustc::ty;

use rustc_target::spec::abi::Abi;

impl<'a, 'gcx, 'tcx> Builder<'a, 'gcx, 'tcx> {
    /// Compile `expr`, storing the result into `destination`, which
    /// is assumed to be uninitialized.
    pub fn into_expr(
        &mut self,
        destination: &Place<'tcx>,
        mut block: BasicBlock,
        expr: Expr<'tcx>,
    ) -> BlockAnd<()> {
        debug!(
            "into_expr(destination={:?}, block={:?}, expr={:?})",
            destination, block, expr
        );

        // since we frequently have to reference `self` from within a
        // closure, where `self` would be shadowed, it's easier to
        // just use the name `this` uniformly
        let this = self;
        let expr_span = expr.span;
        let source_info = this.source_info(expr_span);

        let expr_is_block_or_scope = match expr.kind {
            ExprKind::Block { .. } => true,
            ExprKind::Scope { .. } => true,
            _ => false,
        };

        if !expr_is_block_or_scope {
            this.block_context.push(BlockFrame::SubExpr);
        }

        let block_and = match expr.kind {
            ExprKind::Scope {
                region_scope,
                lint_level,
                value,
            } => {
                let region_scope = (region_scope, source_info);
                this.in_scope(region_scope, lint_level, block, |this| {
                    this.into(destination, block, value)
                })
            }
            ExprKind::Block { body: ast_block } => {
                this.ast_block(destination, block, ast_block, source_info)
            }
            ExprKind::Match { discriminant, arms } => {
                this.match_expr(destination, expr_span, block, discriminant, arms)
            }
            ExprKind::NeverToAny { source } => {
                let source = this.hir.mirror(source);
                let is_call = match source.kind {
                    ExprKind::Call { .. } => true,
                    _ => false,
                };

                unpack!(block = this.as_local_rvalue(block, source));

                // This is an optimization. If the expression was a call then we already have an
                // unreachable block. Don't bother to terminate it and create a new one.
                if is_call {
                    block.unit()
                } else {
                    this.cfg
                        .terminate(block, source_info, TerminatorKind::Unreachable);
                    let end_block = this.cfg.start_new_block();
                    end_block.unit()
                }
            }
            ExprKind::If {
                condition: cond_expr,
                then: then_expr,
                otherwise: else_expr,
            } => {
                let operand = unpack!(block = this.as_local_operand(block, cond_expr));

                let mut then_block = this.cfg.start_new_block();
                let mut else_block = this.cfg.start_new_block();
                let term = TerminatorKind::if_(this.hir.tcx(), operand, then_block, else_block);
                this.cfg.terminate(block, source_info, term);

                unpack!(then_block = this.into(destination, then_block, then_expr));
                else_block = if let Some(else_expr) = else_expr {
                    unpack!(this.into(destination, else_block, else_expr))
                } else {
                    // Body of the `if` expression without an `else` clause must return `()`, thus
                    // we implicitly generate a `else {}` if it is not specified.
                    this.cfg
                        .push_assign_unit(else_block, source_info, destination);
                    else_block
                };

                let join_block = this.cfg.start_new_block();
                this.cfg.terminate(
                    then_block,
                    source_info,
                    TerminatorKind::Goto { target: join_block },
                );
                this.cfg.terminate(
                    else_block,
                    source_info,
                    TerminatorKind::Goto { target: join_block },
                );

                join_block.unit()
            }
            ExprKind::LogicalOp { op, lhs, rhs } => {
                // And:
                //
                // [block: If(lhs)] -true-> [else_block: If(rhs)] -true-> [true_block]
                //        |                          | (false)
                //        +----------false-----------+------------------> [false_block]
                //
                // Or:
                //
                // [block: If(lhs)] -false-> [else_block: If(rhs)] -true-> [true_block]
                //        | (true)                   | (false)
                //  [true_block]               [false_block]

                let (true_block, false_block, mut else_block, join_block) = (
                    this.cfg.start_new_block(),
                    this.cfg.start_new_block(),
                    this.cfg.start_new_block(),
                    this.cfg.start_new_block(),
                );

                let lhs = unpack!(block = this.as_local_operand(block, lhs));
                let blocks = match op {
                    LogicalOp::And => (else_block, false_block),
                    LogicalOp::Or => (true_block, else_block),
                };
                let term = TerminatorKind::if_(this.hir.tcx(), lhs, blocks.0, blocks.1);
                this.cfg.terminate(block, source_info, term);

                let rhs = unpack!(else_block = this.as_local_operand(else_block, rhs));
                let term = TerminatorKind::if_(this.hir.tcx(), rhs, true_block, false_block);
                this.cfg.terminate(else_block, source_info, term);

                this.cfg.push_assign_constant(
                    true_block,
                    source_info,
                    destination,
                    Constant {
                        span: expr_span,
                        ty: this.hir.bool_ty(),
                        user_ty: None,
                        literal: this.hir.true_literal(),
                    },
                );

                this.cfg.push_assign_constant(
                    false_block,
                    source_info,
                    destination,
                    Constant {
                        span: expr_span,
                        ty: this.hir.bool_ty(),
                        user_ty: None,
                        literal: this.hir.false_literal(),
                    },
                );

                this.cfg.terminate(
                    true_block,
                    source_info,
                    TerminatorKind::Goto { target: join_block },
                );
                this.cfg.terminate(
                    false_block,
                    source_info,
                    TerminatorKind::Goto { target: join_block },
                );

                join_block.unit()
            }
            ExprKind::Loop {
                condition: opt_cond_expr,
                body,
            } => {
                // [block] --> [loop_block] -/eval. cond./-> [loop_block_end] -1-> [exit_block]
                //                  ^                               |
                //                  |                               0
                //                  |                               |
                //                  |                               v
                //           [body_block_end] <-/eval. body/-- [body_block]
                //
                // If `opt_cond_expr` is `None`, then the graph is somewhat simplified:
                //
                // [block]
                //    |
                //   [loop_block] -> [body_block] -/eval. body/-> [body_block_end]
                //    |        ^                                         |
                // false link  |                                         |
                //    |        +-----------------------------------------+
                //    +-> [diverge_cleanup]
                // The false link is required to make sure borrowck considers unwinds through the
                // body, even when the exact code in the body cannot unwind

                let loop_block = this.cfg.start_new_block();
                let exit_block = this.cfg.start_new_block();

                // start the loop
                this.cfg.terminate(
                    block,
                    source_info,
                    TerminatorKind::Goto { target: loop_block },
                );

                this.in_breakable_scope(
                    Some(loop_block),
                    exit_block,
                    destination.clone(),
                    move |this| {
                        // conduct the test, if necessary
                        let body_block;
                        if let Some(cond_expr) = opt_cond_expr {
                            let loop_block_end;
                            let cond = unpack!(
                                loop_block_end = this.as_local_operand(loop_block, cond_expr)
                            );
                            body_block = this.cfg.start_new_block();
                            let term =
                                TerminatorKind::if_(this.hir.tcx(), cond, body_block, exit_block);
                            this.cfg.terminate(loop_block_end, source_info, term);

                            // if the test is false, there's no `break` to assign `destination`, so
                            // we have to do it; this overwrites any `break`-assigned value but it's
                            // always `()` anyway
                            this.cfg
                                .push_assign_unit(exit_block, source_info, destination);
                        } else {
                            body_block = this.cfg.start_new_block();
                            let diverge_cleanup = this.diverge_cleanup();
                            this.cfg.terminate(
                                loop_block,
                                source_info,
                                TerminatorKind::FalseUnwind {
                                    real_target: body_block,
                                    unwind: Some(diverge_cleanup),
                                },
                            )
                        }

                        // The “return” value of the loop body must always be an unit. We therefore
                        // introduce a unit temporary as the destination for the loop body.
                        let tmp = this.get_unit_temp();
                        // Execute the body, branching back to the test.
                        let body_block_end = unpack!(this.into(&tmp, body_block, body));
                        this.cfg.terminate(
                            body_block_end,
                            source_info,
                            TerminatorKind::Goto { target: loop_block },
                        );
                    },
                );
                exit_block.unit()
            }
            ExprKind::Call { ty, fun, args, from_hir_call } => {
                // FIXME(canndrew): This is_never should probably be an is_uninhabited
                let diverges = expr.ty.is_never();
                let intrinsic = match ty.sty {
                    ty::FnDef(def_id, _) => {
                        let f = ty.fn_sig(this.hir.tcx());
                        if f.abi() == Abi::RustIntrinsic || f.abi() == Abi::PlatformIntrinsic {
                            Some(this.hir.tcx().item_name(def_id).as_str())
                        } else {
                            None
                        }
                    }
                    _ => None,
                };
                let intrinsic = intrinsic.as_ref().map(|s| &s[..]);
                let fun = unpack!(block = this.as_local_operand(block, fun));
                if intrinsic == Some("move_val_init") {
                    // `move_val_init` has "magic" semantics - the second argument is
                    // always evaluated "directly" into the first one.

                    let mut args = args.into_iter();
                    let ptr = args.next().expect("0 arguments to `move_val_init`");
                    let val = args.next().expect("1 argument to `move_val_init`");
                    assert!(args.next().is_none(), ">2 arguments to `move_val_init`");

                    let ptr = this.hir.mirror(ptr);
                    let ptr_ty = ptr.ty;
                    // Create an *internal* temp for the pointer, so that unsafety
                    // checking won't complain about the raw pointer assignment.
                    let ptr_temp = this.local_decls.push(LocalDecl {
                        mutability: Mutability::Mut,
                        ty: ptr_ty,
                        user_ty: UserTypeProjections::none(),
                        name: None,
                        source_info,
                        visibility_scope: source_info.scope,
                        internal: true,
                        is_user_variable: None,
                        is_block_tail: None,
                    });
                    let ptr_temp = Place::Local(ptr_temp);
                    let block = unpack!(this.into(&ptr_temp, block, ptr));
                    this.into(&ptr_temp.deref(), block, val)
                } else {
                    let args: Vec<_> = args
                        .into_iter()
                        .map(|arg| unpack!(block = this.as_local_operand(block, arg)))
                        .collect();

                    let success = this.cfg.start_new_block();
                    let cleanup = this.diverge_cleanup();
                    this.cfg.terminate(
                        block,
                        source_info,
                        TerminatorKind::Call {
                            func: fun,
                            args,
                            cleanup: Some(cleanup),
                            destination: if diverges {
                                None
                            } else {
                                Some((destination.clone(), success))
                            },
                            from_hir_call,
                        },
                    );
                    success.unit()
                }
            }

            // These cases don't actually need a destination
            ExprKind::Assign { .. }
            | ExprKind::AssignOp { .. }
            | ExprKind::Continue { .. }
            | ExprKind::Break { .. }
            | ExprKind::InlineAsm { .. }
            | ExprKind::Return { .. } => {
                unpack!(block = this.stmt_expr(block, expr, None));
                this.cfg.push_assign_unit(block, source_info, destination);
                block.unit()
            }

            // Avoid creating a temporary
            ExprKind::VarRef { .. } |
            ExprKind::SelfRef |
            ExprKind::StaticRef { .. } |
            ExprKind::PlaceTypeAscription { .. } |
            ExprKind::ValueTypeAscription { .. } => {
                debug_assert!(Category::of(&expr.kind) == Some(Category::Place));

                let place = unpack!(block = this.as_place(block, expr));
                let rvalue = Rvalue::Use(this.consume_by_copy_or_move(place));
                this.cfg
                    .push_assign(block, source_info, destination, rvalue);
                block.unit()
            }
            ExprKind::Index { .. } | ExprKind::Deref { .. } | ExprKind::Field { .. } => {
                debug_assert!(Category::of(&expr.kind) == Some(Category::Place));

                // Create a "fake" temporary variable so that we check that the
                // value is Sized. Usually, this is caught in type checking, but
                // in the case of box expr there is no such check.
                if let Place::Projection(..) = destination {
                    this.local_decls
                        .push(LocalDecl::new_temp(expr.ty, expr.span));
                }

                debug_assert!(Category::of(&expr.kind) == Some(Category::Place));

                let place = unpack!(block = this.as_place(block, expr));
                let rvalue = Rvalue::Use(this.consume_by_copy_or_move(place));
                this.cfg
                    .push_assign(block, source_info, destination, rvalue);
                block.unit()
            }

            // these are the cases that are more naturally handled by some other mode
            ExprKind::Unary { .. }
            | ExprKind::Binary { .. }
            | ExprKind::Box { .. }
            | ExprKind::Cast { .. }
            | ExprKind::Use { .. }
            | ExprKind::ReifyFnPointer { .. }
            | ExprKind::ClosureFnPointer { .. }
            | ExprKind::UnsafeFnPointer { .. }
            | ExprKind::Unsize { .. }
            | ExprKind::Repeat { .. }
            | ExprKind::Borrow { .. }
            | ExprKind::Array { .. }
            | ExprKind::Tuple { .. }
            | ExprKind::Adt { .. }
            | ExprKind::Closure { .. }
            | ExprKind::Literal { .. }
            | ExprKind::Yield { .. } => {
                debug_assert!(match Category::of(&expr.kind).unwrap() {
                    // should be handled above
                    Category::Rvalue(RvalueFunc::Into) => false,

                    // must be handled above or else we get an
                    // infinite loop in the builder; see
                    // e.g., `ExprKind::VarRef` above
                    Category::Place => false,

                    _ => true,
                });

                let rvalue = unpack!(block = this.as_local_rvalue(block, expr));
                this.cfg
                    .push_assign(block, source_info, destination, rvalue);
                block.unit()
            }
        };

        if !expr_is_block_or_scope {
            let popped = this.block_context.pop();
            assert!(popped.is_some());
        }

        block_and
    }
}
<commit_msg>Simplify MIR generation for logical ops<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! See docs in build/expr/mod.rs

use build::expr::category::{Category, RvalueFunc};
use build::{BlockAnd, BlockAndExtension, BlockFrame, Builder};
use hair::*;
use rustc::mir::*;
use rustc::ty;

use rustc_target::spec::abi::Abi;

impl<'a, 'gcx, 'tcx> Builder<'a, 'gcx, 'tcx> {
    /// Compile `expr`, storing the result into `destination`, which
    /// is assumed to be uninitialized.
    pub fn into_expr(
        &mut self,
        destination: &Place<'tcx>,
        mut block: BasicBlock,
        expr: Expr<'tcx>,
    ) -> BlockAnd<()> {
        debug!(
            "into_expr(destination={:?}, block={:?}, expr={:?})",
            destination, block, expr
        );

        // since we frequently have to reference `self` from within a
        // closure, where `self` would be shadowed, it's easier to
        // just use the name `this` uniformly
        let this = self;
        let expr_span = expr.span;
        let source_info = this.source_info(expr_span);

        let expr_is_block_or_scope = match expr.kind {
            ExprKind::Block { .. } => true,
            ExprKind::Scope { .. } => true,
            _ => false,
        };

        if !expr_is_block_or_scope {
            this.block_context.push(BlockFrame::SubExpr);
        }

        let block_and = match expr.kind {
            ExprKind::Scope {
                region_scope,
                lint_level,
                value,
            } => {
                let region_scope = (region_scope, source_info);
                this.in_scope(region_scope, lint_level, block, |this| {
                    this.into(destination, block, value)
                })
            }
            ExprKind::Block { body: ast_block } => {
                this.ast_block(destination, block, ast_block, source_info)
            }
            ExprKind::Match { discriminant, arms } => {
                this.match_expr(destination, expr_span, block, discriminant, arms)
            }
            ExprKind::NeverToAny { source } => {
                let source = this.hir.mirror(source);
                let is_call = match source.kind {
                    ExprKind::Call { .. } => true,
                    _ => false,
                };

                unpack!(block = this.as_local_rvalue(block, source));

                // This is an optimization. If the expression was a call then we already have an
                // unreachable block. Don't bother to terminate it and create a new one.
                if is_call {
                    block.unit()
                } else {
                    this.cfg
                        .terminate(block, source_info, TerminatorKind::Unreachable);
                    let end_block = this.cfg.start_new_block();
                    end_block.unit()
                }
            }
            ExprKind::If {
                condition: cond_expr,
                then: then_expr,
                otherwise: else_expr,
            } => {
                let operand = unpack!(block = this.as_local_operand(block, cond_expr));

                let mut then_block = this.cfg.start_new_block();
                let mut else_block = this.cfg.start_new_block();
                let term = TerminatorKind::if_(this.hir.tcx(), operand, then_block, else_block);
                this.cfg.terminate(block, source_info, term);

                unpack!(then_block = this.into(destination, then_block, then_expr));
                else_block = if let Some(else_expr) = else_expr {
                    unpack!(this.into(destination, else_block, else_expr))
                } else {
                    // Body of the `if` expression without an `else` clause must return `()`, thus
                    // we implicitly generate a `else {}` if it is not specified.
                    this.cfg
                        .push_assign_unit(else_block, source_info, destination);
                    else_block
                };

                let join_block = this.cfg.start_new_block();
                this.cfg.terminate(
                    then_block,
                    source_info,
                    TerminatorKind::Goto { target: join_block },
                );
                this.cfg.terminate(
                    else_block,
                    source_info,
                    TerminatorKind::Goto { target: join_block },
                );

                join_block.unit()
            }
            ExprKind::LogicalOp { op, lhs, rhs } => {
                // And:
                //
                // [block: If(lhs)] -true-> [else_block: dest = (rhs)]
                //        | (false)
                //  [shortcurcuit_block: dest = false]
                //
                // Or:
                //
                // [block: If(lhs)] -false-> [else_block: dest = (rhs)]
                //        | (true)
                //  [shortcurcuit_block: dest = true]

                let (shortcircuit_block, mut else_block, join_block) = (
                    this.cfg.start_new_block(),
                    this.cfg.start_new_block(),
                    this.cfg.start_new_block(),
                );

                let lhs = unpack!(block = this.as_local_operand(block, lhs));
                let blocks = match op {
                    LogicalOp::And => (else_block, shortcircuit_block),
                    LogicalOp::Or => (shortcircuit_block, else_block),
                };
                let term = TerminatorKind::if_(this.hir.tcx(), lhs, blocks.0, blocks.1);
                this.cfg.terminate(block, source_info, term);

                this.cfg.push_assign_constant(
                    shortcircuit_block,
                    source_info,
                    destination,
                    Constant {
                        span: expr_span,
                        ty: this.hir.bool_ty(),
                        user_ty: None,
                        literal: match op {
                            LogicalOp::And => this.hir.false_literal(),
                            LogicalOp::Or => this.hir.true_literal(),
                        },
                    },
                );
                this.cfg.terminate(
                    shortcircuit_block,
                    source_info,
                    TerminatorKind::Goto { target: join_block },
                );

                let rhs = unpack!(else_block = this.as_local_operand(else_block, rhs));
                this.cfg.push_assign(
                    else_block,
                    source_info,
                    destination,
                    Rvalue::Use(rhs),
                );
                this.cfg.terminate(
                    else_block,
                    source_info,
                    TerminatorKind::Goto { target: join_block },
                );

                join_block.unit()
            }
            ExprKind::Loop {
                condition: opt_cond_expr,
                body,
            } => {
                // [block] --> [loop_block] -/eval. cond./-> [loop_block_end] -1-> [exit_block]
                //                  ^                               |
                //                  |                               0
                //                  |                               |
                //                  |                               v
                //           [body_block_end] <-/eval. body/-- [body_block]
                //
                // If `opt_cond_expr` is `None`, then the graph is somewhat simplified:
                //
                // [block]
                //    |
                //   [loop_block] -> [body_block] -/eval. body/-> [body_block_end]
                //    |        ^                                         |
                // false link  |                                         |
                //    |        +-----------------------------------------+
                //    +-> [diverge_cleanup]
                // The false link is required to make sure borrowck considers unwinds through the
                // body, even when the exact code in the body cannot unwind

                let loop_block = this.cfg.start_new_block();
                let exit_block = this.cfg.start_new_block();

                // start the loop
                this.cfg.terminate(
                    block,
                    source_info,
                    TerminatorKind::Goto { target: loop_block },
                );

                this.in_breakable_scope(
                    Some(loop_block),
                    exit_block,
                    destination.clone(),
                    move |this| {
                        // conduct the test, if necessary
                        let body_block;
                        if let Some(cond_expr) = opt_cond_expr {
                            let loop_block_end;
                            let cond = unpack!(
                                loop_block_end = this.as_local_operand(loop_block, cond_expr)
                            );
                            body_block = this.cfg.start_new_block();
                            let term =
                                TerminatorKind::if_(this.hir.tcx(), cond, body_block, exit_block);
                            this.cfg.terminate(loop_block_end, source_info, term);

                            // if the test is false, there's no `break` to assign `destination`, so
                            // we have to do it; this overwrites any `break`-assigned value but it's
                            // always `()` anyway
                            this.cfg
                                .push_assign_unit(exit_block, source_info, destination);
                        } else {
                            body_block = this.cfg.start_new_block();
                            let diverge_cleanup = this.diverge_cleanup();
                            this.cfg.terminate(
                                loop_block,
                                source_info,
                                TerminatorKind::FalseUnwind {
                                    real_target: body_block,
                                    unwind: Some(diverge_cleanup),
                                },
                            )
                        }

                        // The “return” value of the loop body must always be an unit. We therefore
                        // introduce a unit temporary as the destination for the loop body.
                        let tmp = this.get_unit_temp();
                        // Execute the body, branching back to the test.
                        let body_block_end = unpack!(this.into(&tmp, body_block, body));
                        this.cfg.terminate(
                            body_block_end,
                            source_info,
                            TerminatorKind::Goto { target: loop_block },
                        );
                    },
                );
                exit_block.unit()
            }
            ExprKind::Call { ty, fun, args, from_hir_call } => {
                // FIXME(canndrew): This is_never should probably be an is_uninhabited
                let diverges = expr.ty.is_never();
                let intrinsic = match ty.sty {
                    ty::FnDef(def_id, _) => {
                        let f = ty.fn_sig(this.hir.tcx());
                        if f.abi() == Abi::RustIntrinsic || f.abi() == Abi::PlatformIntrinsic {
                            Some(this.hir.tcx().item_name(def_id).as_str())
                        } else {
                            None
                        }
                    }
                    _ => None,
                };
                let intrinsic = intrinsic.as_ref().map(|s| &s[..]);
                let fun = unpack!(block = this.as_local_operand(block, fun));
                if intrinsic == Some("move_val_init") {
                    // `move_val_init` has "magic" semantics - the second argument is
                    // always evaluated "directly" into the first one.

                    let mut args = args.into_iter();
                    let ptr = args.next().expect("0 arguments to `move_val_init`");
                    let val = args.next().expect("1 argument to `move_val_init`");
                    assert!(args.next().is_none(), ">2 arguments to `move_val_init`");

                    let ptr = this.hir.mirror(ptr);
                    let ptr_ty = ptr.ty;
                    // Create an *internal* temp for the pointer, so that unsafety
                    // checking won't complain about the raw pointer assignment.
                    let ptr_temp = this.local_decls.push(LocalDecl {
                        mutability: Mutability::Mut,
                        ty: ptr_ty,
                        user_ty: UserTypeProjections::none(),
                        name: None,
                        source_info,
                        visibility_scope: source_info.scope,
                        internal: true,
                        is_user_variable: None,
                        is_block_tail: None,
                    });
                    let ptr_temp = Place::Local(ptr_temp);
                    let block = unpack!(this.into(&ptr_temp, block, ptr));
                    this.into(&ptr_temp.deref(), block, val)
                } else {
                    let args: Vec<_> = args
                        .into_iter()
                        .map(|arg| unpack!(block = this.as_local_operand(block, arg)))
                        .collect();

                    let success = this.cfg.start_new_block();
                    let cleanup = this.diverge_cleanup();
                    this.cfg.terminate(
                        block,
                        source_info,
                        TerminatorKind::Call {
                            func: fun,
                            args,
                            cleanup: Some(cleanup),
                            destination: if diverges {
                                None
                            } else {
                                Some((destination.clone(), success))
                            },
                            from_hir_call,
                        },
                    );
                    success.unit()
                }
            }

            // These cases don't actually need a destination
            ExprKind::Assign { .. }
            | ExprKind::AssignOp { .. }
            | ExprKind::Continue { .. }
            | ExprKind::Break { .. }
            | ExprKind::InlineAsm { .. }
            | ExprKind::Return { .. } => {
                unpack!(block = this.stmt_expr(block, expr, None));
                this.cfg.push_assign_unit(block, source_info, destination);
                block.unit()
            }

            // Avoid creating a temporary
            ExprKind::VarRef { .. } |
            ExprKind::SelfRef |
            ExprKind::StaticRef { .. } |
            ExprKind::PlaceTypeAscription { .. } |
            ExprKind::ValueTypeAscription { .. } => {
                debug_assert!(Category::of(&expr.kind) == Some(Category::Place));

                let place = unpack!(block = this.as_place(block, expr));
                let rvalue = Rvalue::Use(this.consume_by_copy_or_move(place));
                this.cfg
                    .push_assign(block, source_info, destination, rvalue);
                block.unit()
            }
            ExprKind::Index { .. } | ExprKind::Deref { .. } | ExprKind::Field { .. } => {
                debug_assert!(Category::of(&expr.kind) == Some(Category::Place));

                // Create a "fake" temporary variable so that we check that the
                // value is Sized. Usually, this is caught in type checking, but
                // in the case of box expr there is no such check.
                if let Place::Projection(..) = destination {
                    this.local_decls
                        .push(LocalDecl::new_temp(expr.ty, expr.span));
                }

                debug_assert!(Category::of(&expr.kind) == Some(Category::Place));

                let place = unpack!(block = this.as_place(block, expr));
                let rvalue = Rvalue::Use(this.consume_by_copy_or_move(place));
                this.cfg
                    .push_assign(block, source_info, destination, rvalue);
                block.unit()
            }

            // these are the cases that are more naturally handled by some other mode
            ExprKind::Unary { .. }
            | ExprKind::Binary { .. }
            | ExprKind::Box { .. }
            | ExprKind::Cast { .. }
            | ExprKind::Use { .. }
            | ExprKind::ReifyFnPointer { .. }
            | ExprKind::ClosureFnPointer { .. }
            | ExprKind::UnsafeFnPointer { .. }
            | ExprKind::Unsize { .. }
            | ExprKind::Repeat { .. }
            | ExprKind::Borrow { .. }
            | ExprKind::Array { .. }
            | ExprKind::Tuple { .. }
            | ExprKind::Adt { .. }
            | ExprKind::Closure { .. }
            | ExprKind::Literal { .. }
            | ExprKind::Yield { .. } => {
                debug_assert!(match Category::of(&expr.kind).unwrap() {
                    // should be handled above
                    Category::Rvalue(RvalueFunc::Into) => false,

                    // must be handled above or else we get an
                    // infinite loop in the builder; see
                    // e.g., `ExprKind::VarRef` above
                    Category::Place => false,

                    _ => true,
                });

                let rvalue = unpack!(block = this.as_local_rvalue(block, expr));
                this.cfg
                    .push_assign(block, source_info, destination, rvalue);
                block.unit()
            }
        };

        if !expr_is_block_or_scope {
            let popped = this.block_context.pop();
            assert!(popped.is_some());
        }

        block_and
    }
}
<|endoftext|>


<commit_before>#![deny(trivial_casts, trivial_numeric_casts, unused_import_braces)]

use std::collections::BTreeMap;
use std::fmt::{self, Display, Formatter};
use std::io;
use std::net::IpAddr;
use std::path::Path;

use serde::{de, Deserialize};

#[cfg(feature = "mmap")]
use memmap2::Mmap;
#[cfg(feature = "mmap")]
use memmap2::MmapOptions;
#[cfg(feature = "mmap")]
use std::fs::File;

#[derive(Debug, PartialEq)]
pub enum MaxMindDBError {
    AddressNotFoundError(String),
    InvalidDatabaseError(String),
    IoError(String),
    MapError(String),
    DecodingError(String),
}

impl From<io::Error> for MaxMindDBError {
    fn from(err: io::Error) -> MaxMindDBError {
        // clean up and clean up MaxMindDBError generally
        MaxMindDBError::IoError(err.to_string())
    }
}

impl Display for MaxMindDBError {
    fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), fmt::Error> {
        match self {
            MaxMindDBError::AddressNotFoundError(msg) => {
                write!(fmt, "AddressNotFoundError: {}", msg)?
            }
            MaxMindDBError::InvalidDatabaseError(msg) => {
                write!(fmt, "InvalidDatabaseError: {}", msg)?
            }
            MaxMindDBError::IoError(msg) => write!(fmt, "IoError: {}", msg)?,
            MaxMindDBError::MapError(msg) => write!(fmt, "MapError: {}", msg)?,
            MaxMindDBError::DecodingError(msg) => write!(fmt, "DecodingError: {}", msg)?,
        }
        Ok(())
    }
}

// Use default implementation for `std::error::Error`
impl std::error::Error for MaxMindDBError {}

impl de::Error for MaxMindDBError {
    fn custom<T: Display>(msg: T) -> Self {
        MaxMindDBError::DecodingError(format!("{}", msg))
    }
}

#[derive(Deserialize, Debug)]
pub struct Metadata {
    pub binary_format_major_version: u16,
    pub binary_format_minor_version: u16,
    pub build_epoch: u64,
    pub database_type: String,
    pub description: BTreeMap<String, String>,
    pub ip_version: u16,
    pub languages: Vec<String>,
    pub node_count: u32,
    pub record_size: u16,
}

/// A reader for the MaxMind DB format. The lifetime `'data` is tied to the lifetime of the underlying buffer holding the contents of the database file.
pub struct Reader<S: AsRef<[u8]>> {
    buf: S,
    pub metadata: Metadata,
    ipv4_start: usize,
    pointer_base: usize,
}

#[cfg(feature = "mmap")]
impl<'de> Reader<Mmap> {
    /// Open a MaxMind DB database file by memory mapping it.
    ///
    /// # Example
    ///
    /// ```
    /// let reader = maxminddb::Reader::open_mmap("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
    /// ```
    pub fn open_mmap<P: AsRef<Path>>(database: P) -> Result<Reader<Mmap>, MaxMindDBError> {
        let file_read = File::open(database)?;
        let mmap = unsafe { MmapOptions::new().map(&file_read) }?;
        Reader::from_source(mmap)
    }
}

impl<'de> Reader<Vec<u8>> {
    /// Open a MaxMind DB database file by loading it into memory.
    ///
    /// # Example
    ///
    /// ```
    /// let reader = maxminddb::Reader::open_readfile("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
    /// ```
    pub fn open_readfile<P: AsRef<Path>>(database: P) -> Result<Reader<Vec<u8>>, MaxMindDBError> {
        use std::fs;

        let buf: Vec<u8> = fs::read(&database)?;
        Reader::from_source(buf)
    }
}

impl<'de, S: AsRef<[u8]>> Reader<S> {
    /// Open a MaxMind DB database from anything that implements AsRef<[u8]>
    ///
    /// # Example
    ///
    /// ```
    /// use std::fs;
    /// let buf = fs::read("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
    /// let reader = maxminddb::Reader::from_source(buf).unwrap();
    /// ```
    pub fn from_source(buf: S) -> Result<Reader<S>, MaxMindDBError> {
        let data_section_separator_size = 16;

        let metadata_start = find_metadata_start(buf.as_ref())?;
        let mut type_decoder = decoder::Decoder::new(&buf.as_ref()[metadata_start..], 0);
        let metadata = Metadata::deserialize(&mut type_decoder)?;

        let search_tree_size = (metadata.node_count as usize) * (metadata.record_size as usize) / 4;

        let mut reader = Reader {
            buf,
            pointer_base: search_tree_size + data_section_separator_size,
            metadata,
            ipv4_start: 0,
        };
        reader.ipv4_start = reader.find_ipv4_start()?;

        Ok(reader)
    }

    /// Lookup the socket address in the opened MaxMind DB
    ///
    /// Example:
    ///
    /// ```
    /// use maxminddb::geoip2;
    /// use std::net::IpAddr;
    /// use std::str::FromStr;
    ///
    /// let reader = maxminddb::Reader::open_readfile("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
    ///
    /// let ip: IpAddr = FromStr::from_str("89.160.20.128").unwrap();
    /// let city: geoip2::City = reader.lookup(ip).unwrap();
    /// print!("{:?}", city);
    /// ```
    pub fn lookup<T>(&'de self, address: IpAddr) -> Result<T, MaxMindDBError>
    where
        T: Deserialize<'de>,
    {
        let ip_bytes = ip_to_bytes(address);
        let pointer = self.find_address_in_tree(&ip_bytes)?;
        if pointer == 0 {
            return Err(MaxMindDBError::AddressNotFoundError(
                "Address not found in database".to_owned(),
            ));
        }
        let rec = self.resolve_data_pointer(pointer)?;
        let mut decoder = decoder::Decoder::new(&self.buf.as_ref()[self.pointer_base..], rec);

        T::deserialize(&mut decoder)
    }

    fn find_address_in_tree(&self, ip_address: &[u8]) -> Result<usize, MaxMindDBError> {
        let bit_count = ip_address.len() * 8;
        let mut node = self.start_node(bit_count);

        let node_count = self.metadata.node_count as usize;

        for i in 0..bit_count {
            if node >= node_count {
                break;
            }
            let bit = 1 & (ip_address[i >> 3] >> (7 - (i % 8)));

            node = self.read_node(node, bit as usize)?;
        }
        match node_count {
            n if n == node => Ok(0),
            n if node > n => Ok(node),
            _ => Err(MaxMindDBError::InvalidDatabaseError(
                "invalid node in search tree".to_owned(),
            )),
        }
    }

    fn start_node(&self, length: usize) -> usize {
        if length == 128 {
            0
        } else {
            self.ipv4_start
        }
    }

    fn find_ipv4_start(&self) -> Result<usize, MaxMindDBError> {
        if self.metadata.ip_version != 6 {
            return Ok(0);
        }

        // We are looking up an IPv4 address in an IPv6 tree. Skip over the
        // first 96 nodes.
        let mut node: usize = 0_usize;
        for _ in 0_u8..96 {
            if node >= self.metadata.node_count as usize {
                break;
            }
            node = self.read_node(node, 0)?;
        }
        Ok(node)
    }

    fn read_node(&self, node_number: usize, index: usize) -> Result<usize, MaxMindDBError> {
        let buf = self.buf.as_ref();
        let base_offset = node_number * (self.metadata.record_size as usize) / 4;

        let val = match self.metadata.record_size {
            24 => {
                let offset = base_offset + index * 3;
                to_usize(0, &buf[offset..offset + 3])
            }
            28 => {
                let mut middle = buf[base_offset + 3];
                if index != 0 {
                    middle &= 0x0F
                } else {
                    middle = (0xF0 & middle) >> 4
                }
                let offset = base_offset + index * 4;
                to_usize(middle, &buf[offset..offset + 3])
            }
            32 => {
                let offset = base_offset + index * 4;
                to_usize(0, &buf[offset..offset + 4])
            }
            s => {
                return Err(MaxMindDBError::InvalidDatabaseError(format!(
                    "unknown record size: \
                     {:?}",
                    s
                )))
            }
        };
        Ok(val)
    }

    fn resolve_data_pointer(&self, pointer: usize) -> Result<usize, MaxMindDBError> {
        let resolved = pointer - (self.metadata.node_count as usize) - 16;

        if resolved > self.buf.as_ref().len() {
            return Err(MaxMindDBError::InvalidDatabaseError(
                "the MaxMind DB file's search tree \
                 is corrupt"
                    .to_owned(),
            ));
        }

        Ok(resolved)
    }
}

// I haven't moved all patterns of this form to a generic function as
// the FromPrimitive trait is unstable
fn to_usize(base: u8, bytes: &[u8]) -> usize {
    bytes
        .iter()
        .fold(base as usize, |acc, &b| (acc << 8) | b as usize)
}

fn ip_to_bytes(address: IpAddr) -> Vec<u8> {
    match address {
        IpAddr::V4(a) => a.octets().to_vec(),
        IpAddr::V6(a) => a.octets().to_vec(),
    }
}

fn find_metadata_start(buf: &[u8]) -> Result<usize, MaxMindDBError> {
    const METADATA_START_MARKER: &[u8] = b"\xab\xcd\xefMaxMind.com";

    memchr::memmem::rfind(&buf, METADATA_START_MARKER)
        .map(|x| x + METADATA_START_MARKER.len())
        .ok_or_else(|| {
            MaxMindDBError::InvalidDatabaseError(
                "Could not find MaxMind DB metadata in file.".to_owned(),
            )
        })
}

mod decoder;
pub mod geoip2;

#[cfg(test)]
mod reader_test;

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

    #[test]
    fn test_error_display() {
        assert_eq!(
            format!(
                "{}",
                MaxMindDBError::AddressNotFoundError("something went wrong".to_owned())
            ),
            "AddressNotFoundError: something went wrong".to_owned(),
        );
        assert_eq!(
            format!(
                "{}",
                MaxMindDBError::InvalidDatabaseError("something went wrong".to_owned())
            ),
            "InvalidDatabaseError: something went wrong".to_owned(),
        );
        assert_eq!(
            format!(
                "{}",
                MaxMindDBError::IoError("something went wrong".to_owned())
            ),
            "IoError: something went wrong".to_owned(),
        );
        assert_eq!(
            format!(
                "{}",
                MaxMindDBError::MapError("something went wrong".to_owned())
            ),
            "MapError: something went wrong".to_owned(),
        );
        assert_eq!(
            format!(
                "{}",
                MaxMindDBError::DecodingError("something went wrong".to_owned())
            ),
            "DecodingError: something went wrong".to_owned(),
        );
    }
}
<commit_msg>Reexport `Mmap`<commit_after>#![deny(trivial_casts, trivial_numeric_casts, unused_import_braces)]

use std::collections::BTreeMap;
use std::fmt::{self, Display, Formatter};
use std::io;
use std::net::IpAddr;
use std::path::Path;

use serde::{de, Deserialize};

#[cfg(feature = "mmap")]
pub use memmap2::Mmap;
#[cfg(feature = "mmap")]
use memmap2::MmapOptions;
#[cfg(feature = "mmap")]
use std::fs::File;

#[derive(Debug, PartialEq)]
pub enum MaxMindDBError {
    AddressNotFoundError(String),
    InvalidDatabaseError(String),
    IoError(String),
    MapError(String),
    DecodingError(String),
}

impl From<io::Error> for MaxMindDBError {
    fn from(err: io::Error) -> MaxMindDBError {
        // clean up and clean up MaxMindDBError generally
        MaxMindDBError::IoError(err.to_string())
    }
}

impl Display for MaxMindDBError {
    fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), fmt::Error> {
        match self {
            MaxMindDBError::AddressNotFoundError(msg) => {
                write!(fmt, "AddressNotFoundError: {}", msg)?
            }
            MaxMindDBError::InvalidDatabaseError(msg) => {
                write!(fmt, "InvalidDatabaseError: {}", msg)?
            }
            MaxMindDBError::IoError(msg) => write!(fmt, "IoError: {}", msg)?,
            MaxMindDBError::MapError(msg) => write!(fmt, "MapError: {}", msg)?,
            MaxMindDBError::DecodingError(msg) => write!(fmt, "DecodingError: {}", msg)?,
        }
        Ok(())
    }
}

// Use default implementation for `std::error::Error`
impl std::error::Error for MaxMindDBError {}

impl de::Error for MaxMindDBError {
    fn custom<T: Display>(msg: T) -> Self {
        MaxMindDBError::DecodingError(format!("{}", msg))
    }
}

#[derive(Deserialize, Debug)]
pub struct Metadata {
    pub binary_format_major_version: u16,
    pub binary_format_minor_version: u16,
    pub build_epoch: u64,
    pub database_type: String,
    pub description: BTreeMap<String, String>,
    pub ip_version: u16,
    pub languages: Vec<String>,
    pub node_count: u32,
    pub record_size: u16,
}

/// A reader for the MaxMind DB format. The lifetime `'data` is tied to the lifetime of the underlying buffer holding the contents of the database file.
pub struct Reader<S: AsRef<[u8]>> {
    buf: S,
    pub metadata: Metadata,
    ipv4_start: usize,
    pointer_base: usize,
}

#[cfg(feature = "mmap")]
impl<'de> Reader<Mmap> {
    /// Open a MaxMind DB database file by memory mapping it.
    ///
    /// # Example
    ///
    /// ```
    /// let reader = maxminddb::Reader::open_mmap("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
    /// ```
    pub fn open_mmap<P: AsRef<Path>>(database: P) -> Result<Reader<Mmap>, MaxMindDBError> {
        let file_read = File::open(database)?;
        let mmap = unsafe { MmapOptions::new().map(&file_read) }?;
        Reader::from_source(mmap)
    }
}

impl<'de> Reader<Vec<u8>> {
    /// Open a MaxMind DB database file by loading it into memory.
    ///
    /// # Example
    ///
    /// ```
    /// let reader = maxminddb::Reader::open_readfile("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
    /// ```
    pub fn open_readfile<P: AsRef<Path>>(database: P) -> Result<Reader<Vec<u8>>, MaxMindDBError> {
        use std::fs;

        let buf: Vec<u8> = fs::read(&database)?;
        Reader::from_source(buf)
    }
}

impl<'de, S: AsRef<[u8]>> Reader<S> {
    /// Open a MaxMind DB database from anything that implements AsRef<[u8]>
    ///
    /// # Example
    ///
    /// ```
    /// use std::fs;
    /// let buf = fs::read("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
    /// let reader = maxminddb::Reader::from_source(buf).unwrap();
    /// ```
    pub fn from_source(buf: S) -> Result<Reader<S>, MaxMindDBError> {
        let data_section_separator_size = 16;

        let metadata_start = find_metadata_start(buf.as_ref())?;
        let mut type_decoder = decoder::Decoder::new(&buf.as_ref()[metadata_start..], 0);
        let metadata = Metadata::deserialize(&mut type_decoder)?;

        let search_tree_size = (metadata.node_count as usize) * (metadata.record_size as usize) / 4;

        let mut reader = Reader {
            buf,
            pointer_base: search_tree_size + data_section_separator_size,
            metadata,
            ipv4_start: 0,
        };
        reader.ipv4_start = reader.find_ipv4_start()?;

        Ok(reader)
    }

    /// Lookup the socket address in the opened MaxMind DB
    ///
    /// Example:
    ///
    /// ```
    /// use maxminddb::geoip2;
    /// use std::net::IpAddr;
    /// use std::str::FromStr;
    ///
    /// let reader = maxminddb::Reader::open_readfile("test-data/test-data/GeoIP2-City-Test.mmdb").unwrap();
    ///
    /// let ip: IpAddr = FromStr::from_str("89.160.20.128").unwrap();
    /// let city: geoip2::City = reader.lookup(ip).unwrap();
    /// print!("{:?}", city);
    /// ```
    pub fn lookup<T>(&'de self, address: IpAddr) -> Result<T, MaxMindDBError>
    where
        T: Deserialize<'de>,
    {
        let ip_bytes = ip_to_bytes(address);
        let pointer = self.find_address_in_tree(&ip_bytes)?;
        if pointer == 0 {
            return Err(MaxMindDBError::AddressNotFoundError(
                "Address not found in database".to_owned(),
            ));
        }
        let rec = self.resolve_data_pointer(pointer)?;
        let mut decoder = decoder::Decoder::new(&self.buf.as_ref()[self.pointer_base..], rec);

        T::deserialize(&mut decoder)
    }

    fn find_address_in_tree(&self, ip_address: &[u8]) -> Result<usize, MaxMindDBError> {
        let bit_count = ip_address.len() * 8;
        let mut node = self.start_node(bit_count);

        let node_count = self.metadata.node_count as usize;

        for i in 0..bit_count {
            if node >= node_count {
                break;
            }
            let bit = 1 & (ip_address[i >> 3] >> (7 - (i % 8)));

            node = self.read_node(node, bit as usize)?;
        }
        match node_count {
            n if n == node => Ok(0),
            n if node > n => Ok(node),
            _ => Err(MaxMindDBError::InvalidDatabaseError(
                "invalid node in search tree".to_owned(),
            )),
        }
    }

    fn start_node(&self, length: usize) -> usize {
        if length == 128 {
            0
        } else {
            self.ipv4_start
        }
    }

    fn find_ipv4_start(&self) -> Result<usize, MaxMindDBError> {
        if self.metadata.ip_version != 6 {
            return Ok(0);
        }

        // We are looking up an IPv4 address in an IPv6 tree. Skip over the
        // first 96 nodes.
        let mut node: usize = 0_usize;
        for _ in 0_u8..96 {
            if node >= self.metadata.node_count as usize {
                break;
            }
            node = self.read_node(node, 0)?;
        }
        Ok(node)
    }

    fn read_node(&self, node_number: usize, index: usize) -> Result<usize, MaxMindDBError> {
        let buf = self.buf.as_ref();
        let base_offset = node_number * (self.metadata.record_size as usize) / 4;

        let val = match self.metadata.record_size {
            24 => {
                let offset = base_offset + index * 3;
                to_usize(0, &buf[offset..offset + 3])
            }
            28 => {
                let mut middle = buf[base_offset + 3];
                if index != 0 {
                    middle &= 0x0F
                } else {
                    middle = (0xF0 & middle) >> 4
                }
                let offset = base_offset + index * 4;
                to_usize(middle, &buf[offset..offset + 3])
            }
            32 => {
                let offset = base_offset + index * 4;
                to_usize(0, &buf[offset..offset + 4])
            }
            s => {
                return Err(MaxMindDBError::InvalidDatabaseError(format!(
                    "unknown record size: \
                     {:?}",
                    s
                )))
            }
        };
        Ok(val)
    }

    fn resolve_data_pointer(&self, pointer: usize) -> Result<usize, MaxMindDBError> {
        let resolved = pointer - (self.metadata.node_count as usize) - 16;

        if resolved > self.buf.as_ref().len() {
            return Err(MaxMindDBError::InvalidDatabaseError(
                "the MaxMind DB file's search tree \
                 is corrupt"
                    .to_owned(),
            ));
        }

        Ok(resolved)
    }
}

// I haven't moved all patterns of this form to a generic function as
// the FromPrimitive trait is unstable
fn to_usize(base: u8, bytes: &[u8]) -> usize {
    bytes
        .iter()
        .fold(base as usize, |acc, &b| (acc << 8) | b as usize)
}

fn ip_to_bytes(address: IpAddr) -> Vec<u8> {
    match address {
        IpAddr::V4(a) => a.octets().to_vec(),
        IpAddr::V6(a) => a.octets().to_vec(),
    }
}

fn find_metadata_start(buf: &[u8]) -> Result<usize, MaxMindDBError> {
    const METADATA_START_MARKER: &[u8] = b"\xab\xcd\xefMaxMind.com";

    memchr::memmem::rfind(&buf, METADATA_START_MARKER)
        .map(|x| x + METADATA_START_MARKER.len())
        .ok_or_else(|| {
            MaxMindDBError::InvalidDatabaseError(
                "Could not find MaxMind DB metadata in file.".to_owned(),
            )
        })
}

mod decoder;
pub mod geoip2;

#[cfg(test)]
mod reader_test;

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

    #[test]
    fn test_error_display() {
        assert_eq!(
            format!(
                "{}",
                MaxMindDBError::AddressNotFoundError("something went wrong".to_owned())
            ),
            "AddressNotFoundError: something went wrong".to_owned(),
        );
        assert_eq!(
            format!(
                "{}",
                MaxMindDBError::InvalidDatabaseError("something went wrong".to_owned())
            ),
            "InvalidDatabaseError: something went wrong".to_owned(),
        );
        assert_eq!(
            format!(
                "{}",
                MaxMindDBError::IoError("something went wrong".to_owned())
            ),
            "IoError: something went wrong".to_owned(),
        );
        assert_eq!(
            format!(
                "{}",
                MaxMindDBError::MapError("something went wrong".to_owned())
            ),
            "MapError: something went wrong".to_owned(),
        );
        assert_eq!(
            format!(
                "{}",
                MaxMindDBError::DecodingError("something went wrong".to_owned())
            ),
            "DecodingError: something went wrong".to_owned(),
        );
    }
}
<|endoftext|>