<commit_before>use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::{env, str};
use std::error::Error;
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use toml;

/// A Crate Dependency
pub type Dependency = (String, toml::Value);

/// Enumeration of errors which can occur when working with a rust manifest.
quick_error! {
    #[derive(Debug)]
    pub enum ManifestError {
        /// Cargo.toml could not be found.
        MissingManifest {
            description("missing manifest")
            display("Your Cargo.toml is missing.")
        }
        /// The TOML table could not be found.
        NonExistentTable {
            description("non existent table")
            display("The table could not be found.")
        }
        /// The dependency could not be found.
        NonExistentDependency {
            description("non existent dependency")
            display("The dependency could not be found.")
        }
    }
}

enum CargoFile {
    Config,
    Lock,
}

/// A Cargo Manifest
#[derive(Debug, PartialEq)]
pub struct Manifest {
    /// Manifest contents as TOML data
    pub data: toml::Table,
}

/// If a manifest is specified, return that one, otherise perform a manifest search starting from
/// the current directory.
/// If a manifest is specified, return that one. If a path is specified, perform a manifest search
/// starting from there. If nothing is specified, start searching from the current directory
/// (`cwd`).
fn find(specified: &Option<&str>, file: CargoFile) -> Result<PathBuf, Box<Error>> {
    let file_path = specified.map(PathBuf::from);

    if let Some(path) = file_path {
        if try!(fs::metadata(&path)).is_file() {
            Ok(path)
        } else {
            search(&path, file).map_err(From::from)
        }
    } else {
        env::current_dir()
            .map_err(From::from)
            .and_then(|ref dir| search(dir, file).map_err(From::from))
    }
}

/// Search for Cargo.toml in this directory and recursively up the tree until one is found.
fn search(dir: &Path, file: CargoFile) -> Result<PathBuf, ManifestError> {
    let manifest = match file {
        CargoFile::Config => dir.join("Cargo.toml"),
        CargoFile::Lock => dir.join("Cargo.lock"),
    };

    fs::metadata(&manifest)
        .map(|_| manifest)
        .or(dir.parent().ok_or(ManifestError::MissingManifest).and_then(|dir| search(dir, file)))
}

impl Manifest {
    /// Look for a `Cargo.toml` file
    ///
    /// Starts at the given path an goes into its parent directories until the manifest file is
    /// found. If no path is given, the process's working directory is used as a starting point.
    pub fn find_file(path: &Option<&str>) -> Result<File, Box<Error>> {
        find(path, CargoFile::Config).and_then(|path| {
            OpenOptions::new()
                .read(true)
                .write(true)
                .open(path)
                .map_err(From::from)
        })
    }

    /// Look for a `Cargo.lock` file
    ///
    /// Starts at the given path an goes into its parent directories until the manifest file is
    /// found. If no path is given, the process' working directory is used as a starting point.
    pub fn find_lock_file(path: &Option<&str>) -> Result<File, Box<Error>> {
        find(path, CargoFile::Lock).and_then(|path| {
            OpenOptions::new()
                .read(true)
                .write(true)
                .open(path)
                .map_err(From::from)
        })
    }

    /// Open the `Cargo.toml` for a path (or the process' `cwd`)
    pub fn open(path: &Option<&str>) -> Result<Manifest, Box<Error>> {
        let mut file = try!(Manifest::find_file(path));
        let mut data = String::new();
        try!(file.read_to_string(&mut data));

        data.parse()
    }

    /// Open the `Cargo.lock` for a path (or the process' `cwd`)
    pub fn open_lock_file(path: &Option<&str>) -> Result<Manifest, Box<Error>> {
        let mut file = try!(Manifest::find_lock_file(path));
        let mut data = String::new();
        try!(file.read_to_string(&mut data));

        data.parse()
    }

    /// Overwrite a file with TOML data.
    pub fn write_to_file<T: Seek + Write>(&self, file: &mut T) -> Result<(), Box<Error>> {
        try!(file.seek(SeekFrom::Start(0)));
        let mut toml = self.data.clone();

        let (proj_header, proj_data) = try!(toml.remove("package")
                                                .map(|data| ("package", data))
                                                .or_else(|| {
                                                    toml.remove("project")
                                                        .map(|data| ("project", data))
                                                })
                                                .ok_or(ManifestError::MissingManifest));
        write!(file,
               "[{}]\n{}{}",
               proj_header,
               proj_data,
               toml::Value::Table(toml))
            .map_err(From::from)
    }

    /// Add entry to a Cargo.toml.
    #[cfg_attr(feature = "dev", allow(toplevel_ref_arg))]
    pub fn insert_into_table(&mut self,
                             table: &str,
                             &(ref name, ref data): &Dependency)
                             -> Result<(), ManifestError> {
        let ref mut manifest = self.data;
        let entry = manifest.entry(String::from(table))
                            .or_insert(toml::Value::Table(BTreeMap::new()));
        match *entry {
            toml::Value::Table(ref mut deps) => {
                deps.insert(name.clone(), data.clone());
                Ok(())
            }
            _ => Err(ManifestError::NonExistentTable),
        }
    }

    /// Remove entry from a Cargo.toml.
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate cargo_edit;
    /// # extern crate toml;
    /// # fn main() {
    ///     use cargo_edit::Manifest;
    ///     use toml;
    ///
    ///     let mut manifest = Manifest { data: toml::Table::new() };
    ///     let dep = ("cargo-edit".to_owned(), toml::Value::String("0.1.0".to_owned()));
    ///     let _ = manifest.insert_into_table("dependencies", &dep);
    ///     assert!(manifest.remove_from_table("dependencies", &dep.0).is_ok());
    ///     assert!(manifest.remove_from_table("dependencies", &dep.0).is_err());
    /// # }
    /// ```
    #[cfg_attr(feature = "dev", allow(toplevel_ref_arg))]
    pub fn remove_from_table(&mut self,
                             table: &str,
                             name: &str)
                             -> Result<(), ManifestError> {
        let ref mut manifest = self.data;
        let entry = manifest.entry(String::from(table));

        match entry {
            Entry::Vacant(_) => Err(ManifestError::NonExistentTable),
            Entry::Occupied(entry) => {
                match *entry.into_mut() {
                    toml::Value::Table(ref mut deps) => {
                        deps.remove(name).map(|_| ()).ok_or(ManifestError::NonExistentDependency)
                    }
                    _ => Err(ManifestError::NonExistentTable)
                }
            }
        }
    }

    /// Add multiple dependencies to manifest
    pub fn add_deps(&mut self, table: &str, deps: &[Dependency]) -> Result<(), Box<Error>> {
        deps.iter()
            .map(|dep| self.insert_into_table(table, &dep))
            .collect::<Result<Vec<_>, _>>()
            .map_err(From::from)
            .map(|_| ())
    }
}

impl str::FromStr for Manifest {
    type Err = Box<Error>;

    /// Read manifest data from string
    fn from_str(input: &str) -> Result<Self, Self::Err> {
        let mut parser = toml::Parser::new(&input);

        parser.parse()
              .ok_or(parser.errors.pop())
              .map_err(Option::unwrap)
              .map_err(From::from)
              .map(|data| Manifest { data: data })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::BTreeMap;
    use toml;

    #[test]
    fn add_remove_dependency() {
        let mut manifest = Manifest { data: toml::Table::new() };
        // Create a copy containing empty "dependencies" table because removing
        //   the last entry in a table does not remove the section.
        let mut copy = Manifest { data: toml::Table::new() };
        copy.data.insert("dependencies".to_owned(), toml::Value::Table(BTreeMap::new()));
        let dep = ("cargo-edit".to_owned(), toml::Value::String("0.1.0".to_owned()));
        let _ = manifest.insert_into_table("dependencies", &dep);
        assert!(manifest.remove_from_table("dependencies", &dep.0).is_ok());
        assert_eq!(manifest, copy);
    }

    #[test]
    fn remove_dependency_no_section() {
        let mut manifest = Manifest { data: toml::Table::new() };
        let dep = ("cargo-edit".to_owned(), toml::Value::String("0.1.0".to_owned()));
        assert!(manifest.remove_from_table("dependencies", &dep.0).is_err());
    }

    #[test]
    fn remove_dependency_non_existent() {
        let mut manifest = Manifest { data: toml::Table::new() };
        let dep = ("cargo-edit".to_owned(), toml::Value::String("0.1.0".to_owned()));
        let other_dep = ("other-dep".to_owned(), toml::Value::String("0.1.0".to_owned()));
        let _ = manifest.insert_into_table("dependencies", &other_dep);
        assert!(manifest.remove_from_table("dependencies", &dep.0).is_err());

    }
}
<commit_msg>Nice Manifest Errors<commit_after>use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::{env, str};
use std::error::Error;
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use toml;

/// A Crate Dependency
pub type Dependency = (String, toml::Value);

/// Enumeration of errors which can occur when working with a rust manifest.
quick_error! {
    #[derive(Debug)]
    pub enum ManifestError {
        /// Cargo.toml could not be found.
        MissingManifest {
            description("missing manifest")
            display("Your Cargo.toml is missing.")
        }
        /// The TOML table could not be found.
        NonExistentTable(name: String) {
            description("non existent table")
            display("The table `{}` could not be found.", name)
        }
        /// The dependency could not be found.
        NonExistentDependency(name: String) {
            description("non existent dependency")
            display("The dependency `{}` could not be found.", name)
        }
    }
}

enum CargoFile {
    Config,
    Lock,
}

/// A Cargo Manifest
#[derive(Debug, PartialEq)]
pub struct Manifest {
    /// Manifest contents as TOML data
    pub data: toml::Table,
}

/// If a manifest is specified, return that one, otherise perform a manifest search starting from
/// the current directory.
/// If a manifest is specified, return that one. If a path is specified, perform a manifest search
/// starting from there. If nothing is specified, start searching from the current directory
/// (`cwd`).
fn find(specified: &Option<&str>, file: CargoFile) -> Result<PathBuf, Box<Error>> {
    let file_path = specified.map(PathBuf::from);

    if let Some(path) = file_path {
        if try!(fs::metadata(&path)).is_file() {
            Ok(path)
        } else {
            search(&path, file).map_err(From::from)
        }
    } else {
        env::current_dir()
            .map_err(From::from)
            .and_then(|ref dir| search(dir, file).map_err(From::from))
    }
}

/// Search for Cargo.toml in this directory and recursively up the tree until one is found.
fn search(dir: &Path, file: CargoFile) -> Result<PathBuf, ManifestError> {
    let manifest = match file {
        CargoFile::Config => dir.join("Cargo.toml"),
        CargoFile::Lock => dir.join("Cargo.lock"),
    };

    fs::metadata(&manifest)
        .map(|_| manifest)
        .or(dir.parent().ok_or(ManifestError::MissingManifest).and_then(|dir| search(dir, file)))
}

impl Manifest {
    /// Look for a `Cargo.toml` file
    ///
    /// Starts at the given path an goes into its parent directories until the manifest file is
    /// found. If no path is given, the process's working directory is used as a starting point.
    pub fn find_file(path: &Option<&str>) -> Result<File, Box<Error>> {
        find(path, CargoFile::Config).and_then(|path| {
            OpenOptions::new()
                .read(true)
                .write(true)
                .open(path)
                .map_err(From::from)
        })
    }

