<commit_before>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Support for the 8259 Programmable Interrupt Controller.
//!
//! The 8259 PIC controls all of our IO IRQs. The PIC recieves IRQ requests
//! and feeds them to the CPU in order. An x86 system typically has two PICs,
//! PIC1 and PIC2, each of which provides 8 IRQs. Two PICs provide us with 15
//! unique IRQs, because one interrupt on the leader PIC is linked to interrupts
//! on the follower PIC.
//!
//! Mistakes were made, and Intel boneheadedly decided that it was a good idea
//! to map the vector offset of PIC1 to 8, so that it maps to interrupts
//! 8 ... 15 in the IDT. This conflicts with some of the interrupt numbers
//! used by CPU exceptions. Therefore, we must remap the PIC vectors so that
//! PIC1 starts at 32 and PIC2 at 40.

use super::super::Port;
use spin::Mutex;
use core::mem::transmute;

/// Starting offset for PIC1
const OFFSET: u8 = 0x20;

/// Commands to send to the PIC
#[repr(u8)]
#[derive(Eq, PartialEq, Ord, PartialOrd, Copy, Clone)]
enum Command { Mode8086 = 0x01
             , Init     = 0x11
             , EndIRQ   = 0x20
             , ReadIRR  = 0x0a
             , ReadISR  = 0x0b
             }

/// List of IRQs on the x86.
///
/// See here for more info:
/// https://en.wikibooks.org/wiki/X86_Assembly/Programmable_Interrupt_Controller
#[repr(u8)]
#[derive(Eq, PartialEq, Ord, PartialOrd, Copy, Clone)]
pub enum IRQ { /// System timer IRQ
               Timer        = 0 + OFFSET
             , /// PS/2 keyboard controller
               PS2Keyboard  = 1 + OFFSET
             , /// PIC2 cascade IRQ
               Cascade      = 2 + OFFSET
             , /// COM2 serial port
               COM2         = 3 + OFFSET
             , /// COM1 serial port
               COM1         = 4 + OFFSET
             , /// Line printer 2
               LPT2         = 5 + OFFSET
             , /// Floppy disc controller
               Floppy       = 6 + OFFSET
             , /// Line printer 1
               LPT1         = 7 + OFFSET
             , /// CMOS clock
               RTCTimer     = 8 + OFFSET
             , /// PS/2 mouse controller
               PS2Mouse     = 12 + OFFSET
             , /// Floating-point Coprocessor
               FPU          = 13 + OFFSET
             , /// ATA channel 1
               PrimaryATA   = 14 + OFFSET
             , /// ATA channel 2
               SecondaryATA = 15 + OFFSET
             }

struct PIC {
    /// The base offset to which interrupts on this PIC are mapped
    offset: u8
  , /// The port on the CPU that sends commands to this PIC.
    command_port: Port
  , /// The port that sends and recieves data from the PIC
    data_port: Port
}

impl PIC {

    const fn leader() -> PIC {
        unsafe {
            PIC { offset: OFFSET
                , command_port: Port::new(0x20)
                , data_port: Port::new(0x21)
                }
        }
    }

    const fn follower() -> PIC {
        unsafe {
            PIC { offset: OFFSET + 8
                , command_port: Port::new(0xA0)
                , data_port: Port::new(0xA1)
                }
        }
    }

    #[inline]
    fn is_leader(&self) -> bool {
        self.offset == OFFSET
    }

    #[inline]
    fn send_command(&self, command: Command) {
        unsafe {
            self.command_port
                .out8(command as u8)
        }
    }

    #[inline]
    fn send_data(&self, data: u8) {
        unsafe {
            self.data_port
                .out8(data)
        }
    }

    #[inline]
    fn initialize(&self) {
        self.send_command(Command::Init)
    }

    #[inline]
    fn read_isr(&self) -> u8 {
        self.send_command(Command::ReadISR);
        unsafe { self.data_port.in8() }
    }

    #[inline]
    fn read_irr(&self) -> u8 {
        self.send_command(Command::ReadIRR);
        unsafe { self.data_port.in8() }
    }

}

trait IRQHandler {
    /// Returns whether or not this handler handles the given IRQ
    fn handles(&self, irq: IRQ) -> bool;
    /// End an interrupt request
    fn end_interrupt(&self, irq: IRQ);
}

impl IRQHandler for PIC {

    fn handles(&self, irq: IRQ) -> bool {
        self.offset <= (irq as u8) && (irq as u8) < self.offset + 8
    }

    #[allow(unused_must_use)]
    fn end_interrupt(&self, _: IRQ) {
        self.send_command(Command::EndIRQ)
    }
}

/// A pair of PICs in cascade mode.
///
/// This is the standard configuration on all modern x86 systems.
struct BothPICs (PIC, PIC);

impl BothPICs {
    const fn new() -> Self {
        BothPICs (PIC::leader(), PIC::follower())
    }

    /// Initialize the system's PICs.
    pub unsafe fn initialize(&mut self) {
        let wait_port = Port::new(0x80);
        let wait = || { wait_port.out8(0); };
        // helper macro to avoid writing repetitive code
        macro_rules! send {
            (pic0 => $data:expr) => {
                self.0.send_data($data)
                wait();
            };
            (pic1 => $data:expr) => {
                self.1.send_data($data)
                wait();
            };
        }

        // Read the default interrupt masks from PIC1 and PIC2
        let (saved_mask1, saved_mask2)
            = (self.0.data_port.in8(), self.1.data_port.in8());

        // Send both PICs the 'initialize' command.
        self.0.initialize(); wait();
        self.1.initialize(); wait();

        // Each PIC then expects us to send it the following:
        // 1. the PIC's new vector offset
        send!(pic0 => self.0.offset);
        send!(pic1 => self.1.offset);
        // 2. number to configure PIC cascading
        send!(pic0 => 0x04);
        send!(pic1 => 0x02);
        // 3. command for 8086 mode
        send!(pic0 => Command::Mode8086 as u8);
        send!(pic1 => Command::Mode8086 as u8);
        // 4. finally, the mask we saved earlier
        send!(pic0 => saved_mask1);
        send!(pic1 => saved_mask2);
    }
}

impl IRQHandler for BothPICs {

    fn handles(&self, irq: IRQ) -> bool {
        self.0.handles(irq) ||
        self.1.handles(irq)
    }

    fn end_interrupt(&self, irq: IRQ) {
        if self.1.handles(irq) {
            self.1.end_interrupt(irq);
        }
        self.0.end_interrupt(irq);
    }

}

/// Global PIC instance and mutex
static PICS: Mutex<BothPICs>
    = Mutex::new(BothPICs::new());

/// Initialize the system's Programmable Interrupt Controller
pub unsafe fn initialize() {
    PICS.lock()
        .initialize();
    println!("{:<38}{:>40}", " . . Initializing PICs", "[ OKAY ]")

}

/// If an interrupt is being handled by the PICs, end that interrupt.
///
/// This is called by the interrupt handler at the end of all interrupts.
/// If the interrupt is not a PIC interrupt, it silently does nothing.
pub unsafe fn end_pic_interrupt(interrupt_id: u8) {
    let pics = PICS.lock();
    let irq: IRQ = transmute(interrupt_id);

    if pics.handles(irq) {
        pics.end_interrupt(irq)
    }
}
<commit_msg>[elf] Remove some elided explicit lifetimes<commit_after>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Support for the 8259 Programmable Interrupt Controller.
//!
//! The 8259 PIC controls all of our IO IRQs. The PIC recieves IRQ requests
//! and feeds them to the CPU in order. An x86 system typically has two PICs,
//! PIC1 and PIC2, each of which provides 8 IRQs. Two PICs provide us with 15
//! unique IRQs, because one interrupt on the leader PIC is linked to interrupts
//! on the follower PIC.
//!
//! Mistakes were made, and Intel boneheadedly decided that it was a good idea
//! to map the vector offset of PIC1 to 8, so that it maps to interrupts
//! 8 ... 15 in the IDT. This conflicts with some of the interrupt numbers
//! used by CPU exceptions. Therefore, we must remap the PIC vectors so that
//! PIC1 starts at 32 and PIC2 at 40.

use super::super::Port;
use spin::Mutex;
use core::mem::transmute;

/// Starting offset for PIC1
const OFFSET: u8 = 0x20;

/// Commands to send to the PIC
#[repr(u8)]
#[derive(Eq, PartialEq, Ord, PartialOrd, Copy, Clone)]
enum Command { Mode8086 = 0x01
             , Init     = 0x11
             , EndIRQ   = 0x20
             , ReadIRR  = 0x0a
             , ReadISR  = 0x0b
             }

/// List of IRQs on the x86.
///
/// See here for more info:
/// https://en.wikibooks.org/wiki/X86_Assembly/Programmable_Interrupt_Controller
#[repr(u8)]
#[derive(Eq, PartialEq, Ord, PartialOrd, Copy, Clone)]
pub enum IRQ { /// System timer IRQ
               Timer        = OFFSET
             , /// PS/2 keyboard controller
               PS2Keyboard  = 1 + OFFSET
             , /// PIC2 cascade IRQ
               Cascade      = 2 + OFFSET
             , /// COM2 serial port
               COM2         = 3 + OFFSET
             , /// COM1 serial port
               COM1         = 4 + OFFSET
             , /// Line printer 2
               LPT2         = 5 + OFFSET
             , /// Floppy disc controller
               Floppy       = 6 + OFFSET
             , /// Line printer 1
               LPT1         = 7 + OFFSET
             , /// CMOS clock
               RTCTimer     = 8 + OFFSET
             , /// PS/2 mouse controller
               PS2Mouse     = 12 + OFFSET
             , /// Floating-point Coprocessor
               FPU          = 13 + OFFSET
             , /// ATA channel 1
               PrimaryATA   = 14 + OFFSET
             , /// ATA channel 2
               SecondaryATA = 15 + OFFSET
             }

