<commit_before>#![feature(tcp)]

use std::net::{TcpListener, TcpStream};
use std::thread;
use std::io;
use std::env;
use std::mem;
use std::str::FromStr;
use std::io::Read;
use std::io::Write;

fn main() {
    let port_ : Option<String> = env::args().skip(2).next();
    if let None = port_ {
        println!("usage: {} -p port", env::args().next().unwrap());
        return;
    }

    let port = u16::from_str(&port_.unwrap());
    if let Err(ref e) = port {
        println!("invalid port number given: {}", e);
        return;
    }

    let listener = TcpListener::bind(("127.0.0.1", port.unwrap()));
    if let Err(ref e) = listener {
        println!("failed to listen on port: {}", e);
        return;
    }

    for stream in listener.unwrap().incoming() {
        match stream {
            Ok(stream) => {
                thread::spawn(move|| {
                    handle_client(stream)
                });
            }
            Err(e) => {
                println!("failed to accept connection: {}", e);
            }
        }
    }
}

fn handle_client(mut stream: TcpStream) {
    let mut buf = [0u8; 4];
    let mut challenge = 0u32;
    let mut nread;

    let _ = stream.set_nodelay(true);

    loop {
        nread = 0;
        while nread < buf.len() {
            match stream.read(&mut buf[nread..]) {
                Ok(n) if n == 0 => return,
                Ok(n) => nread += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => {
                    println!("{}", e);
                    return;
                }
            }
        }

        challenge = unsafe { mem::transmute(buf) };
        challenge = u32::from_be(challenge);
        if challenge == 0 {
            std::process::exit(0);
        }
        challenge = u32::to_be(challenge + 1);
        buf = unsafe { mem::transmute(challenge) };

        let mut nwritten = 0;
        while nwritten < buf.len() {
            match stream.write(&buf[nwritten..]) {
                Ok(n) => nwritten += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => {
                    println!("{}", e);
                    return;
                }
            }
        }
    }
}
<commit_msg>No need to initialize challenge to zero<commit_after>#![feature(tcp)]

use std::net::{TcpListener, TcpStream};
use std::thread;
use std::io;
use std::env;
use std::mem;
use std::str::FromStr;
use std::io::Read;
use std::io::Write;

fn main() {
    let port_ : Option<String> = env::args().skip(2).next();
    if let None = port_ {
        println!("usage: {} -p port", env::args().next().unwrap());
        return;
    }

    let port = u16::from_str(&port_.unwrap());
    if let Err(ref e) = port {
        println!("invalid port number given: {}", e);
        return;
    }

    let listener = TcpListener::bind(("127.0.0.1", port.unwrap()));
    if let Err(ref e) = listener {
        println!("failed to listen on port: {}", e);
        return;
    }

    for stream in listener.unwrap().incoming() {
        match stream {
            Ok(stream) => {
                thread::spawn(move|| {
                    handle_client(stream)
                });
            }
            Err(e) => {
                println!("failed to accept connection: {}", e);
            }
        }
    }
}

fn handle_client(mut stream: TcpStream) {
    let mut buf = [0u8; 4];
    let mut challenge;
    let mut nread;

    let _ = stream.set_nodelay(true);

    loop {
        nread = 0;
        while nread < buf.len() {
            match stream.read(&mut buf[nread..]) {
                Ok(n) if n == 0 => return,
                Ok(n) => nread += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => {
                    println!("{}", e);
                    return;
                }
            }
        }

        challenge = unsafe { mem::transmute(buf) };
        challenge = u32::from_be(challenge);
        if challenge == 0 {
            std::process::exit(0);
        }
        challenge = u32::to_be(challenge + 1);
        buf = unsafe { mem::transmute(challenge) };

        let mut nwritten = 0;
        while nwritten < buf.len() {
            match stream.write(&buf[nwritten..]) {
                Ok(n) => nwritten += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => {
                    println!("{}", e);
                    return;
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Interrupt description and set-up code.

use core::fmt;

/// x86 Exception description (see also Intel Vol. 3a Chapter 6).
#[derive(Debug)]
pub struct InterruptDescription {
    pub vector: u8,
    pub mnemonic: &'static str,
    pub description: &'static str,
    pub irqtype: &'static str,
    pub source: &'static str,
}

impl fmt::Display for InterruptDescription {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "{} ({}, vec={}) {}",
               self.mnemonic,
               self.irqtype,
               self.vector,
               self.description)
    }
}

/// x86 External Interrupts (1-16).
pub static EXCEPTIONS: [InterruptDescription; 15] = [InterruptDescription {
                                                         vector: 0,
                                                         mnemonic: "#DE",
                                                         description: "Divide Error",
                                                         irqtype: "Fault",
                                                         source: "DIV and IDIV instructions.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 1,
                                                         mnemonic: "#DB",
                                                         description: "RESERVED",
                                                         irqtype: "Fault/ Trap",
                                                         source: "For Intel use only.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 2,
                                                         mnemonic: "NMI",
                                                         description: "Interrupt",
                                                         irqtype: "Interrupt",
                                                         source: "Nonmaskable external interrupt.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 3,
                                                         mnemonic: "#BP",
                                                         description: "Breakpoint",
                                                         irqtype: "Trap",
                                                         source: "INT 3 instruction.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 4,
                                                         mnemonic: "#OF",
                                                         description: "Overflow",
                                                         irqtype: "Trap",
                                                         source: "INTO instruction.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 5,
                                                         mnemonic: "#BR",
                                                         description: "BOUND Range Exceeded",
                                                         irqtype: "Fault",
                                                         source: "BOUND instruction.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 6,
                                                         mnemonic: "#UD",
                                                         description: "Invalid Opcode (Undefined \
                                                                       Opcode)",
                                                         irqtype: "Fault",
                                                         source: "UD2 instruction or reserved \
                                                                  opcode.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 7,
                                                         mnemonic: "#NM",
                                                         description: "Device Not Available (No \
                                                                       Math Coprocessor)",
                                                         irqtype: "Fault",
                                                         source: "Floating-point or WAIT/FWAIT \
                                                                  instruction.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 8,
                                                         mnemonic: "#DF",
                                                         description: "Double Fault",
                                                         irqtype: "Abort",
                                                         source: "Any instruction that can \
                                                                  generate an exception, an NMI, \
                                                                  or an INTR.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 9,
                                                         mnemonic: "",
                                                         description: "Coprocessor Segment Overrun",
                                                         irqtype: "Fault",
                                                         source: "Floating-point instruction.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 10,
                                                         mnemonic: "#TS",
                                                         description: "Invalid TSS",
                                                         irqtype: "Fault",
                                                         source: "Task switch or TSS access.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 11,
                                                         mnemonic: "#NP",
                                                         description: "Segment Not Present",
                                                         irqtype: "Fault",
                                                         source: "Loading segment registers or \
                                                                  accessing system segments.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 12,
                                                         mnemonic: "#SS",
                                                         description: "Stack-Segment Fault",
                                                         irqtype: "Fault",
                                                         source: "Stack operations and SS register \
                                                                  loads.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 13,
                                                         mnemonic: "#GP",
                                                         description: "General Protection",
                                                         irqtype: "Fault",
                                                         source: "Any memory reference and other \
                                                                  protection checks.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 14,
                                                         mnemonic: "#PF",
                                                         description: "Page Fault",
                                                         irqtype: "Fault",
                                                         source: "Any memory reference.",
                                                     }];

/// Enable Interrupts.
pub unsafe fn enable() {
    asm!("sti");
}

/// Disable Interrupts.
pub unsafe fn disable() {
    asm!("cli");
}

/// Generate a software interrupt.
/// This is a macro argument needs to be an immediate.
#[macro_export]
macro_rules! int {
    ( $x:expr ) => {
        {
            asm!("int $0" :: "N" ($x));
        }
    };
}

/// A struct describing an interrupt gate.  See the Intel manual mentioned
/// above for details, specifically, the section "6.14.1 64-Bit Mode IDT"
/// and "Table 3-2. System-Segment and Gate-Descriptor Types".
#[derive(Debug, Clone, Copy)]
#[repr(C, packed)]
pub struct IdtEntry {
    /// Lower 16 bits of ISR.
    pub base_lo: u16,
    /// Segment selector.
    pub sel: u16,
    /// This must always be zero.
    pub res0: u8,
    /// Flags.
    pub flags: u8,
    /// The upper 48 bits of ISR (the last 16 bits must be zero).
    pub base_hi: u64,
    /// Must be zero.
    pub res1: u16,
}

impl IdtEntry {
    /// Create a "missing" IdtEntry.  This is a `const` function, so we can
    /// call it at compile time to initialize static variables.
    ///
    /// If the CPU tries to invoke a missing interrupt, it will instead
    /// send a General Protection fault (13), with the interrupt number and
    /// some other data stored in the error code.
    pub const fn missing() -> IdtEntry {
        IdtEntry {
            base_lo: 0,
            sel: 0,
            res0: 0,
            flags: 0,
            base_hi: 0,
            res1: 0,
        }
    }

    /// Create a new IdtEntry pointing at `handler`, which must be a
    /// function with interrupt calling conventions.  (This must be
    /// currently defined in assembly language.)  The `gdt_code_selector`
    /// value must be the offset of code segment entry in the GDT.
    ///
    /// Create an interrupt gate with the "Present" flag set, which is the
    /// most common case.  If you need something else, you can construct it
    /// manually.
    pub fn interrupt_gate(gdt_code_selector: u16, handler: *const u8) -> IdtEntry {
        IdtEntry {
            base_lo: ((handler as u64) & 0xFFFF) as u16,
            sel: gdt_code_selector,
            res0: 0,
            // Bit 7: "Present" flag set.
            // Bits 0-4: This is an interrupt gate.
            flags: 0b1000_1110,
            base_hi: (handler as u64) >> 16,
            res1: 0,
        }
    }
}

bitflags!{
    // Taken from Intel Manual Section 4.7 Page-Fault Exceptions.
    flags PageFaultError: u32 {
        /// 0: The fault was caused by a non-present page.
        /// 1: The fault was caused by a page-level protection violation
        const PFAULT_ERROR_P = bit!(0),

        /// 0: The access causing the fault was a read.
        /// 1: The access causing the fault was a write.
        const PFAULT_ERROR_WR = bit!(1),

        /// 0: The access causing the fault originated when the processor
        /// was executing in supervisor mode.
        /// 1: The access causing the fault originated when the processor
        /// was executing in user mode.
        const PFAULT_ERROR_US = bit!(2),

        /// 0: The fault was not caused by reserved bit violation.
        /// 1: The fault was caused by reserved bits set to 1 in a page directory.
        const PFAULT_ERROR_RSVD = bit!(3),

        /// 0: The fault was not caused by an instruction fetch.
        /// 1: The fault was caused by an instruction fetch.
        const PFAULT_ERROR_ID = bit!(4),

        /// 0: The fault was not by protection keys.
        /// 1: There was a protection key violation.
        const PFAULT_ERROR_PK = bit!(5),
    }
}

impl fmt::Debug for PageFaultError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let p = match self.contains(PFAULT_ERROR_P) {
            false => "The fault was caused by a non-present page.",
            true => "The fault was caused by a page-level protection violation.",
        };
        let wr = match self.contains(PFAULT_ERROR_WR) {
            false => "The access causing the fault was a read.",
            true => "The access causing the fault was a write.",
        };
        let us = match self.contains(PFAULT_ERROR_US) {
            false => {
                "The access causing the fault originated when the processor was executing in \
                 supervisor mode."
            }
            true => {
                "The access causing the fault originated when the processor was executing in user \
                 mode."
            }
        };
        let rsvd = match self.contains(PFAULT_ERROR_RSVD) {
            false => "The fault was not caused by reserved bit violation.",
            true => "The fault was caused by reserved bits set to 1 in a page directory.",
        };
        let id = match self.contains(PFAULT_ERROR_ID) {
            false => "The fault was not caused by an instruction fetch.",
            true => "The fault was caused by an instruction fetch.",
        };

        write!(f, "{}\n{}\n{}\n{}\n{}", p, wr, us, rsvd, id)
    }
}

#[test]
fn bit_macro() {
    assert!(PFAULT_ERROR_PK.bits() == 0b100000);
    assert!(PFAULT_ERROR_ID.bits() == 0b10000);
    assert!(PFAULT_ERROR_RSVD.bits() == 0b1000);
    assert!(PFAULT_ERROR_US.bits() == 0b100);
    assert!(PFAULT_ERROR_WR.bits() == 0b10);
    assert!(PFAULT_ERROR_P.bits() == 0b1);
}
<commit_msg>Modify irq::interrupt_gate() signature<commit_after>//! Interrupt description and set-up code.

use core::fmt;
use paging::VAddr;

/// x86 Exception description (see also Intel Vol. 3a Chapter 6).
#[derive(Debug)]
pub struct InterruptDescription {
    pub vector: u8,
    pub mnemonic: &'static str,
    pub description: &'static str,
    pub irqtype: &'static str,
    pub source: &'static str,
}

impl fmt::Display for InterruptDescription {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "{} ({}, vec={}) {}",
               self.mnemonic,
               self.irqtype,
               self.vector,
               self.description)
    }
}

