<commit_before>// Copyright 2012 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.

// Detecting language items.
//
// Language items are items that represent concepts intrinsic to the language
// itself. Examples are:
//
// * Traits that specify "kinds"; e.g. "Freeze", "Send".
//
// * Traits that represent operators; e.g. "Add", "Sub", "Index".
//
// * Functions called by the compiler itself.

use driver::session::Session;
use metadata::csearch::each_lang_item;
use metadata::cstore::iter_crate_data;
use middle::ty::{BuiltinBound, BoundFreeze, BoundSend, BoundSized};
use syntax::ast;
use syntax::ast_util::local_def;
use syntax::attr::AttrMetaMethods;
use syntax::visit;
use syntax::visit::Visitor;

use std::hashmap::HashMap;
use std::iter::Enumerate;
use std::vec;

pub enum LangItem {
    FreezeTraitLangItem,               // 0
    SendTraitLangItem,                 // 1
    SizedTraitLangItem,                // 2

    DropTraitLangItem,                 // 3

    AddTraitLangItem,                  // 4
    SubTraitLangItem,                  // 5
    MulTraitLangItem,                  // 6
    DivTraitLangItem,                  // 7
    RemTraitLangItem,                  // 8
    NegTraitLangItem,                  // 9
    NotTraitLangItem,                  // 10
    BitXorTraitLangItem,               // 11
    BitAndTraitLangItem,               // 12
    BitOrTraitLangItem,                // 13
    ShlTraitLangItem,                  // 14
    ShrTraitLangItem,                  // 15
    IndexTraitLangItem,                // 16

    EqTraitLangItem,                   // 17
    OrdTraitLangItem,                  // 18

    StrEqFnLangItem,                   // 19
    UniqStrEqFnLangItem,               // 20
    FailFnLangItem,                    // 21
    FailBoundsCheckFnLangItem,         // 22
    ExchangeMallocFnLangItem,          // 23
    ClosureExchangeMallocFnLangItem,   // 24
    ExchangeFreeFnLangItem,            // 25
    MallocFnLangItem,                  // 26
    FreeFnLangItem,                    // 27
    BorrowAsImmFnLangItem,             // 28
    BorrowAsMutFnLangItem,             // 29
    ReturnToMutFnLangItem,             // 30
    CheckNotBorrowedFnLangItem,        // 31
    StrDupUniqFnLangItem,              // 32
    RecordBorrowFnLangItem,            // 33
    UnrecordBorrowFnLangItem,          // 34

    StartFnLangItem,                   // 35

    TyDescStructLangItem,              // 36
    TyVisitorTraitLangItem,            // 37
    OpaqueStructLangItem,              // 38

    EventLoopFactoryLangItem,          // 39

    TypeIdLangItem,                    // 40
}

pub struct LanguageItems {
    items: [Option<ast::DefId>, ..41]
}

impl LanguageItems {
    pub fn new() -> LanguageItems {
        LanguageItems {
            items: [ None, ..41 ]
        }
    }

    pub fn items<'a>(&'a self) -> Enumerate<vec::VecIterator<'a, Option<ast::DefId>>> {
        self.items.iter().enumerate()
    }

    pub fn item_name(index: uint) -> &'static str {
        match index {
            0  => "freeze",
            1  => "send",
            2  => "sized",

            3  => "drop",

            4  => "add",
            5  => "sub",
            6  => "mul",
            7  => "div",
            8  => "rem",
            9  => "neg",
            10 => "not",
            11 => "bitxor",
            12 => "bitand",
            13 => "bitor",
            14 => "shl",
            15 => "shr",
            16 => "index",
            17 => "eq",
            18 => "ord",

            19 => "str_eq",
            20 => "uniq_str_eq",
            21 => "fail_",
            22 => "fail_bounds_check",
            23 => "exchange_malloc",
            24 => "closure_exchange_malloc",
            25 => "exchange_free",
            26 => "malloc",
            27 => "free",
            28 => "borrow_as_imm",
            29 => "borrow_as_mut",
            30 => "return_to_mut",
            31 => "check_not_borrowed",
            32 => "strdup_uniq",
            33 => "record_borrow",
            34 => "unrecord_borrow",

            35 => "start",

            36 => "ty_desc",
            37 => "ty_visitor",
            38 => "opaque",

            39 => "event_loop_factory",

            40 => "type_id",

            _ => "???"
        }
    }

    // FIXME #4621: Method macros sure would be nice here.

    pub fn require(&self, it: LangItem) -> Result<ast::DefId, ~str> {
        match self.items[it as uint] {
            Some(id) => Ok(id),
            None => Err(format!("requires `{}` lang_item",
                             LanguageItems::item_name(it as uint)))
        }
    }

    pub fn to_builtin_kind(&self, id: ast::DefId) -> Option<BuiltinBound> {
        if Some(id) == self.freeze_trait() {
            Some(BoundFreeze)
        } else if Some(id) == self.send_trait() {
            Some(BoundSend)
        } else if Some(id) == self.sized_trait() {
            Some(BoundSized)
        } else {
            None
        }
    }

    pub fn freeze_trait(&self) -> Option<ast::DefId> {
        self.items[FreezeTraitLangItem as uint]
    }
    pub fn send_trait(&self) -> Option<ast::DefId> {
        self.items[SendTraitLangItem as uint]
    }
    pub fn sized_trait(&self) -> Option<ast::DefId> {
        self.items[SizedTraitLangItem as uint]
    }

    pub fn drop_trait(&self) -> Option<ast::DefId> {
        self.items[DropTraitLangItem as uint]
    }

    pub fn add_trait(&self) -> Option<ast::DefId> {
        self.items[AddTraitLangItem as uint]
    }
    pub fn sub_trait(&self) -> Option<ast::DefId> {
        self.items[SubTraitLangItem as uint]
    }
    pub fn mul_trait(&self) -> Option<ast::DefId> {
        self.items[MulTraitLangItem as uint]
    }
    pub fn div_trait(&self) -> Option<ast::DefId> {
        self.items[DivTraitLangItem as uint]
    }
    pub fn rem_trait(&self) -> Option<ast::DefId> {
        self.items[RemTraitLangItem as uint]
    }
    pub fn neg_trait(&self) -> Option<ast::DefId> {
        self.items[NegTraitLangItem as uint]
    }
    pub fn not_trait(&self) -> Option<ast::DefId> {
        self.items[NotTraitLangItem as uint]
    }
    pub fn bitxor_trait(&self) -> Option<ast::DefId> {
        self.items[BitXorTraitLangItem as uint]
    }
    pub fn bitand_trait(&self) -> Option<ast::DefId> {
        self.items[BitAndTraitLangItem as uint]
    }
    pub fn bitor_trait(&self) -> Option<ast::DefId> {
        self.items[BitOrTraitLangItem as uint]
    }
    pub fn shl_trait(&self) -> Option<ast::DefId> {
        self.items[ShlTraitLangItem as uint]
    }
    pub fn shr_trait(&self) -> Option<ast::DefId> {
        self.items[ShrTraitLangItem as uint]
    }
    pub fn index_trait(&self) -> Option<ast::DefId> {
        self.items[IndexTraitLangItem as uint]
    }

    pub fn eq_trait(&self) -> Option<ast::DefId> {
        self.items[EqTraitLangItem as uint]
    }
    pub fn ord_trait(&self) -> Option<ast::DefId> {
        self.items[OrdTraitLangItem as uint]
    }

    pub fn str_eq_fn(&self) -> Option<ast::DefId> {
        self.items[StrEqFnLangItem as uint]
    }
    pub fn uniq_str_eq_fn(&self) -> Option<ast::DefId> {
        self.items[UniqStrEqFnLangItem as uint]
    }
    pub fn fail_fn(&self) -> Option<ast::DefId> {
        self.items[FailFnLangItem as uint]
    }
    pub fn fail_bounds_check_fn(&self) -> Option<ast::DefId> {
        self.items[FailBoundsCheckFnLangItem as uint]
    }
    pub fn exchange_malloc_fn(&self) -> Option<ast::DefId> {
        self.items[ExchangeMallocFnLangItem as uint]
    }
    pub fn closure_exchange_malloc_fn(&self) -> Option<ast::DefId> {
        self.items[ClosureExchangeMallocFnLangItem as uint]
    }
    pub fn exchange_free_fn(&self) -> Option<ast::DefId> {
        self.items[ExchangeFreeFnLangItem as uint]
    }
    pub fn malloc_fn(&self) -> Option<ast::DefId> {
        self.items[MallocFnLangItem as uint]
    }
    pub fn free_fn(&self) -> Option<ast::DefId> {
        self.items[FreeFnLangItem as uint]
    }
    pub fn borrow_as_imm_fn(&self) -> Option<ast::DefId> {
        self.items[BorrowAsImmFnLangItem as uint]
    }
    pub fn borrow_as_mut_fn(&self) -> Option<ast::DefId> {
        self.items[BorrowAsMutFnLangItem as uint]
    }
    pub fn return_to_mut_fn(&self) -> Option<ast::DefId> {
        self.items[ReturnToMutFnLangItem as uint]
    }
    pub fn check_not_borrowed_fn(&self) -> Option<ast::DefId> {
        self.items[CheckNotBorrowedFnLangItem as uint]
    }
    pub fn strdup_uniq_fn(&self) -> Option<ast::DefId> {
        self.items[StrDupUniqFnLangItem as uint]
    }
    pub fn record_borrow_fn(&self) -> Option<ast::DefId> {
        self.items[RecordBorrowFnLangItem as uint]
    }
    pub fn unrecord_borrow_fn(&self) -> Option<ast::DefId> {
        self.items[UnrecordBorrowFnLangItem as uint]
    }
    pub fn start_fn(&self) -> Option<ast::DefId> {
        self.items[StartFnLangItem as uint]
    }
    pub fn ty_desc(&self) -> Option<ast::DefId> {
        self.items[TyDescStructLangItem as uint]
    }
    pub fn ty_visitor(&self) -> Option<ast::DefId> {
        self.items[TyVisitorTraitLangItem as uint]
    }
    pub fn opaque(&self) -> Option<ast::DefId> {
        self.items[OpaqueStructLangItem as uint]
    }
    pub fn event_loop_factory(&self) -> Option<ast::DefId> {
        self.items[EventLoopFactoryLangItem as uint]
    }
    pub fn type_id(&self) -> Option<ast::DefId> {
        self.items[TypeIdLangItem as uint]
    }
}

struct LanguageItemCollector {
    items: LanguageItems,

    session: Session,

    item_refs: HashMap<&'static str, uint>,
}

struct LanguageItemVisitor<'self> {
    this: &'self mut LanguageItemCollector,
}

impl<'self> Visitor<()> for LanguageItemVisitor<'self> {
    fn visit_item(&mut self, item: @ast::item, _: ()) {
        match extract(item.attrs) {
            Some(value) => {
                let item_index = self.this.item_refs.find_equiv(&value).map(|x| *x);

                match item_index {
                    Some(item_index) => {
                        self.this.collect_item(item_index, local_def(item.id))
                    }
                    None => {}
                }
            }
            None => {}
        }

        visit::walk_item(self, item, ());
    }
}