    /// Look for a `Cargo.lock` file
    ///
    /// Starts at the given path an goes into its parent directories until the manifest file is
    /// found. If no path is given, the process' working directory is used as a starting point.
    pub fn find_lock_file(path: &Option<&str>) -> Result<File, Box<Error>> {
        find(path, CargoFile::Lock).and_then(|path| {
            OpenOptions::new()
                .read(true)
                .write(true)
                .open(path)
                .map_err(From::from)
        })
    }

    /// Open the `Cargo.toml` for a path (or the process' `cwd`)
    pub fn open(path: &Option<&str>) -> Result<Manifest, Box<Error>> {
        let mut file = try!(Manifest::find_file(path));
        let mut data = String::new();
        try!(file.read_to_string(&mut data));

        data.parse()
    }

    /// Open the `Cargo.lock` for a path (or the process' `cwd`)
    pub fn open_lock_file(path: &Option<&str>) -> Result<Manifest, Box<Error>> {
        let mut file = try!(Manifest::find_lock_file(path));
        let mut data = String::new();
        try!(file.read_to_string(&mut data));

        data.parse()
    }

    /// Overwrite a file with TOML data.
    pub fn write_to_file<T: Seek + Write>(&self, file: &mut T) -> Result<(), Box<Error>> {
        try!(file.seek(SeekFrom::Start(0)));
        let mut toml = self.data.clone();

        let (proj_header, proj_data) = try!(toml.remove("package")
                                                .map(|data| ("package", data))
                                                .or_else(|| {
                                                    toml.remove("project")
                                                        .map(|data| ("project", data))
                                                })
                                                .ok_or(ManifestError::MissingManifest));
        write!(file,
               "[{}]\n{}{}",
               proj_header,
               proj_data,
               toml::Value::Table(toml))
            .map_err(From::from)
    }

    /// Add entry to a Cargo.toml.
    #[cfg_attr(feature = "dev", allow(toplevel_ref_arg))]
    pub fn insert_into_table(&mut self,
                             table: &str,
                             &(ref name, ref data): &Dependency)
                             -> Result<(), ManifestError> {
        let ref mut manifest = self.data;
        let entry = manifest.entry(String::from(table))
                            .or_insert(toml::Value::Table(BTreeMap::new()));
        match *entry {
            toml::Value::Table(ref mut deps) => {
                deps.insert(name.clone(), data.clone());
                Ok(())
            }
            _ => Err(ManifestError::NonExistentTable(table.into())),
        }
    }

    /// Remove entry from a Cargo.toml.
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate cargo_edit;
    /// # extern crate toml;
    /// # fn main() {
    ///     use cargo_edit::Manifest;
    ///     use toml;
    ///
    ///     let mut manifest = Manifest { data: toml::Table::new() };
    ///     let dep = ("cargo-edit".to_owned(), toml::Value::String("0.1.0".to_owned()));
    ///     let _ = manifest.insert_into_table("dependencies", &dep);
    ///     assert!(manifest.remove_from_table("dependencies", &dep.0).is_ok());
    ///     assert!(manifest.remove_from_table("dependencies", &dep.0).is_err());
    /// # }
    /// ```
    pub fn remove_from_table(&mut self,
                             table: &str,
                             name: &str)
                             -> Result<(), ManifestError> {
        let ref mut manifest = self.data;
        let entry = manifest.entry(String::from(table));

        match entry {
            Entry::Vacant(_) => Err(ManifestError::NonExistentTable(table.into())),
            Entry::Occupied(entry) => {
                match *entry.into_mut() {
                    toml::Value::Table(ref mut deps) => {
                        deps.remove(name).map(|_| ()).ok_or(ManifestError::NonExistentDependency(name.into()))
                    }
                    _ => Err(ManifestError::NonExistentTable(table.into()))
                }
            }
        }
    }

    /// Add multiple dependencies to manifest
    pub fn add_deps(&mut self, table: &str, deps: &[Dependency]) -> Result<(), Box<Error>> {
        deps.iter()
            .map(|dep| self.insert_into_table(table, &dep))
            .collect::<Result<Vec<_>, _>>()
            .map_err(From::from)
            .map(|_| ())
    }
}

impl str::FromStr for Manifest {
    type Err = Box<Error>;

    /// Read manifest data from string
    fn from_str(input: &str) -> Result<Self, Self::Err> {
        let mut parser = toml::Parser::new(&input);

        parser.parse()
              .ok_or(parser.errors.pop())
              .map_err(Option::unwrap)
              .map_err(From::from)
              .map(|data| Manifest { data: data })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::BTreeMap;
    use toml;

    #[test]
    fn add_remove_dependency() {
        let mut manifest = Manifest { data: toml::Table::new() };
        // Create a copy containing empty "dependencies" table because removing
        //   the last entry in a table does not remove the section.
        let mut copy = Manifest { data: toml::Table::new() };
        copy.data.insert("dependencies".to_owned(), toml::Value::Table(BTreeMap::new()));
        let dep = ("cargo-edit".to_owned(), toml::Value::String("0.1.0".to_owned()));
        let _ = manifest.insert_into_table("dependencies", &dep);
        assert!(manifest.remove_from_table("dependencies", &dep.0).is_ok());
        assert_eq!(manifest, copy);
    }

    #[test]
    fn remove_dependency_no_section() {
        let mut manifest = Manifest { data: toml::Table::new() };
        let dep = ("cargo-edit".to_owned(), toml::Value::String("0.1.0".to_owned()));
        assert!(manifest.remove_from_table("dependencies", &dep.0).is_err());
    }

    #[test]
    fn remove_dependency_non_existent() {
        let mut manifest = Manifest { data: toml::Table::new() };
        let dep = ("cargo-edit".to_owned(), toml::Value::String("0.1.0".to_owned()));
        let other_dep = ("other-dep".to_owned(), toml::Value::String("0.1.0".to_owned()));
        let _ = manifest.insert_into_table("dependencies", &other_dep);
        assert!(manifest.remove_from_table("dependencies", &dep.0).is_err());

    }
}
<|endoftext|>


<commit_before>//! This module defines the `Bits`, `Dibits`, and `Tribits` iterators as well as the
//! wrapper types `Bit`, `Dibit`, and `Tribit`, for working with sub-byte values.
//!
//! The wrapped values of `Bit`, `Dibit`, and `Tribit, are guaranteed to have only one,
//! two, or three bits, respectively.

use std;

/// Iterate over individual bits of a byte source, MSB to LSB.
pub type Bits<T> = SubByteIter<BitParams, T>;
/// Iterate over the dibits of a byte source, MSB to LSB.
pub type Dibits<T> = SubByteIter<DibitParams, T>;

/// Defines parameters needed for (power of two) sub-byte iterators.
pub trait IterParams {
    /// Type to yield at each iteration.
    type IterType;

    /// Number of bits to consume at each iteration.
    fn bits() -> u8;
    /// Wrap the given bits in container type.
    fn wrap(bits: u8) -> Self::IterType;

    /// Number of iterations needed for each byte.
    fn iterations() -> u8 { 8 / Self::bits() }

    /// Verify the parameters are supported.
    fn validate() {
        // Only powers of two are valid because there can be no "leftovers."
        assert!(Self::bits().is_power_of_two());
    }
}

/// A single bit.
#[derive(Copy, Clone, Default, Debug, PartialEq, Eq)]
pub struct Bit(u8);

impl Bit {
    /// Construct a new `Bit` with the given bit in the LSB position.
    pub fn new(bits: u8) -> Bit {
        assert!(bits >> 1 == 0);
        Bit(bits)
    }

    /// Get the wrapped bit value.
    pub fn bit(&self) -> u8 { self.0 }
}

/// Parameters for `Bits` iterator.
pub struct BitParams;

impl IterParams for BitParams {
    type IterType = Bit;
    fn bits() -> u8 { 1 }
    fn wrap(bits: u8) -> Bit { Bit::new(bits) }
}

/// Two bits.
#[derive(Copy, Clone, Default, Debug, PartialEq, Eq)]
pub struct Dibit(u8);

impl Dibit {
    /// Construct a new `Dibit` with the two given bits in the LSB position.
    pub fn new(bits: u8) -> Dibit {
        assert!(bits >> 2 == 0);
        Dibit(bits)
    }

    /// Get the wrapped dibit.
    pub fn bits(&self) -> u8 { self.0 }
}

/// Parameters for `Dibits` iterator.
pub struct DibitParams;

impl IterParams for DibitParams {
    type IterType = Dibit;
    fn bits() -> u8 { 2 }
    fn wrap(bits: u8) -> Dibit { Dibit::new(bits) }
}

/// Three bits.
#[derive(Copy, Clone, Default, Debug, PartialEq, Eq)]
pub struct Tribit(u8);

impl Tribit {
    /// Construct a new `Tribit` with the three given bits in the LSB position.
    pub fn new(bits: u8) -> Tribit {
        assert!(bits >> 3 == 0);
        Tribit(bits)
    }

    /// Get the wrapped tribit.
    pub fn bits(&self) -> u8 { self.0 }
}

/// An iterator for sub-byte (bit-level) values.
struct SubByteIter<P, T> where
    P: IterParams, T: Iterator<Item = u8>
{
    params: std::marker::PhantomData<P>,
    /// Source of bytes.
    src: T,
    /// Current bit-level index into the current byte.
    idx: u8,
    /// Current byte in the source.
    byte: u8,
}

impl<P, T> SubByteIter<P, T> where
    P: IterParams, T: Iterator<Item = u8>
{
    /// Construct a new `SubByteIter` over the given byte source. All bits are iterated
    /// over, so the number of bits must be a byte multiple.
    pub fn new(src: T) -> SubByteIter<P, T> {
        SubByteIter {
            params: std::marker::PhantomData,
            src: src,
            byte: 0,
            idx: 0,
        }
    }
}

impl<P, T> Iterator for SubByteIter<P, T> where
    P: IterParams, T: Iterator<Item = u8>
{
    type Item = P::IterType;

    fn next(&mut self) -> Option<Self::Item> {
        if self.idx == 0 {
            self.byte = match self.src.next() {
                Some(b) => b,
                None => return None,
            };
        }

        // Extract MSBs.
        let bits = self.byte >> (8 - P::bits());

        // Strip off the MSBs for the next iteration.
        self.byte <<= P::bits();

        // Move to the next item and reset after all have been visited.
        self.idx += 1;
        self.idx %= P::iterations();

        Some(P::wrap(bits))
    }
}

/// Iterates over the tribits in a byte source.
pub struct Tribits<T: Iterator<Item = u8>> {
    /// The source of bytes.
    src: T,
    /// Current bit buffer, containing either 6 or 3 bits (in MSB position) or 0.
    bits: u8,
    /// Tribit index into `bits` (0 or 1).
    idx: usize,
    /// Buffered bits from current source byte, to be added to `bits`.
    buf: u8,
    /// Number of bits in `buf`, either 2, 4, 6, or 0.
    buf_bits: usize,
}

impl<T: Iterator<Item = u8>> Tribits<T> {
    /// Construct a new `Tribits` from the given source of bytes. The number of bytes must
    /// be a multiple of 3 (a multiple of 24 bits).
    pub fn new(src: T) -> Tribits<T> {
        Tribits {
            src: src,
            bits: 0,
            idx: 0,
            buf: 0,
            buf_bits: 0,
        }
    }
}

impl<T: Iterator<Item = u8>> Iterator for Tribits<T> {
    type Item = Tribit;