/// x86 External Interrupts (1-16).
pub static EXCEPTIONS: [InterruptDescription; 15] = [InterruptDescription {
                                                         vector: 0,
                                                         mnemonic: "#DE",
                                                         description: "Divide Error",
                                                         irqtype: "Fault",
                                                         source: "DIV and IDIV instructions.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 1,
                                                         mnemonic: "#DB",
                                                         description: "RESERVED",
                                                         irqtype: "Fault/ Trap",
                                                         source: "For Intel use only.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 2,
                                                         mnemonic: "NMI",
                                                         description: "Interrupt",
                                                         irqtype: "Interrupt",
                                                         source: "Nonmaskable external interrupt.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 3,
                                                         mnemonic: "#BP",
                                                         description: "Breakpoint",
                                                         irqtype: "Trap",
                                                         source: "INT 3 instruction.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 4,
                                                         mnemonic: "#OF",
                                                         description: "Overflow",
                                                         irqtype: "Trap",
                                                         source: "INTO instruction.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 5,
                                                         mnemonic: "#BR",
                                                         description: "BOUND Range Exceeded",
                                                         irqtype: "Fault",
                                                         source: "BOUND instruction.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 6,
                                                         mnemonic: "#UD",
                                                         description: "Invalid Opcode (Undefined \
                                                                       Opcode)",
                                                         irqtype: "Fault",
                                                         source: "UD2 instruction or reserved \
                                                                  opcode.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 7,
                                                         mnemonic: "#NM",
                                                         description: "Device Not Available (No \
                                                                       Math Coprocessor)",
                                                         irqtype: "Fault",
                                                         source: "Floating-point or WAIT/FWAIT \
                                                                  instruction.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 8,
                                                         mnemonic: "#DF",
                                                         description: "Double Fault",
                                                         irqtype: "Abort",
                                                         source: "Any instruction that can \
                                                                  generate an exception, an NMI, \
                                                                  or an INTR.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 9,
                                                         mnemonic: "",
                                                         description: "Coprocessor Segment Overrun",
                                                         irqtype: "Fault",
                                                         source: "Floating-point instruction.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 10,
                                                         mnemonic: "#TS",
                                                         description: "Invalid TSS",
                                                         irqtype: "Fault",
                                                         source: "Task switch or TSS access.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 11,
                                                         mnemonic: "#NP",
                                                         description: "Segment Not Present",
                                                         irqtype: "Fault",
                                                         source: "Loading segment registers or \
                                                                  accessing system segments.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 12,
                                                         mnemonic: "#SS",
                                                         description: "Stack-Segment Fault",
                                                         irqtype: "Fault",
                                                         source: "Stack operations and SS register \
                                                                  loads.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 13,
                                                         mnemonic: "#GP",
                                                         description: "General Protection",
                                                         irqtype: "Fault",
                                                         source: "Any memory reference and other \
                                                                  protection checks.",
                                                     },
                                                     InterruptDescription {
                                                         vector: 14,
                                                         mnemonic: "#PF",
                                                         description: "Page Fault",
                                                         irqtype: "Fault",
                                                         source: "Any memory reference.",
                                                     }];

/// Enable Interrupts.
pub unsafe fn enable() {
    asm!("sti");
}

/// Disable Interrupts.
pub unsafe fn disable() {
    asm!("cli");
}

/// Generate a software interrupt.
/// This is a macro argument needs to be an immediate.
#[macro_export]
macro_rules! int {
    ( $x:expr ) => {
        {
            asm!("int $0" :: "N" ($x));
        }
    };
}

/// A struct describing an interrupt gate.  See the Intel manual mentioned
/// above for details, specifically, the section "6.14.1 64-Bit Mode IDT"
/// and "Table 3-2. System-Segment and Gate-Descriptor Types".
#[derive(Debug, Clone, Copy)]
#[repr(C, packed)]
pub struct IdtEntry {
    /// Lower 16 bits of ISR.
    pub base_lo: u16,
    /// Segment selector.
    pub sel: u16,
    /// This must always be zero.
    pub res0: u8,
    /// Flags.
    pub flags: u8,
    /// The upper 48 bits of ISR (the last 16 bits must be zero).
    pub base_hi: u64,
    /// Must be zero.
    pub res1: u16,
}

impl IdtEntry {
    /// Create a "missing" IdtEntry.  This is a `const` function, so we can
    /// call it at compile time to initialize static variables.
    ///
    /// If the CPU tries to invoke a missing interrupt, it will instead
    /// send a General Protection fault (13), with the interrupt number and
    /// some other data stored in the error code.
    pub const fn missing() -> IdtEntry {
        IdtEntry {
            base_lo: 0,
            sel: 0,
            res0: 0,
            flags: 0,
            base_hi: 0,
            res1: 0,
        }
    }

    /// Create a new IdtEntry pointing at `handler`, which must be a
    /// function with interrupt calling conventions.  (This must be
    /// currently defined in assembly language.)  The `gdt_code_selector`
    /// value must be the offset of code segment entry in the GDT.
    ///
    /// Create an interrupt gate with the "Present" flag set, which is the
    /// most common case.  If you need something else, you can construct it
    /// manually.
    pub const fn interrupt_gate(gdt_code_selector: u16, handler: VAddr) -> IdtEntry {
        IdtEntry {
            base_lo: ((handler.as_usize() as u64) & 0xFFFF) as u16,
            sel: gdt_code_selector,
            res0: 0,
            // Bit 7: "Present" flag set.
            // Bits 0-4: This is an interrupt gate.
            flags: 0b1000_1110,
            base_hi: handler.as_usize() as u64 >> 16,
            res1: 0,
        }
    }
}