impl LanguageItemCollector {
    pub fn new(session: Session) -> LanguageItemCollector {
        let mut item_refs = HashMap::new();

        item_refs.insert("freeze", FreezeTraitLangItem as uint);
        item_refs.insert("send", SendTraitLangItem as uint);
        item_refs.insert("sized", SizedTraitLangItem as uint);

        item_refs.insert("drop", DropTraitLangItem as uint);

        item_refs.insert("add", AddTraitLangItem as uint);
        item_refs.insert("sub", SubTraitLangItem as uint);
        item_refs.insert("mul", MulTraitLangItem as uint);
        item_refs.insert("div", DivTraitLangItem as uint);
        item_refs.insert("rem", RemTraitLangItem as uint);
        item_refs.insert("neg", NegTraitLangItem as uint);
        item_refs.insert("not", NotTraitLangItem as uint);
        item_refs.insert("bitxor", BitXorTraitLangItem as uint);
        item_refs.insert("bitand", BitAndTraitLangItem as uint);
        item_refs.insert("bitor", BitOrTraitLangItem as uint);
        item_refs.insert("shl", ShlTraitLangItem as uint);
        item_refs.insert("shr", ShrTraitLangItem as uint);
        item_refs.insert("index", IndexTraitLangItem as uint);

        item_refs.insert("eq", EqTraitLangItem as uint);
        item_refs.insert("ord", OrdTraitLangItem as uint);

        item_refs.insert("str_eq", StrEqFnLangItem as uint);
        item_refs.insert("uniq_str_eq", UniqStrEqFnLangItem as uint);
        item_refs.insert("fail_", FailFnLangItem as uint);
        item_refs.insert("fail_bounds_check",
                         FailBoundsCheckFnLangItem as uint);
        item_refs.insert("exchange_malloc", ExchangeMallocFnLangItem as uint);
        item_refs.insert("closure_exchange_malloc", ClosureExchangeMallocFnLangItem as uint);
        item_refs.insert("exchange_free", ExchangeFreeFnLangItem as uint);
        item_refs.insert("malloc", MallocFnLangItem as uint);
        item_refs.insert("free", FreeFnLangItem as uint);
        item_refs.insert("borrow_as_imm", BorrowAsImmFnLangItem as uint);
        item_refs.insert("borrow_as_mut", BorrowAsMutFnLangItem as uint);
        item_refs.insert("return_to_mut", ReturnToMutFnLangItem as uint);
        item_refs.insert("check_not_borrowed",
                         CheckNotBorrowedFnLangItem as uint);
        item_refs.insert("strdup_uniq", StrDupUniqFnLangItem as uint);
        item_refs.insert("record_borrow", RecordBorrowFnLangItem as uint);
        item_refs.insert("unrecord_borrow", UnrecordBorrowFnLangItem as uint);
        item_refs.insert("start", StartFnLangItem as uint);
        item_refs.insert("ty_desc", TyDescStructLangItem as uint);
        item_refs.insert("ty_visitor", TyVisitorTraitLangItem as uint);
        item_refs.insert("opaque", OpaqueStructLangItem as uint);
        item_refs.insert("event_loop_factory", EventLoopFactoryLangItem as uint);
        item_refs.insert("type_id", TypeIdLangItem as uint);

        LanguageItemCollector {
            session: session,
            items: LanguageItems::new(),
            item_refs: item_refs
        }
    }

    pub fn collect_item(&mut self, item_index: uint, item_def_id: ast::DefId) {
        // Check for duplicates.
        match self.items.items[item_index] {
            Some(original_def_id) if original_def_id != item_def_id => {
                self.session.err(format!("duplicate entry for `{}`",
                                      LanguageItems::item_name(item_index)));
            }
            Some(_) | None => {
                // OK.
            }
        }

        // Matched.
        self.items.items[item_index] = Some(item_def_id);
    }

    pub fn collect_local_language_items(&mut self, crate: &ast::Crate) {
        let mut v = LanguageItemVisitor { this: self };
        visit::walk_crate(&mut v, crate, ());
    }

    pub fn collect_external_language_items(&mut self) {
        let crate_store = self.session.cstore;
        iter_crate_data(crate_store, |crate_number, _crate_metadata| {
            each_lang_item(crate_store, crate_number, |node_id, item_index| {
                let def_id = ast::DefId { crate: crate_number, node: node_id };
                self.collect_item(item_index, def_id);
                true
            });
        })
    }

    pub fn collect(&mut self, crate: &ast::Crate) {
        self.collect_local_language_items(crate);
        self.collect_external_language_items();
    }
}

pub fn extract(attrs: &[ast::Attribute]) -> Option<@str> {
    for attribute in attrs.iter() {
        match attribute.name_str_pair() {
            Some((key, value)) if "lang" == key => {
                return Some(value);
            }
            Some(..) | None => {}
        }
    }

    return None;
}

pub fn collect_language_items(crate: &ast::Crate,
                              session: Session)
                           -> LanguageItems {
    let mut collector = LanguageItemCollector::new(session);
    collector.collect(crate);
    let LanguageItemCollector { items, .. } = collector;
    session.abort_if_errors();
    items
}
<commit_msg>Use a macro in lang_items to remove duplication.<commit_after>// Copyright 2012 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.

// Detecting language items.
//
// Language items are items that represent concepts intrinsic to the language
// itself. Examples are:
//
// * Traits that specify "kinds"; e.g. "Freeze", "Send".
//
// * Traits that represent operators; e.g. "Add", "Sub", "Index".
//
// * Functions called by the compiler itself.

use driver::session::Session;
use metadata::csearch::each_lang_item;
use metadata::cstore::iter_crate_data;
use middle::ty::{BuiltinBound, BoundFreeze, BoundSend, BoundSized};
use syntax::ast;
use syntax::ast_util::local_def;
use syntax::attr::AttrMetaMethods;
use syntax::visit;
use syntax::visit::Visitor;

use std::hashmap::HashMap;
use std::iter::Enumerate;
use std::vec;

// The actual lang items defined come at the end of this file in one handy table.
// So you probably just want to nip down to the end.
macro_rules! lets_do_this {
    (
        There are $num_lang_items:expr lang items.
        $( $num:pat, $variant:ident, $name:expr, $method:ident; )*
    ) => {

pub enum LangItem {
    $($variant),*
}

pub struct LanguageItems {
    items: [Option<ast::DefId>, ..$num_lang_items]
}

impl LanguageItems {
    pub fn new() -> LanguageItems {
        LanguageItems {
            items: [ None, ..$num_lang_items ]
        }
    }

    pub fn items<'a>(&'a self) -> Enumerate<vec::VecIterator<'a, Option<ast::DefId>>> {
        self.items.iter().enumerate()
    }

    pub fn item_name(index: uint) -> &'static str {
        match index {
            $( $num => $name, )*
            _ => "???"
        }
    }

    pub fn require(&self, it: LangItem) -> Result<ast::DefId, ~str> {
        match self.items[it as uint] {
            Some(id) => Ok(id),
            None => Err(format!("requires `{}` lang_item",
                             LanguageItems::item_name(it as uint)))
        }
    }

    pub fn to_builtin_kind(&self, id: ast::DefId) -> Option<BuiltinBound> {
        if Some(id) == self.freeze_trait() {
            Some(BoundFreeze)
        } else if Some(id) == self.send_trait() {
            Some(BoundSend)
        } else if Some(id) == self.sized_trait() {
            Some(BoundSized)
        } else {
            None
        }
    }

    $(
        pub fn $method(&self) -> Option<ast::DefId> {
            self.items[$variant as uint]
        }
    )*
}

struct LanguageItemCollector {
    items: LanguageItems,

    session: Session,

    item_refs: HashMap<&'static str, uint>,
}

struct LanguageItemVisitor<'self> {
    this: &'self mut LanguageItemCollector,
}

impl<'self> Visitor<()> for LanguageItemVisitor<'self> {
    fn visit_item(&mut self, item: @ast::item, _: ()) {
        match extract(item.attrs) {
            Some(value) => {
                let item_index = self.this.item_refs.find_equiv(&value).map(|x| *x);

                match item_index {
                    Some(item_index) => {
                        self.this.collect_item(item_index, local_def(item.id))
                    }
                    None => {}
                }
            }
            None => {}
        }

        visit::walk_item(self, item, ());
    }
}

impl LanguageItemCollector {
    pub fn new(session: Session) -> LanguageItemCollector {
        let mut item_refs = HashMap::new();

        $( item_refs.insert($name, $variant as uint); )*

        LanguageItemCollector {
            session: session,
            items: LanguageItems::new(),
            item_refs: item_refs
        }
    }

    pub fn collect_item(&mut self, item_index: uint, item_def_id: ast::DefId) {
        // Check for duplicates.
        match self.items.items[item_index] {
            Some(original_def_id) if original_def_id != item_def_id => {
                self.session.err(format!("duplicate entry for `{}`",
                                      LanguageItems::item_name(item_index)));
            }
            Some(_) | None => {
                // OK.
            }
        }

        // Matched.
        self.items.items[item_index] = Some(item_def_id);
    }

    pub fn collect_local_language_items(&mut self, crate: &ast::Crate) {
        let mut v = LanguageItemVisitor { this: self };
        visit::walk_crate(&mut v, crate, ());
    }

    pub fn collect_external_language_items(&mut self) {
        let crate_store = self.session.cstore;
        iter_crate_data(crate_store, |crate_number, _crate_metadata| {
            each_lang_item(crate_store, crate_number, |node_id, item_index| {
                let def_id = ast::DefId { crate: crate_number, node: node_id };
                self.collect_item(item_index, def_id);
                true
            });
        })
    }

    pub fn collect(&mut self, crate: &ast::Crate) {
        self.collect_local_language_items(crate);
        self.collect_external_language_items();
    }
}

pub fn extract(attrs: &[ast::Attribute]) -> Option<@str> {
    for attribute in attrs.iter() {
        match attribute.name_str_pair() {
            Some((key, value)) if "lang" == key => {
                return Some(value);
            }
            Some(..) | None => {}
        }
    }

    return None;
}

pub fn collect_language_items(crate: &ast::Crate,
                              session: Session)
                           -> LanguageItems {
    let mut collector = LanguageItemCollector::new(session);
    collector.collect(crate);
    let LanguageItemCollector { items, .. } = collector;
    session.abort_if_errors();
    items
}

// End of the macro
    }
}

