<commit_before>//! Client-side Proc-Macro crate
//!
//! We separate proc-macro expanding logic to an extern program to allow
//! different implementations (e.g. wasm or dylib loading). And this crate
//! is used to provide basic infrastructure for communication between two
//! processes: Client (RA itself), Server (the external program)

pub mod msg;
mod process;
mod rpc;

use base_db::{Env, ProcMacro};
use std::{
    ffi::OsStr,
    fs::File,
    io::{self, Read},
    path::{Path, PathBuf},
    sync::Arc,
};

use tt::{SmolStr, Subtree};

use crate::process::{ProcMacroProcessSrv, ProcMacroProcessThread};

pub use rpc::{ExpansionResult, ExpansionTask, ListMacrosResult, ListMacrosTask, ProcMacroKind};

use memmap::Mmap;
use object::read::{File as BinaryFile, Object, ObjectSection};
use snap::read::FrameDecoder as SnapDecoder;

#[derive(Debug, Clone)]
struct ProcMacroProcessExpander {
    process: Arc<ProcMacroProcessSrv>,
    dylib_path: PathBuf,
    name: SmolStr,
}

impl Eq for ProcMacroProcessExpander {}
impl PartialEq for ProcMacroProcessExpander {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name
            && self.dylib_path == other.dylib_path
            && Arc::ptr_eq(&self.process, &other.process)
    }
}

impl base_db::ProcMacroExpander for ProcMacroProcessExpander {
    fn expand(
        &self,
        subtree: &Subtree,
        attr: Option<&Subtree>,
        env: &Env,
    ) -> Result<Subtree, tt::ExpansionError> {
        let task = ExpansionTask {
            macro_body: subtree.clone(),
            macro_name: self.name.to_string(),
            attributes: attr.cloned(),
            lib: self.dylib_path.to_path_buf(),
            env: env.iter().map(|(k, v)| (k.to_string(), v.to_string())).collect(),
        };

        let result: ExpansionResult = self.process.send_task(msg::Request::ExpansionMacro(task))?;
        Ok(result.expansion)
    }
}

#[derive(Debug)]
pub struct ProcMacroClient {
    process: Arc<ProcMacroProcessSrv>,
    thread: ProcMacroProcessThread,
}

impl ProcMacroClient {
    pub fn extern_process(
        process_path: PathBuf,
        args: impl IntoIterator<Item = impl AsRef<OsStr>>,
    ) -> io::Result<ProcMacroClient> {
        let (thread, process) = ProcMacroProcessSrv::run(process_path, args)?;
        Ok(ProcMacroClient {
            process: Arc::new(process),
            thread,
        })
    }

    pub fn by_dylib_path(&self, dylib_path: &Path) -> Vec<ProcMacro> {
        let macros = match self.process.find_proc_macros(dylib_path) {
            Err(err) => {
                eprintln!("Failed to find proc macros. Error: {:#?}", err);
                return vec![];
            }
            Ok(macros) => macros,
        };

        macros
            .into_iter()
            .map(|(name, kind)| {
                let name = SmolStr::new(&name);
                let kind = match kind {
                    ProcMacroKind::CustomDerive => base_db::ProcMacroKind::CustomDerive,
                    ProcMacroKind::FuncLike => base_db::ProcMacroKind::FuncLike,
                    ProcMacroKind::Attr => base_db::ProcMacroKind::Attr,
                };
                let expander = Arc::new(ProcMacroProcessExpander {
                    process: self.process.clone(),
                    name: name.clone(),
                    dylib_path: dylib_path.into(),
                });

                ProcMacro {
                    name,
                    kind,
                    expander,
                }
            })
            .collect()
    }

    // This is used inside self.read_version() to locate the ".rustc" section
    // from a proc macro crate's binary file.
    fn read_section<'a>(&self, dylib_binary: &'a [u8], section_name: &str) -> io::Result<&'a [u8]> {
        BinaryFile::parse(dylib_binary)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
            .section_by_name(section_name)
            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "section read error"))?
            .data()
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
    }

    // Check the version of rustc that was used to compile a proc macro crate's
    // binary file.
    // A proc macro crate binary's ".rustc" section has following byte layout:
    // * [b'r',b'u',b's',b't',0,0,0,5] is the first 8 bytes
    // * ff060000 734e6150 is followed, it's the snappy format magic bytes,
    //   means bytes from here(including this sequence) are compressed in
    //   snappy compression format. Version info is inside here, so decompress
    //   this.
    // The bytes you get after decompressing the snappy format portion has
    // following layout:
    // * [b'r',b'u',b's',b't',0,0,0,5] is the first 8 bytes(again)
    // * [crate root bytes] next 4 bytes is to store crate root position,
    //   according to rustc's source code comment
    // * [length byte] next 1 byte tells us how many bytes we should read next
    //   for the version string's utf8 bytes
    // * [version string bytes encoded in utf8] <- GET THIS BOI
    // * [some more bytes that we don really care but still there] :-)
    // Check this issue for more about the bytes layout:
    // https://github.com/rust-analyzer/rust-analyzer/issues/6174
    #[allow(unused)]
    fn read_version(&self, dylib_path: &Path) -> io::Result<String> {
        let dylib_file = File::open(dylib_path)?;
        let dylib_mmaped = unsafe { Mmap::map(&dylib_file) }?;

        let dot_rustc = self.read_section(&dylib_mmaped, ".rustc")?;

        let header = &dot_rustc[..8];
        const EXPECTED_HEADER: [u8; 8] = [b'r', b'u', b's', b't', 0, 0, 0, 5];
        // check if header is valid
        if header != EXPECTED_HEADER {
            return Err(io::Error::new(io::ErrorKind::InvalidData, format!(".rustc section should start with header {:?}; header {:?} is actually presented.",EXPECTED_HEADER ,header)));
        }

        let snappy_portion = &dot_rustc[8..];

        let mut snappy_decoder = SnapDecoder::new(snappy_portion);

        // the bytes before version string bytes, so this basically is:
        // 8 bytes for [b'r',b'u',b's',b't',0,0,0,5]
        // 4 bytes for [crate root bytes]
        // 1 byte for length of version string
        // so 13 bytes in total, and we should check the 13th byte
        // to know the length
        let mut bytes_before_version = [0u8; 13];
        snappy_decoder.read_exact(&mut bytes_before_version)?;
        let length = bytes_before_version[12]; // what? can't use -1 indexing?

        let mut version_string_utf8 = vec![0u8; length as usize];
        snappy_decoder.read_exact(&mut version_string_utf8)?;
        let version_string = String::from_utf8(version_string_utf8);
        version_string.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
    }
}
<commit_msg>Revise error message regarding metadata version<commit_after>//! Client-side Proc-Macro crate
//!
//! We separate proc-macro expanding logic to an extern program to allow
//! different implementations (e.g. wasm or dylib loading). And this crate
//! is used to provide basic infrastructure for communication between two
//! processes: Client (RA itself), Server (the external program)

pub mod msg;
mod process;
mod rpc;

use base_db::{Env, ProcMacro};
use std::{
    ffi::OsStr,
    fs::File,
    io::{self, Read},
    path::{Path, PathBuf},
    sync::Arc,
};

use tt::{SmolStr, Subtree};

use crate::process::{ProcMacroProcessSrv, ProcMacroProcessThread};

pub use rpc::{ExpansionResult, ExpansionTask, ListMacrosResult, ListMacrosTask, ProcMacroKind};

use memmap::Mmap;
use object::read::{File as BinaryFile, Object, ObjectSection};
use snap::read::FrameDecoder as SnapDecoder;

#[derive(Debug, Clone)]
struct ProcMacroProcessExpander {
    process: Arc<ProcMacroProcessSrv>,
    dylib_path: PathBuf,
    name: SmolStr,
}

impl Eq for ProcMacroProcessExpander {}
impl PartialEq for ProcMacroProcessExpander {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name
            && self.dylib_path == other.dylib_path
            && Arc::ptr_eq(&self.process, &other.process)
    }
}

impl base_db::ProcMacroExpander for ProcMacroProcessExpander {
    fn expand(
        &self,
        subtree: &Subtree,
        attr: Option<&Subtree>,
        env: &Env,
    ) -> Result<Subtree, tt::ExpansionError> {
        let task = ExpansionTask {
            macro_body: subtree.clone(),
            macro_name: self.name.to_string(),
            attributes: attr.cloned(),
            lib: self.dylib_path.to_path_buf(),
            env: env.iter().map(|(k, v)| (k.to_string(), v.to_string())).collect(),
        };

        let result: ExpansionResult = self.process.send_task(msg::Request::ExpansionMacro(task))?;
        Ok(result.expansion)
    }
}

#[derive(Debug)]
pub struct ProcMacroClient {
    process: Arc<ProcMacroProcessSrv>,
    thread: ProcMacroProcessThread,
}

impl ProcMacroClient {
    pub fn extern_process(
        process_path: PathBuf,
        args: impl IntoIterator<Item = impl AsRef<OsStr>>,
    ) -> io::Result<ProcMacroClient> {
        let (thread, process) = ProcMacroProcessSrv::run(process_path, args)?;
        Ok(ProcMacroClient { process: Arc::new(process), thread })
    }

    pub fn by_dylib_path(&self, dylib_path: &Path) -> Vec<ProcMacro> {
        let macros = match self.process.find_proc_macros(dylib_path) {
            Err(err) => {
                eprintln!("Failed to find proc macros. Error: {:#?}", err);
                return vec![];
            }
            Ok(macros) => macros,
        };

        macros
            .into_iter()
            .map(|(name, kind)| {
                let name = SmolStr::new(&name);
                let kind = match kind {
                    ProcMacroKind::CustomDerive => base_db::ProcMacroKind::CustomDerive,
                    ProcMacroKind::FuncLike => base_db::ProcMacroKind::FuncLike,
                    ProcMacroKind::Attr => base_db::ProcMacroKind::Attr,
                };
                let expander = Arc::new(ProcMacroProcessExpander {
                    process: self.process.clone(),
                    name: name.clone(),
                    dylib_path: dylib_path.into(),
                });

                ProcMacro { name, kind, expander }
            })
            .collect()
    }