    fn next(&mut self) -> Option<Self::Item> {
        // If on the first tribit, it's time to flush the buffer and (maybe) load another
        // byte.
        if self.idx == 0 {
            // Flush and reset the buffer.
            self.bits = self.buf;
            self.buf = 0;

            // Calculate the number of bits to buffer for the next iteration.
            self.buf_bits += 2;
            self.buf_bits %= 8;

            // Only load a new byte if bits need to be buffered.
            if self.buf_bits != 0 {
                let next = match self.src.next() {
                    Some(b) => b,
                    None => if self.buf_bits == 2 {
                        // In this case we've covered 8 tribits = 24 bits = 3 bytes
                        // exactly, so it's fine if there are no more bytes.
                        return None;
                    } else {
                        panic!("incomplete tribit");
                    }
                };

                // Add in some source bits after the MSBs.
                self.bits |= next >> self.buf_bits << 2;
                // Buffer the rest of the bits.
                self.buf = next << (8 - self.buf_bits);
            }
        }

        // Extract the 3 MSBs and strip them off for next time.
        let bits = self.bits >> 5;
        self.bits <<= 3;

        self.idx += 1;
        self.idx %= 2;

        Some(Tribit::new(bits))
    }
}

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

    #[test]
    fn validate_params() {
        BitParams::validate();
        DibitParams::validate();
    }

    #[test]
    fn test_bits() {
        const BITS: &'static [u8] = &[
            0b00011011,
            0b11001100,
        ];

        {
            let mut d = Dibits::new(BITS.iter().cloned());

            assert!(d.next().unwrap().0 == 0b00);
            assert!(d.next().unwrap().0 == 0b01);
            assert!(d.next().unwrap().0 == 0b10);
            assert!(d.next().unwrap().0 == 0b11);
            assert!(d.next().unwrap().0 == 0b11);
            assert!(d.next().unwrap().0 == 0b00);
            assert!(d.next().unwrap().0 == 0b11);
            assert!(d.next().unwrap().0 == 0b00);
            assert!(d.next().is_none());
        }

        {
            let mut b = Bits::new(BITS.iter().cloned());

            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 0);
        }
    }

    #[test]
    fn test_tribits() {
        let bytes = [
            0b00101001,
            0b11001011,
            0b10111000,
            0b00101001,
            0b11001011,
            0b10111000,
        ];
        let mut t = Tribits::new(bytes.iter().cloned());

        assert_eq!(t.next().unwrap().bits(), 0b001);
        assert_eq!(t.next().unwrap().bits(), 0b010);
        assert_eq!(t.next().unwrap().bits(), 0b011);
        assert_eq!(t.next().unwrap().bits(), 0b100);
        assert_eq!(t.next().unwrap().bits(), 0b101);
        assert_eq!(t.next().unwrap().bits(), 0b110);
        assert_eq!(t.next().unwrap().bits(), 0b111);
        assert_eq!(t.next().unwrap().bits(), 0b000);
        assert_eq!(t.next().unwrap().bits(), 0b001);
        assert_eq!(t.next().unwrap().bits(), 0b010);
        assert_eq!(t.next().unwrap().bits(), 0b011);
        assert_eq!(t.next().unwrap().bits(), 0b100);
        assert_eq!(t.next().unwrap().bits(), 0b101);
        assert_eq!(t.next().unwrap().bits(), 0b110);
        assert_eq!(t.next().unwrap().bits(), 0b111);
        assert_eq!(t.next().unwrap().bits(), 0b000);
        assert!(t.next().is_none());
    }

    #[test]
    #[should_panic]
    fn test_tribits_panic() {
        let bytes = [1, 2, 3, 4];
        let t = Tribits::new(bytes.iter().cloned());

        for _ in t {}
    }
}
<commit_msg>Greatly simplify Tribits iterator<commit_after>//! This module defines the `Bits`, `Dibits`, and `Tribits` iterators as well as the
//! wrapper types `Bit`, `Dibit`, and `Tribit`, for working with sub-byte values.
//!
//! The wrapped values of `Bit`, `Dibit`, and `Tribit, are guaranteed to have only one,
//! two, or three bits, respectively.

use std;

/// Iterate over individual bits of a byte source, MSB to LSB.
pub type Bits<T> = SubByteIter<BitParams, T>;
/// Iterate over the dibits of a byte source, MSB to LSB.
pub type Dibits<T> = SubByteIter<DibitParams, T>;
/// Iterates over the tribits in a byte source, MSB to LSB.
pub type Tribits<T> = SubByteIter<TribitParams, T>;

/// Defines parameters needed for (power of two) sub-byte iterators.
pub trait IterParams {
    /// Type to yield at each iteration.
    type IterType;

    /// Number of bits to consume at each iteration.
    fn bits() -> usize;

    /// Number of bytes to buffer, where the number of bits contained should be the lcm of
    /// 8 and the number of bits per iteration.
    fn buffer() -> usize { 1 }

    /// Amount to shift buffer after loading in bytes.
    fn buffer_shift() -> usize { 32 - 8 * Self::buffer() }

    /// Number of iterations needed for each byte.
    fn iterations() -> usize { 8 * Self::buffer() / Self::bits() }

    /// Wrap the given bits in container type.
    fn wrap(bits: u8) -> Self::IterType;

    /// Verify the parameters are supported.
    fn validate() {
        // Maximum buffer size is currently 32 bits.
        assert!(Self::buffer() <= 4);
    }
}

/// A single bit.
#[derive(Copy, Clone, Default, Debug, PartialEq, Eq)]
pub struct Bit(u8);

impl Bit {
    /// Construct a new `Bit` with the given bit in the LSB position.
    pub fn new(bits: u8) -> Bit {
        assert!(bits >> 1 == 0);
        Bit(bits)
    }

    /// Get the wrapped bit value.
    pub fn bit(&self) -> u8 { self.0 }
}

/// Parameters for `Bits` iterator.
pub struct BitParams;

impl IterParams for BitParams {
    type IterType = Bit;

    fn bits() -> usize { 1 }
    fn wrap(bits: u8) -> Bit { Bit::new(bits) }
}

/// Two bits.
#[derive(Copy, Clone, Default, Debug, PartialEq, Eq)]
pub struct Dibit(u8);

impl Dibit {
    /// Construct a new `Dibit` with the two given bits in the LSB position.
    pub fn new(bits: u8) -> Dibit {
        assert!(bits >> 2 == 0);
        Dibit(bits)
    }

    /// Get the wrapped dibit.
    pub fn bits(&self) -> u8 { self.0 }
}

/// Parameters for `Dibits` iterator.
pub struct DibitParams;

impl IterParams for DibitParams {
    type IterType = Dibit;

    fn bits() -> usize { 2 }
    fn wrap(bits: u8) -> Dibit { Dibit::new(bits) }
}

/// Three bits.
#[derive(Copy, Clone, Default, Debug, PartialEq, Eq)]
pub struct Tribit(u8);

impl Tribit {
    /// Construct a new `Tribit` with the three given bits in the LSB position.
    pub fn new(bits: u8) -> Tribit {
        assert!(bits >> 3 == 0);
        Tribit(bits)
    }

    /// Get the wrapped tribit.
    pub fn bits(&self) -> u8 { self.0 }
}

/// Parameters for `Tribits` iterator.
pub struct TribitParams;

impl IterParams for TribitParams {
    type IterType = Tribit;

    fn bits() -> usize { 3 }
    fn buffer() -> usize { 3 }
    fn wrap(bits: u8) -> Tribit { Tribit::new(bits) }
}

/// An iterator for sub-byte (bit-level) values.
struct SubByteIter<P, T> where
    P: IterParams, T: Iterator<Item = u8>
{
    params: std::marker::PhantomData<P>,
    /// Source of bytes.
    src: T,
    /// Current buffered bits.
    buf: u32,
    /// Current bit-level index into the current byte.
    idx: u8,
}

impl<P, T> SubByteIter<P, T> where
    P: IterParams, T: Iterator<Item = u8>
{
    /// Construct a new `SubByteIter` over the given byte source. All bits are iterated
    /// over, so the number of bits must be a byte multiple.
    pub fn new(src: T) -> SubByteIter<P, T> {
        SubByteIter {
            params: std::marker::PhantomData,
            src: src,
            buf: 0,
            idx: 0,
        }
    }

    /// Consume one or more bytes to create a buffer of bits, filled starting from the
    /// MSB.
    fn buffer(&mut self) -> Option<u32> {
        let (buf, added) = (&mut self.src)
            .take(P::buffer())
            .fold((0, 0), |(buf, added), byte| {
                (buf << 8 | byte as u32, added + 1)
            });

        // It's okay if there are no more source bits here, because we're on a safe
        // boundary.
        if added == 0 {
            return None;
        }

        assert!(added == P::buffer(), "incomplete source");

        Some(buf << P::buffer_shift())
    }
}

impl<P, T> Iterator for SubByteIter<P, T> where
    P: IterParams, T: Iterator<Item = u8>
{
    type Item = P::IterType;

    fn next(&mut self) -> Option<Self::Item> {
        if self.idx == 0 {
            self.buf = match self.buffer() {
                Some(b) => b,
                None => return None,
            };
        }

        // Extract MSBs.
        let bits = self.buf >> (32 - P::bits());

        // Strip off the MSBs for the next iteration.
        self.buf <<= P::bits();

        // Move to the next item and reset after all have been visited.
        self.idx += 1;
        self.idx %= P::iterations() as u8;

        Some(P::wrap(bits as u8))
    }
}

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

    #[test]
    fn validate_params() {
        BitParams::validate();
        DibitParams::validate();
        TribitParams::validate();
    }

    #[test]
    fn test_bits() {
        const BITS: &'static [u8] = &[
            0b00011011,
            0b11001100,
        ];

        {
            let mut d = Dibits::new(BITS.iter().cloned());

            assert!(d.next().unwrap().0 == 0b00);
            assert!(d.next().unwrap().0 == 0b01);
            assert!(d.next().unwrap().0 == 0b10);
            assert!(d.next().unwrap().0 == 0b11);
            assert!(d.next().unwrap().0 == 0b11);
            assert!(d.next().unwrap().0 == 0b00);
            assert!(d.next().unwrap().0 == 0b11);
            assert!(d.next().unwrap().0 == 0b00);
            assert!(d.next().is_none());
        }

        {
            let mut b = Bits::new(BITS.iter().cloned());

            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 1);
            assert!(b.next().unwrap().0 == 0);
            assert!(b.next().unwrap().0 == 0);
        }
    }

    #[test]
    fn test_tribits() {
        let bytes = [
            0b00101001,
            0b11001011,
            0b10111000,
            0b00101001,
            0b11001011,
            0b10111000,
        ];
        let mut t = Tribits::new(bytes.iter().cloned());

        assert_eq!(t.next().unwrap().bits(), 0b001);
        assert_eq!(t.next().unwrap().bits(), 0b010);
        assert_eq!(t.next().unwrap().bits(), 0b011);
        assert_eq!(t.next().unwrap().bits(), 0b100);
        assert_eq!(t.next().unwrap().bits(), 0b101);
        assert_eq!(t.next().unwrap().bits(), 0b110);
        assert_eq!(t.next().unwrap().bits(), 0b111);
        assert_eq!(t.next().unwrap().bits(), 0b000);
        assert_eq!(t.next().unwrap().bits(), 0b001);
        assert_eq!(t.next().unwrap().bits(), 0b010);
        assert_eq!(t.next().unwrap().bits(), 0b011);
        assert_eq!(t.next().unwrap().bits(), 0b100);
        assert_eq!(t.next().unwrap().bits(), 0b101);
        assert_eq!(t.next().unwrap().bits(), 0b110);
        assert_eq!(t.next().unwrap().bits(), 0b111);
        assert_eq!(t.next().unwrap().bits(), 0b000);
        assert!(t.next().is_none());
    }

    #[test]
    #[should_panic]
    fn test_tribits_panic() {
        let bytes = [1, 2, 3, 4];
        let t = Tribits::new(bytes.iter().cloned());

        for _ in t {}
    }
}
<|endoftext|>