lets_do_this! {
    There are 41 lang items.

//  ID, Variant name,                    Name,                      Method name;
    0,  FreezeTraitLangItem,             "freeze",                  freeze_trait;
    1,  SendTraitLangItem,               "send",                    send_trait;
    2,  SizedTraitLangItem,              "sized",                   sized_trait;

    3,  DropTraitLangItem,               "drop",                    drop_trait;

    4,  AddTraitLangItem,                "add",                     add_trait;
    5,  SubTraitLangItem,                "sub",                     sub_trait;
    6,  MulTraitLangItem,                "mul",                     mul_trait;
    7,  DivTraitLangItem,                "div",                     div_trait;
    8,  RemTraitLangItem,                "rem",                     rem_trait;
    9,  NegTraitLangItem,                "neg",                     neg_trait;
    10, NotTraitLangItem,                "not",                     not_trait;
    11, BitXorTraitLangItem,             "bitxor",                  bitxor_trait;
    12, BitAndTraitLangItem,             "bitand",                  bitand_trait;
    13, BitOrTraitLangItem,              "bitor",                   bitor_trait;
    14, ShlTraitLangItem,                "shl",                     shl_trait;
    15, ShrTraitLangItem,                "shr",                     shr_trait;
    16, IndexTraitLangItem,              "index",                   index_trait;

    17, EqTraitLangItem,                 "eq",                      eq_trait;
    18, OrdTraitLangItem,                "ord",                     ord_trait;

    19, StrEqFnLangItem,                 "str_eq",                  str_eq_fn;
    20, UniqStrEqFnLangItem,             "uniq_str_eq",             uniq_str_eq_fn;
    21, FailFnLangItem,                  "fail_",                   fail_fn;
    22, FailBoundsCheckFnLangItem,       "fail_bounds_check",       fail_bounds_check_fn;
    23, ExchangeMallocFnLangItem,        "exchange_malloc",         exchange_malloc_fn;
    24, ClosureExchangeMallocFnLangItem, "closure_exchange_malloc", closure_exchange_malloc_fn;
    25, ExchangeFreeFnLangItem,          "exchange_free",           exchange_free_fn;
    26, MallocFnLangItem,                "malloc",                  malloc_fn;
    27, FreeFnLangItem,                  "free",                    free_fn;
    28, BorrowAsImmFnLangItem,           "borrow_as_imm",           borrow_as_imm_fn;
    29, BorrowAsMutFnLangItem,           "borrow_as_mut",           borrow_as_mut_fn;
    30, ReturnToMutFnLangItem,           "return_to_mut",           return_to_mut_fn;
    31, CheckNotBorrowedFnLangItem,      "check_not_borrowed",      check_not_borrowed_fn;
    32, StrDupUniqFnLangItem,            "strdup_uniq",             strdup_uniq_fn;
    33, RecordBorrowFnLangItem,          "record_borrow",           record_borrow_fn;
    34, UnrecordBorrowFnLangItem,        "unrecord_borrow",         unrecord_borrow_fn;

    35, StartFnLangItem,                 "start",                   start_fn;

    36, TyDescStructLangItem,            "ty_desc",                 ty_desc;
    37, TyVisitorTraitLangItem,          "ty_visitor",              ty_visitor;
    38, OpaqueStructLangItem,            "opaque",                  opaque;

    39, EventLoopFactoryLangItem,        "event_loop_factory",      event_loop_factory;

    40, TypeIdLangItem,                  "type_id",                 type_id;
}
<|endoftext|>


<commit_before>// Code that generates a test runner to run all the tests in a crate

use log::debug;
use smallvec::{smallvec, SmallVec};
use syntax::ast::{self, Ident};
use syntax::attr;
use syntax::entry::{self, EntryPointType};
use syntax::ext::base::{ExtCtxt, Resolver};
use syntax::ext::expand::{AstFragment, ExpansionConfig};
use syntax::feature_gate::Features;
use syntax::mut_visit::{*, ExpectOne};
use syntax::parse::ParseSess;
use syntax::ptr::P;
use syntax::source_map::respan;
use syntax::symbol::{sym, Symbol};
use syntax_pos::{Span, DUMMY_SP};
use syntax_pos::hygiene::{AstPass, SyntaxContext, Transparency};

use std::{iter, mem};

struct Test {
    span: Span,
    ident: Ident,
}

struct TestCtxt<'a> {
    ext_cx: ExtCtxt<'a>,
    test_cases: Vec<Test>,
    reexport_test_harness_main: Option<Symbol>,
    test_runner: Option<ast::Path>,
}

// Traverse the crate, collecting all the test functions, eliding any
// existing main functions, and synthesizing a main test harness
pub fn inject(
    sess: &ParseSess,
    resolver: &mut dyn Resolver,
    should_test: bool,
    krate: &mut ast::Crate,
    span_diagnostic: &errors::Handler,
    features: &Features,
) {
    // Check for #![reexport_test_harness_main = "some_name"] which gives the
    // main test function the name `some_name` without hygiene. This needs to be
    // unconditional, so that the attribute is still marked as used in
    // non-test builds.
    let reexport_test_harness_main =
        attr::first_attr_value_str_by_name(&krate.attrs, sym::reexport_test_harness_main);

    // Do this here so that the test_runner crate attribute gets marked as used
    // even in non-test builds
    let test_runner = get_test_runner(span_diagnostic, &krate);

    if should_test {
        generate_test_harness(sess, resolver, reexport_test_harness_main,
                              krate, features, test_runner)
    }
}

struct TestHarnessGenerator<'a> {
    cx: TestCtxt<'a>,
    tests: Vec<Test>,
}

impl<'a> MutVisitor for TestHarnessGenerator<'a> {
    fn visit_crate(&mut self, c: &mut ast::Crate) {
        noop_visit_crate(c, self);

        // Create a main function to run our tests
        c.module.items.push(mk_main(&mut self.cx));
    }

    fn flat_map_item(&mut self, i: P<ast::Item>) -> SmallVec<[P<ast::Item>; 1]> {
        let mut item = i.into_inner();
        if is_test_case(&item) {
            debug!("this is a test item");

            let test = Test {
                span: item.span,
                ident: item.ident,
            };
            self.tests.push(test);
        }

        // We don't want to recurse into anything other than mods, since
        // mods or tests inside of functions will break things
        if let ast::ItemKind::Mod(mut module) = item.node {
            let tests = mem::take(&mut self.tests);
            noop_visit_mod(&mut module, self);
            let mut tests = mem::replace(&mut self.tests, tests);

            if !tests.is_empty() {
                let parent = if item.id == ast::DUMMY_NODE_ID {
                    ast::CRATE_NODE_ID
                } else {
                    item.id
                };
                // Create an identifier that will hygienically resolve the test
                // case name, even in another module.
                let expn_id = self.cx.ext_cx.resolver.expansion_for_ast_pass(
                    module.inner,
                    AstPass::TestHarness,
                    &[],
                    Some(parent),
                );
                for test in &mut tests {
                    // See the comment on `mk_main` for why we're using
                    // `apply_mark` directly.
                    test.ident.span = test.ident.span.apply_mark(expn_id, Transparency::Opaque);
                }
                self.cx.test_cases.extend(tests);
            }
            item.node = ast::ItemKind::Mod(module);
        }
        smallvec![P(item)]
    }

    fn visit_mac(&mut self, _mac: &mut ast::Mac) {
        // Do nothing.
    }
}

/// A folder used to remove any entry points (like fn main) because the harness
/// generator will provide its own
struct EntryPointCleaner {
    // Current depth in the ast
    depth: usize,
}

impl MutVisitor for EntryPointCleaner {
    fn flat_map_item(&mut self, i: P<ast::Item>) -> SmallVec<[P<ast::Item>; 1]> {
        self.depth += 1;
        let item = noop_flat_map_item(i, self).expect_one("noop did something");
        self.depth -= 1;

        // Remove any #[main] or #[start] from the AST so it doesn't
        // clash with the one we're going to add, but mark it as
        // #[allow(dead_code)] to avoid printing warnings.
        let item = match entry::entry_point_type(&item, self.depth) {
            EntryPointType::MainNamed |
            EntryPointType::MainAttr |
            EntryPointType::Start =>
                item.map(|ast::Item {id, ident, attrs, node, vis, span, tokens}| {
                    let allow_ident = Ident::with_dummy_span(sym::allow);
                    let dc_nested = attr::mk_nested_word_item(Ident::from_str("dead_code"));
                    let allow_dead_code_item = attr::mk_list_item(allow_ident, vec![dc_nested]);
                    let allow_dead_code = attr::mk_attr_outer(allow_dead_code_item);

                    ast::Item {
                        id,
                        ident,
                        attrs: attrs.into_iter()
                            .filter(|attr| {
                                !attr.check_name(sym::main) && !attr.check_name(sym::start)
                            })
                            .chain(iter::once(allow_dead_code))
                            .collect(),
                        node,
                        vis,
                        span,
                        tokens,
                    }
                }),
            EntryPointType::None |
            EntryPointType::OtherMain => item,
        };

        smallvec![item]
    }

    fn visit_mac(&mut self, _mac: &mut ast::Mac) {
        // Do nothing.
    }
}

/// Crawl over the crate, inserting test reexports and the test main function
fn generate_test_harness(sess: &ParseSess,
                         resolver: &mut dyn Resolver,
                         reexport_test_harness_main: Option<Symbol>,
                         krate: &mut ast::Crate,
                         features: &Features,
                         test_runner: Option<ast::Path>) {
    // Remove the entry points
    let mut cleaner = EntryPointCleaner { depth: 0 };
    cleaner.visit_crate(krate);

    let mut econfig = ExpansionConfig::default("test".to_string());
    econfig.features = Some(features);

    let cx = TestCtxt {
        ext_cx: ExtCtxt::new(sess, econfig, resolver),
        test_cases: Vec::new(),
        reexport_test_harness_main,
        test_runner
    };

    TestHarnessGenerator {
        cx,
        tests: Vec::new(),
    }.visit_crate(krate);
}

/// Creates a function item for use as the main function of a test build.
/// This function will call the `test_runner` as specified by the crate attribute
fn mk_main(cx: &mut TestCtxt<'_>) -> P<ast::Item> {
    // Writing this out by hand:
    //        pub fn main() {
    //            #![main]
    //            test::test_main_static(&[..tests]);
    //        }
    let expn_id = cx.ext_cx.resolver.expansion_for_ast_pass(
        DUMMY_SP,
        AstPass::TestHarness,
        &[sym::main, sym::test, sym::rustc_attrs],
        None,
    );
    let sp = DUMMY_SP.with_def_site_ctxt(expn_id);
    let ecx = &cx.ext_cx;
    let test_id = Ident::new(sym::test, sp);

    // test::test_main_static(...)
    let mut test_runner = cx.test_runner.clone().unwrap_or(
        ecx.path(sp, vec![
            test_id, ecx.ident_of("test_main_static")
        ]));

    test_runner.span = sp;

    let test_main_path_expr = ecx.expr_path(test_runner);
    let call_test_main = ecx.expr_call(sp, test_main_path_expr,
                                       vec![mk_tests_slice(cx, sp)]);
    let call_test_main = ecx.stmt_expr(call_test_main);

    // #![main]
    let main_meta = ecx.meta_word(sp, sym::main);
    let main_attr = ecx.attribute(main_meta);

    // extern crate test
    let test_extern_stmt = ecx.stmt_item(sp, ecx.item(sp,
        test_id,
        vec![],
        ast::ItemKind::ExternCrate(None)
    ));

    // pub fn main() { ... }
    let main_ret_ty = ecx.ty(sp, ast::TyKind::Tup(vec![]));

    // If no test runner is provided we need to import the test crate
    let main_body = if cx.test_runner.is_none() {
        ecx.block(sp, vec![test_extern_stmt, call_test_main])
    } else {
        ecx.block(sp, vec![call_test_main])
    };

    let main = ast::ItemKind::Fn(ecx.fn_decl(vec![], ast::FunctionRetTy::Ty(main_ret_ty)),
                           ast::FnHeader::default(),
                           ast::Generics::default(),
                           main_body);

    // Honor the reexport_test_harness_main attribute
    let main_id = match cx.reexport_test_harness_main {
        Some(sym) => Ident::new(sym, sp.with_ctxt(SyntaxContext::root())),
        None => Ident::from_str_and_span("main", sp),
    };

    let main = P(ast::Item {
        ident: main_id,
        attrs: vec![main_attr],
        id: ast::DUMMY_NODE_ID,
        node: main,
        vis: respan(sp, ast::VisibilityKind::Public),
        span: sp,
        tokens: None,
    });

    // Integrate the new item into existing module structures.
    let main = AstFragment::Items(smallvec![main]);
    cx.ext_cx.monotonic_expander().fully_expand_fragment(main).make_items().pop().unwrap()
}