    // This is used inside self.read_version() to locate the ".rustc" section
    // from a proc macro crate's binary file.
    fn read_section<'a>(&self, dylib_binary: &'a [u8], section_name: &str) -> io::Result<&'a [u8]> {
        BinaryFile::parse(dylib_binary)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
            .section_by_name(section_name)
            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "section read error"))?
            .data()
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
    }

    // Check the version of rustc that was used to compile a proc macro crate's
    // binary file.
    // A proc macro crate binary's ".rustc" section has following byte layout:
    // * [b'r',b'u',b's',b't',0,0,0,5] is the first 8 bytes
    // * ff060000 734e6150 is followed, it's the snappy format magic bytes,
    //   means bytes from here(including this sequence) are compressed in
    //   snappy compression format. Version info is inside here, so decompress
    //   this.
    // The bytes you get after decompressing the snappy format portion has
    // following layout:
    // * [b'r',b'u',b's',b't',0,0,0,5] is the first 8 bytes(again)
    // * [crate root bytes] next 4 bytes is to store crate root position,
    //   according to rustc's source code comment
    // * [length byte] next 1 byte tells us how many bytes we should read next
    //   for the version string's utf8 bytes
    // * [version string bytes encoded in utf8] <- GET THIS BOI
    // * [some more bytes that we don really care but still there] :-)
    // Check this issue for more about the bytes layout:
    // https://github.com/rust-analyzer/rust-analyzer/issues/6174
    #[allow(unused)]
    fn read_version(&self, dylib_path: &Path) -> io::Result<String> {
        let dylib_file = File::open(dylib_path)?;
        let dylib_mmaped = unsafe { Mmap::map(&dylib_file) }?;

        let dot_rustc = self.read_section(&dylib_mmaped, ".rustc")?;

        let header = &dot_rustc[..8];
        const EXPECTED_HEADER: [u8; 8] = [b'r', b'u', b's', b't', 0, 0, 0, 5];
        // check if header is valid
        if header != EXPECTED_HEADER {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("only metadata version 5 is supported, section header was: {:?}", header),
            ));
        }

        let snappy_portion = &dot_rustc[8..];

        let mut snappy_decoder = SnapDecoder::new(snappy_portion);

        // the bytes before version string bytes, so this basically is:
        // 8 bytes for [b'r',b'u',b's',b't',0,0,0,5]
        // 4 bytes for [crate root bytes]
        // 1 byte for length of version string
        // so 13 bytes in total, and we should check the 13th byte
        // to know the length
        let mut bytes_before_version = [0u8; 13];
        snappy_decoder.read_exact(&mut bytes_before_version)?;
        let length = bytes_before_version[12]; // what? can't use -1 indexing?

        let mut version_string_utf8 = vec![0u8; length as usize];
        snappy_decoder.read_exact(&mut version_string_utf8)?;
        let version_string = String::from_utf8(version_string_utf8);
        version_string.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
    }
}
<|endoftext|>


<commit_before>use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::io::{Read, Write};

/// Represents an opcode in the MongoDB Wire Protocol.
#[derive(Clone)]
pub enum OpCode {
    Reply = 1,
    Update = 2001,
    Insert = 2002,
    Query = 2004,
    GetMore = 2005,
}

impl OpCode {
    /// Maps integer values to OpCodes
    ///
    /// # Arguments
    ///
    /// `i` - The integer to map.
    ///
    /// # Return value
    ///
    /// Returns the matching opcode, or `None` if the integer isn't a valid
    /// opcode.
    pub fn from_i32(i: i32) -> Option<OpCode> {
        match i {
            1 => Some(OpCode::Reply),
            2001 => Some(OpCode::Update),
            2002 => Some(OpCode::Insert),
            2004 => Some(OpCode::Query),
            2005 => Some(OpCode::GetMore),
            _ => None
        }
    }
}

impl ToString for OpCode {
    fn to_string(&self) -> String {
        match self {
            &OpCode::Reply => "OP_REPLY".to_owned(),
            &OpCode::Update => "OP_UPDATE".to_owned(),
            &OpCode::Insert => "OP_INSERT".to_owned(),
            &OpCode::Query => "OP_QUERY".to_owned(),
            &OpCode::GetMore => "OP_GET_MORE".to_owned(),
        }
    }
}

/// Represents a header in the MongoDB Wire Protocol.
///
/// # Fields
///
/// `message_length` - The length of the entire message in bytes.
/// `request_id` - Identifies the request being sent. This should be `0` in a
///                response from the server.
/// `response_to` - Identifies which response the message is a response to. This
///                 should be `0` in a request from the client.
/// `op_code`     - Identifies which type of message is being sent.
// #[derive(Clone)]
pub struct Header {
    pub message_length: i32,
    pub request_id: i32,
    response_to: i32,
    pub op_code: OpCode,
}

impl Header {
    /// Constructs a new Header.
    ///
    /// # Arguments
    ///
    /// `message_length` - The length of the message in bytes.
    /// `request_id` - Identifier for the request, or `0` if the the message
    ///                is a response.
    /// `response_to` - Identifies which request the message is in response to,
    ///                or `0` if the the message is a request.
    /// `op_code`    - Identifies which type of message is being sent.
    ///
    /// # Return value
    ///
    /// Returns the newly-created Header.
    pub fn new(message_length: i32, request_id: i32, response_to: i32,
           op_code: OpCode) -> Header {
        Header { message_length: message_length, request_id: request_id,
                 response_to: response_to, op_code: op_code }
    }

    /// Construcs a new Header for a request.
    ///
    /// # Arguments
    ///
    /// `message_length` - The length of the message in bytes.
    /// `request_id` - Identifier for the request, or `0` if the the message
    ///                is a response.
    /// `op_code`    - Identifies which type of message is being sent.
    ///
    /// # Return value
    ///
    /// Returns a new Header with `response_to` set to 0.
    fn new_request(message_length: i32, request_id: i32,
                   op_code: OpCode) -> Header {
        Header::new(message_length, request_id, 0, op_code)
    }

    fn new_reply(message_length: i32, response_to: i32,
                 op_code: OpCode) -> Header {
        Header::new(message_length, 0, response_to, op_code)
    }

    pub fn with_update(message_length: i32, request_id: i32) -> Header {
        Header::new_request(message_length, request_id, OpCode::Update)
    }

    pub fn with_insert(message_length: i32, request_id: i32) -> Header {
        Header::new_request(message_length, request_id, OpCode::Insert)
    }

    /// Constructs a new Header for a query.
    ///
    /// # Arguments
    ///
    /// `message_length` - The length of the message in bytes.
    /// `request_id` - Identifier for the request, or `0` if the the message
    ///                is a response.
    /// # Return value
    ///
    /// Returns a new Header with `response_to` set to 0 and `op_code`
    /// set to `OpQuery`.
    pub fn with_query(message_length: i32, request_id: i32) -> Header {
        Header::new_request(message_length, request_id, OpCode::Query)
    }

    pub fn with_get_more(message_length: i32, request_id: i32) -> Header {
        Header::new_request(message_length, request_id, OpCode::GetMore)
    }

    /// Writes the serialized Header to a buffer.
    ///
    /// # Arguments
    ///
    /// `buffer` - The buffer to write to.
    ///
    /// # Return value
    ///
    /// Returns nothing on success, or an error string on failure.
    pub fn write(&self, buffer: &mut Write) -> Result<(), String> {
        if buffer.write_i32::<LittleEndian>(self.message_length).is_err() {
            return Err("Unable to write message_length".to_owned())
        }

        if buffer.write_i32::<LittleEndian>(self.request_id).is_err() {
            return Err("Unable to write request_id".to_owned())
        }

        if buffer.write_i32::<LittleEndian>(self.response_to).is_err() {
            return Err("Unable to write response_to".to_owned())
        }

        if buffer.write_i32::<LittleEndian>(self.op_code.clone() as i32).is_err() {
            return Err("Unable to write op_code".to_owned())
        }

        let _ = buffer.flush();

        Ok(())
    }

    /// Reads a serialized Header from a buffer.
    ///
    /// # Arguments
    ///
    /// `buffer` - The buffer to read from.
    ///
    /// # Return value
    ///
    /// Returns the parsed Header on success, or an error string on failure.
    pub fn read(buffer: &mut Read) -> Result<Header, String> {
        let message_length = match buffer.read_i32::<LittleEndian>() {
            Ok(i) => i,
            Err(_) => return Err("unable to read length".to_owned())
        };

        let request_id = match buffer.read_i32::<LittleEndian>() {
            Ok(i) => i,
            Err(_) => return Err("unable to read request_id".to_owned())
        };

        let response_to = match buffer.read_i32::<LittleEndian>() {
            Ok(i) => i,
            Err(_) => return Err("unable to read response_to".to_owned())
        };

        let op_code = match buffer.read_i32::<LittleEndian>() {
            Ok(i) => match OpCode::from_i32(i) {
                Some(op) => op,
                None => return Err(format!("invalid opcode: {}", i))
            },
            Err(_) => return Err("unable to read op_code".to_owned())
        };

        Ok(Header::new(message_length, request_id, response_to, op_code))
    }
}
<commit_msg>added documentation to the header module<commit_after>use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::io::{Read, Write};

/// Represents an opcode in the MongoDB Wire Protocol.
#[derive(Clone)]
pub enum OpCode {
    Reply = 1,
    Update = 2001,
    Insert = 2002,
    Query = 2004,
    GetMore = 2005,
}

impl OpCode {
    /// Maps integer values to OpCodes
    ///
    /// # Arguments
    ///
    /// `i` - The integer to map.
    ///
    /// # Return value
    ///
    /// Returns the matching opcode, or `None` if the integer isn't a valid
    /// opcode.
    pub fn from_i32(i: i32) -> Option<OpCode> {
        match i {
            1 => Some(OpCode::Reply),
            2001 => Some(OpCode::Update),
            2002 => Some(OpCode::Insert),
            2004 => Some(OpCode::Query),
            2005 => Some(OpCode::GetMore),
            _ => None
        }
    }
}

impl ToString for OpCode {
    /// Gets the string representation of an opcode.
    ///
    /// # Return value
    ///
    /// Returns the string represetnation of the opcode.
    fn to_string(&self) -> String {
        match self {
            &OpCode::Reply => "OP_REPLY".to_owned(),
            &OpCode::Update => "OP_UPDATE".to_owned(),
            &OpCode::Insert => "OP_INSERT".to_owned(),
            &OpCode::Query => "OP_QUERY".to_owned(),
            &OpCode::GetMore => "OP_GET_MORE".to_owned(),
        }
    }
}