bitflags!{
    // Taken from Intel Manual Section 4.7 Page-Fault Exceptions.
    flags PageFaultError: u32 {
        /// 0: The fault was caused by a non-present page.
        /// 1: The fault was caused by a page-level protection violation
        const PFAULT_ERROR_P = bit!(0),

        /// 0: The access causing the fault was a read.
        /// 1: The access causing the fault was a write.
        const PFAULT_ERROR_WR = bit!(1),

        /// 0: The access causing the fault originated when the processor
        /// was executing in supervisor mode.
        /// 1: The access causing the fault originated when the processor
        /// was executing in user mode.
        const PFAULT_ERROR_US = bit!(2),

        /// 0: The fault was not caused by reserved bit violation.
        /// 1: The fault was caused by reserved bits set to 1 in a page directory.
        const PFAULT_ERROR_RSVD = bit!(3),

        /// 0: The fault was not caused by an instruction fetch.
        /// 1: The fault was caused by an instruction fetch.
        const PFAULT_ERROR_ID = bit!(4),

        /// 0: The fault was not by protection keys.
        /// 1: There was a protection key violation.
        const PFAULT_ERROR_PK = bit!(5),
    }
}

impl fmt::Debug for PageFaultError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let p = match self.contains(PFAULT_ERROR_P) {
            false => "The fault was caused by a non-present page.",
            true => "The fault was caused by a page-level protection violation.",
        };
        let wr = match self.contains(PFAULT_ERROR_WR) {
            false => "The access causing the fault was a read.",
            true => "The access causing the fault was a write.",
        };
        let us = match self.contains(PFAULT_ERROR_US) {
            false => {
                "The access causing the fault originated when the processor was executing in \
                 supervisor mode."
            }
            true => {
                "The access causing the fault originated when the processor was executing in user \
                 mode."
            }
        };
        let rsvd = match self.contains(PFAULT_ERROR_RSVD) {
            false => "The fault was not caused by reserved bit violation.",
            true => "The fault was caused by reserved bits set to 1 in a page directory.",
        };
        let id = match self.contains(PFAULT_ERROR_ID) {
            false => "The fault was not caused by an instruction fetch.",
            true => "The fault was caused by an instruction fetch.",
        };

        write!(f, "{}\n{}\n{}\n{}\n{}", p, wr, us, rsvd, id)
    }
}

#[test]
fn bit_macro() {
    assert!(PFAULT_ERROR_PK.bits() == 0b100000);
    assert!(PFAULT_ERROR_ID.bits() == 0b10000);
    assert!(PFAULT_ERROR_RSVD.bits() == 0b1000);
    assert!(PFAULT_ERROR_US.bits() == 0b100);
    assert!(PFAULT_ERROR_WR.bits() == 0b10);
    assert!(PFAULT_ERROR_P.bits() == 0b1);
}
<|endoftext|>


<commit_before>//! API to generate `.rs` files.
//!
//! This API does not require `protoc` command present in `$PATH`.
//!
//! ```
//! extern crate protoc_rust;
//!
//! fn main() {
//!     protobuf_codegen_pure::Codegen::new()
//!         .out_dir("src/protos")
//!         .inputs(&["protos/a.proto", "protos/b.proto"])
//!         .include("protos")
//!         .run()
//!         .expect("Codegen failed.");
//! }
//! ```
//!
//! It is advisable that `protobuf-codegen-pure` build-dependecy version be the same as
//! `protobuf` dependency.
//!
//! The alternative is to use `protoc-rust` crate.

#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]

extern crate protobuf;
extern crate protobuf_codegen;

use std::fmt;
use std::fs;
use std::io;
use std::path::Path;
use std::path::PathBuf;
use std::process;
use std::str;

mod linked_hash_map;
mod model;
mod parse_dependencies;
mod parser;
mod path;
mod rel_path;

use linked_hash_map::LinkedHashMap;
use protobuf::descriptor::FileDescriptorProto;
pub use protobuf_codegen::Customize;
use protobuf_codegen::ProtoPath;
use protobuf_codegen::ProtoPathBuf;

use crate::parser::ParserErrorWithLocation;

mod convert;
#[cfg(test)]
mod test_against_protobuf_protos;

/// Invoke pure rust codegen. See [crate docs](crate) for example.
// TODO: merge with protoc-rust def
#[derive(Debug, Default)]
pub struct Codegen {
    /// --lang_out= param
    out_dir: PathBuf,
    /// -I args
    includes: Vec<PathBuf>,
    /// List of .proto files to compile
    inputs: Vec<PathBuf>,
    /// Customize code generation
    customize: Customize,
}

impl Codegen {
    /// Fresh new codegen object.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the output directory for codegen.
    pub fn out_dir(&mut self, out_dir: impl AsRef<Path>) -> &mut Self {
        self.out_dir = out_dir.as_ref().to_owned();
        self
    }

    /// Add an include directory.
    pub fn include(&mut self, include: impl AsRef<Path>) -> &mut Self {
        self.includes.push(include.as_ref().to_owned());
        self
    }

    /// Add include directories.
    pub fn includes(&mut self, includes: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for include in includes {
            self.include(include);
        }
        self
    }

    /// Add an input (`.proto` file).
    pub fn input(&mut self, input: impl AsRef<Path>) -> &mut Self {
        self.inputs.push(input.as_ref().to_owned());
        self
    }

    /// Add inputs (`.proto` files).
    pub fn inputs(&mut self, inputs: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for input in inputs {
            self.input(input);
        }
        self
    }

    /// Specify generated code [`Customize`] object.
    pub fn customize(&mut self, customize: Customize) -> &mut Self {
        self.customize = customize;
        self
    }

    /// Like `protoc --rust_out=...` but without requiring `protoc` or `protoc-gen-rust`
    /// commands in `$PATH`.
    pub fn run(&self) -> anyhow::Result<()> {
        let p = parse_and_typecheck(&self.includes, &self.inputs)?;

        protobuf_codegen::gen_and_write(
            &p.file_descriptors,
            &format!("protobuf-codegen-pure={}", env!("CARGO_PKG_VERSION")),
            &p.relative_paths,
            &self.out_dir,
            &self.customize,
        )
    }

    /// Similar to `run`, but prints the message to stderr and exits the process on error.
    pub fn run_from_script(&self) {
        if let Err(e) = self.run() {
            eprintln!("pure rust codegen failed: {}", e);
            process::exit(1);
        }
    }
}

#[derive(Clone)]
pub(crate) struct FileDescriptorPair {
    parsed: model::FileDescriptor,
    descriptor: protobuf::descriptor::FileDescriptorProto,
}

#[derive(Debug, thiserror::Error)]
#[error("error in `{file}`: {error}")]
struct WithFileError {
    file: String,
    #[source]
    error: anyhow::Error,
}

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("protobuf path `{0}` is not found in import path {1}")]
    FileNotFoundInImportPath(String, String),
    #[error("file `{0}` must reside in include path {1}")]
    FileMustResideInImportPath(String, String),
    #[error("could not read file `{0}`: {1}")]
    CouldNotReadFile(String, io::Error),
    #[error("file `{0}` content is not UTF-8")]
    FileContentIsNotUtf8(String),
}

/// Resolve `.proto` files. `Display` is used for error messages.
trait ProtoPathResolver: fmt::Display {
    fn resolve(&self, path: &ProtoPath) -> anyhow::Result<Option<ResolvedProtoFile>>;
}

struct Run<R>
where
    R: ProtoPathResolver,
{
    parsed_files: LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    resolver: R,
}