struct PIC {
    /// The base offset to which interrupts on this PIC are mapped
    offset: u8
  , /// The port on the CPU that sends commands to this PIC.
    command_port: Port
  , /// The port that sends and recieves data from the PIC
    data_port: Port
}

impl PIC {

    pub const fn leader() -> PIC {
        unsafe {
            PIC { offset: OFFSET
                , command_port: Port::new(0x20)
                , data_port: Port::new(0x21)
                }
        }
    }

    pub const fn follower() -> PIC {
        unsafe {
            PIC { offset: OFFSET + 8
                , command_port: Port::new(0xA0)
                , data_port: Port::new(0xA1)
                }
        }
    }

    #[inline]
    pub fn is_leader(&self) -> bool {
        self.offset == OFFSET
    }

    #[inline]
    pub fn send_command(&self, command: Command) {
        unsafe {
            self.command_port
                .out8(command as u8)
        }
    }

    #[inline]
    pub fn send_data(&self, data: u8) {
        unsafe {
            self.data_port
                .out8(data)
        }
    }

    #[inline]
    pub fn initialize(&self) {
        self.send_command(Command::Init)
    }

    #[inline]
    pub fn read_isr(&self) -> u8 {
        self.send_command(Command::ReadISR);
        unsafe { self.data_port.in8() }
    }

    #[inline]
    fn read_irr(&self) -> u8 {
        self.send_command(Command::ReadIRR);
        unsafe { self.data_port.in8() }
    }

}

trait IRQHandler {
    /// Returns whether or not this handler handles the given IRQ
    fn handles(&self, irq: IRQ) -> bool;
    /// End an interrupt request
    fn end_interrupt(&self, irq: IRQ);
}

impl IRQHandler for PIC {

    fn handles(&self, irq: IRQ) -> bool {
        self.offset <= (irq as u8) && (irq as u8) < self.offset + 8
    }

    #[allow(unused_must_use)]
    fn end_interrupt(&self, _: IRQ) {
        self.send_command(Command::EndIRQ)
    }
}

/// A pair of PICs in cascade mode.
///
/// This is the standard configuration on all modern x86 systems.
struct BothPICs (PIC, PIC);

impl BothPICs {
    const fn new() -> Self {
        BothPICs (PIC::leader(), PIC::follower())
    }

    /// Initialize the system's PICs.
    pub unsafe fn initialize(&mut self) {
        let wait_port = Port::new(0x80);
        let wait = || { wait_port.out8(0); };
        // helper macro to avoid writing repetitive code
        macro_rules! send {
            (pic0 => $data:expr) => {
                self.0.send_data($data)
                wait();
            };
            (pic1 => $data:expr) => {
                self.1.send_data($data)
                wait();
            };
        }

        // Read the default interrupt masks from PIC1 and PIC2
        let (saved_mask1, saved_mask2)
            = (self.0.data_port.in8(), self.1.data_port.in8());

        // Send both PICs the 'initialize' command.
        self.0.initialize(); wait();
        self.1.initialize(); wait();

        // Each PIC then expects us to send it the following:
        // 1. the PIC's new vector offset
        send!(pic0 => self.0.offset);
        send!(pic1 => self.1.offset);
        // 2. number to configure PIC cascading
        send!(pic0 => 0x04);
        send!(pic1 => 0x02);
        // 3. command for 8086 mode
        send!(pic0 => Command::Mode8086 as u8);
        send!(pic1 => Command::Mode8086 as u8);
        // 4. finally, the mask we saved earlier
        send!(pic0 => saved_mask1);
        send!(pic1 => saved_mask2);
    }
}

impl IRQHandler for BothPICs {

    fn handles(&self, irq: IRQ) -> bool {
        self.0.handles(irq) ||
        self.1.handles(irq)
    }

    fn end_interrupt(&self, irq: IRQ) {
        if self.1.handles(irq) {
            self.1.end_interrupt(irq);
        }
        self.0.end_interrupt(irq);
    }

}

/// Global PIC instance and mutex
static PICS: Mutex<BothPICs>
    = Mutex::new(BothPICs::new());

/// Initialize the system's Programmable Interrupt Controller
pub unsafe fn initialize() {
    PICS.lock()
        .initialize();
    println!("{:<38}{:>40}", " . . Initializing PICs", "[ OKAY ]")

}

/// If an interrupt is being handled by the PICs, end that interrupt.
///
/// This is called by the interrupt handler at the end of all interrupts.
/// If the interrupt is not a PIC interrupt, it silently does nothing.
pub unsafe fn end_pic_interrupt(interrupt_id: u8) {
    let pics = PICS.lock();
    let irq: IRQ = transmute(interrupt_id);

    if pics.handles(irq) {
        pics.end_interrupt(irq)
    }
}
<|endoftext|>


<commit_before>use super::super::backend::Backend;
use super::super::buckets::Buckets;
use std::net::{Ipv4Addr, SocketAddrV4, TcpStream};
use std::str::FromStr;
use std::fmt::Write;
use std::io::Write as IoWrite;
use time;

#[derive(Debug)]
pub struct Graphite {
    addr: SocketAddrV4,
    last_flush_time: u64,
    last_flush_length: u64
}

impl Graphite {
    /// Create a Graphite formatter
    ///
    /// # Examples
    ///
    /// ```
    /// let graph = Graphite::new(host, port);
    /// ```
    pub fn new(host: &str, port: u16) -> Graphite {
        let ip = Ipv4Addr::from_str(&host).unwrap();
        let addr = SocketAddrV4::new(ip, port);
        Graphite {
            addr: addr,
            last_flush_time: 0,
            last_flush_length: 0,
        }
    }

    /// Convert the buckets into a String that
    /// can be sent to graphite's newline API
    pub fn format_stats(&self, buckets: &Buckets) -> String {
        let start = time::get_time().sec;
        let mut stats = String::new();

        write!(
            stats,
            "{} {} {}\n",
            "statsd.bad_messages",
            buckets.bad_messages(),
            start).unwrap();
        write!(
            stats,
            "{} {} {}\n",
            "statsd.total_messages",
            buckets.total_messages(),
            start).unwrap();

        for (key, value) in buckets.counters().iter() {
            write!(
                stats,
                "{} {} {} \n",
                key,
                value,
                start).unwrap();
        }

        for (key, value) in buckets.gauges().iter() {
            write!(
                stats,
                "{} {} {} \n",
                key,
                value,
                start).unwrap();
        }

        // The raw timer data is not sent to graphite.
        for (key, value) in buckets.timer_data().iter() {
            write!(
                stats,
                "{} {} {} \n",
                key,
                value,
                start).unwrap();
        }
        stats
    }
}

impl Backend for Graphite {
    fn flush_buckets(&mut self, buckets: &Buckets) {
        let stats = self.format_stats(&buckets);

        let mut stream = TcpStream::connect(self.addr).unwrap();
        let _ = stream.write(stats.as_bytes());
    }
}

#[cfg(test)]
mod test {
    use super::super::super::metric::{Metric, MetricKind};
    use super::super::super::buckets::Buckets;
    use super::super::super::metric_processor::process;
    use super::*;

    fn make_buckets() -> Buckets {
        let mut buckets = Buckets::new();
        let m1 = Metric::new("test.counter", 1.0, MetricKind::Counter(1.0));
        let m2 = Metric::new("test.gauge", 3.211, MetricKind::Gauge);

        let m3 = Metric::new("test.timer", 12.101, MetricKind::Timer);
        let m4 = Metric::new("test.timer", 1.101, MetricKind::Timer);
        let m5 = Metric::new("test.timer", 3.101, MetricKind::Timer);
        buckets.add(&m1);
        buckets.add(&m2);
        buckets.add(&m3);
        buckets.add(&m4);
        buckets.add(&m5);
        buckets
    }

    #[test]
    fn test_format_buckets_no_timers() {
        let buckets = make_buckets();
        let graphite = Graphite::new("127.0.0.1", 2003);
        let result = graphite.format_stats(&buckets);
        let mut lines: Vec<&str> = result.lines().collect();

        assert_eq!(4, lines.len());
        assert!(lines[0].contains("statsd.bad_messages 0"));
        assert!(lines[1].contains("statsd.total_messages 5"));
        assert!(lines[2].contains("test.counter 1"));
        assert!(lines[3].contains("test.gauge 3.211"));
    }
}
<commit_msg>Add tests for computed timer metrics & graphite.<commit_after>use super::super::backend::Backend;
use super::super::buckets::Buckets;
use std::net::{Ipv4Addr, SocketAddrV4, TcpStream};
use std::str::FromStr;
use std::fmt::Write;
use std::io::Write as IoWrite;
use time;

#[derive(Debug)]
pub struct Graphite {
    addr: SocketAddrV4,
    last_flush_time: u64,
    last_flush_length: u64
}

impl Graphite {
    /// Create a Graphite formatter
    ///
    /// # Examples
    ///
    /// ```
    /// let graph = Graphite::new(host, port);
    /// ```
    pub fn new(host: &str, port: u16) -> Graphite {
        let ip = Ipv4Addr::from_str(&host).unwrap();
        let addr = SocketAddrV4::new(ip, port);
        Graphite {
            addr: addr,
            last_flush_time: 0,
            last_flush_length: 0,
        }
    }