/// Represents a header in the MongoDB Wire Protocol.
///
/// # Fields
///
/// `message_length` - The length of the entire message in bytes.
/// `request_id` - Identifies the request being sent. This should be `0` in a
///                response from the server.
/// `response_to` - Identifies which response the message is a response to. This
///                 should be `0` in a request from the client.
/// `op_code`     - Identifies which type of message is being sent.
// #[derive(Clone)]
pub struct Header {
    pub message_length: i32,
    pub request_id: i32,
    response_to: i32,
    pub op_code: OpCode,
}

impl Header {
    /// Constructs a new Header.
    ///
    /// # Arguments
    ///
    /// `message_length` - The length of the message in bytes.
    /// `request_id` - Identifier for the request, or `0` if the the message
    ///                is a response.
    /// `response_to` - Identifies which request the message is in response to,
    ///                or `0` if the the message is a request.
    /// `op_code`    - Identifies which type of message is being sent.
    ///
    /// # Return value
    ///
    /// Returns the newly-created Header.
    pub fn new(message_length: i32, request_id: i32, response_to: i32,
           op_code: OpCode) -> Header {
        Header { message_length: message_length, request_id: request_id,
                 response_to: response_to, op_code: op_code }
    }

    /// Construcs a new Header for a request.
    ///
    /// # Arguments
    ///
    /// `message_length` - The length of the message in bytes.
    /// `request_id` - Identifier for the request, or `0` if the the message
    ///                is a response.
    /// `op_code`    - Identifies which type of message is being sent.
    ///
    /// # Return value
    ///
    /// Returns a new Header with `response_to` set to 0.
    fn new_request(message_length: i32, request_id: i32,
                   op_code: OpCode) -> Header {
        Header::new(message_length, request_id, 0, op_code)
    }

    /// Construcs a new Header for a reply.
    ///
    /// # Arguments
    ///
    /// `message_length` - The length of the message in bytes.
    /// `request_id` - Identifier for the request, or `0` if the the message
    ///                is a response.
    /// `op_code`    - Identifies which type of message is being sent.
    ///
    /// # Return value
    ///
    /// Returns a new Header with `request_id` set to 0.
    fn new_reply(message_length: i32, response_to: i32,
                 op_code: OpCode) -> Header {
        Header::new(message_length, 0, response_to, op_code)
    }

    /// Constructs a new Header for an OP_UPDATE.
    ///
    /// # Arguments
    ///
    /// `message_length` - The length of the message in bytes.
    /// `request_id` - Identifier for the request, or `0` if the the message
    ///                is a response.
    /// # Return value
    ///
    /// Returns a new Header with `response_to` set to 0 and `op_code`
    /// set to `Update`.
    pub fn with_update(message_length: i32, request_id: i32) -> Header {
        Header::new_request(message_length, request_id, OpCode::Update)
    }

    /// Constructs a new Header for an OP_INSERT.
    ///
    /// # Arguments
    ///
    /// `message_length` - The length of the message in bytes.
    /// `request_id` - Identifier for the request, or `0` if the the message
    ///                is a response.
    /// # Return value
    ///
    /// Returns a new Header with `response_to` set to 0 and `op_code`
    /// set to `Insert`.
    pub fn with_insert(message_length: i32, request_id: i32) -> Header {
        Header::new_request(message_length, request_id, OpCode::Insert)
    }

    /// Constructs a new Header for an OP_QUERY.
    ///
    /// # Arguments
    ///
    /// `message_length` - The length of the message in bytes.
    /// `request_id` - Identifier for the request, or `0` if the the message
    ///                is a response.
    /// # Return value
    ///
    /// Returns a new Header with `response_to` set to 0 and `op_code`
    /// set to `Query`.
    pub fn with_query(message_length: i32, request_id: i32) -> Header {
        Header::new_request(message_length, request_id, OpCode::Query)
    }

    /// Constructs a new Header for an OP_GET_MORE.
    ///
    /// # Arguments
    ///
    /// `message_length` - The length of the message in bytes.
    /// `request_id` - Identifier for the request, or `0` if the the message
    ///                is a response.
    /// # Return value
    ///
    /// Returns a new Header with `response_to` set to 0 and `op_code`
    /// set to `GetMore`.
    pub fn with_get_more(message_length: i32, request_id: i32) -> Header {
        Header::new_request(message_length, request_id, OpCode::GetMore)
    }

    /// Writes the serialized Header to a buffer.
    ///
    /// # Arguments
    ///
    /// `buffer` - The buffer to write to.
    ///
    /// # Return value
    ///
    /// Returns nothing on success, or an error string on failure.
    pub fn write(&self, buffer: &mut Write) -> Result<(), String> {
        if buffer.write_i32::<LittleEndian>(self.message_length).is_err() {
            return Err("Unable to write message_length".to_owned())
        }

        if buffer.write_i32::<LittleEndian>(self.request_id).is_err() {
            return Err("Unable to write request_id".to_owned())
        }

        if buffer.write_i32::<LittleEndian>(self.response_to).is_err() {
            return Err("Unable to write response_to".to_owned())
        }

        if buffer.write_i32::<LittleEndian>(self.op_code.clone() as i32).is_err() {
            return Err("Unable to write op_code".to_owned())
        }

        let _ = buffer.flush();

        Ok(())
    }

    /// Reads a serialized Header from a buffer.
    ///
    /// # Arguments
    ///
    /// `buffer` - The buffer to read from.
    ///
    /// # Return value
    ///
    /// Returns the parsed Header on success, or an error string on failure.
    pub fn read(buffer: &mut Read) -> Result<Header, String> {
        let message_length = match buffer.read_i32::<LittleEndian>() {
            Ok(i) => i,
            Err(_) => return Err("unable to read length".to_owned())
        };

        let request_id = match buffer.read_i32::<LittleEndian>() {
            Ok(i) => i,
            Err(_) => return Err("unable to read request_id".to_owned())
        };

        let response_to = match buffer.read_i32::<LittleEndian>() {
            Ok(i) => i,
            Err(_) => return Err("unable to read response_to".to_owned())
        };

        let op_code = match buffer.read_i32::<LittleEndian>() {
            Ok(i) => match OpCode::from_i32(i) {
                Some(op) => op,
                None => return Err(format!("invalid opcode: {}", i))
            },
            Err(_) => return Err("unable to read op_code".to_owned())
        };

        Ok(Header::new(message_length, request_id, response_to, op_code))
    }
}
<|endoftext|>


<commit_before>//! This module contains functions for retrieve the original AST from lowered
//! `hir`.

#![deny(clippy::missing_docs_in_private_items)]

use crate::utils::{is_expn_of, match_def_path, match_qpath, paths};
use if_chain::if_chain;
use rustc::lint::LateContext;
use rustc::{hir, ty};
use syntax::ast;

/// Converts a hir binary operator to the corresponding `ast` type.
#[must_use]
pub fn binop(op: hir::BinOpKind) -> ast::BinOpKind {
    match op {
        hir::BinOpKind::Eq => ast::BinOpKind::Eq,
        hir::BinOpKind::Ge => ast::BinOpKind::Ge,
        hir::BinOpKind::Gt => ast::BinOpKind::Gt,
        hir::BinOpKind::Le => ast::BinOpKind::Le,
        hir::BinOpKind::Lt => ast::BinOpKind::Lt,
        hir::BinOpKind::Ne => ast::BinOpKind::Ne,
        hir::BinOpKind::Or => ast::BinOpKind::Or,
        hir::BinOpKind::Add => ast::BinOpKind::Add,
        hir::BinOpKind::And => ast::BinOpKind::And,
        hir::BinOpKind::BitAnd => ast::BinOpKind::BitAnd,
        hir::BinOpKind::BitOr => ast::BinOpKind::BitOr,
        hir::BinOpKind::BitXor => ast::BinOpKind::BitXor,
        hir::BinOpKind::Div => ast::BinOpKind::Div,
        hir::BinOpKind::Mul => ast::BinOpKind::Mul,
        hir::BinOpKind::Rem => ast::BinOpKind::Rem,
        hir::BinOpKind::Shl => ast::BinOpKind::Shl,
        hir::BinOpKind::Shr => ast::BinOpKind::Shr,
        hir::BinOpKind::Sub => ast::BinOpKind::Sub,
    }
}

/// Represent a range akin to `ast::ExprKind::Range`.
#[derive(Debug, Copy, Clone)]
pub struct Range<'a> {
    /// The lower bound of the range, or `None` for ranges such as `..X`.
    pub start: Option<&'a hir::Expr>,
    /// The upper bound of the range, or `None` for ranges such as `X..`.
    pub end: Option<&'a hir::Expr>,
    /// Whether the interval is open or closed.
    pub limits: ast::RangeLimits,
}