<commit_before>use std::cell::RefCell;

use itertools::Itertools;
use hir::{source_binder, diagnostics::{Diagnostic as _, DiagnosticSink}};
use ra_db::SourceDatabase;
use ra_syntax::{
    Location, SourceFile, SyntaxKind, TextRange, SyntaxNode,
    ast::{self, AstNode},
};
use ra_text_edit::{TextEdit, TextEditBuilder};

use crate::{Diagnostic, FileId, FileSystemEdit, SourceChange, SourceFileEdit, db::RootDatabase};

#[derive(Debug, Copy, Clone)]
pub enum Severity {
    Error,
    WeakWarning,
}

pub(crate) fn diagnostics(db: &RootDatabase, file_id: FileId) -> Vec<Diagnostic> {
    let source_file = db.parse(file_id);
    let mut res = Vec::new();

    syntax_errors(&mut res, &source_file);

    for node in source_file.syntax().descendants() {
        check_unnecessary_braces_in_use_statement(&mut res, file_id, node);
        check_struct_shorthand_initialization(&mut res, file_id, node);
    }
    let res = RefCell::new(res);
    let mut sink = DiagnosticSink::new(|d| {
        res.borrow_mut().push(Diagnostic {
            message: d.message(),
            range: d.highlight_range(),
            severity: Severity::Error,
            fix: None,
        })
    })
    .on::<hir::diagnostics::UnresolvedModule, _>(|d| {
        let source_root = db.file_source_root(d.file().original_file(db));
        let create_file = FileSystemEdit::CreateFile { source_root, path: d.candidate.clone() };
        let fix = SourceChange {
            label: "create module".to_string(),
            source_file_edits: Vec::new(),
            file_system_edits: vec![create_file],
            cursor_position: None,
        };
        res.borrow_mut().push(Diagnostic {
            range: d.highlight_range(),
            message: d.message(),
            severity: Severity::Error,
            fix: Some(fix),
        })
    });
    if let Some(m) = source_binder::module_from_file_id(db, file_id) {
        m.diagnostics(db, &mut sink);
    };
    drop(sink);
    res.into_inner()
}

fn syntax_errors(acc: &mut Vec<Diagnostic>, source_file: &SourceFile) {
    fn location_to_range(location: Location) -> TextRange {
        match location {
            Location::Offset(offset) => TextRange::offset_len(offset, 1.into()),
            Location::Range(range) => range,
        }
    }

    acc.extend(source_file.errors().into_iter().map(|err| Diagnostic {
        range: location_to_range(err.location()),
        message: format!("Syntax Error: {}", err),
        severity: Severity::Error,
        fix: None,
    }));
}

fn check_unnecessary_braces_in_use_statement(
    acc: &mut Vec<Diagnostic>,
    file_id: FileId,
    node: &SyntaxNode,
) -> Option<()> {
    let use_tree_list = ast::UseTreeList::cast(node)?;
    if let Some((single_use_tree,)) = use_tree_list.use_trees().collect_tuple() {
        let range = use_tree_list.syntax().range();
        let edit =
            text_edit_for_remove_unnecessary_braces_with_self_in_use_statement(single_use_tree)
                .unwrap_or_else(|| {
                    let to_replace = single_use_tree.syntax().text().to_string();
                    let mut edit_builder = TextEditBuilder::default();
                    edit_builder.delete(range);
                    edit_builder.insert(range.start(), to_replace);
                    edit_builder.finish()
                });

        acc.push(Diagnostic {
            range,
            message: format!("Unnecessary braces in use statement"),
            severity: Severity::WeakWarning,
            fix: Some(SourceChange::source_file_edit(
                "Remove unnecessary braces",
                SourceFileEdit { file_id, edit },
            )),
        });
    }

    Some(())
}

fn text_edit_for_remove_unnecessary_braces_with_self_in_use_statement(
    single_use_tree: &ast::UseTree,
) -> Option<TextEdit> {
    let use_tree_list_node = single_use_tree.syntax().parent()?;
    if single_use_tree.path()?.segment()?.syntax().first_child()?.kind() == SyntaxKind::SELF_KW {
        let start = use_tree_list_node.prev_sibling()?.range().start();
        let end = use_tree_list_node.range().end();
        let range = TextRange::from_to(start, end);
        let mut edit_builder = TextEditBuilder::default();
        edit_builder.delete(range);
        return Some(edit_builder.finish());
    }
    None
}

fn check_struct_shorthand_initialization(
    acc: &mut Vec<Diagnostic>,
    file_id: FileId,
    node: &SyntaxNode,
) -> Option<()> {
    let struct_lit = ast::StructLit::cast(node)?;
    let named_field_list = struct_lit.named_field_list()?;
    for named_field in named_field_list.fields() {
        if let (Some(name_ref), Some(expr)) = (named_field.name_ref(), named_field.expr()) {
            let field_name = name_ref.syntax().text().to_string();
            let field_expr = expr.syntax().text().to_string();
            if field_name == field_expr {
                let mut edit_builder = TextEditBuilder::default();
                edit_builder.delete(named_field.syntax().range());
                edit_builder.insert(named_field.syntax().range().start(), field_name);
                let edit = edit_builder.finish();

                acc.push(Diagnostic {
                    range: named_field.syntax().range(),
                    message: format!("Shorthand struct initialization"),
                    severity: Severity::WeakWarning,
                    fix: Some(SourceChange::source_file_edit(
                        "use struct shorthand initialization",
                        SourceFileEdit { file_id, edit },
                    )),
                });
            }
        }
    }
    Some(())
}

#[cfg(test)]
mod tests {
    use test_utils::assert_eq_text;

    use super::*;

    type DiagnosticChecker = fn(&mut Vec<Diagnostic>, FileId, &SyntaxNode) -> Option<()>;

    fn check_not_applicable(code: &str, func: DiagnosticChecker) {
        let file = SourceFile::parse(code);
        let mut diagnostics = Vec::new();
        for node in file.syntax().descendants() {
            func(&mut diagnostics, FileId(0), node);
        }
        assert!(diagnostics.is_empty());
    }

    fn check_apply(before: &str, after: &str, func: DiagnosticChecker) {
        let file = SourceFile::parse(before);
        let mut diagnostics = Vec::new();
        for node in file.syntax().descendants() {
            func(&mut diagnostics, FileId(0), node);
        }
        let diagnostic =
            diagnostics.pop().unwrap_or_else(|| panic!("no diagnostics for:\n{}\n", before));
        let mut fix = diagnostic.fix.unwrap();
        let edit = fix.source_file_edits.pop().unwrap().edit;
        let actual = edit.apply(&before);
        assert_eq_text!(after, &actual);
    }

    #[test]
    fn test_check_unnecessary_braces_in_use_statement() {
        check_not_applicable(
            "
            use a;
            use a::{c, d::e};
        ",
            check_unnecessary_braces_in_use_statement,
        );
        check_apply("use {b};", "use b;", check_unnecessary_braces_in_use_statement);
        check_apply("use a::{c};", "use a::c;", check_unnecessary_braces_in_use_statement);
        check_apply("use a::{self};", "use a;", check_unnecessary_braces_in_use_statement);
        check_apply(
            "use a::{c, d::{e}};",
            "use a::{c, d::e};",
            check_unnecessary_braces_in_use_statement,
        );
    }

    #[test]
    fn test_check_struct_shorthand_initialization() {
        check_not_applicable(
            r#"
            struct A {
                a: &'static str
            }

            fn main() {
                A {
                    a: "hello"
                }
            }
        "#,
            check_struct_shorthand_initialization,
        );

        check_apply(
            r#"
struct A {
    a: &'static str
}

fn main() {
    let a = "haha";
    A {
        a: a
    }
}
        "#,
            r#"
struct A {
    a: &'static str
}

fn main() {
    let a = "haha";
    A {
        a
    }
}
        "#,
            check_struct_shorthand_initialization,
        );

        check_apply(
            r#"
struct A {
    a: &'static str,
    b: &'static str
}

fn main() {
    let a = "haha";
    let b = "bb";
    A {
        a: a,
        b
    }
}
        "#,
            r#"
struct A {
    a: &'static str,
    b: &'static str
}

fn main() {
    let a = "haha";
    let b = "bb";
    A {
        a,
        b
    }
}
        "#,
            check_struct_shorthand_initialization,
        );
    }
}
<commit_msg>simplify<commit_after>use std::cell::RefCell;

use itertools::Itertools;
use hir::{source_binder, diagnostics::{Diagnostic as _, DiagnosticSink}};
use ra_db::SourceDatabase;
use ra_syntax::{
    Location, SourceFile, SyntaxKind, TextRange, SyntaxNode,
    ast::{self, AstNode},
};
use ra_text_edit::{TextEdit, TextEditBuilder};

use crate::{Diagnostic, FileId, FileSystemEdit, SourceChange, SourceFileEdit, db::RootDatabase};

#[derive(Debug, Copy, Clone)]
pub enum Severity {
    Error,
    WeakWarning,
}

pub(crate) fn diagnostics(db: &RootDatabase, file_id: FileId) -> Vec<Diagnostic> {
    let source_file = db.parse(file_id);
    let mut res = Vec::new();

    syntax_errors(&mut res, &source_file);

    for node in source_file.syntax().descendants() {
        check_unnecessary_braces_in_use_statement(&mut res, file_id, node);
        check_struct_shorthand_initialization(&mut res, file_id, node);
    }
    let res = RefCell::new(res);
    let mut sink = DiagnosticSink::new(|d| {
        res.borrow_mut().push(Diagnostic {
            message: d.message(),
            range: d.highlight_range(),
            severity: Severity::Error,
            fix: None,
        })
    })
    .on::<hir::diagnostics::UnresolvedModule, _>(|d| {
        let source_root = db.file_source_root(d.file().original_file(db));
        let create_file = FileSystemEdit::CreateFile { source_root, path: d.candidate.clone() };
        let fix = SourceChange::file_system_edit("create module", create_file);
        res.borrow_mut().push(Diagnostic {
            range: d.highlight_range(),
            message: d.message(),
            severity: Severity::Error,
            fix: Some(fix),
        })
    });
    if let Some(m) = source_binder::module_from_file_id(db, file_id) {
        m.diagnostics(db, &mut sink);
    };
    drop(sink);
    res.into_inner()
}