    /// Convert the buckets into a String that
    /// can be sent to graphite's newline API
    pub fn format_stats(&self, buckets: &Buckets) -> String {
        let start = time::get_time().sec;
        let mut stats = String::new();

        write!(
            stats,
            "{} {} {}\n",
            "statsd.bad_messages",
            buckets.bad_messages(),
            start).unwrap();
        write!(
            stats,
            "{} {} {}\n",
            "statsd.total_messages",
            buckets.total_messages(),
            start).unwrap();

        for (key, value) in buckets.counters().iter() {
            write!(
                stats,
                "{} {} {} \n",
                key,
                value,
                start).unwrap();
        }

        for (key, value) in buckets.gauges().iter() {
            write!(
                stats,
                "{} {} {} \n",
                key,
                value,
                start).unwrap();
        }

        // The raw timer data is not sent to graphite.
        for (key, value) in buckets.timer_data().iter() {
            write!(
                stats,
                "{} {} {} \n",
                key,
                value,
                start).unwrap();
        }
        stats
    }
}

impl Backend for Graphite {
    fn flush_buckets(&mut self, buckets: &Buckets) {
        let stats = self.format_stats(&buckets);

        let mut stream = TcpStream::connect(self.addr).unwrap();
        let _ = stream.write(stats.as_bytes());
    }
}

#[cfg(test)]
mod test {
    use super::super::super::metric::{Metric, MetricKind};
    use super::super::super::buckets::Buckets;
    use super::super::super::metric_processor::process;
    use super::*;

    fn make_buckets() -> Buckets {
        let mut buckets = Buckets::new();
        let m1 = Metric::new("test.counter", 1.0, MetricKind::Counter(1.0));
        let m2 = Metric::new("test.gauge", 3.211, MetricKind::Gauge);

        let m3 = Metric::new("test.timer", 12.101, MetricKind::Timer);
        let m4 = Metric::new("test.timer", 1.101, MetricKind::Timer);
        let m5 = Metric::new("test.timer", 3.101, MetricKind::Timer);
        buckets.add(&m1);
        buckets.add(&m2);
        buckets.add(&m3);
        buckets.add(&m4);
        buckets.add(&m5);
        buckets
    }

    #[test]
    fn test_format_buckets_no_timers() {
        let buckets = make_buckets();
        let graphite = Graphite::new("127.0.0.1", 2003);
        let result = graphite.format_stats(&buckets);
        let mut lines: Vec<&str> = result.lines().collect();

        assert_eq!(4, lines.len());
        assert!(lines[0].contains("statsd.bad_messages 0"));
        assert!(lines[1].contains("statsd.total_messages 5"));
        assert!(lines[2].contains("test.counter 1"));
        assert!(lines[3].contains("test.gauge 3.211"));
    }

    #[test]
    fn test_format_buckets_timers() {
        let mut buckets = make_buckets();
        process(&mut buckets);

        let graphite = Graphite::new("127.0.0.1", 2003);
        let result = graphite.format_stats(&buckets);
        let mut lines: Vec<&str> = result.lines().collect();

        assert_eq!(12, lines.len());

        assert!(result.contains("test.timer.max 12.101"));
        assert!(result.contains("test.timer.min 1.101"));
        assert!(result.contains("test.timer.count 3"));
    }
}
<|endoftext|>


<commit_before>#![feature(integer_atomics)]

extern crate abra;
extern crate rocksdb;
extern crate rustc_serialize;
#[macro_use]
extern crate maplit;
extern crate byteorder;

use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicU32, Ordering};
use std::collections::BTreeMap;

use rocksdb::{DB, Writable, Options, MergeOperands};
use abra::{Term, Document};
use abra::schema::{Schema, FieldType, FieldRef, AddFieldError};
use rustc_serialize::{json, Encodable};
use byteorder::{ByteOrder, BigEndian};

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
struct TermRef(u32);