/// Creates a slice containing every test like so:
/// &[test1, test2]
fn mk_tests_slice(cx: &TestCtxt<'_>, sp: Span) -> P<ast::Expr> {
    debug!("building test vector from {} tests", cx.test_cases.len());
    let ref ecx = cx.ext_cx;

    ecx.expr_vec_slice(sp,
        cx.test_cases.iter().map(|test| {
            ecx.expr_addr_of(test.span,
                ecx.expr_path(ecx.path(test.span, vec![test.ident])))
        }).collect())
}

fn is_test_case(i: &ast::Item) -> bool {
    attr::contains_name(&i.attrs, sym::rustc_test_marker)
}

fn get_test_runner(sd: &errors::Handler, krate: &ast::Crate) -> Option<ast::Path> {
    let test_attr = attr::find_by_name(&krate.attrs, sym::test_runner)?;
    test_attr.meta_item_list().map(|meta_list| {
        if meta_list.len() != 1 {
            sd.span_fatal(test_attr.span,
                "`#![test_runner(..)]` accepts exactly 1 argument").raise()
        }
        match meta_list[0].meta_item() {
            Some(meta_item) if meta_item.is_word() => meta_item.path.clone(),
            _ => sd.span_fatal(test_attr.span, "`test_runner` argument must be a path").raise()
        }
    })
}
<commit_msg>Document test harness generation<commit_after>// Code that generates a test runner to run all the tests in a crate

use log::debug;
use smallvec::{smallvec, SmallVec};
use syntax::ast::{self, Ident};
use syntax::attr;
use syntax::entry::{self, EntryPointType};
use syntax::ext::base::{ExtCtxt, Resolver};
use syntax::ext::expand::{AstFragment, ExpansionConfig};
use syntax::feature_gate::Features;
use syntax::mut_visit::{*, ExpectOne};
use syntax::parse::ParseSess;
use syntax::ptr::P;
use syntax::source_map::respan;
use syntax::symbol::{sym, Symbol};
use syntax_pos::{Span, DUMMY_SP};
use syntax_pos::hygiene::{AstPass, SyntaxContext, Transparency};

use std::{iter, mem};

struct Test {
    span: Span,
    ident: Ident,
}

struct TestCtxt<'a> {
    ext_cx: ExtCtxt<'a>,
    def_site: Span,
    test_cases: Vec<Test>,
    reexport_test_harness_main: Option<Symbol>,
    test_runner: Option<ast::Path>,
}

// Traverse the crate, collecting all the test functions, eliding any
// existing main functions, and synthesizing a main test harness
pub fn inject(
    sess: &ParseSess,
    resolver: &mut dyn Resolver,
    should_test: bool,
    krate: &mut ast::Crate,
    span_diagnostic: &errors::Handler,
    features: &Features,
) {
    // Check for #![reexport_test_harness_main = "some_name"] which gives the
    // main test function the name `some_name` without hygiene. This needs to be
    // unconditional, so that the attribute is still marked as used in
    // non-test builds.
    let reexport_test_harness_main =
        attr::first_attr_value_str_by_name(&krate.attrs, sym::reexport_test_harness_main);

    // Do this here so that the test_runner crate attribute gets marked as used
    // even in non-test builds
    let test_runner = get_test_runner(span_diagnostic, &krate);

    if should_test {
        generate_test_harness(sess, resolver, reexport_test_harness_main,
                              krate, features, test_runner)
    }
}

struct TestHarnessGenerator<'a> {
    cx: TestCtxt<'a>,
    tests: Vec<Test>,
}

impl<'a> MutVisitor for TestHarnessGenerator<'a> {
    fn visit_crate(&mut self, c: &mut ast::Crate) {
        noop_visit_crate(c, self);

        // Create a main function to run our tests
        c.module.items.push(mk_main(&mut self.cx));
    }

    fn flat_map_item(&mut self, i: P<ast::Item>) -> SmallVec<[P<ast::Item>; 1]> {
        let mut item = i.into_inner();
        if is_test_case(&item) {
            debug!("this is a test item");

            let test = Test {
                span: item.span,
                ident: item.ident,
            };
            self.tests.push(test);
        }

        // We don't want to recurse into anything other than mods, since
        // mods or tests inside of functions will break things
        if let ast::ItemKind::Mod(mut module) = item.node {
            let tests = mem::take(&mut self.tests);
            noop_visit_mod(&mut module, self);
            let mut tests = mem::replace(&mut self.tests, tests);

            if !tests.is_empty() {
                let parent = if item.id == ast::DUMMY_NODE_ID {
                    ast::CRATE_NODE_ID
                } else {
                    item.id
                };
                // Create an identifier that will hygienically resolve the test
                // case name, even in another module.
                let expn_id = self.cx.ext_cx.resolver.expansion_for_ast_pass(
                    module.inner,
                    AstPass::TestHarness,
                    &[],
                    Some(parent),
                );
                for test in &mut tests {
                    // See the comment on `mk_main` for why we're using
                    // `apply_mark` directly.
                    test.ident.span = test.ident.span.apply_mark(expn_id, Transparency::Opaque);
                }
                self.cx.test_cases.extend(tests);
            }
            item.node = ast::ItemKind::Mod(module);
        }
        smallvec![P(item)]
    }

    fn visit_mac(&mut self, _mac: &mut ast::Mac) {
        // Do nothing.
    }
}

/// A folder used to remove any entry points (like fn main) because the harness
/// generator will provide its own
struct EntryPointCleaner {
    // Current depth in the ast
    depth: usize,
    def_site: Span,
}

impl MutVisitor for EntryPointCleaner {
    fn flat_map_item(&mut self, i: P<ast::Item>) -> SmallVec<[P<ast::Item>; 1]> {
        self.depth += 1;
        let item = noop_flat_map_item(i, self).expect_one("noop did something");
        self.depth -= 1;

        // Remove any #[main] or #[start] from the AST so it doesn't
        // clash with the one we're going to add, but mark it as
        // #[allow(dead_code)] to avoid printing warnings.
        let item = match entry::entry_point_type(&item, self.depth) {
            EntryPointType::MainNamed |
            EntryPointType::MainAttr |
            EntryPointType::Start =>
                item.map(|ast::Item {id, ident, attrs, node, vis, span, tokens}| {
                    let allow_ident = Ident::new(sym::allow, self.def_site);
                    let dc_nested = attr::mk_nested_word_item(
                        Ident::from_str_and_span("dead_code", self.def_site),
                    );
                    let allow_dead_code_item = attr::mk_list_item(allow_ident, vec![dc_nested]);
                    let allow_dead_code = attr::mk_attr_outer(allow_dead_code_item);

                    ast::Item {
                        id,
                        ident,
                        attrs: attrs.into_iter()
                            .filter(|attr| {
                                !attr.check_name(sym::main) && !attr.check_name(sym::start)
                            })
                            .chain(iter::once(allow_dead_code))
                            .collect(),
                        node,
                        vis,
                        span,
                        tokens,
                    }
                }),
            EntryPointType::None |
            EntryPointType::OtherMain => item,
        };

        smallvec![item]
    }

    fn visit_mac(&mut self, _mac: &mut ast::Mac) {
        // Do nothing.
    }
}

/// Crawl over the crate, inserting test reexports and the test main function
fn generate_test_harness(sess: &ParseSess,
                         resolver: &mut dyn Resolver,
                         reexport_test_harness_main: Option<Symbol>,
                         krate: &mut ast::Crate,
                         features: &Features,
                         test_runner: Option<ast::Path>) {
    let mut econfig = ExpansionConfig::default("test".to_string());
    econfig.features = Some(features);

    let ext_cx = ExtCtxt::new(sess, econfig, resolver);

    let expn_id = ext_cx.resolver.expansion_for_ast_pass(
        DUMMY_SP,
        AstPass::TestHarness,
        &[sym::main, sym::test, sym::rustc_attrs],
        None,
    );
    let def_site = DUMMY_SP.with_def_site_ctxt(expn_id);

    // Remove the entry points
    let mut cleaner = EntryPointCleaner { depth: 0, def_site };
    cleaner.visit_crate(krate);

    let cx = TestCtxt {
        ext_cx,
        def_site,
        test_cases: Vec::new(),
        reexport_test_harness_main,
        test_runner
    };

    TestHarnessGenerator {
        cx,
        tests: Vec::new(),
    }.visit_crate(krate);
}