/// Higher a `hir` range to something similar to `ast::ExprKind::Range`.
pub fn range<'a, 'b, 'tcx>(cx: &LateContext<'a, 'tcx>, expr: &'b hir::Expr) -> Option<Range<'b>> {
    /// Finds the field named `name` in the field. Always return `Some` for
    /// convenience.
    fn get_field<'c>(name: &str, fields: &'c [hir::Field]) -> Option<&'c hir::Expr> {
        let expr = &fields.iter().find(|field| field.ident.name.as_str() == name)?.expr;

        Some(expr)
    }

    let def_path = match cx.tables.expr_ty(expr).kind {
        ty::Adt(def, _) => cx.tcx.def_path(def.did),
        _ => return None,
    };

    // sanity checks for std::ops::RangeXXXX
    if def_path.data.len() != 3 {
        return None;
    }
    if def_path.data.get(0)?.data.as_interned_str().as_symbol() != sym!(ops) {
        return None;
    }
    if def_path.data.get(1)?.data.as_interned_str().as_symbol() != sym!(range) {
        return None;
    }
    let type_name = def_path.data.get(2)?.data.as_interned_str();
    let range_types = [
        "RangeFrom",
        "RangeFull",
        "RangeInclusive",
        "Range",
        "RangeTo",
        "RangeToInclusive",
    ];
    if !range_types.contains(&&*type_name.as_str()) {
        return None;
    }

    // The range syntax is expanded to literal paths starting with `core` or `std`
    // depending on
    // `#[no_std]`. Testing both instead of resolving the paths.

    match expr.kind {
        hir::ExprKind::Path(ref path) => {
            if match_qpath(path, &paths::RANGE_FULL_STD) || match_qpath(path, &paths::RANGE_FULL) {
                Some(Range {
                    start: None,
                    end: None,
                    limits: ast::RangeLimits::HalfOpen,
                })
            } else {
                None
            }
        },
        hir::ExprKind::Call(ref path, ref args) => {
            if let hir::ExprKind::Path(ref path) = path.kind {
                if match_qpath(path, &paths::RANGE_INCLUSIVE_STD_NEW) || match_qpath(path, &paths::RANGE_INCLUSIVE_NEW)
                {
                    Some(Range {
                        start: Some(&args[0]),
                        end: Some(&args[1]),
                        limits: ast::RangeLimits::Closed,
                    })
                } else {
                    None
                }
            } else {
                None
            }
        },
        hir::ExprKind::Struct(ref path, ref fields, None) => {
            if match_qpath(path, &paths::RANGE_FROM_STD) || match_qpath(path, &paths::RANGE_FROM) {
                Some(Range {
                    start: Some(get_field("start", fields)?),
                    end: None,
                    limits: ast::RangeLimits::HalfOpen,
                })
            } else if match_qpath(path, &paths::RANGE_STD) || match_qpath(path, &paths::RANGE) {
                Some(Range {
                    start: Some(get_field("start", fields)?),
                    end: Some(get_field("end", fields)?),
                    limits: ast::RangeLimits::HalfOpen,
                })
            } else if match_qpath(path, &paths::RANGE_TO_INCLUSIVE_STD) || match_qpath(path, &paths::RANGE_TO_INCLUSIVE)
            {
                Some(Range {
                    start: None,
                    end: Some(get_field("end", fields)?),
                    limits: ast::RangeLimits::Closed,
                })
            } else if match_qpath(path, &paths::RANGE_TO_STD) || match_qpath(path, &paths::RANGE_TO) {
                Some(Range {
                    start: None,
                    end: Some(get_field("end", fields)?),
                    limits: ast::RangeLimits::HalfOpen,
                })
            } else {
                None
            }
        },
        _ => None,
    }
}

/// Checks if a `let` statement is from a `for` loop desugaring.
pub fn is_from_for_desugar(local: &hir::Local) -> bool {
    // This will detect plain for-loops without an actual variable binding:
    //
    // ```
    // for x in some_vec {
    //     // do stuff
    // }
    // ```
    if_chain! {
        if let Some(ref expr) = local.init;
        if let hir::ExprKind::Match(_, _, hir::MatchSource::ForLoopDesugar) = expr.kind;
        then {
            return true;
        }
    }

    // This detects a variable binding in for loop to avoid `let_unit_value`
    // lint (see issue #1964).
    //
    // ```
    // for _ in vec![()] {
    //     // anything
    // }
    // ```
    if let hir::LocalSource::ForLoopDesugar = local.source {
        return true;
    }

    false
}

/// Recover the essential nodes of a desugared for loop:
/// `for pat in arg { body }` becomes `(pat, arg, body)`.
pub fn for_loop(expr: &hir::Expr) -> Option<(&hir::Pat, &hir::Expr, &hir::Expr)> {
    if_chain! {
        if let hir::ExprKind::Match(ref iterexpr, ref arms, hir::MatchSource::ForLoopDesugar) = expr.kind;
        if let hir::ExprKind::Call(_, ref iterargs) = iterexpr.kind;
        if iterargs.len() == 1 && arms.len() == 1 && arms[0].guard.is_none();
        if let hir::ExprKind::Loop(ref block, _, _) = arms[0].body.kind;
        if block.expr.is_none();
        if let [ _, _, ref let_stmt, ref body ] = *block.stmts;
        if let hir::StmtKind::Local(ref local) = let_stmt.kind;
        if let hir::StmtKind::Expr(ref expr) = body.kind;
        then {
            return Some((&*local.pat, &iterargs[0], expr));
        }
    }
    None
}

/// Recover the essential nodes of a desugared while loop:
/// `while cond { body }` becomes `(cond, body)`.
pub fn while_loop(expr: &hir::Expr) -> Option<(&hir::Expr, &hir::Expr)> {
    if_chain! {
        if let hir::ExprKind::Loop(block, _, hir::LoopSource::While) = &expr.kind;
        if let hir::Block { expr: Some(expr), .. } = &**block;
        if let hir::ExprKind::Match(cond, arms, hir::MatchSource::WhileDesugar) = &expr.kind;
        if let hir::ExprKind::DropTemps(cond) = &cond.kind;
        if let [arm, ..] = &arms[..];
        if let hir::Arm { body, .. } = arm;
        then {
            return Some((cond, body));
        }
    }
    None
}

/// Recover the essential nodes of a desugared if block
/// `if cond { then } else { els }` becomes `(cond, then, Some(els))`
pub fn if_block(expr: &hir::Expr) -> Option<(&hir::Expr, &hir::Expr, Option<&hir::Expr>)> {
    if let hir::ExprKind::Match(ref cond, ref arms, hir::MatchSource::IfDesugar { contains_else_clause }) = expr.kind {
        let cond = if let hir::ExprKind::DropTemps(ref cond) = cond.kind {
            cond
        } else {
            panic!("If block desugar must contain DropTemps");
        };
        let then = &arms[0].body;
        let els = if contains_else_clause {
            Some(&*arms[1].body)
        } else {
            None
        };
        Some((cond, then, els))
    } else {
        None
    }
}

/// Represent the pre-expansion arguments of a `vec!` invocation.
pub enum VecArgs<'a> {
    /// `vec![elem; len]`
    Repeat(&'a hir::Expr, &'a hir::Expr),
    /// `vec![a, b, c]`
    Vec(&'a [hir::Expr]),
}

/// Returns the arguments of the `vec!` macro if this expression was expanded
/// from `vec!`.
pub fn vec_macro<'e>(cx: &LateContext<'_, '_>, expr: &'e hir::Expr) -> Option<VecArgs<'e>> {
    if_chain! {
        if let hir::ExprKind::Call(ref fun, ref args) = expr.kind;
        if let hir::ExprKind::Path(ref qpath) = fun.kind;
        if is_expn_of(fun.span, "vec").is_some();
        if let Some(fun_def_id) = cx.tables.qpath_res(qpath, fun.hir_id).opt_def_id();
        then {
            return if match_def_path(cx, fun_def_id, &paths::VEC_FROM_ELEM) && args.len() == 2 {
                // `vec![elem; size]` case
                Some(VecArgs::Repeat(&args[0], &args[1]))
            }
            else if match_def_path(cx, fun_def_id, &paths::SLICE_INTO_VEC) && args.len() == 1 {
                // `vec![a, b, c]` case
                if_chain! {
                    if let hir::ExprKind::Box(ref boxed) = args[0].kind;
                    if let hir::ExprKind::Array(ref args) = boxed.kind;
                    then {
                        return Some(VecArgs::Vec(&*args));
                    }
                }

                None
            }
            else {
                None
            };
        }
    }

    None
}
<commit_msg>Rustup to rust-lang/rust#65657<commit_after>//! This module contains functions for retrieve the original AST from lowered
//! `hir`.

#![deny(clippy::missing_docs_in_private_items)]

use crate::utils::{is_expn_of, match_def_path, match_qpath, paths};
use if_chain::if_chain;
use rustc::lint::LateContext;
use rustc::{hir, ty};
use syntax::ast;

/// Converts a hir binary operator to the corresponding `ast` type.
#[must_use]
pub fn binop(op: hir::BinOpKind) -> ast::BinOpKind {
    match op {
        hir::BinOpKind::Eq => ast::BinOpKind::Eq,
        hir::BinOpKind::Ge => ast::BinOpKind::Ge,
        hir::BinOpKind::Gt => ast::BinOpKind::Gt,
        hir::BinOpKind::Le => ast::BinOpKind::Le,
        hir::BinOpKind::Lt => ast::BinOpKind::Lt,
        hir::BinOpKind::Ne => ast::BinOpKind::Ne,
        hir::BinOpKind::Or => ast::BinOpKind::Or,
        hir::BinOpKind::Add => ast::BinOpKind::Add,
        hir::BinOpKind::And => ast::BinOpKind::And,
        hir::BinOpKind::BitAnd => ast::BinOpKind::BitAnd,
        hir::BinOpKind::BitOr => ast::BinOpKind::BitOr,
        hir::BinOpKind::BitXor => ast::BinOpKind::BitXor,
        hir::BinOpKind::Div => ast::BinOpKind::Div,
        hir::BinOpKind::Mul => ast::BinOpKind::Mul,
        hir::BinOpKind::Rem => ast::BinOpKind::Rem,
        hir::BinOpKind::Shl => ast::BinOpKind::Shl,
        hir::BinOpKind::Shr => ast::BinOpKind::Shr,
        hir::BinOpKind::Sub => ast::BinOpKind::Sub,
    }
}

/// Represent a range akin to `ast::ExprKind::Range`.
#[derive(Debug, Copy, Clone)]
pub struct Range<'a> {
    /// The lower bound of the range, or `None` for ranges such as `..X`.
    pub start: Option<&'a hir::Expr>,
    /// The upper bound of the range, or `None` for ranges such as `X..`.
    pub end: Option<&'a hir::Expr>,
    /// Whether the interval is open or closed.
    pub limits: ast::RangeLimits,
}