impl<R> Run<R>
where
    R: ProtoPathResolver,
{
    fn get_file_and_all_deps_already_parsed(
        &self,
        protobuf_path: &ProtoPath,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        if let Some(_) = result.get(protobuf_path) {
            return;
        }

        let pair = self
            .parsed_files
            .get(protobuf_path)
            .expect("must be already parsed");
        result.insert(protobuf_path.to_proto_path_buf(), pair.clone());

        self.get_all_deps_already_parsed(&pair.parsed, result);
    }

    fn get_all_deps_already_parsed(
        &self,
        parsed: &model::FileDescriptor,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        for import in &parsed.imports {
            self.get_file_and_all_deps_already_parsed(&import.path, result);
        }
    }

    fn add_file_content(
        &mut self,
        protobuf_path: &ProtoPath,
        resolved: &ResolvedProtoFile,
    ) -> anyhow::Result<()> {
        let content = str::from_utf8(&resolved.content)
            .map_err(|_| Error::FileContentIsNotUtf8(protobuf_path.to_string()))?;

        let parsed = model::FileDescriptor::parse(&content).map_err(|e| WithFileError {
            file: resolved.path.clone(),
            error: e.into(),
        })?;

        for import in &parsed.imports {
            self.add_imported_file(&import.path)?;
        }

        let mut this_file_deps = LinkedHashMap::new();
        self.get_all_deps_already_parsed(&parsed, &mut this_file_deps);

        let this_file_deps: Vec<_> = this_file_deps.into_iter().map(|(_, v)| v).collect();

        let descriptor = convert::file_descriptor(protobuf_path, &parsed, &this_file_deps)
            .map_err(|e| WithFileError {
                file: resolved.path.clone(),
                error: e.into(),
            })?;

        self.parsed_files.insert(
            protobuf_path.to_proto_path_buf(),
            FileDescriptorPair { parsed, descriptor },
        );

        Ok(())
    }

    fn add_imported_file(&mut self, protobuf_path: &ProtoPath) -> anyhow::Result<()> {
        if let Some(_) = self.parsed_files.get(protobuf_path) {
            return Ok(());
        }

        let resolved = self.resolver.resolve(protobuf_path)?;
        if let Some(resolved) = resolved {
            return self.add_file_content(protobuf_path, &resolved);
        }

        let embedded = match protobuf_path.to_str() {
            "rustproto.proto" => Some(RUSTPROTO_PROTO),
            "google/protobuf/any.proto" => Some(ANY_PROTO),
            "google/protobuf/api.proto" => Some(API_PROTO),
            "google/protobuf/descriptor.proto" => Some(DESCRIPTOR_PROTO),
            "google/protobuf/duration.proto" => Some(DURATION_PROTO),
            "google/protobuf/empty.proto" => Some(EMPTY_PROTO),
            "google/protobuf/field_mask.proto" => Some(FIELD_MASK_PROTO),
            "google/protobuf/source_context.proto" => Some(SOURCE_CONTEXT_PROTO),
            "google/protobuf/struct.proto" => Some(STRUCT_PROTO),
            "google/protobuf/timestamp.proto" => Some(TIMESTAMP_PROTO),
            "google/protobuf/type.proto" => Some(TYPE_PROTO),
            "google/protobuf/wrappers.proto" => Some(WRAPPERS_PROTO),
            _ => None,
        };

        match embedded {
            Some(content) => self.add_file_content(
                protobuf_path,
                &ResolvedProtoFile {
                    path: protobuf_path.to_string(),
                    content: content.as_bytes().to_vec(),
                },
            ),
            None => Err(Error::FileNotFoundInImportPath(
                protobuf_path.to_string(),
                format!("{}", self.resolver),
            )
            .into()),
        }
    }
}

/// Result of parsing `.proto` files.
#[doc(hidden)]
pub struct ParsedAndTypechecked {
    /// One entry for each input `.proto` file.
    pub relative_paths: Vec<ProtoPathBuf>,
    /// All parsed `.proto` files including dependencies of input files.
    pub file_descriptors: Vec<protobuf::descriptor::FileDescriptorProto>,
}

fn path_to_proto_path(path: &Path, includes: &[PathBuf]) -> anyhow::Result<ProtoPathBuf> {
    for include in includes {
        if include == Path::new(".") && path.is_relative() {
            // Special handling of `.` to allow using `.` as an include path
            // and `foo.proto` as input.
            return ProtoPathBuf::from_path(path);
        }
        match path.strip_prefix(include) {
            Ok(stripped) => return ProtoPathBuf::from_path(stripped),
            Err(_) => continue,
        }
    }
    Err(
        Error::FileMustResideInImportPath(path.display().to_string(), format!("{:?}", includes))
            .into(),
    )
}

struct ResolvedProtoFile {
    /// For error reporting.
    pub path: String,
    /// File content.
    pub content: Vec<u8>,
}

fn fs_resolver(includes: &[PathBuf]) -> impl ProtoPathResolver {
    struct Impl {
        includes: Vec<PathBuf>,
    }

    impl fmt::Display for Impl {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "{:?}", self.includes)
        }
    }

    impl ProtoPathResolver for Impl {
        fn resolve(&self, proto_path: &ProtoPath) -> anyhow::Result<Option<ResolvedProtoFile>> {
            for include_dir in &self.includes {
                let fs_path = include_dir.join(proto_path.to_path());
                match fs::read_to_string(&fs_path) {
                    Ok(content) => {
                        return Ok(Some(ResolvedProtoFile {
                            path: fs_path.display().to_string(),
                            content: content.into_bytes(),
                        }))
                    }
                    Err(e) if e.kind() == io::ErrorKind::NotFound => continue,
                    Err(e) => {
                        return Err(Error::CouldNotReadFile(fs_path.display().to_string(), e).into())
                    }
                }
            }
            Ok(None)
        }
    }

    Impl {
        includes: includes.to_vec(),
    }
}

#[doc(hidden)]
pub fn parse_and_typecheck(
    includes: &[PathBuf],
    input: &[PathBuf],
) -> anyhow::Result<ParsedAndTypechecked> {
    let mut run = Run {
        parsed_files: LinkedHashMap::new(),
        resolver: fs_resolver(includes),
    };

    let relative_paths = input
        .iter()
        .map(|input| Ok((path_to_proto_path(input, includes)?, input)))
        .collect::<anyhow::Result<Vec<_>>>()?;

    for (proto_path, path) in &relative_paths {
        let content = fs::read_to_string(path)
            .map_err(|e| Error::CouldNotReadFile(path.display().to_string(), e))?;
        run.add_file_content(
            proto_path,
            &ResolvedProtoFile {
                path: path.display().to_string(),
                content: content.into_bytes(),
            },
        )?;
    }

    let file_descriptors: Vec<_> = run
        .parsed_files
        .into_iter()
        .map(|(_, v)| v.descriptor)
        .collect();

    Ok(ParsedAndTypechecked {
        relative_paths: relative_paths.into_iter().map(|(p, _)| p).collect(),
        file_descriptors,
    })
}

/// Parse imports from a `.proto` file.
///
/// The result is [`FileDescriptorProto`] object with only `*dependency` fields filled.
#[doc(hidden)]
pub fn parse_dependencies(content: &str) -> Result<FileDescriptorProto, ParserErrorWithLocation> {
    parse_dependencies::parse_dependencies(content)
}

// TODO: these include don't work when publishing to crates
const RUSTPROTO_PROTO: &str = include_str!("proto/rustproto.proto");
const ANY_PROTO: &str = include_str!("proto/google/protobuf/any.proto");
const API_PROTO: &str = include_str!("proto/google/protobuf/api.proto");
const DESCRIPTOR_PROTO: &str = include_str!("proto/google/protobuf/descriptor.proto");
const DURATION_PROTO: &str = include_str!("proto/google/protobuf/duration.proto");
const EMPTY_PROTO: &str = include_str!("proto/google/protobuf/empty.proto");
const FIELD_MASK_PROTO: &str = include_str!("proto/google/protobuf/field_mask.proto");
const SOURCE_CONTEXT_PROTO: &str = include_str!("proto/google/protobuf/source_context.proto");
const STRUCT_PROTO: &str = include_str!("proto/google/protobuf/struct.proto");
const TIMESTAMP_PROTO: &str = include_str!("proto/google/protobuf/timestamp.proto");
const TYPE_PROTO: &str = include_str!("proto/google/protobuf/type.proto");
const WRAPPERS_PROTO: &str = include_str!("proto/google/protobuf/wrappers.proto");
<commit_msg>parse_and_typecheck with custom resolver<commit_after>//! API to generate `.rs` files.
//!
//! This API does not require `protoc` command present in `$PATH`.
//!
//! ```
//! extern crate protoc_rust;
//!
//! fn main() {
//!     protobuf_codegen_pure::Codegen::new()
//!         .out_dir("src/protos")
//!         .inputs(&["protos/a.proto", "protos/b.proto"])
//!         .include("protos")
//!         .run()
//!         .expect("Codegen failed.");
//! }
//! ```
//!
//! It is advisable that `protobuf-codegen-pure` build-dependecy version be the same as
//! `protobuf` dependency.
//!
//! The alternative is to use `protoc-rust` crate.

#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]

extern crate protobuf;
extern crate protobuf_codegen;

use std::fmt;
use std::fs;
use std::io;
use std::path::Path;
use std::path::PathBuf;
use std::process;
use std::str;

mod linked_hash_map;
mod model;
mod parse_dependencies;
mod parser;
mod path;
mod rel_path;

use linked_hash_map::LinkedHashMap;
use protobuf::descriptor::FileDescriptorProto;
pub use protobuf_codegen::Customize;
use protobuf_codegen::ProtoPath;
use protobuf_codegen::ProtoPathBuf;

use crate::parser::ParserErrorWithLocation;

mod convert;
#[cfg(test)]
mod test_against_protobuf_protos;

/// Invoke pure rust codegen. See [crate docs](crate) for example.
// TODO: merge with protoc-rust def
#[derive(Debug, Default)]
pub struct Codegen {
    /// --lang_out= param
    out_dir: PathBuf,
    /// -I args
    includes: Vec<PathBuf>,
    /// List of .proto files to compile
    inputs: Vec<PathBuf>,
    /// Customize code generation
    customize: Customize,
}

impl Codegen {
    /// Fresh new codegen object.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the output directory for codegen.
    pub fn out_dir(&mut self, out_dir: impl AsRef<Path>) -> &mut Self {
        self.out_dir = out_dir.as_ref().to_owned();
        self
    }

    /// Add an include directory.
    pub fn include(&mut self, include: impl AsRef<Path>) -> &mut Self {
        self.includes.push(include.as_ref().to_owned());
        self
    }

    /// Add include directories.
    pub fn includes(&mut self, includes: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for include in includes {
            self.include(include);
        }
        self
    }

    /// Add an input (`.proto` file).
    pub fn input(&mut self, input: impl AsRef<Path>) -> &mut Self {
        self.inputs.push(input.as_ref().to_owned());
        self
    }