/// Creates a function item for use as the main function of a test build.
/// This function will call the `test_runner` as specified by the crate attribute
///
/// By default this expands to
///
/// #[main]
/// pub fn main() {
///     extern crate test;
///     test::test_main_static(&[
///         &test_const1,
///         &test_const2,
///         &test_const3,
///     ]);
/// }
///
/// Most of the Ident have the usual def-site hygiene for the AST pass. The
/// exception is the `test_const`s. These have a syntax context that has two
/// opaque marks: one from the expansion of `test` or `test_case`, and one
/// generated  in `TestHarnessGenerator::flat_map_item`. When resolving this
/// identifier after failing to find a matching identifier in the root module
/// we remove the outer mark, and try resolving at its def-site, which will
/// then resolve to `test_const`.
///
/// The expansion here can be controlled by two attributes:
///
/// `reexport_test_harness_main` provides a different name for the `main`
/// function and `test_runner` provides a path that replaces
/// `test::test_main_static`.
fn mk_main(cx: &mut TestCtxt<'_>) -> P<ast::Item> {
    let sp = cx.def_site;
    let ecx = &cx.ext_cx;
    let test_id = Ident::new(sym::test, sp);

    // test::test_main_static(...)
    let mut test_runner = cx.test_runner.clone().unwrap_or(
        ecx.path(sp, vec![test_id, Ident::from_str_and_span("test_main_static", sp)]));

    test_runner.span = sp;

    let test_main_path_expr = ecx.expr_path(test_runner);
    let call_test_main = ecx.expr_call(sp, test_main_path_expr,
                                       vec![mk_tests_slice(cx, sp)]);
    let call_test_main = ecx.stmt_expr(call_test_main);

    // extern crate test
    let test_extern_stmt = ecx.stmt_item(sp, ecx.item(sp,
        test_id,
        vec![],
        ast::ItemKind::ExternCrate(None)
    ));

    // #[main]
    let main_meta = ecx.meta_word(sp, sym::main);
    let main_attr = ecx.attribute(main_meta);

    // pub fn main() { ... }
    let main_ret_ty = ecx.ty(sp, ast::TyKind::Tup(vec![]));

    // If no test runner is provided we need to import the test crate
    let main_body = if cx.test_runner.is_none() {
        ecx.block(sp, vec![test_extern_stmt, call_test_main])
    } else {
        ecx.block(sp, vec![call_test_main])
    };

    let main = ast::ItemKind::Fn(ecx.fn_decl(vec![], ast::FunctionRetTy::Ty(main_ret_ty)),
                           ast::FnHeader::default(),
                           ast::Generics::default(),
                           main_body);

    // Honor the reexport_test_harness_main attribute
    let main_id = match cx.reexport_test_harness_main {
        Some(sym) => Ident::new(sym, sp.with_ctxt(SyntaxContext::root())),
        None => Ident::from_str_and_span("main", sp),
    };

    let main = P(ast::Item {
        ident: main_id,
        attrs: vec![main_attr],
        id: ast::DUMMY_NODE_ID,
        node: main,
        vis: respan(sp, ast::VisibilityKind::Public),
        span: sp,
        tokens: None,
    });

    // Integrate the new item into existing module structures.
    let main = AstFragment::Items(smallvec![main]);
    cx.ext_cx.monotonic_expander().fully_expand_fragment(main).make_items().pop().unwrap()
}

/// Creates a slice containing every test like so:
/// &[&test1, &test2]
fn mk_tests_slice(cx: &TestCtxt<'_>, sp: Span) -> P<ast::Expr> {
    debug!("building test vector from {} tests", cx.test_cases.len());
    let ref ecx = cx.ext_cx;

    ecx.expr_vec_slice(sp,
        cx.test_cases.iter().map(|test| {
            ecx.expr_addr_of(test.span,
                ecx.expr_path(ecx.path(test.span, vec![test.ident])))
        }).collect())
}

fn is_test_case(i: &ast::Item) -> bool {
    attr::contains_name(&i.attrs, sym::rustc_test_marker)
}

fn get_test_runner(sd: &errors::Handler, krate: &ast::Crate) -> Option<ast::Path> {
    let test_attr = attr::find_by_name(&krate.attrs, sym::test_runner)?;
    test_attr.meta_item_list().map(|meta_list| {
        if meta_list.len() != 1 {
            sd.span_fatal(test_attr.span,
                "`#![test_runner(..)]` accepts exactly 1 argument").raise()
        }
        match meta_list[0].meta_item() {
            Some(meta_item) if meta_item.is_word() => meta_item.path.clone(),
            _ => sd.span_fatal(test_attr.span, "`test_runner` argument must be a path").raise()
        }
    })
}
<|endoftext|>


<commit_before>use Result;
use DocId;
use std::io;
use schema::Schema;
use schema::Term;
use core::Segment;
use core::SerializableSegment;
use fastfield::FastFieldsWriter;
use schema::Field;
use schema::FieldValue;
use schema::FieldType;
use indexer::segment_serializer::SegmentSerializer;
use datastruct::stacker::Heap;
use indexer::index_writer::MARGIN_IN_BYTES;
use super::operation::AddOperation;
use postings::MultiFieldPostingsWriter;

/// A `SegmentWriter` is in charge of creating segment index from a
/// documents.
///
/// They creates the postings list in anonymous memory.
/// The segment is layed on disk when the segment gets `finalized`.
pub struct SegmentWriter<'a> {
    heap: &'a Heap,
    max_doc: DocId,
    multifield_postings: MultiFieldPostingsWriter<'a>,
    segment_serializer: SegmentSerializer,
    fast_field_writers: FastFieldsWriter,
    fieldnorms_writer: FastFieldsWriter,
    doc_opstamps: Vec<u64>,
}

fn create_fieldnorms_writer(schema: &Schema) -> FastFieldsWriter {
    let u64_fields: Vec<Field> = schema
        .fields()
        .iter()
        .enumerate()
        .filter(|&(_, field_entry)| field_entry.is_indexed())
        .map(|(field_id, _)| Field(field_id as u32))
        .collect();
    FastFieldsWriter::new(u64_fields)
}

impl<'a> SegmentWriter<'a> {
    /// Creates a new `SegmentWriter`
    ///
    /// The arguments are defined as follows
    ///
    /// - heap: most of the segment writer data (terms, and postings lists recorders)
    /// is stored in a user-defined heap object. This makes it possible for the user to define
    /// the flushing behavior as a buffer limit
    /// - segment: The segment being written
    /// - schema
    pub fn for_segment(heap: &'a Heap,
                       table_bits: usize,
                       mut segment: Segment,
                       schema: &Schema)
                       -> Result<SegmentWriter<'a>> {
        let segment_serializer = SegmentSerializer::for_segment(&mut segment)?;
        let multifield_postings = MultiFieldPostingsWriter::new(schema, table_bits, heap);
        Ok(SegmentWriter {
               heap: heap,
               max_doc: 0,
               multifield_postings: multifield_postings,
               fieldnorms_writer: create_fieldnorms_writer(schema),
               segment_serializer: segment_serializer,
               fast_field_writers: FastFieldsWriter::from_schema(schema),
               doc_opstamps: Vec::with_capacity(1_000),
           })
    }

    /// Lay on disk the current content of the `SegmentWriter`
    ///
    /// Finalize consumes the `SegmentWriter`, so that it cannot
    /// be used afterwards.
    pub fn finalize(self) -> Result<Vec<u64>> {
        write(&self.multifield_postings,
              &self.fast_field_writers,
              &self.fieldnorms_writer,
              self.segment_serializer)?;
        Ok(self.doc_opstamps)
    }

    /// Returns true iff the segment writer's buffer has reached capacity.
    ///
    /// The limit is defined as `the user defined heap size - an arbitrary margin of 10MB`
    /// The `Segment` is `finalize`d when the buffer gets full.
    ///
    /// Because, we cannot cut through a document, the margin is there to ensure that we rarely
    /// exceeds the heap size.
    pub fn is_buffer_full(&self) -> bool {
        self.heap.num_free_bytes() <= MARGIN_IN_BYTES
    }

    /// Return true if the term dictionary hashmap is reaching capacity.
    /// It is one of the condition that triggers a `SegmentWriter` to
    /// be finalized.
    pub(crate) fn is_term_saturated(&self) -> bool {
        self.multifield_postings.is_term_saturated()
    }

    /// Indexes a new document
    ///
    /// As a user, you should rather use `IndexWriter`'s add_document.
    pub fn add_document(&mut self,
                        add_operation: &AddOperation,
                        schema: &Schema)
                        -> io::Result<()> {
        let doc_id = self.max_doc;
        let doc = &add_operation.document;
        self.doc_opstamps.push(add_operation.opstamp);
        for (field, field_values) in doc.get_sorted_field_values() {
            let field_options = schema.get_field_entry(field);
            match *field_options.field_type() {
                FieldType::Str(ref text_options) => {
                    let num_tokens: u32 = if text_options.get_indexing_options().is_tokenized() {
                        self.multifield_postings
                            .index_text(doc_id, field, &field_values)
                    } else {
                        let num_field_values = field_values.len() as u32;
                        for field_value in field_values {
                            let term = Term::from_field_text(field, field_value.value().text());
                            self.multifield_postings.suscribe(doc_id, &term);
                        }
                        num_field_values
                    };
                    self.fieldnorms_writer
                        .get_field_writer(field)
                        .map(|field_norms_writer| field_norms_writer.add_val(num_tokens as u64));
                }
                FieldType::U64(ref int_option) => {
                    if int_option.is_indexed() {
                        for field_value in field_values {
                            let term = Term::from_field_u64(field_value.field(),
                                                            field_value.value().u64_value());
                            self.multifield_postings.suscribe(doc_id, &term);
                        }
                    }
                }
                FieldType::I64(ref int_option) => {
                    if int_option.is_indexed() {
                        for field_value in field_values {
                            let term = Term::from_field_i64(field_value.field(),
                                                            field_value.value().i64_value());
                            self.multifield_postings.suscribe(doc_id, &term);
                        }
                    }
                }
            }
        }
        self.fieldnorms_writer.fill_val_up_to(doc_id);
        self.fast_field_writers.add_document(doc);
        let stored_fieldvalues: Vec<&FieldValue> = doc.field_values()
            .iter()
            .filter(|field_value| schema.get_field_entry(field_value.field()).is_stored())
            .collect();
        let doc_writer = self.segment_serializer.get_store_writer();
        try!(doc_writer.store(&stored_fieldvalues));
        self.max_doc += 1;
        Ok(())
    }

    /// Max doc is
    /// - the number of documents in the segment assuming there is no deletes
    /// - the maximum document id (including deleted documents) + 1
    ///
    /// Currently, **tantivy** does not handle deletes anyway,
    /// so `max_doc == num_docs`
    pub fn max_doc(&self) -> u32 {
        self.max_doc
    }

    /// Number of documents in the index.
    /// Deleted documents are not counted.
    ///
    /// Currently, **tantivy** does not handle deletes anyway,
    /// so `max_doc == num_docs`
    #[allow(dead_code)]
    pub fn num_docs(&self) -> u32 {
        self.max_doc
    }
}

// This method is used as a trick to workaround the borrow checker
fn write(multifield_postings: &MultiFieldPostingsWriter,
         fast_field_writers: &FastFieldsWriter,
         fieldnorms_writer: &FastFieldsWriter,
         mut serializer: SegmentSerializer)
         -> Result<()> {

    try!(multifield_postings.serialize(serializer.get_postings_serializer()));
    try!(fast_field_writers.serialize(serializer.get_fast_field_serializer()));
    try!(fieldnorms_writer.serialize(serializer.get_fieldnorms_serializer()));
    try!(serializer.close());

    Ok(())
}

impl<'a> SerializableSegment for SegmentWriter<'a> {
    fn write(&self, serializer: SegmentSerializer) -> Result<u32> {
        let max_doc = self.max_doc;
        write(&self.multifield_postings,
              &self.fast_field_writers,
              &self.fieldnorms_writer,
              serializer)?;
        Ok(max_doc)
    }
}
<commit_msg>Closes #199. Unindexed fields are indexed as untokenized<commit_after>use Result;
use DocId;
use std::io;
use schema::Schema;
use schema::Term;
use core::Segment;
use core::SerializableSegment;
use fastfield::FastFieldsWriter;
use schema::Field;
use schema::FieldValue;
use schema::FieldType;
use indexer::segment_serializer::SegmentSerializer;
use datastruct::stacker::Heap;
use indexer::index_writer::MARGIN_IN_BYTES;
use super::operation::AddOperation;
use postings::MultiFieldPostingsWriter;