/// Higher a `hir` range to something similar to `ast::ExprKind::Range`.
pub fn range<'a, 'b, 'tcx>(cx: &LateContext<'a, 'tcx>, expr: &'b hir::Expr) -> Option<Range<'b>> {
    /// Finds the field named `name` in the field. Always return `Some` for
    /// convenience.
    fn get_field<'c>(name: &str, fields: &'c [hir::Field]) -> Option<&'c hir::Expr> {
        let expr = &fields.iter().find(|field| field.ident.name.as_str() == name)?.expr;

        Some(expr)
    }

    let def_path = match cx.tables.expr_ty(expr).kind {
        ty::Adt(def, _) => cx.tcx.def_path(def.did),
        _ => return None,
    };

    // sanity checks for std::ops::RangeXXXX
    if def_path.data.len() != 3 {
        return None;
    }
    if def_path.data.get(0)?.data.as_symbol() != sym!(ops) {
        return None;
    }
    if def_path.data.get(1)?.data.as_symbol() != sym!(range) {
        return None;
    }
    let type_name = def_path.data.get(2)?.data.as_symbol();
    let range_types = [
        "RangeFrom",
        "RangeFull",
        "RangeInclusive",
        "Range",
        "RangeTo",
        "RangeToInclusive",
    ];
    if !range_types.contains(&&*type_name.as_str()) {
        return None;
    }

    // The range syntax is expanded to literal paths starting with `core` or `std`
    // depending on
    // `#[no_std]`. Testing both instead of resolving the paths.

    match expr.kind {
        hir::ExprKind::Path(ref path) => {
            if match_qpath(path, &paths::RANGE_FULL_STD) || match_qpath(path, &paths::RANGE_FULL) {
                Some(Range {
                    start: None,
                    end: None,
                    limits: ast::RangeLimits::HalfOpen,
                })
            } else {
                None
            }
        },
        hir::ExprKind::Call(ref path, ref args) => {
            if let hir::ExprKind::Path(ref path) = path.kind {
                if match_qpath(path, &paths::RANGE_INCLUSIVE_STD_NEW) || match_qpath(path, &paths::RANGE_INCLUSIVE_NEW)
                {
                    Some(Range {
                        start: Some(&args[0]),
                        end: Some(&args[1]),
                        limits: ast::RangeLimits::Closed,
                    })
                } else {
                    None
                }
            } else {
                None
            }
        },
        hir::ExprKind::Struct(ref path, ref fields, None) => {
            if match_qpath(path, &paths::RANGE_FROM_STD) || match_qpath(path, &paths::RANGE_FROM) {
                Some(Range {
                    start: Some(get_field("start", fields)?),
                    end: None,
                    limits: ast::RangeLimits::HalfOpen,
                })
            } else if match_qpath(path, &paths::RANGE_STD) || match_qpath(path, &paths::RANGE) {
                Some(Range {
                    start: Some(get_field("start", fields)?),
                    end: Some(get_field("end", fields)?),
                    limits: ast::RangeLimits::HalfOpen,
                })
            } else if match_qpath(path, &paths::RANGE_TO_INCLUSIVE_STD) || match_qpath(path, &paths::RANGE_TO_INCLUSIVE)
            {
                Some(Range {
                    start: None,
                    end: Some(get_field("end", fields)?),
                    limits: ast::RangeLimits::Closed,
                })
            } else if match_qpath(path, &paths::RANGE_TO_STD) || match_qpath(path, &paths::RANGE_TO) {
                Some(Range {
                    start: None,
                    end: Some(get_field("end", fields)?),
                    limits: ast::RangeLimits::HalfOpen,
                })
            } else {
                None
            }
        },
        _ => None,
    }
}

/// Checks if a `let` statement is from a `for` loop desugaring.
pub fn is_from_for_desugar(local: &hir::Local) -> bool {
    // This will detect plain for-loops without an actual variable binding:
    //
    // ```
    // for x in some_vec {
    //     // do stuff
    // }
    // ```
    if_chain! {
        if let Some(ref expr) = local.init;
        if let hir::ExprKind::Match(_, _, hir::MatchSource::ForLoopDesugar) = expr.kind;
        then {
            return true;
        }
    }

    // This detects a variable binding in for loop to avoid `let_unit_value`
    // lint (see issue #1964).
    //
    // ```
    // for _ in vec![()] {
    //     // anything
    // }
    // ```
    if let hir::LocalSource::ForLoopDesugar = local.source {
        return true;
    }

    false
}

/// Recover the essential nodes of a desugared for loop:
/// `for pat in arg { body }` becomes `(pat, arg, body)`.
pub fn for_loop(expr: &hir::Expr) -> Option<(&hir::Pat, &hir::Expr, &hir::Expr)> {
    if_chain! {
        if let hir::ExprKind::Match(ref iterexpr, ref arms, hir::MatchSource::ForLoopDesugar) = expr.kind;
        if let hir::ExprKind::Call(_, ref iterargs) = iterexpr.kind;
        if iterargs.len() == 1 && arms.len() == 1 && arms[0].guard.is_none();
        if let hir::ExprKind::Loop(ref block, _, _) = arms[0].body.kind;
        if block.expr.is_none();
        if let [ _, _, ref let_stmt, ref body ] = *block.stmts;
        if let hir::StmtKind::Local(ref local) = let_stmt.kind;
        if let hir::StmtKind::Expr(ref expr) = body.kind;
        then {
            return Some((&*local.pat, &iterargs[0], expr));
        }
    }
    None
}

/// Recover the essential nodes of a desugared while loop:
/// `while cond { body }` becomes `(cond, body)`.
pub fn while_loop(expr: &hir::Expr) -> Option<(&hir::Expr, &hir::Expr)> {
    if_chain! {
        if let hir::ExprKind::Loop(block, _, hir::LoopSource::While) = &expr.kind;
        if let hir::Block { expr: Some(expr), .. } = &**block;
        if let hir::ExprKind::Match(cond, arms, hir::MatchSource::WhileDesugar) = &expr.kind;
        if let hir::ExprKind::DropTemps(cond) = &cond.kind;
        if let [arm, ..] = &arms[..];
        if let hir::Arm { body, .. } = arm;
        then {
            return Some((cond, body));
        }
    }
    None
}

/// Recover the essential nodes of a desugared if block
/// `if cond { then } else { els }` becomes `(cond, then, Some(els))`
pub fn if_block(expr: &hir::Expr) -> Option<(&hir::Expr, &hir::Expr, Option<&hir::Expr>)> {
    if let hir::ExprKind::Match(ref cond, ref arms, hir::MatchSource::IfDesugar { contains_else_clause }) = expr.kind {
        let cond = if let hir::ExprKind::DropTemps(ref cond) = cond.kind {
            cond
        } else {
            panic!("If block desugar must contain DropTemps");
        };
        let then = &arms[0].body;
        let els = if contains_else_clause {
            Some(&*arms[1].body)
        } else {
            None
        };
        Some((cond, then, els))
    } else {
        None
    }
}

/// Represent the pre-expansion arguments of a `vec!` invocation.
pub enum VecArgs<'a> {
    /// `vec![elem; len]`
    Repeat(&'a hir::Expr, &'a hir::Expr),
    /// `vec![a, b, c]`
    Vec(&'a [hir::Expr]),
}

/// Returns the arguments of the `vec!` macro if this expression was expanded
/// from `vec!`.
pub fn vec_macro<'e>(cx: &LateContext<'_, '_>, expr: &'e hir::Expr) -> Option<VecArgs<'e>> {
    if_chain! {
        if let hir::ExprKind::Call(ref fun, ref args) = expr.kind;
        if let hir::ExprKind::Path(ref qpath) = fun.kind;
        if is_expn_of(fun.span, "vec").is_some();
        if let Some(fun_def_id) = cx.tables.qpath_res(qpath, fun.hir_id).opt_def_id();
        then {
            return if match_def_path(cx, fun_def_id, &paths::VEC_FROM_ELEM) && args.len() == 2 {
                // `vec![elem; size]` case
                Some(VecArgs::Repeat(&args[0], &args[1]))
            }
            else if match_def_path(cx, fun_def_id, &paths::SLICE_INTO_VEC) && args.len() == 1 {
                // `vec![a, b, c]` case
                if_chain! {
                    if let hir::ExprKind::Box(ref boxed) = args[0].kind;
                    if let hir::ExprKind::Array(ref args) = boxed.kind;
                    then {
                        return Some(VecArgs::Vec(&*args));
                    }
                }

                None
            }
            else {
                None
            };
        }
    }

    None
}
<|endoftext|>


<commit_before>// Lumol, an extensible molecular simulation engine
// Copyright (C) Lumol's contributors — BSD license

use special::Error;
use std::f64::consts::PI;

use sys::System;
use types::{Matrix3, Vector3D, Zero};
use consts::ELCC;
use energy::{PairRestriction, RestrictionInfo};
use parallel::prelude::*;
use parallel::ThreadLocalStore;

use super::{GlobalPotential, CoulombicPotential, GlobalCache};

/// Wolf summation for coulombic interactions.
///
/// This is a fast, direct, pairwise summation for coulombic potential
/// [Wolf1999].
///
/// # Examples
///
/// ```
/// use lumol::energy::Wolf;
/// use lumol::units;
///
/// // A relatively large cutoff is needed for Wolf summation
/// let wolf = Wolf::new(12.0);
///
/// use lumol::sys::System;
/// use lumol::sys::Particle;
/// use lumol::sys::UnitCell;
/// use lumol::types::Vector3D;
///
/// // Setup a system containing a NaCl pair
/// let mut system = System::new();
/// system.set_cell(UnitCell::cubic(10.0));
///
/// let mut na = Particle::new("Na");
/// na.charge = 1.0;
/// na.position = Vector3D::new(0.0, 0.0, 0.0);
///
/// let mut cl = Particle::new("Cl");
/// cl.charge = -1.0;
/// cl.position = Vector3D::new(2.0, 0.0, 0.0);
///
/// system.add_particle(na);
/// system.add_particle(cl);
///
/// // Use Wolf summation for electrostatic interactions
/// system.interactions_mut().set_coulomb(Box::new(wolf));
///
/// assert_eq!(system.potential_energy(), -0.07292902695393541);
/// ```
///
/// [Wolf1999]: Wolf, D. et al. J. Chem. Phys. 110, 8254 (1999).
#[derive(Clone)]
pub struct Wolf {
    /// Splitting parameter
    alpha: f64,
    /// Cutoff radius in real space
    cutoff: f64,
    /// Energy constant for wolf computation
    energy_cst: f64,
    /// Force constant for wolf computation
    force_cst: f64,
    /// Restriction scheme
    restriction: PairRestriction,
}