fn syntax_errors(acc: &mut Vec<Diagnostic>, source_file: &SourceFile) {
    fn location_to_range(location: Location) -> TextRange {
        match location {
            Location::Offset(offset) => TextRange::offset_len(offset, 1.into()),
            Location::Range(range) => range,
        }
    }

    acc.extend(source_file.errors().into_iter().map(|err| Diagnostic {
        range: location_to_range(err.location()),
        message: format!("Syntax Error: {}", err),
        severity: Severity::Error,
        fix: None,
    }));
}

fn check_unnecessary_braces_in_use_statement(
    acc: &mut Vec<Diagnostic>,
    file_id: FileId,
    node: &SyntaxNode,
) -> Option<()> {
    let use_tree_list = ast::UseTreeList::cast(node)?;
    if let Some((single_use_tree,)) = use_tree_list.use_trees().collect_tuple() {
        let range = use_tree_list.syntax().range();
        let edit =
            text_edit_for_remove_unnecessary_braces_with_self_in_use_statement(single_use_tree)
                .unwrap_or_else(|| {
                    let to_replace = single_use_tree.syntax().text().to_string();
                    let mut edit_builder = TextEditBuilder::default();
                    edit_builder.delete(range);
                    edit_builder.insert(range.start(), to_replace);
                    edit_builder.finish()
                });

        acc.push(Diagnostic {
            range,
            message: format!("Unnecessary braces in use statement"),
            severity: Severity::WeakWarning,
            fix: Some(SourceChange::source_file_edit(
                "Remove unnecessary braces",
                SourceFileEdit { file_id, edit },
            )),
        });
    }

    Some(())
}

fn text_edit_for_remove_unnecessary_braces_with_self_in_use_statement(
    single_use_tree: &ast::UseTree,
) -> Option<TextEdit> {
    let use_tree_list_node = single_use_tree.syntax().parent()?;
    if single_use_tree.path()?.segment()?.syntax().first_child()?.kind() == SyntaxKind::SELF_KW {
        let start = use_tree_list_node.prev_sibling()?.range().start();
        let end = use_tree_list_node.range().end();
        let range = TextRange::from_to(start, end);
        let mut edit_builder = TextEditBuilder::default();
        edit_builder.delete(range);
        return Some(edit_builder.finish());
    }
    None
}

fn check_struct_shorthand_initialization(
    acc: &mut Vec<Diagnostic>,
    file_id: FileId,
    node: &SyntaxNode,
) -> Option<()> {
    let struct_lit = ast::StructLit::cast(node)?;
    let named_field_list = struct_lit.named_field_list()?;
    for named_field in named_field_list.fields() {
        if let (Some(name_ref), Some(expr)) = (named_field.name_ref(), named_field.expr()) {
            let field_name = name_ref.syntax().text().to_string();
            let field_expr = expr.syntax().text().to_string();
            if field_name == field_expr {
                let mut edit_builder = TextEditBuilder::default();
                edit_builder.delete(named_field.syntax().range());
                edit_builder.insert(named_field.syntax().range().start(), field_name);
                let edit = edit_builder.finish();

                acc.push(Diagnostic {
                    range: named_field.syntax().range(),
                    message: format!("Shorthand struct initialization"),
                    severity: Severity::WeakWarning,
                    fix: Some(SourceChange::source_file_edit(
                        "use struct shorthand initialization",
                        SourceFileEdit { file_id, edit },
                    )),
                });
            }
        }
    }
    Some(())
}

#[cfg(test)]
mod tests {
    use test_utils::assert_eq_text;

    use super::*;

    type DiagnosticChecker = fn(&mut Vec<Diagnostic>, FileId, &SyntaxNode) -> Option<()>;

    fn check_not_applicable(code: &str, func: DiagnosticChecker) {
        let file = SourceFile::parse(code);
        let mut diagnostics = Vec::new();
        for node in file.syntax().descendants() {
            func(&mut diagnostics, FileId(0), node);
        }
        assert!(diagnostics.is_empty());
    }

    fn check_apply(before: &str, after: &str, func: DiagnosticChecker) {
        let file = SourceFile::parse(before);
        let mut diagnostics = Vec::new();
        for node in file.syntax().descendants() {
            func(&mut diagnostics, FileId(0), node);
        }
        let diagnostic =
            diagnostics.pop().unwrap_or_else(|| panic!("no diagnostics for:\n{}\n", before));
        let mut fix = diagnostic.fix.unwrap();
        let edit = fix.source_file_edits.pop().unwrap().edit;
        let actual = edit.apply(&before);
        assert_eq_text!(after, &actual);
    }

    #[test]
    fn test_check_unnecessary_braces_in_use_statement() {
        check_not_applicable(
            "
            use a;
            use a::{c, d::e};
        ",
            check_unnecessary_braces_in_use_statement,
        );
        check_apply("use {b};", "use b;", check_unnecessary_braces_in_use_statement);
        check_apply("use a::{c};", "use a::c;", check_unnecessary_braces_in_use_statement);
        check_apply("use a::{self};", "use a;", check_unnecessary_braces_in_use_statement);
        check_apply(
            "use a::{c, d::{e}};",
            "use a::{c, d::e};",
            check_unnecessary_braces_in_use_statement,
        );
    }

    #[test]
    fn test_check_struct_shorthand_initialization() {
        check_not_applicable(
            r#"
            struct A {
                a: &'static str
            }

            fn main() {
                A {
                    a: "hello"
                }
            }
        "#,
            check_struct_shorthand_initialization,
        );

        check_apply(
            r#"
struct A {
    a: &'static str
}

fn main() {
    let a = "haha";
    A {
        a: a
    }
}
        "#,
            r#"
struct A {
    a: &'static str
}

fn main() {
    let a = "haha";
    A {
        a
    }
}
        "#,
            check_struct_shorthand_initialization,
        );

        check_apply(
            r#"
struct A {
    a: &'static str,
    b: &'static str
}

fn main() {
    let a = "haha";
    let b = "bb";
    A {
        a: a,
        b
    }
}
        "#,
            r#"
struct A {
    a: &'static str,
    b: &'static str
}

fn main() {
    let a = "haha";
    let b = "bb";
    A {
        a,
        b
    }
}
        "#,
            check_struct_shorthand_initialization,
        );
    }
}
<|endoftext|>


<commit_before>//! Driver for rust-analyzer.
//!
//! Based on cli flags, either spawns an LSP server, or runs a batch analysis
mod args;
mod logger;

use std::{convert::TryFrom, env, fs, path::PathBuf, process};

use lsp_server::Connection;
use project_model::ProjectManifest;
use rust_analyzer::{cli, config::Config, from_json, Result};
use vfs::AbsPathBuf;

#[cfg(all(feature = "mimalloc"))]
#[global_allocator]
static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;

#[cfg(all(feature = "jemalloc", not(target_env = "msvc")))]
#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

fn main() {
    if let Err(err) = try_main() {
        eprintln!("{}", err);
        process::exit(101);
    }
}

fn try_main() -> Result<()> {
    let args = args::Args::parse()?;
    setup_logging(args.log_file)?;
    match args.command {
        args::Command::RunServer => run_server()?,
        args::Command::PrintConfigSchema => {
            println!("{:#}", Config::json_schema());
        }
        args::Command::ProcMacro => proc_macro_srv::cli::run()?,

        args::Command::Parse { no_dump } => cli::parse(no_dump)?,
        args::Command::Symbols => cli::symbols()?,
        args::Command::Highlight { rainbow } => cli::highlight(rainbow)?,
        args::Command::AnalysisStats(cmd) => cmd.run(args.verbosity)?,
        args::Command::Bench(cmd) => cmd.run(args.verbosity)?,
        args::Command::Diagnostics { path, load_output_dirs, with_proc_macro } => {
            cli::diagnostics(path.as_ref(), load_output_dirs, with_proc_macro)?
        }
        args::Command::Ssr { rules } => {
            cli::apply_ssr_rules(rules)?;
        }
        args::Command::StructuredSearch { patterns, debug_snippet } => {
            cli::search_for_patterns(patterns, debug_snippet)?;
        }
        args::Command::Version => println!("rust-analyzer {}", env!("REV")),
        args::Command::Help => {}
    }
    Ok(())
}

fn setup_logging(log_file: Option<PathBuf>) -> Result<()> {
    env::set_var("RUST_BACKTRACE", "short");

    let log_file = match log_file {
        Some(path) => {
            if let Some(parent) = path.parent() {
                let _ = fs::create_dir_all(parent);
            }
            Some(fs::File::create(path)?)
        }
        None => None,
    };
    let filter = env::var("RA_LOG").ok();
    logger::Logger::new(log_file, filter.as_deref()).install();

    tracing_setup::setup_tracing()?;

    profile::init();

    if !cfg!(debug_assertions) {
        stdx::set_assert_hook(|loc, args| log::error!("assertion failed at {}: {}", loc, args));
    }

    Ok(())
}

mod tracing_setup {
    use tracing::subscriber;
    use tracing_subscriber::layer::SubscriberExt;
    use tracing_subscriber::EnvFilter;
    use tracing_subscriber::Registry;
    use tracing_tree::HierarchicalLayer;

    pub(crate) fn setup_tracing() -> super::Result<()> {
        let filter = EnvFilter::from_env("CHALK_DEBUG");
        let layer = HierarchicalLayer::default()
            .with_indent_lines(true)
            .with_ansi(false)
            .with_indent_amount(2)
            .with_writer(std::io::stderr);
        let subscriber = Registry::default().with(filter).with(layer);
        subscriber::set_global_default(subscriber)?;
        Ok(())
    }
}