/// A `SegmentWriter` is in charge of creating segment index from a
/// documents.
///
/// They creates the postings list in anonymous memory.
/// The segment is layed on disk when the segment gets `finalized`.
pub struct SegmentWriter<'a> {
    heap: &'a Heap,
    max_doc: DocId,
    multifield_postings: MultiFieldPostingsWriter<'a>,
    segment_serializer: SegmentSerializer,
    fast_field_writers: FastFieldsWriter,
    fieldnorms_writer: FastFieldsWriter,
    doc_opstamps: Vec<u64>,
}

fn create_fieldnorms_writer(schema: &Schema) -> FastFieldsWriter {
    let u64_fields: Vec<Field> = schema
        .fields()
        .iter()
        .enumerate()
        .filter(|&(_, field_entry)| field_entry.is_indexed())
        .map(|(field_id, _)| Field(field_id as u32))
        .collect();
    FastFieldsWriter::new(u64_fields)
}

impl<'a> SegmentWriter<'a> {
    /// Creates a new `SegmentWriter`
    ///
    /// The arguments are defined as follows
    ///
    /// - heap: most of the segment writer data (terms, and postings lists recorders)
    /// is stored in a user-defined heap object. This makes it possible for the user to define
    /// the flushing behavior as a buffer limit
    /// - segment: The segment being written
    /// - schema
    pub fn for_segment(heap: &'a Heap,
                       table_bits: usize,
                       mut segment: Segment,
                       schema: &Schema)
                       -> Result<SegmentWriter<'a>> {
        let segment_serializer = SegmentSerializer::for_segment(&mut segment)?;
        let multifield_postings = MultiFieldPostingsWriter::new(schema, table_bits, heap);
        Ok(SegmentWriter {
               heap: heap,
               max_doc: 0,
               multifield_postings: multifield_postings,
               fieldnorms_writer: create_fieldnorms_writer(schema),
               segment_serializer: segment_serializer,
               fast_field_writers: FastFieldsWriter::from_schema(schema),
               doc_opstamps: Vec::with_capacity(1_000),
           })
    }

    /// Lay on disk the current content of the `SegmentWriter`
    ///
    /// Finalize consumes the `SegmentWriter`, so that it cannot
    /// be used afterwards.
    pub fn finalize(self) -> Result<Vec<u64>> {
        write(&self.multifield_postings,
              &self.fast_field_writers,
              &self.fieldnorms_writer,
              self.segment_serializer)?;
        Ok(self.doc_opstamps)
    }

    /// Returns true iff the segment writer's buffer has reached capacity.
    ///
    /// The limit is defined as `the user defined heap size - an arbitrary margin of 10MB`
    /// The `Segment` is `finalize`d when the buffer gets full.
    ///
    /// Because, we cannot cut through a document, the margin is there to ensure that we rarely
    /// exceeds the heap size.
    pub fn is_buffer_full(&self) -> bool {
        self.heap.num_free_bytes() <= MARGIN_IN_BYTES
    }

    /// Return true if the term dictionary hashmap is reaching capacity.
    /// It is one of the condition that triggers a `SegmentWriter` to
    /// be finalized.
    pub(crate) fn is_term_saturated(&self) -> bool {
        self.multifield_postings.is_term_saturated()
    }

    /// Indexes a new document
    ///
    /// As a user, you should rather use `IndexWriter`'s add_document.
    pub fn add_document(&mut self,
                        add_operation: &AddOperation,
                        schema: &Schema)
                        -> io::Result<()> {
        let doc_id = self.max_doc;
        let doc = &add_operation.document;
        self.doc_opstamps.push(add_operation.opstamp);
        for (field, field_values) in doc.get_sorted_field_values() {
            let field_options = schema.get_field_entry(field);
            if !field_options.is_indexed() {
                continue;
            }
            match *field_options.field_type() {
                FieldType::Str(ref text_options) => {
                    let num_tokens: u32 = if text_options.get_indexing_options().is_tokenized() {
                        self.multifield_postings
                            .index_text(doc_id, field, &field_values)
                    } else {
                        let num_field_values = field_values.len() as u32;
                        for field_value in field_values {
                            let term = Term::from_field_text(field, field_value.value().text());
                            self.multifield_postings.suscribe(doc_id, &term);
                        }
                        num_field_values
                    };
                    self.fieldnorms_writer
                        .get_field_writer(field)
                        .map(|field_norms_writer| field_norms_writer.add_val(num_tokens as u64));
                }
                FieldType::U64(ref int_option) => {
                    if int_option.is_indexed() {
                        for field_value in field_values {
                            let term = Term::from_field_u64(field_value.field(),
                                                            field_value.value().u64_value());
                            self.multifield_postings.suscribe(doc_id, &term);
                        }
                    }
                }
                FieldType::I64(ref int_option) => {
                    if int_option.is_indexed() {
                        for field_value in field_values {
                            let term = Term::from_field_i64(field_value.field(),
                                                            field_value.value().i64_value());
                            self.multifield_postings.suscribe(doc_id, &term);
                        }
                    }
                }
            }
        }
        self.fieldnorms_writer.fill_val_up_to(doc_id);
        self.fast_field_writers.add_document(doc);
        let stored_fieldvalues: Vec<&FieldValue> = doc.field_values()
            .iter()
            .filter(|field_value| schema.get_field_entry(field_value.field()).is_stored())
            .collect();
        let doc_writer = self.segment_serializer.get_store_writer();
        try!(doc_writer.store(&stored_fieldvalues));
        self.max_doc += 1;
        Ok(())
    }

    /// Max doc is
    /// - the number of documents in the segment assuming there is no deletes
    /// - the maximum document id (including deleted documents) + 1
    ///
    /// Currently, **tantivy** does not handle deletes anyway,
    /// so `max_doc == num_docs`
    pub fn max_doc(&self) -> u32 {
        self.max_doc
    }

    /// Number of documents in the index.
    /// Deleted documents are not counted.
    ///
    /// Currently, **tantivy** does not handle deletes anyway,
    /// so `max_doc == num_docs`
    #[allow(dead_code)]
    pub fn num_docs(&self) -> u32 {
        self.max_doc
    }
}

// This method is used as a trick to workaround the borrow checker
fn write(multifield_postings: &MultiFieldPostingsWriter,
         fast_field_writers: &FastFieldsWriter,
         fieldnorms_writer: &FastFieldsWriter,
         mut serializer: SegmentSerializer)
         -> Result<()> {

    try!(multifield_postings.serialize(serializer.get_postings_serializer()));
    try!(fast_field_writers.serialize(serializer.get_fast_field_serializer()));
    try!(fieldnorms_writer.serialize(serializer.get_fieldnorms_serializer()));
    try!(serializer.close());

    Ok(())
}

impl<'a> SerializableSegment for SegmentWriter<'a> {
    fn write(&self, serializer: SegmentSerializer) -> Result<u32> {
        let max_doc = self.max_doc;
        write(&self.multifield_postings,
              &self.fast_field_writers,
              &self.fieldnorms_writer,
              serializer)?;
        Ok(max_doc)
    }
}
<|endoftext|>


<commit_before>// Copyright 2014-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.

use prelude::v1::*;

use alloc::boxed::FnBox;
use cmp;
#[cfg(not(any(target_env = "newlib", target_os = "solaris")))]
use ffi::CStr;
use io;
use libc;
use mem;
use ptr;
use sys::os;
use time::Duration;

use sys_common::thread::*;

pub struct Thread {
    id: libc::pthread_t,
}

// Some platforms may have pthread_t as a pointer in which case we still want
// a thread to be Send/Sync
unsafe impl Send for Thread {}
unsafe impl Sync for Thread {}

impl Thread {
    pub unsafe fn new<'a>(stack: usize, p: Box<FnBox() + 'a>)
                          -> io::Result<Thread> {
        let p = box p;
        let mut native: libc::pthread_t = mem::zeroed();
        let mut attr: libc::pthread_attr_t = mem::zeroed();
        assert_eq!(libc::pthread_attr_init(&mut attr), 0);

        let stack_size = cmp::max(stack, min_stack_size(&attr));
        match libc::pthread_attr_setstacksize(&mut attr,
                                              stack_size as libc::size_t) {
            0 => {}
            n => {
                assert_eq!(n, libc::EINVAL);
                // EINVAL means |stack_size| is either too small or not a
                // multiple of the system page size.  Because it's definitely
                // >= PTHREAD_STACK_MIN, it must be an alignment issue.
                // Round up to the nearest page and try again.
                let page_size = os::page_size();
                let stack_size = (stack_size + page_size - 1) &
                                 (-(page_size as isize - 1) as usize - 1);
                let stack_size = stack_size as libc::size_t;
                assert_eq!(libc::pthread_attr_setstacksize(&mut attr,
                                                           stack_size), 0);
            }
        };

        let ret = libc::pthread_create(&mut native, &attr, thread_start,
                                       &*p as *const _ as *mut _);
        assert_eq!(libc::pthread_attr_destroy(&mut attr), 0);

        return if ret != 0 {
            Err(io::Error::from_raw_os_error(ret))
        } else {
            mem::forget(p); // ownership passed to pthread_create
            Ok(Thread { id: native })
        };

        extern fn thread_start(main: *mut libc::c_void) -> *mut libc::c_void {
            unsafe { start_thread(main); }
            ptr::null_mut()
        }
    }

    pub fn yield_now() {
        let ret = unsafe { libc::sched_yield() };
        debug_assert_eq!(ret, 0);
    }

    #[cfg(any(target_os = "linux",
              target_os = "android",
              target_os = "emscripten"))]
    pub fn set_name(name: &CStr) {
        const PR_SET_NAME: libc::c_int = 15;
        // pthread wrapper only appeared in glibc 2.12, so we use syscall
        // directly.
        unsafe {
            libc::prctl(PR_SET_NAME, name.as_ptr() as libc::c_ulong, 0, 0, 0);
        }
    }

    #[cfg(any(target_os = "freebsd",
              target_os = "dragonfly",
              target_os = "bitrig",
              target_os = "openbsd"))]
    pub fn set_name(name: &CStr) {
        unsafe {
            libc::pthread_set_name_np(libc::pthread_self(), name.as_ptr());
        }
    }

    #[cfg(any(target_os = "macos", target_os = "ios"))]
    pub fn set_name(name: &CStr) {
        unsafe {
            libc::pthread_setname_np(name.as_ptr());
        }
    }

    #[cfg(target_os = "netbsd")]
    pub fn set_name(name: &CStr) {
        use ffi::CString;
        let cname = CString::new(&b"%s"[..]).unwrap();
        unsafe {
            libc::pthread_setname_np(libc::pthread_self(), cname.as_ptr(),
                                     name.as_ptr() as *mut libc::c_void);
        }
    }
    #[cfg(any(target_env = "newlib", target_os = "solaris"))]
    pub fn set_name(_name: &CStr) {
        // Newlib and Illumos has no way to set a thread name.
    }

    pub fn sleep(dur: Duration) {
        let mut ts = libc::timespec {
            tv_sec: dur.as_secs() as libc::time_t,
            tv_nsec: dur.subsec_nanos() as libc::c_long,
        };

        // If we're awoken with a signal then the return value will be -1 and
        // nanosleep will fill in `ts` with the remaining time.
        unsafe {
            while libc::nanosleep(&ts, &mut ts) == -1 {
                assert_eq!(os::errno(), libc::EINTR);
            }
        }
    }

    pub fn join(self) {
        unsafe {
            let ret = libc::pthread_join(self.id, ptr::null_mut());
            mem::forget(self);
            debug_assert_eq!(ret, 0);
        }
    }

    pub fn id(&self) -> libc::pthread_t { self.id }

    pub fn into_id(self) -> libc::pthread_t {
        let id = self.id;
        mem::forget(self);
        id
    }
}