    /// Add inputs (`.proto` files).
    pub fn inputs(&mut self, inputs: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for input in inputs {
            self.input(input);
        }
        self
    }

    /// Specify generated code [`Customize`] object.
    pub fn customize(&mut self, customize: Customize) -> &mut Self {
        self.customize = customize;
        self
    }

    /// Like `protoc --rust_out=...` but without requiring `protoc` or `protoc-gen-rust`
    /// commands in `$PATH`.
    pub fn run(&self) -> anyhow::Result<()> {
        let p = parse_and_typecheck(&self.includes, &self.inputs)?;

        protobuf_codegen::gen_and_write(
            &p.file_descriptors,
            &format!("protobuf-codegen-pure={}", env!("CARGO_PKG_VERSION")),
            &p.relative_paths,
            &self.out_dir,
            &self.customize,
        )
    }

    /// Similar to `run`, but prints the message to stderr and exits the process on error.
    pub fn run_from_script(&self) {
        if let Err(e) = self.run() {
            eprintln!("pure rust codegen failed: {}", e);
            process::exit(1);
        }
    }
}

#[derive(Clone)]
pub(crate) struct FileDescriptorPair {
    parsed: model::FileDescriptor,
    descriptor: protobuf::descriptor::FileDescriptorProto,
}

#[derive(Debug, thiserror::Error)]
#[error("error in `{file}`: {error}")]
struct WithFileError {
    file: String,
    #[source]
    error: anyhow::Error,
}

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("protobuf path `{0}` is not found in import path {1}")]
    FileNotFoundInImportPath(String, String),
    #[error("file `{0}` must reside in include path {1}")]
    FileMustResideInImportPath(String, String),
    #[error("could not read file `{0}`: {1}")]
    CouldNotReadFile(String, io::Error),
    #[error("file `{0}` content is not UTF-8")]
    FileContentIsNotUtf8(String),
}

/// Resolve `.proto` files. `Display` is used for error messages.
pub trait ProtoPathResolver: fmt::Display {
    /// Resolve a `.proto` file.
    ///
    /// Return `None` if a path is unknown, and if a path is a built-in protobuf file,
    /// like `google/protobuf/descriptor.proto`, it will be handled by the library.
    fn resolve(&self, path: &ProtoPath) -> anyhow::Result<Option<ResolvedProtoFile>>;
}

struct Run<R>
where
    R: ProtoPathResolver,
{
    parsed_files: LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    resolver: R,
}

impl<R> Run<R>
where
    R: ProtoPathResolver,
{
    fn get_file_and_all_deps_already_parsed(
        &self,
        protobuf_path: &ProtoPath,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        if let Some(_) = result.get(protobuf_path) {
            return;
        }

        let pair = self
            .parsed_files
            .get(protobuf_path)
            .expect("must be already parsed");
        result.insert(protobuf_path.to_proto_path_buf(), pair.clone());

        self.get_all_deps_already_parsed(&pair.parsed, result);
    }

    fn get_all_deps_already_parsed(
        &self,
        parsed: &model::FileDescriptor,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        for import in &parsed.imports {
            self.get_file_and_all_deps_already_parsed(&import.path, result);
        }
    }

    fn add_file_content(
        &mut self,
        protobuf_path: &ProtoPath,
        resolved: &ResolvedProtoFile,
    ) -> anyhow::Result<()> {
        let content = str::from_utf8(&resolved.content)
            .map_err(|_| Error::FileContentIsNotUtf8(protobuf_path.to_string()))?;

        let parsed = model::FileDescriptor::parse(&content).map_err(|e| WithFileError {
            file: resolved.path.clone(),
            error: e.into(),
        })?;

        for import in &parsed.imports {
            self.add_imported_file(&import.path)?;
        }

        let mut this_file_deps = LinkedHashMap::new();
        self.get_all_deps_already_parsed(&parsed, &mut this_file_deps);

        let this_file_deps: Vec<_> = this_file_deps.into_iter().map(|(_, v)| v).collect();

        let descriptor = convert::file_descriptor(protobuf_path, &parsed, &this_file_deps)
            .map_err(|e| WithFileError {
                file: resolved.path.clone(),
                error: e.into(),
            })?;

        self.parsed_files.insert(
            protobuf_path.to_proto_path_buf(),
            FileDescriptorPair { parsed, descriptor },
        );

        Ok(())
    }

    fn add_imported_file(&mut self, protobuf_path: &ProtoPath) -> anyhow::Result<()> {
        if let Some(_) = self.parsed_files.get(protobuf_path) {
            return Ok(());
        }

        let resolved = self.resolver.resolve(protobuf_path)?;
        if let Some(resolved) = resolved {
            return self.add_file_content(protobuf_path, &resolved);
        }

        let embedded = match protobuf_path.to_str() {
            "rustproto.proto" => Some(RUSTPROTO_PROTO),
            "google/protobuf/any.proto" => Some(ANY_PROTO),
            "google/protobuf/api.proto" => Some(API_PROTO),
            "google/protobuf/descriptor.proto" => Some(DESCRIPTOR_PROTO),
            "google/protobuf/duration.proto" => Some(DURATION_PROTO),
            "google/protobuf/empty.proto" => Some(EMPTY_PROTO),
            "google/protobuf/field_mask.proto" => Some(FIELD_MASK_PROTO),
            "google/protobuf/source_context.proto" => Some(SOURCE_CONTEXT_PROTO),
            "google/protobuf/struct.proto" => Some(STRUCT_PROTO),
            "google/protobuf/timestamp.proto" => Some(TIMESTAMP_PROTO),
            "google/protobuf/type.proto" => Some(TYPE_PROTO),
            "google/protobuf/wrappers.proto" => Some(WRAPPERS_PROTO),
            _ => None,
        };

        match embedded {
            Some(content) => self.add_file_content(
                protobuf_path,
                &ResolvedProtoFile {
                    path: protobuf_path.to_string(),
                    content: content.as_bytes().to_vec(),
                },
            ),
            None => Err(Error::FileNotFoundInImportPath(
                protobuf_path.to_string(),
                format!("{}", self.resolver),
            )
            .into()),
        }
    }
}

/// Result of parsing `.proto` files.
#[doc(hidden)]
pub struct ParsedAndTypechecked {
    /// One entry for each input `.proto` file.
    pub relative_paths: Vec<ProtoPathBuf>,
    /// All parsed `.proto` files including dependencies of input files.
    pub file_descriptors: Vec<protobuf::descriptor::FileDescriptorProto>,
}

fn path_to_proto_path(path: &Path, includes: &[PathBuf]) -> anyhow::Result<ProtoPathBuf> {
    for include in includes {
        if include == Path::new(".") && path.is_relative() {
            // Special handling of `.` to allow using `.` as an include path
            // and `foo.proto` as input.
            return ProtoPathBuf::from_path(path);
        }
        match path.strip_prefix(include) {
            Ok(stripped) => return ProtoPathBuf::from_path(stripped),
            Err(_) => continue,
        }
    }
    Err(
        Error::FileMustResideInImportPath(path.display().to_string(), format!("{:?}", includes))
            .into(),
    )
}

/// `.proto` file result provided from the [`ProtoPathResolver`].
pub struct ResolvedProtoFile {
    /// For error reporting.
    pub path: String,
    /// File content.
    pub content: Vec<u8>,
}

fn fs_resolver(includes: &[PathBuf]) -> impl ProtoPathResolver {
    struct Impl {
        includes: Vec<PathBuf>,
    }

    impl fmt::Display for Impl {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "{:?}", self.includes)
        }
    }

    impl ProtoPathResolver for Impl {
        fn resolve(&self, proto_path: &ProtoPath) -> anyhow::Result<Option<ResolvedProtoFile>> {
            for include_dir in &self.includes {
                let fs_path = include_dir.join(proto_path.to_path());
                match fs::read_to_string(&fs_path) {
                    Ok(content) => {
                        return Ok(Some(ResolvedProtoFile {
                            path: fs_path.display().to_string(),
                            content: content.into_bytes(),
                        }))
                    }
                    Err(e) if e.kind() == io::ErrorKind::NotFound => continue,
                    Err(e) => {
                        return Err(Error::CouldNotReadFile(fs_path.display().to_string(), e).into())
                    }
                }
            }
            Ok(None)
        }
    }

    Impl {
        includes: includes.to_vec(),
    }
}

#[doc(hidden)]
pub fn parse_and_typecheck(
    includes: &[PathBuf],
    input: &[PathBuf],
) -> anyhow::Result<ParsedAndTypechecked> {
    let mut run = Run {
        parsed_files: LinkedHashMap::new(),
        resolver: fs_resolver(includes),
    };

    let relative_paths = input
        .iter()
        .map(|input| Ok((path_to_proto_path(input, includes)?, input)))
        .collect::<anyhow::Result<Vec<_>>>()?;

    for (proto_path, path) in &relative_paths {
        let content = fs::read_to_string(path)
            .map_err(|e| Error::CouldNotReadFile(path.display().to_string(), e))?;
        run.add_file_content(
            proto_path,
            &ResolvedProtoFile {
                path: path.display().to_string(),
                content: content.into_bytes(),
            },
        )?;
    }

    let file_descriptors: Vec<_> = run
        .parsed_files
        .into_iter()
        .map(|(_, v)| v.descriptor)
        .collect();

    Ok(ParsedAndTypechecked {
        relative_paths: relative_paths.into_iter().map(|(p, _)| p).collect(),
        file_descriptors,
    })
}