impl Wolf {
    /// Create a new Wolf summation, using a real-space cutoff of `cutoff`.
    pub fn new(cutoff: f64) -> Wolf {
        assert!(cutoff > 0.0, "Got a negative cutoff in Wolf summation");
        let alpha = PI/cutoff;
        let e_cst = f64::erfc(alpha*cutoff)/cutoff;
        let f_cst = f64::erfc(alpha*cutoff)/(cutoff*cutoff) + 2.0*alpha/f64::sqrt(PI) * f64::exp(-alpha*alpha*cutoff*cutoff)/cutoff;
        Wolf{
            alpha: alpha,
            cutoff: cutoff,
            energy_cst: e_cst,
            force_cst: f_cst,
            restriction: PairRestriction::None,
        }
    }

    /// Compute the energy for the pair of particles with charge `qi` and `qj`,
    /// at the distance of `rij`. The `scaling` parameter comes from the
    /// restriction associated with this potential.
    #[inline]
    fn energy_pair(&self, info: RestrictionInfo, qi: f64, qj: f64, rij: f64) -> f64 {
        if rij > self.cutoff || info.excluded {
            return 0.0;
        }
        info.scaling * qi * qj * (f64::erfc(self.alpha*rij)/rij - self.energy_cst) / ELCC
    }

    /// Compute the energy for self interaction of a particle with charge `qi`
    #[inline]
    fn energy_self(&self, qi: f64) -> f64 {
        qi * qi * (self.energy_cst/2.0 + self.alpha/f64::sqrt(PI)) / ELCC
    }

    /// Compute the force for self the pair of particles with charge `qi` and
    /// `qj`, at the distance of `rij`. The `scaling` parameter comes from the
    /// restriction associated with this potential.
    #[inline]
    fn force_pair(&self, info: RestrictionInfo, qi: f64, qj: f64, rij: Vector3D) -> Vector3D {
        let d = rij.norm();
        if d > self.cutoff || info.excluded {
            return Vector3D::zero();
        }
        let factor = f64::erfc(self.alpha*d)/(d*d) + 2.0*self.alpha/f64::sqrt(PI) * f64::exp(-self.alpha*self.alpha*d*d)/d;
        return info.scaling * qi * qj * (factor - self.force_cst) * rij.normalized() / ELCC;
    }
}

impl GlobalCache for Wolf {
    fn move_particles_cost(&self, system: &System, idxes: &[usize], newpos: &[Vector3D]) -> f64 {
        let mut e_old = 0.0;
        let mut e_new = 0.0;

        // Iterate over all interactions between a moved particle and a
        // particle not moved
        for (idx, &i) in idxes.iter().enumerate() {
            let qi = system[i].charge;
            if qi == 0.0 {continue;}
            for j in (0..system.size()).filter(|x| !idxes.contains(x)) {
                let qj = system[j].charge;
                if qj == 0.0 {continue;}

                let r_old = system.cell().distance(&system[i].position, &system[j].position);
                let r_new = system.cell().distance(&newpos[idx], &system[j].position);

                let distance = system.bond_distance(i, j);
                let info = self.restriction.information(distance);

                e_old += self.energy_pair(info, qi, qj, r_old);
                e_new += self.energy_pair(info, qi, qj, r_new);
            }
        }

        // Iterate over all interactions between two moved particles
        for (idx, &i) in idxes.iter().enumerate() {
            let qi = system[i].charge;
            if qi == 0.0 {continue;}
            for (jdx, &j) in idxes.iter().enumerate().skip(idx + 1) {
                let qj = system[j].charge;
                if qj == 0.0 {continue;}

                let r_old = system.distance(i, j);
                let r_new = system.cell().distance(&newpos[idx], &newpos[jdx]);

                let distance = system.bond_distance(i, j);
                let info = self.restriction.information(distance);

                e_old += self.energy_pair(info, qi, qj, r_old);
                e_new += self.energy_pair(info, qi, qj, r_new);
            }
        }

        return e_new - e_old;
    }

    fn update(&self) {
        // Nothing to do
    }
}

impl GlobalPotential for Wolf {
    fn cutoff(&self) -> Option<f64> {
        Some(self.cutoff)
    }

    fn energy(&self, system: &System) -> f64 {
        let natoms = system.size();
        let mut res = 0.0;
        for i in 0..natoms {
            let qi = system[i].charge;
            if qi == 0.0 {continue}
            for j in i+1..natoms {
                let qj = system[j].charge;
                if qj == 0.0 {continue}

                let distance = system.bond_distance(i, j);
                let info = self.restriction.information(distance);

                let rij = system.distance(i, j);
                res += self.energy_pair(info, qi, qj, rij);
            }
            // Remove self term
            res -= self.energy_self(qi);
        }
        return res;
    }

    fn forces(&self, system: &System) -> Vec<Vector3D> {
        // To avoid race conditions, each thread needs its
        // own local forces Vec
        let natoms = system.size();
        let thread_forces_store = ThreadLocalStore::new(|| vec![Vector3D::zero(); natoms]);

        (0..natoms).into_par_iter().for_each(|i| {
            /// Get the thread local forces Vec
            let mut thread_forces = thread_forces_store.borrow_mut();

            let qi = system[i].charge;
            if qi == 0.0 { return; }
            for j in i + 1..natoms {
                let qj = system[j].charge;
                if qj == 0.0 { continue }

                let distance = system.bond_distance(i, j);
                let info = self.restriction.information(distance);

                let rij = system.nearest_image(i, j);
                let force = self.force_pair(info, qi, qj, rij);
                thread_forces[i] += force;
                thread_forces[j] -= force;
            }
        });

        // At this point all the forces are computed, but the
        // results are scattered across all thread local Vecs,
        // here we gather them.
        let mut forces = vec![Vector3D::zero(); natoms];
        thread_forces_store.sum_local_values(&mut forces);

        return forces;
    }

    fn virial(&self, system: &System) -> Matrix3 {
        let natoms = system.size();
        let mut res = Matrix3::zero();
        for i in 0..natoms {
            let qi = system[i].charge;
            if qi == 0.0 {continue}
            for j in i+1..natoms {
                let qj = system[j].charge;
                if qj == 0.0 {continue}

                let distance = system.bond_distance(i, j);
                let info = self.restriction.information(distance);

                let rij = system.nearest_image(i, j);
                let force = self.force_pair(info, qi, qj, rij);
                res += force.tensorial(&rij);
            }
        }
        return res;
    }
}

impl CoulombicPotential for Wolf {
    fn set_restriction(&mut self, restriction: PairRestriction) {
        self.restriction = restriction;
    }
}

#[cfg(test)]
mod tests {
    pub use super::*;
    use sys::{System, UnitCell, Particle};
    use types::{Vector3D, Zero};
    use energy::GlobalPotential;

    const E_BRUTE_FORCE: f64 = -0.09262397663346732;

    pub fn testing_system() -> System {
        let mut system = System::from_cell(UnitCell::cubic(20.0));

        system.add_particle(Particle::new("Cl"));
        system[0].charge = -1.0;
        system[0].position = Vector3D::zero();

        system.add_particle(Particle::new("Na"));
        system[1].charge = 1.0;
        system[1].position = Vector3D::new(1.5, 0.0, 0.0);

        return system;
    }

    #[test]
    fn energy() {
        let system = testing_system();
        let wolf = Wolf::new(8.0);

        let e = wolf.energy(&system);
        // Wolf is not very good for heterogeneous systems
        assert_ulps_eq!(e, E_BRUTE_FORCE, epsilon=1e-2);
    }

    #[test]
    fn forces() {
        let mut system = testing_system();
        let wolf = Wolf::new(8.0);

        let forces = wolf.forces(&system);
        let norm = (forces[0] + forces[1]).norm();
        // Total force should be null
        assert_ulps_eq!(norm, 0.0);

        // Finite difference computation of the force
        let e = wolf.energy(&system);
        let eps = 1e-9;
        system[0].position[0] += eps;

        let e1 = wolf.energy(&system);
        let force = wolf.forces(&system)[0][0];
        assert_relative_eq!((e - e1) / eps, force, epsilon=1e-6);
    }

    mod cache {
        use super::*;
        use sys::System;
        use types::Vector3D;
        use energy::{GlobalPotential, PairRestriction, CoulombicPotential, GlobalCache};

        pub fn testing_system() -> System {
            use utils::system_from_xyz;
            let mut system = system_from_xyz("6
            bonds
            O  0.0  0.0  0.0
            H -0.7 -0.7  0.3
            H  0.3 -0.3 -0.8
            O  2.0  2.0  0.0
            H  1.3  1.3  0.3
            H  2.3  1.7 -0.8
            ");
            system.set_cell(UnitCell::cubic(20.0));
            assert!(system.molecules().len() == 2);

            for particle in &mut system {
                if particle.name() == "O" {
                    particle.charge = -0.8476;
                } else if particle.name() == "H" {
                    particle.charge = 0.4238;
                }
            }
            return system;
        }

        #[test]
        fn move_atoms() {
            let mut system = testing_system();
            let mut wolf = Wolf::new(8.0);
            wolf.set_restriction(PairRestriction::InterMolecular);

            let check = wolf.clone();

            let old_e = check.energy(&system);
            let idxes = &[0, 1];
            let newpos = &[Vector3D::new(0.0, 0.0, 0.5), Vector3D::new(-0.7, 0.2, 1.5)];

            let cost = wolf.move_particles_cost(&system, idxes, newpos);

            system[0].position = newpos[0];
            system[1].position = newpos[1];
            let new_e = check.energy(&system);
            assert_ulps_eq!(cost, new_e - old_e);
        }
    }
}
<commit_msg>Parallelize Wolf::energy<commit_after>// Lumol, an extensible molecular simulation engine
// Copyright (C) Lumol's contributors — BSD license

use special::Error;
use std::f64::consts::PI;

use rayon::prelude::*;

use sys::System;
use types::{Matrix3, Vector3D, Zero};
use consts::ELCC;
use energy::{PairRestriction, RestrictionInfo};
use parallel::prelude::*;
use parallel::ThreadLocalStore;

use super::{GlobalPotential, CoulombicPotential, GlobalCache};

/// Wolf summation for coulombic interactions.
///
/// This is a fast, direct, pairwise summation for coulombic potential
/// [Wolf1999].
///
/// # Examples
///
/// ```
/// use lumol::energy::Wolf;
/// use lumol::units;
///
/// // A relatively large cutoff is needed for Wolf summation
/// let wolf = Wolf::new(12.0);
///
/// use lumol::sys::System;
/// use lumol::sys::Particle;
/// use lumol::sys::UnitCell;
/// use lumol::types::Vector3D;
///
/// // Setup a system containing a NaCl pair
/// let mut system = System::new();
/// system.set_cell(UnitCell::cubic(10.0));
///
/// let mut na = Particle::new("Na");
/// na.charge = 1.0;
/// na.position = Vector3D::new(0.0, 0.0, 0.0);
///
/// let mut cl = Particle::new("Cl");
/// cl.charge = -1.0;
/// cl.position = Vector3D::new(2.0, 0.0, 0.0);
///
/// system.add_particle(na);
/// system.add_particle(cl);
///
/// // Use Wolf summation for electrostatic interactions
/// system.interactions_mut().set_coulomb(Box::new(wolf));
///
/// assert_eq!(system.potential_energy(), -0.07292902695393541);
/// ```
///
/// [Wolf1999]: Wolf, D. et al. J. Chem. Phys. 110, 8254 (1999).
#[derive(Clone)]
pub struct Wolf {
    /// Splitting parameter
    alpha: f64,
    /// Cutoff radius in real space
    cutoff: f64,
    /// Energy constant for wolf computation
    energy_cst: f64,
    /// Force constant for wolf computation
    force_cst: f64,
    /// Restriction scheme
    restriction: PairRestriction,
}

impl Wolf {
    /// Create a new Wolf summation, using a real-space cutoff of `cutoff`.
    pub fn new(cutoff: f64) -> Wolf {
        assert!(cutoff > 0.0, "Got a negative cutoff in Wolf summation");
        let alpha = PI/cutoff;
        let e_cst = f64::erfc(alpha*cutoff)/cutoff;
        let f_cst = f64::erfc(alpha*cutoff)/(cutoff*cutoff) + 2.0*alpha/f64::sqrt(PI) * f64::exp(-alpha*alpha*cutoff*cutoff)/cutoff;
        Wolf{
            alpha: alpha,
            cutoff: cutoff,
            energy_cst: e_cst,
            force_cst: f_cst,
            restriction: PairRestriction::None,
        }
    }