impl Drop for Thread {
    fn drop(&mut self) {
        let ret = unsafe { libc::pthread_detach(self.id) };
        debug_assert_eq!(ret, 0);
    }
}

#[cfg(all(not(all(target_os = "linux", not(target_env = "musl"))),
          not(target_os = "freebsd"),
          not(target_os = "macos"),
          not(target_os = "bitrig"),
          not(all(target_os = "netbsd", not(target_vendor = "rumprun"))),
          not(target_os = "openbsd"),
          not(target_os = "solaris")))]
#[cfg_attr(test, allow(dead_code))]
pub mod guard {
    pub unsafe fn current() -> Option<usize> { None }
    pub unsafe fn init() -> Option<usize> { None }
}

#[cfg(any(all(target_os = "linux", not(target_env = "musl")),
          target_os = "freebsd",
          target_os = "macos",
          target_os = "bitrig",
          all(target_os = "netbsd", not(target_vendor = "rumprun")),
          target_os = "openbsd",
          target_os = "solaris"))]
#[cfg_attr(test, allow(dead_code))]
pub mod guard {
    use prelude::v1::*;

    use libc;
    use libc::mmap;
    use libc::{PROT_NONE, MAP_PRIVATE, MAP_ANON, MAP_FAILED, MAP_FIXED};
    use sys::os;

    #[cfg(any(target_os = "macos",
              target_os = "bitrig",
              target_os = "openbsd",
              target_os = "solaris"))]
    unsafe fn get_stack_start() -> Option<*mut libc::c_void> {
        current().map(|s| s as *mut libc::c_void)
    }

    #[cfg(target_os = "freebsd")]
    unsafe fn get_stack_start() -> Option<*mut libc::c_void> {
        let mut ret = None;
        let mut attr: libc::pthread_attr_t = ::mem::zeroed();
        assert_eq!(libc::pthread_attr_init(&mut attr), 0);
        if libc::pthread_attr_get_np(libc::pthread_self(), &mut attr) == 0 {
            let mut stackaddr = ::ptr::null_mut();
            let mut stacksize = 0;
            assert_eq!(libc::pthread_attr_getstack(&attr, &mut stackaddr,
                                                   &mut stacksize), 0);
            ret = Some(stackaddr);
        }
        assert_eq!(libc::pthread_attr_destroy(&mut attr), 0);
        ret
    }

    #[cfg(any(target_os = "linux", target_os = "android", target_os = "netbsd"))]
    unsafe fn get_stack_start() -> Option<*mut libc::c_void> {
        let mut ret = None;
        let mut attr: libc::pthread_attr_t = ::mem::zeroed();
        assert_eq!(libc::pthread_attr_init(&mut attr), 0);
        if libc::pthread_getattr_np(libc::pthread_self(), &mut attr) == 0 {
            let mut stackaddr = ::ptr::null_mut();
            let mut stacksize = 0;
            assert_eq!(libc::pthread_attr_getstack(&attr, &mut stackaddr,
                                                   &mut stacksize), 0);
            ret = Some(stackaddr);
        }
        assert_eq!(libc::pthread_attr_destroy(&mut attr), 0);
        ret
    }

    pub unsafe fn init() -> Option<usize> {
        let psize = os::page_size();
        let mut stackaddr = match get_stack_start() {
            Some(addr) => addr,
            None => return None,
        };

        // Ensure stackaddr is page aligned! A parent process might
        // have reset RLIMIT_STACK to be non-page aligned. The
        // pthread_attr_getstack() reports the usable stack area
        // stackaddr < stackaddr + stacksize, so if stackaddr is not
        // page-aligned, calculate the fix such that stackaddr <
        // new_page_aligned_stackaddr < stackaddr + stacksize
        let remainder = (stackaddr as usize) % psize;
        if remainder != 0 {
            stackaddr = ((stackaddr as usize) + psize - remainder)
                as *mut libc::c_void;
        }

        // Rellocate the last page of the stack.
        // This ensures SIGBUS will be raised on
        // stack overflow.
        let result = mmap(stackaddr,
                          psize as libc::size_t,
                          PROT_NONE,
                          MAP_PRIVATE | MAP_ANON | MAP_FIXED,
                          -1,
                          0);

        if result != stackaddr || result == MAP_FAILED {
            panic!("failed to allocate a guard page");
        }

        let offset = if cfg!(any(target_os = "linux", target_os = "freebsd")) {
            2
        } else {
            1
        };

        Some(stackaddr as usize + offset * psize)
    }

    #[cfg(target_os = "solaris")]
    pub unsafe fn current() -> Option<usize> {
        let mut current_stack: libc::stack_t = ::mem::zeroed();
        assert_eq!(libc::stack_getbounds(&mut current_stack), 0);
        Some(current_stack.ss_sp as usize)
    }

    #[cfg(target_os = "macos")]
    pub unsafe fn current() -> Option<usize> {
        Some((libc::pthread_get_stackaddr_np(libc::pthread_self()) as libc::size_t -
              libc::pthread_get_stacksize_np(libc::pthread_self())) as usize)
    }

    #[cfg(any(target_os = "openbsd", target_os = "bitrig"))]
    pub unsafe fn current() -> Option<usize> {
        let mut current_stack: libc::stack_t = ::mem::zeroed();
        assert_eq!(libc::pthread_stackseg_np(libc::pthread_self(),
                                             &mut current_stack), 0);

        let extra = if cfg!(target_os = "bitrig") {3} else {1} * os::page_size();
        Some(if libc::pthread_main_np() == 1 {
            // main thread
            current_stack.ss_sp as usize - current_stack.ss_size as usize + extra
        } else {
            // new thread
            current_stack.ss_sp as usize - current_stack.ss_size as usize
        })
    }

    #[cfg(target_os = "freebsd")]
    pub unsafe fn current() -> Option<usize> {
        let mut ret = None;
        let mut attr: libc::pthread_attr_t = ::mem::zeroed();
        assert_eq!(libc::pthread_attr_init(&mut attr), 0);
        if libc::pthread_attr_get_np(libc::pthread_self(), &mut attr) == 0 {
            let mut guardsize = 0;
            assert_eq!(libc::pthread_attr_getguardsize(&attr, &mut guardsize), 0);
            if guardsize == 0 {
                panic!("there is no guard page");
            }
            let mut stackaddr = ::ptr::null_mut();
            let mut size = 0;
            assert_eq!(libc::pthread_attr_getstack(&attr, &mut stackaddr,
                                                   &mut size), 0);
            ret = Some(stackaddr as usize - guardsize as usize);
        }
        assert_eq!(libc::pthread_attr_destroy(&mut attr), 0);
        ret
    }

    #[cfg(any(target_os = "linux", target_os = "android", target_os = "netbsd"))]
    pub unsafe fn current() -> Option<usize> {
        let mut ret = None;
        let mut attr: libc::pthread_attr_t = ::mem::zeroed();
        assert_eq!(libc::pthread_attr_init(&mut attr), 0);
        if libc::pthread_getattr_np(libc::pthread_self(), &mut attr) == 0 {
            let mut guardsize = 0;
            assert_eq!(libc::pthread_attr_getguardsize(&attr, &mut guardsize), 0);
            if guardsize == 0 {
                panic!("there is no guard page");
            }
            let mut stackaddr = ::ptr::null_mut();
            let mut size = 0;
            assert_eq!(libc::pthread_attr_getstack(&attr, &mut stackaddr,
                                                   &mut size), 0);

            ret = if cfg!(target_os = "netbsd") {
                Some(stackaddr as usize)
            } else {
                Some(stackaddr as usize + guardsize as usize)
            };
        }
        assert_eq!(libc::pthread_attr_destroy(&mut attr), 0);
        ret
    }
}

// glibc >= 2.15 has a __pthread_get_minstack() function that returns
// PTHREAD_STACK_MIN plus however many bytes are needed for thread-local
// storage.  We need that information to avoid blowing up when a small stack
// is created in an application with big thread-local storage requirements.
// See #6233 for rationale and details.
#[cfg(target_os = "linux")]
#[allow(deprecated)]
fn min_stack_size(attr: *const libc::pthread_attr_t) -> usize {
    weak!(fn __pthread_get_minstack(*const libc::pthread_attr_t) -> libc::size_t);

    match __pthread_get_minstack.get() {
        None => libc::PTHREAD_STACK_MIN as usize,
        Some(f) => unsafe { f(attr) as usize },
    }
}

// No point in looking up __pthread_get_minstack() on non-glibc
// platforms.
#[cfg(all(not(target_os = "linux"),
          not(target_os = "netbsd")))]
fn min_stack_size(_: *const libc::pthread_attr_t) -> usize {
    libc::PTHREAD_STACK_MIN as usize
}

#[cfg(target_os = "netbsd")]
fn min_stack_size(_: *const libc::pthread_attr_t) -> usize {
    2048 // just a guess
}
<commit_msg>Reduce code duplication in thread.rs<commit_after>// Copyright 2014-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.

use prelude::v1::*;

use alloc::boxed::FnBox;
use cmp;
#[cfg(not(any(target_env = "newlib", target_os = "solaris")))]
use ffi::CStr;
use io;
use libc;
use mem;
use ptr;
use sys::os;
use time::Duration;

use sys_common::thread::*;

pub struct Thread {
    id: libc::pthread_t,
}

// Some platforms may have pthread_t as a pointer in which case we still want
// a thread to be Send/Sync
unsafe impl Send for Thread {}
unsafe impl Sync for Thread {}