#[doc(hidden)]
pub fn parse_and_typecheck_custom(
    input: &[ProtoPathBuf],
    resolver: impl ProtoPathResolver,
) -> anyhow::Result<Vec<FileDescriptorProto>> {
    let mut run = Run {
        parsed_files: LinkedHashMap::new(),
        resolver,
    };

    for proto_path in input {
        run.add_imported_file(proto_path)?;
    }

    Ok(run
        .parsed_files
        .into_iter()
        .map(|(_, v)| v.descriptor)
        .collect())
}

/// Parse imports from a `.proto` file.
///
/// The result is [`FileDescriptorProto`] object with only `*dependency` fields filled.
#[doc(hidden)]
pub fn parse_dependencies(content: &str) -> Result<FileDescriptorProto, ParserErrorWithLocation> {
    parse_dependencies::parse_dependencies(content)
}

// TODO: these include don't work when publishing to crates
const RUSTPROTO_PROTO: &str = include_str!("proto/rustproto.proto");
const ANY_PROTO: &str = include_str!("proto/google/protobuf/any.proto");
const API_PROTO: &str = include_str!("proto/google/protobuf/api.proto");
const DESCRIPTOR_PROTO: &str = include_str!("proto/google/protobuf/descriptor.proto");
const DURATION_PROTO: &str = include_str!("proto/google/protobuf/duration.proto");
const EMPTY_PROTO: &str = include_str!("proto/google/protobuf/empty.proto");
const FIELD_MASK_PROTO: &str = include_str!("proto/google/protobuf/field_mask.proto");
const SOURCE_CONTEXT_PROTO: &str = include_str!("proto/google/protobuf/source_context.proto");
const STRUCT_PROTO: &str = include_str!("proto/google/protobuf/struct.proto");
const TIMESTAMP_PROTO: &str = include_str!("proto/google/protobuf/timestamp.proto");
const TYPE_PROTO: &str = include_str!("proto/google/protobuf/type.proto");
const WRAPPERS_PROTO: &str = include_str!("proto/google/protobuf/wrappers.proto");
<|endoftext|>


<commit_before>use collections::slice::{mod, AsSlice, SlicePrelude};
use collections::vec::{mod, Vec};
use core::mem;
use core::iter::{mod, AdditiveIterator, Iterator};
use core::option::{mod, Some, None, Option};

use iobuf::Iobuf;

/// A span over potentially many Iobufs. This is useful as a "string" type where
/// the contents of the string can come from multiple IObufs, and you want to
/// avoid copying the buffer contents unnecessarily.
///
/// As an optimization, pushing an Iobuf that points to data immediately after
/// the range represented by the last Iobuf pushed will result in just expanding
/// the held Iobuf's range. This prevents allocating lots of unnecessary
/// intermediate buffers, while still maintaining the illusion of "pushing lots
/// of buffers" while incrementally parsing.
///
/// A `BufSpan` is internally represented as either an `Iobuf` or a `Vec<Iobuf>`,
/// depending on how many different buffers were used.
pub enum BufSpan<Buf> {
  Empty,
  One (Buf),
  Many(Vec<Buf>),
}

impl<Buf: Iobuf> BufSpan<Buf> {
  /// Creates a new, empty `Bufspan`.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf};
  ///
  /// let s: BufSpan<ROIobuf<'static>> = BufSpan::new();
  /// assert!(s.is_empty());
  /// ```
  #[inline]
  pub fn new() -> BufSpan<Buf> {
    Empty
  }

  /// Creates a new `BufSpan` from an Iobuf.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf};
  ///
  /// let s = BufSpan::from_buf(ROIobuf::from_slice(b"hello"));
  /// assert_eq!(s.iter().count(), 1);
  /// assert_eq!(s.count_bytes(), 5);
  /// ```
  #[inline]
  pub fn from_buf(b: Buf) -> BufSpan<Buf> {
    One(b)
  }

  /// Returns `true` iff the span is over an empty range.
  ///
  /// ```
  /// use iobuf::{BufSpan, Iobuf, ROIobuf};
  ///
  /// let mut s = BufSpan::new();
  ///
  /// assert!(s.is_empty());
  ///
  /// s.push(ROIobuf::from_str(""));
  /// assert!(s.is_empty());
  ///
  /// s.push(ROIobuf::from_str("hello, world!"));
  /// assert!(!s.is_empty());
  /// ```
  #[inline]
  pub fn is_empty(&self) -> bool {
    match *self {
      Empty => true,
      _     => false,
    }
  }

  /// The fast path during pushing -- either fills in the first buffer, or
  /// extends an existing one.
  ///
  /// Returns `None` if the fast path was taken and nothing more needs to be
  /// done. Returns `Some` if we need to do a slow push.
  fn try_to_extend(&mut self, b: Buf) -> Option<Buf> {
    let len = b.len();
    if len == 0 { return None; }

    match *self {
      Empty => {},
      One(ref mut b0) => {
        unsafe {
          if b0.is_extended_by(&b) {
            b0.unsafe_extend(len);
            return None;
          } else {
            return Some(b); // Upgrade the `One` into a `Many` in the slow path.
          }
        }
      }
      Many(ref mut v) => {
        // I wish this wouldn't unwind, and just abort... :(
        let last = v.last_mut().unwrap();
        unsafe {
          if last.is_extended_by(&b) {
            last.unsafe_extend(len);
            return None;
          } else {
            return Some(b);
          }
        }
      }
    }

    // Handle this case in the fast path, instead of leaving it for slow_push
    // to clean up.
    *self = One(b);
    None
  }

  /// Appends a buffer to a `BufSpan`. If the buffer is an extension of the
  /// previously pushed buffer, the range will be extended. Otherwise, the new
  /// non-extension buffer will be added to the end of a vector.
  pub fn push(&mut self, b: Buf) {
    match self.try_to_extend(b) {
      None    => {},
      Some(b) => self.slow_push(b),
    }
  }

  /// The slow path during a push. This is only taken if a `BufSpan` must span
  /// multiple backing buffers.
  #[cold]
  fn slow_push(&mut self, b: Buf) {
    let this = mem::replace(self, Empty);
    *self =
      match this {
        Empty   => One(b),
        One(b0) => {
          let mut v = Vec::with_capacity(2);
          v.push(b0);
          v.push(b);
          Many(v)
        },
        Many(mut bs) => { bs.push(b); Many(bs) }
      };
  }

  /// Returns an iterator over references to the buffers inside the `BufSpan`.
  #[inline]
  pub fn iter<'a>(&'a self) -> SpanIter<'a, Buf> {
    match *self {
      Empty       => Opt(None.into_iter()),
      One (ref b) => Opt(Some(b).into_iter()),
      Many(ref v) => Lot(v.as_slice().iter()),
    }
  }

  /// Returns a moving iterator over the buffers inside the `BufSpan`.
  #[inline]
  pub fn into_iter(self) -> SpanMoveIter<Buf> {
    match self {
      Empty   => MoveOpt(None.into_iter()),
      One (b) => MoveOpt(Some(b).into_iter()),
      Many(v) => MoveLot(v.into_iter()),
    }
  }

  /// Returns an iterator over the bytes in the `BufSpan`.
  #[inline]
  pub fn iter_bytes<'a>(&'a self) -> ByteIter<'a, Buf> {
    self.iter()
        .flat_map(|buf| unsafe { buf.as_window_slice().iter() })
        .map(|&b| b)
  }

  /// Returns `true` iff the bytes in this `BufSpan` are the same as the bytes
  /// in the other `BufSpan`.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf, RWIobuf};
  ///
  /// let a = BufSpan::from_buf(ROIobuf::from_str("hello"));
  /// let b = BufSpan::from_buf(RWIobuf::from_string("hello".into_string()));
  ///
  /// assert!(a.byte_equal(&b));
  ///
  /// let mut c = BufSpan::from_buf(ROIobuf::from_str("hel"));
  /// c.push(ROIobuf::from_str("lo"));
  ///
  /// assert!(a.byte_equal(&c)); assert!(c.byte_equal(&a));
  ///
  /// let d = BufSpan::from_buf(ROIobuf::from_str("helo"));
  /// assert!(!a.byte_equal(&d));
  /// ```
  #[inline]
  pub fn byte_equal<Buf2: Iobuf>(&self, other: &BufSpan<Buf2>) -> bool {
    self.count_bytes() == other.count_bytes()
    && self.iter_bytes().zip(other.iter_bytes()).all(|(a, b)| a == b)
  }

  /// A more efficient version of byte_equal, specialized to work exclusively on
  /// slices.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf};
  ///
  /// let a = BufSpan::from_buf(ROIobuf::from_str("hello"));
  ///
  /// assert!(a.byte_equal_slice(b"hello"));
  /// assert!(!a.byte_equal_slice(b"helo"));
  /// ```
  #[inline]
  pub fn byte_equal_slice(&self, other: &[u8]) -> bool {
    self.count_bytes() as uint == other.len()
    && self.iter_bytes().zip(other.iter()).all(|(a, &b)| a == b)
  }

  /// Counts the number of bytes this `BufSpan` is over. This is
  /// `O(self.iter().len())`.
  #[inline]
  pub fn count_bytes(&self) -> u32 {
    // `self.iter().map(|b| b.len()).sum()` would be shorter, but I like to
    // specialize for the much more common case of empty or singular `BufSpan`s.
    match *self {
      Empty       => 0,
      One (ref b) => b.len(),
      Many(ref v) => v.iter().map(|b| b.len()).sum(),
    }
  }
}