impl TermRef {
    pub fn ord(&self) -> u32 {
        self.0
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
struct DocRef(u32, u16);

impl DocRef {
    pub fn chunk(&self) -> u32 {
        self.0
    }

    pub fn ord(&self) -> u16 {
        self.1
    }
}

fn merge_keys(key: &[u8], existing_val: Option<&[u8]>, operands: &mut MergeOperands) -> Vec<u8> {
    match key[0] {
        b'd' => {
            // Directory (sequence of two byte document ids)

            // Allocate vec for new Value
            let new_size = match existing_val {
                Some(existing_val) => existing_val.len(),
                None => 0,
            } + operands.size_hint().0 * 2;

            let mut new_val = Vec::with_capacity(new_size);

            // Push existing value
            existing_val.map(|v| {
                for b in v {
                    new_val.push(*b);
                }
            });

            // Append new entries
            for op in operands {
                for b in op {
                    new_val.push(*b);
                }
            }

            new_val
        }
        _ => {
            // Unrecognised key, fallback to emulating a put operation (by taking the last value)
            operands.last().unwrap().iter().cloned().collect()
        }
    }
}

pub struct RocksDBIndexStore {
    schema: Arc<Schema>,
    db: DB,
    next_term_ref: AtomicU32,
    term_dictionary: RwLock<BTreeMap<Vec<u8>, TermRef>>,
    next_chunk: AtomicU32,
}

impl RocksDBIndexStore {
    pub fn create(path: &str) -> Result<RocksDBIndexStore, String> {
        let mut opts = Options::default();
        opts.add_merge_operator("merge operator", merge_keys);
        opts.create_if_missing(true);
        let db = try!(DB::open(&opts, path));

        // Schema
        let schema = Schema::new();
        db.put(b".schema", json::encode(&schema).unwrap().as_bytes());

        // Next term ref
        db.put(b".next_term_ref", b"1");

        // Next chunk
        db.put(b".next_chunk", b"1");

        Ok(RocksDBIndexStore {
            schema: Arc::new(schema),
            db: db,
            next_term_ref: AtomicU32::new(1),
            term_dictionary: RwLock::new(BTreeMap::new()),
            next_chunk: AtomicU32::new(1),
        })
    }

    pub fn open(path: &str) -> Result<RocksDBIndexStore, String> {
        let mut opts = Options::default();
        opts.add_merge_operator("merge operator", merge_keys);
        let db = try!(DB::open(&opts, path));

        let schema = match db.get(b".schema") {
            Ok(Some(schema)) => {
                let schema = schema.to_utf8().unwrap().to_string();
                json::decode(&schema).unwrap()
            }
            Ok(None) => Schema::new(),  // TODO: error
            Err(_) => Schema::new(),  // TODO: error
        };

        let next_term_ref = match db.get(b".next_term_ref") {
            Ok(Some(next_term_ref)) => {
                next_term_ref.to_utf8().unwrap().parse::<u32>().unwrap()
            }
            Ok(None) => 1,  // TODO: error
            Err(_) => 1,  // TODO: error
        };

        let next_chunk = match db.get(b".next_chunk") {
            Ok(Some(next_chunk)) => {
                next_chunk.to_utf8().unwrap().parse::<u32>().unwrap()
            }
            Ok(None) => 1,  // TODO: error
            Err(_) => 1,  // TODO: error
        };

        Ok(RocksDBIndexStore {
            schema: Arc::new(schema),
            db: db,
            next_term_ref: AtomicU32::new(next_term_ref),
            term_dictionary: RwLock::new(BTreeMap::new()),
            next_chunk: AtomicU32::new(next_chunk),
        })
    }

    pub fn add_field(&mut self, name: String, field_type: FieldType) -> Result<FieldRef, AddFieldError> {
        let mut schema_copy = (*self.schema).clone();
        let field_ref = try!(schema_copy.add_field(name, field_type));
        self.schema = Arc::new(schema_copy);

        self.db.put(b".schema", json::encode(&self.schema).unwrap().as_bytes());

        Ok(field_ref)
    }

    pub fn remove_field(&mut self, field_ref: &FieldRef) -> bool {
        let mut schema_copy = (*self.schema).clone();
        let field_removed = schema_copy.remove_field(field_ref);

        if field_removed {
            self.schema = Arc::new(schema_copy);
            self.db.put(b".schema", json::encode(&self.schema).unwrap().as_bytes());
        }

        field_removed
    }

    fn get_or_create_term(&mut self, term: &Term) -> TermRef {
        let term_bytes = term.to_bytes();

        if let Some(term_ref) = self.term_dictionary.read().unwrap().get(&term_bytes) {
            return *term_ref;
        }

        // Term doesn't exist in the term dictionary

        // Increment next_term_ref
        let next_term_ref = self.next_term_ref.fetch_add(1, Ordering::SeqCst);
        self.db.put(b".next_term_ref", (next_term_ref + 1).to_string().as_bytes());

        // Create term ref
        let term_ref = TermRef(next_term_ref);

        // Get exclusive lock to term dictionary
        let mut term_dictionary = self.term_dictionary.write().unwrap();

        // It's possible that another thread has written the term to the dictionary
        // since we checked earlier. If this is the case, We should forget about
        // writing our TermRef and use the one that has been inserted already.
        if let Some(term_ref) = term_dictionary.get(&term_bytes) {
            return *term_ref;
        }

        // Write it to the on-disk term dictionary
        let mut key = Vec::with_capacity(1 + term_bytes.len());
        key.push(b't');
        for byte in term_bytes.iter() {
            key.push(*byte);
        }
        self.db.put(&key, next_term_ref.to_string().as_bytes());

        // Write it to the term dictionary
        term_dictionary.insert(term_bytes, term_ref);

        term_ref
    }

    pub fn insert_or_update_document(&mut self, doc: Document) {
        // Allocate a new chunk for the document

        // Chunk merges are very slow so we should avoid doing them at runtime
        // which is why each new document is created in a fresh chunk.
        // Later on, a background process will come and merge any small chunks
        // together.

        // For best performance, documents should be inserted in batches.

        // Increment next_chunk
        let next_chunk = self.next_chunk.fetch_add(1, Ordering::SeqCst);
        self.db.put(b".next_chunk", (next_chunk + 1).to_string().as_bytes());

        // Create doc ref
        let doc_ref = DocRef(next_chunk, 0);

        // Insert contents
        for (field_ref, tokens) in doc.fields.iter() {
            for token in tokens.iter() {
                let term_ref = self.get_or_create_term(&token.term);

                let mut key = Vec::with_capacity(20);
                key.push(b'd');
                for byte in field_ref.ord().to_string().as_bytes() {
                    key.push(*byte);
                }
                key.push(b'/');
                for byte in term_ref.ord().to_string().as_bytes() {
                    key.push(*byte);
                }
                key.push(b'/');
                for byte in doc_ref.chunk().to_string().as_bytes() {
                    key.push(*byte);
                }

                let mut doc_id_bytes = [0; 2];
                BigEndian::write_u16(&mut doc_id_bytes, doc_ref.ord());
                self.db.merge(&key, &doc_id_bytes);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use rocksdb::{DB, Options, IteratorMode};
    use abra::{Term, Token, Document};
    use abra::schema::{Schema, FieldType, FieldRef};

    use super::RocksDBIndexStore;

    fn clean_test_indices() {
        use std::fs::remove_dir_all;

        remove_dir_all("test_indices");
    }

    #[test]
    fn test_create() {
        clean_test_indices();

        let store = RocksDBIndexStore::create("test_indices/test_create");
        assert!(store.is_ok());
    }

    #[test]
    fn test_open() {
        clean_test_indices();

        let store = RocksDBIndexStore::open("test_indices/test_open");
        assert!(store.is_err());

        // Create DB
        let mut opts = Options::default();
        opts.create_if_missing(true);
        DB::open(&opts, "test_indices/test_open").unwrap();

        let store = RocksDBIndexStore::open("test_indices/test_open");
        assert!(store.is_ok());
    }

    fn make_test_store(path: &str) -> RocksDBIndexStore {
        let mut store = RocksDBIndexStore::create(path).unwrap();
        let mut title_field = store.add_field("title".to_string(), FieldType::Text).unwrap();
        let mut body_field = store.add_field("body".to_string(), FieldType::Text).unwrap();

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("hello".to_string()), position: 1 },
                    Token { term: Term::String("world".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("howdy".to_string()), position: 1 },
                    Token { term: Term::String("partner".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store
    }

    pub fn print_keys(db: &DB) {
        fn bytes_to_string(bytes: &Box<[u8]>) -> String {
            use std::char;

            let mut string = String::new();

            for byte in bytes.iter() {
                if *byte < 128 {
                    // ASCII character
                    string.push(char::from_u32(*byte as u32).unwrap());
                } else {
                    string.push('?');
                }
            }

            string
        }

        for (key, value) in db.iterator(IteratorMode::Start) {
            println!("{} = {:?}", bytes_to_string(&key), value);
        }
    }

    #[test]
    fn test() {
        clean_test_indices();

        let store = make_test_store("test_indices/test");

        print_keys(&store.db);

    }
}
<commit_msg>Atomic writes of new documents<commit_after>#![feature(integer_atomics)]

extern crate abra;
extern crate rocksdb;
extern crate rustc_serialize;
#[macro_use]
extern crate maplit;
extern crate byteorder;

use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicU32, Ordering};
use std::collections::BTreeMap;

use rocksdb::{DB, WriteBatch, Writable, Options, MergeOperands};
use abra::{Term, Document};
use abra::schema::{Schema, FieldType, FieldRef, AddFieldError};
use rustc_serialize::{json, Encodable};
use byteorder::{ByteOrder, BigEndian};

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
struct TermRef(u32);

impl TermRef {
    pub fn ord(&self) -> u32 {
        self.0
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
struct DocRef(u32, u16);

impl DocRef {
    pub fn chunk(&self) -> u32 {
        self.0
    }

    pub fn ord(&self) -> u16 {
        self.1
    }
}

fn merge_keys(key: &[u8], existing_val: Option<&[u8]>, operands: &mut MergeOperands) -> Vec<u8> {
    match key[0] {
        b'd' => {
            // Directory (sequence of two byte document ids)

            // Allocate vec for new Value
            let new_size = match existing_val {
                Some(existing_val) => existing_val.len(),
                None => 0,
            } + operands.size_hint().0 * 2;

            let mut new_val = Vec::with_capacity(new_size);

            // Push existing value
            existing_val.map(|v| {
                for b in v {
                    new_val.push(*b);
                }
            });

            // Append new entries
            for op in operands {
                for b in op {
                    new_val.push(*b);
                }
            }

            new_val
        }
        _ => {
            // Unrecognised key, fallback to emulating a put operation (by taking the last value)
            operands.last().unwrap().iter().cloned().collect()
        }
    }
}

pub struct RocksDBIndexStore {
    schema: Arc<Schema>,
    db: DB,
    next_term_ref: AtomicU32,
    term_dictionary: RwLock<BTreeMap<Vec<u8>, TermRef>>,
    next_chunk: AtomicU32,
}

impl RocksDBIndexStore {
    pub fn create(path: &str) -> Result<RocksDBIndexStore, String> {
        let mut opts = Options::default();
        opts.add_merge_operator("merge operator", merge_keys);
        opts.create_if_missing(true);
        let db = try!(DB::open(&opts, path));

        // Schema
        let schema = Schema::new();
        db.put(b".schema", json::encode(&schema).unwrap().as_bytes());

        // Next term ref
        db.put(b".next_term_ref", b"1");

        // Next chunk
        db.put(b".next_chunk", b"1");

        Ok(RocksDBIndexStore {
            schema: Arc::new(schema),
            db: db,
            next_term_ref: AtomicU32::new(1),
            term_dictionary: RwLock::new(BTreeMap::new()),
            next_chunk: AtomicU32::new(1),
        })
    }

    pub fn open(path: &str) -> Result<RocksDBIndexStore, String> {
        let mut opts = Options::default();
        opts.add_merge_operator("merge operator", merge_keys);
        let db = try!(DB::open(&opts, path));

        let schema = match db.get(b".schema") {
            Ok(Some(schema)) => {
                let schema = schema.to_utf8().unwrap().to_string();
                json::decode(&schema).unwrap()
            }
            Ok(None) => Schema::new(),  // TODO: error
            Err(_) => Schema::new(),  // TODO: error
        };

        let next_term_ref = match db.get(b".next_term_ref") {
            Ok(Some(next_term_ref)) => {
                next_term_ref.to_utf8().unwrap().parse::<u32>().unwrap()
            }
            Ok(None) => 1,  // TODO: error
            Err(_) => 1,  // TODO: error
        };

        let next_chunk = match db.get(b".next_chunk") {
            Ok(Some(next_chunk)) => {
                next_chunk.to_utf8().unwrap().parse::<u32>().unwrap()
            }
            Ok(None) => 1,  // TODO: error
            Err(_) => 1,  // TODO: error
        };

        Ok(RocksDBIndexStore {
            schema: Arc::new(schema),
            db: db,
            next_term_ref: AtomicU32::new(next_term_ref),
            term_dictionary: RwLock::new(BTreeMap::new()),
            next_chunk: AtomicU32::new(next_chunk),
        })
    }

    pub fn add_field(&mut self, name: String, field_type: FieldType) -> Result<FieldRef, AddFieldError> {
        let mut schema_copy = (*self.schema).clone();
        let field_ref = try!(schema_copy.add_field(name, field_type));
        self.schema = Arc::new(schema_copy);

        self.db.put(b".schema", json::encode(&self.schema).unwrap().as_bytes());

        Ok(field_ref)
    }

    pub fn remove_field(&mut self, field_ref: &FieldRef) -> bool {
        let mut schema_copy = (*self.schema).clone();
        let field_removed = schema_copy.remove_field(field_ref);

        if field_removed {
            self.schema = Arc::new(schema_copy);
            self.db.put(b".schema", json::encode(&self.schema).unwrap().as_bytes());
        }

        field_removed
    }

    fn get_or_create_term(&mut self, term: &Term) -> TermRef {
        let term_bytes = term.to_bytes();

        if let Some(term_ref) = self.term_dictionary.read().unwrap().get(&term_bytes) {
            return *term_ref;
        }

        // Term doesn't exist in the term dictionary

        // Increment next_term_ref
        let next_term_ref = self.next_term_ref.fetch_add(1, Ordering::SeqCst);
        self.db.put(b".next_term_ref", (next_term_ref + 1).to_string().as_bytes());

        // Create term ref
        let term_ref = TermRef(next_term_ref);

        // Get exclusive lock to term dictionary
        let mut term_dictionary = self.term_dictionary.write().unwrap();

        // It's possible that another thread has written the term to the dictionary
        // since we checked earlier. If this is the case, We should forget about
        // writing our TermRef and use the one that has been inserted already.
        if let Some(term_ref) = term_dictionary.get(&term_bytes) {
            return *term_ref;
        }

        // Write it to the on-disk term dictionary
        let mut key = Vec::with_capacity(1 + term_bytes.len());
        key.push(b't');
        for byte in term_bytes.iter() {
            key.push(*byte);
        }
        self.db.put(&key, next_term_ref.to_string().as_bytes());

        // Write it to the term dictionary
        term_dictionary.insert(term_bytes, term_ref);

        term_ref
    }

    pub fn insert_or_update_document(&mut self, doc: Document) {
        // Allocate a new chunk for the document

        // Chunk merges are very slow so we should avoid doing them at runtime
        // which is why each new document is created in a fresh chunk.
        // Later on, a background process will come and merge any small chunks
        // together.

        // For best performance, documents should be inserted in batches.

        // Increment next_chunk
        let next_chunk = self.next_chunk.fetch_add(1, Ordering::SeqCst);
        self.db.put(b".next_chunk", (next_chunk + 1).to_string().as_bytes());

        // Create doc ref
        let doc_ref = DocRef(next_chunk, 0);

        // Start write batch
        let mut write_batch = WriteBatch::default();

        // Insert contents
        for (field_ref, tokens) in doc.fields.iter() {
            for token in tokens.iter() {
                let term_ref = self.get_or_create_term(&token.term);

                let mut key = Vec::with_capacity(20);
                key.push(b'd');
                for byte in field_ref.ord().to_string().as_bytes() {
                    key.push(*byte);
                }
                key.push(b'/');
                for byte in term_ref.ord().to_string().as_bytes() {
                    key.push(*byte);
                }
                key.push(b'/');
                for byte in doc_ref.chunk().to_string().as_bytes() {
                    key.push(*byte);
                }

                let mut doc_id_bytes = [0; 2];
                BigEndian::write_u16(&mut doc_id_bytes, doc_ref.ord());
                write_batch.merge(&key, &doc_id_bytes);
            }
        }

        // Write
        self.db.write(write_batch);
    }
}

#[cfg(test)]
mod tests {
    use rocksdb::{DB, Options, IteratorMode};
    use abra::{Term, Token, Document};
    use abra::schema::{Schema, FieldType, FieldRef};

    use super::RocksDBIndexStore;

    fn clean_test_indices() {
        use std::fs::remove_dir_all;

        remove_dir_all("test_indices");
    }

    #[test]
    fn test_create() {
        clean_test_indices();

        let store = RocksDBIndexStore::create("test_indices/test_create");
        assert!(store.is_ok());
    }

    #[test]
    fn test_open() {
        clean_test_indices();

        let store = RocksDBIndexStore::open("test_indices/test_open");
        assert!(store.is_err());

        // Create DB
        let mut opts = Options::default();
        opts.create_if_missing(true);
        DB::open(&opts, "test_indices/test_open").unwrap();

        let store = RocksDBIndexStore::open("test_indices/test_open");
        assert!(store.is_ok());
    }

    fn make_test_store(path: &str) -> RocksDBIndexStore {
        let mut store = RocksDBIndexStore::create(path).unwrap();
        let mut title_field = store.add_field("title".to_string(), FieldType::Text).unwrap();
        let mut body_field = store.add_field("body".to_string(), FieldType::Text).unwrap();

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("hello".to_string()), position: 1 },
                    Token { term: Term::String("world".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("howdy".to_string()), position: 1 },
                    Token { term: Term::String("partner".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store
    }

    pub fn print_keys(db: &DB) {
        fn bytes_to_string(bytes: &Box<[u8]>) -> String {
            use std::char;

            let mut string = String::new();

            for byte in bytes.iter() {
                if *byte < 128 {
                    // ASCII character
                    string.push(char::from_u32(*byte as u32).unwrap());
                } else {
                    string.push('?');
                }
            }

            string
        }

        for (key, value) in db.iterator(IteratorMode::Start) {
            println!("{} = {:?}", bytes_to_string(&key), value);
        }
    }

    #[test]
    fn test() {
        clean_test_indices();

        let store = make_test_store("test_indices/test");

        print_keys(&store.db);

    }
}
<|endoftext|>


<commit_before>// Copyright 2018 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 lib features (i.e. features that are not lang features).
//
// These are declared using stability attributes (e.g. `#[stable (..)]`
// and `#[unstable (..)]`), but are not declared in one single location
// (unlike lang features), which means we need to collect them instead.

use ty::TyCtxt;
use syntax::symbol::Symbol;
use syntax::ast::{Attribute, MetaItem, MetaItemKind};
use syntax_pos::{Span, DUMMY_SP};
use hir;
use hir::itemlikevisit::ItemLikeVisitor;
use rustc_data_structures::fx::{FxHashSet, FxHashMap};
use errors::DiagnosticId;

pub struct LibFeatures {
    // A map from feature to stabilisation version.
    pub stable: FxHashMap<Symbol, Symbol>,
    pub unstable: FxHashSet<Symbol>,
}

impl LibFeatures {
    fn new() -> LibFeatures {
        LibFeatures {
            stable: FxHashMap(),
            unstable: FxHashSet(),
        }
    }

    pub fn to_vec(&self) -> Vec<(Symbol, Option<Symbol>)> {
        let mut all_features: Vec<_> = self.stable.iter().map(|(f, s)| (*f, Some(*s)))
            .chain(self.unstable.iter().map(|f| (*f, None)))
            .collect();
        all_features.sort_unstable_by_key(|f| f.0.as_str());
        all_features
    }
}

pub struct LibFeatureCollector<'a, 'tcx: 'a> {
    tcx: TyCtxt<'a, 'tcx, 'tcx>,
    lib_features: LibFeatures,
}

impl<'a, 'tcx> LibFeatureCollector<'a, 'tcx> {
    fn new(tcx: TyCtxt<'a, 'tcx, 'tcx>) -> LibFeatureCollector<'a, 'tcx> {
        LibFeatureCollector {
            tcx,
            lib_features: LibFeatures::new(),
        }
    }

    fn extract(&self, attrs: &[Attribute]) -> Vec<(Symbol, Option<Symbol>, Span)> {
        let stab_attrs = vec!["stable", "unstable", "rustc_const_unstable"];
        let mut features = vec![];

        for attr in attrs {
            // Find a stability attribute (i.e. `#[stable (..)]`, `#[unstable (..)]`,
            // `#[rustc_const_unstable (..)]`).
            if let Some(stab_attr) = stab_attrs.iter().find(|stab_attr| {
                attr.check_name(stab_attr)
            }) {
                let meta_item = attr.meta();
                if let Some(MetaItem { node: MetaItemKind::List(ref metas), .. }) = meta_item {
                    let mut feature = None;
                    let mut since = None;
                    for meta in metas {
                        if let Some(mi) = meta.meta_item() {
                            // Find the `feature = ".."` meta-item.
                            match (&*mi.name().as_str(), mi.value_str()) {
                                ("feature", val) => feature = val,
                                ("since", val) => since = val,
                                _ => {}
                            }
                        }
                    }
                    if let Some(feature) = feature {
                        // This additional check for stability is to make sure we
                        // don't emit additional, irrelevant errors for malformed
                        // attributes.
                        if *stab_attr != "stable" || since.is_some() {
                            features.push((feature, since, attr.span));
                        }
                    }
                    // We need to iterate over the other attributes, because
                    // `rustc_const_unstable` is not mutually exclusive with
                    // the other stability attributes, so we can't just `break`
                    // here.
                }
            }
        }

        features
    }

    fn collect_feature(&mut self, feature: Symbol, since: Option<Symbol>, span: Span) {
        let already_in_stable = self.lib_features.stable.contains_key(&feature);
        let already_in_unstable = self.lib_features.unstable.contains(&feature);

        match (since, already_in_stable, already_in_unstable) {
            (Some(since), _, false) => {
                if let Some(prev_since) = self.lib_features.stable.get(&feature) {
                    if *prev_since != since {
                        let msg = format!(
                            "feature `{}` is declared stable since {}, \
                             but was previously declared stable since {}",
                            feature,
                            since,
                            prev_since,
                        );
                        self.tcx.sess.struct_span_err_with_code(span, &msg,
                            DiagnosticId::Error("E0711".into())).emit();
                        return;
                    }
                }

                self.lib_features.stable.insert(feature, since);
            }
            (None, false, _) => {
                self.lib_features.unstable.insert(feature);
            }
            (Some(_), _, true) | (None, true, _) => {
                let msg = format!(
                    "feature `{}` is declared {}, but was previously declared {}",
                    feature,
                    if since.is_some() { "stable"} else { "unstable" },
                    if since.is_none() { "stable"} else { "unstable" },
                );
                self.tcx.sess.struct_span_err_with_code(span, &msg,
                    DiagnosticId::Error("E0711".into())).emit();
            }
        }
    }

    fn collect_from_attrs(&mut self, attrs: &[Attribute]) {
        for (feature, stable, span) in self.extract(attrs) {
            self.collect_feature(feature, stable, span);
        }
    }
}

impl<'a, 'v, 'tcx> ItemLikeVisitor<'v> for LibFeatureCollector<'a, 'tcx> {
    fn visit_item(&mut self, item: &hir::Item) {
        self.collect_from_attrs(&item.attrs);
    }

    fn visit_trait_item(&mut self, trait_item: &hir::TraitItem) {
        self.collect_from_attrs(&trait_item.attrs);
    }

    fn visit_impl_item(&mut self, impl_item: &hir::ImplItem) {
        self.collect_from_attrs(&impl_item.attrs);
    }
}

pub fn collect<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>) -> LibFeatures {
    let mut collector = LibFeatureCollector::new(tcx);
    for &cnum in tcx.crates().iter() {
        for &(feature, since) in tcx.defined_lib_features(cnum).iter() {
            collector.collect_feature(feature, since, DUMMY_SP);
        }
    }
    collector.collect_from_attrs(&tcx.hir.krate().attrs);
    tcx.hir.krate().visit_all_item_likes(&mut collector);
    for exported_macro in &tcx.hir.krate().exported_macros {
        collector.collect_from_attrs(&exported_macro.attrs);
    }
    collector.lib_features
}
<commit_msg>Visit all attributes for feature collection<commit_after>// Copyright 2018 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 lib features (i.e. features that are not lang features).
//
// These are declared using stability attributes (e.g. `#[stable (..)]`
// and `#[unstable (..)]`), but are not declared in one single location
// (unlike lang features), which means we need to collect them instead.

use ty::TyCtxt;
use syntax::symbol::Symbol;
use syntax::ast::{Attribute, MetaItem, MetaItemKind};
use syntax_pos::{Span, DUMMY_SP};
use hir::intravisit::{self, NestedVisitorMap, Visitor};
use rustc_data_structures::fx::{FxHashSet, FxHashMap};
use errors::DiagnosticId;

pub struct LibFeatures {
    // A map from feature to stabilisation version.
    pub stable: FxHashMap<Symbol, Symbol>,
    pub unstable: FxHashSet<Symbol>,
}

impl LibFeatures {
    fn new() -> LibFeatures {
        LibFeatures {
            stable: FxHashMap(),
            unstable: FxHashSet(),
        }
    }

    pub fn to_vec(&self) -> Vec<(Symbol, Option<Symbol>)> {
        let mut all_features: Vec<_> = self.stable.iter().map(|(f, s)| (*f, Some(*s)))
            .chain(self.unstable.iter().map(|f| (*f, None)))
            .collect();
        all_features.sort_unstable_by_key(|f| f.0.as_str());
        all_features
    }
}

pub struct LibFeatureCollector<'a, 'tcx: 'a> {
    tcx: TyCtxt<'a, 'tcx, 'tcx>,
    lib_features: LibFeatures,
}

impl<'a, 'tcx> LibFeatureCollector<'a, 'tcx> {
    fn new(tcx: TyCtxt<'a, 'tcx, 'tcx>) -> LibFeatureCollector<'a, 'tcx> {
        LibFeatureCollector {
            tcx,
            lib_features: LibFeatures::new(),
        }
    }

    fn extract(&self, attr: &Attribute) -> Option<(Symbol, Option<Symbol>, Span)> {
        let stab_attrs = vec!["stable", "unstable", "rustc_const_unstable"];

        // Find a stability attribute (i.e. `#[stable (..)]`, `#[unstable (..)]`,
        // `#[rustc_const_unstable (..)]`).
        if let Some(stab_attr) = stab_attrs.iter().find(|stab_attr| {
            attr.check_name(stab_attr)
        }) {
            let meta_item = attr.meta();
            if let Some(MetaItem { node: MetaItemKind::List(ref metas), .. }) = meta_item {
                let mut feature = None;
                let mut since = None;
                for meta in metas {
                    if let Some(mi) = meta.meta_item() {
                        // Find the `feature = ".."` meta-item.
                        match (&*mi.name().as_str(), mi.value_str()) {
                            ("feature", val) => feature = val,
                            ("since", val) => since = val,
                            _ => {}
                        }
                    }
                }
                if let Some(feature) = feature {
                    // This additional check for stability is to make sure we
                    // don't emit additional, irrelevant errors for malformed
                    // attributes.
                    if *stab_attr != "stable" || since.is_some() {
                        return Some((feature, since, attr.span));
                    }
                }
                // We need to iterate over the other attributes, because
                // `rustc_const_unstable` is not mutually exclusive with
                // the other stability attributes, so we can't just `break`
                // here.
            }
        }

        None
    }

    fn collect_feature(&mut self, feature: Symbol, since: Option<Symbol>, span: Span) {
        let already_in_stable = self.lib_features.stable.contains_key(&feature);
        let already_in_unstable = self.lib_features.unstable.contains(&feature);

        match (since, already_in_stable, already_in_unstable) {
            (Some(since), _, false) => {
                if let Some(prev_since) = self.lib_features.stable.get(&feature) {
                    if *prev_since != since {
                        let msg = format!(
                            "feature `{}` is declared stable since {}, \
                             but was previously declared stable since {}",
                            feature,
                            since,
                            prev_since,
                        );
                        self.tcx.sess.struct_span_err_with_code(span, &msg,
                            DiagnosticId::Error("E0711".into())).emit();
                        return;
                    }
                }

                self.lib_features.stable.insert(feature, since);
            }
            (None, false, _) => {
                self.lib_features.unstable.insert(feature);
            }
            (Some(_), _, true) | (None, true, _) => {
                let msg = format!(
                    "feature `{}` is declared {}, but was previously declared {}",
                    feature,
                    if since.is_some() { "stable"} else { "unstable" },
                    if since.is_none() { "stable"} else { "unstable" },
                );
                self.tcx.sess.struct_span_err_with_code(span, &msg,
                    DiagnosticId::Error("E0711".into())).emit();
            }
        }
    }
}

impl<'a, 'tcx> Visitor<'tcx> for LibFeatureCollector<'a, 'tcx> {
    fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
        NestedVisitorMap::All(&self.tcx.hir)
    }

    fn visit_attribute(&mut self, attr: &'tcx Attribute) {
        if let Some((feature, stable, span)) = self.extract(attr) {
            self.collect_feature(feature, stable, span);
        }
    }
}

pub fn collect<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>) -> LibFeatures {
    let mut collector = LibFeatureCollector::new(tcx);
    for &cnum in tcx.crates().iter() {
        for &(feature, since) in tcx.defined_lib_features(cnum).iter() {
            collector.collect_feature(feature, since, DUMMY_SP);
        }
    }
    intravisit::walk_crate(&mut collector, tcx.hir.krate());
    collector.lib_features
}
<|endoftext|>


<commit_before>use super::diagnostics::Error;
use super::item::ItemInfo;
use super::Parser;

use crate::{new_sub_parser_from_file, DirectoryOwnership};

use rustc_ast::ast::{self, Attribute, Crate, Ident, ItemKind, Mod};
use rustc_ast::attr;
use rustc_ast::token::{self, TokenKind};
use rustc_errors::PResult;
use rustc_span::source_map::{FileName, SourceMap, Span, DUMMY_SP};
use rustc_span::symbol::sym;

use std::path::{self, Path, PathBuf};

/// Information about the path to a module.
// Public for rustfmt usage.
pub struct ModulePath {
    name: String,
    path_exists: bool,
    pub result: Result<ModulePathSuccess, Error>,
}

// Public for rustfmt usage.
pub struct ModulePathSuccess {
    pub path: PathBuf,
    pub directory_ownership: DirectoryOwnership,
}

impl<'a> Parser<'a> {
    /// Parses a source module as a crate. This is the main entry point for the parser.
    pub fn parse_crate_mod(&mut self) -> PResult<'a, Crate> {
        let lo = self.token.span;
        let krate = Ok(ast::Crate {
            attrs: self.parse_inner_attributes()?,
            module: self.parse_mod_items(&token::Eof, lo)?,
            span: lo.to(self.token.span),
            // Filled in by proc_macro_harness::inject()
            proc_macros: Vec::new(),
        });
        krate
    }

    /// Parses a `mod <foo> { ... }` or `mod <foo>;` item.
    pub(super) fn parse_item_mod(&mut self, attrs: &mut Vec<Attribute>) -> PResult<'a, ItemInfo> {
        let in_cfg = crate::config::process_configure_mod(self.sess, self.cfg_mods, attrs);

        let id_span = self.token.span;
        let id = self.parse_ident()?;
        let (module, mut inner_attrs) = if self.eat(&token::Semi) {
            if in_cfg && self.recurse_into_file_modules {
                // This mod is in an external file. Let's go get it!
                let ModulePathSuccess { path, directory_ownership } =
                    self.submod_path(id, &attrs, id_span)?;
                self.eval_src_mod(path, directory_ownership, id.to_string(), id_span)?
            } else {
                (ast::Mod { inner: DUMMY_SP, items: Vec::new(), inline: false }, Vec::new())
            }
        } else {
            let old_directory = self.directory.clone();
            self.push_directory(id, &attrs);

            self.expect(&token::OpenDelim(token::Brace))?;
            let mod_inner_lo = self.token.span;
            let inner_attrs = self.parse_inner_attributes()?;
            let module = self.parse_mod_items(&token::CloseDelim(token::Brace), mod_inner_lo)?;

            self.directory = old_directory;
            (module, inner_attrs)
        };
        attrs.append(&mut inner_attrs);
        Ok((id, ItemKind::Mod(module)))
    }

    /// Given a termination token, parses all of the items in a module.
    fn parse_mod_items(&mut self, term: &TokenKind, inner_lo: Span) -> PResult<'a, Mod> {
        let mut items = vec![];
        while let Some(item) = self.parse_item()? {
            items.push(item);
            self.maybe_consume_incorrect_semicolon(&items);
        }

        if !self.eat(term) {
            let token_str = super::token_descr(&self.token);
            if !self.maybe_consume_incorrect_semicolon(&items) {
                let msg = &format!("expected item, found {}", token_str);
                let mut err = self.struct_span_err(self.token.span, msg);
                err.span_label(self.token.span, "expected item");
                return Err(err);
            }
        }

        let hi = if self.token.span.is_dummy() { inner_lo } else { self.prev_token.span };

        Ok(Mod { inner: inner_lo.to(hi), items, inline: true })
    }

    fn submod_path(
        &mut self,
        id: ast::Ident,
        outer_attrs: &[Attribute],
        id_sp: Span,
    ) -> PResult<'a, ModulePathSuccess> {
        if let Some(path) = Parser::submod_path_from_attr(outer_attrs, &self.directory.path) {
            return Ok(ModulePathSuccess {
                directory_ownership: match path.file_name().and_then(|s| s.to_str()) {
                    // All `#[path]` files are treated as though they are a `mod.rs` file.
                    // This means that `mod foo;` declarations inside `#[path]`-included
                    // files are siblings,
                    //
                    // Note that this will produce weirdness when a file named `foo.rs` is
                    // `#[path]` included and contains a `mod foo;` declaration.
                    // If you encounter this, it's your own darn fault :P
                    Some(_) => DirectoryOwnership::Owned { relative: None },
                    _ => DirectoryOwnership::UnownedViaMod,
                },
                path,
            });
        }

        let relative = match self.directory.ownership {
            DirectoryOwnership::Owned { relative } => relative,
            DirectoryOwnership::UnownedViaBlock | DirectoryOwnership::UnownedViaMod => None,
        };
        let paths =
            Parser::default_submod_path(id, relative, &self.directory.path, self.sess.source_map());

        match self.directory.ownership {
            DirectoryOwnership::Owned { .. } => {
                paths.result.map_err(|err| self.span_fatal_err(id_sp, err))
            }
            DirectoryOwnership::UnownedViaBlock => {
                let msg = "Cannot declare a non-inline module inside a block \
                    unless it has a path attribute";
                let mut err = self.struct_span_err(id_sp, msg);
                if paths.path_exists {
                    let msg = format!(
                        "Maybe `use` the module `{}` instead of redeclaring it",
                        paths.name
                    );
                    err.span_note(id_sp, &msg);
                }
                Err(err)
            }
            DirectoryOwnership::UnownedViaMod => {
                let mut err =
                    self.struct_span_err(id_sp, "cannot declare a new module at this location");
                if !id_sp.is_dummy() {
                    let src_path = self.sess.source_map().span_to_filename(id_sp);
                    if let FileName::Real(src_path) = src_path {
                        if let Some(stem) = src_path.file_stem() {
                            let mut dest_path = src_path.clone();
                            dest_path.set_file_name(stem);
                            dest_path.push("mod.rs");
                            err.span_note(
                                id_sp,
                                &format!(
                                    "maybe move this module `{}` to its own \
                                                directory via `{}`",
                                    src_path.display(),
                                    dest_path.display()
                                ),
                            );
                        }
                    }
                }
                if paths.path_exists {
                    err.span_note(
                        id_sp,
                        &format!(
                            "... or maybe `use` the module `{}` instead \
                                            of possibly redeclaring it",
                            paths.name
                        ),
                    );
                }
                Err(err)
            }
        }
    }