fn run_server() -> Result<()> {
    log::info!("server will start");

    let (connection, io_threads) = Connection::stdio();

    let (initialize_id, initialize_params) = connection.initialize_start()?;
    log::info!("InitializeParams: {}", initialize_params);
    let initialize_params =
        from_json::<lsp_types::InitializeParams>("InitializeParams", initialize_params)?;

    let server_capabilities = rust_analyzer::server_capabilities(&initialize_params.capabilities);

    let initialize_result = lsp_types::InitializeResult {
        capabilities: server_capabilities,
        server_info: Some(lsp_types::ServerInfo {
            name: String::from("rust-analyzer"),
            version: Some(String::from(env!("REV"))),
        }),
    };

    let initialize_result = serde_json::to_value(initialize_result).unwrap();

    connection.initialize_finish(initialize_id, initialize_result)?;

    if let Some(client_info) = initialize_params.client_info {
        log::info!("Client '{}' {}", client_info.name, client_info.version.unwrap_or_default());
    }

    let config = {
        let root_path = match initialize_params
            .root_uri
            .and_then(|it| it.to_file_path().ok())
            .and_then(|it| AbsPathBuf::try_from(it).ok())
        {
            Some(it) => it,
            None => {
                let cwd = env::current_dir()?;
                AbsPathBuf::assert(cwd)
            }
        };

        let mut config = Config::new(root_path, initialize_params.capabilities);
        if let Some(json) = initialize_params.initialization_options {
            config.update(json);
        }

        if config.linked_projects().is_empty() {
            let workspace_roots = initialize_params
                .workspace_folders
                .map(|workspaces| {
                    workspaces
                        .into_iter()
                        .filter_map(|it| it.uri.to_file_path().ok())
                        .filter_map(|it| AbsPathBuf::try_from(it).ok())
                        .collect::<Vec<_>>()
                })
                .filter(|workspaces| !workspaces.is_empty())
                .unwrap_or_else(|| vec![config.root_path.clone()]);

            let discovered = ProjectManifest::discover_all(&workspace_roots);
            log::info!("discovered projects: {:?}", discovered);
            if discovered.is_empty() {
                log::error!("failed to find any projects in {:?}", workspace_roots);
            }

            config.discovered_projects = Some(discovered);
        }

        config
    };

    rust_analyzer::main_loop(config, connection)?;

    io_threads.join()?;
    log::info!("server did shut down");
    Ok(())
}
<commit_msg>Make assertion failures more annoying for maintainers<commit_after>//! Driver for rust-analyzer.
//!
//! Based on cli flags, either spawns an LSP server, or runs a batch analysis
mod args;
mod logger;

use std::{convert::TryFrom, env, fs, path::PathBuf, process};

use lsp_server::Connection;
use project_model::ProjectManifest;
use rust_analyzer::{cli, config::Config, from_json, Result};
use vfs::AbsPathBuf;

#[cfg(all(feature = "mimalloc"))]
#[global_allocator]
static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;

#[cfg(all(feature = "jemalloc", not(target_env = "msvc")))]
#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

fn main() {
    if let Err(err) = try_main() {
        eprintln!("{}", err);
        process::exit(101);
    }
}

fn try_main() -> Result<()> {
    let args = args::Args::parse()?;
    setup_logging(args.log_file)?;
    match args.command {
        args::Command::RunServer => run_server()?,
        args::Command::PrintConfigSchema => {
            println!("{:#}", Config::json_schema());
        }
        args::Command::ProcMacro => proc_macro_srv::cli::run()?,

        args::Command::Parse { no_dump } => cli::parse(no_dump)?,
        args::Command::Symbols => cli::symbols()?,
        args::Command::Highlight { rainbow } => cli::highlight(rainbow)?,
        args::Command::AnalysisStats(cmd) => cmd.run(args.verbosity)?,
        args::Command::Bench(cmd) => cmd.run(args.verbosity)?,
        args::Command::Diagnostics { path, load_output_dirs, with_proc_macro } => {
            cli::diagnostics(path.as_ref(), load_output_dirs, with_proc_macro)?
        }
        args::Command::Ssr { rules } => {
            cli::apply_ssr_rules(rules)?;
        }
        args::Command::StructuredSearch { patterns, debug_snippet } => {
            cli::search_for_patterns(patterns, debug_snippet)?;
        }
        args::Command::Version => println!("rust-analyzer {}", env!("REV")),
        args::Command::Help => {}
    }
    Ok(())
}

fn setup_logging(log_file: Option<PathBuf>) -> Result<()> {
    env::set_var("RUST_BACKTRACE", "short");

    let log_file = match log_file {
        Some(path) => {
            if let Some(parent) = path.parent() {
                let _ = fs::create_dir_all(parent);
            }
            Some(fs::File::create(path)?)
        }
        None => None,
    };
    let filter = env::var("RA_LOG").ok();
    logger::Logger::new(log_file, filter.as_deref()).install();

    tracing_setup::setup_tracing()?;

    profile::init();

    if !cfg!(debug_assertions) {
        stdx::set_assert_hook(|loc, args| {
            if env::var("RA_PROFILE").is_ok() {
                panic!("assertion failed at {}: {}", loc, args)
            }
            log::error!("assertion failed at {}: {}", loc, args)
        });
    }

    Ok(())
}

mod tracing_setup {
    use tracing::subscriber;
    use tracing_subscriber::layer::SubscriberExt;
    use tracing_subscriber::EnvFilter;
    use tracing_subscriber::Registry;
    use tracing_tree::HierarchicalLayer;

    pub(crate) fn setup_tracing() -> super::Result<()> {
        let filter = EnvFilter::from_env("CHALK_DEBUG");
        let layer = HierarchicalLayer::default()
            .with_indent_lines(true)
            .with_ansi(false)
            .with_indent_amount(2)
            .with_writer(std::io::stderr);
        let subscriber = Registry::default().with(filter).with(layer);
        subscriber::set_global_default(subscriber)?;
        Ok(())
    }
}

fn run_server() -> Result<()> {
    log::info!("server will start");

    let (connection, io_threads) = Connection::stdio();

    let (initialize_id, initialize_params) = connection.initialize_start()?;
    log::info!("InitializeParams: {}", initialize_params);
    let initialize_params =
        from_json::<lsp_types::InitializeParams>("InitializeParams", initialize_params)?;

    let server_capabilities = rust_analyzer::server_capabilities(&initialize_params.capabilities);

    let initialize_result = lsp_types::InitializeResult {
        capabilities: server_capabilities,
        server_info: Some(lsp_types::ServerInfo {
            name: String::from("rust-analyzer"),
            version: Some(String::from(env!("REV"))),
        }),
    };

    let initialize_result = serde_json::to_value(initialize_result).unwrap();

    connection.initialize_finish(initialize_id, initialize_result)?;

    if let Some(client_info) = initialize_params.client_info {
        log::info!("Client '{}' {}", client_info.name, client_info.version.unwrap_or_default());
    }

    let config = {
        let root_path = match initialize_params
            .root_uri
            .and_then(|it| it.to_file_path().ok())
            .and_then(|it| AbsPathBuf::try_from(it).ok())
        {
            Some(it) => it,
            None => {
                let cwd = env::current_dir()?;
                AbsPathBuf::assert(cwd)
            }
        };

        let mut config = Config::new(root_path, initialize_params.capabilities);
        if let Some(json) = initialize_params.initialization_options {
            config.update(json);
        }

        if config.linked_projects().is_empty() {
            let workspace_roots = initialize_params
                .workspace_folders
                .map(|workspaces| {
                    workspaces
                        .into_iter()
                        .filter_map(|it| it.uri.to_file_path().ok())
                        .filter_map(|it| AbsPathBuf::try_from(it).ok())
                        .collect::<Vec<_>>()
                })
                .filter(|workspaces| !workspaces.is_empty())
                .unwrap_or_else(|| vec![config.root_path.clone()]);

            let discovered = ProjectManifest::discover_all(&workspace_roots);
            log::info!("discovered projects: {:?}", discovered);
            if discovered.is_empty() {
                log::error!("failed to find any projects in {:?}", workspace_roots);
            }

            config.discovered_projects = Some(discovered);
        }

        config
    };

    rust_analyzer::main_loop(config, connection)?;

    io_threads.join()?;
    log::info!("server did shut down");
    Ok(())
}
<|endoftext|>


<commit_before>use crate::rotate::*;
use crate::spiral::Spiral;
use oppai_field::cell::Cell;
use oppai_field::field::{Field, Pos};
use oppai_field::player::Player;
use oppai_patterns::dfa::{Dfa, DfaState};
use rand::Rng;
use std::{
  cmp,
  collections::{HashMap, HashSet},
  fs::File,
  io::Read,
  str::FromStr,
};
use tar::Archive;

const PATTERNS_STR: &str = "patterns";

#[derive(Clone, Debug)]
struct Move {
  x: i32,
  y: i32,
  p: f64, // priority
}

pub struct Patterns {
  min_size: u32,
  dfa: Dfa<Move>,
}

impl Patterns {
  pub fn empty() -> Patterns {
    Patterns {
      min_size: u32::max_value(),
      dfa: Dfa::empty(),
    }
  }