/// An iterator over the bytes in a `BufSpan`.
pub type ByteIter<'a, Buf> =
  iter::Map<'static, &'a u8, u8,
    iter::FlatMap<'static, &'a Buf,
      SpanIter<'a, Buf>,
      slice::Items<'a, u8>>>;

/// An iterator over references to buffers inside a `BufSpan`.
pub enum SpanIter<'a, Buf: 'a> {
  Opt(option::Item<&'a Buf>),
  Lot(slice::Items<'a, Buf>),
}

impl<'a, Buf: Iobuf> Iterator<&'a Buf> for SpanIter<'a, Buf> {
  #[inline(always)]
  fn next(&mut self) -> Option<&'a Buf> {
    // I'm couting on this match getting lifted out of the loop with
    // loop-invariant code motion.
    match *self {
      Opt(ref mut iter) => iter.next(),
      Lot(ref mut iter) => iter.next(),
    }
  }

  #[inline(always)]
  fn size_hint(&self) -> (uint, Option<uint>) {
    match *self {
      Opt(ref iter) => iter.size_hint(),
      Lot(ref iter) => iter.size_hint(),
    }
  }
}

/// A moving iterator over buffers inside a `BufSpan`.
pub enum SpanMoveIter<Buf> {
  MoveOpt(option::Item<Buf>),
  MoveLot(vec::MoveItems<Buf>),
}

impl<Buf: Iobuf> Iterator<Buf> for SpanMoveIter<Buf> {
  #[inline(always)]
  fn next(&mut self) -> Option<Buf> {
    // I'm couting on this match getting lifted out of the loop with
    // loop-invariant code motion.
    match *self {
      MoveOpt(ref mut iter) => iter.next(),
      MoveLot(ref mut iter) => iter.next(),
    }
  }

  #[inline(always)]
  fn size_hint(&self) -> (uint, Option<uint>) {
    match *self {
      MoveOpt(ref iter) => iter.size_hint(),
      MoveLot(ref iter) => iter.size_hint(),
    }
  }
}
<commit_msg>impls for BufSpan<commit_after>use collections::slice::{mod, AsSlice, SlicePrelude};
use collections::vec::{mod, Vec};
use core::clone::Clone;
use core::fmt;
use core::mem;
use core::iter::{mod, AdditiveIterator, Iterator};
use core::option::{mod, Some, None, Option};
use core::result::{Ok, Err};

use iobuf::Iobuf;

/// A span over potentially many Iobufs. This is useful as a "string" type where
/// the contents of the string can come from multiple IObufs, and you want to
/// avoid copying the buffer contents unnecessarily.
///
/// As an optimization, pushing an Iobuf that points to data immediately after
/// the range represented by the last Iobuf pushed will result in just expanding
/// the held Iobuf's range. This prevents allocating lots of unnecessary
/// intermediate buffers, while still maintaining the illusion of "pushing lots
/// of buffers" while incrementally parsing.
///
/// A `BufSpan` is internally represented as either an `Iobuf` or a `Vec<Iobuf>`,
/// depending on how many different buffers were used.
pub enum BufSpan<Buf> {
  Empty,
  One (Buf),
  Many(Vec<Buf>),
}

impl<Buf: Iobuf> Clone for BufSpan<Buf> {
  #[inline]
  fn clone(&self) -> BufSpan<Buf> {
    match *self {
      Empty       => Empty,
      One(ref b)  => One ((*b).clone()),
      Many(ref v) => Many((*v).clone()),
    }
  }
}

impl<Buf: Iobuf> fmt::Show for BufSpan<Buf> {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    let mut first_time = true;

    for b in self.iter() {
      if !first_time {
        try!(write!(f, "\n"));
      }

      try!(b.fmt(f));

      first_time = false;
    }

    Ok(())
  }
}

impl<Buf: Iobuf> BufSpan<Buf> {
  /// Creates a new, empty `Bufspan`.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf};
  ///
  /// let s: BufSpan<ROIobuf<'static>> = BufSpan::new();
  /// assert!(s.is_empty());
  /// ```
  #[inline]
  pub fn new() -> BufSpan<Buf> {
    Empty
  }

  /// Creates a new `BufSpan` from an Iobuf.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf};
  ///
  /// let s = BufSpan::from_buf(ROIobuf::from_slice(b"hello"));
  /// assert_eq!(s.iter().count(), 1);
  /// assert_eq!(s.count_bytes(), 5);
  /// ```
  #[inline]
  pub fn from_buf(b: Buf) -> BufSpan<Buf> {
    One(b)
  }

  /// Returns `true` iff the span is over an empty range.
  ///
  /// ```
  /// use iobuf::{BufSpan, Iobuf, ROIobuf};
  ///
  /// let mut s = BufSpan::new();
  ///
  /// assert!(s.is_empty());
  ///
  /// s.push(ROIobuf::from_str(""));
  /// assert!(s.is_empty());
  ///
  /// s.push(ROIobuf::from_str("hello, world!"));
  /// assert!(!s.is_empty());
  /// ```
  #[inline]
  pub fn is_empty(&self) -> bool {
    match *self {
      Empty => true,
      _     => false,
    }
  }

  /// The fast path during pushing -- either fills in the first buffer, or
  /// extends an existing one.
  ///
  /// Returns `None` if the fast path was taken and nothing more needs to be
  /// done. Returns `Some` if we need to do a slow push.
  fn try_to_extend(&mut self, b: Buf) -> Option<Buf> {
    let len = b.len();
    if len == 0 { return None; }

    match *self {
      Empty => {},
      One(ref mut b0) => {
        unsafe {
          if b0.is_extended_by(&b) {
            b0.unsafe_extend(len);
            return None;
          } else {
            return Some(b); // Upgrade the `One` into a `Many` in the slow path.
          }
        }
      }
      Many(ref mut v) => {
        // I wish this wouldn't unwind, and just abort... :(
        let last = v.last_mut().unwrap();
        unsafe {
          if last.is_extended_by(&b) {
            last.unsafe_extend(len);
            return None;
          } else {
            return Some(b);
          }
        }
      }
    }

    // Handle this case in the fast path, instead of leaving it for slow_push
    // to clean up.
    *self = One(b);
    None
  }

  /// Appends a buffer to a `BufSpan`. If the buffer is an extension of the
  /// previously pushed buffer, the range will be extended. Otherwise, the new
  /// non-extension buffer will be added to the end of a vector.
  ///
  /// ```
  /// use iobuf::{BufSpan, Iobuf, ROIobuf};
  ///
  /// let mut s = BufSpan::new();
  ///
  /// s.push(ROIobuf::from_str("he"));
  /// s.push(ROIobuf::from_str("llo"));
  ///
  /// assert_eq!(s.count_bytes() as uint, "hello".len());
  /// assert_eq!(s.iter().count(), 2);
  ///
  /// let mut b0 = ROIobuf::from_str(" world");
  /// let mut b1 = b0.clone();
  ///
  /// b0.resize(2).unwrap();
  /// b1.advance(2).unwrap();
  ///
  /// s.push(b0);
  /// s.push(b1);
  ///
  /// // b0 and b1 are immediately after each other, and from the same buffer,
  /// // so get merged into one Iobuf.
  /// assert_eq!(s.count_bytes() as uint, "hello world".len());
  /// assert_eq!(s.iter().count(), 3);
  /// ```
  pub fn push(&mut self, b: Buf) {
    match self.try_to_extend(b) {
      None    => {},
      Some(b) => self.slow_push(b),
    }
  }

  /// The slow path during a push. This is only taken if a `BufSpan` must span
  /// multiple backing buffers.
  #[cold]
  fn slow_push(&mut self, b: Buf) {
    let this = mem::replace(self, Empty);
    *self =
      match this {
        Empty   => One(b),
        One(b0) => {
          let mut v = Vec::with_capacity(2);
          v.push(b0);
          v.push(b);
          Many(v)
        },
        Many(mut bs) => { bs.push(b); Many(bs) }
      };
  }