    // Public for rustfmt usage.
    pub fn submod_path_from_attr(attrs: &[Attribute], dir_path: &Path) -> Option<PathBuf> {
        if let Some(s) = attr::first_attr_value_str_by_name(attrs, sym::path) {
            let s = s.as_str();

            // On windows, the base path might have the form
            // `\\?\foo\bar` in which case it does not tolerate
            // mixed `/` and `\` separators, so canonicalize
            // `/` to `\`.
            #[cfg(windows)]
            let s = s.replace("/", "\\");
            Some(dir_path.join(&*s))
        } else {
            None
        }
    }

    /// Returns a path to a module.
    // Public for rustfmt usage.
    pub fn default_submod_path(
        id: ast::Ident,
        relative: Option<ast::Ident>,
        dir_path: &Path,
        source_map: &SourceMap,
    ) -> ModulePath {
        // If we're in a foo.rs file instead of a mod.rs file,
        // we need to look for submodules in
        // `./foo/<id>.rs` and `./foo/<id>/mod.rs` rather than
        // `./<id>.rs` and `./<id>/mod.rs`.
        let relative_prefix_string;
        let relative_prefix = if let Some(ident) = relative {
            relative_prefix_string = format!("{}{}", ident.name, path::MAIN_SEPARATOR);
            &relative_prefix_string
        } else {
            ""
        };

        let mod_name = id.name.to_string();
        let default_path_str = format!("{}{}.rs", relative_prefix, mod_name);
        let secondary_path_str =
            format!("{}{}{}mod.rs", relative_prefix, mod_name, path::MAIN_SEPARATOR);
        let default_path = dir_path.join(&default_path_str);
        let secondary_path = dir_path.join(&secondary_path_str);
        let default_exists = source_map.file_exists(&default_path);
        let secondary_exists = source_map.file_exists(&secondary_path);

        let result = match (default_exists, secondary_exists) {
            (true, false) => Ok(ModulePathSuccess {
                path: default_path,
                directory_ownership: DirectoryOwnership::Owned { relative: Some(id) },
            }),
            (false, true) => Ok(ModulePathSuccess {
                path: secondary_path,
                directory_ownership: DirectoryOwnership::Owned { relative: None },
            }),
            (false, false) => {
                Err(Error::FileNotFoundForModule { mod_name: mod_name.clone(), default_path })
            }
            (true, true) => Err(Error::DuplicatePaths {
                mod_name: mod_name.clone(),
                default_path: default_path_str,
                secondary_path: secondary_path_str,
            }),
        };

        ModulePath { name: mod_name, path_exists: default_exists || secondary_exists, result }
    }