  fn parse_header(string: &str) -> Result<(u32, u32, f64), &'static str> {
    let mut split = string.split_whitespace().fuse();
    let width = split
      .next()
      .and_then(|s| u32::from_str(s).ok())
      .ok_or("Expected width of type u32.")?;
    if width < 2 {
      return Err("Minimum allowed width is 2");
    }
    let height = split
      .next()
      .and_then(|s| u32::from_str(s).ok())
      .ok_or("Expected height of type u32.")?;
    if height < 2 {
      return Err("Minimum allowed height is 2");
    }
    let priority = split
      .next()
      .and_then(|s| f64::from_str(s).ok())
      .ok_or("Expected priority of type f64.")?;
    Ok((width, height, priority))
  }

  fn parse_move(string: &str) -> Result<Move, &'static str> {
    let mut split = string.split(' ').fuse();
    let x = split
      .next()
      .and_then(|s| u32::from_str(s).ok())
      .ok_or("Expected x coordinate of type u32.")?;
    let y = split
      .next()
      .and_then(|s| u32::from_str(s).ok())
      .ok_or("Expected y coordinate of type u32.")?;
    let p = split
      .next()
      .and_then(|s| f64::from_str(s).ok())
      .ok_or("Expected probability of type f64.")?;
    Ok(Move {
      x: x as i32,
      y: y as i32,
      p,
    })
  }

  fn covering_spiral_length(side_of_square: u32) -> u32 {
    let x = side_of_square / 2 + 1;
    let y = (1 - side_of_square % 2) * side_of_square * 2;
    (8 * x - 13) * x + 6 - y
  }

  fn build_dfa(width: u32, height: u32, moves: &[Move], rotation: u32, s: &str) -> Result<Dfa<Move>, &'static str> {
    let (rotated_width, rotated_height) = rotate_sizes(width, height, rotation);
    let center_x = (rotated_width - 1) / 2;
    let center_y = (rotated_height - 1) / 2;
    let spiral_length = Patterns::covering_spiral_length(cmp::max(width, height)) as usize;
    let mut states = Vec::with_capacity(spiral_length + 2);
    let fs = spiral_length; // "Found" state.
    let nfs = spiral_length + 1; // "Not found" state.
    let s_bytes = s.as_bytes();
    for (i, (shift_x, shift_y)) in Spiral::new().take(spiral_length).enumerate() {
      let nxt = i + 1;
      let rotated_x = center_x as i32 + shift_x;
      let rotated_y = center_y as i32 + shift_y;
      let state =
        if rotated_x >= 0 && rotated_x < rotated_width as i32 && rotated_y >= 0 && rotated_y < rotated_height as i32 {
          let (x, y) = rotate_back(
            rotated_width,
            rotated_height,
            rotated_x as u32,
            rotated_y as u32,
            rotation,
          );
          let pos = y * width + x;
          match s_bytes[pos as usize] as char {
            '.' | '+' => DfaState::new(nxt, nfs, nfs, nfs, false, Vec::with_capacity(0)),
            '?' => DfaState::new(nxt, nxt, nxt, nfs, false, Vec::with_capacity(0)),
            '*' => DfaState::new(nxt, nxt, nxt, nxt, false, Vec::with_capacity(0)),
            'X' => DfaState::new(nfs, nxt, nfs, nfs, false, Vec::with_capacity(0)),
            'O' => DfaState::new(nfs, nfs, nxt, nfs, false, Vec::with_capacity(0)),
            'x' => DfaState::new(nxt, nxt, nfs, nfs, false, Vec::with_capacity(0)),
            'o' => DfaState::new(nxt, nfs, nxt, nfs, false, Vec::with_capacity(0)),
            '#' => DfaState::new(nfs, nfs, nfs, nxt, false, Vec::with_capacity(0)),
            _ => return Err("Invalid character in the pattern."),
          }
        } else {
          DfaState::new(nxt, nxt, nxt, nxt, false, Vec::with_capacity(0))
        };
      states.push(state);
    }
    let mut c = 0;
    for state in states.iter().rev().take(spiral_length - 1) {
      if state.empty == nfs || state.red == nfs || state.black == nfs || state.bad == nfs {
        break;
      }
      c += 1;
    }
    let new_fs = fs - c;
    let new_nfs = nfs - c;
    if c > 0 {
      states.truncate(spiral_length - c);
      for state in &mut states {
        if state.empty == nfs {
          state.empty = new_nfs;
        }
        if state.red == nfs {
          state.red = new_nfs;
        }
        if state.black == nfs {
          state.black = new_nfs;
        }
        if state.bad == nfs {
          state.bad = new_nfs;
        }
      }
    }
    let rotated_moves = moves
      .iter()
      .map(|m| {
        let (x, y) = rotate(width, height, m.x as u32, m.y as u32, rotation);
        Move {
          x: x as i32 - center_x as i32,
          y: y as i32 - center_y as i32,
          p: m.p,
        }
      })
      .collect();
    states.push(DfaState::new(new_fs, new_fs, new_fs, new_fs, true, rotated_moves));
    states.push(DfaState::new(
      new_nfs,
      new_nfs,
      new_nfs,
      new_nfs,
      true,
      Vec::with_capacity(0),
    ));
    Ok(Dfa::new(states))
  }

  fn get_pattern_moves(width: u32, s: &str) -> Result<HashSet<(u32, u32)>, &'static str> {
    let mut pattern_moves = HashSet::new();
    for i in s.chars().enumerate().filter(|&(_, c)| c == '+').map(|(i, _)| i as u32) {
      pattern_moves.insert((i % width, i / width));
    }
    if pattern_moves.is_empty() {
      Err("Moves are not defined.")
    } else {
      Ok(pattern_moves)
    }
  }

  pub fn union(&self, patterns: &Patterns) -> Patterns {
    Patterns {
      dfa: self.dfa.product(&patterns.dfa),
      min_size: cmp::min(self.min_size, patterns.min_size),
    }
  }

  pub fn from_str(string: &str) -> Result<Patterns, &'static str> {
    let mut split = string.split('\n').map(str::trim).filter(|line| !line.is_empty());
    if let Some(header_str) = split.next() {
      // Read header from input string.
      let (width, height, priority) = Patterns::parse_header(header_str)?;
      // Read pattern from input string.
      let mut pattern_str = String::new();
      for _ in 0..height {
        let s = split.next().ok_or("Unexpected end of pattern.")?;
        assert_eq!(s.len() as u32, width);
        pattern_str.push_str(s);
      }
      // Get moves set from pattern.
      let pattern_moves = Patterns::get_pattern_moves(width, &pattern_str)?;
      // Read and verify moves from input string.
      let mut moves = Vec::with_capacity(pattern_moves.len());
      let mut moves_set = HashSet::new();
      let mut priorities_sum = 0f64;
      for s in split {
        let m = Patterns::parse_move(s)?;
        priorities_sum += m.p;
        moves_set.insert((m.x as u32, m.y as u32));
        moves.push(m);
      }
      if moves_set != pattern_moves {
        return Err("Moves list does not match moves in the pattern.");
      }
      for m in &mut moves {
        m.p = m.p * priority / priorities_sum;
      }
      // Build DFA for each rotation.
      let mut dfa = Dfa::empty();
      for rotation in 0..8 {
        let cur_dfa = Patterns::build_dfa(width, height, &moves, rotation, &pattern_str)?;
        dfa = dfa.product(&cur_dfa);
      }
      Ok(Patterns {
        dfa,
        min_size: cmp::min(width, height),
      })
    } else {
      Err("Empty pattern.")
    }
  }

  fn from_strings(strings: &[(String, String)]) -> Patterns {
    let len = strings.len();
    if let [(ref name, ref pattern_str)] = *strings {
      match Patterns::from_str(pattern_str) {
        Ok(patterns) => patterns,
        Err(e) => {
          error!(target: PATTERNS_STR, "Failed to parse pattern {}: {}", name, e);
          Patterns::empty()
        }
      }
    } else {
      let split_idx = len / 2;
      let left = &strings[0..split_idx];
      let right = &strings[split_idx..len];
      let (left_patterns, right_patterns) =
        rayon::join(|| Patterns::from_strings(left), || Patterns::from_strings(right));
      left_patterns.union(&right_patterns)
    }
  }

  pub fn from_tar(file: File) -> Patterns {
    let mut archive = Archive::new(file);
    let mut strings = Vec::new();
    let iter = archive
      .entries()
      .expect("Reading of tar archive is failed.")
      .map(|file| file.expect("Reading of file in tar archive is failed."));
    for mut file in iter.filter(|file| file.header().entry_type().is_file()) {
      let name = file
        .header()
        .path()
        .ok()
        .and_then(|path| path.to_str().map(ToOwned::to_owned))
        .unwrap_or_else(|| "<unknown>".to_owned());
      info!(target: PATTERNS_STR, "Loading pattern '{}'", name);
      let mut s = String::new();
      file.read_to_string(&mut s).ok();
      strings.push((name, s));
    }
    let patterns = Patterns::from_strings(&strings);
    info!(target: PATTERNS_STR, "DFA total size: {}.", patterns.dfa.states_count());
    patterns
  }

  pub fn find(&self, field: &Field, player: Player, first_match: bool) -> Vec<(Pos, f64)> {
    if self.dfa.is_empty() || field.width() < self.min_size - 2 || field.height() < self.min_size - 2 {
      return Vec::with_capacity(0);
    }
    let mut priorities_sum = 0f64;
    let mut matched = Vec::new();
    let left_border = (self.min_size as i32 - 1) / 2 - 1;
    let right_border = self.min_size as i32 / 2 - 1;
    let inv_color = player == Player::Black;
    for y in left_border..field.height() as i32 - right_border {
      for x in left_border..field.width() as i32 - right_border {
        let moves = self.dfa.run(
          &mut Spiral::new().map(|(shift_x, shift_y)| {
            let cur_x = x + shift_x;
            let cur_y = y + shift_y;
            if cur_x >= 0 && cur_x < field.width() as i32 && cur_y >= 0 && cur_y < field.height() as i32 {
              let pos = field.to_pos(cur_x as u32, cur_y as u32);
              field.cell(pos)
            } else {
              Cell::new(true)
            }
          }),
          inv_color,
          first_match,
        );
        for m in moves {
          let move_x = (x as i32 + m.x) as u32;
          let move_y = (y as i32 + m.y) as u32;
          priorities_sum += m.p;
          matched.push((field.to_pos(move_x, move_y), m.p));
        }
      }
    }
    for &mut (_, ref mut p) in &mut matched {
      *p /= priorities_sum;
    }
    info!(
      target: PATTERNS_STR,
      "Found moves: {:?}.",
      matched
        .iter()
        .map(|&(pos, p)| (field.to_x(pos), field.to_y(pos), p))
        .collect::<Vec<(u32, u32, f64)>>()
    );
    matched
  }

  pub fn find_sorted(&self, field: &Field, player: Player, first_match: bool) -> Vec<(Pos, f64)> {
    let moves = self.find(field, player, first_match);
    let mut map = HashMap::with_capacity(moves.len());
    for (pos, p) in moves {
      let sum_p = map.get(&pos).cloned().unwrap_or(0f64) + p;
      map.insert(pos, sum_p);
    }
    let mut result = map.into_iter().collect::<Vec<(Pos, f64)>>();
    result.sort_by(|&(_, a), &(_, b)| b.partial_cmp(&a).expect("Cann't compare f64 types."));
    info!(
      target: PATTERNS_STR,
      "Found sorted moves: {:?}.",
      result
        .iter()
        .map(|&(pos, p)| (field.to_x(pos), field.to_y(pos), p))
        .collect::<Vec<(u32, u32, f64)>>()
    );
    result
  }

  pub fn find_foreground(&self, field: &Field, player: Player, first_match: bool) -> Option<Pos> {
    self
      .find_sorted(field, player, first_match)
      .first()
      .map(|&(pos, _)| pos)
  }

  pub fn find_rand<T: Rng>(&self, field: &Field, player: Player, first_match: bool, rng: &mut T) -> Option<Pos> {
    let moves = self.find_sorted(field, player, first_match);
    if moves.is_empty() {
      return None;
    }
    let rand = rng.gen();
    let mut sum = 0f64;
    let mut idx = 0;
    while sum < rand && idx < moves.len() {
      sum += moves[idx].1;
      idx += 1;
    }
    Some(moves[idx - 1].0)
  }
}
<commit_msg>Don't pass pattern name to parsing function.<commit_after>use crate::rotate::*;
use crate::spiral::Spiral;
use oppai_field::cell::Cell;
use oppai_field::field::{Field, Pos};
use oppai_field::player::Player;
use oppai_patterns::dfa::{Dfa, DfaState};
use rand::Rng;
use std::{
  cmp,
  collections::{HashMap, HashSet},
  fs::File,
  io::Read,
  str::FromStr,
};
use tar::Archive;

const PATTERNS_STR: &str = "patterns";

#[derive(Clone, Debug)]
struct Move {
  x: i32,
  y: i32,
  p: f64, // priority
}

pub struct Patterns {
  min_size: u32,
  dfa: Dfa<Move>,
}

impl Patterns {
  pub fn empty() -> Patterns {
    Patterns {
      min_size: u32::max_value(),
      dfa: Dfa::empty(),
    }
  }