  /// Returns an iterator over references to the buffers inside the `BufSpan`.
  #[inline]
  pub fn iter<'a>(&'a self) -> SpanIter<'a, Buf> {
    match *self {
      Empty       => Opt(None.into_iter()),
      One (ref b) => Opt(Some(b).into_iter()),
      Many(ref v) => Lot(v.as_slice().iter()),
    }
  }

  /// Returns a moving iterator over the buffers inside the `BufSpan`.
  #[inline]
  pub fn into_iter(self) -> SpanMoveIter<Buf> {
    match self {
      Empty   => MoveOpt(None.into_iter()),
      One (b) => MoveOpt(Some(b).into_iter()),
      Many(v) => MoveLot(v.into_iter()),
    }
  }

  /// Returns an iterator over the bytes in the `BufSpan`.
  #[inline]
  pub fn iter_bytes<'a>(&'a self) -> ByteIter<'a, Buf> {
    self.iter()
        .flat_map(|buf| unsafe { buf.as_window_slice().iter() })
        .map(|&b| b)
  }

  /// Returns `true` iff the bytes in this `BufSpan` are the same as the bytes
  /// in the other `BufSpan`.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf, RWIobuf};
  ///
  /// let a = BufSpan::from_buf(ROIobuf::from_str("hello"));
  /// let b = BufSpan::from_buf(RWIobuf::from_string("hello".into_string()));
  ///
  /// assert!(a.byte_equal(&b));
  ///
  /// let mut c = BufSpan::from_buf(ROIobuf::from_str("hel"));
  /// c.push(ROIobuf::from_str("lo"));
  ///
  /// assert!(a.byte_equal(&c)); assert!(c.byte_equal(&a));
  ///
  /// let d = BufSpan::from_buf(ROIobuf::from_str("helo"));
  /// assert!(!a.byte_equal(&d));
  /// ```
  #[inline]
  pub fn byte_equal<Buf2: Iobuf>(&self, other: &BufSpan<Buf2>) -> bool {
    self.count_bytes() == other.count_bytes()
    && self.iter_bytes().zip(other.iter_bytes()).all(|(a, b)| a == b)
  }

  /// A more efficient version of byte_equal, specialized to work exclusively on
  /// slices.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf};
  ///
  /// let a = BufSpan::from_buf(ROIobuf::from_str("hello"));
  ///
  /// assert!(a.byte_equal_slice(b"hello"));
  /// assert!(!a.byte_equal_slice(b"helo"));
  /// ```
  #[inline]
  pub fn byte_equal_slice(&self, other: &[u8]) -> bool {
    self.count_bytes() as uint == other.len()
    && self.iter_bytes().zip(other.iter()).all(|(a, &b)| a == b)
  }

  /// Counts the number of bytes this `BufSpan` is over. This is
  /// `O(self.iter().len())`.
  #[inline]
  pub fn count_bytes(&self) -> u32 {
    // `self.iter().map(|b| b.len()).sum()` would be shorter, but I like to
    // specialize for the much more common case of empty or singular `BufSpan`s.
    match *self {
      Empty       => 0,
      One (ref b) => b.len(),
      Many(ref v) => v.iter().map(|b| b.len()).sum(),
    }
  }
}

/// An iterator over the bytes in a `BufSpan`.
pub type ByteIter<'a, Buf> =
  iter::Map<'static, &'a u8, u8,
    iter::FlatMap<'static, &'a Buf,
      SpanIter<'a, Buf>,
      slice::Items<'a, u8>>>;

/// An iterator over references to buffers inside a `BufSpan`.
pub enum SpanIter<'a, Buf: 'a> {
  Opt(option::Item<&'a Buf>),
  Lot(slice::Items<'a, Buf>),
}

impl<'a, Buf: Iobuf> Iterator<&'a Buf> for SpanIter<'a, Buf> {
  #[inline(always)]
  fn next(&mut self) -> Option<&'a Buf> {
    // I'm couting on this match getting lifted out of the loop with
    // loop-invariant code motion.
    match *self {
      Opt(ref mut iter) => iter.next(),
      Lot(ref mut iter) => iter.next(),
    }
  }

  #[inline(always)]
  fn size_hint(&self) -> (uint, Option<uint>) {
    match *self {
      Opt(ref iter) => iter.size_hint(),
      Lot(ref iter) => iter.size_hint(),
    }
  }
}

/// A moving iterator over buffers inside a `BufSpan`.
pub enum SpanMoveIter<Buf> {
  MoveOpt(option::Item<Buf>),
  MoveLot(vec::MoveItems<Buf>),
}

impl<Buf: Iobuf> Iterator<Buf> for SpanMoveIter<Buf> {
  #[inline(always)]
  fn next(&mut self) -> Option<Buf> {
    // I'm couting on this match getting lifted out of the loop with
    // loop-invariant code motion.
    match *self {
      MoveOpt(ref mut iter) => iter.next(),
      MoveLot(ref mut iter) => iter.next(),
    }
  }

  #[inline(always)]
  fn size_hint(&self) -> (uint, Option<uint>) {
    match *self {
      MoveOpt(ref iter) => iter.size_hint(),
      MoveLot(ref iter) => iter.size_hint(),
    }
  }
}
<|endoftext|>


<commit_before>#![feature(plugin_registrar, rustc_private)]

extern crate rustc;
extern crate rustc_plugin;
extern crate syntax;

use rustc_plugin::Registry;
use syntax::tokenstream::TokenTree;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, MacResult};
use syntax::parse::token::{self, DelimToken, Token};

mod parser_any_macro;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("interpolate_idents", interpolate_idents);
}

fn interpolate_idents<'a>(cx: &'a mut ExtCtxt,
              _: Span,
              tts: &[TokenTree]) -> Box<MacResult + 'a> {
    fn concat_idents(tts: &Vec<TokenTree>, delim: DelimToken) -> Option<TokenTree> {
        match delim {
            DelimToken::Bracket => {
                let mut new_ident = String::new();
                let mut new_span: Option<Span> = None;

                for token in tts.iter() {
                    match token {
                        &TokenTree::Token(ref span, Token::Ident(ref ident)) => {
                            match new_span {
                                Some(ref mut s) => { s.hi = span.hi; },
                                None => { new_span = Some(span.clone()); },
                            }
                            new_ident.push_str(&ident.name.as_str());
                        },
                        _ => return None,
                    }
                }

                match new_span {
                    Some(s) => {
                        let new_ident = token::str_to_ident(&new_ident[..]);
                        Some(TokenTree::Token(s, Token::Ident(new_ident)))
                    },
                    None => None
                }
            },
            _ => None,
        }
    }

    fn map_tts(tts: &[TokenTree]) -> Vec<TokenTree> {
        tts.iter().map(|t| {
            match t {
                &TokenTree::Delimited(ref s, ref d) => {
                    match concat_idents(&d.tts, d.delim) {
                        Some(t) => t,
                        None => {
                            TokenTree::Delimited(s.clone(), syntax::tokenstream::Delimited {
                                delim: d.delim,
                                open_span: d.open_span,
                                tts: map_tts(&*d.tts),
                                close_span: d.close_span,
                            })
                        },
                    }
                },
                &TokenTree::Sequence(ref s, ref d) => {
                    TokenTree::Sequence(s.clone(), syntax::tokenstream::SequenceRepetition {
                        tts: map_tts(&*d.tts),
                        separator: d.separator.clone(),
                        op: d.op,
                        num_captures: d.num_captures,
                    })
                },
                _ => t.clone(),
            }
        }).collect()
    }

    let parser = cx.new_parser_from_tts(&map_tts(tts)[..]);
    Box::new(parser_any_macro::ParserAnyMacro::new(parser)) as Box<MacResult>
}
<commit_msg>Fixed for 2016-07-09 nightly<commit_after>#![feature(plugin_registrar, rustc_private)]

extern crate rustc;
extern crate rustc_plugin;
extern crate syntax;

use rustc_plugin::Registry;
use std::rc::Rc;
use syntax::tokenstream::TokenTree;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, MacResult};
use syntax::parse::token::{self, DelimToken, Token};

mod parser_any_macro;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("interpolate_idents", interpolate_idents);
}

fn interpolate_idents<'a>(cx: &'a mut ExtCtxt,
              _: Span,
              tts: &[TokenTree]) -> Box<MacResult + 'a> {
    fn concat_idents(tts: &Vec<TokenTree>, delim: DelimToken) -> Option<TokenTree> {
        match delim {
            DelimToken::Bracket => {
                let mut new_ident = String::new();
                let mut new_span: Option<Span> = None;

                for token in tts.iter() {
                    match token {
                        &TokenTree::Token(ref span, Token::Ident(ref ident)) => {
                            match new_span {
                                Some(ref mut s) => { s.hi = span.hi; },
                                None => { new_span = Some(span.clone()); },
                            }
                            new_ident.push_str(&ident.name.as_str());
                        },
                        _ => return None,
                    }
                }

                match new_span {
                    Some(s) => {
                        let new_ident = token::str_to_ident(&new_ident[..]);
                        Some(TokenTree::Token(s, Token::Ident(new_ident)))
                    },
                    None => None
                }
            },
            _ => None,
        }
    }

    fn map_tts(tts: &[TokenTree]) -> Vec<TokenTree> {
        tts.iter().map(|t| {
            match t {
                &TokenTree::Delimited(ref s, ref d) => {
                    match concat_idents(&d.tts, d.delim) {
                        Some(t) => t,
                        None => {
                            TokenTree::Delimited(s.clone(), Rc::new(syntax::tokenstream::Delimited {
                                delim: d.delim,
                                open_span: d.open_span,
                                tts: map_tts(&*d.tts),
                                close_span: d.close_span,
                            }))
                        },
                    }
                },
                &TokenTree::Sequence(ref s, ref d) => {
                    TokenTree::Sequence(s.clone(), Rc::new(syntax::tokenstream::SequenceRepetition {
                        tts: map_tts(&*d.tts),
                        separator: d.separator.clone(),
                        op: d.op,
                        num_captures: d.num_captures,
                    }))
                },
                _ => t.clone(),
            }
        }).collect()
    }

    let parser = cx.new_parser_from_tts(&map_tts(tts)[..]);
    Box::new(parser_any_macro::ParserAnyMacro::new(parser)) as Box<MacResult>
}
<|endoftext|>