    /// Reads a module from a source file.
    fn eval_src_mod(
        &mut self,
        path: PathBuf,
        directory_ownership: DirectoryOwnership,
        name: String,
        id_sp: Span,
    ) -> PResult<'a, (Mod, Vec<Attribute>)> {
        let mut included_mod_stack = self.sess.included_mod_stack.borrow_mut();
        if let Some(i) = included_mod_stack.iter().position(|p| *p == path) {
            let mut err = String::from("circular modules: ");
            let len = included_mod_stack.len();
            for p in &included_mod_stack[i..len] {
                err.push_str(&p.to_string_lossy());
                err.push_str(" -> ");
            }
            err.push_str(&path.to_string_lossy());
            return Err(self.struct_span_err(id_sp, &err[..]));
        }
        included_mod_stack.push(path.clone());
        drop(included_mod_stack);

        let mut p0 =
            new_sub_parser_from_file(self.sess, &path, directory_ownership, Some(name), id_sp);
        p0.cfg_mods = self.cfg_mods;
        let mod_inner_lo = p0.token.span;
        let mod_attrs = p0.parse_inner_attributes()?;
        let mut m0 = p0.parse_mod_items(&token::Eof, mod_inner_lo)?;
        m0.inline = false;
        self.sess.included_mod_stack.borrow_mut().pop();
        Ok((m0, mod_attrs))
    }

    fn push_directory(&mut self, id: Ident, attrs: &[Attribute]) {
        if let Some(path) = attr::first_attr_value_str_by_name(attrs, sym::path) {
            self.directory.path.push(&*path.as_str());
            self.directory.ownership = DirectoryOwnership::Owned { relative: None };
        } else {
            // We have to push on the current module name in the case of relative
            // paths in order to ensure that any additional module paths from inline
            // `mod x { ... }` come after the relative extension.
            //
            // For example, a `mod z { ... }` inside `x/y.rs` should set the current
            // directory path to `/x/y/z`, not `/x/z` with a relative offset of `y`.
            if let DirectoryOwnership::Owned { relative } = &mut self.directory.ownership {
                if let Some(ident) = relative.take() {
                    // remove the relative offset
                    self.directory.path.push(&*ident.as_str());
                }
            }
            self.directory.path.push(&*id.as_str());
        }
    }
}
<commit_msg>submod_path_from_attr: simplify & document<commit_after>use super::diagnostics::Error;
use super::item::ItemInfo;
use super::Parser;