impl Thread {
    pub unsafe fn new<'a>(stack: usize, p: Box<FnBox() + 'a>)
                          -> io::Result<Thread> {
        let p = box p;
        let mut native: libc::pthread_t = mem::zeroed();
        let mut attr: libc::pthread_attr_t = mem::zeroed();
        assert_eq!(libc::pthread_attr_init(&mut attr), 0);

        let stack_size = cmp::max(stack, min_stack_size(&attr));
        match libc::pthread_attr_setstacksize(&mut attr,
                                              stack_size as libc::size_t) {
            0 => {}
            n => {
                assert_eq!(n, libc::EINVAL);
                // EINVAL means |stack_size| is either too small or not a
                // multiple of the system page size.  Because it's definitely
                // >= PTHREAD_STACK_MIN, it must be an alignment issue.
                // Round up to the nearest page and try again.
                let page_size = os::page_size();
                let stack_size = (stack_size + page_size - 1) &
                                 (-(page_size as isize - 1) as usize - 1);
                let stack_size = stack_size as libc::size_t;
                assert_eq!(libc::pthread_attr_setstacksize(&mut attr,
                                                           stack_size), 0);
            }
        };

        let ret = libc::pthread_create(&mut native, &attr, thread_start,
                                       &*p as *const _ as *mut _);
        assert_eq!(libc::pthread_attr_destroy(&mut attr), 0);

        return if ret != 0 {
            Err(io::Error::from_raw_os_error(ret))
        } else {
            mem::forget(p); // ownership passed to pthread_create
            Ok(Thread { id: native })
        };

        extern fn thread_start(main: *mut libc::c_void) -> *mut libc::c_void {
            unsafe { start_thread(main); }
            ptr::null_mut()
        }
    }

    pub fn yield_now() {
        let ret = unsafe { libc::sched_yield() };
        debug_assert_eq!(ret, 0);
    }

    #[cfg(any(target_os = "linux",
              target_os = "android",
              target_os = "emscripten"))]
    pub fn set_name(name: &CStr) {
        const PR_SET_NAME: libc::c_int = 15;
        // pthread wrapper only appeared in glibc 2.12, so we use syscall
        // directly.
        unsafe {
            libc::prctl(PR_SET_NAME, name.as_ptr() as libc::c_ulong, 0, 0, 0);
        }
    }

    #[cfg(any(target_os = "freebsd",
              target_os = "dragonfly",
              target_os = "bitrig",
              target_os = "openbsd"))]
    pub fn set_name(name: &CStr) {
        unsafe {
            libc::pthread_set_name_np(libc::pthread_self(), name.as_ptr());
        }
    }

    #[cfg(any(target_os = "macos", target_os = "ios"))]
    pub fn set_name(name: &CStr) {
        unsafe {
            libc::pthread_setname_np(name.as_ptr());
        }
    }

    #[cfg(target_os = "netbsd")]
    pub fn set_name(name: &CStr) {
        use ffi::CString;
        let cname = CString::new(&b"%s"[..]).unwrap();
        unsafe {
            libc::pthread_setname_np(libc::pthread_self(), cname.as_ptr(),
                                     name.as_ptr() as *mut libc::c_void);
        }
    }
    #[cfg(any(target_env = "newlib", target_os = "solaris"))]
    pub fn set_name(_name: &CStr) {
        // Newlib and Illumos has no way to set a thread name.
    }

    pub fn sleep(dur: Duration) {
        let mut ts = libc::timespec {
            tv_sec: dur.as_secs() as libc::time_t,
            tv_nsec: dur.subsec_nanos() as libc::c_long,
        };

        // If we're awoken with a signal then the return value will be -1 and
        // nanosleep will fill in `ts` with the remaining time.
        unsafe {
            while libc::nanosleep(&ts, &mut ts) == -1 {
                assert_eq!(os::errno(), libc::EINTR);
            }
        }
    }

    pub fn join(self) {
        unsafe {
            let ret = libc::pthread_join(self.id, ptr::null_mut());
            mem::forget(self);
            debug_assert_eq!(ret, 0);
        }
    }

    pub fn id(&self) -> libc::pthread_t { self.id }

    pub fn into_id(self) -> libc::pthread_t {
        let id = self.id;
        mem::forget(self);
        id
    }
}

impl Drop for Thread {
    fn drop(&mut self) {
        let ret = unsafe { libc::pthread_detach(self.id) };
        debug_assert_eq!(ret, 0);
    }
}

#[cfg(all(not(all(target_os = "linux", not(target_env = "musl"))),
          not(target_os = "freebsd"),
          not(target_os = "macos"),
          not(target_os = "bitrig"),
          not(all(target_os = "netbsd", not(target_vendor = "rumprun"))),
          not(target_os = "openbsd"),
          not(target_os = "solaris")))]
#[cfg_attr(test, allow(dead_code))]
pub mod guard {
    pub unsafe fn current() -> Option<usize> { None }
    pub unsafe fn init() -> Option<usize> { None }
}

#[cfg(any(all(target_os = "linux", not(target_env = "musl")),
          target_os = "freebsd",
          target_os = "macos",
          target_os = "bitrig",
          all(target_os = "netbsd", not(target_vendor = "rumprun")),
          target_os = "openbsd",
          target_os = "solaris"))]
#[cfg_attr(test, allow(dead_code))]
pub mod guard {
    use prelude::v1::*;

    use libc;
    use libc::mmap;
    use libc::{PROT_NONE, MAP_PRIVATE, MAP_ANON, MAP_FAILED, MAP_FIXED};
    use sys::os;

    #[cfg(any(target_os = "macos",
              target_os = "bitrig",
              target_os = "openbsd",
              target_os = "solaris"))]
    unsafe fn get_stack_start() -> Option<*mut libc::c_void> {
        current().map(|s| s as *mut libc::c_void)
    }

    #[cfg(any(target_os = "android", target_os = "freebsd",
              target_os = "linux", target_os = "netbsd"))]
    unsafe fn get_stack_start() -> Option<*mut libc::c_void> {
        let mut ret = None;
        let mut attr: libc::pthread_attr_t = ::mem::zeroed();
        assert_eq!(libc::pthread_attr_init(&mut attr), 0);
        #[cfg(target_os = "freebsd")]
            let e = libc::pthread_attr_get_np(libc::pthread_self(), &mut attr);
        #[cfg(not(target_os = "freebsd"))]
            let e = libc::pthread_getattr_np(libc::pthread_self(), &mut attr);
        if e == 0 {
            let mut stackaddr = ::ptr::null_mut();
            let mut stacksize = 0;
            assert_eq!(libc::pthread_attr_getstack(&attr, &mut stackaddr,
                                                   &mut stacksize), 0);
            ret = Some(stackaddr);
        }
        assert_eq!(libc::pthread_attr_destroy(&mut attr), 0);
        ret
    }

    pub unsafe fn init() -> Option<usize> {
        let psize = os::page_size();
        let mut stackaddr = match get_stack_start() {
            Some(addr) => addr,
            None => return None,
        };

        // Ensure stackaddr is page aligned! A parent process might
        // have reset RLIMIT_STACK to be non-page aligned. The
        // pthread_attr_getstack() reports the usable stack area
        // stackaddr < stackaddr + stacksize, so if stackaddr is not
        // page-aligned, calculate the fix such that stackaddr <
        // new_page_aligned_stackaddr < stackaddr + stacksize
        let remainder = (stackaddr as usize) % psize;
        if remainder != 0 {
            stackaddr = ((stackaddr as usize) + psize - remainder)
                as *mut libc::c_void;
        }

        // Rellocate the last page of the stack.
        // This ensures SIGBUS will be raised on
        // stack overflow.
        let result = mmap(stackaddr,
                          psize as libc::size_t,
                          PROT_NONE,
                          MAP_PRIVATE | MAP_ANON | MAP_FIXED,
                          -1,
                          0);

        if result != stackaddr || result == MAP_FAILED {
            panic!("failed to allocate a guard page");
        }

        let offset = if cfg!(any(target_os = "linux", target_os = "freebsd")) {
            2
        } else {
            1
        };

        Some(stackaddr as usize + offset * psize)
    }

    #[cfg(target_os = "solaris")]
    pub unsafe fn current() -> Option<usize> {
        let mut current_stack: libc::stack_t = ::mem::zeroed();
        assert_eq!(libc::stack_getbounds(&mut current_stack), 0);
        Some(current_stack.ss_sp as usize)
    }

    #[cfg(target_os = "macos")]
    pub unsafe fn current() -> Option<usize> {
        Some((libc::pthread_get_stackaddr_np(libc::pthread_self()) as libc::size_t -
              libc::pthread_get_stacksize_np(libc::pthread_self())) as usize)
    }

    #[cfg(any(target_os = "openbsd", target_os = "bitrig"))]
    pub unsafe fn current() -> Option<usize> {
        let mut current_stack: libc::stack_t = ::mem::zeroed();
        assert_eq!(libc::pthread_stackseg_np(libc::pthread_self(),
                                             &mut current_stack), 0);

        let extra = if cfg!(target_os = "bitrig") {3} else {1} * os::page_size();
        Some(if libc::pthread_main_np() == 1 {
            // main thread
            current_stack.ss_sp as usize - current_stack.ss_size as usize + extra
        } else {
            // new thread
            current_stack.ss_sp as usize - current_stack.ss_size as usize
        })
    }

    #[cfg(any(target_os = "android", target_os = "freebsd",
              target_os = "linux", target_os = "netbsd"))]
    pub unsafe fn current() -> Option<usize> {
        let mut ret = None;
        let mut attr: libc::pthread_attr_t = ::mem::zeroed();
        assert_eq!(libc::pthread_attr_init(&mut attr), 0);
        #[cfg(target_os = "freebsd")]
            let e = libc::pthread_attr_get_np(libc::pthread_self(), &mut attr);
        #[cfg(not(target_os = "freebsd"))]
            let e = libc::pthread_getattr_np(libc::pthread_self(), &mut attr);
        if e == 0 {
        //if libc::pthread_getattr_np(libc::pthread_self(), &mut attr) == 0 {
            let mut guardsize = 0;
            assert_eq!(libc::pthread_attr_getguardsize(&attr, &mut guardsize), 0);
            if guardsize == 0 {
                panic!("there is no guard page");
            }
            let mut stackaddr = ::ptr::null_mut();
            let mut size = 0;
            assert_eq!(libc::pthread_attr_getstack(&attr, &mut stackaddr,
                                                   &mut size), 0);

            ret = if cfg!(target_os = "freebsd") {
                Some(stackaddr as usize - guardsize as usize)
            } else if cfg!(target_os = "netbsd") {
                Some(stackaddr as usize)
            } else {
                Some(stackaddr as usize + guardsize as usize)
            };
        }
        assert_eq!(libc::pthread_attr_destroy(&mut attr), 0);
        ret
    }
}

// glibc >= 2.15 has a __pthread_get_minstack() function that returns
// PTHREAD_STACK_MIN plus however many bytes are needed for thread-local
// storage.  We need that information to avoid blowing up when a small stack
// is created in an application with big thread-local storage requirements.
// See #6233 for rationale and details.
#[cfg(target_os = "linux")]
#[allow(deprecated)]
fn min_stack_size(attr: *const libc::pthread_attr_t) -> usize {
    weak!(fn __pthread_get_minstack(*const libc::pthread_attr_t) -> libc::size_t);

    match __pthread_get_minstack.get() {
        None => libc::PTHREAD_STACK_MIN as usize,
        Some(f) => unsafe { f(attr) as usize },
    }
}

// No point in looking up __pthread_get_minstack() on non-glibc
// platforms.
#[cfg(all(not(target_os = "linux"),
          not(target_os = "netbsd")))]
fn min_stack_size(_: *const libc::pthread_attr_t) -> usize {
    libc::PTHREAD_STACK_MIN as usize
}

#[cfg(target_os = "netbsd")]
fn min_stack_size(_: *const libc::pthread_attr_t) -> usize {
    2048 // just a guess
}
<|endoftext|>