  fn parse_header(string: &str) -> Result<(u32, u32, f64), &'static str> {
    let mut split = string.split_whitespace().fuse();
    let width = split
      .next()
      .and_then(|s| u32::from_str(s).ok())
      .ok_or("Expected width of type u32.")?;
    if width < 2 {
      return Err("Minimum allowed width is 2");
    }
    let height = split
      .next()
      .and_then(|s| u32::from_str(s).ok())
      .ok_or("Expected height of type u32.")?;
    if height < 2 {
      return Err("Minimum allowed height is 2");
    }
    let priority = split
      .next()
      .and_then(|s| f64::from_str(s).ok())
      .ok_or("Expected priority of type f64.")?;
    Ok((width, height, priority))
  }

  fn parse_move(string: &str) -> Result<Move, &'static str> {
    let mut split = string.split(' ').fuse();
    let x = split
      .next()
      .and_then(|s| u32::from_str(s).ok())
      .ok_or("Expected x coordinate of type u32.")?;
    let y = split
      .next()
      .and_then(|s| u32::from_str(s).ok())
      .ok_or("Expected y coordinate of type u32.")?;
    let p = split
      .next()
      .and_then(|s| f64::from_str(s).ok())
      .ok_or("Expected probability of type f64.")?;
    Ok(Move {
      x: x as i32,
      y: y as i32,
      p,
    })
  }

  fn covering_spiral_length(side_of_square: u32) -> u32 {
    let x = side_of_square / 2 + 1;
    let y = (1 - side_of_square % 2) * side_of_square * 2;
    (8 * x - 13) * x + 6 - y
  }

  fn build_dfa(width: u32, height: u32, moves: &[Move], rotation: u32, s: &str) -> Result<Dfa<Move>, &'static str> {
    let (rotated_width, rotated_height) = rotate_sizes(width, height, rotation);
    let center_x = (rotated_width - 1) / 2;
    let center_y = (rotated_height - 1) / 2;
    let spiral_length = Patterns::covering_spiral_length(cmp::max(width, height)) as usize;
    let mut states = Vec::with_capacity(spiral_length + 2);
    let fs = spiral_length; // "Found" state.
    let nfs = spiral_length + 1; // "Not found" state.
    let s_bytes = s.as_bytes();
    for (i, (shift_x, shift_y)) in Spiral::new().take(spiral_length).enumerate() {
      let nxt = i + 1;
      let rotated_x = center_x as i32 + shift_x;
      let rotated_y = center_y as i32 + shift_y;
      let state =
        if rotated_x >= 0 && rotated_x < rotated_width as i32 && rotated_y >= 0 && rotated_y < rotated_height as i32 {
          let (x, y) = rotate_back(
            rotated_width,
            rotated_height,
            rotated_x as u32,
            rotated_y as u32,
            rotation,
          );
          let pos = y * width + x;
          match s_bytes[pos as usize] as char {
            '.' | '+' => DfaState::new(nxt, nfs, nfs, nfs, false, Vec::with_capacity(0)),
            '?' => DfaState::new(nxt, nxt, nxt, nfs, false, Vec::with_capacity(0)),
            '*' => DfaState::new(nxt, nxt, nxt, nxt, false, Vec::with_capacity(0)),
            'X' => DfaState::new(nfs, nxt, nfs, nfs, false, Vec::with_capacity(0)),
            'O' => DfaState::new(nfs, nfs, nxt, nfs, false, Vec::with_capacity(0)),
            'x' => DfaState::new(nxt, nxt, nfs, nfs, false, Vec::with_capacity(0)),
            'o' => DfaState::new(nxt, nfs, nxt, nfs, false, Vec::with_capacity(0)),
            '#' => DfaState::new(nfs, nfs, nfs, nxt, false, Vec::with_capacity(0)),
            _ => return Err("Invalid character in the pattern."),
          }
        } else {
          DfaState::new(nxt, nxt, nxt, nxt, false, Vec::with_capacity(0))
        };
      states.push(state);
    }
    let mut c = 0;
    for state in states.iter().rev().take(spiral_length - 1) {
      if state.empty == nfs || state.red == nfs || state.black == nfs || state.bad == nfs {
        break;
      }
      c += 1;
    }
    let new_fs = fs - c;
    let new_nfs = nfs - c;
    if c > 0 {
      states.truncate(spiral_length - c);
      for state in &mut states {
        if state.empty == nfs {
          state.empty = new_nfs;
        }
        if state.red == nfs {
          state.red = new_nfs;
        }
        if state.black == nfs {
          state.black = new_nfs;
        }
        if state.bad == nfs {
          state.bad = new_nfs;
        }
      }
    }
    let rotated_moves = moves
      .iter()
      .map(|m| {
        let (x, y) = rotate(width, height, m.x as u32, m.y as u32, rotation);
        Move {
          x: x as i32 - center_x as i32,
          y: y as i32 - center_y as i32,
          p: m.p,
        }
      })
      .collect();
    states.push(DfaState::new(new_fs, new_fs, new_fs, new_fs, true, rotated_moves));
    states.push(DfaState::new(
      new_nfs,
      new_nfs,
      new_nfs,
      new_nfs,
      true,
      Vec::with_capacity(0),
    ));
    Ok(Dfa::new(states))
  }

  fn get_pattern_moves(width: u32, s: &str) -> Result<HashSet<(u32, u32)>, &'static str> {
    let mut pattern_moves = HashSet::new();
    for i in s.chars().enumerate().filter(|&(_, c)| c == '+').map(|(i, _)| i as u32) {
      pattern_moves.insert((i % width, i / width));
    }
    if pattern_moves.is_empty() {
      Err("Moves are not defined.")
    } else {
      Ok(pattern_moves)
    }
  }

  pub fn union(&self, patterns: &Patterns) -> Patterns {
    Patterns {
      dfa: self.dfa.product(&patterns.dfa),
      min_size: cmp::min(self.min_size, patterns.min_size),
    }
  }

  pub fn from_str(string: &str) -> Result<Patterns, &'static str> {
    let mut split = string.split('\n').map(str::trim).filter(|line| !line.is_empty());
    if let Some(header_str) = split.next() {
      // Read header from input string.
      let (width, height, priority) = Patterns::parse_header(header_str)?;
      // Read pattern from input string.
      let mut pattern_str = String::new();
      for _ in 0..height {
        let s = split.next().ok_or("Unexpected end of pattern.")?;
        assert_eq!(s.len() as u32, width);
        pattern_str.push_str(s);
      }
      // Get moves set from pattern.
      let pattern_moves = Patterns::get_pattern_moves(width, &pattern_str)?;
      // Read and verify moves from input string.
      let mut moves = Vec::with_capacity(pattern_moves.len());
      let mut moves_set = HashSet::new();
      let mut priorities_sum = 0f64;
      for s in split {
        let m = Patterns::parse_move(s)?;
        priorities_sum += m.p;
        moves_set.insert((m.x as u32, m.y as u32));
        moves.push(m);
      }
      if moves_set != pattern_moves {
        return Err("Moves list does not match moves in the pattern.");
      }
      for m in &mut moves {
        m.p = m.p * priority / priorities_sum;
      }
      // Build DFA for each rotation.
      let mut dfa = Dfa::empty();
      for rotation in 0..8 {
        let cur_dfa = Patterns::build_dfa(width, height, &moves, rotation, &pattern_str)?;
        dfa = dfa.product(&cur_dfa);
      }
      Ok(Patterns {
        dfa,
        min_size: cmp::min(width, height),
      })
    } else {
      Err("Empty pattern.")
    }
  }

  fn from_strings(strings: &[String]) -> Patterns {
    let len = strings.len();
    if let [ref pattern_str] = *strings {
      match Patterns::from_str(pattern_str) {
        Ok(patterns) => patterns,
        Err(e) => {
          error!(target: PATTERNS_STR, "Failed to parse pattern: {}\n{}", e, pattern_str);
          Patterns::empty()
        }
      }
    } else {
      let split_idx = len / 2;
      let left = &strings[0..split_idx];
      let right = &strings[split_idx..len];
      let (left_patterns, right_patterns) =
        rayon::join(|| Patterns::from_strings(left), || Patterns::from_strings(right));
      left_patterns.union(&right_patterns)
    }
  }

  pub fn from_tar(file: File) -> Patterns {
    let mut archive = Archive::new(file);
    let mut strings = Vec::new();
    let iter = archive
      .entries()
      .expect("Reading of tar archive is failed.")
      .map(|file| file.expect("Reading of file in tar archive is failed."));
    for mut file in iter.filter(|file| file.header().entry_type().is_file()) {
      let name = file
        .header()
        .path()
        .ok()
        .and_then(|path| path.to_str().map(ToOwned::to_owned))
        .unwrap_or_else(|| "<unknown>".to_owned());
      info!(target: PATTERNS_STR, "Loading pattern '{}'", name);
      let mut s = String::new();
      file.read_to_string(&mut s).ok();
      strings.push(s);
    }
    let patterns = Patterns::from_strings(&strings);
    info!(target: PATTERNS_STR, "DFA total size: {}.", patterns.dfa.states_count());
    patterns
  }

  pub fn find(&self, field: &Field, player: Player, first_match: bool) -> Vec<(Pos, f64)> {
    if self.dfa.is_empty() || field.width() < self.min_size - 2 || field.height() < self.min_size - 2 {
      return Vec::with_capacity(0);
    }
    let mut priorities_sum = 0f64;
    let mut matched = Vec::new();
    let left_border = (self.min_size as i32 - 1) / 2 - 1;
    let right_border = self.min_size as i32 / 2 - 1;
    let inv_color = player == Player::Black;
    for y in left_border..field.height() as i32 - right_border {
      for x in left_border..field.width() as i32 - right_border {
        let moves = self.dfa.run(
          &mut Spiral::new().map(|(shift_x, shift_y)| {
            let cur_x = x + shift_x;
            let cur_y = y + shift_y;
            if cur_x >= 0 && cur_x < field.width() as i32 && cur_y >= 0 && cur_y < field.height() as i32 {
              let pos = field.to_pos(cur_x as u32, cur_y as u32);
              field.cell(pos)
            } else {
              Cell::new(true)
            }
          }),
          inv_color,
          first_match,
        );
        for m in moves {
          let move_x = (x as i32 + m.x) as u32;
          let move_y = (y as i32 + m.y) as u32;
          priorities_sum += m.p;
          matched.push((field.to_pos(move_x, move_y), m.p));
        }
      }
    }
    for &mut (_, ref mut p) in &mut matched {
      *p /= priorities_sum;
    }
    info!(
      target: PATTERNS_STR,
      "Found moves: {:?}.",
      matched
        .iter()
        .map(|&(pos, p)| (field.to_x(pos), field.to_y(pos), p))
        .collect::<Vec<(u32, u32, f64)>>()
    );
    matched
  }

  pub fn find_sorted(&self, field: &Field, player: Player, first_match: bool) -> Vec<(Pos, f64)> {
    let moves = self.find(field, player, first_match);
    let mut map = HashMap::with_capacity(moves.len());
    for (pos, p) in moves {
      let sum_p = map.get(&pos).cloned().unwrap_or(0f64) + p;
      map.insert(pos, sum_p);
    }
    let mut result = map.into_iter().collect::<Vec<(Pos, f64)>>();
    result.sort_by(|&(_, a), &(_, b)| b.partial_cmp(&a).expect("Cann't compare f64 types."));
    info!(
      target: PATTERNS_STR,
      "Found sorted moves: {:?}.",
      result
        .iter()
        .map(|&(pos, p)| (field.to_x(pos), field.to_y(pos), p))
        .collect::<Vec<(u32, u32, f64)>>()
    );
    result
  }

  pub fn find_foreground(&self, field: &Field, player: Player, first_match: bool) -> Option<Pos> {
    self
      .find_sorted(field, player, first_match)
      .first()
      .map(|&(pos, _)| pos)
  }

  pub fn find_rand<T: Rng>(&self, field: &Field, player: Player, first_match: bool, rng: &mut T) -> Option<Pos> {
    let moves = self.find_sorted(field, player, first_match);
    if moves.is_empty() {
      return None;
    }
    let rand = rng.gen();
    let mut sum = 0f64;
    let mut idx = 0;
    while sum < rand && idx < moves.len() {
      sum += moves[idx].1;
      idx += 1;
    }
    Some(moves[idx - 1].0)
  }
}
<|endoftext|>