use crate::{new_sub_parser_from_file, DirectoryOwnership};

use rustc_ast::ast::{self, Attribute, Crate, Ident, ItemKind, Mod};
use rustc_ast::attr;
use rustc_ast::token::{self, TokenKind};
use rustc_errors::PResult;
use rustc_span::source_map::{FileName, SourceMap, Span, DUMMY_SP};
use rustc_span::symbol::sym;

use std::path::{self, Path, PathBuf};

/// Information about the path to a module.
// Public for rustfmt usage.
pub struct ModulePath {
    name: String,
    path_exists: bool,
    pub result: Result<ModulePathSuccess, Error>,
}

// Public for rustfmt usage.
pub struct ModulePathSuccess {
    pub path: PathBuf,
    pub directory_ownership: DirectoryOwnership,
}

impl<'a> Parser<'a> {
    /// Parses a source module as a crate. This is the main entry point for the parser.
    pub fn parse_crate_mod(&mut self) -> PResult<'a, Crate> {
        let lo = self.token.span;
        let krate = Ok(ast::Crate {
            attrs: self.parse_inner_attributes()?,
            module: self.parse_mod_items(&token::Eof, lo)?,
            span: lo.to(self.token.span),
            // Filled in by proc_macro_harness::inject()
            proc_macros: Vec::new(),
        });
        krate
    }

    /// Parses a `mod <foo> { ... }` or `mod <foo>;` item.
    pub(super) fn parse_item_mod(&mut self, attrs: &mut Vec<Attribute>) -> PResult<'a, ItemInfo> {
        let in_cfg = crate::config::process_configure_mod(self.sess, self.cfg_mods, attrs);

        let id_span = self.token.span;
        let id = self.parse_ident()?;
        let (module, mut inner_attrs) = if self.eat(&token::Semi) {
            if in_cfg && self.recurse_into_file_modules {
                // This mod is in an external file. Let's go get it!
                let ModulePathSuccess { path, directory_ownership } =
                    self.submod_path(id, &attrs, id_span)?;
                self.eval_src_mod(path, directory_ownership, id.to_string(), id_span)?
            } else {
                (ast::Mod { inner: DUMMY_SP, items: Vec::new(), inline: false }, Vec::new())
            }
        } else {
            let old_directory = self.directory.clone();
            self.push_directory(id, &attrs);

            self.expect(&token::OpenDelim(token::Brace))?;
            let mod_inner_lo = self.token.span;
            let inner_attrs = self.parse_inner_attributes()?;
            let module = self.parse_mod_items(&token::CloseDelim(token::Brace), mod_inner_lo)?;

            self.directory = old_directory;
            (module, inner_attrs)
        };
        attrs.append(&mut inner_attrs);
        Ok((id, ItemKind::Mod(module)))
    }

    /// Given a termination token, parses all of the items in a module.
    fn parse_mod_items(&mut self, term: &TokenKind, inner_lo: Span) -> PResult<'a, Mod> {
        let mut items = vec![];
        while let Some(item) = self.parse_item()? {
            items.push(item);
            self.maybe_consume_incorrect_semicolon(&items);
        }

        if !self.eat(term) {
            let token_str = super::token_descr(&self.token);
            if !self.maybe_consume_incorrect_semicolon(&items) {
                let msg = &format!("expected item, found {}", token_str);
                let mut err = self.struct_span_err(self.token.span, msg);
                err.span_label(self.token.span, "expected item");
                return Err(err);
            }
        }

        let hi = if self.token.span.is_dummy() { inner_lo } else { self.prev_token.span };

        Ok(Mod { inner: inner_lo.to(hi), items, inline: true })
    }

    fn submod_path(
        &mut self,
        id: ast::Ident,
        outer_attrs: &[Attribute],
        id_sp: Span,
    ) -> PResult<'a, ModulePathSuccess> {
        if let Some(path) = Parser::submod_path_from_attr(outer_attrs, &self.directory.path) {
            return Ok(ModulePathSuccess {
                directory_ownership: match path.file_name().and_then(|s| s.to_str()) {
                    // All `#[path]` files are treated as though they are a `mod.rs` file.
                    // This means that `mod foo;` declarations inside `#[path]`-included
                    // files are siblings,
                    //
                    // Note that this will produce weirdness when a file named `foo.rs` is
                    // `#[path]` included and contains a `mod foo;` declaration.
                    // If you encounter this, it's your own darn fault :P
                    Some(_) => DirectoryOwnership::Owned { relative: None },
                    _ => DirectoryOwnership::UnownedViaMod,
                },
                path,
            });
        }

        let relative = match self.directory.ownership {
            DirectoryOwnership::Owned { relative } => relative,
            DirectoryOwnership::UnownedViaBlock | DirectoryOwnership::UnownedViaMod => None,
        };
        let paths =
            Parser::default_submod_path(id, relative, &self.directory.path, self.sess.source_map());

        match self.directory.ownership {
            DirectoryOwnership::Owned { .. } => {
                paths.result.map_err(|err| self.span_fatal_err(id_sp, err))
            }
            DirectoryOwnership::UnownedViaBlock => {
                let msg = "Cannot declare a non-inline module inside a block \
                    unless it has a path attribute";
                let mut err = self.struct_span_err(id_sp, msg);
                if paths.path_exists {
                    let msg = format!(
                        "Maybe `use` the module `{}` instead of redeclaring it",
                        paths.name
                    );
                    err.span_note(id_sp, &msg);
                }
                Err(err)
            }
            DirectoryOwnership::UnownedViaMod => {
                let mut err =
                    self.struct_span_err(id_sp, "cannot declare a new module at this location");
                if !id_sp.is_dummy() {
                    let src_path = self.sess.source_map().span_to_filename(id_sp);
                    if let FileName::Real(src_path) = src_path {
                        if let Some(stem) = src_path.file_stem() {
                            let mut dest_path = src_path.clone();
                            dest_path.set_file_name(stem);
                            dest_path.push("mod.rs");
                            err.span_note(
                                id_sp,
                                &format!(
                                    "maybe move this module `{}` to its own \
                                                directory via `{}`",
                                    src_path.display(),
                                    dest_path.display()
                                ),
                            );
                        }
                    }
                }
                if paths.path_exists {
                    err.span_note(
                        id_sp,
                        &format!(
                            "... or maybe `use` the module `{}` instead \
                                            of possibly redeclaring it",
                            paths.name
                        ),
                    );
                }
                Err(err)
            }
        }
    }

    /// Derive a submodule path from the first found `#[path = "path_string"]`.
    /// The provided `dir_path` is joined with the `path_string`.
    // Public for rustfmt usage.
    pub fn submod_path_from_attr(attrs: &[Attribute], dir_path: &Path) -> Option<PathBuf> {
        // Extract path string from first `#[path = "path_string"]` attribute.
        let path_string = attr::first_attr_value_str_by_name(attrs, sym::path)?;
        let path_string = path_string.as_str();

        // On windows, the base path might have the form
        // `\\?\foo\bar` in which case it does not tolerate
        // mixed `/` and `\` separators, so canonicalize
        // `/` to `\`.
        #[cfg(windows)]
        let path_string = path_string.replace("/", "\\");

        Some(dir_path.join(&*path_string))
    }

    /// Returns a path to a module.
    // Public for rustfmt usage.
    pub fn default_submod_path(
        id: ast::Ident,
        relative: Option<ast::Ident>,
        dir_path: &Path,
        source_map: &SourceMap,
    ) -> ModulePath {
        // If we're in a foo.rs file instead of a mod.rs file,
        // we need to look for submodules in
        // `./foo/<id>.rs` and `./foo/<id>/mod.rs` rather than
        // `./<id>.rs` and `./<id>/mod.rs`.
        let relative_prefix_string;
        let relative_prefix = if let Some(ident) = relative {
            relative_prefix_string = format!("{}{}", ident.name, path::MAIN_SEPARATOR);
            &relative_prefix_string
        } else {
            ""
        };

        let mod_name = id.name.to_string();
        let default_path_str = format!("{}{}.rs", relative_prefix, mod_name);
        let secondary_path_str =
            format!("{}{}{}mod.rs", relative_prefix, mod_name, path::MAIN_SEPARATOR);
        let default_path = dir_path.join(&default_path_str);
        let secondary_path = dir_path.join(&secondary_path_str);
        let default_exists = source_map.file_exists(&default_path);
        let secondary_exists = source_map.file_exists(&secondary_path);

        let result = match (default_exists, secondary_exists) {
            (true, false) => Ok(ModulePathSuccess {
                path: default_path,
                directory_ownership: DirectoryOwnership::Owned { relative: Some(id) },
            }),
            (false, true) => Ok(ModulePathSuccess {
                path: secondary_path,
                directory_ownership: DirectoryOwnership::Owned { relative: None },
            }),
            (false, false) => {
                Err(Error::FileNotFoundForModule { mod_name: mod_name.clone(), default_path })
            }
            (true, true) => Err(Error::DuplicatePaths {
                mod_name: mod_name.clone(),
                default_path: default_path_str,
                secondary_path: secondary_path_str,
            }),
        };

        ModulePath { name: mod_name, path_exists: default_exists || secondary_exists, result }
    }

    /// Reads a module from a source file.
    fn eval_src_mod(
        &mut self,
        path: PathBuf,
        directory_ownership: DirectoryOwnership,
        name: String,
        id_sp: Span,
    ) -> PResult<'a, (Mod, Vec<Attribute>)> {
        let mut included_mod_stack = self.sess.included_mod_stack.borrow_mut();
        if let Some(i) = included_mod_stack.iter().position(|p| *p == path) {
            let mut err = String::from("circular modules: ");
            let len = included_mod_stack.len();
            for p in &included_mod_stack[i..len] {
                err.push_str(&p.to_string_lossy());
                err.push_str(" -> ");
            }
            err.push_str(&path.to_string_lossy());
            return Err(self.struct_span_err(id_sp, &err[..]));
        }
        included_mod_stack.push(path.clone());
        drop(included_mod_stack);

        let mut p0 =
            new_sub_parser_from_file(self.sess, &path, directory_ownership, Some(name), id_sp);
        p0.cfg_mods = self.cfg_mods;
        let mod_inner_lo = p0.token.span;
        let mod_attrs = p0.parse_inner_attributes()?;
        let mut m0 = p0.parse_mod_items(&token::Eof, mod_inner_lo)?;
        m0.inline = false;
        self.sess.included_mod_stack.borrow_mut().pop();
        Ok((m0, mod_attrs))
    }

    fn push_directory(&mut self, id: Ident, attrs: &[Attribute]) {
        if let Some(path) = attr::first_attr_value_str_by_name(attrs, sym::path) {
            self.directory.path.push(&*path.as_str());
            self.directory.ownership = DirectoryOwnership::Owned { relative: None };
        } else {
            // We have to push on the current module name in the case of relative
            // paths in order to ensure that any additional module paths from inline
            // `mod x { ... }` come after the relative extension.
            //
            // For example, a `mod z { ... }` inside `x/y.rs` should set the current
            // directory path to `/x/y/z`, not `/x/z` with a relative offset of `y`.
            if let DirectoryOwnership::Owned { relative } = &mut self.directory.ownership {
                if let Some(ident) = relative.take() {
                    // remove the relative offset
                    self.directory.path.push(&*ident.as_str());
                }
            }
            self.directory.path.push(&*id.as_str());
        }
    }
}
<|endoftext|>