    /// Compute the energy for the pair of particles with charge `qi` and `qj`,
    /// at the distance of `rij`. The `scaling` parameter comes from the
    /// restriction associated with this potential.
    #[inline]
    fn energy_pair(&self, info: RestrictionInfo, qi: f64, qj: f64, rij: f64) -> f64 {
        if rij > self.cutoff || info.excluded {
            return 0.0;
        }
        info.scaling * qi * qj * (f64::erfc(self.alpha*rij)/rij - self.energy_cst) / ELCC
    }

    /// Compute the energy for self interaction of a particle with charge `qi`
    #[inline]
    fn energy_self(&self, qi: f64) -> f64 {
        qi * qi * (self.energy_cst/2.0 + self.alpha/f64::sqrt(PI)) / ELCC
    }

    /// Compute the force for self the pair of particles with charge `qi` and
    /// `qj`, at the distance of `rij`. The `scaling` parameter comes from the
    /// restriction associated with this potential.
    #[inline]
    fn force_pair(&self, info: RestrictionInfo, qi: f64, qj: f64, rij: Vector3D) -> Vector3D {
        let d = rij.norm();
        if d > self.cutoff || info.excluded {
            return Vector3D::zero();
        }
        let factor = f64::erfc(self.alpha*d)/(d*d) + 2.0*self.alpha/f64::sqrt(PI) * f64::exp(-self.alpha*self.alpha*d*d)/d;
        return info.scaling * qi * qj * (factor - self.force_cst) * rij.normalized() / ELCC;
    }
}

impl GlobalCache for Wolf {
    fn move_particles_cost(&self, system: &System, idxes: &[usize], newpos: &[Vector3D]) -> f64 {
        let mut e_old = 0.0;
        let mut e_new = 0.0;

        // Iterate over all interactions between a moved particle and a
        // particle not moved
        for (idx, &i) in idxes.iter().enumerate() {
            let qi = system[i].charge;
            if qi == 0.0 {continue;}
            for j in (0..system.size()).filter(|x| !idxes.contains(x)) {
                let qj = system[j].charge;
                if qj == 0.0 {continue;}

                let r_old = system.cell().distance(&system[i].position, &system[j].position);
                let r_new = system.cell().distance(&newpos[idx], &system[j].position);

                let distance = system.bond_distance(i, j);
                let info = self.restriction.information(distance);

                e_old += self.energy_pair(info, qi, qj, r_old);
                e_new += self.energy_pair(info, qi, qj, r_new);
            }
        }

        // Iterate over all interactions between two moved particles
        for (idx, &i) in idxes.iter().enumerate() {
            let qi = system[i].charge;
            if qi == 0.0 {continue;}
            for (jdx, &j) in idxes.iter().enumerate().skip(idx + 1) {
                let qj = system[j].charge;
                if qj == 0.0 {continue;}

                let r_old = system.distance(i, j);
                let r_new = system.cell().distance(&newpos[idx], &newpos[jdx]);

                let distance = system.bond_distance(i, j);
                let info = self.restriction.information(distance);

                e_old += self.energy_pair(info, qi, qj, r_old);
                e_new += self.energy_pair(info, qi, qj, r_new);
            }
        }

        return e_new - e_old;
    }

    fn update(&self) {
        // Nothing to do
    }
}

impl GlobalPotential for Wolf {
    fn cutoff(&self) -> Option<f64> {
        Some(self.cutoff)
    }

    fn energy(&self, system: &System) -> f64 {
        let natoms = system.size();

        (0..natoms).into_par_iter().map(|i|{

            let mut energy = 0.0;
            let qi = system[i].charge;
            if qi == 0.0 { return 0.0; }

            for j in i+1..natoms {
                let qj = system[j].charge;
                if qj == 0.0 {continue}

                let distance = system.bond_distance(i, j);
                let info = self.restriction.information(distance);

                let rij = system.distance(i, j);
                energy  += self.energy_pair(info, qi, qj, rij);
            }

            energy - self.energy_self(qi)
        }).sum()
    }

    fn forces(&self, system: &System) -> Vec<Vector3D> {
        // To avoid race conditions, each thread needs its
        // own local forces Vec
        let natoms = system.size();
        let thread_forces_store = ThreadLocalStore::new(|| vec![Vector3D::zero(); natoms]);

        (0..natoms).into_par_iter().for_each(|i| {
            /// Get the thread local forces Vec
            let mut thread_forces = thread_forces_store.borrow_mut();

            let qi = system[i].charge;
            if qi == 0.0 { return; }
            for j in i + 1..natoms {
                let qj = system[j].charge;
                if qj == 0.0 { continue }

                let distance = system.bond_distance(i, j);
                let info = self.restriction.information(distance);

                let rij = system.nearest_image(i, j);
                let force = self.force_pair(info, qi, qj, rij);
                thread_forces[i] += force;
                thread_forces[j] -= force;
            }
        });

        // At this point all the forces are computed, but the
        // results are scattered across all thread local Vecs,
        // here we gather them.
        let mut forces = vec![Vector3D::zero(); natoms];
        thread_forces_store.sum_local_values(&mut forces);

        return forces;
    }

    fn virial(&self, system: &System) -> Matrix3 {
        let natoms = system.size();
        let mut res = Matrix3::zero();
        for i in 0..natoms {
            let qi = system[i].charge;
            if qi == 0.0 {continue}
            for j in i+1..natoms {
                let qj = system[j].charge;
                if qj == 0.0 {continue}

                let distance = system.bond_distance(i, j);
                let info = self.restriction.information(distance);

                let rij = system.nearest_image(i, j);
                let force = self.force_pair(info, qi, qj, rij);
                res += force.tensorial(&rij);
            }
        }
        return res;
    }
}

impl CoulombicPotential for Wolf {
    fn set_restriction(&mut self, restriction: PairRestriction) {
        self.restriction = restriction;
    }
}

#[cfg(test)]
mod tests {
    pub use super::*;
    use sys::{System, UnitCell, Particle};
    use types::{Vector3D, Zero};
    use energy::GlobalPotential;

    const E_BRUTE_FORCE: f64 = -0.09262397663346732;

    pub fn testing_system() -> System {
        let mut system = System::from_cell(UnitCell::cubic(20.0));

        system.add_particle(Particle::new("Cl"));
        system[0].charge = -1.0;
        system[0].position = Vector3D::zero();

        system.add_particle(Particle::new("Na"));
        system[1].charge = 1.0;
        system[1].position = Vector3D::new(1.5, 0.0, 0.0);

        return system;
    }

    #[test]
    fn energy() {
        let system = testing_system();
        let wolf = Wolf::new(8.0);

        let e = wolf.energy(&system);
        // Wolf is not very good for heterogeneous systems
        assert_ulps_eq!(e, E_BRUTE_FORCE, epsilon=1e-2);
    }

    #[test]
    fn forces() {
        let mut system = testing_system();
        let wolf = Wolf::new(8.0);

        let forces = wolf.forces(&system);
        let norm = (forces[0] + forces[1]).norm();
        // Total force should be null
        assert_ulps_eq!(norm, 0.0);

        // Finite difference computation of the force
        let e = wolf.energy(&system);
        let eps = 1e-9;
        system[0].position[0] += eps;

        let e1 = wolf.energy(&system);
        let force = wolf.forces(&system)[0][0];
        assert_relative_eq!((e - e1) / eps, force, epsilon=1e-6);
    }

    mod cache {
        use super::*;
        use sys::System;
        use types::Vector3D;
        use energy::{GlobalPotential, PairRestriction, CoulombicPotential, GlobalCache};

        pub fn testing_system() -> System {
            use utils::system_from_xyz;
            let mut system = system_from_xyz("6
            bonds
            O  0.0  0.0  0.0
            H -0.7 -0.7  0.3
            H  0.3 -0.3 -0.8
            O  2.0  2.0  0.0
            H  1.3  1.3  0.3
            H  2.3  1.7 -0.8
            ");
            system.set_cell(UnitCell::cubic(20.0));
            assert!(system.molecules().len() == 2);

            for particle in &mut system {
                if particle.name() == "O" {
                    particle.charge = -0.8476;
                } else if particle.name() == "H" {
                    particle.charge = 0.4238;
                }
            }
            return system;
        }

        #[test]
        fn move_atoms() {
            let mut system = testing_system();
            let mut wolf = Wolf::new(8.0);
            wolf.set_restriction(PairRestriction::InterMolecular);

            let check = wolf.clone();

            let old_e = check.energy(&system);
            let idxes = &[0, 1];
            let newpos = &[Vector3D::new(0.0, 0.0, 0.5), Vector3D::new(-0.7, 0.2, 1.5)];

            let cost = wolf.move_particles_cost(&system, idxes, newpos);

            system[0].position = newpos[0];
            system[1].position = newpos[1];
            let new_e = check.energy(&system);
            assert_ulps_eq!(cost, new_e - old_e);
        }
    }
}
<|endoftext|>


<commit_before>// Generate bindings: bindgen --constified-enum "nvml.+" --ctypes-prefix raw
// --no-doc-comments --raw-line "#![allow(non_upper_case_globals)]" --raw-line
// "#![allow(non_camel_case_types)]" --raw-line "#![allow(non_snake_case)]"
// --raw-line "#![allow(dead_code)]"  --raw-line "use std::os::raw;" -o
// genned_bindings.rs nvml.h

#[cfg(target_os = "windows")]
fn main() {
    println!("cargo:rustc-link-lib=nvml");
    println!("cargo:rustc-link-search=C:\\Program Files\\NVIDIA Corporation\\NVSMI");
}

#[cfg(target_os = "linux")]
fn main() {
    let paths = std::fs::read_dir("/usr/lib").unwrap();

    for path in paths {
        let entry = path.unwrap().path();
        if entry.is_dir() {
            let entry_string = entry.to_string_lossy();
            if entry_string.contains("nvidia") && !entry_string.ends_with("prime") {
                println!("cargo:rustc-link-search=native={}", entry_string);
                break;
            }
        }
    }

    println!("cargo:rustc-link-lib=nvidia-ml");
}
<commit_msg>Output a better error message upon attempted macOS compilation<commit_after>// Generate bindings: bindgen --constified-enum "nvml.+" --ctypes-prefix raw
// --no-doc-comments --raw-line "#![allow(non_upper_case_globals)]" --raw-line
// "#![allow(non_camel_case_types)]" --raw-line "#![allow(non_snake_case)]"
// --raw-line "#![allow(dead_code)]"  --raw-line "use std::os::raw;" -o
// genned_bindings.rs nvml.h

#[cfg(target_os = "windows")]
fn main() {
    println!("cargo:rustc-link-lib=nvml");
    println!("cargo:rustc-link-search=C:\\Program Files\\NVIDIA Corporation\\NVSMI");
}

#[cfg(target_os = "linux")]
fn main() {
    let paths = std::fs::read_dir("/usr/lib").unwrap();

    for path in paths {
        let entry = path.unwrap().path();
        if entry.is_dir() {
            let entry_string = entry.to_string_lossy();
            if entry_string.contains("nvidia") && !entry_string.ends_with("prime") {
                println!("cargo:rustc-link-search=native={}", entry_string);
                break;
            }
        }
    }

    println!("cargo:rustc-link-lib=nvidia-ml");
}

#[cfg(target_os = "macos")]
fn main() {
    compile_error!("NVML is not supported on macOS and therefore this crate cannot compile.");
}
<|endoftext|>


<commit_before>use crate::{ChainSpec, Epoch};
use serde_derive::{Deserialize, Serialize};
use std::convert::TryFrom;
use std::fmt::{self, Display, Formatter};
use std::str::FromStr;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(try_from = "String")]
#[serde(into = "String")]
pub enum ForkName {
    Base,
    Altair,
}

impl ForkName {
    pub fn list_all() -> Vec<ForkName> {
        vec![ForkName::Base, ForkName::Altair]
    }

    /// Set the activation slots in the given `ChainSpec` so that the fork named by `self`
    /// is the only fork in effect from genesis.
    pub fn make_genesis_spec(&self, mut spec: ChainSpec) -> ChainSpec {
        // Assumes GENESIS_EPOCH = 0, which is safe because it's a constant.
        match self {
            ForkName::Base => {
                spec.altair_fork_epoch = None;
                spec
            }
            ForkName::Altair => {
                spec.altair_fork_epoch = Some(Epoch::new(0));
                spec
            }
        }
    }

    /// Return the name of the fork immediately prior to the current one.
    ///
    /// If `self` is `ForkName::Base` then `Base` is returned.
    pub fn previous_fork(self) -> Option<ForkName> {
        match self {
            ForkName::Base => None,
            ForkName::Altair => Some(ForkName::Base),
        }
    }

    /// Return the name of the fork immediately after the current one.
    ///
    /// If `self` is the last known fork and has no successor, `None` is returned.
    pub fn next_fork(self) -> Option<ForkName> {
        match self {
            ForkName::Base => Some(ForkName::Altair),
            ForkName::Altair => None,
        }
    }
}

impl FromStr for ForkName {
    type Err = ();

    fn from_str(fork_name: &str) -> Result<Self, ()> {
        Ok(match fork_name {
            "phase0" | "base" => ForkName::Base,
            "altair" => ForkName::Altair,
            _ => return Err(()),
        })
    }
}

impl Display for ForkName {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {
        match self {
            ForkName::Base => "phase0".fmt(f),
            ForkName::Altair => "altair".fmt(f),
        }
    }
}

impl From<ForkName> for String {
    fn from(fork: ForkName) -> String {
        fork.to_string()
    }
}

impl TryFrom<String> for ForkName {
    type Error = String;

    fn try_from(s: String) -> Result<Self, Self::Error> {
        Self::from_str(&s).map_err(|()| format!("Invalid fork name: {}", s))
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct InconsistentFork {
    pub fork_at_slot: ForkName,
    pub object_fork: ForkName,
}

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

    #[test]
    fn previous_and_next_fork_consistent() {
        assert_eq!(ForkName::Altair.next_fork(), None);
        assert_eq!(ForkName::Base.previous_fork(), None);

        for (prev_fork, fork) in ForkName::list_all().into_iter().tuple_windows() {
            assert_eq!(prev_fork.next_fork(), Some(fork));
            assert_eq!(fork.previous_fork(), Some(prev_fork));
        }
    }
}
<commit_msg>Case insensitive match for ForkName (#2587)<commit_after>use crate::{ChainSpec, Epoch};
use serde_derive::{Deserialize, Serialize};
use std::convert::TryFrom;
use std::fmt::{self, Display, Formatter};
use std::str::FromStr;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(try_from = "String")]
#[serde(into = "String")]
pub enum ForkName {
    Base,
    Altair,
}

impl ForkName {
    pub fn list_all() -> Vec<ForkName> {
        vec![ForkName::Base, ForkName::Altair]
    }

    /// Set the activation slots in the given `ChainSpec` so that the fork named by `self`
    /// is the only fork in effect from genesis.
    pub fn make_genesis_spec(&self, mut spec: ChainSpec) -> ChainSpec {
        // Assumes GENESIS_EPOCH = 0, which is safe because it's a constant.
        match self {
            ForkName::Base => {
                spec.altair_fork_epoch = None;
                spec
            }
            ForkName::Altair => {
                spec.altair_fork_epoch = Some(Epoch::new(0));
                spec
            }
        }
    }

    /// Return the name of the fork immediately prior to the current one.
    ///
    /// If `self` is `ForkName::Base` then `Base` is returned.
    pub fn previous_fork(self) -> Option<ForkName> {
        match self {
            ForkName::Base => None,
            ForkName::Altair => Some(ForkName::Base),
        }
    }

    /// Return the name of the fork immediately after the current one.
    ///
    /// If `self` is the last known fork and has no successor, `None` is returned.
    pub fn next_fork(self) -> Option<ForkName> {
        match self {
            ForkName::Base => Some(ForkName::Altair),
            ForkName::Altair => None,
        }
    }
}

impl FromStr for ForkName {
    type Err = ();

    fn from_str(fork_name: &str) -> Result<Self, ()> {
        Ok(match fork_name.to_lowercase().as_ref() {
            "phase0" | "base" => ForkName::Base,
            "altair" => ForkName::Altair,
            _ => return Err(()),
        })
    }
}

impl Display for ForkName {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {
        match self {
            ForkName::Base => "phase0".fmt(f),
            ForkName::Altair => "altair".fmt(f),
        }
    }
}

impl From<ForkName> for String {
    fn from(fork: ForkName) -> String {
        fork.to_string()
    }
}

impl TryFrom<String> for ForkName {
    type Error = String;

    fn try_from(s: String) -> Result<Self, Self::Error> {
        Self::from_str(&s).map_err(|()| format!("Invalid fork name: {}", s))
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct InconsistentFork {
    pub fork_at_slot: ForkName,
    pub object_fork: ForkName,
}

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

    #[test]
    fn previous_and_next_fork_consistent() {
        assert_eq!(ForkName::Altair.next_fork(), None);
        assert_eq!(ForkName::Base.previous_fork(), None);

        for (prev_fork, fork) in ForkName::list_all().into_iter().tuple_windows() {
            assert_eq!(prev_fork.next_fork(), Some(fork));
            assert_eq!(fork.previous_fork(), Some(prev_fork));
        }
    }

    #[test]
    fn fork_name_case_insensitive_match() {
        assert_eq!(ForkName::from_str("BASE"), Ok(ForkName::Base));
        assert_eq!(ForkName::from_str("BaSe"), Ok(ForkName::Base));
        assert_eq!(ForkName::from_str("base"), Ok(ForkName::Base));

        assert_eq!(ForkName::from_str("PHASE0"), Ok(ForkName::Base));
        assert_eq!(ForkName::from_str("PhAsE0"), Ok(ForkName::Base));
        assert_eq!(ForkName::from_str("phase0"), Ok(ForkName::Base));

        assert_eq!(ForkName::from_str("ALTAIR"), Ok(ForkName::Altair));
        assert_eq!(ForkName::from_str("AlTaIr"), Ok(ForkName::Altair));
        assert_eq!(ForkName::from_str("altair"), Ok(ForkName::Altair));

        assert_eq!(ForkName::from_str("NO_NAME"), Err(()));
        assert_eq!(ForkName::from_str("no_name"), Err(()));
    }
}
<|endoftext|>